[SPARK-49547][SQL][PYTHON] Add iterator of RecordBatch API to applyInArrow

[Kimahriman]

What changes were proposed in this pull request?

Add the option to applyInArrow to take a function that takes an iterator of RecordBatch and returns an iterator of RecordBatch. A new eval type is added SQL_GROUPED_MAP_ARROW_ITER_UDF, and is detected via type hints on the function.

Why are the changes needed?

Having a single Table as input and a single Table as output requires collecting all inputs and outputs in memory for a single batch. This can require excessive memory for certain edge cases with large groups. Inputs and outputs already get serialized as record batches, so simply expose this lazy iterator directly instead of forcing materialization into a table.

Does this PR introduce any user-facing change?

Yes, a new function signature supported by applyInArrow.

Example:

import pyarrow as pa
import pyarrow.compute as pc
def sum_func(key: Tuple[pa.Scalar, ...], batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
    total = 0
    for batch in batches:
        total += pc.sum(batch.column("v")).as_py()
    yield pyarrow.RecordBatch.from_pydict({"id": [key[0].as_py()], "v": [total]})

df.groupby("id").applyInArrow(sum_func, schema="id long, v double").show()

+---+----+
| id|   v|
+---+----+
|  1| 3.0|
|  2|18.0|
+---+----+

How was this patch tested?

Updated existing UTs to test both Table signatures and RecordBatch signatures

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

No

[Kimahriman]

@zhengruifeng @HyukjinKwon

[zhengruifeng]

Looks pretty good!
only a few minor comments.

also cc @HyukjinKwon

[zhengruifeng]

shall we add a dedicated function def wrap_grouped_map_arrow_iter_udf for the new eval type?

[Kimahriman]

yep it is cleaner that way, done

[zhengruifeng]

please also add a simple example in the PR description, section Does this PR introduce any user-facing change?

[zhengruifeng]

let's add

    .. versionchanged:: 4.1.0
        Supports iterator API ...

[zhengruifeng]

lets also add two simple examples (w/o key) in the Examples section

[Kimahriman]

I added one with and without key, since the key is more relevant for this API I think

[zhengruifeng]

thank you @Kimahriman so much for spending time on this.
merged to master for 4.1