viirya opened a new pull request, #59581: URL: https://github.com/apache/spark/pull/59581
### What changes were proposed in this pull request? SPARK-50235 made the generated code of `ColumnarToRowExec` release each consumed batch with `closeIfFreeable`, and left the non-codegen path for a follow-up. This PR makes the remaining consumers of columnar batches release them: - Add `ColumnarToRowIterator`, which flattens batches into `UnsafeRow`s and releases each batch with `closeIfFreeable` once all its rows have been returned, before it fetches the next one, as `ColumnarToRowExec.advanceBatch` does. It is used by the non-codegen `ColumnarToRowEvaluatorFactory`, and by `ArrowEvalPythonExec.doExecute`, which flattens the output batches of the columnar evaluator for a row-based parent. - In `ColumnarArrowEvalPythonEvaluatorFactory`, the row queue and row fallback paths copy what they need out of each input batch, so they release the batch once the Python writer asks for the next one. Releasing the output batches of `ArrowEvalPythonExec` must not free the result columns, which the Python runner loads into the same Arrow vectors for every batch. So this PR also adds an internal marker interface, `ReusedColumnVector`, for vectors that their producer may release or overwrite after returning them. A consumer must not hold such a vector past the next batch, and its `closeIfFreeable` is a no-op; a vector without the marker may be held, and is released by its consumer. - `WritableColumnVector` and `ConstantColumnVector`, whose `closeIfFreeable` is already a no-op, implement it. - `ReusedArrowColumnVector` is an `ArrowColumnVector` carrying it, with a no-op `closeIfFreeable`. The Arrow output processors of the Python runners now return it. The marker is checked on the top-level columns of a batch: the child vectors of a struct, array or map column share its lifetime. The marker is also the protocol for the follow-up to SPARK-60294 (#59551): the Arrow columnar path of Python UDFs holds the input columns until their results come back from Python, which is only safe for vectors without the marker. That change is left to a separate PR. ### Why are the changes needed? - When batches whose release is left to their consumer, e.g. the fresh Arrow vectors of a DataSource V2 reader, are consumed by the non-codegen `ColumnarToRowExec` or by an Arrow Python UDF, their vectors stay allocated until the task completes, so the memory grows with the size of the partition. - Closing the Arrow vectors of the Python runner's output root, as `closeIfFreeable` on them would do, clears the child columns of struct results for all the following batches. ### Does this PR introduce _any_ user-facing change? No. A `ColumnVector` that its producer reuses but whose `closeIfFreeable` is not a no-op is now closed by these consumers as well, as it already is by the generated code of `ColumnarToRowExec` since SPARK-50235. ### How was this patch tested? Added a test to `ArrowColumnarPythonUDFSuite` that reads 500 batches of fresh Arrow vectors with whole-stage codegen enabled and disabled, and checks that the consumer releases every batch before the task completes, with at most one batch unreleased at a time. 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. Ran `ArrowColumnarPythonUDFSuite`, `ArrowCachedBatchSerializerSuite`, `ColumnarBatchSuite`, `ColumnVectorSuite`, `ConstantColumnVectorSuite`, `InMemoryColumnarQuerySuite`, `PythonUDFSuite`, `FileMetadataStructSuite`, `ColumnarRulesSuite` and `SparkPlanSuite`. ### 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]
