viirya opened a new pull request, #59786:
URL: https://github.com/apache/spark/pull/59786
### What changes were proposed in this pull request?
This PR fixes how `SliceBytesArrowOutputProcessorImpl`, used when the
internal config `spark.sql.execution.arrow.maxBytesPerOutputBatch` is set,
measures an output batch and compares it with the limit. It addresses the six
items in SPARK-60395, raised in the review of #59616:
1. Add `SliceBytesArrowOutputProcessorImpl.bufferSize(vector): Long` and use
it to measure batches. `ValueVector.getBufferSize` computes the size in int, so
it wraps around over 2 GiB: a struct vector sums its children and returns
`(int) buffer`, and fixed- and variable-width vectors compute `valueCount *
typeWidth` and their offsets in int. The new function computes fixed- and
variable-width vectors as Long, and sums the children of a nested vector as
Long. The nested vector's own buffers (validity, offsets, type ids) are what
remains of its `getBufferSize` after its children's `getBufferSize` are
subtracted in int, which cancels out any wrap-around without listing every
nested type.
2. Both checks now use `<=`. This matches the config doc ("equal or less
than this value") and the records processor.
3. The shrink search starts from the row count of the last slice that fit in
the current batch, instead of from all the remaining rows for every output
batch. When the whole batch does not fit, the first probe is `0.9 * rowCount`,
which skips the `slice(0, rowCount)` probe that always measured the same as the
root.
4. The one-row fallback is now `root.slice(currentRowIdx, 1)`, without the
always-true `if (sliceBatchCount == 0)` and the unused initial `rootSlice`.
5. `SQLConf.arrowMaxBytesPerOutputBatch` now returns `math.max(value,
0L).toInt`. Before, `PythonArrowOutput` narrowed the Long with `.toInt`, so a
value below `Int.MinValue` wrapped around to a positive limit.
6. Drop the redundant `rowCount > 0 &&` in
`BaseSliceArrowOutputProcessor.loadBatch`.
### Why are the changes needed?
- The outputs of mapInArrow, mapInPandas, applyInPandas and Python UDTFs are
one top-level struct vector. A batch over 2 GiB could measure negative and be
output without slicing. For example, three int64 columns of 2^27 rows measure
as -1,006,632,960 bytes.
- A batch whose size equals the limit was sliced, although the doc says it
is allowed.
- Every probe that does not start at row 0 copies its offset buffers.
Restarting the search from all the remaining rows made splitting one batch of R
rows into chunks of c rows copy O(R^2/c) offsets.
- A negative value, which the doc describes as no limit, could become a
limit, e.g. `-3g` became 1 GiB and `-4294967295` became 1 byte. If
`maxRecordsPerOutputBatch` was also set, the check that both limits are not set
at once then threw.
### Does this PR introduce _any_ user-facing change?
Only for the internal config
`spark.sql.execution.arrow.maxBytesPerOutputBatch`:
- Output batches over 2 GiB are now sliced.
- A batch whose size equals the limit is no longer sliced.
- Values below `Int.MinValue` now mean no limit, as the doc says.
Within one Python batch, slices no longer grow back after a smaller one. If
the earlier rows of a batch are larger than the later ones, the later slices
can be smaller than before, but still within the limit.
### How was this patch tested?
New tests in `PythonArrowOutputSuite`:
- A batch whose size equals the limit is not sliced, and one byte less
slices it.
- The number of size measurements stays linear in the number of slices.
Splitting 2000 rows into about 50 slices took 1513 measurements before this
change.
- `bufferSize` matches `getBufferSize` for nested and leaf vectors (boolean,
decimal, binary, array of struct, map of array) on roots and slices below 2 GiB.
- `bufferSize` does not wrap around for a struct vector whose children
report 3 GiB in total, while `StructVector.getBufferSize` is negative. This
uses vectors that report a large size, so nothing over 2 GiB is allocated.
- `maxBytesPerOutputBatch` values of `-3g`, `-4294967295`, `-1` and `0` mean
no limit.
The first, second and last tests fail without the fix. Also ran
`pyspark.sql.tests.arrow.test_arrow_map`,
`pyspark.sql.tests.pandas.test_pandas_grouped_map_with_state`,
`pyspark.sql.tests.arrow.test_arrow_python_udf_cached` and
`pyspark.sql.tests.test_udtf`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]
