From 9c4563dfe57183cb50e1bacac5258dddff0f6ab0 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 19 Oct 2026 15:08:46 -0700 Subject: [PATCH 1/4] [SPARK-60367][SQL] Fix duplicated rows and closed reader root in sliced Arrow output processors SliceRecordsArrowOutputProcessorImpl returned the whole reader root for the last chunk of a batch it had already sliced, so the rows already output were output again. BaseSliceArrowOutputProcessor also closed the previous output root when a different root was produced, including the reader root itself, which is reused for the next input batches; the next read then failed. Output the whole root only when no row of the batch has been output yet, and never close the reader root when releasing the previous output batch. Co-authored-by: Isaac --- .../execution/python/PythonArrowOutput.scala | 25 +-- .../python/PythonArrowOutputSuite.scala | 159 ++++++++++++++++++ 2 files changed, 174 insertions(+), 10 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala 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..859427043b356 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 @@ -227,10 +227,7 @@ abstract class BaseSliceArrowOutputProcessor( val batchRoot = getRoot if (batchRoot != prevRoot) { - if (prevRoot != null) { - prevVectors.foreach(_.close()) - prevRoot.close() - } + closePrevSlice() prevRoot = batchRoot } @@ -248,11 +245,17 @@ abstract class BaseSliceArrowOutputProcessor( }.toArray[ColumnVector] } - override def close(): Unit = { - if (prevRoot != null) { + // Releases the previous output batch if it is a slice. The reader root is reused for every + // input batch, so it must not be closed here; `ArrowOutputProcessorImpl.close` closes it. + private def closePrevSlice(): Unit = { + if (prevRoot != null && (prevRoot ne root)) { prevVectors.foreach(_.close()) prevRoot.close() } + } + + override def close(): Unit = { + closePrevSlice() super.close() } } @@ -264,7 +267,8 @@ class SliceBytesArrowOutputProcessorImpl( extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) { protected override def getRoot: VectorSchemaRoot = { - if (getBatchBytes(root) < arrowMaxBytesPerOutputBatch) { + // Output the whole root only if no row of this batch has been output yet. + if (currentRowIdx == 0 && getBatchBytes(root) < arrowMaxBytesPerOutputBatch) { currentRowIdx = currentRowIdx + root.getRowCount root } else { @@ -314,10 +318,11 @@ class SliceRecordsArrowOutputProcessorImpl( protected override def getRoot: VectorSchemaRoot = { val remainingRows = rowCount - currentRowIdx - val rootSlice = if (remainingRows > arrowMaxRecordsPerOutputBatch) { - root.slice(currentRowIdx, arrowMaxRecordsPerOutputBatch) - } else { + // Output the whole root only if no row of this batch has been output yet. + val rootSlice = if (currentRowIdx == 0 && remainingRows <= arrowMaxRecordsPerOutputBatch) { root + } else { + root.slice(currentRowIdx, math.min(remainingRows, arrowMaxRecordsPerOutputBatch)) } currentRowIdx = currentRowIdx + rootSlice.getRowCount diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala new file mode 100644 index 0000000000000..d3e92adea25de --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala @@ -0,0 +1,159 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.unsafe.types.UTF8String + +class PythonArrowOutputSuite extends SparkFunSuite { + + private val schema = new StructType() + .add("id", LongType) + .add("s", StringType) + .add("st", new StructType().add("a", IntegerType).add("b", StringType)) + + private type Value = (Long, String, Option[(Int, String)]) + + private def value(i: Long): Value = { + val s = if (i % 3 == 0) null else s"str_$i" + val st = if (i % 5 == 0) None else Some((i.toInt * 2, if (i % 2 == 0) null else s"b_$i")) + (i, s, st) + } + + private def toRow(v: Value): InternalRow = { + val st = v._3.map { case (a, b) => InternalRow(a, UTF8String.fromString(b)) }.orNull + InternalRow(v._1, UTF8String.fromString(v._2), st) + } + + private def fromRow(row: InternalRow): Value = { + val s = if (row.isNullAt(1)) null else row.getUTF8String(1).toString + val st = if (row.isNullAt(2)) { + None + } else { + val struct = row.getStruct(2, 2) + val b = if (struct.isNullAt(1)) null else struct.getUTF8String(1).toString + Some((struct.getInt(0), b)) + } + (row.getLong(0), s, st) + } + + /** Writes an Arrow IPC stream with one record batch per entry of `batchSizes`. */ + private def writeStream(batchSizes: Seq[Int]): Array[Byte] = { + val arrowSchema = ArrowUtils.toArrowSchema(schema, "UTC", true, false) + val allocator = ArrowUtils.rootAllocator.newChildAllocator("writer", 0, Long.MaxValue) + val root = VectorSchemaRoot.create(arrowSchema, allocator) + val out = new ByteArrayOutputStream() + try { + val writer = ArrowWriter.create(root) + val streamWriter = new ArrowStreamWriter(root, null, out) + streamWriter.start() + var next = 0L + batchSizes.foreach { size => + (0 until size).foreach { _ => + writer.write(toRow(value(next))) + next += 1 + } + writer.finish() + streamWriter.writeBatch() + writer.reset() + } + streamWriter.end() + } finally { + root.close() + allocator.close() + } + out.toByteArray + } + + /** + * Drives `processor` the way `PythonArrowOutput.newReaderIterator.read()` does, and returns + * the values of every output batch. Rows are read before the next batch is produced, since + * the processor releases the previous output batch when it produces the next one. + */ + private def readAll( + batchSizes: Seq[Int], + createProcessor: (ArrowStreamReader, Map[String, SQLMetric]) => ArrowOutputProcessor) + : Seq[Seq[Value]] = { + val allocator: BufferAllocator = + ArrowUtils.rootAllocator.newChildAllocator("reader", 0, Long.MaxValue) + val reader = new ArrowStreamReader(new ByteArrayInputStream(writeStream(batchSizes)), allocator) + val metrics = Map( + "pythonNumRowsReceived" -> new SQLMetric("sum"), + "pythonDataReceived" -> new SQLMetric("size")) + val processor = createProcessor(reader, metrics) + val output = ArrayBuffer.empty[Seq[Value]] + while (processor.loadBatch()) { + val batch = processor.produceBatch() + output += batch.rowIterator().asScala.map(fromRow).toSeq + } + processor.close() + reader.close(false) + // Fails if any root or slice handed out by the processor was leaked. + allocator.close() + assert(metrics("pythonNumRowsReceived").value === batchSizes.sum) + output.toSeq + } + + private def checkAllRowsOnceInOrder(batchSizes: Seq[Int], output: Seq[Seq[Value]]): Unit = { + assert(output.flatten === (0L until batchSizes.sum).map(value)) + } + + Seq( + (Seq(3, 3, 3), 2), + (Seq(4, 4, 4), 2), + (Seq(2, 5, 2, 7), 2), + (Seq(1, 0, 6, 1, 3), 3), + (Seq(10), 3) + ).foreach { case (batchSizes, maxRecords) => + test(s"slice by records: batches $batchSizes, max $maxRecords records per batch") { + val output = readAll(batchSizes, (reader, metrics) => + new SliceRecordsArrowOutputProcessorImpl(reader, metrics, maxRecords)) + checkAllRowsOnceInOrder(batchSizes, output) + assert(output.forall(_.size <= maxRecords)) + } + } + + Seq( + Seq(1, 100, 1), + Seq(100, 1, 100), + Seq(1, 0, 57, 2, 33) + ).foreach { batchSizes => + test(s"slice by bytes: batches $batchSizes") { + val maxBytes = 512 + val output = readAll(batchSizes, (reader, metrics) => + new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes)) + checkAllRowsOnceInOrder(batchSizes, output) + // The large input batches must have been sliced. + assert(output.size > batchSizes.size) + } + } +} From c0490ad088dc0ca6932d13fb108e1335b0d8c202 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 20 Oct 2026 11:07:40 -0700 Subject: [PATCH 2/4] Address review comments Disable output batch slicing in the runners whose consumers depend on the batch boundaries: ApplyInPandasWithStatePythonRunner, ArrowPythonUDTFRunner and ColumnarArrowPythonWithNamedArgumentRunner. Simplify BaseSliceArrowOutputProcessor.produceBatch and SliceRecordsArrowOutputProcessorImpl.getRoot. Check the byte bound and the one-row fallback in PythonArrowOutputSuite, make test_different_output_length exact, and add UDTF and applyInPandasWithState tests under the slicing configs. Co-authored-by: Isaac --- .../pyspark/sql/tests/arrow/test_arrow_map.py | 3 +- .../test_pandas_grouped_map_with_state.py | 25 ++++++ python/pyspark/sql/tests/test_udtf.py | 21 +++++ .../python/ArrowPythonUDTFRunner.scala | 5 ++ .../python/ColumnarArrowPythonRunner.scala | 5 ++ .../execution/python/PythonArrowOutput.scala | 17 ++-- .../ApplyInPandasWithStatePythonRunner.scala | 5 ++ .../python/PythonArrowOutputSuite.scala | 80 ++++++++----------- 8 files changed, 101 insertions(+), 60 deletions(-) diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py b/python/pyspark/sql/tests/arrow/test_arrow_map.py index 853ba731f1df8..622c085cae3e3 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_map.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py @@ -151,7 +151,8 @@ def func(iterator): df = self.spark.range(10) actual = df.repartition(1).mapInArrow(func, "a long").collect() - self.assertEqual(set((r.a for r in actual)), set(range(100))) + # Each input batch produces 0..99 once, in order, since there is one partition. + self.assertEqual([r.a for r in actual], list(range(100)) * (len(actual) // 100)) def test_other_than_recordbatch_iter(self): with self.quiet(): diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py index 38f0933c35302..1893192610092 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py @@ -194,6 +194,31 @@ def check_results(batch_df, _): self._test_apply_in_pandas_with_state_basic(func, check_results) + def test_apply_in_pandas_with_state_output_batch_slicing(self): + # The output batch of each group has 5 data rows and 1 state row. It must not be sliced + # by the configs below, since its data and state counts cover the whole batch. + def func(key, pdf_iter, state): + for _ in pdf_iter: + pass + state.update((1,)) + yield pd.DataFrame( + {"key": [f"{key[0]}_{i}" for i in range(5)], "countAsString": list("01234")} + ) + + expected = sorted( + Row(key=f"{k}_{i}", countAsString=str(i)) for k in ["hello", "this"] for i in range(5) + ) + + def check_results(batch_df, _): + assert sorted(batch_df.collect()) == expected + + for conf in [ + {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}, + {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"}, + ]: + with self.subTest(conf=conf), self.sql_conf(conf): + self._test_apply_in_pandas_with_state_basic(func, check_results) + def test_apply_in_pandas_with_state_python_worker_random_failure(self): input_path = tempfile.mkdtemp() output_path = tempfile.mkdtemp() diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py index 5e3df713ca3c7..5889f800ee777 100644 --- a/python/pyspark/sql/tests/test_udtf.py +++ b/python/pyspark/sql/tests/test_udtf.py @@ -3786,6 +3786,27 @@ def eval(self): ): udtf(TestUDTF, returnType=ret_type)().collect() + def test_udtf_lateral_join_with_output_batch_slicing(self): + # Each output batch is joined with one input row, so it must not be sliced by the + # configs below. + @udtf(returnType="x: int") + class TestUDTF: + def eval(self, n: int): + for i in range(5): + yield (n * 10 + i,) + + self.spark.udtf.register("test_udtf_slicing", TestUDTF) + expected = [Row(id=n, x=n * 10 + i) for n in range(3) for i in range(5)] + for conf in [ + {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}, + {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"}, + ]: + with self.subTest(conf=conf), self.sql_conf(conf): + actual = self.spark.sql( + "SELECT t.id, u.x FROM range(0, 3, 1, 1) t, LATERAL test_udtf_slicing(t.id) u" + ).collect() + self.assertEqual(sorted(actual), expected) + class LegacyUDTFArrowTests(LegacyUDTFArrowTestsMixin, ReusedSQLTestCase): @classmethod diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala index deceed02abb9a..f63a28effb7c5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala @@ -51,6 +51,11 @@ class ArrowPythonUDTFRunner( with BasicPythonArrowOutput { ArrowUtils.failDuplicatedFieldNames(schema) + // `EvalPythonUDTFExec` joins each output batch with one input row, so the batches returned + // from Python must not be sliced. + override protected def arrowMaxRecordsPerOutputBatch: Int = -1 + override protected def arrowMaxBytesPerOutputBatch: Int = -1 + override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf override protected def evalConf: Map[String, String] = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala index de19b43d32c1d..714ce62b3050e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala @@ -52,6 +52,11 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner( with ColumnarArrowPythonInput with BasicPythonArrowOutput { + // `ColumnarArrowEvalPythonEvaluatorFactory` pairs each output batch with one input batch, so + // the batches returned from Python must not be sliced. + override protected def arrowMaxRecordsPerOutputBatch: Int = -1 + override protected def arrowMaxBytesPerOutputBatch: Int = -1 + override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf // The input schema (and per-UDF nesting depth) travel in evalConf, exactly like 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 859427043b356..bed9dda441c42 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 @@ -225,11 +225,8 @@ abstract class BaseSliceArrowOutputProcessor( override def produceBatch(): ColumnarBatch = { val batchRoot = getRoot - - if (batchRoot != prevRoot) { - closePrevSlice() - prevRoot = batchRoot - } + closePrevSlice() + prevRoot = batchRoot val vectors = getVectors(batchRoot) prevVectors = vectors @@ -317,13 +314,9 @@ class SliceRecordsArrowOutputProcessorImpl( extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) { protected override def getRoot: VectorSchemaRoot = { - val remainingRows = rowCount - currentRowIdx - // Output the whole root only if no row of this batch has been output yet. - val rootSlice = if (currentRowIdx == 0 && remainingRows <= arrowMaxRecordsPerOutputBatch) { - root - } else { - root.slice(currentRowIdx, math.min(remainingRows, arrowMaxRecordsPerOutputBatch)) - } + val numRows = math.min(rowCount - currentRowIdx, arrowMaxRecordsPerOutputBatch) + // Output the whole root only if this chunk is the whole batch. + val rootSlice = if (numRows == rowCount) root else root.slice(currentRowIdx, numRows) currentRowIdx = currentRowIdx + rootSlice.getRowCount diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala index b512c99f7b6be..f58138c4b68e0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala @@ -72,6 +72,11 @@ class ApplyInPandasWithStatePythonRunner( with PythonArrowOutput[OutType] { ArrowUtils.failDuplicatedFieldNames(inputSchema) + // The data and state counts of each batch returned from Python cover the whole batch and are + // read from its first row, so the batch must not be sliced. + override protected def arrowMaxRecordsPerOutputBatch: Int = -1 + override protected def arrowMaxBytesPerOutputBatch: Int = -1 + // Installed at worker launch. This runner is constructed per task, so each launch reads the // configuration that task carries. A running streaming query carries a snapshot: // `StreamExecution` runs its batches on `sparkSession.cloneSession()`, whose conf is copied at diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala index d3e92adea25de..0ade7ce546a06 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala @@ -22,8 +22,6 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream} import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ -import org.apache.arrow.memory.BufferAllocator -import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} import org.apache.spark.SparkFunSuite @@ -32,6 +30,7 @@ import org.apache.spark.sql.execution.arrow.ArrowWriter import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types._ import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.ArrowColumnVector import org.apache.spark.unsafe.types.UTF8String class PythonArrowOutputSuite extends SparkFunSuite { @@ -41,45 +40,27 @@ class PythonArrowOutputSuite extends SparkFunSuite { .add("s", StringType) .add("st", new StructType().add("a", IntegerType).add("b", StringType)) - private type Value = (Long, String, Option[(Int, String)]) - - private def value(i: Long): Value = { - val s = if (i % 3 == 0) null else s"str_$i" - val st = if (i % 5 == 0) None else Some((i.toInt * 2, if (i % 2 == 0) null else s"b_$i")) - (i, s, st) - } - - private def toRow(v: Value): InternalRow = { - val st = v._3.map { case (a, b) => InternalRow(a, UTF8String.fromString(b)) }.orNull - InternalRow(v._1, UTF8String.fromString(v._2), st) - } - - private def fromRow(row: InternalRow): Value = { - val s = if (row.isNullAt(1)) null else row.getUTF8String(1).toString - val st = if (row.isNullAt(2)) { - None + private def row(i: Long): InternalRow = { + val s = if (i % 3 == 0) null else UTF8String.fromString(s"str_$i") + val st = if (i % 5 == 0) { + null } else { - val struct = row.getStruct(2, 2) - val b = if (struct.isNullAt(1)) null else struct.getUTF8String(1).toString - Some((struct.getInt(0), b)) + InternalRow(i.toInt * 2, if (i % 2 == 0) null else UTF8String.fromString(s"b_$i")) } - (row.getLong(0), s, st) + InternalRow(i, s, st) } /** Writes an Arrow IPC stream with one record batch per entry of `batchSizes`. */ private def writeStream(batchSizes: Seq[Int]): Array[Byte] = { - val arrowSchema = ArrowUtils.toArrowSchema(schema, "UTC", true, false) - val allocator = ArrowUtils.rootAllocator.newChildAllocator("writer", 0, Long.MaxValue) - val root = VectorSchemaRoot.create(arrowSchema, allocator) + val writer = ArrowWriter.create(schema, "UTC") val out = new ByteArrayOutputStream() try { - val writer = ArrowWriter.create(root) - val streamWriter = new ArrowStreamWriter(root, null, out) + val streamWriter = new ArrowStreamWriter(writer.root, null, out) streamWriter.start() var next = 0L batchSizes.foreach { size => (0 until size).foreach { _ => - writer.write(toRow(value(next))) + writer.write(row(next)) next += 1 } writer.finish() @@ -88,32 +69,33 @@ class PythonArrowOutputSuite extends SparkFunSuite { } streamWriter.end() } finally { - root.close() - allocator.close() + writer.root.close() } out.toByteArray } /** * Drives `processor` the way `PythonArrowOutput.newReaderIterator.read()` does, and returns - * the values of every output batch. Rows are read before the next batch is produced, since - * the processor releases the previous output batch when it produces the next one. + * the rows and the buffer size in bytes of every output batch. Rows are copied before the + * next batch is produced, since the processor releases the previous output batch then. */ private def readAll( batchSizes: Seq[Int], createProcessor: (ArrowStreamReader, Map[String, SQLMetric]) => ArrowOutputProcessor) - : Seq[Seq[Value]] = { - val allocator: BufferAllocator = - ArrowUtils.rootAllocator.newChildAllocator("reader", 0, Long.MaxValue) + : Seq[(Seq[InternalRow], Long)] = { + val allocator = ArrowUtils.rootAllocator.newChildAllocator("reader", 0, Long.MaxValue) val reader = new ArrowStreamReader(new ByteArrayInputStream(writeStream(batchSizes)), allocator) val metrics = Map( "pythonNumRowsReceived" -> new SQLMetric("sum"), "pythonDataReceived" -> new SQLMetric("size")) val processor = createProcessor(reader, metrics) - val output = ArrayBuffer.empty[Seq[Value]] + val output = ArrayBuffer.empty[(Seq[InternalRow], Long)] while (processor.loadBatch()) { val batch = processor.produceBatch() - output += batch.rowIterator().asScala.map(fromRow).toSeq + val bytes = (0 until batch.numCols()).map { i => + batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector.getBufferSize.toLong + }.sum + output += ((batch.rowIterator().asScala.map(_.copy()).toSeq, bytes)) } processor.close() reader.close(false) @@ -123,8 +105,10 @@ class PythonArrowOutputSuite extends SparkFunSuite { output.toSeq } - private def checkAllRowsOnceInOrder(batchSizes: Seq[Int], output: Seq[Seq[Value]]): Unit = { - assert(output.flatten === (0L until batchSizes.sum).map(value)) + private def checkAllRowsOnceInOrder( + batchSizes: Seq[Int], + output: Seq[(Seq[InternalRow], Long)]): Unit = { + assert(output.flatMap(_._1) === (0L until batchSizes.sum).map(row)) } Seq( @@ -138,22 +122,24 @@ class PythonArrowOutputSuite extends SparkFunSuite { val output = readAll(batchSizes, (reader, metrics) => new SliceRecordsArrowOutputProcessorImpl(reader, metrics, maxRecords)) checkAllRowsOnceInOrder(batchSizes, output) - assert(output.forall(_.size <= maxRecords)) + assert(output.forall(_._1.size <= maxRecords)) } } Seq( - Seq(1, 100, 1), - Seq(100, 1, 100), - Seq(1, 0, 57, 2, 33) - ).foreach { batchSizes => - test(s"slice by bytes: batches $batchSizes") { - val maxBytes = 512 + (Seq(1, 100, 1), 512), + (Seq(100, 1, 100), 512), + (Seq(1, 0, 57, 2, 33), 512), + // Even a single row is larger than the limit, so every output batch has one row. + (Seq(2, 0, 5, 1), 3) + ).foreach { case (batchSizes, maxBytes) => + test(s"slice by bytes: batches $batchSizes, max $maxBytes bytes per batch") { val output = readAll(batchSizes, (reader, metrics) => new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes)) checkAllRowsOnceInOrder(batchSizes, output) // The large input batches must have been sliced. assert(output.size > batchSizes.size) + assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes < maxBytes }) } } } From eced9b0474ca33bfbd47b349e975967fdd74a115 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 20 Oct 2026 17:13:38 -0700 Subject: [PATCH 3/4] Address second round of review comments Add a columnar-input Arrow Python UDF test under the slicing configs, say in the docs of both configs which cases ignore them, narrow the comment in ColumnarArrowPythonWithNamedArgumentRunner to evalArrowColumnar, pin the output sizes of the records tests, and fix several test comments and checks. Co-authored-by: Isaac --- python/pyspark/sql/tests/arrow/test_arrow_map.py | 2 +- .../tests/arrow/test_arrow_python_udf_cached.py | 16 ++++++++++++++++ .../pandas/test_pandas_grouped_map_with_state.py | 7 +++---- python/pyspark/sql/tests/test_udtf.py | 6 +++--- .../org/apache/spark/sql/internal/SQLConf.scala | 7 +++++-- .../python/ColumnarArrowPythonRunner.scala | 4 ++-- .../sql/execution/python/PythonArrowOutput.scala | 2 +- .../python/PythonArrowOutputSuite.scala | 10 +++++++--- 8 files changed, 38 insertions(+), 16 deletions(-) diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py b/python/pyspark/sql/tests/arrow/test_arrow_map.py index 622c085cae3e3..e08b65e759dd5 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_map.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py @@ -152,7 +152,7 @@ def func(iterator): df = self.spark.range(10) actual = df.repartition(1).mapInArrow(func, "a long").collect() # Each input batch produces 0..99 once, in order, since there is one partition. - self.assertEqual([r.a for r in actual], list(range(100)) * (len(actual) // 100)) + self.assertEqual([r.a for r in actual], list(range(100)) * max(1, len(actual) // 100)) def test_other_than_recordbatch_iter(self): with self.quiet(): diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py index ba14a9b740757..eb435414f2b75 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py @@ -107,6 +107,22 @@ def _assert_udf_on_cached_strings(self): finally: df.unpersist() + def test_udf_on_cached_input_with_output_batch_slicing(self): + # The columnar path pairs each result batch with one input batch, so the 10-row result + # of the single cached batch must not be sliced under the configs below. + df = self.spark.range(0, 10, 1, 1).cache() + try: + for conf in [ + {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}, + {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"}, + ]: + with self.subTest(conf=conf), self.sql_conf(conf): + result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id"))) + self._assert_columnar_arrow_eval(result, expected_nodes=1) + assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)]) + finally: + df.unpersist() + def test_udf_on_cached_strings(self): # Cache-scan vectors use standard 32-bit var-width offsets, matching the default # interchange schema: the columnar fast path forwards them verbatim. diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py index 1893192610092..f6fe70a438c1e 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py @@ -195,11 +195,10 @@ def check_results(batch_df, _): self._test_apply_in_pandas_with_state_basic(func, check_results) def test_apply_in_pandas_with_state_output_batch_slicing(self): - # The output batch of each group has 5 data rows and 1 state row. It must not be sliced - # by the configs below, since its data and state counts cover the whole batch. + # The worker packs the data rows (5 per group) and state rows (1 per group) of the groups + # in a partition into one output batch. Its data and state counts cover the whole batch, + # so it must not be sliced by the configs below. def func(key, pdf_iter, state): - for _ in pdf_iter: - pass state.update((1,)) yield pd.DataFrame( {"key": [f"{key[0]}_{i}" for i in range(5)], "countAsString": list("01234")} diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py index 5889f800ee777..24cfe5e9369e6 100644 --- a/python/pyspark/sql/tests/test_udtf.py +++ b/python/pyspark/sql/tests/test_udtf.py @@ -3802,10 +3802,10 @@ def eval(self, n: int): {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"}, ]: with self.subTest(conf=conf), self.sql_conf(conf): - actual = self.spark.sql( + df = self.spark.sql( "SELECT t.id, u.x FROM range(0, 3, 1, 1) t, LATERAL test_udtf_slicing(t.id) u" - ).collect() - self.assertEqual(sorted(actual), expected) + ) + assertDataFrameEqual(df, expected) class LegacyUDTFArrowTests(LegacyUDTFArrowTestsMixin, ReusedSQLTestCase): diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 6af6dfd9cd6cf..ba1d0667f515e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -5452,7 +5452,9 @@ object SQLConf { "there is no limit. Note that the complete ArrowRecordBatch is actually created but " + "the number of records is limited when sending it to the downstream operator. This is " + "used to avoid large batches being sent to the downstream operator including " + - "the columnar-based operator implemented by third-party libraries.") + "the columnar-based operator implemented by third-party libraries. " + + "applyInPandasWithState, Arrow Python UDTFs and Arrow Python UDFs that take " + + "columnar input directly ignore this config.") .version("4.1.0") .internal() .intConf @@ -5468,7 +5470,8 @@ object SQLConf { "the columnar-based operator implemented by third-party libraries. Spark will try to " + "create batches with the size equal or less than this value. Normally it should not " + "happen, but if in extreme case that even one record is still larger than this value, " + - "Spark will create a batch with one record.") + "Spark will create a batch with one record. applyInPandasWithState, Arrow Python " + + "UDTFs and Arrow Python UDFs that take columnar input directly ignore this config.") .version("4.1.0") .internal() .bytesConf(ByteUnit.BYTE) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala index 714ce62b3050e..81f78399784bf 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala @@ -52,8 +52,8 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner( with ColumnarArrowPythonInput with BasicPythonArrowOutput { - // `ColumnarArrowEvalPythonEvaluatorFactory` pairs each output batch with one input batch, so - // the batches returned from Python must not be sliced. + // `ColumnarArrowEvalPythonEvaluatorFactory.evalArrowColumnar` pairs each output batch with + // one input batch, so the batches returned from Python must not be sliced. override protected def arrowMaxRecordsPerOutputBatch: Int = -1 override protected def arrowMaxBytesPerOutputBatch: Int = -1 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 bed9dda441c42..ec9345e70a0f5 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 @@ -318,7 +318,7 @@ class SliceRecordsArrowOutputProcessorImpl( // Output the whole root only if this chunk is the whole batch. val rootSlice = if (numRows == rowCount) root else root.slice(currentRowIdx, numRows) - currentRowIdx = currentRowIdx + rootSlice.getRowCount + currentRowIdx = currentRowIdx + numRows rootSlice } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala index 0ade7ce546a06..83799adf4e2bf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala @@ -99,7 +99,7 @@ class PythonArrowOutputSuite extends SparkFunSuite { } processor.close() reader.close(false) - // Fails if any root or slice handed out by the processor was leaked. + // Fails if any slice made by the processor was leaked. allocator.close() assert(metrics("pythonNumRowsReceived").value === batchSizes.sum) output.toSeq @@ -122,7 +122,11 @@ class PythonArrowOutputSuite extends SparkFunSuite { val output = readAll(batchSizes, (reader, metrics) => new SliceRecordsArrowOutputProcessorImpl(reader, metrics, maxRecords)) checkAllRowsOnceInOrder(batchSizes, output) - assert(output.forall(_._1.size <= maxRecords)) + val expectedSizes = batchSizes.flatMap { n => + val sizes = Seq.fill(n / maxRecords)(maxRecords) ++ Seq(n % maxRecords).filter(_ > 0) + if (sizes.isEmpty) Seq(0) else sizes + } + assert(output.map(_._1.size) === expectedSizes) } } @@ -130,7 +134,7 @@ class PythonArrowOutputSuite extends SparkFunSuite { (Seq(1, 100, 1), 512), (Seq(100, 1, 100), 512), (Seq(1, 0, 57, 2, 33), 512), - // Even a single row is larger than the limit, so every output batch has one row. + // Even a single row is larger than the limit, so every non-empty output batch has one row. (Seq(2, 0, 5, 1), 3) ).foreach { case (batchSizes, maxBytes) => test(s"slice by bytes: batches $batchSizes, max $maxBytes bytes per batch") { From 9656977968b270c201054086d2e8a43fb9d0937f Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 23 Oct 2026 00:05:10 -0700 Subject: [PATCH 4/4] Address third round of review comments Document when a runner has to override the output slicing limits, drop prevVectors and reuse the base column vectors for whole-root outputs in BaseSliceArrowOutputProcessor, add an Arrow Python UDF test on the row path under the records config, close the processor twice in PythonArrowOutputSuite as production does, stop the query in the applyInPandasWithState test helper, and move the cached-input test next to related tests. Co-authored-by: Isaac --- .../sql/tests/arrow/test_arrow_python_udf.py | 8 +++++ .../arrow/test_arrow_python_udf_cached.py | 32 +++++++++---------- .../test_pandas_grouped_map_with_state.py | 12 ++++--- .../execution/python/PythonArrowOutput.scala | 28 ++++++++-------- .../python/PythonArrowOutputSuite.scala | 4 ++- 5 files changed, 50 insertions(+), 34 deletions(-) diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py index 417f1082cdc45..43622f61b18b5 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py @@ -385,6 +385,14 @@ def create_struct_with_interval(interval_val, name_val): self.assertIsNotNone(result["interval_field"]) self.assertEqual(result["name"], "test_name") + def test_udf_with_output_batch_slicing(self): + # Each result row is joined with the next queued input row, so the sliced 10-row result + # batch must be output once and in order. + df = self.spark.range(0, 10, 1, 1) + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}): + result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id"))) + assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)]) + @unittest.skipIf( not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py index eb435414f2b75..196c172de54e8 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py @@ -107,22 +107,6 @@ def _assert_udf_on_cached_strings(self): finally: df.unpersist() - def test_udf_on_cached_input_with_output_batch_slicing(self): - # The columnar path pairs each result batch with one input batch, so the 10-row result - # of the single cached batch must not be sliced under the configs below. - df = self.spark.range(0, 10, 1, 1).cache() - try: - for conf in [ - {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}, - {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"}, - ]: - with self.subTest(conf=conf), self.sql_conf(conf): - result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id"))) - self._assert_columnar_arrow_eval(result, expected_nodes=1) - assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)]) - finally: - df.unpersist() - def test_udf_on_cached_strings(self): # Cache-scan vectors use standard 32-bit var-width offsets, matching the default # interchange schema: the columnar fast path forwards them verbatim. @@ -149,6 +133,22 @@ def test_udf_on_cached_numeric_column_alongside_unselected_string(self): finally: df.unpersist() + def test_udf_on_cached_input_with_output_batch_slicing(self): + # The columnar path pairs each result batch with one input batch, so the 10-row result + # of the single cached batch must not be sliced under the configs below. + df = self.spark.range(0, 10, 1, 1).cache() + try: + for conf in [ + {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}, + {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"}, + ]: + with self.subTest(conf=conf), self.sql_conf(conf): + result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id"))) + self._assert_columnar_arrow_eval(result, expected_nodes=1) + assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)]) + finally: + df.unpersist() + def test_udf_on_cached_map_column(self): # Canonical map vectors from the cache pass the gate (their entry-struct children carry # the canonical "key"/"value" names in order) and ride the columnar path verbatim. The diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py index f6fe70a438c1e..dd58f04064bd4 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py @@ -93,10 +93,14 @@ def prepare_test_resource(): .start() ) - self.assertEqual(q.name, "this_query") - self.assertTrue(q.isActive) - q.processAllAvailable() - self.assertTrue(q.exception() is None) + try: + self.assertEqual(q.name, "this_query") + self.assertTrue(q.isActive) + q.processAllAvailable() + self.assertTrue(q.exception() is None) + finally: + q.stop() + shutil.rmtree(input_path, ignore_errors=True) def test_apply_in_pandas_with_state_basic(self): def func(key, pdf_iter, state): 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 ec9345e70a0f5..0fbf3240bf1bb 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 @@ -44,6 +44,9 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] { self: BasePythonRunner[ protected def deserializeColumnarBatch(batch: ColumnarBatch, schema: StructType): OUT + // The limits for slicing the batches returned from Python. A runner whose consumer cannot take + // a batch split into several, e.g. one that pairs each batch with one input row, overrides both + // to -1, as `ArrowPythonUDTFRunner` does. protected def arrowMaxRecordsPerOutputBatch: Int = SQLConf.get.arrowMaxRecordsPerOutputBatch protected def arrowMaxBytesPerOutputBatch: Int = SQLConf.get.arrowMaxBytesPerOutputBatch.toInt @@ -210,8 +213,7 @@ abstract class BaseSliceArrowOutputProcessor( extends ArrowOutputProcessorImpl(reader, pythonMetrics) { protected var currentRowIdx = -1 - protected var prevRoot: VectorSchemaRoot = null - protected var prevVectors: Array[ColumnVector] = _ + protected var prevRoot: VectorSchemaRoot = root override def loadBatch(): Boolean = { if (rowCount > 0 && currentRowIdx < rowCount) { @@ -228,25 +230,25 @@ abstract class BaseSliceArrowOutputProcessor( closePrevSlice() prevRoot = batchRoot - val vectors = getVectors(batchRoot) - prevVectors = vectors - - val batch = new ColumnarBatch(vectors) + val batch = new ColumnarBatch(getVectors(batchRoot)) batch.setNumRows(batchRoot.getRowCount) batch } protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = { - root.getFieldVectors.asScala.map { vector => - new ArrowColumnVector(vector) - }.toArray[ColumnVector] + if (root eq this.root) { + super.getVectors(root) + } else { + root.getFieldVectors.asScala.map { vector => + new ArrowColumnVector(vector) + }.toArray[ColumnVector] + } } - // Releases the previous output batch if it is a slice. The reader root is reused for every - // input batch, so it must not be closed here; `ArrowOutputProcessorImpl.close` closes it. + // Releases the previous output batch if it is a slice. The reader loads every batch from Python + // into the same root, so it must not be closed here; `ArrowOutputProcessorImpl.close` closes it. private def closePrevSlice(): Unit = { - if (prevRoot != null && (prevRoot ne root)) { - prevVectors.foreach(_.close()) + if (prevRoot ne root) { prevRoot.close() } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala index 83799adf4e2bf..6ee8fcba9db97 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala @@ -101,6 +101,8 @@ class PythonArrowOutputSuite extends SparkFunSuite { reader.close(false) // Fails if any slice made by the processor was leaked. allocator.close() + // `PythonArrowOutput` closes the processor again in its task completion listener. + processor.close() assert(metrics("pythonNumRowsReceived").value === batchSizes.sum) output.toSeq } @@ -141,7 +143,7 @@ class PythonArrowOutputSuite extends SparkFunSuite { val output = readAll(batchSizes, (reader, metrics) => new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes)) checkAllRowsOnceInOrder(batchSizes, output) - // The large input batches must have been sliced. + // Guards the choice of cases: each has a batch that does not fit, which must be sliced. assert(output.size > batchSizes.size) assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes < maxBytes }) }