HyukjinKwon opened a new pull request, #59551: URL: https://github.com/apache/spark/pull/59551
### What changes were proposed in this pull request? The Arrow columnar path of `ColumnarArrowEvalPythonEvaluatorFactory`, added by SPARK-56350, keeps the pass-through `ColumnVector`s of each input batch in a queue until the Python result of that batch comes back, which can be many batches later. Its comment assumes that this is "safe because Arrow vectors are independently allocated per batch". A columnar `PartitionReader` does not promise that, and `ColumnarBatch` "is meant to be reused during the entire data loading process": a reader may refill the same vectors for every batch, or close them once its partition ends, as `DataSourceRDD` does as soon as a reader is exhausted. This PR: - copies the pass-through columns of each batch into vectors of the task, with Arrow's `VectorAppender`, which copies their buffers, and releases them once the output batch that holds them is no longer used: when the next output batch is produced, or at the end of the task; - keeps them in a `ConcurrentLinkedQueue` instead of an `ArrayDeque`: in pipelined mode (`spark.python.udf.pipelined.enabled`), the writer thread fills the queue while the task thread drains it; - only takes this path when all the columns of a batch are Arrow-backed, instead of only the first one, since the others cannot be copied this way. Other batches take the row queue path, as non-Arrow batches already do. ### Why are the changes needed? With the default configuration, a scalar pandas or Arrow UDF on a column of a source whose Arrow vectors are refilled or released fails with `IndexOutOfBoundsException: index: 0, length: 1 (expected: range(0, 0))`, or pairs the pass-through values of a later batch with the UDF results of an earlier one. It was reported by @FelixYBW and reproduced by @sarutak in #59424 with native data sources, whose reader imports every batch into the same vectors, and also without them, with `ArrowCachedBatchSerializer`: `spark.range(0, 2000000, 1, 1).cache()` with `spark.sql.inMemoryColumnarStorage.batchSize=1000` and the same UDF fails with the same exception. cc @viirya ### Does this PR introduce _any_ user-facing change? Yes, it fixes a bug of 4.2.0: Python UDFs on the columns of an Arrow-backed source that reuses or releases its vectors, such as the Arrow cache, no longer fail or return wrong results. ### How was this patch tested? A new test in `ArrowColumnarPythonUDFSuite`, with a new option of the test source `ArrowBackedDataSourceV2`, `reuseVectors`, under which its reader refills the same vectors for every batch and closes them when it is closed, as columnar readers may. Its existing reader allocates new vectors for every batch and keeps them all until it is closed. The test runs a UDF on a string column, so that the cast of the test UDF is removed and the batches take the Arrow columnar path, over 2 batches and over 2000 batches of 100 rows, and checks the pass-through columns of every row against its UDF result. Before this change, it fails with the `IndexOutOfBoundsException` above in both cases. `ArrowColumnarPythonUDFSuite` passes (7 tests), as does `scalastyle`. The copy is a copy of the buffers of the pass-through columns, much cheaper than the row conversion that this path avoids. I did not rerun the benchmark of SPARK-56350. ### 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]
