szehon-ho commented on code in PR #59683: URL: https://github.com/apache/spark/pull/59683#discussion_r4364465205
########## docs/sql-v2-data-sources.md: ########## @@ -466,6 +512,219 @@ The same `Table` and `Scan` abstractions support streaming queries. A `Table` th A `Table` that declares `STREAMING_WRITE` supports streaming writes through `Write.toStreaming()`, which returns a `StreamingWrite`. +## Building a Connector + +Spark's own test connectors make a working example of each step: +[`JavaSimpleDataSourceV2`](https://github.com/apache/spark/blob/master/sql/core/src/test/java/test/org/apache/spark/sql/connector/JavaSimpleDataSourceV2.java) +for a minimal read, +[`JavaAdvancedDataSourceV2WithV2Filter`](https://github.com/apache/spark/blob/master/sql/core/src/test/java/test/org/apache/spark/sql/connector/JavaAdvancedDataSourceV2WithV2Filter.java) +for pushdown, and +[`JavaSimpleWritableDataSource`](https://github.com/apache/spark/blob/master/sql/core/src/test/java/test/org/apache/spark/sql/connector/JavaSimpleWritableDataSource.java) +for writes, exercised by +[`DataSourceV2Suite`](https://github.com/apache/spark/blob/master/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala). +The three share a few test helpers. `TestingV2Source` is a test trait over `TableProvider`: it +implements `inferSchema` (two `int` columns, `i` and `j`) and forwards +`getTable(schema, partitioning, properties)` to a one-argument `getTable(options)`, which is what +the excerpts override; a connector implements the three-argument form. `JavaSimpleBatchTable` is a +`Table` that declares `BATCH_READ` and returns that schema (through the deprecated `schema()`; a +new `Table` implements `columns()`); `JavaRangeInputPartition` carries a `start` and `end`; and +`JavaSimpleScanBuilder` is a `ScanBuilder` that is also its own `Scan` and `Batch`. + +### Step 1: A Minimal Batch Read + +The smallest connector is a `TableProvider` that returns a `Table`, a `Table` that declares +`BATCH_READ` and returns a `ScanBuilder`, and a `ScanBuilder` whose `Batch` plans partitions and +creates readers. `JavaSimpleDataSourceV2` is exactly that: + +{% highlight java %} +public class JavaSimpleDataSourceV2 implements TestingV2Source { + + static class MyScanBuilder extends JavaSimpleScanBuilder { + @Override + public InputPartition[] planInputPartitions() { + InputPartition[] partitions = new InputPartition[2]; + partitions[0] = new JavaRangeInputPartition(0, 5); + partitions[1] = new JavaRangeInputPartition(5, 10); + return partitions; + } + } + + @Override + public Table getTable(CaseInsensitiveStringMap options) { + return new JavaSimpleBatchTable() { + @Override + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) { + return new MyScanBuilder(); + } + }; + } +} +{% endhighlight %} + +`planInputPartitions()` runs once on the driver and returns one `InputPartition` per task. Each +partition is serialized to an executor, so it should hold only what identifies the split (here a +range of integers; for a file source, a path and a byte range). The `PartitionReaderFactory` is +serialized the same way and creates a `PartitionReader` on the executor. The reader is a cursor: +`next()` advances, `get()` returns the current row, `close()` releases any handle. The one in +`JavaSimpleReaderFactory` counts through the partition's range: + +{% highlight java %} +return new PartitionReader<InternalRow>() { + private int current = p.start - 1; + + @Override + public boolean next() { + current += 1; + return current < p.end; + } + + @Override + public InternalRow get() { + return new GenericInternalRow(new Object[] {current, -current}); + } + + @Override + public void close() { + } +}; +{% endhighlight %} + +A `PartitionReader<InternalRow>` produces `InternalRow`, which holds each column +in Spark's internal representation rather than the external one a `Row` carries: a boxed `int` +is fine for an `IntegerType` column as above, but a `StringType` column holds `UTF8String` (not +`java.lang.String`), `DecimalType` holds `Decimal`, `DateType` holds an `int` of days, +`TimestampType` a `long` of microseconds, `BinaryType` a `byte[]`, and nested types hold +`ArrayData`, `MapData` and `InternalRow`. `InternalRow` and `GenericInternalRow` live in Spark's +internal `catalyst` package and have no API docs, but they are the row types the connector API +uses. The accessor names on `InternalRow` (`getUTF8String`, +`getDecimal`, `getInt`, `getLong`, `getBinary`, `getArray`, `getMap`, `getStruct`) are the +reference for each type. `GenericInternalRow` takes `Object[]`, so a value of the wrong +representation compiles and fails on the executor with a `ClassCastException`. + +With this much, `spark.read.format("...JavaSimpleDataSourceV2").load()` returns ten rows, and +the `select` and `filter` in the "simplest implementation" test work too: Spark applies the +projection and the filter itself, above the scan. + +### Step 2: Projection and Filter Pushdown (Optional) + +A connector that can skip columns or rows adds the pushdown mix-ins to its `ScanBuilder`. The +list, the order Spark applies them in, and the residual-filter contract are under +[ScanBuilder](#scanbuilder). `JavaAdvancedDataSourceV2WithV2Filter` accepts only +`i > <int literal>` and returns everything else for Spark to evaluate. Unlike the test, which +asserts the shapes it expects, the version below checks them, because Spark can also push `i > j` +as a `>` predicate; it also matches on the public `NamedReference` rather than the test's +`FieldReference`. Replacing `predicates` on each call is right while `supportsIterativePushdown()` +returns `false` (the default); an iterative source accumulates them instead: + +{% highlight java %} +@Override +public Predicate[] pushPredicates(Predicate[] predicates) { + List<Predicate> supported = new ArrayList<>(); + List<Predicate> unsupported = new ArrayList<>(); + for (Predicate p : predicates) { + if (isIGreaterThanInt(p)) { + supported.add(p); + } else { + unsupported.add(p); + } + } + this.predicates = supported.toArray(new Predicate[0]); + return unsupported.toArray(new Predicate[0]); +} + +private static boolean isIGreaterThanInt(Predicate p) { + return p.name().equals(">") + && p.children()[0] instanceof NamedReference column + && column.fieldNames().length == 1 + && column.fieldNames()[0].equals("i") + && p.children()[1] instanceof Literal<?> value + && value.value() instanceof Integer; Review Comment: Could we restrict accepted thresholds or show partition planning that clamps to the original data range? This helper accepts `i > -2`, but the linked `AdvancedBatchWithV2Filter.planInputPartitions()` starts its first partition at `lowerBound + 1`, so it produces an extra `(-1, 1)` row even though the unfiltered source contains only `i = 0..9`. The test connector's planner is simplified for its test inputs; accepting every integer threshold in the walkthrough carries that incorrect behavior into the example. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
