viirya opened a new pull request, #59582: URL: https://github.com/apache/spark/pull/59582
This PR is stacked on #59581 (SPARK-60306): only the last commit, "Hold only unmarked vectors in the Arrow columnar path of Python UDFs", is new here. It will also enable `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` by default again once #59551 (SPARK-60294) is merged, so it stays a draft until then. ### What changes were proposed in this pull request? The Arrow columnar path of `ColumnarArrowEvalPythonEvaluatorFactory` (SPARK-56350) keeps the pass-through columns of each input batch until its result comes back from Python, many batches later. #59551 disables the columnar input by default because a producer may release or overwrite the vectors it returned before then. #59581 adds `ReusedColumnVector`, an internal marker for such vectors. This PR makes the Arrow columnar path follow it: - The Arrow columnar path is taken only when every column of the first batch is an `ArrowColumnVector` without the marker; otherwise the batches take the row queue path. The held columns are carried to the consumer of the output batch, which releases them with `closeIfFreeable`. If a later batch has marked vectors, its pass-through columns are copied into on-heap vectors, and the batch is released once it has been written to Python. - The queue of pass-through columns is now a `ConcurrentLinkedQueue`: in pipelined mode, the writer thread adds to it while the task thread polls it. - `ArrowCachedBatchToColumnarBatchIterator` already loads every cached batch into a root of its own, but closed the previous root when the next batch was read. It now leaves the release to the consumer: a root is closed once all the columns of its batch are closed, and the roots left unreleased are closed when the task completes. So the Arrow cache stays on the columnar path, rather than marking its vectors and making cache + UDF take the slower row queue path. The other Arrow producers that reuse a root (`PythonStreamingSourceRunner`, `ArrowRRunner`, `ExecuteExternalUDFExec`, `ArrowConverters`, the cogroup runner) only return rows, so they cannot reach this path. The output of a stacked `ArrowEvalPythonExec` is marked by #59581. TODO once #59551 is merged: rebase, enable `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` by default, and rewrite its documentation. ### Why are the changes needed? With the columnar input enabled, Python UDFs fail with `IndexOutOfBoundsException`, or pair the pass-through values of a later batch with the results of an earlier one, when the source of the batches reuses or releases its vectors, e.g. the Arrow cache. With the columnar input disabled, Arrow-backed sources pay the ColumnarToRow and ArrowWriter conversions that it avoids. ### Does this PR introduce _any_ user-facing change? Yes, once the configuration is enabled by default again: Python UDFs over Arrow-backed batches whose vectors are not marked as reused take the Arrow columnar path, and batches whose vectors are marked as reused take the row queue path. ### How was this patch tested? Added tests to `ArrowColumnarPythonUDFSuite`: - a reader that refills the same marked vectors for every batch takes the row queue path and returns correct results; - a reader with fresh vectors takes the Arrow columnar path without a copy, returns correct results, and its 500 batches are all released before the task completes; - a UDF over 2,000,000 rows cached with `ArrowCachedBatchSerializer` in batches of 1,000 takes the Arrow columnar path, and the pass-through columns match the UDF results; - two stacked `ArrowEvalPythonExec`s, where the outer one takes the row queue path; - `ArrowColumnarPythonUDFPipelinedSuite` runs the suite with `spark.python.udf.pipelined.enabled`. Reverting the change to the Arrow cache makes the cache test fail with `IndexOutOfBoundsException: index: 0, length: 1 (expected: range(0, 0))`. Ran `ArrowColumnarPythonUDFSuite`, `ArrowColumnarPythonUDFPipelinedSuite`, `ArrowCachedBatchSerializerSuite`, `InMemoryColumnarQuerySuite`, `PythonUDFSuite` and `ColumnarBatchSuite`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac (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]
