[viirya]

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 UnsafeRows 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, closes its struct vectors, which drops their children, so loading the next batch with a struct result fails.

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.

Added a test that runs a pandas UDF over 500 batches of fresh Arrow vectors, on a column (the Arrow columnar path) and on an expression (the row fallback path), and checks the values of every row and that every input batch is released before the task completes. Added ArrowOutputProcessorSuite, which reads a struct column over several batches through ArrowOutputProcessorImpl and releases each output batch with closeIfFreeable. The new tests fail if the result columns are plain ArrowColumnVectors, or if ColumnarToRowIterator stops releasing batches.

Ran ArrowOutputProcessorSuite, ArrowCachedPythonUDFSuite, 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)

[HyukjinKwon]

Review summary

The release mechanics themselves look sound: ColumnarToRowIterator mirrors ColumnarToRowExec.advanceBatch (each batch is released only after its last row has been projected into the projection's own buffer, and before the next fetch), and on Paths 2 and 3 both Python writer implementations finish serializing a batch before asking for the next one, so releasing on the next hasNext/next is safe. The holding behavior of the Arrow columnar path (Path 1) for reused upstream vectors is unchanged by this PR and is declared out of scope.

Two findings remain. The new WrittenBatchReleasingIterator has a Boolean release predicate that neither caller ever sets to false, which can be simplified to an unconditional release. More importantly, none of the new Python UDF behaviors (releasing input and pass-through batches, and preserving the reused result vectors across batches) is covered by a multi-batch test, so a regression there could silently corrupt UDF results or retain memory. Separately, there is a design question about the default the new ReusedColumnVector marker applies to unmarked external vectors.

Findings

2 total: 0 P0, 0 P1, 1 P2, 1 P3.

Non-blocking (P2)

• No test exercises the new release and result-preservation behavior on the Python UDF paths — sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala:117 — see inline.

Nit (P3)

• WrittenBatchReleasingIterator's Boolean release predicate is never false — sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:355 — see inline.

Decision challenges

ReusedColumnVector treats every unmarked (including external) vector as safe to hold

ReusedColumnVector makes the vector's type the lifetime protocol, and its default is that any vector without the marker "may be held by the consumer past the next batch without a copy". Only Spark's own WritableColumnVector, ConstantColumnVector and ReusedArrowColumnVector carry it, and it lives in org.apache.spark.sql.execution.vectorized, which is excluded from unidoc. Meanwhile the public @Evolving ColumnVector javadoc still says vectors are "expected to be reused during the entire data loading process", and describes closeIfFreeable as a no-op only for writable and constant vectors.

Nothing branches on the marker in this PR, and today's Arrow columnar path already holds every ArrowColumnVector input, so the planned SPARK-60294 follow-up would not make third-party sources worse than they are now. But it would codify holding their reused vectors as correct: a DSv2 reader that refills one VectorSchemaRoot per batch and wraps it in the public ArrowColumnVector (or in its own subclass with a no-op closeIfFreeable, as the public doc suggests for reused vectors) would still be classified as holdable, and its only way to opt out would be to implement an undocumented internal type.

Is that default intended? Alternatives could be to also treat a vector whose closeIfFreeable is a no-op as reused, to provide a supported public way to declare reuse, or at least to document the hold/copy protocol on the public ColumnVector.closeIfFreeable.

[HyukjinKwon]

Nit (P3): Both call sites always return true here: Path 2's callback copies the rows and returns true, and Path 3 passes _ => true. Since the class is private[python] with no other callers, the if (enqueue(batch)) false arm is unreachable, and readers have to reason about a "keep this batch unreleased" mode that never occurs. Could this take a ColumnarBatch => Unit callback (or none, for Path 3) and always set written = batch, with the class doc updated to match?

Verification:

• Inspection: The iterator's API exposes no release predicate, and both call sites compile against the simplified signature.
• Behavior: Arrow Python UDF results stay correct on the simple-column and complex-expression row-queue paths.

[viirya]

Thanks. In this PR both callers release every batch, but #59582, which is stacked on this one, uses the false case: on the Arrow columnar path, a batch whose columns are held is not released once it is written, but by the consumer of the output batch. I kept the predicate here so that its signature does not change again there. Happy to simplify it here and add it back in #59582 if you prefer.

[HyukjinKwon]

Non-blocking (P2): The new release-counting fixture is only exercised by a query without a Python UDF, and the existing UDF tests in ArrowColumnarPythonUDFSuite read 10-100 rows with the default 1024-row batch size, so each partition has a single batch. As a result, none of the Python-side changes in this PR has a failure signal: if ArrowOutputProcessorImpl went back to plain ArrowColumnVector (re-introducing the struct-result clearing the description mentions), if WrittenBatchReleasingIterator released a batch before the writer drained it, or if this ColumnarToRowIterator stopped releasing output batches, the suite would stay green. Could you add multi-batch UDF tests (gated by shouldTestPandasUDFs) using the new batchSize option: a UDF over a plain column (Arrow columnar path, flattened here for a row parent), a UDF over a complex expression (row fallback path), and a struct-returning UDF consumed with whole-stage codegen on and off, asserting values plus the source's release counters?

Verification:

• Regression: Multi-batch Arrow Python UDF queries on the Arrow columnar path and the row fallback path return correct values, and every input batch is released before task completion with at most one unreleased at a time.
• Regression: A struct-returning Arrow Python UDF over multiple batches returns correct struct fields with whole-stage codegen enabled and disabled.

[viirya]

Thanks, added in 2b49870:

• Arrow Python UDF over many batches releases each input batch runs a pandas UDF over 500 batches of fresh Arrow vectors, on a column (the Arrow columnar path, flattened by doExecute for a row-based parent) and on an expression (the row fallback path). It checks the values of every row and that every input batch is released before the task completes. On the row fallback path at most one batch is unreleased at a time; on the Arrow columnar path the writer runs ahead of the results, so the batches it holds are bounded by how far ahead it is rather than by one.
• The test UDFs of IntegratedUDFTestUtils only return strings, so for struct results ArrowOutputProcessorSuite reads a struct column over several batches from an Arrow stream through ArrowOutputProcessorImpl, and releases each output batch with closeIfFreeable. It does not depend on whole-stage codegen, since the consumer only calls closeIfFreeable. With plain ArrowColumnVectors it fails: closing the struct vector of the root drops its children, and loading the next batch fails with should have as many children as in the schema: found 0 expected 2.

The new tests fail if the result columns are plain ArrowColumnVectors, or if ColumnarToRowIterator stops releasing batches.

[dongjoon-hyun]

Could you resolve the conflicts, @viirya ?

[dongjoon-hyun]

Thank you for the fix, @viirya. I reviewed the latest commit, 2b49870, and left 15 inline comments. Here is a summary, ordered by priority.

Correctness

1. ReusedColumnVector L29: the Arrow cache scan frees the vectors of a batch once the next one is read but leaves them unmarked, so the new contract lets a consumer hold them, and the follow-up to SPARK-60294 would bring back the #59551 failure.
2. ArrowEvalPythonExec L112: pre-existing since SPARK-60294, doExecute re-reads the columnar-input conf at execution time, so a plan built with it enabled fails with a ClassCastException once it is disabled.
3. BaseSliceArrowOutputProcessor L252: the new comment misses that getVectors also wraps the reader's own root, which produceBatch (L236-241, pre-existing) closes itself on a root-to-slice transition.
4. SliceRecordsArrowOutputProcessorImpl L323-334: pre-existing, the last chunk of a sliced Python batch returns the whole root, so the rows already returned as slices come out again.

Tests

5. ArrowOutputProcessorSuite L79: the change to BaseSliceArrowOutputProcessor.getVectors has no test, and the struct-returning UDF of the earlier test request is still open.
6. ArrowColumnarPythonUDFSuite L158: the name iteration does not assert that it takes the Arrow columnar path.
7. ArrowBackedPartitionReader L280: the cleanup listener registered from the constructor can run before the pipelined Python writer stops.
8. ArrowColumnarPythonUDFSuite L120: the new tests read 500 batches where 20 check the same invariants.

API and design

9. ArrowEvalPythonExec L120: with the default supportsRowBased, the planner would insert the standard ColumnarToRowExec above this node, and this private flatten could go.
10. ReusedColumnVector L36: nothing reads the marker yet, so it could land with the follow-up that uses it.

Docs and comments

11. PR description: the testing section does not mention ArrowOutputProcessorSuite or the new UDF test.
12. ArrowBackedPartitionReader L250: the column list omits data.

Cleanup

13. ColumnarToRowIterator L61: the existing flatMap with CompletionIterator gives the same release order without a new class.
14. ReleaseCountingArrowColumnVector L294: counting in closeIfFreeable makes closed and taskCompleted unnecessary.
15. ArrowColumnarPythonUDFSuite L143: the suite already enables columnar input.

Items 1 and 2 matter most. Item 2 is pre-existing, but it is in the branch this PR edits. Items 3 and 4 are pre-existing bugs in the output-slicing processors next to this change, so a separate JIRA may fit them better.

[dongjoon-hyun]

The Arrow cache scan frees the vectors of a batch once the next one is read, as L22 describes, but its vectors are not marked, so L28-30 let a consumer hold them past the next batch. ArrowCachedBatchToColumnarBatchIterator.next() closes previousRoot (ArrowCachedBatchSerializer.scala:1332-1337), while it wraps the vectors in plain ArrowColumnVectors (ArrowCachedBatchSerializer.scala:1356 and 1359). In #59551 you listed this iterator next to ArrowOutputProcessorImpl and wrote that it "should only close a root once none of its vectors is still in use, or mark its vectors as reused", but only ArrowOutputProcessorImpl is marked here. Once the follow-up holds unmarked inputs on the Arrow columnar path, the Arrow cache case from #59551 comes back: with spark.sql.cache.serializer=org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer and columnar input enabled, a pandas UDF over spark.range(0, 2000000, 1, 1).cache() takes that path, and next() for batch N+1 closes the root of batch N while its pass-through columns are still held. Could you wrap the cache vectors in ReusedArrowColumnVector in this PR too? vectorTypes (ArrowCachedBatchSerializer.scala:141), the isArrow check and the zero-copy re-cache path all still accept it. The iterator would then also need to close previousRoot once it is exhausted, since a consumer's closeIfFreeable would no longer free the last root.

[dongjoon-hyun]

Pre-existing since SPARK-60294, but in the branch this PR edits: L112 chooses between the columnar and row paths at execution time, from the live session conf. supportsColumnar (L104-105) reads conf.arrowPySparkUDFColumnarInputEnabled, and SparkPlan.conf is session.sessionState.conf (SparkPlan.scala:70-76), while the transitions were inserted with the value at planning time. A Dataset reuses its executedPlan across actions (Dataset.scala:2334-2339). So enable spark.sql.execution.arrow.pythonUDF.columnarInput.enabled, collect() a pandas UDF over an Arrow-backed DSv2 scan, disable the conf, and collect() the same Dataset again: the plan still has ArrowEvalPythonExec directly over the columnar-only BatchScanExec, but L112 is now false, so super.doExecute() calls child.execute().map(_.copy()) (EvalPythonExec.scala:73). BatchScanExec.doExecute (DataSourceV2ScanExecBase.scala:210-216) then hands the ColumnarBatches of its columnar DataSourceRDD over as rows, which should end in a ClassCastException. A .cache()d plan that is materialized after the conf changes takes the same route. Before SPARK-60294, this check was the plan-time child.supportsColumnar. Could the choice be fixed when the plan is built, for example with a flag set at planning time, or through the planner as in my comment on L120?

[dongjoon-hyun]

This comment covers only the slices, but getVectors also wraps the reader's own root, and produceBatch closes that root itself. getRoot returns the unsliced root for a batch that fits (L275-277) and for the last chunk of every batch in SliceRecordsArrowOutputProcessorImpl (L327-328). That is the case where the no-op closeIfFreeable from L255 matters, since closing the struct vector of the root removes its children for every later batch. Separately, and pre-existing since SPARK-51769: when the previous chunk was root and the next getRoot returns a slice, L236-241 close prevVectors and prevRoot, i.e. the reader's root that holds the batch being sliced. With spark.sql.execution.arrow.maxRecordsPerOutputBatch=3 and a mapInArrow partition whose Python output is two 10-row batches, the next root.slice(3, 3) of the second batch would throw IllegalArgumentException (valueCount 0). With maxBytesPerOutputBatch, a small batch followed by a large one would make getRoot return the cleared root with its old row count. ReusedArrowColumnVector cannot prevent this, because L238 calls close(). Could the comment say that these vectors are either the reader's root, into which the next batch is loaded, or a slice of it that is closed once the next one is produced? And could produceBatch close prevRoot and prevVectors only when prevRoot ne root, here or in a separate JIRA?

[dongjoon-hyun]

Pre-existing since SPARK-51769: the batch returned here repeats rows at the end of every Python batch that SliceRecordsArrowOutputProcessorImpl slices. Once remainingRows <= arrowMaxRecordsPerOutputBatch, its getRoot (L323-334) returns the whole root instead of the remaining rows, and L248 sets numRows to the full row count. With maxRecordsPerOutputBatch=3, a 10-row batch comes out as rows 0-2, 3-5 and 6-8, then 0-9 again: 19 rows instead of 10. test_different_output_length in test_arrow_map.py, run by MapInArrowWithOutputArrowBatchSlicingRecordsTests, would get 199 rows for 100, but it compares sets, so it passes. The loop in ArrowOutputProcessorSuite with new SliceRecordsArrowOutputProcessorImpl(reader, metrics(), 2) would show it as rows 0, 1, 0, 1, 2 for the first batch. Could getRoot return root.slice(currentRowIdx, remainingRows) once currentRowIdx > 0, here or in a separate JIRA?

[dongjoon-hyun]

Only ArrowOutputProcessorImpl is tested here, so the change to BaseSliceArrowOutputProcessor.getVectors (PythonArrowOutput.scala:255) has no test. The slicing tests in test_arrow_map.py go through mapInArrow, which never releases its output batches, so nothing would fail if PythonArrowOutput.scala:255 went back to ArrowColumnVector. That line matters on the Arrow columnar path when output slicing is enabled but no batch is split, the only case the passThruRows == numRows assertion (ColumnarArrowEvalPythonEvaluatorFactory.scala:223-225) allows: getRoot then returns the reader's root, which ColumnarToRowIterator in ArrowEvalPythonExec.doExecute releases. The struct-returning UDF of the earlier test request (#59581 (comment)) is also still open: TestScalarPandasUDF always returns strings, so the new many-batch UDF test cannot notice cleared struct children. Could this test also run the loop with new SliceRecordsArrowOutputProcessorImpl(reader, metrics(), rowsPerBatch) and new SliceBytesArrowOutputProcessorImpl(reader, metrics(), Int.MaxValue), i.e. with limits that never split a batch? The whole-stage codegen on/off part of that request makes no difference here: ArrowEvalPythonExec.supportsRowBased is true, so both modes go through ColumnarToRowIterator.

[dongjoon-hyun]

Minor: the PR description does not mention this suite or the new UDF test. "How was this patch tested?" still describes only ColumnarToRow releases each batch once its rows are consumed, and the "Ran ..." list does not include ArrowOutputProcessorSuite. The template asks "If tests were added, say they were added here.", and dev/merge_spark_pr.py copies the description into the commit message. Also, "The marker is checked on the top-level columns of a batch" describes a check that no code performs yet. Could you update the description?

[dongjoon-hyun]

Nit: this list still has only id, name and value, while next() also builds a data column of 100 random characters (L317-318, L404), which the new UDF test checks (ArrowColumnarPythonUDFSuite.scala:167). Could you add - data (STRING): 100 random characters?

[dongjoon-hyun]

Minor: the existing flatMap with Spark's CompletionIterator gives the same release order without a new class. Scala 2.13's Iterator.flatMap calls hasNext of the inner iterator, which runs the completion, before hasNext/next() of the outer one, so each batch is released after its last row and before the next one is fetched, as ColumnarToRowExec.advanceBatch does:

inputs.head.flatMap { input =>
  numInputBatches += 1
  numOutputRows += input.numRows()
  CompletionIterator[InternalRow, Iterator[InternalRow]](
    input.rowIterator().asScala.map(toUnsafe), input.closeIfFreeable())
}

ArrowEvalPythonExec.doExecute can do the same, and FlatMapGroupsInPandasWithStateExec already uses this form. That removes this class, its NoSuchElementException path and the onNextBatch default that only one caller uses. Could you use it instead?

[dongjoon-hyun]

Nit: counting in close() needs both closed and taskCompleted to skip the closes of the completion listener. Every consumer under test releases through closeIfFreeable, so counting there counts exactly what L83-84 describe, and both flags can go:

override def closeIfFreeable(): Unit = {
  allBatches.synchronized {
    numReleased += 1
    ArrowBackedDataSourceV2.releasedBatches.incrementAndGet()
  }
  super.closeIfFreeable()
}

A double release would then show up as releasedBatches > numBatches instead of being hidden by closed. Could you count it this way?

[dongjoon-hyun]

Nit: sparkConf (L38-40) already enables ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED for every test in this suite, and the other columnar tests set only ARROW_PYSPARK_EXECUTION_ENABLED. Could this be withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true")?

[viirya]

Thanks for raising this. Nothing in this PR branches on the marker; it matters for #59582, where the Arrow columnar path decides which input columns it may hold.

For #59582, we could also treat a vector as reused when its class overrides closeIfFreeable, which is already the public signal: the doc of ColumnVector.closeIfFreeable says that writable and constant vectors, which are reused across batches, override it to do nothing. It can be checked once per class, e.g. cls.getMethod("closeIfFreeable").getDeclaringClass != classOf[ColumnVector]. The check cannot tell whether the override is a no-op, but it errs on the safe side: a vector wrongly treated as reused only takes the slower row queue path. A third-party Arrow vector that follows the public doc would then not be held, and we would document the hold/copy protocol on ColumnVector.closeIfFreeable.

Taken further, the internal marker could go away. The in-tree vectors carrying it (WritableColumnVector, ConstantColumnVector and ReusedArrowColumnVector) are exactly those whose closeIfFreeable is overridden as a no-op, so the override alone could be the protocol, with a single public rule. The marker states the intent explicitly and is a plain type check, while the override check relies on reflection. Which would you prefer: keeping the marker and also honoring the override, or only the override?

A source that wraps a reused root in a plain ArrowColumnVector without overriding closeIfFreeable would still be held. Its vectors are already closed by the generated code of ColumnarToRowExec since SPARK-50235, though, so it is broken today regardless of this path.