viirya commented on code in PR #59551:
URL: https://github.com/apache/spark/pull/59551#discussion_r4299582736
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -188,13 +194,26 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
outputTypes: Seq[DataType],
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
- val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ // The pass-through columns of a batch are needed until its result comes
back from Python,
+ // which can be many batches later, but the producer may reuse or
release the vectors of a
+ // batch once it is asked for the next one, as ColumnarBatch allows. So
they are copied, and
+ // released once the output batch that holds them is no longer used:
when the next output
+ // batch is produced, or at the end of the task. In pipelined mode,
add() runs in the writer
+ // thread and poll() in the task thread.
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"pass-through columns of Python UDFs in task
${context.taskAttemptId()}",
+ 0, Long.MaxValue)
+ val passThruQueue = new ConcurrentLinkedQueue[(Array[ColumnVector],
Int)]()
+ var outputPassThruCols = Array.empty[ColumnVector]
+ context.addTaskCompletionListener[Unit] { _ =>
+ outputPassThruCols.foreach(_.close())
+ passThruQueue.forEach(_._1.foreach(_.close()))
+ passThruQueue.clear()
+ allocator.close()
+ }
val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
+ passThruQueue.add((copyColumns(batch, allocator), batch.numRows()))
batch
}
Review Comment:
Now that the pass-through columns are copied, the input batch is no longer
needed once it has been written to Python, so this path could also release it,
as `ColumnarToRow` does with `closeIfFreeable` (SPARK-50235). Today nothing
does: a row-based parent calls `ArrowEvalPythonExec.doExecute`, since it is
also `supportsRowBased`, and that flattens the batches without
`closeIfFreeable`. So a source that allocates new vectors for every batch and
relies on the consumer to release them keeps all of them until the end of the
task. I checked it with such a test source over 50000 rows in batches of 100:
none of the 500 batches was released, both on 9c9a900 and on the merge base. So
this is not a regression of this PR. With the change below, all 500 were
released as they were consumed, the peak of the source's allocator dropped from
73.7 MB to 147 KB, and the new `reuseVectors` test still passes:
```scala
val bufferedIter = new Iterator[ColumnarBatch] {
private var previous: ColumnarBatch = null
override def hasNext: Boolean = {
if (previous != null) {
previous.closeIfFreeable()
previous = null
}
inputIter.hasNext
}
override def next(): ColumnarBatch = {
val batch = inputIter.next()
passThruQueue.add((copyColumns(batch, allocator), batch.numRows()))
previous = batch
batch
}
}
```
The previous batch is closed before the producer is asked for the next one,
so a reader that refills the same vectors is not affected. Paths 2 and 3 have
the same gap, so a separate JIRA for all three paths would also work.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -216,10 +235,52 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
assert(passThruRows == numRows,
s"Batch size mismatch: pass-through has " +
s"$passThruRows rows but UDF result has $numRows rows.")
+ // The previous output batch is no longer used once the next one is
asked for.
+ outputPassThruCols.foreach(_.close())
+ outputPassThruCols = passThruCols
Review Comment:
nit: when the parent consumes columnar batches, such as the Arrow cache, it
may already have closed these columns with `closeIfFreeable`. Closing them
again is harmless, but saying so here would save the next reader from checking.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -119,9 +122,13 @@ private class ArrowBackedScanBuilder(options:
CaseInsensitiveStringMap)
private val columnar =
options.getBoolean("columnar", true)
+ private val reuseVectors =
+ options.getBoolean("reuseVectors", false)
+ private val batchSize =
+ options.getInt("batchSize", ArrowBackedDataSourceV2.BATCH_SIZE)
Review Comment:
nit: `batchSize` only applies when `reuseVectors` is true.
`ArrowBackedPartitionReader` still uses `BATCH_SIZE`, so
`readArrowSource(batchSize = ...)` without `reuseVectors` ignores it without
saying so. Could the existing reader take it too, or could the option say that
it only applies to the reusing reader?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -188,13 +194,26 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
outputTypes: Seq[DataType],
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
- val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ // The pass-through columns of a batch are needed until its result comes
back from Python,
+ // which can be many batches later, but the producer may reuse or
release the vectors of a
+ // batch once it is asked for the next one, as ColumnarBatch allows. So
they are copied, and
+ // released once the output batch that holds them is no longer used:
when the next output
+ // batch is produced, or at the end of the task. In pipelined mode,
add() runs in the writer
+ // thread and poll() in the task thread.
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"pass-through columns of Python UDFs in task
${context.taskAttemptId()}",
+ 0, Long.MaxValue)
+ val passThruQueue = new ConcurrentLinkedQueue[(Array[ColumnVector],
Int)]()
+ var outputPassThruCols = Array.empty[ColumnVector]
+ context.addTaskCompletionListener[Unit] { _ =>
Review Comment:
The root cause here, `get()` returning right after `cancel(true)` without
waiting for the writer, is tracked in SPARK-60113, which makes the runner's
listener wait for the writer to exit. Once that lands, this path is covered
without its own lock. This listener is registered before `pyRunner.compute`, so
it runs after the runner's listener, by which time the writer has finished
`copyColumns` and `add`. The same holds for stacked `ArrowEvalPythonExec`s in
pipelined mode, where the outer writer thread drives the inner `resultIter.map`
and the assignment of `outputPassThruCols`.
@HyukjinKwon, could you add a comment here saying that this listener must be
registered before `compute` for that reason? Nothing else enforces the order.
--
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]
