sarutak commented on code in PR #59551:
URL: https://github.com/apache/spark/pull/59551#discussion_r4298230010
##########
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
new ColumnarBatch(passThruCols ++ resultCols, numRows)
}
}
+ /** Copies the pass-through columns of a batch into vectors allocated from
`allocator`. */
+ private def copyColumns(
+ batch: ColumnarBatch,
+ allocator: BufferAllocator): Array[ColumnVector] = {
+ val copies = ArrayBuffer.empty[ColumnVector]
+ try {
+ childOutput.indices.foreach { i =>
+ copies += copyColumn(batch.column(i), allocator)
+ }
+ copies.toArray
+ } catch {
+ case e: Throwable =>
+ copies.foreach(_.close())
+ throw e
+ }
+ }
+
+ private def copyColumn(column: ColumnVector, allocator: BufferAllocator):
ColumnVector = {
+ column match {
+ case arrow: ArrowColumnVector =>
+ val source = arrow.getValueVector
+ val copy = source.getField.createVector(allocator)
+ try {
+ copy.setInitialCapacity(source.getValueCount)
+ copy.allocateNew()
+ // Appends the values of the source to the empty copy, copying the
buffers.
+ source.accept(new VectorAppender(copy), null)
Review Comment:
`VectorAppender` does not support `ListViewVector`, which
`ArrowColumnVector` reads since SPARK-57929. With Arrow 19.0.0, I copied a
vector of each type that `ArrowColumnVector` accepts this way: all of them work
(including `ViewVarCharVector`, `ViewVarBinaryVector`, map, list and struct),
except `ListViewVector`, which fails with `UnsupportedOperationException:
VectorVisitor for ListViewVector is not supported.`
Since every pass-through column is copied, a query over a source that
returns a `ListViewVector` in any column now fails on this path, even when the
UDF does not read that column. Before this PR, the column was only referenced,
so it worked. Could the check at L157 also send such batches to the row queue
path (for example, by requiring that no column is a `ListViewVector`), and
could a test cover it?
##########
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:
In pipelined mode, this listener can run while the writer thread is still in
`copyColumns` or before its `passThruQueue.add` at L216. Then the copy is not
closed, and `allocator.close()` fails with "Memory was leaked".
The root cause is not in this PR but in the listener of
`BasePythonRunner.createPipelinedDataIn`, which runs before this one. It calls
`writerFuture.cancel(true)` and then `writerFuture.get()`, but `get()` on a
cancelled `FutureTask` throws `CancellationException` right away, without
waiting for the thread. A small program confirmed this: `get()` returned while
the task was still running, and the task finished 500 ms later. So the comment
at `PythonRunner.scala` L545 ("Wait for the writer to actually exit ...") does
not hold. This PR adds a resource that the writer thread allocates and that a
listener frees.
To check it, I forced this order on 9c9a900 without changing its non-test
code, in pipelined mode with the `reuseVectors` source:
1. The task thread reads 10 rows and then stops reading
(`executedPlan.execute().mapPartitions(it => Iterator(it.take(10).size))`,
since `limit` is pushed below the UDF).
2. The writer thread, in `copyColumns` of the next batch, blocks without
being interruptible in `getValueVector` of the last column. I did this with a
test-only subclass of `ArrowColumnVector`, to stand in for the copy of a wide
column. At that point, the copies of the other three columns are allocated.
3. The task ends, and this listener runs while the writer is blocked there.
(A listener that the test reader registers, and that runs after this one,
confirms the order.)
All 10 runs failed with `TaskCompletionListenerException: Memory was leaked
by query. Memory leaked: (3072)` from `Allocator(pass-through columns of Python
UDFs in task ...)`. Also, the memory allocated in `ArrowUtils.rootAllocator`
grew by 3072 bytes with each run, and was still allocated 2 seconds after the
job ended, so the copies were not freed. Without forcing the order, I did not
see it in 100 runs on the narrow columns of the test source. A task stops early
this way when it is killed or fails, or when an operator above it stops reading.
Could the copy and `add` be guarded against the listener here, for example
with a lock and a `closed` flag that both check, so that a copy made after the
listener ran is closed right away? Fixing the runner so that its listener
really waits for the writer (for example, with a latch released at the end of
`PipelinedWriterRunnable.run`) would also cover this, and the existing cases
that the comment there describes, but that seems better as a separate JIRA.
##########
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
new ColumnarBatch(passThruCols ++ resultCols, numRows)
}
}
+ /** Copies the pass-through columns of a batch into vectors allocated from
`allocator`. */
+ private def copyColumns(
+ batch: ColumnarBatch,
+ allocator: BufferAllocator): Array[ColumnVector] = {
+ val copies = ArrayBuffer.empty[ColumnVector]
+ try {
+ childOutput.indices.foreach { i =>
+ copies += copyColumn(batch.column(i), allocator)
+ }
+ copies.toArray
+ } catch {
+ case e: Throwable =>
+ copies.foreach(_.close())
+ throw e
+ }
+ }
+
+ private def copyColumn(column: ColumnVector, allocator: BufferAllocator):
ColumnVector = {
+ column match {
+ case arrow: ArrowColumnVector =>
+ val source = arrow.getValueVector
+ val copy = source.getField.createVector(allocator)
+ try {
+ copy.setInitialCapacity(source.getValueCount)
Review Comment:
The copy takes `source.getValueCount` rows, while the output batch has
`batch.numRows()` rows. Before this PR, a reader that set the number of rows of
the batch but not the value count of its vectors worked on this path, since
`ArrowColumnVector` does not read the value count and `writeArrowDirect` sets
it on the inputs of the UDF. Now the copy is shorter, and reading the missing
rows throws `IndexOutOfBoundsException` (I checked this with a value count of
0) or, within the allocated capacity, may read as nulls. The readers in Spark
set it, so this is minor, but a check that `source.getValueCount >=
batch.numRows()`, failing with a clear error, would make this case explicit.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -103,6 +109,36 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("Arrow-backed source that reuses its vectors: pass-through columns
match the results") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_reuse_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // Spark writes batches to Python before the results of the earlier ones
come back. With
+ // two batches, the reader is closed once it is exhausted, before their
results come back.
+ // With many, the reader refills its vectors with the next batches
before then.
+ Seq(200, 200000).foreach { numRows =>
+ // The test UDF casts its input to string: on a string column, the
cast is removed, so
+ // the input is a column, and the batches take the Arrow columnar path.
+ val result = readArrowSource(numRows = numRows, reuseVectors = true,
batchSize = 100)
+ .selectExpr("id", "name", "value", "data", "arrow_reuse_udf(name) as
udf_name")
+ val arrowExec =
collectNodes[ArrowEvalPythonExec](result.queryExecution.executedPlan)
+ assert(arrowExec.nonEmpty && arrowExec.head.child.supportsColumnar)
Review Comment:
Two suggestions for the test:
- `child.supportsColumnar` is also true for the other two paths. The test
takes the Arrow columnar path only because the cast of the test UDF is removed,
so if that changes, it would still pass without checking anything. Could it
also assert that all the inputs of the UDF are column references, for example
`arrowExec.head.udfs.flatMap(_.children).forall(_.isInstanceOf[AttributeReference])`?
- No test sets `spark.python.udf.pipelined.enabled=true`, so the change to
`ConcurrentLinkedQueue` is not covered, and no test has a source with both
Arrow and non-Arrow columns, which the new check at L157 is for. As the
pipelined mode is a `SparkConf`, the first one would need a suite that sets it
in `sparkConf`.
--
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]
