HyukjinKwon opened a new pull request, #59934: URL: https://github.com/apache/spark/pull/59934
### What changes were proposed in this pull request? A follow-up to SPARK-60101 (#59424), which @zhouyuan asked for in https://github.com/apache/spark/pull/59424#issuecomment-6196064644. The design questions were listed in the reply https://github.com/apache/spark/pull/59424#issuecomment-6204557249. The change is a developer API to read a partition of a native scan as an Arrow C stream: ```java @DeveloperApi public interface SupportsArrowStreamRead extends PartitionReaderFactory { AutoCloseable openArrowStream(InputPartition partition, long streamAddress); } ``` The reader factory of the batch scans of native data sources implements it, with the stream that `read` of the library initializes. This PR is based on #59931 (version 1 of the interface), whose commit it includes until that PR is merged. The questions of the reply are settled as follows: - **Which object.** The engine gets the reader factory of the scan from its plan, `BatchScanExec.readerFactory`, and the partitions with it. The factory now only holds the location of the library, the name of the data source, the read schema and the state of the reader. - **Who owns the task context.** `openArrowStream` is called in the task that reads the partition. It returns a handle that the engine closes once it released the stream, which frees the `SparkDsTaskContext` that the library may use until then. Spark also closes it when the task completes, so the stream must be released before. The stream has to be zero-initialized, and if `openArrowStream` fails, the stream is released. - **What the engine checks.** Spark does not check the arrays of such a stream, while it checks those that it imports itself, because what to check depends on the importer. Arrow Java ignores the offsets of the arrays, for example, while Arrow C++ and arrow-rs do not. Spark still checks that the library initialized the stream, with non-NULL callbacks. The callbacks can be called from any thread, but never concurrently, as the header already says. `NativePartitionReader` opens its streams with the same code. Closing a `NativeLibrary.Task` is now synchronized, because an engine may close it on its own thread while the end of the task closes it. The docs describe the API under "Native Execution Engines". Micro-batch scans are left out: their partitions are read with the state of the stream reader of each micro-batch. ### Why are the changes needed? Without it, a query on a native data source under a native engine such as Gluten imports the data into Arrow Java. A `ColumnarToRowExec` then converts it to rows, which the engine converts back to its own columnar format: ``` read -> ArrowArrayStream -> Arrow Java -> ArrowColumnVector -> ColumnarToRow -> RowToVeloxColumnar ``` With the API, the engine imports the stream of the library directly. Library loading, the checksums of the packages and `spark_ds_init` stay in Spark, so the same library runs on Spark alone and under the engine. ### Does this PR introduce _any_ user-facing change? It adds a developer API for engines, `SupportsArrowStreamRead`, in 4.4.0. Nothing changes for queries. ### How was this patch tested? "open the partitions of a batch scan as Arrow C streams" in `NativeDataSourceSuite` acts as an engine: in each task it opens the stream of a partition, imports it on another thread, and checks the data and the task context that the library got. It also checks that a failing `read`, and a library that does not initialize the stream, raise `NATIVE_DATA_SOURCE_ERROR` and leave the stream released. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) This pull request and its description were written by Isaac. -- 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]
