This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c2c80ff57851 [SPARK-60306][SQL] Release consumed columnar batches in
non-codegen ColumnarToRow and Arrow Python UDFs
c2c80ff57851 is described below
commit c2c80ff57851de034dbc183f5df2e16b9b21a819
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Mon Oct 26 18:43:34 2026 -0700
[SPARK-60306][SQL] Release consumed columnar batches in non-codegen
ColumnarToRow and Arrow Python UDFs
### What changes were proposed in this pull request?
SPARK-50235 made the generated code of `ColumnarToRowExec` release each
consumed batch with `closeIfFreeable`, and left the non-codegen path for a
follow-up. This PR makes the remaining consumers of columnar batches release
them:
- The non-codegen `ColumnarToRowEvaluatorFactory` releases each batch with
a `CompletionIterator` once all its rows have been returned, before it fetches
the next one, as `ColumnarToRowExec.advanceBatch` does.
- `ArrowEvalPythonExec` no longer declares `supportsRowBased`. When it
takes columnar input, it only outputs columnar batches, so the planner inserts
a `ColumnarToRowExec` above it for a row-based parent, which releases the
batches, instead of `doExecute` flattening them. `doExecute` is now only the
row path, so whether the node takes columnar input is fixed when the plan is
built.
- In `ColumnarArrowEvalPythonEvaluatorFactory`, the row queue and row
fallback paths copy what they need out of each input batch, so they release the
batch once the Python writer asks for the next one.
Releasing the output batches of `ArrowEvalPythonExec` must not free the
result columns, which the Python runner loads into the same Arrow vectors for
every batch, or slices of them. So the Arrow output processors of the Python
runners now return `ReusedArrowColumnVector`, a new internal
`ArrowColumnVector` whose `closeIfFreeable` is a no-op, like that of
`WritableColumnVector` and `ConstantColumnVector`.
This is also a step toward the follow-up to SPARK-60294 (#59551), #59582,
in which the Arrow columnar path of Python UDFs holds only the input columns
that their producer does not reuse.
### Why are the changes needed?
- When batches whose release is left to their consumer, e.g. the fresh
Arrow vectors of a DataSource V2 reader, are consumed by the non-codegen
`ColumnarToRowExec` or by an Arrow Python UDF, their vectors stay allocated
until the task completes, so the memory grows with the size of the partition.
- Closing the Arrow vectors of the Python runner's output root, as
`closeIfFreeable` on them would do, closes its struct vectors, which drops
their children, so loading the next batch with a struct result fails.
- `ArrowEvalPythonExec.doExecute` chose between the columnar and row paths
from the configuration at execution time, while the transitions were inserted
at planning time. A plan built with
`spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` and first executed
after it was disabled read the batches of a columnar-only scan as rows, and
failed with `ClassCastException: class
org.apache.spark.sql.vectorized.ColumnarBatch cannot be cast to class
org.apache.spark.sql.catalyst.InternalRow`.
### Does this PR introduce _any_ user-facing change?
Yes, for Data Source V2 sources that reuse their `ColumnVector`s across
batches without overriding `closeIfFreeable` to do nothing. The conversion of
columnar batches to rows now releases each batch with `closeIfFreeable` also
when it is not whole-stage code-generated, e.g. when the query has more than
`spark.sql.codegen.maxFields` fields or `spark.sql.codegen.wholeStage` is
`false`, as the generated code of `ColumnarToRowExec` has done since
SPARK-50235 (3.5.4 and 4.0). Such a source [...]
Otherwise, this fixes the `ClassCastException` above, in unreleased
versions only. With columnar input enabled, EXPLAIN now shows a `ColumnarToRow`
above `ArrowEvalPython`.
### How was this patch tested?
New tests:
- `ArrowColumnarPythonUDFSuite`: `SPARK-60306: ColumnarToRow releases each
batch once its rows are consumed` reads 20 batches of fresh Arrow vectors with
whole-stage codegen enabled and disabled, and checks that the consumer releases
every batch, with at most one unreleased at a time. `SPARK-60306: Arrow Python
UDF over many batches releases each input batch` runs a pandas UDF over 20
batches, on a column (the Arrow columnar path), on a column of batches whose
first column is not an ` [...]
- `ArrowOutputProcessorSuite`: `SPARK-60306: Struct output columns stay
valid after a consumer releases each batch` reads a struct column over several
batches through `ArrowOutputProcessorImpl`, and through the slicing processors
with limits that never split a batch, and releases each output batch with
`closeIfFreeable`.
The new tests fail if the result columns are plain `ArrowColumnVector`s, if
the non-codegen `ColumnarToRowExec` stops releasing batches, or if the row
queue path stops releasing them or releases them before they are written to
Python.
Ran `ArrowOutputProcessorSuite`, `PythonArrowOutputSuite`,
`ArrowColumnarPythonUDFSuite`, `ArrowCachedPythonUDFSuite`,
`ArrowCachedBatchSerializerSuite`, `ColumnarRulesSuite`, `SparkPlanSuite`,
`PythonUDFSuite`, `ExtractPythonUDFsSuite`, `InMemoryColumnarQuerySuite`,
`ColumnarBatchSuite` and `ConstantColumnVectorSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac (Claude Opus 5.5)
Closes #59581 from viirya/SPARK-60306.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
docs/sql-migration-guide.md | 1 +
.../vectorized/ReusedArrowColumnVector.java | 45 ++++++++
.../sql/execution/ColumnarEvaluatorFactory.scala | 7 +-
.../sql/execution/python/ArrowEvalPythonExec.scala | 25 +----
.../ColumnarArrowEvalPythonEvaluatorFactory.scala | 55 ++++++++-
.../sql/execution/python/PythonArrowOutput.scala | 16 ++-
.../execution/python/ArrowBackedDataSourceV2.scala | 124 ++++++++++++++++++---
.../python/ArrowColumnarPythonUDFSuite.scala | 123 +++++++++++++++++++-
.../python/ArrowOutputProcessorSuite.scala | 111 ++++++++++++++++++
9 files changed, 461 insertions(+), 46 deletions(-)
diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md
index f385cfac3696..8d62c4a2dfac 100644
--- a/docs/sql-migration-guide.md
+++ b/docs/sql-migration-guide.md
@@ -30,6 +30,7 @@ license: |
- Since Spark 4.4, quoted Variant path keys decode `\\`, `\/`, `\b`, `\f`,
`\n`, `\r`, `\t`, `\uXXXX` (including valid surrogate pairs), and escapes of
the enclosing quote (`\'` or `\"`). Previously, backslashes were literal, so
existing paths can match different keys or become invalid. For example, use
`$['a\\']` instead of `$['a\']`. Unknown escapes and unpaired surrogate escapes
stay literal. Dot-notation keys are unchanged. To keep the old behavior, set
`spark.sql.legacy.variantPathB [...]
- Since Spark 4.4, `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`
is `false` by default: Arrow-based Python UDFs read rows from an Arrow-backed
upstream operator, such as a Data Source V2 scan or the Arrow cache, as before
Spark 4.2. With columnar input, they kept the columns that they pass through
until the results of their batch came back from Python, and failed or returned
wrong results when the operator reused or released the vectors of the batch in
the meantime. To enab [...]
+- Since Spark 4.4, the conversion of columnar batches to rows releases each
batch with `ColumnVector.closeIfFreeable` once its rows are consumed also when
it is not whole-stage code-generated, e.g. when the query has more than
`spark.sql.codegen.maxFields` fields or `spark.sql.codegen.wholeStage` is
`false`, as the generated code has done since Spark 3.5.4 and 4.0. A
`ColumnVector` that a Data Source V2 source reuses across batches, such as an
`ArrowColumnVector` over a `VectorSchemaRoot [...]
- Since Spark 4.4, `spark.sql.charVarchar.standardSemantics.enabled` (default
`false`) is a preview feature that may change in future releases. When true, it
enables first-class `CHAR(n)` and `VARCHAR(n)` semantics throughout schemas and
query plans: `CHAR` -> `VARCHAR` -> `STRING` form a type-precedence family,
explicit casts can return constrained types, least-common-type resolution for
`COALESCE`/`CASE`/`UNION` can return `CHAR` or `VARCHAR`, and string
expressions that transform thei [...]
- Since Spark 4.4, when `spark.sql.charVarchar.standardSemantics.enabled` is
true, `from_json` length-checks JSON object names used as `MAP<CHAR(n), _>` or
`MAP<VARCHAR(n), _>` keys without padding or trimming them. A `CHAR(n)` key
must already be exactly `n` characters, and a `VARCHAR(n)` key must already be
at most `n` characters; duplicate names are kept, exactly as for STRING keys,
and `spark.sql.mapKeyDedupPolicy` is not applied. A key that fails the check
makes the enclosing map a [...]
- Since Spark 4.4, `schema_of_json`, `schema_of_csv`, and `schema_of_xml`
accept CHAR/VARCHAR document arguments when CHAR/VARCHAR are first-class. In
Spark 4.3, with `spark.sql.preserveCharVarcharTypeInfo` on, they rejected those
arguments because the child type was not STRING.
diff --git
a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java
b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java
new file mode 100644
index 000000000000..5224bb1ef86c
--- /dev/null
+++
b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java
@@ -0,0 +1,45 @@
+/*
+ * 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()}. This covers the vector itself: the vectors that {@link
#getChild(int)}
+ * returns for the fields of a struct are plain {@link ArrowColumnVector}s,
which consumers do not
+ * release on their own.
+ */
+public class ReusedArrowColumnVector extends ArrowColumnVector {
+
+ public ReusedArrowColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ @Override
+ public void closeIfFreeable() {
+ // no-op: the producer owns the vector and frees it in `close`.
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala
index 2ed94c36436f..debc2a6f3e06 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala
@@ -27,6 +27,7 @@ import
org.apache.spark.sql.execution.vectorized.{OffHeapColumnVector, OnHeapCol
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.CompletionIterator
class ColumnarToRowEvaluatorFactory(
childOutput: Seq[Attribute],
@@ -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]](
+ input.rowIterator().asScala.map(toUnsafe), input.closeIfFreeable())
}
}
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
index ddfa84466272..19b2b6a7ed3f 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
+++
b/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
+ // batch once its rows are consumed. Otherwise it reads rows from its child,
+ // even from one that also supports columnar, such as the scan of the Arrow
+ // cache. Either way, the choice is fixed when the plan is built.
override def supportsColumnar: Boolean =
child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled
- override def supportsRowBased: Boolean = true
-
- override protected def doExecute(): RDD[InternalRow] = {
- // Checks the configuration too, through `supportsColumnar`: a child that
also returns rows,
- // such as the scan of the Arrow cache, gets no ColumnarToRowExec when
columnar input is
- // disabled, and still supports columnar.
- if (supportsColumnar) {
- // 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.
- doExecuteColumnar().mapPartitionsInternal { batchIter =>
- val toUnsafe = UnsafeProjection.create(schema)
- batchIter.flatMap(_.rowIterator().asScala.map(toUnsafe))
- }
- } else {
- // Row-based path: unchanged.
- super.doExecute()
- }
- }
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
val inputRDD = child.executeColumnar()
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
index a0895ec61f4f..ef8a9ac38143 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
@@ -50,7 +50,12 @@ import org.apache.spark.util.Utils
* reuse or release the vectors of a batch once it is asked for the next
* one. So `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`
* is disabled by default. Output is produced by columnar combining:
- * passThruCols ++ resultCols -> ColumnarBatch.
+ * passThruCols ++ resultCols -> ColumnarBatch, so the consumer of the
+ * output releases the pass-through columns with closeIfFreeable on the
+ * task thread. With `spark.python.udf.pipelined.enabled`, the Python
+ * writer reads the next input batches on its own thread meanwhile, so
+ * the upstream operator must allow its vectors to be released while it
+ * returns the next batches.
*
* 2. '''Non-Arrow columnar path''' (UDF inputs are simple column refs
* BUT columns are NOT [[ArrowColumnVector]]): Non-Arrow columnar
@@ -63,6 +68,9 @@ import org.apache.spark.util.Utils
* [[BasicPythonArrowInput]]. Uses [[HybridRowQueue]] for pass-through
* and converts joined rows back to [[ColumnarBatch]].
*
+ * Paths 2 and 3 copy what they need from an input batch, so they release
+ * it with closeIfFreeable once it has been written to Python.
+ *
* TODO: Add a physical plan rule that inserts a ProjectExec before
* ArrowEvalPythonExec to pre-evaluate complex UDF input expressions
* into simple column references. This would eliminate path 3.
@@ -255,12 +263,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 =>
batch.rowIterator().asScala.foreach { row =>
queue.add(unsafeProj(row).copy())
}
- batch
- }
+ true
+ })
val pyRunner =
new ColumnarArrowPythonWithNamedArgumentRunner(
pyFuncs, evalType, argMetas, udfInputSchema,
@@ -274,7 +282,9 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val projection = MutableProjection.create(
allInputs, childOutput)
projection.initialize(context.partitionId())
- val batchIter = inputIter.map { batch =>
+ // The rows of a batch are written to Python before the writer asks
+ // for the next batch, which releases it.
+ val batchIter = new WrittenBatchReleasingIterator(inputIter, _ =>
true).map { batch =>
batch.rowIterator().asScala.map { row =>
queue.add(unsafeProj(row).copy())
projection(row)
@@ -343,6 +353,41 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
}
}
+/**
+ * Hands the input batches to the Python writer. `enqueue` is called on each
+ * batch before the writer gets it, and returns whether the batch should be
+ * released with closeIfFreeable once it has been written: the writer is done
+ * with a batch when it asks for the next one or reaches the end.
+ */
+private[python] class WrittenBatchReleasingIterator(
+ input: Iterator[ColumnarBatch],
+ enqueue: ColumnarBatch => Boolean)
+ extends Iterator[ColumnarBatch] {
+
+ private var written: ColumnarBatch = _
+
+ private def releaseWritten(): Unit = {
+ if (written != null) {
+ written.closeIfFreeable()
+ written = null
+ }
+ }
+
+ override def hasNext: Boolean = {
+ releaseWritten()
+ input.hasNext
+ }
+
+ override def next(): ColumnarBatch = {
+ releaseWritten()
+ val batch = input.next()
+ if (enqueue(batch)) {
+ written = batch
+ }
+ batch
+ }
+}
+
/** Iterator wrapper that allows peeking at the first element. */
private[python] class PeekableIterator[T](underlying: Iterator[T])
extends Iterator[T] {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
index a39dbb97825d..6df6e8904985 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
@@ -27,10 +27,11 @@ import org.apache.arrow.vector.ipc.ArrowStreamReader
import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.api.python.{BasePythonRunner, PythonWorker,
PythonWorkerHandle, SpecialLengths}
import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.util.ArrowUtils
-import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
/**
* A trait that can be mixed-in with [[BasePythonRunner]]. It implements the
logic from
@@ -170,12 +171,18 @@ trait ArrowOutputProcessor {
def close(): Unit
}
+/**
+ * Loads every batch from Python into the reader's root and returns batches
over the same vectors,
+ * so the vectors of a batch are only valid until the next one is loaded. They
are
+ * [[ReusedArrowColumnVector]]s: a consumer does not hold them past the next
batch, and its
+ * `closeIfFreeable` leaves them to this processor, which closes them in
`close`.
+ */
class ArrowOutputProcessorImpl(reader: ArrowStreamReader, pythonMetrics:
Map[String, SQLMetric])
extends ArrowOutputProcessor {
protected val root = reader.getVectorSchemaRoot()
protected val schema: StructType =
ArrowUtils.fromArrowSchema(root.getSchema())
private val vectors: Array[ColumnVector] =
root.getFieldVectors().asScala.map { vector =>
- new ArrowColumnVector(vector)
+ new ReusedArrowColumnVector(vector)
}.toArray[ColumnVector]
protected var rowCount = -1
@@ -235,12 +242,15 @@ abstract class BaseSliceArrowOutputProcessor(
batch
}
+ // The vectors are either those of the reader's root, into which the next
batch is loaded, or
+ // those of a slice of it, which is closed once the next one is produced, so
both are reused
+ // vectors.
protected override def getVectors(root: VectorSchemaRoot):
Array[ColumnVector] = {
if (root eq this.root) {
super.getVectors(root)
} else {
root.getFieldVectors.asScala.map { vector =>
- new ArrowColumnVector(vector)
+ new ReusedArrowColumnVector(vector)
}.toArray[ColumnVector]
}
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala
index da6a423a39bc..69139c892ced 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala
@@ -18,12 +18,14 @@ package org.apache.spark.sql.execution.python
import java.nio.charset.StandardCharsets.UTF_8
import java.util
+import java.util.concurrent.atomic.AtomicInteger
import scala.jdk.CollectionConverters._
import org.apache.arrow.memory.BufferAllocator
import org.apache.arrow.vector._
+import org.apache.spark.TaskContext
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.connector.catalog.{SupportsRead, Table,
TableCapability, TableProvider}
import org.apache.spark.sql.connector.expressions.Transform
@@ -31,7 +33,8 @@ import org.apache.spark.sql.connector.read._
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.ArrowUtils
import org.apache.spark.sql.util.CaseInsensitiveStringMap
-import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarArray,
ColumnarBatch, ColumnarMap, ColumnVector}
+import org.apache.spark.unsafe.types.UTF8String
/**
* A test-only DataSource V2 implementation that produces Arrow-backed
@@ -77,6 +80,19 @@ object ArrowBackedDataSourceV2 {
val DEFAULT_NUM_ROWS: Int = 10000
val DEFAULT_NUM_PARTITIONS: Int = 2
val BATCH_SIZE: Int = 1024
+
+ // The batches that ArrowBackedPartitionReaders returned, and the ones among
them that their
+ // consumer released with closeIfFreeable.
+ val returnedBatches = new AtomicInteger
+ val releasedBatches = new AtomicInteger
+ // The most batches of one reader returned and not released yet at a time.
+ val maxUnreleasedBatches = new AtomicInteger
+
+ def resetBatchCounters(): Unit = {
+ returnedBatches.set(0)
+ releasedBatches.set(0)
+ maxUnreleasedBatches.set(0)
+ }
}
/** Table that supports batch read with columnar output. */
@@ -126,9 +142,11 @@ private class ArrowBackedScanBuilder(options:
CaseInsensitiveStringMap)
options.getBoolean("reuseVectors", false)
private val batchSize =
options.getInt("batchSize", ArrowBackedDataSourceV2.BATCH_SIZE)
+ private val arrowIdColumn =
+ options.getBoolean("arrowIdColumn", true)
override def createReaderFactory(): PartitionReaderFactory = {
- new ArrowBackedReaderFactory(columnar, reuseVectors, batchSize)
+ new ArrowBackedReaderFactory(columnar, reuseVectors, batchSize,
arrowIdColumn)
}
}
@@ -143,11 +161,14 @@ private case class ArrowPartition(start: Int, end: Int)
* InternalRow output by iterating the ColumnarBatch row by row.
* This allows benchmarking the same data source with and without
* the columnar optimization.
+ * @param arrowIdColumn if false, the `id` column of the batches is not an
+ * [[ArrowColumnVector]], so that a batch is not entirely Arrow-backed.
*/
private class ArrowBackedReaderFactory(
columnar: Boolean,
reuseVectors: Boolean,
- batchSize: Int)
+ batchSize: Int,
+ arrowIdColumn: Boolean)
extends PartitionReaderFactory {
override def supportColumnarReads(
@@ -156,7 +177,7 @@ private class ArrowBackedReaderFactory(
override def createReader(
partition: InputPartition): PartitionReader[InternalRow] = {
val ArrowPartition(start, end) = partition
- val inner = new ArrowBackedPartitionReader(start, end, batchSize)
+ val inner = new ArrowBackedPartitionReader(start, end, batchSize,
arrowIdColumn)
// Wrap columnar reader as row reader.
new PartitionReader[InternalRow] {
private var rowIter: java.util.Iterator[InternalRow] = _
@@ -177,7 +198,7 @@ private class ArrowBackedReaderFactory(
if (reuseVectors) {
new ReusingArrowPartitionReader(start, end, batchSize)
} else {
- new ArrowBackedPartitionReader(start, end, batchSize)
+ new ArrowBackedPartitionReader(start, end, batchSize, arrowIdColumn)
}
}
}
@@ -232,13 +253,26 @@ 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:
* - id (INT): sequential integer starting from partition start
* - name (STRING): "row_<id>"
* - value (DOUBLE): id * 0.1
+ * - data (STRING): 100 random characters
+ *
+ * Every batch has vectors of its own, which are not reused vectors: the
consumer may
+ * hold them past the next batch, even past the close of the reader, and
releases them with
+ * closeIfFreeable. The ones it does not release are closed when the task
completes. Their
+ * release is counted in [[ArrowBackedDataSourceV2]].
+ *
+ * If `arrowIdColumn` is false, the `id` column is a [[ColumnVector]] over its
Arrow vector that is
+ * not an [[ArrowColumnVector]].
*/
-private class ArrowBackedPartitionReader(start: Int, end: Int, batchSize: Int)
- extends PartitionReader[ColumnarBatch] {
+private class ArrowBackedPartitionReader(
+ start: Int,
+ end: Int,
+ batchSize: Int,
+ arrowIdColumn: Boolean = true)
+ extends PartitionReader[ColumnarBatch] {
private var current = start
private var currentBatch: ColumnarBatch = _
@@ -246,14 +280,67 @@ 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
+
+ // The consumer may still hold the returned batches, even past the close of
the reader, so they
+ // are closed when the task completes. A pipelined Python writer that reads
them is stopped by a
+ // pre-completion listener, which runs before this one.
+ TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+ allBatches.synchronized {
+ allBatches.forEach(_.close())
+ allBatches.clear()
+ }
+ allocator.close()
+ }
+
+ private def countRelease(): Unit = allBatches.synchronized {
+ numReleased += 1
+ ArrowBackedDataSourceV2.releasedBatches.incrementAndGet()
+ }
+
+ /** The first column of a batch, which counts the releases of the batch by
its consumer. */
+ private class ReleaseCountingArrowColumnVector(vector: ValueVector)
+ extends ArrowColumnVector(vector) {
+ override def closeIfFreeable(): Unit = {
+ countRelease()
+ super.closeIfFreeable()
+ }
+ }
+
+ /** Like [[ReleaseCountingArrowColumnVector]], but not an
[[ArrowColumnVector]]. */
+ private class ReleaseCountingColumnVector(vector: ArrowColumnVector)
+ extends ColumnVector(vector.dataType()) {
+ override def closeIfFreeable(): Unit = {
+ countRelease()
+ close()
+ }
+ override def close(): Unit = vector.close()
+ override def hasNull(): Boolean = vector.hasNull()
+ override def numNulls(): Int = vector.numNulls()
+ override def isNullAt(rowId: Int): Boolean = vector.isNullAt(rowId)
+ override def getBoolean(rowId: Int): Boolean = vector.getBoolean(rowId)
+ override def getByte(rowId: Int): Byte = vector.getByte(rowId)
+ override def getShort(rowId: Int): Short = vector.getShort(rowId)
+ override def getInt(rowId: Int): Int = vector.getInt(rowId)
+ override def getLong(rowId: Int): Long = vector.getLong(rowId)
+ override def getFloat(rowId: Int): Float = vector.getFloat(rowId)
+ override def getDouble(rowId: Int): Double = vector.getDouble(rowId)
+ override def getArray(rowId: Int): ColumnarArray = vector.getArray(rowId)
+ override def getMap(ordinal: Int): ColumnarMap = vector.getMap(ordinal)
+ override def getDecimal(rowId: Int, precision: Int, scale: Int): Decimal =
+ vector.getDecimal(rowId, precision, scale)
+ override def getUTF8String(rowId: Int): UTF8String =
vector.getUTF8String(rowId)
+ override def getBinary(rowId: Int): Array[Byte] = vector.getBinary(rowId)
+ override def getChild(ordinal: Int): ColumnVector =
vector.getChild(ordinal)
+ }
override def next(): Boolean = {
if (current >= end) return false
@@ -269,24 +356,31 @@ private class ArrowBackedPartitionReader(start: Int, end:
Int, batchSize: Int)
// Wrap in ArrowColumnVector
val columns: Array[ColumnVector] = Array(
- new ArrowColumnVector(idVector),
+ if (arrowIdColumn) {
+ new ReleaseCountingArrowColumnVector(idVector)
+ } else {
+ new ReleaseCountingColumnVector(new ArrowColumnVector(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, which are closed when
the task completes.
override def close(): Unit = {
- allBatches.forEach(_.close())
- allBatches.clear()
currentBatch = null
- allocator.close()
}
private def createIntVector(
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala
index 06f162a237c7..badfa6a694bd 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala
@@ -19,7 +19,9 @@ package org.apache.spark.sql.execution.python
import org.apache.spark.SparkConf
import org.apache.spark.sql.IntegratedUDFTestUtils
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan,
WholeStageCodegenExec}
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.StringType
@@ -46,16 +48,30 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
numRows: Int = 100,
numPartitions: Int = 1,
reuseVectors: Boolean = false,
- batchSize: Int = ArrowBackedDataSourceV2.BATCH_SIZE) = {
+ batchSize: Int = ArrowBackedDataSourceV2.BATCH_SIZE,
+ arrowIdColumn: Boolean = true) = {
spark.read
.format(arrowSource)
.option("numRows", numRows.toString)
.option("numPartitions", numPartitions.toString)
.option("reuseVectors", reuseVectors.toString)
.option("batchSize", batchSize.toString)
+ .option("arrowIdColumn", arrowIdColumn.toString)
.load()
}
+ /**
+ * Whether `exec` takes the Arrow columnar path, which outputs its input
columns as they are,
+ * while the other paths output on-heap vectors.
+ */
+ private def takesArrowColumnarPath(exec: ArrowEvalPythonExec): Boolean = {
+ val columnClasses = exec.executeColumnar().mapPartitionsInternal { batches
=>
+ batches.map(_.column(0).getClass.getName)
+ }.collect().toSet
+ assert(columnClasses.size == 1, s"Expected one path for all batches:
$columnClasses")
+ columnClasses.head != classOf[OnHeapColumnVector].getName
+ }
+
private def collectNodes[T <: SparkPlan](
plan: SparkPlan)(implicit tag: reflect.ClassTag[T]): Seq[T] = {
plan.collect {
@@ -114,6 +130,109 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60306: 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("SPARK-60306: 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
of the Arrow-backed
+ // source take the Arrow columnar path, whose output batches carry the
input columns to the
+ // ColumnarToRowExec above ArrowEvalPythonExec, which releases them. If
the `id` column is not
+ // an ArrowColumnVector, they take the row queue path, and on an
expression, the row fallback
+ // path. Both release each batch once it has been written to Python.
+ for {
+ codegen <- Seq(true, false)
+ (input, suffix, arrowIdColumn) <- Seq(
+ ("name", "", true), ("name", "", false), ("concat(name, '_')", "_",
true))
+ } withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString) {
+ val result = readArrowSource(
+ numRows = numRows, batchSize = numRows / numBatches, arrowIdColumn
= arrowIdColumn)
+ .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")
+ val arrowColumnarPath = suffix.isEmpty && arrowIdColumn
+ assert(takesArrowColumnarPath(arrowExec.head) == arrowColumnarPath)
+
+ 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 queue
and row fallback
+ // paths release a batch before they read 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 (!arrowColumnarPath) {
+ 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("SPARK-60306: A plan built with columnar input runs as built after the
configuration " +
+ "changes") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_replan_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // The physical plan is built with columnar input enabled, over the
columnar-only scan, and
+ // first executed after columnar input is disabled.
+ val numRows = 100
+ val result = readArrowSource(numRows = numRows)
+ .selectExpr("id", "name", "value", "data", "arrow_replan_udf(name) as
udf_name")
+ val plan = result.queryExecution.executedPlan
+
assert(collectNodes[ArrowEvalPythonExec](plan).exists(_.child.supportsColumnar))
+ withSQLConf(SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key ->
"false") {
+ val rows = result.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ assert(rows.forall(row => row.getString(4) == s"row_${row.getInt(0)}"))
+ }
+ }
+ }
+
test("Arrow-backed source that reuses its vectors: correct with the default
configuration") {
assume(shouldTestPandasUDFs)
// The Arrow columnar path keeps the pass-through columns of a batch until
its result comes
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowOutputProcessorSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowOutputProcessorSuite.scala
new file mode 100644
index 000000000000..7d2bb4dd0313
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowOutputProcessorSuite.scala
@@ -0,0 +1,111 @@
+/*
+ * 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.ReusedArrowColumnVector
+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()
+ }
+ }
+
+ // The output processors load 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. The slicing processors are
given limits that
+ // never split a batch, so that they return the reader's root as the
whole-batch processor does.
+ Seq[(String, ArrowStreamReader => ArrowOutputProcessor)](
+ "whole batches" -> (new ArrowOutputProcessorImpl(_, metrics())),
+ "records limit" -> (new SliceRecordsArrowOutputProcessorImpl(_, metrics(),
rowsPerBatch)),
+ "bytes limit" -> (new SliceBytesArrowOutputProcessorImpl(_, metrics(),
Int.MaxValue))
+ ).foreach { case (name, newProcessor) =>
+ test(s"SPARK-60306: Struct output columns stay valid after a consumer
releases each " +
+ s"batch: $name") {
+ val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator("arrow-output-reader", 0,
Long.MaxValue)
+ val reader = new ArrowStreamReader(new
ByteArrayInputStream(arrowStream()), allocator)
+ val processor = newProcessor(reader)
+ try {
+ val values = ArrayBuffer[(Int, String)]()
+ while (processor.loadBatch()) {
+ val batch = processor.produceBatch()
+ val column = batch.column(0)
+ assert(column.isInstanceOf[ReusedArrowColumnVector])
+ (0 until batch.numRows()).foreach { i =>
+ val struct = column.getStruct(i)
+ values += ((struct.getInt(0), struct.getUTF8String(1).toString))
+ }
+ batch.closeIfFreeable()
+ }
+ assert(values.toSeq == (0 until numBatches * rowsPerBatch).map(v =>
(v, s"v$v")))
+ } finally {
+ processor.close()
+ reader.close()
+ allocator.close()
+ }
+ }
+ }
+
+ private def metrics(): Map[String, SQLMetric] = Map(
+ "pythonNumRowsReceived" -> new SQLMetric("sum"),
+ "pythonDataReceived" -> new SQLMetric("size"))
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
