gengliangwang commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6480975990
Following up from the Rust SDK side (#59520), after checking this PR at
`1f40484`. Three suggestions for the C ABI that would remove code from every
SDK without adding much to the header. They change existing signatures or the
meaning of a flag, so they need to land before version 1 ships in 4.4.0, but
they can be follow-ups to this PR.
1. **Pass the schema, and whether the operation is batch or streaming, to
the executors.**
[`read`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L540-L542)
and
[`create_data_writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L591-L592)
get neither. Spark has the schema in both places:
[`NativePartitionReaderFactory`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L295-L304)
holds `readSchema`, and a write knows its input schema. So the Rust SDK
serializes an Arrow schema into its own state
([`ReaderState`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L458-L461),
[`WriterState`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sq
l/native/rust/src/write.rs#L59-L62)), and tags that state with an enum
([`Kind`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L464-L469))
to tell batch from streaming, since both share these two entry points. A C++
SDK would have to do the same, with Arrow IPC or an encoding of its own.
Could `read` and `create_data_writer` take the schema (the read schema
after pruning, and the input schema of the write) as a `struct ArrowSchema*`,
and `read` a batch/streaming flag? For writes, `epoch_id == -1` already tells
them apart. With the schema in `create_data_writer`,
[`write`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L600)
could also drop its `ArrowSchema` argument, or the header could say that it is
always that schema: the SDK imports it and compares its fields for every batch
today
([`import_batch`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/write.rs#L392-L414)).
2. **Split `SPARK_DS_TRUNCATE` into a flag for batch writes and one for
streaming writes.**
[`SPARK_DS_TRUNCATE`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L362)
covers both the save mode "overwrite" and the output mode "complete". A source
that supports only one of them either declares it and fails late, in
`create_writer`, as the Rust SDK does
([`effective`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L122-L131),
[the checks of
`create_writer`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L286-L305)),
or does not declare it and loses the one that it supports. With two flags,
such as `SPARK_DS_BATCH_TRUNCATE` and `SPARK_DS_STREAMING_TRUNCATE`, the table
reports `TRUNCATE` from the batch flag while Spark analyzes a query, and Spark
rejects the output mode "complete" from the streaming flag when the query
starts, as it does today.
Since this changes the meaning of an existing flag rather than appending one,
it cannot wait until after the release.
3. **Return the serialized state from the planning call, instead of a
separate serialize call.**
Spark always calls
[`partitions`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L463)
and then
[`serialize_reader`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L471)
([`plan`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L279-L291)),
[`stream_partitions`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L500-L501)
and then
[`serialize_stream_reader`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L509),
and [`create_
writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L428-L429)
and then
[`serialize_writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L554)
([`writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L127-L133)),
back to back. Because they are separate calls, the header needs rules for
their order, and the Rust SDK keeps a state machine and caches to return the
state of the plan that it already made
([`Planner`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L440-L450),
[`DriverWriter::serialized`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439
ac/sql/native/rust/src/write.rs#L156-L187)). Could `partitions` and
`stream_partitions` take a second sink, `set_state`, and `create_writer` a sink
for the state of the writer? That removes three functions from the table and
the ordering rules from the header.
I will update #59520 to whatever shape lands.
--
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]
