Repository navigation
Conversation
…erwrite There is no agreement between the producer and the consumer of a ColumnarBatch, so the type of its vectors is the protocol, as for closeIfFreeable (SPARK-50235). Add ReusedColumnVector, an internal marker with no methods: a vector carrying it may be released or overwritten by its producer after it is returned, so a consumer must not hold it past the next batch, and its closeIfFreeable is a no-op. A vector without it may be held, and the consumer releases it with closeIfFreeable. WritableColumnVector and ConstantColumnVector, whose closeIfFreeable is already a no-op, implement the marker: readers reuse them across batches and free them when they are closed. ReusedArrowColumnVector is an ArrowColumnVector carrying it, with a no-op closeIfFreeable, for producers that reuse Arrow vectors. The child vectors of a struct, array or map column share the lifetime of the column, so the marker is only checked on the top-level columns of a batch. The Arrow output processors of the Python runners load every batch into the same root, or close the slices of a batch once the next one is produced, so their columns are now ReusedArrowColumnVectors. This also makes closeIfFreeable on them safe: closing the shared ArrowColumnVectors of the root would clear their child columns for the next batches. Co-authored-by: Isaac <no-reply@databricks.com>
…n ColumnarToRow The consumer of a batch releases its vectors with closeIfFreeable once it is done with them (SPARK-50235), which the generated code of ColumnarToRowExec does in advanceBatch. Some consumers never did, so the fresh vectors of every batch stayed live until the end of the task: - the non-codegen ColumnarToRowEvaluatorFactory; - ArrowEvalPythonExec.doExecute, which flattens the output batches of the columnar evaluator for a row-based parent; - the row queue and row fallback paths of the columnar evaluator, which copy the rows and write the input batches to Python. Add ColumnarToRowIterator, which releases a batch once all its rows have been returned and before the next one is fetched, as advanceBatch does, and use it in both row conversions. The rows are projected into a buffer of their own, so the last row of a batch stays valid. The result columns of the Python runners are reused vectors, so releasing an output batch leaves them to the runner. The row paths release each input batch once the Python writer asks for the next one. The test data source no longer frees the returned batches when its reader is closed, which DataSourceRDD does as soon as the reader is exhausted, but when the task completes, and counts the batches its consumer releases before then. Co-authored-by: Isaac <no-reply@databricks.com>
…f Python UDFs The Arrow columnar path of ColumnarArrowEvalPythonEvaluatorFactory (SPARK-56350) keeps the pass-through columns of a batch until its result comes back from Python, many batches later. This breaks when the producer releases or overwrites the vectors it returned: IndexOutOfBoundsException, or pass-through values of a later batch paired with the results of an earlier one. - Take the path only when every column of the first batch is an ArrowColumnVector that is not a ReusedColumnVector; otherwise take the row queue path. The output batch carries the held columns to its consumer, which releases them with closeIfFreeable. If a later batch has reused vectors, its pass-through columns are copied into on-heap vectors and the batch is released once written. - The queue of pass-through columns is a ConcurrentLinkedQueue: in pipelined mode the writer thread adds to it while the task thread polls it. - ArrowCachedBatchToColumnarBatchIterator already loads every batch into a root of its own, but closed the previous root in next(). It now leaves the release to the consumer, which closes the root once all columns of its batch are closed, and closes the roots left when the task completes. This keeps the Arrow cache on the columnar path, rather than marking its vectors and making cache + UDF take the slower row queue path. The tests cover a reader that reuses its vectors (row queue path, correct), a reader with fresh vectors (columnar path, no copy, each batch released before the end of the task), the Arrow cache with 2000 batches, two stacked ArrowEvalPythonExecs, and pipelined mode. Co-authored-by: Isaac <no-reply@databricks.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.enabledby 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 addsReusedColumnVector, an internal marker for such vectors. This PR makes the Arrow columnar path follow it:ArrowColumnVectorwithout 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 withcloseIfFreeable. 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.ConcurrentLinkedQueue: in pipelined mode, the writer thread adds to it while the task thread polls it.ArrowCachedBatchToColumnarBatchIteratoralready 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 stackedArrowEvalPythonExecis marked by #59581.TODO once #59551 is merged: rebase, enable
spark.sql.execution.arrow.pythonUDF.columnarInput.enabledby 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:ArrowCachedBatchSerializerin batches of 1,000 takes the Arrow columnar path, and the pass-through columns match the UDF results;ArrowEvalPythonExecs, where the outer one takes the row queue path;ArrowColumnarPythonUDFPipelinedSuiteruns the suite withspark.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,PythonUDFSuiteandColumnarBatchSuite.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac (Claude Opus 5.5)