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/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/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/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..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,23 +44,30 @@ 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 * [[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. @@ -86,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]) = { @@ -144,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) { @@ -175,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], @@ -188,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, @@ -251,12 +280,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 +299,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) @@ -300,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 = { @@ -339,6 +379,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 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] } 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..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,15 +16,21 @@ */ 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 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 @@ -74,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 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. */ @@ -119,9 +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) + new ArrowBackedReaderFactory(columnar, reuseVectors, batchSize) } } @@ -137,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) - extends PartitionReaderFactory { +private class ArrowBackedReaderFactory( + columnar: Boolean, + reuseVectors: Boolean, + batchSize: Int) + extends PartitionReaderFactory { override def supportColumnarReads( partition: InputPartition): Boolean = columnar @@ -146,7 +172,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 +190,121 @@ private class ArrowBackedReaderFactory(columnar: Boolean) override def createColumnarReader( partition: InputPartition): PartitionReader[ColumnarBatch] = { val ArrowPartition(start, end) = partition - new ArrowBackedPartitionReader(start, end) + 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() } } /** * 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 [[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]]. */ -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 +320,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..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.SparkPlan -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, WholeStageCodegenExec} +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,17 +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) = { + 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 { @@ -103,6 +151,159 @@ 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) { + 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") { @@ -179,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") +}