This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 5a83232af30b [SPARK-60627][SQL] Do not slice an Arrow batch whose
slice positions overflow int in `SliceRecordsArrowOutputProcessorImpl`
5a83232af30b is described below
commit 5a83232af30be1fd4c13ae067bd29c2e473358fb
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Wed Oct 28 18:31:55 2026 -0700
[SPARK-60627][SQL] Do not slice an Arrow batch whose slice positions
overflow int in `SliceRecordsArrowOutputProcessorImpl`
### What changes were proposed in this pull request?
This PR fixes `SliceRecordsArrowOutputProcessorImpl`, which slices the
Arrow batches returned from Python when the internal config
`spark.sql.execution.arrow.maxRecordsPerOutputBatch` is set, so that it outputs
a batch whole when Arrow cannot slice it, instead of failing the task.
- Add `BaseSliceArrowOutputProcessor.canSlice`, which returns whether
`VectorSchemaRoot.slice` can slice every vector of a batch, including nested
ones. It rejects a fixed-width vector of 2 GiB or more, a view vector with 2^27
rows or more, and a list (or map) vector with 2^29 - 1 rows or more.
- `SliceRecordsArrowOutputProcessorImpl.getRoot` checks it before any row
of a batch is output, and outputs the batch whole if the check fails, as it
does for a batch within the limit.
- The doc of `spark.sql.execution.arrow.maxRecordsPerOutputBatch` now says
that a batch that Arrow cannot slice, e.g. one with a fixed-width column of 2
GiB or more, is sent as a whole.
The check lives in `BaseSliceArrowOutputProcessor` so that
`SliceBytesArrowOutputProcessorImpl` can share it. SPARK-60395 proposes the
same check for the bytes slicing.
- https://github.com/apache/spark/pull/59786
### Why are the changes needed?
In Arrow Java 19.0.0 (and 18.3.0 in `branch-4.1`), the transfer pairs that
`VectorSchemaRoot.slice` uses compute some buffer positions in int:
- `BaseFixedWidthVector.splitAndTransferValueBuffer`: `startIndex *
typeWidth` and `length * typeWidth`
- `ListVector`'s transfer pair: `startIndex * OFFSET_WIDTH` and
`(startIndex + length) * OFFSET_WIDTH`
- `BaseVariableWidthViewVector.splitAndTransferViewBufferAndDataBuffer`: `i
* ELEMENT_SIZE`
Once a slice reaches 2 GiB into such a buffer, the position wraps around to
a negative value and `ArrowBuf.slice` throws `IndexOutOfBoundsException`.
`VectorSchemaRoot.slice` does not close the vectors it has sliced so far when
this happens, so they keep holding the buffers of the batch, and the batch is
leaked as well.
The records slicing receives such a batch when Python returns one large
batch, e.g., `mapInArrow` writes each batch that its function yields as is. For
example, with `maxRecordsPerOutputBatch=16777216`, one batch of 285,212,672
rows (2^28 + 2^24) of one bigint column (2.28 GB) is output as 16 slices, and
then the slice from row 268,435,456 (2^28) fails:
```
java.lang.IndexOutOfBoundsException: index (-2147483648) must not be
negative
at
org.apache.arrow.util.Preconditions.checkPositionIndex(Preconditions.java:1128)
at
org.apache.arrow.util.Preconditions.checkPositionIndex(Preconditions.java:1111)
at org.apache.arrow.memory.ArrowBuf.slice(ArrowBuf.java:185)
at
org.apache.arrow.vector.BaseFixedWidthVector.splitAndTransferValueBuffer(BaseFixedWidthVector.java:653)
at
org.apache.arrow.vector.BaseFixedWidthVector.splitAndTransferTo(BaseFixedWidthVector.java:644)
at
org.apache.arrow.vector.BigIntVector$TransferImpl.splitAndTransfer(BigIntVector.java:346)
at
org.apache.arrow.vector.VectorSchemaRoot.lambda$slice$1(VectorSchemaRoot.java:347)
```
All 2,317,352,960 bytes of that batch then stay allocated after the
processor is closed.
This has existed since the records slicing was added by SPARK-51769 in
Apache Spark [4.1.0](https://github.com/apache/spark/releases/tag/v4.1.0)
(2025-12-11).
### Does this PR introduce _any_ user-facing change?
Yes, only for users who set the internal config
`spark.sql.execution.arrow.maxRecordsPerOutputBatch`, which is disabled by
default. A batch from Python that Arrow cannot slice is now output whole, with
more records than the limit, instead of failing the task with
`IndexOutOfBoundsException` and leaking the batch.
### How was this patch tested?
- Added a test to `PythonArrowOutputSuite` that runs
`SliceRecordsArrowOutputProcessorImpl` over a struct column, like the output of
`mapInArrow`, holding an array column that reports 2^29 + 2^20 rows without
allocating them, and checks that the batch is output whole instead of sliced. A
real overflow needs buffers over 2 GiB, so the manual runs below cover the
overflow and the leak. It uses a list vector because the fixed-width vectors
that Spark reads are final classes.
- Added a test of `BaseSliceArrowOutputProcessor.canSlice` at the
boundaries of an 8-byte fixed-width vector and of a list vector.
- Manually ran `SliceRecordsArrowOutputProcessorImpl` over a real 2.28 GB
bigint vector of 285,212,672 rows with a limit of 16,777,216 records. Without
this fix, it fails and leaks as shown above. With this fix, the batch is output
whole and the allocator closes without a leak.
- Manually ran the following on `local[1]` with `--conf
spark.sql.execution.arrow.maxRecordsPerOutputBatch=16777216`. With this fix, it
returns the right count, sum and max, and the runner logs `Batches: 1, Data:
2176.00 MB`, i.e., Python sends the 2.28 GB batch as one batch, which is output
whole.
```python
import numpy as np
import pyarrow as pa
rows = (1 << 28) + (1 << 24)
def f(it):
for _ in it:
pass
yield pa.RecordBatch.from_arrays([pa.array(np.arange(rows,
dtype=np.int64))], ["id"])
df = spark.range(1, numPartitions=1).mapInArrow(f, "id long")
df.selectExpr("count(*)", "sum(id)", "max(id)").collect()
# [Row(count(1)=285212672, sum(id)=40673133992083456, max(id)=285212671)]
```
```
build/sbt 'sql/testOnly *PythonArrowOutputSuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
Closes #59871 from dongjoon-hyun/SPARK-60627.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 4 +-
.../sql/execution/python/PythonArrowOutput.scala | 46 +++++++++++---
.../execution/python/PythonArrowOutputSuite.scala | 71 ++++++++++++++++++++++
3 files changed, 113 insertions(+), 8 deletions(-)
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 12b6254bd395..550536cf2a95 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
@@ -5562,7 +5562,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. A
batch that " +
+ "Arrow cannot slice, e.g. one with a fixed-width column of 2 GiB or
more, is sent " +
+ "as a whole. " +
"applyInPandasWithState, Arrow Python UDTFs and Arrow Python UDFs that
take " +
"columnar input directly ignore this config.")
.version("4.1.0")
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 6df6e8904985..07a0939736e9 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
@@ -21,7 +21,8 @@ import java.util.concurrent.atomic.AtomicBoolean
import scala.jdk.CollectionConverters._
-import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.{BaseFixedWidthVector,
BaseVariableWidthViewVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{BaseRepeatedValueVector, ListVector}
import org.apache.arrow.vector.ipc.ArrowStreamReader
import org.apache.spark.{SparkEnv, TaskContext}
@@ -267,6 +268,33 @@ abstract class BaseSliceArrowOutputProcessor(
closePrevSlice()
super.close()
}
+
+ private[python] def canSlice(root: VectorSchemaRoot): Boolean =
+ root.getFieldVectors.asScala.forall(BaseSliceArrowOutputProcessor.canSlice)
+}
+
+object BaseSliceArrowOutputProcessor {
+ /**
+ * Returns whether `VectorSchemaRoot.slice` can slice `vector` and its
children. Arrow computes
+ * some byte positions of a slice in int, so it throws, and leaks the
vectors sliced so far, for
+ * a fixed-width vector of 2 GiB or more, a view vector with 2^27 rows or
more, or a list vector
+ * with 2^29 - 1 rows or more. A batch that holds one is output whole
instead.
+ *
+ * The offset positions of a list vector overflow only from 2^29 rows, but a
slice of `length`
+ * rows also allocates `(length + 1) * 4` bytes of offsets in int, and a
list vector nested in
+ * another one is sliced over all its rows when the slice of its parent
covers them.
+ */
+ private[python] def canSlice(vector: FieldVector): Boolean = {
+ val fits = vector match {
+ case v: BaseFixedWidthVector => v.getValueCount.toLong * v.getTypeWidth
<= Int.MaxValue
+ case v: BaseVariableWidthViewVector =>
+ v.getValueCount.toLong * BaseVariableWidthViewVector.ELEMENT_SIZE <=
Int.MaxValue
+ case v: ListVector =>
+ (v.getValueCount + 1L) * BaseRepeatedValueVector.OFFSET_WIDTH <=
Int.MaxValue
+ case _ => true
+ }
+ fits && vector.getChildrenFromFields.asScala.forall(canSlice)
+ }
}
class SliceBytesArrowOutputProcessorImpl(
@@ -327,11 +355,15 @@ class SliceRecordsArrowOutputProcessorImpl(
protected override def getRoot: VectorSchemaRoot = {
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 + numRows
-
- rootSlice
+ // Output the whole root only if this chunk is the whole batch, or if
Arrow cannot slice the
+ // batch, which is checked before any row of it has been output.
+ if (numRows == rowCount || (currentRowIdx == 0 && !canSlice(root))) {
+ currentRowIdx = rowCount
+ root
+ } else {
+ val rootSlice = root.slice(currentRowIdx, numRows)
+ 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 6ee8fcba9db9..735b5b8f50ed 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,7 +22,10 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
+import org.apache.arrow.vector.{BigIntVector, FieldVector,
FixedSizeBinaryVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{ListVector, StructVector}
import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType}
import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.InternalRow
@@ -132,6 +135,74 @@ class PythonArrowOutputSuite extends SparkFunSuite {
}
}
+ test("SPARK-60627: slice by records: a batch that Arrow cannot slice is
output whole") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("unsliceable",
0, Long.MaxValue)
+ // An array column that reports 2^29 + 2^20 rows without allocating them,
as if its offsets
+ // took over 2 GiB, so that Arrow cannot slice it. Since a real overflow
needs buffers over
+ // 2 GiB, this checks only that the batch is output whole instead of
sliced. The fixed-width
+ // vectors that Spark reads, e.g. of an int64 column over 2 GiB, are final
classes, so a list
+ // stands in.
+ val rows = (1 << 29) + (1 << 20)
+ val list = new ListVector("l", allocator,
FieldType.nullable(ArrowType.List.INSTANCE), null) {
+ override def getValueCount: Int = rows
+ }
+ list.addOrGetVector[BigIntVector](FieldType.nullable(new ArrowType.Int(64,
true)))
+ // Like the output of mapInArrow, the column is in a struct column.
+ val struct = new StructVector(
+ "st", allocator, FieldType.nullable(ArrowType.Struct.INSTANCE), null) {
+ putChild("l", list)
+ }
+ val root = new VectorSchemaRoot(
+ Seq(struct.getField).asJava, Seq[FieldVector](struct).asJava, rows)
+ val reader = new ArrowStreamReader(
+ new ByteArrayInputStream(Array.emptyByteArray), allocator) {
+ private var loaded = false
+ override def getVectorSchemaRoot: VectorSchemaRoot = root
+ override def loadNextBatch(): Boolean = {
+ val hasNext = !loaded
+ loaded = true
+ hasNext
+ }
+ override def bytesRead(): Long = 0L
+ }
+ val metrics = Map(
+ "pythonNumRowsReceived" -> new SQLMetric("sum"),
+ "pythonDataReceived" -> new SQLMetric("size"))
+ val processor = new SliceRecordsArrowOutputProcessorImpl(reader, metrics,
1 << 20)
+ try {
+ assert(processor.loadBatch())
+ assert(processor.produceBatch().numRows() === rows)
+ assert(!processor.loadBatch())
+ } finally {
+ processor.close()
+ allocator.close()
+ }
+ }
+
+ test("SPARK-60627: canSlice rejects vectors whose slice positions overflow
int") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("canSlice", 0,
Long.MaxValue)
+ // Vectors that report many values without allocating them.
+ def fixed(rows: Int): FixedSizeBinaryVector =
+ new FixedSizeBinaryVector("f", allocator, 8) {
+ override def getValueCount: Int = rows
+ }
+ def list(rows: Int): ListVector = new ListVector(
+ "l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null) {
+ override def getValueCount: Int = rows
+ }
+ val vectors =
+ Seq(fixed((1 << 28) - 1), fixed(1 << 28), list((1 << 29) - 2), list((1
<< 29) - 1))
+ try {
+ // `splitAndTransfer` computes `startIndex * 8` and `length * 8` of an
8-byte wide column,
+ // and the offset positions and `(length + 1) * 4` offset bytes of a
list column, in int.
+ assert(vectors.map(BaseSliceArrowOutputProcessor.canSlice) ===
+ Seq(true, false, true, false))
+ } finally {
+ vectors.foreach(_.close())
+ allocator.close()
+ }
+ }
+
Seq(
(Seq(1, 100, 1), 512),
(Seq(100, 1, 100), 512),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
