viirya commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4318668596
##########
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:
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.
##########
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:
Thanks, added in 2b49870c814:
- `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
`ArrowColumnVector`s 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 `ArrowColumnVector`s, or
if `ColumnarToRowIterator` stops releasing batches.
--
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]
