dongjoon-hyun commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4318377906
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -242,9 +249,10 @@ abstract class BaseSliceArrowOutputProcessor(
batch
}
+ // The slices of a batch are closed once the next one is produced, so they
are reused vectors.
Review Comment:
**This comment covers only the slices, but `getVectors` also wraps the
reader's own `root`, and `produceBatch` closes that root itself.** `getRoot`
returns the unsliced `root` for a batch that fits (L275-277) and for the last
chunk of every batch in `SliceRecordsArrowOutputProcessorImpl` (L327-328). That
is the case where the no-op `closeIfFreeable` from L255 matters, since closing
the struct vector of the root removes its children for every later batch.
Separately, and pre-existing since SPARK-51769: when the previous chunk was
`root` and the next `getRoot` returns a slice, L236-241 close `prevVectors` and
`prevRoot`, i.e. the reader's root that holds the batch being sliced. With
`spark.sql.execution.arrow.maxRecordsPerOutputBatch=3` and a `mapInArrow`
partition whose Python output is two 10-row batches, the next `root.slice(3,
3)` of the second batch would throw `IllegalArgumentException` (`valueCount`
0). With `maxBytesPerOutputBatch`, a small batch followed by a large one woul
d make `getRoot` return the cleared root with its old row count.
`ReusedArrowColumnVector` cannot prevent this, because L238 calls `close()`.
Could the comment say that these vectors are either the reader's root, into
which the next batch is loaded, or a slice of it that is closed once the next
one is produced? And could `produceBatch` close `prevRoot` and `prevVectors`
only when `prevRoot ne root`, here or in a separate JIRA?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -114,6 +114,71 @@ 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 = 50000
+ val numBatches = 500
+ 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)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_release_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ val numRows = 50000
+ val numBatches = 500
+ // 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
row conversion of
+ // ArrowEvalPythonExec, which releases them. On an expression, they take
the row fallback
+ // path, which releases each batch once it has been written to Python.
+ Seq("name" -> "", "concat(name, '_')" -> "_").foreach { case (input,
suffix) =>
+ val result = readArrowSource(numRows = numRows, batchSize = numRows /
numBatches)
+ .selectExpr("id", "name", "value", "data",
s"arrow_release_udf($input) as udf_out")
+ val arrowExec =
collectNodes[ArrowEvalPythonExec](result.queryExecution.executedPlan)
+ assert(arrowExec.nonEmpty && arrowExec.head.supportsColumnar)
Review Comment:
Minor: the `name` iteration relies on the comment at L150-153, not on an
assertion, to take the Arrow columnar path. Today `SimplifyCasts` removes the
`Cast(name, StringType)` that `TestScalarPandasUDF` adds
(IntegratedUDFTestUtils.scala:1403), so this iteration does cover the release
in `ArrowEvalPythonExec.doExecute`. But L158 holds on every path, the counters
at L173-174 also reach 500 on the row fallback path, and `maxUnreleasedBatches`
is checked only for the expression (L175-177). If the cast were kept one day,
both iterations would take the row fallback path and the test would still pass.
Could L158 also check the path, e.g.
`assert(arrowExec.head.udfs.forall(_.children.forall(_.isInstanceOf[AttributeReference]))
== suffix.isEmpty)`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -242,9 +249,10 @@ abstract class BaseSliceArrowOutputProcessor(
batch
Review Comment:
**Pre-existing since SPARK-51769: the batch returned here repeats rows at
the end of every Python batch that `SliceRecordsArrowOutputProcessorImpl`
slices.** Once `remainingRows <= arrowMaxRecordsPerOutputBatch`, its `getRoot`
(L323-334) returns the whole `root` instead of the remaining rows, and L248
sets `numRows` to the full row count. With `maxRecordsPerOutputBatch=3`, a
10-row batch comes out as rows 0-2, 3-5 and 6-8, then 0-9 again: 19 rows
instead of 10. `test_different_output_length` in `test_arrow_map.py`, run by
`MapInArrowWithOutputArrowBatchSlicingRecordsTests`, would get 199 rows for
100, but it compares sets, so it passes. The loop in
`ArrowOutputProcessorSuite` with `new
SliceRecordsArrowOutputProcessorImpl(reader, metrics(), 2)` would show it as
rows 0, 1, 0, 1, 2 for the first batch. Could `getRoot` return
`root.slice(currentRowIdx, remainingRows)` once `currentRowIdx > 0`, here or in
a separate JIRA?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala:
##########
@@ -113,10 +113,11 @@ case class ArrowEvalPythonExec(
// Columnar path: delegate to doExecuteColumnar, flatten to
Review Comment:
**Pre-existing since SPARK-60294, but in the branch this PR edits: L112
chooses between the columnar and row paths at execution time, from the live
session conf.** `supportsColumnar` (L104-105) reads
`conf.arrowPySparkUDFColumnarInputEnabled`, and `SparkPlan.conf` is
`session.sessionState.conf` (SparkPlan.scala:70-76), while the transitions were
inserted with the value at planning time. A Dataset reuses its `executedPlan`
across actions (Dataset.scala:2334-2339). So enable
`spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`, `collect()` a
pandas UDF over an Arrow-backed DSv2 scan, disable the conf, and `collect()`
the same Dataset again: the plan still has `ArrowEvalPythonExec` directly over
the columnar-only `BatchScanExec`, but L112 is now false, so
`super.doExecute()` calls `child.execute().map(_.copy())`
(EvalPythonExec.scala:73). `BatchScanExec.doExecute`
(DataSourceV2ScanExecBase.scala:210-216) then hands the `ColumnarBatch`es of
its columnar `DataSourceRDD` over as
rows, which should end in a `ClassCastException`. A `.cache()`d plan that is
materialized after the conf changes takes the same route. Before SPARK-60294,
this check was the plan-time `child.supportsColumnar`. Could the choice be
fixed when the plan is built, for example with a flag set at planning time, or
through the planner as in my comment on L120?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowOutputProcessorSuite.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * 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.python
+
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.vectorized.ReusedColumnVector
+import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.unsafe.types.UTF8String
+
+class ArrowOutputProcessorSuite extends SparkFunSuite {
+
+ private val schema = new StructType()
+ .add("s", new StructType().add("a", IntegerType).add("b", StringType))
+ private val numBatches = 3
+ private val rowsPerBatch = 3
+
+ /** Writes `numBatches` batches of a struct column as an Arrow stream, as a
Python worker does. */
+ private def arrowStream(): Array[Byte] = {
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator("arrow-output-writer", 0,
Long.MaxValue)
+ val root = VectorSchemaRoot.create(
+ ArrowUtils.toArrowSchema(schema, null, errorOnDuplicatedFieldNames =
true,
+ largeVarTypes = false),
+ allocator)
+ try {
+ val out = new ByteArrayOutputStream()
+ val writer = new ArrowStreamWriter(root, null, out)
+ val arrowWriter = ArrowWriter.create(root)
+ writer.start()
+ (0 until numBatches * rowsPerBatch).grouped(rowsPerBatch).foreach {
values =>
+ values.foreach { v =>
+ arrowWriter.write(InternalRow(InternalRow(v,
UTF8String.fromString(s"v$v"))))
+ }
+ arrowWriter.finish()
+ writer.writeBatch()
+ arrowWriter.reset()
+ }
+ writer.end()
+ out.toByteArray
+ } finally {
+ root.close()
+ allocator.close()
+ }
+ }
+
+ test("Struct output columns stay valid after a consumer releases each
batch") {
+ // The output processor loads every batch into the same vectors. A
consumer releases each
+ // batch with closeIfFreeable once it is done with it, which must leave
the vectors, including
+ // the child columns of a struct, to the processor.
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator("arrow-output-reader", 0,
Long.MaxValue)
+ val reader = new ArrowStreamReader(new
ByteArrayInputStream(arrowStream()), allocator)
+ val processor = new ArrowOutputProcessorImpl(reader, metrics())
Review Comment:
**Only `ArrowOutputProcessorImpl` is tested here, so the change to
`BaseSliceArrowOutputProcessor.getVectors` (PythonArrowOutput.scala:255) has no
test.** The slicing tests in `test_arrow_map.py` go through `mapInArrow`, which
never releases its output batches, so nothing would fail if
PythonArrowOutput.scala:255 went back to `ArrowColumnVector`. That line matters
on the Arrow columnar path when output slicing is enabled but no batch is
split, the only case the `passThruRows == numRows` assertion
(ColumnarArrowEvalPythonEvaluatorFactory.scala:223-225) allows: `getRoot` then
returns the reader's root, which `ColumnarToRowIterator` in
`ArrowEvalPythonExec.doExecute` releases. The struct-returning UDF of the
earlier test request
(https://github.com/apache/spark/pull/59581#discussion_r4309307580) is also
still open: `TestScalarPandasUDF` always returns strings, so the new many-batch
UDF test cannot notice cleared struct children. Could this test also run the
loop with `new SliceRecord
sArrowOutputProcessorImpl(reader, metrics(), rowsPerBatch)` and `new
SliceBytesArrowOutputProcessorImpl(reader, metrics(), Int.MaxValue)`, i.e. with
limits that never split a batch? The whole-stage codegen on/off part of that
request makes no difference here: `ArrowEvalPythonExec.supportsRowBased` is
true, so both modes go through `ColumnarToRowIterator`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java:
##########
@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+/**
+ * A marker for a {@link org.apache.spark.sql.vectorized.ColumnVector} that
its producer may
+ * release or overwrite after returning it in a batch: for example, a reader
that refills the same
+ * vectors for every batch, or frees the vectors of a batch once the next one
is read or once it is
+ * closed. A consumer may only read such a vector until it asks the producer
for the next batch,
+ * and must copy what it needs past that point.
+ *
+ * <p>There is no other agreement between the producer and the consumer of a
batch, so the type of
+ * its vectors is the protocol, as for {@code ColumnVector.closeIfFreeable}:
the producer owns the
+ * lifetime of a vector carrying this marker, and its {@code closeIfFreeable}
is a no-op. A vector
+ * without it may be held by the consumer past the next batch without a copy,
and the consumer
Review Comment:
**The Arrow cache scan frees the vectors of a batch once the next one is
read, as L22 describes, but its vectors are not marked, so L28-30 let a
consumer hold them past the next batch.**
`ArrowCachedBatchToColumnarBatchIterator.next()` closes `previousRoot`
(ArrowCachedBatchSerializer.scala:1332-1337), while it wraps the vectors in
plain `ArrowColumnVector`s (ArrowCachedBatchSerializer.scala:1356 and 1359). In
#59551 you listed this iterator next to `ArrowOutputProcessorImpl` and wrote
that it "should only close a root once none of its vectors is still in use, or
mark its vectors as reused", but only `ArrowOutputProcessorImpl` is marked
here. Once the follow-up holds unmarked inputs on the Arrow columnar path, the
Arrow cache case from #59551 comes back: with
`spark.sql.cache.serializer=org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer`
and columnar input enabled, a pandas UDF over `spark.range(0, 2000000, 1,
1).cache()` takes that path, and `next()` for batch N+1
closes the root of batch N while its pass-through columns are still held.
Could you wrap the cache vectors in `ReusedArrowColumnVector` in this PR too?
`vectorTypes` (ArrowCachedBatchSerializer.scala:141), the `isArrow` check and
the zero-copy re-cache path all still accept it. The iterator would then also
need to close `previousRoot` once it is exhausted, since a consumer's
`closeIfFreeable` would no longer free the last root.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -114,6 +114,71 @@ 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 = 50000
Review Comment:
Minor: the two new tests read 50,000 rows in 500 batches, each of them
twice. In the CI run for 2b49870 they took 759 ms and 1.852 s, about half of
the suite's 5.07 s, and the UDF test makes 1,000 sequential round trips to
Python. Their assertions (`returnedBatches == releasedBatches == numBatches`,
`maxUnreleasedBatches == 1`) only need two or more batches. Could this test use
2,000 rows in 20 batches, and the UDF test 20,000 rows in 20 batches? The
latter still gives two output batches on the row fallback path
(`arrowMaxRecordsPerBatch` is 10,000 by default), so the release and refill of
its reused output batch is still exercised.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowOutputProcessorSuite.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * 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.python
+
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.vectorized.ReusedColumnVector
+import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.unsafe.types.UTF8String
+
+class ArrowOutputProcessorSuite extends SparkFunSuite {
Review Comment:
Minor: the PR description does not mention this suite or the new UDF test.
"How was this patch tested?" still describes only `ColumnarToRow releases each
batch once its rows are consumed`, and the "Ran ..." list does not include
`ArrowOutputProcessorSuite`. The template asks "If tests were added, say they
were added here.", and `dev/merge_spark_pr.py` copies the description into the
commit message. Also, "The marker is checked on the top-level columns of a
batch" describes a check that no code performs yet. Could you update the
description?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -246,14 +266,44 @@ private class ArrowBackedPartitionReader(start: Int, end:
Int, batchSize: Int)
// Single allocator for the entire partition reader lifetime.
// Arrow vectors from earlier batches may still be referenced by
// downstream operators (e.g., pass-through column queue), so we
- // must NOT close the allocator until the reader itself is closed.
+ // must NOT close the allocator until the task completes.
private val allocator: BufferAllocator =
ArrowUtils.rootAllocator.newChildAllocator(
s"arrow-test-reader-$start", 0, Long.MaxValue)
- // Track all batches so we can close them in close().
+ // Track all batches so we can close them when the task completes.
private val allBatches =
new java.util.ArrayList[ColumnarBatch]()
+ private var numReleased = 0
+ @volatile private var taskCompleted = false
+
+ TaskContext.get().addTaskCompletionListener[Unit] { _ =>
Review Comment:
Minor: registering the cleanup listener in the constructor can make it run
before the pipelined Python writer stops. `DataSourceRDD` creates the reader on
the first `hasNext` (DataSourceRDD.scala:188-195), and on the row path
(columnar input off, the default) nothing peeks first, so with
`spark.python.udf.pipelined.enabled=true` the writer thread creates this reader
after `createPipelinedDataIn` registered its writer-wait listener
(PythonRunner.scala:553-561). Listeners run in LIFO order, so on a task kill,
or an early stop that is not pushed below the UDF, L280-287 free every batch
and the allocator while the writer may still be reading the current batch. That
is the SPARK-33277 hazard the writer-wait listener exists for. Before this PR,
`close()` freed them, either on the consuming thread or from the listener that
`DataSourceRDD` registers before the runner. No test runs this source in
pipelined mode today. Could `close()` register this listener instead of the
constructor?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -232,10 +247,15 @@ private class ReusingArrowPartitionReader(start: Int,
end: Int, batchSize: Int)
/**
* Partition reader that produces Arrow-backed [[ColumnarBatch]]es.
*
- * Each batch contains up to [[ArrowBackedDataSourceV2.BATCH_SIZE]] rows with:
+ * Each batch contains up to `batchSize` rows with:
Review Comment:
Nit: this list still has only `id`, `name` and `value`, while `next()` also
builds a `data` column of 100 random characters (L317-318, L404), which the new
UDF test checks (ArrowColumnarPythonUDFSuite.scala:167). Could you add `- data
(STRING): 100 random characters`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala:
##########
@@ -113,10 +113,11 @@ case class ArrowEvalPythonExec(
// Columnar path: delegate to doExecuteColumnar, flatten to
// UnsafeRow. ColumnarBatchRow from rowIterator() is NOT
// UnsafeRow, and downstream operators (e.g., outer
- // EvalPythonExec) may cast to UnsafeRow.
+ // EvalPythonExec) may cast to UnsafeRow. Like ColumnarToRowExec,
+ // this releases each batch once its rows are consumed: the
+ // batches of the Arrow columnar path carry the input columns.
doExecuteColumnar().mapPartitionsInternal { batchIter =>
- val toUnsafe = UnsafeProjection.create(schema)
- batchIter.flatMap(_.rowIterator().asScala.map(toUnsafe))
+ new ColumnarToRowIterator(batchIter, UnsafeProjection.create(schema))
Review Comment:
Minor: this private flatten, and now the release in it, are needed only
because L106 declares `supportsRowBased = true`. With the default
(`!supportsColumnar`), `ApplyColumnarRulesAndInsertTransitions` would put a
`ColumnarToRowExec` above this node (Columnar.scala:581-584) and still keep the
scan-to-UDF edge columnar (Columnar.scala:569-570). With columnar input off,
the plan would not change. That `ColumnarToRowExec` already releases batches in
both codegen and non-codegen modes, is fused into the whole-stage codegen of
the parent, and reports `numInputBatches` and `numOutputRows`. `doExecute`
would then need only the row path, and the choice would be fixed at planning
time (see my comment on L113). SPARK-56350 declared `supportsRowBased = true`
"so Spark's transition rules no longer insert ColumnarToRowExec when the child
supports columnar output", which `supportsColumnar` alone already achieves. The
trade-offs are that EXPLAIN would show a ColumnarToRow above ArrowEvalPython w
hen columnar input is enabled, and that the codegen limit cleanup
(Columnar.scala:201-204) calls `close()`, not `closeIfFreeable()`, on the
Python result vectors. Could this be considered, here or in a follow-up?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java:
##########
@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+/**
+ * A marker for a {@link org.apache.spark.sql.vectorized.ColumnVector} that
its producer may
+ * release or overwrite after returning it in a batch: for example, a reader
that refills the same
+ * vectors for every batch, or frees the vectors of a batch once the next one
is read or once it is
+ * closed. A consumer may only read such a vector until it asks the producer
for the next batch,
+ * and must copy what it needs past that point.
+ *
+ * <p>There is no other agreement between the producer and the consumer of a
batch, so the type of
+ * its vectors is the protocol, as for {@code ColumnVector.closeIfFreeable}:
the producer owns the
+ * lifetime of a vector carrying this marker, and its {@code closeIfFreeable}
is a no-op. A vector
+ * without it may be held by the consumer past the next batch without a copy,
and the consumer
+ * releases it with {@code closeIfFreeable} once done with it.
+ *
+ * <p>The marker covers the whole vector, so it is only checked on the
top-level columns of a
+ * batch. The child vectors of a struct, array or map column belong to that
column and share its
+ * lifetime.
+ */
+public interface ReusedColumnVector {
Review Comment:
Minor: nothing in the main code reads this interface yet, and both fixes
here only need the no-op `closeIfFreeable` of `ReusedArrowColumnVector`. Its
only reader is the test assertion at ArrowOutputProcessorSuite.scala:85, so
removing `implements ReusedColumnVector` from the three classes changes no
behavior. My comment on L29 shows why the classification is easier to review
next to the code that relies on it, and in #59551 you planned to add the marker
together with closing the gaps in path 1. This is about where the marker lands,
not about the default for unmarked vectors raised in the earlier review. Could
the interface and the `implements` clauses move to that follow-up, with the
test asserting `isInstanceOf[ReusedArrowColumnVector]` here?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -114,6 +114,71 @@ 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 = 50000
+ val numBatches = 500
+ 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)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true") {
Review Comment:
Nit: `sparkConf` (L38-40) already enables
`ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED` for every test in this suite, and
the other columnar tests set only `ARROW_PYSPARK_EXECUTION_ENABLED`. Could this
be `withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true")`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala:
##########
@@ -44,12 +42,52 @@ class ColumnarToRowEvaluatorFactory(
inputs: Iterator[ColumnarBatch]*): Iterator[InternalRow] = {
assert(inputs.length == 1)
val toUnsafe = UnsafeProjection.create(childOutput, childOutput)
- inputs.head.flatMap { input =>
+ new ColumnarToRowIterator(inputs.head, toUnsafe, { batch =>
numInputBatches += 1
- numOutputRows += input.numRows()
- input.rowIterator().asScala.map(toUnsafe)
+ numOutputRows += batch.numRows()
+ })
+ }
+ }
+}
+
+/**
+ * Flattens columnar batches into `UnsafeRow`s. A batch is released with
closeIfFreeable once all
+ * its rows have been returned, before the next batch is fetched, as the
generated code of
+ * [[ColumnarToRowExec]] does in [[ColumnarToRowExec.advanceBatch]]. The rows
are projected into a
+ * buffer of their own, so the last row of a batch stays valid after the batch
is released.
+ *
+ * @param onNextBatch called on each batch once it is fetched
+ */
+private[sql] class ColumnarToRowIterator(
Review Comment:
Minor: the existing `flatMap` with Spark's `CompletionIterator` gives the
same release order without a new class. Scala 2.13's `Iterator.flatMap` calls
`hasNext` of the inner iterator, which runs the completion, before
`hasNext`/`next()` of the outer one, so each batch is released after its last
row and before the next one is fetched, as `ColumnarToRowExec.advanceBatch`
does:
```scala
inputs.head.flatMap { input =>
numInputBatches += 1
numOutputRows += input.numRows()
CompletionIterator[InternalRow, Iterator[InternalRow]](
input.rowIterator().asScala.map(toUnsafe), input.closeIfFreeable())
}
```
`ArrowEvalPythonExec.doExecute` can do the same, and
`FlatMapGroupsInPandasWithStateExec` already uses this form. That removes this
class, its `NoSuchElementException` path and the `onNextBatch` default that
only one caller uses. Could you use it instead?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -246,14 +266,44 @@ private class ArrowBackedPartitionReader(start: Int, end:
Int, batchSize: Int)
// Single allocator for the entire partition reader lifetime.
// Arrow vectors from earlier batches may still be referenced by
// downstream operators (e.g., pass-through column queue), so we
- // must NOT close the allocator until the reader itself is closed.
+ // must NOT close the allocator until the task completes.
private val allocator: BufferAllocator =
ArrowUtils.rootAllocator.newChildAllocator(
s"arrow-test-reader-$start", 0, Long.MaxValue)
- // Track all batches so we can close them in close().
+ // Track all batches so we can close them when the task completes.
private val allBatches =
new java.util.ArrayList[ColumnarBatch]()
+ private var numReleased = 0
+ @volatile private var taskCompleted = false
+
+ TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+ taskCompleted = true
+ allBatches.synchronized {
+ allBatches.forEach(_.close())
+ allBatches.clear()
+ }
+ allocator.close()
+ }
+
+ /** The first column of a batch, which counts the release of the batch. */
+ private class ReleaseCountingArrowColumnVector(vector: ValueVector)
+ extends ArrowColumnVector(vector) {
+ private var closed = false
+
+ override def close(): Unit = {
Review Comment:
Nit: counting in `close()` needs both `closed` and `taskCompleted` to skip
the closes of the completion listener. Every consumer under test releases
through `closeIfFreeable`, so counting there counts exactly what L83-84
describe, and both flags can go:
```scala
override def closeIfFreeable(): Unit = {
allBatches.synchronized {
numReleased += 1
ArrowBackedDataSourceV2.releasedBatches.incrementAndGet()
}
super.closeIfFreeable()
}
```
A double release would then show up as `releasedBatches > numBatches`
instead of being hidden by `closed`. Could you count it this way?
--
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]
