dongjoon-hyun opened a new pull request, #59871: URL: https://github.com/apache/spark/pull/59871
### 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 is sent as a whole. The check lives in `BaseSliceArrowOutputProcessor` so that `SliceBytesArrowOutputProcessorImpl` can share it. https://github.com/apache/spark/pull/59786 proposes the same check for the bytes slicing. ### 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. It checks that the batch is output whole and nothing is leaked. It fails without this fix, where the processor slices the batch. 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 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
