Yicong-Huang opened a new pull request, #58182: URL: https://github.com/apache/spark/pull/58182
### What changes were proposed in this pull request? Add an internal config `spark.sql.execution.pythonUDF.arrow.maxBytesPerOutputBatch` (bytes, default `-1` = no limit) and a new `ArrowBatchTransformer.resize_batches(Iterator[RecordBatch], max_bytes)` helper. When the config is set, the `applyInPandas` (`SQL_GROUPED_MAP_PANDAS_UDF`) worker slices each output Arrow `RecordBatch` that exceeds the cap into `ceil(nbytes / max_bytes)` row-balanced, zero-copy pieces before sending them to the JVM; a batch already within the cap (or empty) passes through unchanged. The config value is shipped to the worker through `ArrowPythonRunner.getPythonRunnerConfMap`. The config is intentionally general (not `applyInPandas`-specific) so the same output-batch bound can be extended to other Arrow-based Python UDF eval types later. This PR wires it for `applyInPandas`, which produces one Arrow batch per group and is the case most prone to a single oversized output batch. ### Why are the changes needed? `applyInPandas` produces one output Arrow `RecordBatch` per group, so a large group yields a single oversized batch. That strains the JVM and downstream operators (large allocations, and columnar operators that do not handle very large batches well). Bounding the output batch size on the worker side lets users cap the batch produced for a big group without changing the UDF. ### Does this PR introduce _any_ user-facing change? No. The config is internal and defaults to `-1` (no limit), so behavior is unchanged unless it is explicitly set. ### How was this patch tested? New unit tests for `resize_batches` in `test_conversion.py`: a large batch splits into row-balanced pieces preserving row order, and small and empty batches pass through unchanged. ### Was this patch authored or co-authored using generative AI tooling? No -- 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]
