viirya commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4362037317
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -232,10 +247,15 @@ private class ReusingArrowPartitionReader(start: Int,
end: Int, batchSize: Int)
/**
* Partition reader that produces Arrow-backed [[ColumnarBatch]]es.
*
- * Each batch contains up to [[ArrowBackedDataSourceV2.BATCH_SIZE]] rows with:
+ * Each batch contains up to `batchSize` rows with:
Review Comment:
Added in 09bae80321d.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala:
##########
@@ -44,12 +42,52 @@ class ColumnarToRowEvaluatorFactory(
inputs: Iterator[ColumnarBatch]*): Iterator[InternalRow] = {
assert(inputs.length == 1)
val toUnsafe = UnsafeProjection.create(childOutput, childOutput)
- inputs.head.flatMap { input =>
+ new ColumnarToRowIterator(inputs.head, toUnsafe, { batch =>
numInputBatches += 1
- numOutputRows += input.numRows()
- input.rowIterator().asScala.map(toUnsafe)
+ numOutputRows += batch.numRows()
+ })
+ }
+ }
+}
+
+/**
+ * Flattens columnar batches into `UnsafeRow`s. A batch is released with
closeIfFreeable once all
+ * its rows have been returned, before the next batch is fetched, as the
generated code of
+ * [[ColumnarToRowExec]] does in [[ColumnarToRowExec.advanceBatch]]. The rows
are projected into a
+ * buffer of their own, so the last row of a batch stays valid after the batch
is released.
+ *
+ * @param onNextBatch called on each batch once it is fetched
+ */
+private[sql] class ColumnarToRowIterator(
Review Comment:
Thanks, done in 09bae80321d. `ColumnarToRowIterator` is removed:
`ColumnarToRowEvaluatorFactory` uses `CompletionIterator`, and
`ArrowEvalPythonExec.doExecute` no longer flattens batches after your comment
on L120.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -246,14 +266,44 @@ private class ArrowBackedPartitionReader(start: Int, end:
Int, batchSize: Int)
// Single allocator for the entire partition reader lifetime.
// Arrow vectors from earlier batches may still be referenced by
// downstream operators (e.g., pass-through column queue), so we
- // must NOT close the allocator until the reader itself is closed.
+ // must NOT close the allocator until the task completes.
private val allocator: BufferAllocator =
ArrowUtils.rootAllocator.newChildAllocator(
s"arrow-test-reader-$start", 0, Long.MaxValue)
- // Track all batches so we can close them in close().
+ // Track all batches so we can close them when the task completes.
private val allBatches =
new java.util.ArrayList[ColumnarBatch]()
+ private var numReleased = 0
+ @volatile private var taskCompleted = false
+
+ TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+ taskCompleted = true
+ allBatches.synchronized {
+ allBatches.forEach(_.close())
+ allBatches.clear()
+ }
+ allocator.close()
+ }
+
+ /** The first column of a batch, which counts the release of the batch. */
+ private class ReleaseCountingArrowColumnVector(vector: ValueVector)
+ extends ArrowColumnVector(vector) {
+ private var closed = false
+
+ override def close(): Unit = {
Review Comment:
Done in 09bae80321d, and both flags are gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -114,6 +114,71 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("ColumnarToRow releases each batch once its rows are consumed") {
+ Seq(true, false).foreach { codegen =>
+ withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString) {
+ val numRows = 50000
+ val numBatches = 500
+ val df = readArrowSource(numRows = numRows, batchSize = numRows /
numBatches)
+ val plan = df.queryExecution.executedPlan
+ assert(collectNodes[ColumnarToRowExec](plan).nonEmpty)
+ assert(collectNodes[WholeStageCodegenExec](plan).nonEmpty == codegen)
+
+ ArrowBackedDataSourceV2.resetBatchCounters()
+ val rows = df.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ assert(rows.forall(row => row.getString(1) == s"row_${row.getInt(0)}"))
+ // A batch is released before the next one is read.
+ assert(ArrowBackedDataSourceV2.returnedBatches.get == numBatches)
+ assert(ArrowBackedDataSourceV2.releasedBatches.get == numBatches)
+ assert(ArrowBackedDataSourceV2.maxUnreleasedBatches.get == 1)
+ }
+ }
+ }
+
+ test("Arrow Python UDF over many batches releases each input batch") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true") {
Review Comment:
Removed in 09bae80321d.
--
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]
