[SPARK-53614][PYTHON] Add Iterator[pandas.DataFrame] support to applyInPandas

[Yicong-Huang]

What changes were proposed in this pull request?

This PR adds support for the Iterator[pandas.DataFrame] API in groupBy().applyInPandas(), enabling batch-by-batch processing of grouped data for improved memory efficiency and scalability.

Key Changes:

1. New PythonEvalType: Added SQL_GROUPED_MAP_PANDAS_ITER_UDF to distinguish iterator-based UDFs from standard grouped map UDFs

2. Type Inference: Implemented automatic detection of iterator signatures:

   • Iterator[pd.DataFrame] -> Iterator[pd.DataFrame]
   • Tuple[Any, ...], Iterator[pd.DataFrame] -> Iterator[pd.DataFrame]

3. Streaming Serialization: Created GroupPandasIterUDFSerializer that streams results without materializing all DataFrames in memory

4. Configuration Change: Updated FlatMapGroupsInPandasExec which was hardcoding pythonEvalType = 201 instead of extracting it from the UDF expression (mirrored fix from FlatMapGroupsInArrowExec)

Why are the changes needed?

The existing applyInPandas() API loads entire groups into memory as single DataFrames. For large groups, this can cause OOM errors. The iterator API allows:

• Memory Efficiency: Process data batch-by-batch instead of materializing entire groups
• Scalability: Handle arbitrarily large groups that don't fit in memory
• Consistency: Mirrors the existing applyInArrow() iterator API design

Does this PR introduce any user-facing changes?

Yes, this PR adds a new API variant for applyInPandas():

Before (existing API, still supported):

def normalize(pdf: pd.DataFrame) -> pd.DataFrame:
    return pdf.assign(v=(pdf.v - pdf.v.mean()) / pdf.v.std())

df.groupBy("id").applyInPandas(normalize, schema="id long, v double")

After (new iterator API):

from typing import Iterator

def normalize(batches: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    # Process data batch-by-batch
    for batch in batches:
        yield batch.assign(v=(batch.v - batch.v.mean()) / batch.v.std())

df.groupBy("id").applyInPandas(normalize, schema="id long, v double")

With Grouping Keys:

from typing import Iterator, Tuple, Any

def sum_by_key(key: Tuple[Any, ...], batches: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    total = 0
    for batch in batches:
        total += batch['v'].sum()
    yield pd.DataFrame({"id": [key[0]], "total": [total]})

df.groupBy("id").applyInPandas(sum_by_key, schema="id long, total double")

Backward Compatibility: The existing DataFrame-to-DataFrame API is fully preserved and continues to work without changes.

How was this patch tested?

• Added test_apply_in_pandas_iterator_basic - Basic functionality test
• Added test_apply_in_pandas_iterator_with_keys - Test with grouping keys
• Added test_apply_in_pandas_iterator_batch_slicing - Pressure test with 10M rows, 20 columns
• Added test_apply_in_pandas_iterator_with_keys_batch_slicing - Pressure test with keys

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

Yes, tests generated by Cursor.

[zhengruifeng]

LGTM, only a few minor comments

[zhengruifeng]

Suggested change

import warnings

[Yicong-Huang]

removed

[zhengruifeng]

such imports should be move to the head of the file

[Yicong-Huang]

moved

[zhengruifeng]

why we need import builtins?
I think there is no name conflict if we use sf.max/min/sum in this file

[Yicong-Huang]

somehow when I use sum directly it would use column.sum. Do you know the reason? I changed to use builtins to avoid this conflict

[Yicong-Huang]

moved typing import

[zhengruifeng]

we don't have column.sum, do you mean sf.sum?
in some test files, sum is imported, so the builtin sum is overridden

[zhengruifeng]

let's directly compare the rows

self.assertEqual(result, [Row(...), Row(...), ...])

[Yicong-Huang]

updated

[zhengruifeng]

Suggested change

	if is_iterator_dataframe or is_iterator_dataframe_with_keys:
	return PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
	# Default to non-iterator (standard grouped map)
	return PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
	if is_iterator_dataframe or is_iterator_dataframe_with_keys:
	return PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
	# Default to non-iterator (standard grouped map)
	return PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF

this part should match

spark/python/pyspark/sql/pandas/typehints.py

Lines 368 to 379 in 9e12201

	# pa.Table -> pa.Table
	is_table = (
	len(parameters_sig) == 1 and parameters_sig[0] == pa.Table and return_annotation == pa.Table
	)
	# Tuple[pa.Scalar, ...], pa.Table -> pa.Table
	is_table_with_keys = (
	len(parameters_sig) == 2 and parameters_sig[1] == pa.Table and return_annotation == pa.Table
	)
	if is_table or is_table_with_keys:
	return PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF
	return None

we can align it in a followup

[Yicong-Huang]

let's discuss and do it with a follow up if needed.

[zhengruifeng]

thanks, merged to master