dongjoon-hyun commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4379185698
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -269,24 +304,42 @@ private class ArrowBackedPartitionReader(start: Int, end:
Int, batchSize: Int)
// Wrap in ArrowColumnVector
val columns: Array[ColumnVector] = Array(
- new ArrowColumnVector(idVector),
+ new ReleaseCountingArrowColumnVector(idVector),
new ArrowColumnVector(nameVector),
new ArrowColumnVector(valueVector),
new ArrowColumnVector(dataVector))
currentBatch = new ColumnarBatch(columns, count)
- allBatches.add(currentBatch)
+ allBatches.synchronized {
+ allBatches.add(currentBatch)
+ ArrowBackedDataSourceV2.returnedBatches.incrementAndGet()
+ ArrowBackedDataSourceV2.maxUnreleasedBatches.accumulateAndGet(
+ allBatches.size - numReleased, (a, b) => math.max(a, b))
+ }
current += count
true
}
override def get(): ColumnarBatch = currentBatch
+ // The consumer may still hold the returned batches, so they are closed when
the task completes.
+ // The listener is registered here rather than in the constructor: the
reader may be created by
+ // the writer thread of a pipelined Python UDF after the listener that waits
for that writer to
+ // stop, and the listeners run in reverse order, so the batches could be
freed while the writer
Review Comment:
**5. [Nit] Since SPARK-60113, the order of the listeners no longer matters
here**
The listener that stops the pipelined writer is now added with
`addTaskPreCompletionListener` (PythonRunner.scala:226). Such listeners run
before every completion listener, from any thread and even one added later
(TaskContext.scala:150-163), so "the listeners run in reverse order, so the
batches could be freed while the writer still reads them" no longer holds: a
listener added from the constructor, even on the writer thread, would also run
only after the writer has stopped. SPARK-60113 (`ca1f5c5752e`) was already in
this branch through `0dd22061215` when `09bae80321d` moved the registration
here after [my earlier
comment](https://github.com/apache/spark/pull/59581#discussion_r4318377947),
which predates SPARK-60113. Could the comment give the current reason, or could
the registration go back to the constructor, which would also make `closed`
unnecessary?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -255,12 +258,12 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val resultIter = if (inputColumnIndices.isDefined) {
// Path 2: simple column refs, non-Arrow vectors.
- val bufferedIter = inputIter.map { batch =>
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, {
batch =>
Review Comment:
**3. [Low] No test takes path 2, so this release is untested**
The new UDF test covers the Arrow columnar path (`name`) and the row
fallback path (`concat(name, '_')`), but not this one.
`ArrowBackedPartitionReader` always puts an `ArrowColumnVector` first
(ArrowBackedDataSourceV2.scala:306-310), so `isArrow` (L155-158) is true for
every batch of the test source, and simple column references always take path
1. If this iterator released a batch in `next()`, before the writer serializes
it, or stopped releasing at all, every test would still pass. Could
`ArrowBackedPartitionReader` get an option that wraps `id` in a
release-counting `ColumnVector` that is not an `ArrowColumnVector`, delegating
to the Arrow vector, so that `arrow_release_udf(name)` takes this path with the
same value and counter assertions?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala:
##########
@@ -44,10 +45,14 @@ class ColumnarToRowEvaluatorFactory(
inputs: Iterator[ColumnarBatch]*): Iterator[InternalRow] = {
assert(inputs.length == 1)
val toUnsafe = UnsafeProjection.create(childOutput, childOutput)
+ // Releases each batch once all its rows have been returned, before the
next batch is
+ // fetched, as the generated code of ColumnarToRowExec does in
advanceBatch: flatMap checks
+ // the iterator of the current batch, which runs its completion, before
the next batch.
inputs.head.flatMap { input =>
numInputBatches += 1
numOutputRows += input.numRows()
- input.rowIterator().asScala.map(toUnsafe)
+ CompletionIterator[InternalRow, Iterator[InternalRow]](
Review Comment:
**1. [Low] When `ColumnarToRowExec` is not compiled, this now releases
batches that nothing released before**
The description says that a reused vector whose `closeIfFreeable` is not a
no-op "is now closed by these consumers as well, as it already is by the
generated code of `ColumnarToRowExec` since SPARK-50235", and [your
reply](https://github.com/apache/spark/pull/59581#issuecomment-6276911959) says
that such a source "is broken today regardless of this path". That holds only
where `ColumnarToRowExec` is compiled. `CollapseCodegenStages` leaves it out
when its input or output has more than `spark.sql.codegen.maxFields` (100)
fields, nested ones included (WholeStageCodegenExec.scala:1214-1218), and a
stage that fails to compile or exceeds `spark.sql.codegen.hugeMethodLimit`
falls back to `child.execute()` (WholeStageCodegenExec.scala:1034 and 1046). In
those cases, and with `spark.sql.codegen.wholeStage=false`, `doExecute`
(Columnar.scala:84-97) runs this evaluator, which released nothing so far.
So a DSv2 source that refills one `VectorSchemaRoot` per batch and wraps it
in plain `ArrowColumnVector`s works today on a table with more than 100 fields,
and fails after this change once it has a struct column. `closeIfFreeable`
closes the struct vector, which drops its children, so the next
`VectorLoader.load` into that root fails with `should have as many children as
in the schema`, as `ArrowOutputProcessorSuite` shows for the Python output. The
wrapper's `childColumns` is null by then as well
(ArrowColumnVector.java:63-69). Could the user-facing section say that this now
also applies where `ColumnarToRowExec` is not compiled, e.g. for more than
`spark.sql.codegen.maxFields` fields? A line in the SQL migration guide may fit
too.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java:
##########
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.vectorized;
+
+import org.apache.arrow.vector.ValueVector;
+
+import org.apache.spark.sql.vectorized.ArrowColumnVector;
+
+/**
+ * An {@link ArrowColumnVector} whose producer may release or overwrite it
after returning it, for
+ * example by loading every batch into the same Arrow vectors. A consumer must
not hold it past the
+ * next batch.
+ *
+ * <p>Like the vectors that {@link WritableColumnVector} and {@link
ConstantColumnVector} reuse
+ * across batches, the producer releases it: {@link #closeIfFreeable()}, which
a consumer calls
+ * once it is done with a batch, is a no-op, and the producer frees the vector
with
+ * {@link #close()}.
+ */
+public class ReusedArrowColumnVector extends ArrowColumnVector {
+
+ public ReusedArrowColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ @Override
+ public void closeIfFreeable() {
Review Comment:
**4. [Nit] Only the top-level vector is covered, not the child columns of a
struct**
`getChild` returns the plain `ArrowColumnVector`s that `ArrowColumnVector`
builds for the fields of a struct (ArrowColumnVector.java:163 and 257), so
`closeIfFreeable` on one of them still closes a vector of the producer's root.
No consumer in Spark releases a child column on its own today, so this is only
about the Javadoc above, which says that a consumer's `closeIfFreeable` leaves
the vector to its producer. Could the Javadoc say that this covers the
top-level vector, or could `getChild` return reused vectors too?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala:
##########
@@ -117,28 +117,13 @@ case class ArrowEvalPythonExec(
// accept columnar input to avoid the ColumnarToRow -> ArrowWriter
round-trip.
// The Arrow FieldVectors are extracted directly from ArrowColumnVector and
// serialized to IPC, bypassing the row-based ArrowWriter conversion.
+ // This node then outputs only columnar batches, so the planner inserts a
+ // ColumnarToRowExec above it for a row-based parent, which releases each
Review Comment:
**2. [Low] In pipelined mode, the Arrow columnar path now releases its input
batches on another thread than the one reading them**
On this path, each output batch carries the columns of its input batch
(ColumnarArrowEvalPythonEvaluatorFactory.scala:226), so the `ColumnarToRowExec`
described here calls `closeIfFreeable` on them from the task thread. With
`spark.python.udf.pipelined.enabled=true`, the writer pulls the upstream
iterator on its own pool thread (PythonRunner.scala:1577 and 1594), so
`PartitionReader.next()` for the later batches (DataSourceRDD.scala:240) runs
on that thread at the same time. Before this PR nothing released these batches,
and `advanceBatch`, the non-codegen evaluator and the row queue paths all
release on the thread that reads. `ArrowBackedPartitionReader` synchronizes
`allBatches` (ArrowBackedDataSourceV2.scala:285 and 313) because its
`closeIfFreeable` and `next()` can now run on these two threads.
Arrow allocators are thread-safe, but a source whose vectors return their
buffers to a pool of the reader, or update other reader state when closed, now
has to be thread-safe as well. Columnar input and pipelined mode are both off
by default, so documenting it may be enough. Could the class doc of
`ColumnarArrowEvalPythonEvaluatorFactory` (L45-53), which already says that the
batches must stay valid until their results come back, also say that they are
then released from the task thread while the writer may be reading the next
ones?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -114,6 +115,103 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("ColumnarToRow releases each batch once its rows are consumed") {
+ Seq(true, false).foreach { codegen =>
+ withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString) {
+ val numRows = 2000
+ val numBatches = 20
+ val df = readArrowSource(numRows = numRows, batchSize = numRows /
numBatches)
+ val plan = df.queryExecution.executedPlan
+ assert(collectNodes[ColumnarToRowExec](plan).nonEmpty)
+ assert(collectNodes[WholeStageCodegenExec](plan).nonEmpty == codegen)
+
+ ArrowBackedDataSourceV2.resetBatchCounters()
+ val rows = df.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ assert(rows.forall(row => row.getString(1) == s"row_${row.getInt(0)}"))
+ // A batch is released before the next one is read.
+ assert(ArrowBackedDataSourceV2.returnedBatches.get == numBatches)
+ assert(ArrowBackedDataSourceV2.releasedBatches.get == numBatches)
+ assert(ArrowBackedDataSourceV2.maxUnreleasedBatches.get == 1)
+ }
+ }
+ }
+
+ test("Arrow Python UDF over many batches releases each input batch") {
+ assume(shouldTestPandasUDFs)
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_release_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ val numRows = 20000
+ val numBatches = 20
+ // On a string column, the cast of the test UDF is removed, so the batches
take the Arrow
+ // columnar path, whose output batches carry the input columns to the
ColumnarToRowExec above
+ // ArrowEvalPythonExec, which releases them. On an expression, they take
the row fallback
+ // path, which releases each batch once it has been written to Python.
+ for {
+ codegen <- Seq(true, false)
+ (input, suffix) <- Seq("name" -> "", "concat(name, '_')" -> "_")
+ } withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString) {
+ val result = readArrowSource(numRows = numRows, batchSize = numRows /
numBatches)
+ .selectExpr("id", "name", "value", "data",
s"arrow_release_udf($input) as udf_out")
+ val plan = result.queryExecution.executedPlan
+ val arrowExec = collectNodes[ArrowEvalPythonExec](plan)
+ assert(arrowExec.nonEmpty && arrowExec.head.supportsColumnar)
+
assert(arrowExec.head.udfs.forall(_.children.forall(_.isInstanceOf[AttributeReference]))
+ == suffix.isEmpty, s"Unexpected UDF inputs:\n$plan")
+ assert(collectNodes[ColumnarToRowExec](plan).exists(
+ _.find(_.isInstanceOf[ArrowEvalPythonExec]).isDefined), s"No
ColumnarToRowExec:\n$plan")
+
+ ArrowBackedDataSourceV2.resetBatchCounters()
+ val rows = result.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ rows.foreach { row =>
+ val id = row.getInt(0)
+ assert(row.getString(1) == s"row_$id", s"name mismatch for id $id")
+ assert(row.getDouble(2) == id * 0.1, s"value mismatch for id $id")
+ assert(row.getString(3).length == 100, s"data mismatch for id $id")
+ assert(row.getString(4) == s"row_$id$suffix", s"UDF result mismatch
for id $id")
+ }
+ // Every batch is released before the task completes. The row fallback
path releases a
+ // batch before it reads the next one, while the Arrow columnar path
holds the batches
+ // that it has written to Python until their results are consumed.
+ assert(ArrowBackedDataSourceV2.returnedBatches.get == numBatches)
+ assert(ArrowBackedDataSourceV2.releasedBatches.get == numBatches)
+ if (suffix.nonEmpty) {
+ assert(ArrowBackedDataSourceV2.maxUnreleasedBatches.get == 1)
+ }
+
+ // A limit stops consuming the output batches early, and the generated
code of
+ // ColumnarToRowExec closes the current one.
+ assert(result.limit(10).selectExpr("count(*)").head().getLong(0) == 10)
+ }
+ }
+
+ test("A plan built with columnar input runs as built after the configuration
changes") {
Review Comment:
**6. [Nit] Could the new tests carry the JIRA ID?**
This test reproduces the `ClassCastException` that this PR fixes, so it
could be named `SPARK-60306: A plan built with columnar input runs as built
after the configuration changes`, as Spark tests for a JIRA issue usually are.
The same applies to the tests at L118 and L140, and to the one in
`ArrowOutputProcessorSuite`.
--
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]
