lisancao commented on code in PR #59683:
URL: https://github.com/apache/spark/pull/59683#discussion_r4342964848
##########
docs/sql-v2-data-sources.md:
##########
@@ -466,6 +497,507 @@ 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
+
+A minimal batch connector is a `TableProvider`, a `Table` that declares what
it can do, a
+`ScanBuilder` that plans partitions and creates readers, and a `WriteBuilder`
that implements the
+commit protocol. The examples read and write a CSV file in which every column
is a string, so the
+API surface stays visible.
+
+### Declaring the Table and Its Capabilities
+
+The `Table` names the dataset, exposes its columns, declares its capabilities,
and hands out
+builders through `SupportsRead` and `SupportsWrite`. It does no I/O itself.
+
+{% highlight java %}
+public class CsvTable implements SupportsRead, SupportsWrite {
+
+ private static final Set<TableCapability> CAPABILITIES = EnumSet.of(
+ TableCapability.BATCH_READ,
+ TableCapability.BATCH_WRITE,
+ TableCapability.TRUNCATE);
+
+ private final StructType schema;
+ private final CsvFile file;
+
+ public CsvTable(StructType schema, CsvFile file) {
+ this.schema = schema;
+ this.file = file;
+ }
+
+ @Override
+ public String name() {
+ return "csv:" + file.path();
+ }
+
+ @Override
+ public Column[] columns() {
+ return CatalogHelpers.asColumns(schema);
+ }
+
+ @Override
+ public Set<TableCapability> capabilities() {
+ return CAPABILITIES;
+ }
+
+ @Override
+ public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) {
+ return new CsvScanBuilder(schema, file);
+ }
+
+ @Override
+ public WriteBuilder newWriteBuilder(LogicalWriteInfo info) {
+ return new CsvWriteBuilder(info.schema(), file);
+ }
+}
+{% endhighlight %}
+
+Spark checks
[`TableCapability`](api/java/org/apache/spark/sql/connector/catalog/TableCapability.html)
+during analysis, before it calls `newScanBuilder` or `newWriteBuilder`, so the
set decides which
+operations are planned regardless of which mix-ins the table implements:
+
+| Capability | What it enables |
+|------------|-----------------|
+| `BATCH_READ` | Batch reads through `newScanBuilder` and `Scan.toBatch()` |
+| `MICRO_BATCH_READ` / `CONTINUOUS_READ` | Streaming reads through
`Scan.toMicroBatchStream` / `Scan.toContinuousStream` |
+| `BATCH_WRITE` | Append writes through `newWriteBuilder` and
`Write.toBatch()` |
+| `STREAMING_WRITE` | Streaming writes through `Write.toStreaming()` |
+| `TRUNCATE` | Overwriting the whole table (`SaveMode.Overwrite`, `INSERT
OVERWRITE` without a partition spec); the builder also implements
`SupportsTruncate` |
+| `OVERWRITE_BY_FILTER` | Overwriting the rows that match a predicate
(`DataFrameWriterV2.overwrite(condition)`, `INSERT OVERWRITE` with static
partition values); the builder also implements `SupportsOverwriteV2` |
+| `OVERWRITE_DYNAMIC` | Dynamic partition overwrite; the builder also
implements `SupportsDynamicOverwrite` |
+| `ACCEPT_ANY_SCHEMA` | Skips Spark's write-side schema resolution; the table
accepts input of any schema |
+| `AUTOMATIC_SCHEMA_EVOLUTION` | Lets Spark alter the table schema as part of
an operation when the input requires it |
+| `V1_BATCH_WRITE` | Append writes through a
[`V1Write`](api/java/org/apache/spark/sql/connector/write/V1Write.html)
wrapping a V1 `InsertableRelation`; combines with `TRUNCATE` and
`OVERWRITE_BY_FILTER` but not `OVERWRITE_DYNAMIC` |
+| `GENERATE_COLUMN_VALUES_ON_WRITE` | Spark computes missing generated-column
values on by-name writes and validates explicitly supplied ones (since 4.3.0) |
+| `SCAN_MERGING` | Spark may fuse two scans of the table that differ only in
pruned columns or pushed filters into one scan (since 4.3.0) |
+
+A table loaded through a catalog without `BATCH_READ` fails with
+`[UNSUPPORTED_FEATURE.TABLE_OPERATION] ... does not support batch scan`; a
table from a
+`TableProvider` without it is not treated as a V2 source at all, and Spark
falls back to the V1
+lookup for the class name. `SaveMode.Overwrite` needs `BATCH_WRITE` plus
`TRUNCATE` or
+`OVERWRITE_BY_FILTER`, and `SupportsTruncate` on the builder does not stand in
for the
+capability.
+
+### Implementing the Read Path
+
+`Batch.planInputPartitions()` runs once on the driver and returns one
+[`InputPartition`](api/java/org/apache/spark/sql/connector/read/InputPartition.html)
per task.
+Each partition is serialized to an executor, so it holds only what identifies
the split and
+nothing that cannot cross the wire (no connections, file handles, or
references back to the
+`Scan`):
+
+{% highlight java %}
+public class CsvInputPartition implements InputPartition {
+
+ private static final long serialVersionUID = 1L;
+
+ final long startByte;
+ final long endByte;
+
+ public CsvInputPartition(long startByte, long endByte) {
+ this.startByte = startByte;
+ this.endByte = endByte;
+ }
+}
+{% endhighlight %}
+
+Here the split is a byte range computed from the file length, so planning
never reads data.
+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, and
+`close()` releases the handle.
+
+{% highlight java %}
+@Override
+public boolean next() {
+ while (range.next()) {
+ List<String> cells = CsvFile.splitLine(range.get());
+ if (!matchesPushedFilters(cells, fullSchema, pushedFilters)) {
+ continue;
+ }
+ Object[] values = new Object[projection.length];
+ for (int i = 0; i < projection.length; i++) {
+ int src = projection[i];
+ values[i] = src < cells.size() ? UTF8String.fromString(cells.get(src)) :
null;
+ }
+ currentRow = new GenericInternalRow(values);
+ return true;
+ }
+ currentRow = null;
+ return false;
+}
+
+@Override
+public InternalRow get() {
+ return currentRow;
+}
+{% endhighlight %}
+
+A `PartitionReader<InternalRow>` produces
`org.apache.spark.sql.catalyst.InternalRow`, which
+holds each column in Spark's internal representation rather than the external
one a `Row`
+carries: `UTF8String` for strings, `Decimal` for decimals, an `int` of days
for dates, a `long`
+of microseconds for timestamps, `byte[]` for binary, and `ArrayData`,
`MapData`, and
+`InternalRow` for nested types. The accessor names on `InternalRow`
(`getUTF8String`,
+`getDecimal`, `getInt`, `getLong`, `getBinary`, `getArray`, `getMap`,
`getStruct`) are the
+reference for each type. `GenericInternalRow` takes `Object[]`, so a
`java.lang.String` in a
+string column compiles and fails on the executor with
+`ClassCastException: class java.lang.String cannot be cast to class
org.apache.spark.unsafe.types.UTF8String`.
+
+### Pushdown and How to Confirm It
+
+Pushdown is negotiated on the `ScanBuilder` through the mix-ins described under
+[ScanBuilder](#scanbuilder). The builder records what it accepted and carries
it into the
+`Scan` and the reader factory, so the reader applies exactly the filters and
the projection the
+builder claimed:
+
+{% highlight java %}
+public class CsvScanBuilder
+ implements ScanBuilder, Scan, Batch, SupportsPushDownFilters,
SupportsPushDownRequiredColumns {
+
+ private final StructType fullSchema;
+ private final CsvFile file;
+
+ private StructType requiredSchema;
+ private Filter[] pushedFilters = new Filter[0];
+
+ public CsvScanBuilder(StructType fullSchema, CsvFile file) {
+ this.fullSchema = fullSchema;
+ this.file = file;
+ this.requiredSchema = fullSchema;
+ }
+
+ @Override
+ public Filter[] pushFilters(Filter[] filters) {
+ List<Filter> accepted = new ArrayList<>();
+ List<Filter> rejected = new ArrayList<>();
+ for (Filter f : filters) {
+ if (f instanceof EqualTo eq &&
fullSchema.getFieldIndex(eq.attribute()).isDefined()) {
+ accepted.add(f);
+ } else {
+ rejected.add(f);
+ }
+ }
+ this.pushedFilters = accepted.toArray(new Filter[0]);
+ return rejected.toArray(new Filter[0]);
+ }
+
+ @Override
+ public Filter[] pushedFilters() {
+ return pushedFilters;
+ }
+
+ @Override
+ public void pruneColumns(StructType requiredSchema) {
+ this.requiredSchema = requiredSchema;
+ }
+
+ @Override
+ public Scan build() {
+ return this;
+ }
+
+ @Override
+ public Batch toBatch() {
+ return this;
+ }
+
+ @Override
+ public StructType readSchema() {
+ return requiredSchema;
+ }
+
+ @Override
+ public String description() {
+ return "CsvScan(path=" + file.path()
+ + ", columns=" + String.join(",", requiredSchema.fieldNames())
+ + ", pushedFilters=" + pushedFilters.length + ")";
+ }
+
+ @Override
+ public InputPartition[] planInputPartitions() {
+ List<long[]> ranges = file.byteRanges();
+ InputPartition[] partitions = new InputPartition[ranges.size()];
+ for (int i = 0; i < ranges.size(); i++) {
+ partitions[i] = new CsvInputPartition(ranges.get(i)[0],
ranges.get(i)[1]);
+ }
+ return partitions;
+ }
+
+ @Override
+ public PartitionReaderFactory createReaderFactory() {
+ return new CsvReaderFactory(fullSchema, requiredSchema, pushedFilters,
file);
+ }
+}
+{% endhighlight %}
+
+`Scan.description()` renders inside the `BatchScan` node of the physical plan,
after the table
+name and the pruned output columns, which makes it the place to report what
the connector was
+actually given. For `df.select("name").filter(df.dept == "eng")` the plan
shows the pruned
+columns and the pushed filter at the scan, while the residual `isnotnull`
stays in Spark's own
+`Filter` node:
+
+```
+== Physical Plan ==
+*(1) Project [name#6]
++- *(1) Filter isnotnull(dept#7)
+ +- BatchScan csv:/tmp/dsv2lab/people.csv[name#6, dept#7]
CsvScan(path=/tmp/dsv2lab/people.csv, columns=name,dept, pushedFilters=1)
RuntimeFilters: []
+```
+
+The default `description()` is the implementation's class name.
+
+### Implementing the Write Path
+
+The `WriteBuilder` records the write mode and builds a `Write` whose
`toBatch()` returns the
+`BatchWrite`. Implementing `SupportsTruncate` on the builder, together with
the `TRUNCATE`
+capability on the table, is what makes `SaveMode.Overwrite` work:
+
+{% highlight java %}
+public class CsvWriteBuilder implements WriteBuilder, SupportsTruncate {
+
+ private final StructType schema;
+ private final CsvFile file;
+ private boolean truncate = false;
+
+ public CsvWriteBuilder(StructType schema, CsvFile file) {
+ this.schema = schema;
+ this.file = file;
+ }
+
+ @Override
+ public WriteBuilder truncate() {
+ this.truncate = true;
+ return this;
+ }
+
+ @Override
+ public Write build() {
+ return new Write() {
+ @Override
+ public BatchWrite toBatch() {
+ return new CsvBatchWrite(schema, file, truncate);
+ }
+ };
+ }
+}
+{% endhighlight %}
+
+Spark then runs the commit protocol:
+
+```
+driver: BatchWrite.createBatchWriterFactory(info)
+executor: DataWriter.write(row) per row
+ DataWriter.commit() -> WriterCommitMessage (or abort() on
failure)
+driver: BatchWrite.commit(messages[]) (or
abort(messages[]))
+```
+
+Each `DataWriter` streams its rows to its own staging location and returns a
+[`WriterCommitMessage`](api/java/org/apache/spark/sql/connector/write/WriterCommitMessage.html)
+that identifies what it wrote. The messages are collected on the driver, so a
message carries a
+reference (a path, an object key, a transaction id), not the rows:
+
+{% highlight java %}
+static final class CsvCommitMessage implements WriterCommitMessage {
+ private static final long serialVersionUID = 1L;
+ final String partFile;
+
+ CsvCommitMessage(String partFile) {
+ this.partFile = partFile;
+ }
+}
+
+static final class CsvDataWriter implements DataWriter<InternalRow> {
+ private final StructType schema;
+ private final Path partFile;
+ private final BufferedWriter out;
+
+ CsvDataWriter(StructType schema, int partitionId, long taskId) {
+ this.schema = schema;
+ try {
+ this.partFile = Files.createTempFile("csvv2-part-" + partitionId + "-" +
taskId + "-", ".csv");
Review Comment:
replaced with JavaSimpleWritableDataSource's approach: tasks stage under
<path>/_temporary/<queryId> on the table's own filesystem and the driver
renames from there on commit, and the text now says the staging location has to
be reachable from both sides.
Let me know if that works better.
--
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]
