ivoson opened a new pull request, #56649:
URL: https://github.com/apache/spark/pull/56649
### What changes were proposed in this pull request?
Regular batched Python UDFs (`BatchEvalPythonExec`) pickle each batch of
rows into a single contiguous on-heap byte array, and batches are formed by
**row count only** (`spark.sql.execution.python.udf.maxRecordsPerBatch`,
default 100). With wide rows, a single pickled batch can reach hundreds of MB
to over 1 GB of contiguous JVM heap and OOM the executor.
This PR adds an optional byte-size bound on these batches, plus
observability to size and verify the bound. All new behavior is off by default;
the default path is byte-for-byte unchanged.
1. **Byte cap** — new config
`spark.sql.execution.python.udf.maxBytesPerBatch` (internal, default `-1` =
off). When set to a finite value, a batch is cut at the **min** of the
row-count and byte limits. The per-row size is a best-effort estimate of the
*pickled* size of the converted row, accounted by `EvaluatePython.toJava` at
its leaf cases **during conversion itself** (no second traversal):
`UTF8String.numBytes` is the exact UTF-8 byte count pickle writes, decimals are
sized by precision, binary/geo/variant by their byte lengths, and unknown leaf
types get a small positive constant so the cap never goes blind. A single row
larger than the cap still forms a one-row batch. The new
`ByteBoundedAsArrayIterator` performs the byte-bounded grouping.
2. **Observability** — new `SQLMetric`s on `BatchEvalPythonExec`:
- `pythonPeakPickledBatchBytes` — peak per-batch pickled size (size
metric; cross-task peak shown in the UI's min/med/max breakdown, like
`peakMemory`). Always recorded.
- `pythonOversizedBatchCount` — number of batches cut at the byte limit
(only populated when a finite cap is set).
- `pythonEstimatedInputBytes` — sum of the per-row size estimates, to
compare against the measured `pythonDataSent` for estimator-accuracy analysis
(only populated when a cap is set).
The Python UDTF path (`BatchEvalPythonUDTFExec`) calls `getInputIterator`
with no cap and is unaffected.
### Why are the changes needed?
Batching by row count alone gives no bound on the per-batch heap footprint:
with wide rows (large strings/binary/decimal payloads), one pickled batch
becomes a very large contiguous allocation that can OOM the executor even when
the row *count* is small. The byte cap provides a guardrail against that
allocation, and the metrics make peak batch pressure observable so a safe
threshold can be chosen and the estimator's accuracy verified before enforcing.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
New unit tests in `BatchEvalPythonExecSuite`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
--
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]
