viirya commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4362034968
##########
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:
Added in 09bae80321d: the test asserts that the UDF inputs are attribute
references exactly for the `name` iteration, and that a `ColumnarToRowExec`
sits above `ArrowEvalPythonExec`.
##########
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:
Good catch, done in 09bae80321d: `close()` registers the listener. By the
time the reader is closed, either the writer has read every batch, or the task
is completing and the listener that waits for the pipelined writer has already
run.
##########
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:
Done in 09bae80321d: 2,000 rows in 20 batches for the ColumnarToRow test,
and 20,000 rows in 20 batches for the UDF test.
##########
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:
Thanks, done in 09bae80321d: `ArrowEvalPythonExec` no longer overrides
`supportsRowBased`. The planner inserts a `ColumnarToRowExec` above it when it
takes columnar input, and `doExecute` is the row path only, which also fixes
your comment on L113. About the trade-offs: EXPLAIN now shows the
`ColumnarToRow` above `ArrowEvalPython` when columnar input is enabled, and the
UDF test runs a limit with whole-stage codegen on, where the generated code
closes the current batch, including the Python result vectors.
`ArrowOutputProcessorImpl.close` closing them again afterwards is fine.
##########
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:
Agreed, done in 09bae80321d: the interface and the `implements` clauses
moved to #59582, and the test asserts `isInstanceOf[ReusedArrowColumnVector]`.
It also leaves this PR independent of the open question in
https://github.com/apache/spark/pull/59581#issuecomment-6276911959, whether the
follow-up should classify the input with the marker, with an overridden
`closeIfFreeable`, or both: `ReusedArrowColumnVector` is needed either way.
##########
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:
Thanks, I updated the description for 09bae80321d: the testing section lists
the new tests and suites, and the marker is no longer described, since it moved
to #59582.
--
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]
