HyukjinKwon commented on code in PR #59551:
URL: https://github.com/apache/spark/pull/59551#discussion_r4299677395
##########
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:
Thanks, fixed in c6bcd015ef34: the Arrow columnar path is now only taken
when no column has a list view (`ListViewVector` or `LargeListViewVector`) at
any level, so these batches take the row queue path. `Arrow-backed source with
columns that the Arrow columnar path does not copy` reads a list view column,
and fails with `UnsupportedOperationException: VectorVisitor for ListViewVector
is not supported.` without the check.
##########
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 for forcing the order to show it. Fixed in c6bcd015ef34 with a lock,
as you suggested. The listener sets `completed` and releases the queued copies
under it. The writer copies outside of it, then adds the copies, or releases
them right away if the task completed meanwhile. The allocator is closed by
whichever of the two ends last, so a copy that blocks does not block the
listener, and nothing leaks. As @viirya pointed out, the runner itself is
SPARK-60113 (#59332), and the comment here now also says that this listener
must be registered before `compute`.
##########
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:
Thanks, done in c6bcd015ef34:
- The tests now also check that the inputs of the UDF are columns, besides
`child.supportsColumnar`.
- `ArrowColumnarPythonUDFPipelinedSuite` runs the tests of the suite with
`spark.python.udf.pipelined.enabled=true`, and `Arrow-backed source with
columns that the Arrow columnar path does not copy` reads a column that is not
Arrow-backed, and one that is a list view. Without the check at L157, it fails
with an internal error for the first one.
##########
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:
Thanks, done in c6bcd015ef34: a vector with fewer values than the rows of
its batch fails with an internal error that says so, such as `The Arrow vector
of the column name has 0 values, but its batch has 100 rows.`
##########
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:
Thanks for checking it. As you said, it is not a regression of this PR, and
paths 2 and 3 have the same gap, so I opened SPARK-60306 to release the input
batches of all three paths once they are written to Python, with your change
for this path.
##########
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:
Done in c6bcd015ef34: the comment says that a columnar parent, such as the
Arrow cache, may already have closed the columns with `closeIfFreeable`, and
that closing them again has no effect.
##########
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:
Done in c6bcd015ef34: `ArrowBackedPartitionReader` takes `batchSize` too.
--
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]
