HyukjinKwon commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4309307580
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala:
##########
@@ -110,10 +110,11 @@ case class ArrowEvalPythonExec(
// Columnar path: delegate to doExecuteColumnar, flatten to
// UnsafeRow. ColumnarBatchRow from rowIterator() is NOT
// UnsafeRow, and downstream operators (e.g., outer
- // EvalPythonExec) may cast to UnsafeRow.
+ // EvalPythonExec) may cast to UnsafeRow. Like ColumnarToRowExec,
+ // this releases each batch once its rows are consumed: the
+ // batches of the Arrow columnar path carry the input columns.
doExecuteColumnar().mapPartitionsInternal { batchIter =>
- val toUnsafe = UnsafeProjection.create(schema)
- batchIter.flatMap(_.rowIterator().asScala.map(toUnsafe))
+ new ColumnarToRowIterator(batchIter, UnsafeProjection.create(schema))
Review Comment:
**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 s
ource'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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -339,6 +344,41 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
}
}
+/**
+ * Hands the input batches to the Python writer. `enqueue` is called on each
+ * batch before the writer gets it, and returns whether the batch should be
+ * released with closeIfFreeable once it has been written: the writer is done
+ * with a batch when it asks for the next one or reaches the end.
+ */
+private[python] class WrittenBatchReleasingIterator(
+ input: Iterator[ColumnarBatch],
+ enqueue: ColumnarBatch => Boolean)
Review Comment:
**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.
--
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]
