From bc2449ac6ed7cb3f3d2150565d535b6506f8aaf1 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 17 Oct 2026 01:42:49 -0700 Subject: [PATCH 1/3] Add a marker for column vectors that their producer may release or overwrite There is no agreement between the producer and the consumer of a ColumnarBatch, so the type of its vectors is the protocol, as for closeIfFreeable (SPARK-50235). Add ReusedColumnVector, an internal marker with no methods: a vector carrying it may be released or overwritten by its producer after it is returned, so a consumer must not hold it past the next batch, and its closeIfFreeable is a no-op. A vector without it may be held, and the consumer releases it with closeIfFreeable. WritableColumnVector and ConstantColumnVector, whose closeIfFreeable is already a no-op, implement the marker: readers reuse them across batches and free them when they are closed. ReusedArrowColumnVector is an ArrowColumnVector carrying it, with a no-op closeIfFreeable, for producers that reuse Arrow vectors. The child vectors of a struct, array or map column share the lifetime of the column, so the marker is only checked on the top-level columns of a batch. The Arrow output processors of the Python runners load every batch into the same root, or close the slices of a batch once the next one is produced, so their columns are now ReusedArrowColumnVectors. This also makes closeIfFreeable on them safe: closing the shared ArrowColumnVectors of the root would clear their child columns for the next batches. Co-authored-by: Isaac --- .../vectorized/ConstantColumnVector.java | 5 ++- .../vectorized/ReusedArrowColumnVector.java | 40 +++++++++++++++++++ .../vectorized/ReusedColumnVector.java | 37 +++++++++++++++++ .../vectorized/WritableColumnVector.java | 4 +- .../execution/python/PythonArrowOutput.scala | 14 +++++-- 5 files changed, 94 insertions(+), 6 deletions(-) create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ConstantColumnVector.java b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ConstantColumnVector.java index a579abf2565eb..3da05102fa7cd 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ConstantColumnVector.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ConstantColumnVector.java @@ -34,8 +34,11 @@ * which will set the exact same value to all rows. * * Capacity: The vector stores only one copy of the data. + * + * Readers reuse a ConstantColumnVector across batches and free it when they are closed, so it is a + * {@link ReusedColumnVector}. */ -public class ConstantColumnVector extends ColumnVector { +public class ConstantColumnVector extends ColumnVector implements ReusedColumnVector { // The data stored in this ConstantColumnVector, the vector stores only one copy of the data. private byte nullData; 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 0000000000000..0ca87a1e01e53 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java @@ -0,0 +1,40 @@ +/* + * 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. See {@link ReusedColumnVector}. + * + *

The producer releases the vector, so {@link #closeIfFreeable()} is a no-op, and the + * producer frees it with {@link #close()}. + */ +public class ReusedArrowColumnVector extends ArrowColumnVector implements ReusedColumnVector { + + 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/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java new file mode 100644 index 0000000000000..8b37eb90c26dc --- /dev/null +++ b/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. + * + *

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. + * + *

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 { +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java index fab67965d701a..41344f073dfaa 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java @@ -50,9 +50,9 @@ * A WritableColumnVector should be considered immutable once originally created. In other words, * it is not valid to call put APIs after reads until reset() is called. * - * WritableColumnVector are intended to be reused. + * WritableColumnVector are intended to be reused, so they are {@link ReusedColumnVector}s. */ -public abstract class WritableColumnVector extends ColumnVector { +public abstract class WritableColumnVector extends ColumnVector implements ReusedColumnVector { private final byte[] byte8 = new byte[8]; // Lookup table that "expands" each of the 8 bits in a byte into a separate byte within a long. 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 1cf2664f67b35..5c9bb367822ad 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 @@ -167,12 +168,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 @@ -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. protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = { root.getFieldVectors.asScala.map { vector => - new ArrowColumnVector(vector) + new ReusedArrowColumnVector(vector) }.toArray[ColumnVector] } From dcdb133fea9303ba70a9c91cf2b6fcc6868e0e65 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 17 Oct 2026 01:43:32 -0700 Subject: [PATCH 2/3] [SPARK-60306] Release the input batches of Python UDFs and non-codegen ColumnarToRow The consumer of a batch releases its vectors with closeIfFreeable once it is done with them (SPARK-50235), which the generated code of ColumnarToRowExec does in advanceBatch. Some consumers never did, so the fresh vectors of every batch stayed live until the end of the task: - the non-codegen ColumnarToRowEvaluatorFactory; - ArrowEvalPythonExec.doExecute, which flattens the output batches of the columnar evaluator for a row-based parent; - the row queue and row fallback paths of the columnar evaluator, which copy the rows and write the input batches to Python. Add ColumnarToRowIterator, which releases a batch once all its rows have been returned and before the next one is fetched, as advanceBatch does, and use it in both row conversions. The rows are projected into a buffer of their own, so the last row of a batch stays valid. The result columns of the Python runners are reused vectors, so releasing an output batch leaves them to the runner. The row paths release each input batch once the Python writer asks for the next one. The test data source no longer frees the returned batches when its reader is closed, which DataSourceRDD does as soon as the reader is exhausted, but when the task completes, and counts the batches its consumer releases before then. Co-authored-by: Isaac --- .../execution/ColumnarEvaluatorFactory.scala | 48 +++++++++-- .../python/ArrowEvalPythonExec.scala | 9 +- ...umnarArrowEvalPythonEvaluatorFactory.scala | 48 ++++++++++- .../python/ArrowBackedDataSourceV2.scala | 83 +++++++++++++++---- .../python/ArrowColumnarPythonUDFSuite.scala | 30 ++++++- 5 files changed, 189 insertions(+), 29 deletions(-) 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 2ed94c36436ff..3480870e96ecb 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 @@ -17,8 +17,6 @@ package org.apache.spark.sql.execution -import scala.jdk.CollectionConverters._ - import org.apache.spark.{PartitionEvaluator, PartitionEvaluatorFactory, TaskContext} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeProjection} @@ -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( + input: Iterator[ColumnarBatch], + toUnsafe: UnsafeProjection, + onNextBatch: ColumnarBatch => Unit = _ => ()) + extends Iterator[InternalRow] { + + private var batch: ColumnarBatch = _ + private var rows: java.util.Iterator[InternalRow] = java.util.Collections.emptyIterator() + + override def hasNext: Boolean = { + while (!rows.hasNext) { + if (batch != null) { + batch.closeIfFreeable() + batch = null + } + if (!input.hasNext) { + return false } + batch = input.next() + onNextBatch(batch) + rows = batch.rowIterator() + } + true + } + + override def next(): InternalRow = { + if (!hasNext) { + throw new NoSuchElementException("End of the columnar batches") } + toUnsafe(rows.next()) } } 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 0170d4354a6a9..e3b30c7a5a643 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 @@ -25,7 +25,7 @@ import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.{ColumnarToRowIterator, SparkPlan} import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata import org.apache.spark.sql.types.{StructType, UserDefinedType} @@ -110,10 +110,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)) } } else { // Row-based path: unchanged. 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 3981602875adb..2e3ec55a99efe 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 @@ -60,6 +60,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. @@ -251,12 +254,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, @@ -270,7 +273,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) @@ -339,6 +344,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/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 992182871ac90..978d3cf772a63 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 @@ -17,10 +17,12 @@ package org.apache.spark.sql.execution.python import java.util +import java.util.concurrent.atomic.AtomicInteger 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 @@ -74,6 +76,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 before the end of the task. + 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. */ @@ -120,8 +135,11 @@ private class ArrowBackedScanBuilder(options: CaseInsensitiveStringMap) private val columnar = options.getBoolean("columnar", true) + private val batchSize = + options.getInt("batchSize", ArrowBackedDataSourceV2.BATCH_SIZE) + override def createReaderFactory(): PartitionReaderFactory = { - new ArrowBackedReaderFactory(columnar) + new ArrowBackedReaderFactory(columnar, batchSize) } } @@ -137,7 +155,7 @@ private case class ArrowPartition(start: Int, end: Int) * This allows benchmarking the same data source with and without * the columnar optimization. */ -private class ArrowBackedReaderFactory(columnar: Boolean) +private class ArrowBackedReaderFactory(columnar: Boolean, batchSize: Int) extends PartitionReaderFactory { override def supportColumnarReads( @@ -146,7 +164,7 @@ private class ArrowBackedReaderFactory(columnar: Boolean) override def createReader( partition: InputPartition): PartitionReader[InternalRow] = { val ArrowPartition(start, end) = partition - val inner = new ArrowBackedPartitionReader(start, end) + val inner = new ArrowBackedPartitionReader(start, end, batchSize) // Wrap columnar reader as row reader. new PartitionReader[InternalRow] { private var rowIter: java.util.Iterator[InternalRow] = _ @@ -164,36 +182,70 @@ private class ArrowBackedReaderFactory(columnar: Boolean) override def createColumnarReader( partition: InputPartition): PartitionReader[ColumnarBatch] = { val ArrowPartition(start, end) = partition - new ArrowBackedPartitionReader(start, end) + new ArrowBackedPartitionReader(start, end, batchSize) } } /** * 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_" * - value (DOUBLE): id * 0.1 + * + * 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]]. */ -private class ArrowBackedPartitionReader(start: Int, end: Int) +private class ArrowBackedPartitionReader(start: Int, end: Int, batchSize: Int) extends PartitionReader[ColumnarBatch] { - private val batchSize = ArrowBackedDataSourceV2.BATCH_SIZE private var current = start private var currentBatch: ColumnarBatch = _ // 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 = { + allBatches.synchronized { + if (!closed) { + closed = true + if (!taskCompleted) { + numReleased += 1 + ArrowBackedDataSourceV2.releasedBatches.incrementAndGet() + } + } + } + super.close() + } + } override def next(): Boolean = { if (current >= end) return false @@ -209,24 +261,27 @@ private class ArrowBackedPartitionReader(start: Int, end: 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, 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 faf8c77678f3a..7f1ecc0e6707f 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 @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.python import org.apache.spark.sql.IntegratedUDFTestUtils -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.StringType @@ -37,11 +37,15 @@ class ArrowColumnarPythonUDFSuite extends SharedSparkSession { private val arrowSource = "org.apache.spark.sql.execution.python.ArrowBackedDataSourceV2" - private def readArrowSource(numRows: Int = 100, numPartitions: Int = 1) = { + private def readArrowSource( + numRows: Int = 100, + numPartitions: Int = 1, + batchSize: Int = ArrowBackedDataSourceV2.BATCH_SIZE) = { spark.read .format(arrowSource) .option("numRows", numRows.toString) .option("numPartitions", numPartitions.toString) + .option("batchSize", batchSize.toString) .load() } @@ -103,6 +107,28 @@ 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-backed source: multiple UDF columns") { assume(shouldTestPandasUDFs) withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") { From 897f39f1d9034b6e22d20a331c03f094e095f3df Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 17 Oct 2026 01:45:32 -0700 Subject: [PATCH 3/3] [SPARK-60325] Hold only unmarked vectors in the Arrow columnar path of Python UDFs The Arrow columnar path of ColumnarArrowEvalPythonEvaluatorFactory (SPARK-56350) keeps the pass-through columns of a batch until its result comes back from Python, many batches later. This breaks when the producer releases or overwrites the vectors it returned: IndexOutOfBoundsException, or pass-through values of a later batch paired with the results of an earlier one. - Take the path only when every column of the first batch is an ArrowColumnVector that is not a ReusedColumnVector; otherwise take the row queue path. The output batch carries the held columns to its consumer, which releases them with closeIfFreeable. If a later batch has reused vectors, its pass-through columns are copied into on-heap vectors and the batch is released once written. - The queue of pass-through columns is a ConcurrentLinkedQueue: in pipelined mode the writer thread adds to it while the task thread polls it. - ArrowCachedBatchToColumnarBatchIterator already loads every batch into a root of its own, but closed the previous root in next(). It now leaves the release to the consumer, which closes the root once all columns of its batch are closed, and closes the roots left when the task completes. This keeps the Arrow cache on the columnar path, rather than marking its vectors and making cache + UDF take the slower row queue path. The tests cover a reader that reuses its vectors (row queue path, correct), a reader with fresh vectors (columnar path, no copy, each batch released before the end of the task), the Arrow cache with 2000 batches, two stacked ArrowEvalPythonExecs, and pipelined mode. Co-authored-by: Isaac --- .../columnar/ArrowCachedBatchSerializer.scala | 91 ++++++--- ...umnarArrowEvalPythonEvaluatorFactory.scala | 79 +++++--- .../python/ArrowBackedDataSourceV2.scala | 71 ++++++- .../python/ArrowColumnarPythonUDFSuite.scala | 183 +++++++++++++++++- 4 files changed, 371 insertions(+), 53 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala index baba2c1365c82..2a5b06eea4843 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala @@ -1257,6 +1257,12 @@ private class ColumnarBatchToArrowCachedBatchIterator( /** * Iterator that converts ArrowCachedBatch to ColumnarBatch. + * + * Every batch is loaded into its own root, and its columns are plain [[ArrowColumnVector]]s, which + * a consumer may hold past the next batch: the consumer releases a batch with closeIfFreeable once + * it is done with it (see [[org.apache.spark.sql.execution.vectorized.ReusedColumnVector]]), which + * closes its root once all its columns are closed. The roots of the batches that are not released + * that way, e.g. when the consumer stops early, are closed when the task completes. */ private class ArrowCachedBatchToColumnarBatchIterator( batchIter: Iterator[CachedBatch], @@ -1291,14 +1297,16 @@ private class ArrowCachedBatchToColumnarBatchIterator( arrowSchema } - // Track only the previous root to close it when next batch is produced - private var previousRoot: VectorSchemaRoot = null + // The roots of the returned batches that the consumer has not released yet. A consumer may + // release a batch from another thread than the one calling next(), e.g. the writer thread of a + // pipelined Python UDF, so it is guarded by its own lock. + private val unreleasedRoots = java.util.Collections.newSetFromMap( + new java.util.IdentityHashMap[VectorSchemaRoot, java.lang.Boolean]()) // Prefetch support: deserialize the next batch into its own root in a background thread while // the current batch is being consumed. Only the deserialization (IPC read + decompression + - // loading into a fresh root) happens off-thread; closing the previous root stays on the - // consumer thread in next(), so the vectors backing a returned ColumnarBatch are never released - // while the consumer may still read them. + // loading into a fresh root) happens off-thread; the vectors backing a returned ColumnarBatch are + // only released by its consumer, so never while the consumer may still read them. private val prefetchExecutor: ExecutorService = if (prefetchEnabled) { Executors.newSingleThreadExecutor(r => { val t = new Thread(r, "arrow-cache-prefetch") @@ -1310,7 +1318,7 @@ private class ArrowCachedBatchToColumnarBatchIterator( } private var prefetchFuture: Future[VectorSchemaRoot] = _ - // Register cleanup - close remaining root and allocator when task completes + // Register cleanup - close the unreleased roots and the allocator when the task completes Option(TaskContext.get()).foreach { tc => tc.addTaskCompletionListener[Unit] { _ => // Stop the worker and close any root it already produced before closing the allocator. @@ -1319,9 +1327,9 @@ private class ArrowCachedBatchToColumnarBatchIterator( // would then fail with "Memory was leaked by query". prefetchFuture = ArrowCachedBatchSerializer.drainAndClosePrefetch( prefetchExecutor, prefetchFuture) - if (previousRoot != null) { - previousRoot.close() - previousRoot = null + unreleasedRoots.synchronized { + unreleasedRoots.forEach(_.close()) + unreleasedRoots.clear() } allocator.close() } @@ -1330,12 +1338,6 @@ private class ArrowCachedBatchToColumnarBatchIterator( override def hasNext: Boolean = prefetchFuture != null || batchIter.hasNext override def next(): ColumnarBatch = { - // Close the previous root since the consumer has moved on from the batch it backed. - if (previousRoot != null) { - previousRoot.close() - previousRoot = null - } - val root = if (prefetchFuture != null) { val r = try { prefetchFuture.get() @@ -1348,19 +1350,24 @@ private class ArrowCachedBatchToColumnarBatchIterator( deserializeToRoot(batchIter.next().asInstanceOf[ArrowCachedBatch]) } - previousRoot = root - // When projected on load, the root already holds only the selected columns in output order, // so wrap its vectors directly. Otherwise it holds all cache columns and must be selected. - val selectedColumns = if (canProjectOnLoad) { - root.getFieldVectors.asScala.map(v => new ArrowColumnVector(v)).toArray[ColumnVector] + val selectedVectors = if (canProjectOnLoad) { + root.getFieldVectors.asScala.toArray + } else { + columnIndices.map(root.getFieldVectors.get(_)) + } + val numRows = root.getRowCount + val batch = if (selectedVectors.isEmpty) { + // No column can release the root, and it holds no buffer the batch needs. + root.close() + new ColumnarBatch(Array.empty[ColumnVector], numRows) } else { - val allColumns = root.getFieldVectors.asScala.map { vector => - new ArrowColumnVector(vector) - }.toArray[ColumnVector] - columnIndices.map(allColumns(_)) + unreleasedRoots.synchronized(unreleasedRoots.add(root)) + val release = new RootRelease(root, selectedVectors.length) + new ColumnarBatch( + selectedVectors.map(new ReleasingArrowColumnVector(_, release): ColumnVector), numRows) } - val batch = new ColumnarBatch(selectedColumns, root.getRowCount) // Start prefetching the next batch while this one is being consumed. submitPrefetch() @@ -1400,6 +1407,42 @@ private class ArrowCachedBatchToColumnarBatchIterator( } } + /** Closes `root` once all `numColumns` columns of the batch over it have been closed. */ + private class RootRelease(root: VectorSchemaRoot, numColumns: Int) { + private var numOpen = numColumns + + /** + * Closes a column of the batch with `closeColumn`, and the root after the last one. A no-op + * once the root is closed, e.g. by the task completion listener. + */ + def closeColumn(closeColumn: () => Unit): Unit = unreleasedRoots.synchronized { + if (unreleasedRoots.contains(root)) { + closeColumn() + numOpen -= 1 + if (numOpen == 0) { + unreleasedRoots.remove(root) + root.close() + } + } + } + } + + /** A column of a returned batch, which releases the batch's root once all are closed. */ + private class ReleasingArrowColumnVector( + vector: org.apache.arrow.vector.FieldVector, + release: RootRelease) extends ArrowColumnVector(vector) { + private var closed = false + + private def closeVector(): Unit = super.close() + + override def close(): Unit = unreleasedRoots.synchronized { + if (!closed) { + closed = true + release.closeColumn(() => closeVector()) + } + } + } + /** Submit deserialization of the next batch to the background thread, if prefetch is enabled. */ private def submitPrefetch(): Unit = { if (prefetchEnabled && batchIter.hasNext) { 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 2e3ec55a99efe..2efd9a8af9e4c 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 @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.python import java.io.File -import java.util.ArrayDeque +import java.util.concurrent.ConcurrentLinkedQueue import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ @@ -27,11 +27,12 @@ import org.apache.spark.api.python.ChainedPythonFunctions import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.RowToColumnConverter import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata -import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector, ReusedColumnVector, WritableColumnVector} import org.apache.spark.sql.types.{DataType, StructField, StructType, UserDefinedType} import org.apache.spark.sql.types.DataType.equalsIgnoreCompatibleCollation import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} @@ -43,17 +44,21 @@ import org.apache.spark.util.Utils * Three execution paths based on input characteristics: * * 1. '''Arrow columnar path''' (UDF inputs are simple column refs AND - * columns are [[ArrowColumnVector]]): Arrow FieldVectors are extracted - * directly and serialized to IPC. Pass-through columns are kept as - * [[ColumnVector]] references (safe because Arrow vectors are - * independently allocated per batch). Output is produced by columnar - * combining: passThruCols ++ resultCols -> ColumnarBatch. + * columns are [[ArrowColumnVector]]s that are not [[ReusedColumnVector]]s): + * Arrow FieldVectors are extracted directly and serialized to IPC. + * Pass-through columns are kept as [[ColumnVector]] references until the + * result of their batch comes back from Python, which the producer allows + * by not marking them as reused. Output is produced by columnar + * combining: passThruCols ++ resultCols -> ColumnarBatch, so its consumer + * releases the pass-through columns with closeIfFreeable. * * 2. '''Non-Arrow columnar path''' (UDF inputs are simple column refs - * BUT columns are NOT [[ArrowColumnVector]]): Non-Arrow columnar - * readers (Parquet vectorized, InMemoryTableScan) reuse ColumnVector - * objects across batches. Uses [[HybridRowQueue]] for pass-through - * buffering and converts joined rows back to [[ColumnarBatch]]. + * BUT columns are NOT such [[ArrowColumnVector]]s): e.g. Parquet + * vectorized readers reuse [[WritableColumnVector]]s across batches, and + * the Python output of a stacked [[ArrowEvalPythonExec]] is loaded into + * the same vectors for every batch. Uses [[HybridRowQueue]] for + * pass-through buffering and converts joined rows back to + * [[ColumnarBatch]]. * * 3. '''Row fallback path''' (UDF inputs contain complex expressions): * Rows are extracted via [[MutableProjection]] and sent using @@ -89,6 +94,8 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( private class ColumnarArrowEvalPythonPartitionEvaluator extends PartitionEvaluator[ColumnarBatch, ColumnarBatch] { + private lazy val childSchema = DataTypeUtils.fromAttributes(childOutput) + private def collectFunctions( udf: PythonUDF ): ((ChainedPythonFunctions, Long), Seq[Expression]) = { @@ -147,11 +154,17 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( val inputColumnIndices = resolveColumnIndices(allInputs.toSeq) - // Peek at first batch to check if Arrow-backed. + // Peek at the first batch to check whether its columns are Arrow vectors + // that may be held until the result of the batch comes back. val peekIter = new PeekableIterator(inputIter) val isArrow = peekIter.peek().exists { batch => - batch.numCols() > 0 && - batch.column(0).isInstanceOf[ArrowColumnVector] + batch.numCols() > 0 && (0 until batch.numCols()).forall { i => + batch.column(i) match { + case _: ReusedColumnVector => false + case _: ArrowColumnVector => true + case _ => false + } + } } if (inputColumnIndices.isDefined && isArrow) { @@ -178,9 +191,15 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( } /** - * Path 1: Arrow columnar. ColumnVector references are safe - * because Arrow vectors are independently allocated per batch. - * Combines pass-through + UDF result columns into ColumnarBatch. + * Path 1: Arrow columnar. Holding the ColumnVector references is safe + * because the producer does not mark them as reused (see + * [[ReusedColumnVector]]). Combines pass-through + UDF result columns + * into ColumnarBatch. + * + * The path is chosen from the first batch, while a producer could + * return reused vectors in a later batch. The pass-through columns of + * such a batch are copied instead, and the batch is released once it + * has been written to Python. */ private def evalArrowColumnar( inputIter: Iterator[ColumnarBatch], @@ -191,15 +210,22 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( outputTypes: Seq[DataType], columnIndices: Array[Int]): Iterator[ColumnarBatch] = { + // In pipelined mode the batches are added in the writer thread and + // polled in the task thread. val passThruQueue = - new ArrayDeque[(Array[ColumnVector], Int)]() + new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]() - val bufferedIter = inputIter.map { batch => + val bufferedIter = new WrittenBatchReleasingIterator(inputIter, { batch => val passThruCols = childOutput.indices.map( i => batch.column(i)).toArray - passThruQueue.add((passThruCols, batch.numRows())) - batch - } + if (passThruCols.exists(_.isInstanceOf[ReusedColumnVector])) { + passThruQueue.add((copyColumns(batch), batch.numRows())) + true + } else { + passThruQueue.add((passThruCols, batch.numRows())) + false + } + }) val pyRunner = new ColumnarArrowPythonWithNamedArgumentRunner( pyFuncs, evalType, argMetas, udfInputSchema, @@ -305,6 +331,15 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( rowsToColumnarBatches(rowIter, context) } + /** Copies the pass-through columns of `batch` into on-heap vectors. */ + private def copyColumns(batch: ColumnarBatch): Array[ColumnVector] = { + val vectors = OnHeapColumnVector.allocateColumns( + math.max(batch.numRows(), 1), childSchema).toArray[WritableColumnVector] + val converter = new RowToColumnConverter(childSchema) + batch.rowIterator().asScala.foreach(converter.convert(_, vectors)) + vectors.toArray[ColumnVector] + } + private def validateOutputTypes( batch: ColumnarBatch, outputTypes: Seq[DataType]): Unit = { 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 978d3cf772a63..f6b1ffeb339e2 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 @@ -16,9 +16,12 @@ */ 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._ @@ -27,6 +30,7 @@ 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 import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector import org.apache.spark.sql.types._ import org.apache.spark.sql.util.ArrowUtils import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -134,12 +138,13 @@ private class ArrowBackedScanBuilder(options: CaseInsensitiveStringMap) private val columnar = options.getBoolean("columnar", true) - + private val reuseVectors = + options.getBoolean("reuseVectors", false) private val batchSize = options.getInt("batchSize", ArrowBackedDataSourceV2.BATCH_SIZE) override def createReaderFactory(): PartitionReaderFactory = { - new ArrowBackedReaderFactory(columnar, batchSize) + new ArrowBackedReaderFactory(columnar, reuseVectors, batchSize) } } @@ -155,8 +160,11 @@ private case class ArrowPartition(start: Int, end: Int) * This allows benchmarking the same data source with and without * the columnar optimization. */ -private class ArrowBackedReaderFactory(columnar: Boolean, batchSize: Int) - extends PartitionReaderFactory { +private class ArrowBackedReaderFactory( + columnar: Boolean, + reuseVectors: Boolean, + batchSize: Int) + extends PartitionReaderFactory { override def supportColumnarReads( partition: InputPartition): Boolean = columnar @@ -182,7 +190,58 @@ private class ArrowBackedReaderFactory(columnar: Boolean, batchSize: Int) override def createColumnarReader( partition: InputPartition): PartitionReader[ColumnarBatch] = { val ArrowPartition(start, end) = partition - new ArrowBackedPartitionReader(start, end, batchSize) + if (reuseVectors) { + new ReusingArrowPartitionReader(start, end, batchSize) + } else { + new ArrowBackedPartitionReader(start, end, batchSize) + } + } +} + +/** + * Reads the same rows as [[ArrowBackedPartitionReader]], with short values in `data`, but refills + * the same vectors for every batch and closes them once the reader is closed, as columnar readers + * may: a batch is only valid until the next one, so its columns are [[ReusedArrowColumnVector]]s. + */ +private class ReusingArrowPartitionReader(start: Int, end: Int, batchSize: Int) + extends PartitionReader[ColumnarBatch] { + private var current = start + private val allocator = ArrowUtils.rootAllocator.newChildAllocator( + s"reusing-arrow-test-reader-$start", 0, Long.MaxValue) + private val root = VectorSchemaRoot.create( + ArrowUtils.toArrowSchema(ArrowBackedDataSourceV2.SCHEMA, null, + errorOnDuplicatedFieldNames = true, largeVarTypes = false), + allocator) + private val batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ReusedArrowColumnVector(_): ColumnVector).toArray) + + override def next(): Boolean = { + if (current >= end) return false + val count = math.min(batchSize, end - current) + val id = root.getVector("id").asInstanceOf[IntVector] + val name = root.getVector("name").asInstanceOf[VarCharVector] + val value = root.getVector("value").asInstanceOf[Float8Vector] + val data = root.getVector("data").asInstanceOf[VarCharVector] + // Overwrites the values of the previous batch in place. + root.getFieldVectors.forEach(_.reset()) + (0 until count).foreach { i => + val row = current + i + id.setSafe(i, row) + name.setSafe(i, s"row_$row".getBytes(UTF_8)) + value.setSafe(i, row * 0.1) + data.setSafe(i, s"data_$row".getBytes(UTF_8)) + } + root.setRowCount(count) + batch.setNumRows(count) + current += count + true + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = { + root.close() + allocator.close() } } @@ -194,7 +253,7 @@ private class ArrowBackedReaderFactory(columnar: Boolean, batchSize: Int) * - name (STRING): "row_" * - value (DOUBLE): id * 0.1 * - * Every batch has vectors of its own, which are not reused vectors: the consumer may + * Every batch has vectors of its own, which are not [[ReusedArrowColumnVector]]s: 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]]. 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 7f1ecc0e6707f..2691f80962f76 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 @@ -17,9 +17,14 @@ package org.apache.spark.sql.execution.python +import org.apache.spark.SparkConf +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION import org.apache.spark.sql.IntegratedUDFTestUtils import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, WholeStageCodegenExec} -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.execution.columnar.InMemoryRelation +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.StringType @@ -34,21 +39,60 @@ class ArrowColumnarPythonUDFSuite extends SharedSparkSession { import IntegratedUDFTestUtils._ + // The cache serializer is static. The tests of the columnar input path enable it. + override protected def sparkConf: SparkConf = super.sparkConf + .set(StaticSQLConf.SPARK_CACHE_SERIALIZER.key, + "org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer") + .set(SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key, "true") + + // InMemoryRelation caches the serializer in a process-wide field on first use, so reset it to + // pick up the Arrow serializer configured above, and again so as not to leak it to later suites. + override def beforeAll(): Unit = { + super.beforeAll() + InMemoryRelation.clearSerializer() + } + + override def afterAll(): Unit = { + InMemoryRelation.clearSerializer() + super.afterAll() + } + private val arrowSource = "org.apache.spark.sql.execution.python.ArrowBackedDataSourceV2" private def readArrowSource( numRows: Int = 100, numPartitions: Int = 1, + reuseVectors: Boolean = false, batchSize: Int = ArrowBackedDataSourceV2.BATCH_SIZE) = { spark.read .format(arrowSource) .option("numRows", numRows.toString) .option("numPartitions", numPartitions.toString) + .option("reuseVectors", reuseVectors.toString) .option("batchSize", batchSize.toString) .load() } + private def singleArrowEvalPythonExec(plan: SparkPlan): ArrowEvalPythonExec = { + val execs = collectNodes[ArrowEvalPythonExec](plan) + assert(execs.length == 1, s"Expected one ArrowEvalPythonExec in plan:\n$plan") + execs.head + } + + /** + * Whether `exec` takes the Arrow columnar path, which outputs the input columns as they are, + * while the other paths output on-heap vectors. + */ + private def takesArrowColumnarPath(exec: ArrowEvalPythonExec): Boolean = { + assert(exec.supportsColumnar) + 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 { @@ -107,6 +151,137 @@ class ArrowColumnarPythonUDFSuite extends SharedSparkSession { } } + test("Arrow-backed source that reuses its vectors: Arrow columnar path not taken") { + assume(shouldTestPandasUDFs) + withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") { + val pandasUDF = TestScalarPandasUDF( + name = "arrow_reuse_udf", returnType = Some(StringType)) + registerTestUDF(pandasUDF, spark) + + // The Arrow columnar path would keep the pass-through columns of a batch until its result + // comes back from Python. The reader marks its vectors as reused, since it refills them for + // the next batch and closes them once it is closed, as soon as it is exhausted, so the + // batches take the row queue path. With two batches, the reader is closed before their + // results come back. With many, it refills its vectors with the next batches before then. + // The test UDF casts its input to string: on a string column, the cast is removed, so the + // inputs are column references. + Seq(200, 200000).foreach { numRows => + val result = readArrowSource(numRows = numRows, reuseVectors = true, batchSize = 100) + .selectExpr("id", "name", "value", "data", "arrow_reuse_udf(name) as udf_name") + val arrowExec = singleArrowEvalPythonExec(result.queryExecution.executedPlan) + assert(!takesArrowColumnarPath(arrowExec)) + + 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.getString(3) == s"data_$id", s"data mismatch for id $id") + assert(row.getString(4) == s"row_$id", s"UDF result mismatch for id $id") + } + } + } + } + + test("Arrow-backed source with fresh vectors: held without a copy and released once consumed") { + assume(shouldTestPandasUDFs) + withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") { + val pandasUDF = TestScalarPandasUDF( + name = "arrow_fresh_udf", returnType = Some(StringType)) + registerTestUDF(pandasUDF, spark) + + val numRows = 50000 + val numBatches = 500 + val result = readArrowSource(numRows = numRows, batchSize = numRows / numBatches) + .selectExpr("id", "name", "value", "data", "arrow_fresh_udf(name) as udf_name") + // The pass-through columns are output as the reader returned them. + assert(takesArrowColumnarPath( + singleArrowEvalPythonExec(result.queryExecution.executedPlan))) + + 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", s"UDF result mismatch for id $id") + } + // The consumer of the output releases each input batch once its rows are consumed, before + // the task completes. How many are held at a time depends on how far the writer runs + // ahead of the results, which may be all of them for a small input. + assert(ArrowBackedDataSourceV2.returnedBatches.get == numBatches) + assert(ArrowBackedDataSourceV2.releasedBatches.get == numBatches) + } + } + + test("Arrow cache: the pass-through columns match the UDF results") { + assume(shouldTestPandasUDFs) + withSQLConf( + SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true", + SQLConf.COLUMN_BATCH_SIZE.key -> "1000") { + val pandasUDF = TestScalarPandasUDF( + name = "arrow_cache_udf", returnType = Some(StringType)) + registerTestUDF(pandasUDF, spark) + + val numRows = 2000000L + val cached = spark.range(0, numRows, 1, 1).selectExpr("id", "cast(id as string) as s") + .cache() + try { + val result = cached.selectExpr("id", "s", "arrow_cache_udf(s) as u") + // The cache loads every batch into a root of its own and leaves its release to the + // consumer, so the batches take the Arrow columnar path. + val arrowExec = singleArrowEvalPythonExec(result.queryExecution.executedPlan) + assert(arrowExec.child.supportsColumnar) + assert(takesArrowColumnarPath(arrowExec)) + + val counts = result + .selectExpr("count(*)", "count_if(s = cast(id as string) and u = s)") + .head() + assert(counts.getLong(0) == numRows) + assert(counts.getLong(1) == numRows) + } finally { + cached.unpersist(blocking = true) + } + } + } + + test("Stacked Python UDFs over an Arrow-backed source") { + assume(shouldTestPythonUDFs && shouldTestPandasUDFs) + withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") { + // UDFs of different eval types are not chained, so they are evaluated by two stacked + // ArrowEvalPythonExecs. + val innerUDF = TestPythonUDF(name = "arrow_inner_udf", returnType = Some(StringType), + pythonEvalType = PythonEvalType.SQL_ARROW_BATCHED_UDF) + val outerUDF = TestScalarPandasUDF(name = "arrow_outer_udf", returnType = Some(StringType)) + registerTestUDF(innerUDF, spark) + registerTestUDF(outerUDF, spark) + + val numRows = 20000 + val result = readArrowSource(numRows = numRows, batchSize = 100) + .selectExpr("id", "name", "value", "data", "arrow_outer_udf(arrow_inner_udf(name)) as u") + val plan = result.queryExecution.executedPlan + val execs = collectNodes[ArrowEvalPythonExec](plan) + assert(execs.length == 2, s"Expected two ArrowEvalPythonExecs in plan:\n$plan") + val outerExec = execs.find(_.child.isInstanceOf[ArrowEvalPythonExec]).getOrElse( + fail(s"Expected stacked ArrowEvalPythonExecs in plan:\n$plan")) + // The inner one holds the pass-through columns of the reader, while the results of Python + // are loaded into the same vectors for every batch, so the outer one does not hold them. + assert(takesArrowColumnarPath(outerExec.child.asInstanceOf[ArrowEvalPythonExec])) + assert(!takesArrowColumnarPath(outerExec)) + + 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.getString(3).length == 100, s"data mismatch for id $id") + assert(row.getString(4) == s"row_$id", s"UDF result mismatch for id $id") + } + } + } + test("ColumnarToRow releases each batch once its rows are consumed") { Seq(true, false).foreach { codegen => withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString) { @@ -205,3 +380,9 @@ class ArrowColumnarPythonUDFSuite extends SharedSparkSession { } } } + +/** Runs the tests of [[ArrowColumnarPythonUDFSuite]] with pipelined Python UDF execution. */ +class ArrowColumnarPythonUDFPipelinedSuite extends ArrowColumnarPythonUDFSuite { + override protected def sparkConf: SparkConf = + super.sparkConf.set(PYTHON_UDF_PIPELINED_EXECUTION.key, "true") +}