HyukjinKwon commented on code in PR #59551:
URL: https://github.com/apache/spark/pull/59551#discussion_r4299678631
##########
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:
Thanks, I added the comment in c6bcd015ef34: this listener must be
registered before `compute`, so that it runs after the listener of the runner.
Until SPARK-60113 makes that listener wait for the writer, the copies are
guarded against it, as @sarutak suggested: a copy made once the task completed
is released right away, and the allocator is closed by whichever of the two
ends last. Once SPARK-60113 lands, the guard can be removed.
--
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]
