package datasources;
import datasources.utils.DBClientWrapper;
import datasources.utils.DBTableReader;
import edb.common.Split;
import edb.common.UnknownTableException;
import org.apache.log4j.Logger;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.sources.v2.DataSourceOptions;
import org.apache.spark.sql.sources.v2.DataSourceV2;
import org.apache.spark.sql.sources.v2.ReadSupport;
import org.apache.spark.sql.sources.v2.reader.DataReader;
import org.apache.spark.sql.sources.v2.reader.DataReaderFactory;
import org.apache.spark.sql.sources.v2.reader.DataSourceReader;
import org.apache.spark.sql.sources.v2.reader.SupportsReportPartitioning;
import org.apache.spark.sql.sources.v2.reader.partitioning.ClusteredDistribution;
import org.apache.spark.sql.sources.v2.reader.partitioning.Distribution;
import org.apache.spark.sql.sources.v2.reader.partitioning.Partitioning;
import org.apache.spark.sql.types.StructType;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
/**
* This DataSource also supports parallel reads (i.e.: on multiple executors)
* from the ExampleDB.
*
* The interesting feature of this example is that it supports informing the
* Spark SQL optimizer whether the table is partitioned in the right way to avoid shuffles
* in certain queries. One example is grouping queries, where shuffles can be avoided if the
* table is clustered in such a way that each group (cluster) is fully contained in a
* single partition. Since ExampleDB only supports clustered indexes on single columns,
* in practice a shuffle can be avoided if the table is clustered on one of the grouping
* columns. (In ExampleDB clustered tables, splits always respect clustering.)
*
* It gets a table name from its configuration and infers a schema from
* that table. If a number of partitions is specified in properties, it is used. Otherwise,
* the table's default partition count (always 4 in ExampleDB) is used.
*/
public class PartitioningRowDataSource implements DataSourceV2, ReadSupport {
static Logger log = Logger.getLogger(PartitioningRowDataSource.class.getName());
/**
* Spark calls this to create the reader. Notice how it pulls the host and port
* on which ExampleDB is listening, as well as a table name, from the supplied options.
* @param options
* @return
*/
@Override
public DataSourceReader createReader(DataSourceOptions options) {
String host = options.get("host").orElse("localhost");
int port = options.getInt("port", -1);
String table = options.get("table").orElse("unknownTable"); // TODO: throw
int partitions = Integer.parseInt(options.get("partitions").orElse("0"));
return new Reader(host, port, table, partitions);
}
/**
* This is how Spark discovers the source table's schema by requesting it from ExmapleDB,
* and how it obtains the reader factories to be used by the executors to create readers.
* Notice that one factory is created for each partition.
*/
static class Reader implements SupportsReportPartitioning {
static Logger log = Logger.getLogger(Reader.class.getName());
public Reader(String host, int port, String table, int partitions) {
_host = host;
_port = port;
_table = table;
_requestedPartitions = partitions;
}
private String _host;
private int _port;
private String _table;
private int _requestedPartitions;
//
// dynamic properties inferred from database
//
private boolean _initialized = false;
private StructType _schema;
private String _clusteredColumn;
private List<Split> _splits;
private void initialize() {
if (!_initialized) {
log.info("initializing");
DBClientWrapper db = new DBClientWrapper(_host, _port);
db.connect();
try {
_schema = db.getSparkSchema(_table);
_clusteredColumn = db.getClusteredIndexColumn(_table);
if (_requestedPartitions == 0)
_splits = db.getSplits(_table);
else
_splits = db.getSplits(_table, _requestedPartitions);
} catch (UnknownTableException ute) {
throw new RuntimeException(ute);
} finally {
db.disconnect();
}
_initialized = true;
log.info("initialized");
}
}
@Override
public StructType readSchema() {
log.info("schema requested for table [" + _table + "]");
initialize();
return _schema;
}
@Override
public List<DataReaderFactory<Row>> createDataReaderFactories() {
log.info("reader factories requested for table [" + _table + "]");
initialize();
List<DataReaderFactory<Row>> factories = new ArrayList<>();
for (Split split : _splits) {
DataReaderFactory<Row> factory =
new SplitDataReaderFactory(_host, _port, _table, readSchema(), split);
factories.add(factory);
}
return factories;
}
@Override
public Partitioning outputPartitioning() {
log.info("output partitioning requested for table [" + _table + "]");
return new SingleClusteredColumnPartitioning(
_clusteredColumn, _splits.size());
}
}
static class SingleClusteredColumnPartitioning implements Partitioning {
static Logger log = Logger.getLogger(SingleClusteredColumnPartitioning.class.getName());
public SingleClusteredColumnPartitioning(String columnName, int partitions) {
_columnName = columnName;
_partitions = partitions;
}
@Override
public int numPartitions() {
log.info("asked for numPartitions");
return _partitions;
}
@Override
public boolean satisfy(Distribution distribution) {
//
// Since Spark may add other Distribution policies in the future, we can't assume
// it's always a ClusteredDistribution
//
if (distribution instanceof ClusteredDistribution) {
String[] clusteredCols = ((ClusteredDistribution) distribution).clusteredColumns;
StringBuilder logEntryBuilder = new StringBuilder();
logEntryBuilder.append("asked to satisfy ClusteredDistribution on columns ");
if (clusteredCols.length > 0) {
for (String col : clusteredCols) {
logEntryBuilder.append("[");
logEntryBuilder.append(col);
logEntryBuilder.append("] ");
}
}
log.info(logEntryBuilder.toString());
if (_columnName == null) {
log.info("no cluster column so does not satisfy");
return false;
} else {
boolean satisfies = Arrays.asList(clusteredCols).contains(_columnName);
log.info("based on cluster column: " + satisfies);
return satisfies;
}
}
log.info("asked to satisfy unknown distribution of type [" +
distribution.getClass().getCanonicalName() + "]");
return false;
}
private String _columnName;
private int _partitions;
}
/**
* This is used by each executor to read from ExampleDB. It uses the Split to know
* which data to read.
* Also note that when DBClientWrapper's getTableReader() method is called
* it reads ALL the data in its own Split eagerly.
*/
static class TaskDataReader implements DataReader<Row> {
static Logger log = Logger.getLogger(TaskDataReader.class.getName());
public TaskDataReader(String host, int port, String table,
StructType schema, Split split)
throws UnknownTableException {
log.info("Task reading from [" + host + ":" + port + "]" );
_db = new DBClientWrapper(host, port);
_db.connect();
_reader = _db.getTableReader(table, schema.fieldNames(), split);
}
private DBClientWrapper _db;
private DBTableReader _reader;
@Override
public boolean next() {
return _reader.next();
}
@Override
public Row get() {
return _reader.get();
}
@Override
public void close() throws IOException {
_db.disconnect();
}
}
/**
* Note that this has to be serializable. Each instance is sent to an executor,
* which uses it to create a reader for its own use.
*/
static class SplitDataReaderFactory implements DataReaderFactory<Row> {
static Logger log = Logger.getLogger(SplitDataReaderFactory.class.getName());
public SplitDataReaderFactory(String host, int port,
String table, StructType schema,
Split split) {
_host = host;
_port = port;
_table = table;
_schema = schema;
_split = split;
}
private String _host;
private int _port;
private String _table;
private StructType _schema;
private Split _split;
@Override
public DataReader<Row> createDataReader() {
log.info("Factory creating reader for [" + _host + ":" + _port + "]" );
try {
return new TaskDataReader(_host, _port, _table, _schema, _split);
} catch (UnknownTableException ute) {
throw new RuntimeException(ute);
}
}
}
}
