[SPARK-56342][PYTHON] Tighten type hints for refactored eval type functions in worker.py

### What changes were proposed in this pull request?

Tighten the type hints for `def func` signatures in `read_udfs()` for eval types that have been refactored to be self-contained. Specifically:

- Replace `Iterator[Any]` with `Iterator["GroupedBatch"]` for grouped eval types (`SQL_GROUPED_AGG_ARROW_UDF`, `SQL_GROUPED_AGG_ARROW_ITER_UDF`, `SQL_WINDOW_AGG_ARROW_UDF`)
- Rename the `batches` parameter to `data` across all refactored eval types for consistency, since the same call site passes either a flat stream or grouped stream
- Use `for batch in data` for flat streams and `for group in data` for grouped streams
- Define `GroupedBatch` and `CoGroupedBatch` type aliases in `pyspark.sql.pandas._typing`

### Why are the changes needed?

After refactoring eval types to be self-contained in `read_udfs()`, the `def func` signatures used `Iterator[Any]` for grouped eval types, losing type information. Tightening these hints improves readability and enables better static analysis.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Existing tests.

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

No.

Closes #55178 from Yicong-Huang/SPARK-56342/tighten-type-hints.

Authored-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>

Author: Yicong-Huang

python/pyspark/sql/pandas/_typing/__init__.pyi

@@ -69,6 +69,12 @@ ArrowGroupedAggUDFType = Literal[252]
 ArrowWindowAggUDFType = Literal[253]
 ArrowGroupedAggIterUDFType = Literal[254]
 
+# Arrow stream types
+# A single group of Arrow batches (e.g., one key group in groupBy).
+GroupedBatch = Iterator[pyarrow.RecordBatch]
+# A group of two relations for cogroup operations (e.g., cogroupBy).
+CoGroupedBatch = Tuple[Iterator[pyarrow.RecordBatch], Iterator[pyarrow.RecordBatch]]
+
 class ArrowVariadicScalarToScalarFunction(Protocol):
     def __call__(self, *_: pyarrow.Array) -> pyarrow.Array: ...
 

python/pyspark/worker.py

@@ -26,7 +26,10 @@
 import inspect
 import itertools
 import json
-from typing import Any, Callable, Iterable, Iterator, Optional, Tuple, Union
+from typing import Any, Callable, Iterable, Iterator, Optional, Tuple, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from pyspark.sql.pandas._typing import GroupedBatch
 
 from pyspark.accumulators import (
     SpecialAccumulatorIds,
@@ -2573,12 +2576,12 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
         assert num_udfs == 1, "One MAP_ARROW_ITER UDF expected here."
         udf_func: Callable[[Iterator[pa.RecordBatch]], Iterator[pa.RecordBatch]] = udfs[0][0]
 
-        def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+        def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
             """Apply mapInArrow UDF"""
 
             # Pre-processing
             input_batches: Iterator[pa.RecordBatch] = map(
-                ArrowBatchTransformer.flatten_struct, batches
+                ArrowBatchTransformer.flatten_struct, data
             )
 
             # invoke the UDF
@@ -2601,15 +2604,15 @@ def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.Rec
             prefers_large_types=runner_conf.use_large_var_types,
         )
 
-        def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+        def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
             """Apply scalar Arrow UDFs"""
 
-            for input_batch in batches:
+            for batch in data:
                 output_batch = pa.RecordBatch.from_arrays(
                     [
                         udf_func(
-                            *[input_batch.column(o) for o in args_offsets],
-                            **{k: input_batch.column(v) for k, v in kwargs_offsets.items()},
+                            *[batch.column(o) for o in args_offsets],
+                            **{k: batch.column(v) for k, v in kwargs_offsets.items()},
                         )
                         for udf_func, args_offsets, kwargs_offsets, _ in udfs
                     ],
@@ -2618,7 +2621,7 @@ def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.Rec
                 output_batch = ArrowBatchTransformer.enforce_schema(
                     output_batch, combined_arrow_schema
                 )
-                verify_scalar_result(output_batch, input_batch.num_rows)
+                verify_scalar_result(output_batch, batch.num_rows)
                 yield output_batch
 
         # profiling is not supported for UDF
@@ -2635,7 +2638,7 @@ def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.Rec
             return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types
         )
 
-        def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+        def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
             """Apply scalar Arrow iterator UDF"""
 
             num_input_rows = 0
@@ -2647,7 +2650,7 @@ def extract_args(batch: pa.RecordBatch):
                 return args[0] if len(args) == 1 else args
 
             # Extract args from input batches (streaming)
-            args_iter = map(extract_args, batches)
+            args_iter = map(extract_args, data)
 
             # Call UDF and verify result type (iterator of pa.Array)
             verified_iter = verify_result(pa.Array)(udf_func(args_iter))
@@ -2697,9 +2700,11 @@ def process_results():
             prefers_large_types=runner_conf.use_large_var_types,
         )
 
-        def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
-            for group_batches in batches:
-                batch_list = list(group_batches)
+        def grouped_func(
+            split_index: int, data: Iterator["GroupedBatch"]
+        ) -> Iterator[pa.RecordBatch]:
+            for group in data:
+                batch_list = list(group)
                 if not batch_list:
                     continue
                 if hasattr(pa, "concat_batches"):
@@ -2722,7 +2727,7 @@ def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
                 yield ArrowBatchTransformer.enforce_schema(batch, return_schema)
 
         # profiling is not supported for UDF
-        return func, None, ser, ser
+        return grouped_func, None, ser, ser
 
     if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF:
         import pyarrow as pa
@@ -2740,9 +2745,11 @@ def extract_args(batch):
             args = tuple(batch.column(o) for o in args_offsets)
             return args[0] if len(args) == 1 else args
 
-        def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
-            for group_batches in batches:
-                batch_iter = map(extract_args, group_batches)
+        def grouped_func(
+            split_index: int, data: Iterator["GroupedBatch"]
+        ) -> Iterator[pa.RecordBatch]:
+            for group in data:
+                batch_iter = map(extract_args, group)
                 result = udf_func(batch_iter)
                 # Drain remaining batches to maintain stream position
                 for _ in batch_iter:
@@ -2751,7 +2758,7 @@ def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
                 yield ArrowBatchTransformer.enforce_schema(batch, return_schema)
 
         # profiling is not supported for UDF
-        return func, None, ser, ser
+        return grouped_func, None, ser, ser
 
     if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF:
         import pyarrow as pa
@@ -2766,9 +2773,11 @@ def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
             prefers_large_types=runner_conf.use_large_var_types,
         )
 
-        def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
-            for group_batches in batches:
-                batch_list = list(group_batches)
+        def grouped_func(
+            split_index: int, data: Iterator["GroupedBatch"]
+        ) -> Iterator[pa.RecordBatch]:
+            for group in data:
+                batch_list = list(group)
                 if not batch_list:
                     continue
                 if hasattr(pa, "concat_batches"):
@@ -2817,7 +2826,7 @@ def func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
                 yield ArrowBatchTransformer.enforce_schema(batch, return_schema)
 
         # profiling is not supported for UDF
-        return func, None, ser, ser
+        return grouped_func, None, ser, ser
 
     if (
         eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF
@@ -2868,8 +2877,8 @@ def _evaluate_batch_udf(udf_func, rows):
             with ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool:
                 return list(pool.map(lambda row: udf_func(*row), rows))
 
-        def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
-            for input_batch in batches:
+        def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+            for input_batch in data:
                 num_rows = input_batch.num_rows
 
                 # --- Input: Arrow -> Python columns ---
@@ -2953,8 +2962,8 @@ def _evaluate_batch_udf_legacy(udf_func, rows):
             with ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool:
                 return list(pool.map(lambda row: udf_func(*row), rows))
 
-        def func(split_index: int, batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
-            for input_batch in batches:
+        def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+            for input_batch in data:
                 # --- Input: Arrow -> pandas columns ---
                 pandas_columns = ArrowBatchTransformer.to_pandas(
                     input_batch,