HyukjinKwon commented on code in PR #58549:
URL: https://github.com/apache/spark/pull/58549#discussion_r4118165534
##########
python/pyspark/sql/udf.py:
##########
@@ -467,7 +498,10 @@ def _check_return_type(returnType: DataType, evalType:
int) -> None:
"return_type": str(returnType),
},
)
- elif evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF:
+ elif (
+ evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF
+ or evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF
Review Comment:
This also changes behavior for non-CHAR/VARCHAR types.
`SQL_GROUPED_AGG_PANDAS_ITER_UDF` and `SQL_GROUPED_AGG_ARROW_ITER_UDF`
previously had no branch in `_check_return_type`. Now they pick up the existing
StructType fail-fast (pandas) and `to_arrow_type` validation (both) at
definition time, for any return type. Is that intended? If so, maybe mention it
in the PR description or split it out, since it's outside the CHAR/VARCHAR
scope described there.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -189,12 +216,64 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+ val passThruClosed = new AtomicBoolean(false)
+
+ def closeCols(cols: Array[ColumnVector]): Unit = {
+ var i = 0
+ while (i < cols.length) {
+ cols(i).close()
+ i += 1
+ }
+ }
+
+ def drainPassThru(): Unit = {
+ passThruClosed.set(true)
+ var entry = passThruQueue.poll()
+ while (entry != null) {
+ closeCols(entry._1)
+ entry = passThruQueue.poll()
+ }
+ }
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ drainPassThru()
+ }
val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
+ // The input reader owns and may close its vectors as soon as the
Python runner consumes
+ // the input iterator. Create independent vector views whose buffers
remain valid until
+ // the corresponding output batch is closed.
+ val cols = Array.ofDim[ColumnVector](childOutput.length)
+ var allocated = 0
+ try {
+ var i = 0
+ while (i < childOutput.length) {
+ val vector =
batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector
+ val transferPair = vector.getTransferPair(ArrowUtils.rootAllocator)
Review Comment:
Who releases these transferred vectors? They are allocated against
`ArrowUtils.rootAllocator`, and after `passThruQueue.poll()` below the output
`ColumnarBatch` is the only thing holding them. I couldn't find anything
downstream that closes per-batch output: `ArrowEvalPythonExec.doExecute` just
`flatMap`s over `rowIterator()`, and `ColumnarToRowExec` doesn't close each
consumed batch either. The task-completion `drainPassThru()` only reaches
entries still in the queue.
The transfer keeps a reference to the input buffers, so it looks like every
pass-through column of every batch stays allocated in the root allocator after
the task finishes. Before this change they were plain references to
reader-owned vectors, so nothing leaked.
Could we close the previous batch's pass-through columns when the next
result batch is produced (and at task completion), or allocate them from a
task-scoped child allocator that is closed on completion? A test asserting that
`ArrowUtils.rootAllocator.getAllocatedMemory` is unchanged after a Path 1 query
would catch this.
--
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]
