tree 373d74d6ba392af4140480b828f0841074070ee4
parent b005f56928429938812a6efcd0a05b5c2986bfce
author Ruifeng Zheng <ruifengz@apache.org> 1758713467 +0800
committer Ruifeng Zheng <ruifengz@apache.org> 1758713467 +0800

[SPARK-53562][PYTHON] Limit Arrow batch sizes in `applyInArrow` and `applyInPandas`

### What changes were proposed in this pull request?
This PR is derived from utkarsh39's internal optimization to make `applyInArrow` and `applyInPandas` respect the arrow batch size limitation:

- spark.sql.execution.arrow.maxRecordsPerBatch
- spark.sql.execution.arrow.maxBytesPerBatch

~~And this optimization is gated by `spark.sql.execution.arrow.arrowBatchSlicing.enabled`~~

### Why are the changes needed?

When evaluating Python UDFs, Spark buffers the data to be sent to the Python process in a DirectByteBuffer. Spark does NOT buffer the entire input of the Python UDF in the buffer. It buffers the data output by a single call to ArrowWriter.writeNextInputToStream() , a logical equivalent of a Next() call on an iterator. Currently there are are no guarantees on the limit of the data output by a ArrowWriter.writeNextInputToStream() call. As the output of ArrowWriter.writeNextInputToStream() is buffered in a DirectByteBuffer, we run into problems when the output size exceeds a certain threshold.

Consider the following plan

```
PandasUDF(a, b) <= Python UDF
 + GroupingAgg (groupingKey = [a], output = [b])
   + Scan (a, b)
```

When the output of the grouping agg has a large number of rows per grouping key i.e. a large number of rows that have the same value of a. In such an event, for each grouping key value, a, we buffer all the rows [(a, b1), (a, b2), ...] in a DirectByteBuffer before we flush them to the Python process for UDF evaluation. The request to allocate a large DirectByteBuffer to accommodate a large batch of rows fails resulting in an OutOfMemory error.

### Does this PR introduce _any_ user-facing change?
No, internal optimization

### How was this patch tested?
added UTs

### Was this patch authored or co-authored using generative AI tooling?
No

Closes #52303 from zhengruifeng/agg_slice.

Authored-by: Ruifeng Zheng <ruifengz@apache.org>
Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>