lisancao commented on code in PR #59683: URL: https://github.com/apache/spark/pull/59683#discussion_r4343314428
########## docs/sql-v2-data-sources.md: ########## @@ -466,6 +497,246 @@ 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 + +The sections above describe the interfaces one at a time. This section puts them together in +the order a connector author needs them, using the connectors Spark tests itself with: +[`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 `TableProvider` whose schema is two +`int` columns, `i` and `j`; `JavaSimpleBatchTable` is a `Table` that declares `BATCH_READ` and +returns that schema; `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`](api/java/org/apache/spark/sql/catalyst/InternalRow.html), 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`. 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 +full list and order are under [ScanBuilder](#scanbuilder)). The contract that matters most is +for filters: the method receives the candidate predicates, keeps the ones the connector will +apply, and returns the **residual**, the predicates Spark must still evaluate after the scan. +`JavaAdvancedDataSourceV2WithV2Filter` accepts only `i > <int literal>` and hands everything +else back (the test's `assertInstanceOf` checks and the mirror-image `unsupported` stream are +elided here): + +{% highlight java %} +@Override +public Predicate[] pushPredicates(Predicate[] predicates) { + Predicate[] supported = Arrays.stream(predicates).filter(f -> { + if (f.name().equals(">")) { + FieldReference column = (FieldReference) f.children()[0]; + Literal value = (Literal) f.children()[1]; + return column.describe().equals("i") && value.value() instanceof Integer; + } else { + return false; + } + }).toArray(Predicate[]::new); + Predicate[] unsupported = /* the complement of supported */; + this.predicates = supported; + return unsupported; +} + +@Override +public Predicate[] pushedPredicates() { + return predicates; +} + +@Override +public void pruneColumns(StructType requiredSchema) { + this.requiredSchema = requiredSchema; +} +{% endhighlight %} + +Spark evaluates only what `pushPredicates` returns (plus anything it could not translate), so a +predicate that is accepted here and then not applied by the reader is silently dropped from the +query. If the connector can only use a predicate to skip whole files or row groups, report it as +pushed and also return it, so Spark re-checks it row by row. `pushedPredicates()` is what the +query plan reports; the `Batch` in the same class uses the accepted predicate to plan fewer +partitions, and its reader fills only the columns in `requiredSchema`. + +The "advanced implementation with V2 Filter" test shows the contract from the outside by +inspecting the `Batch` Spark planned: `df.filter($"i" > 3)` leaves one pushed predicate on `i` +and reads both columns, `df.select($"i").filter($"i" > 6)` keeps the predicate and prunes the +schema to `i`, and `df.select($"j").filter($"j" < -10)` pushes nothing, because the connector +only accepts predicates on `i`, so Spark evaluates that filter itself. To see the same on a +connector of your own, override `Scan.description()`, which renders inside the `BatchScan` node +of `EXPLAIN` after the output columns; the default is the class name. + +### Step 3: Batch Writes + +A writable table declares `BATCH_WRITE` and implements `SupportsWrite`; it adds `TRUNCATE` and +`SupportsTruncate` on the builder if `SaveMode.Overwrite` should work (the full list is under +[Table](#table) and in +[`TableCapability`](api/java/org/apache/spark/sql/connector/catalog/TableCapability.html)). +Spark checks the capability set during analysis, before calling any builder, so a missing value +fails the query even when the mix-in is implemented. From `JavaSimpleWritableDataSource`: + +{% highlight java %} +@Override +public Set<TableCapability> capabilities() { + return EnumSet.of( + TableCapability.BATCH_READ, + TableCapability.BATCH_WRITE, + TableCapability.TRUNCATE); +} +{% endhighlight %} + +The write is split across the driver and the executors: + +``` +driver: WriteBuilder.build() -> Write.toBatch() -> BatchWrite.createBatchWriterFactory(info) +executor: DataWriterFactory.createWriter(partitionId, taskId) + DataWriter.write(row) per row + DataWriter.commit() -> WriterCommitMessage (or abort() on failure) +driver: BatchWrite.commit(messages[]) (or abort(messages[])) +``` + +Each `DataWriter` writes its task's output to a staging location, and `BatchWrite.commit` on the +driver makes all of it visible in one step. The staging location must be reachable from both Review Comment: Reworded and made the documentation more succinct per your suggestions. I do want to make sure enough is there for agents/llm to be able to investigate further, which I think we have a good balance for with the current state. -- 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]
