feat: introduce iterator API for Arrow grouped agg UDF

Author: Yicong-Huang

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

@@ -74,6 +74,7 @@ private[spark] object PythonEvalType {
   val SQL_SCALAR_ARROW_ITER_UDF = 251
   val SQL_GROUPED_AGG_ARROW_UDF = 252
   val SQL_WINDOW_AGG_ARROW_UDF = 253
+  val SQL_GROUPED_AGG_ARROW_ITER_UDF = 254
 
   val SQL_TABLE_UDF = 300
   val SQL_ARROW_TABLE_UDF = 301
@@ -111,6 +112,7 @@ private[spark] object PythonEvalType {
     case SQL_SCALAR_ARROW_ITER_UDF => "SQL_SCALAR_ARROW_ITER_UDF"
     case SQL_GROUPED_AGG_ARROW_UDF => "SQL_GROUPED_AGG_ARROW_UDF"
     case SQL_WINDOW_AGG_ARROW_UDF => "SQL_WINDOW_AGG_ARROW_UDF"
+    case SQL_GROUPED_AGG_ARROW_ITER_UDF => "SQL_GROUPED_AGG_ARROW_ITER_UDF"
   }
 }
 

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

@@ -66,6 +66,7 @@ ArrowScalarUDFType = Literal[250]
 ArrowScalarIterUDFType = Literal[251]
 ArrowGroupedAggUDFType = Literal[252]
 ArrowWindowAggUDFType = Literal[253]
+ArrowGroupedAggIterUDFType = Literal[254]
 
 class ArrowVariadicScalarToScalarFunction(Protocol):
     def __call__(self, *_: pyarrow.Array) -> pyarrow.Array: ...

python/pyspark/sql/pandas/functions.py

@@ -50,6 +50,8 @@ class ArrowUDFType:
 
     GROUPED_AGG = PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF
 
+    GROUPED_AGG_ITER = PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF
+
 
 def arrow_udf(f=None, returnType=None, functionType=None):
     """
@@ -301,6 +303,66 @@ def calculate(iterator: Iterator[pa.Array]) -> Iterator[pa.Array]:
             Therefore, mutating the input arrays is not allowed and will cause incorrect results.
             For the same reason, users should also not rely on the index of the input arrays.
 
+    * Iterator of Arrays to Scalar
+        `Iterator[pyarrow.Array]` -> `Any`
+
+        The function takes an iterator of `pyarrow.Array` and returns a scalar value. This is
+        useful for grouped aggregations where the UDF can process all batches for a group
+        iteratively, which is more memory-efficient than loading all data at once. The returned
+        scalar can be a python primitive type, a numpy data type, or a `pyarrow.Scalar` instance.
+
+        >>> import pandas as pd
+        >>> from typing import Iterator
+        >>> @arrow_udf("double")
+        ... def arrow_mean(it: Iterator[pa.Array]) -> float:
+        ...     sum_val = 0.0
+        ...     cnt = 0
+        ...     for v in it:
+        ...         assert isinstance(v, pa.Array)
+        ...         sum_val += pa.compute.sum(v).as_py()
+        ...         cnt += len(v)
+        ...     return sum_val / cnt
+        ...
+        >>> df = spark.createDataFrame(
+        ...     [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
+        >>> df.groupby("id").agg(arrow_mean(df['v'])).show()  # doctest: +SKIP
+        +---+---------------+
+        | id|arrow_mean(v) |
+        +---+---------------+
+        |  1|           1.5|
+        |  2|           6.0|
+        +---+---------------+
+
+    * Iterator of Multiple Arrays to Scalar
+        `Iterator[Tuple[pyarrow.Array, ...]]` -> `Any`
+
+        The function takes an iterator of a tuple of multiple `pyarrow.Array` and returns a
+        scalar value. This is useful for grouped aggregations with multiple input columns.
+
+        >>> from typing import Iterator, Tuple
+        >>> import numpy as np
+        >>> @arrow_udf("double")
+        ... def arrow_weighted_mean(it: Iterator[Tuple[pa.Array, pa.Array]]) -> float:
+        ...     weighted_sum = 0.0
+        ...     weight = 0.0
+        ...     for v, w in it:
+        ...         assert isinstance(v, pa.Array)
+        ...         assert isinstance(w, pa.Array)
+        ...         weighted_sum += np.dot(v, w)
+        ...         weight += pa.compute.sum(w).as_py()
+        ...     return weighted_sum / weight
+        ...
+        >>> df = spark.createDataFrame(
+        ...     [(1, 1.0, 1.0), (1, 2.0, 2.0), (2, 3.0, 1.0), (2, 5.0, 2.0), (2, 10.0, 3.0)],
+        ...     ("id", "v", "w"))
+        >>> df.groupby("id").agg(arrow_weighted_mean(df["v"], df["w"])).show()  # doctest: +SKIP
+        +---+---------------------------------+
+        | id|arrow_weighted_mean(v, w)        |
+        +---+---------------------------------+
+        |  1|               1.6666666666666667|
+        |  2|                7.166666666666667|
+        +---+---------------------------------+
+
     Notes
     -----
     The user-defined functions do not support conditional expressions or short circuiting
@@ -720,6 +782,7 @@ def vectorized_udf(
         PythonEvalType.SQL_SCALAR_ARROW_UDF,
         PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
         PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
+        PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
         None,
     ]:  # None means it should infer the type from type hints.
         raise PySparkTypeError(
@@ -768,6 +831,7 @@ def _validate_vectorized_udf(f, evalType, kind: str = "pandas") -> int:
         PythonEvalType.SQL_SCALAR_ARROW_UDF,
         PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
         PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
+        PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
     ]:
         warnings.warn(
             "It is preferred to specify type hints for "

python/pyspark/sql/pandas/serializers.py

@@ -1200,6 +1200,72 @@ def __repr__(self):
         return "ArrowStreamAggArrowUDFSerializer"
 
 
+# Serializer for SQL_GROUPED_AGG_ARROW_ITER_UDF
+class ArrowStreamAggArrowIterUDFSerializer(ArrowStreamArrowUDFSerializer):
+    def __init__(
+        self,
+        timezone,
+        safecheck,
+        assign_cols_by_name,
+        arrow_cast,
+    ):
+        super().__init__(
+            timezone=timezone,
+            safecheck=safecheck,
+            assign_cols_by_name=False,
+            arrow_cast=True,
+        )
+        self._timezone = timezone
+        self._safecheck = safecheck
+        self._assign_cols_by_name = assign_cols_by_name
+        self._arrow_cast = arrow_cast
+
+    def load_stream(self, stream):
+        """
+        Yield column iterators instead of concatenating batches.
+        Each group yields a structure where indexing by column offset gives an iterator of arrays.
+        """
+        import pyarrow as pa
+
+        dataframes_in_group = None
+
+        while dataframes_in_group is None or dataframes_in_group > 0:
+            dataframes_in_group = read_int(stream)
+
+            if dataframes_in_group == 1:
+                batches = list(ArrowStreamSerializer.load_stream(self, stream))
+                # Create a structure that can be indexed by column offset to get column iterators
+                # The mapper will do a[offset] to get each column's iterator
+                if len(batches) > 0:
+                    num_cols = batches[0].num_columns
+
+                    # Create a custom class that can be indexed to get column iterators
+                    class ColumnIterators:
+                        def __init__(self, batches, num_cols):
+                            self._batches = batches
+                            self._num_cols = num_cols
+
+                        def __getitem__(self, col_idx):
+                            return (batch.column(col_idx) for batch in self._batches)
+
+                        def __len__(self):
+                            return self._num_cols
+
+                    yield ColumnIterators(batches, num_cols)
+                else:
+                    # Empty group
+                    yield []
+
+            elif dataframes_in_group != 0:
+                raise PySparkValueError(
+                    errorClass="INVALID_NUMBER_OF_DATAFRAMES_IN_GROUP",
+                    messageParameters={"dataframes_in_group": str(dataframes_in_group)},
+                )
+
+    def __repr__(self):
+        return "ArrowStreamAggArrowIterUDFSerializer"
+
+
 class GroupPandasUDFSerializer(ArrowStreamPandasUDFSerializer):
     def __init__(
         self,

python/pyspark/sql/pandas/typehints.py

@@ -29,6 +29,7 @@
         ArrowScalarUDFType,
         ArrowScalarIterUDFType,
         ArrowGroupedAggUDFType,
+        ArrowGroupedAggIterUDFType,
         ArrowGroupedMapIterUDFType,
         ArrowGroupedMapUDFType,
         ArrowGroupedMapFunction,
@@ -156,7 +157,14 @@ def infer_pandas_eval_type(
 
 def infer_arrow_eval_type(
     sig: Signature, type_hints: Dict[str, Any]
-) -> Optional[Union["ArrowScalarUDFType", "ArrowScalarIterUDFType", "ArrowGroupedAggUDFType"]]:
+) -> Optional[
+    Union[
+        "ArrowScalarUDFType",
+        "ArrowScalarIterUDFType",
+        "ArrowGroupedAggUDFType",
+        "ArrowGroupedAggIterUDFType",
+    ]
+]:
     """
     Infers the evaluation type in :class:`pyspark.util.PythonEvalType` from
     :class:`inspect.Signature` instance and type hints.
@@ -226,6 +234,41 @@ def infer_arrow_eval_type(
     if is_iterator_array:
         return ArrowUDFType.SCALAR_ITER
 
+    # Iterator[Tuple[pa.Array, ...]] -> Any
+    is_iterator_tuple_array_agg = (
+        len(parameters_sig) == 1
+        and check_iterator_annotation(  # Iterator
+            parameters_sig[0],
+            parameter_check_func=lambda a: check_tuple_annotation(  # Tuple
+                a,
+                parameter_check_func=lambda ta: (ta == Ellipsis or ta == pa.Array),
+            ),
+        )
+        and (
+            return_annotation != pa.Array
+            and not check_iterator_annotation(return_annotation)
+            and not check_tuple_annotation(return_annotation)
+        )
+    )
+    if is_iterator_tuple_array_agg:
+        return ArrowUDFType.GROUPED_AGG_ITER
+
+    # Iterator[pa.Array] -> Any
+    is_iterator_array_agg = (
+        len(parameters_sig) == 1
+        and check_iterator_annotation(
+            parameters_sig[0],
+            parameter_check_func=lambda a: a == pa.Array,
+        )
+        and (
+            return_annotation != pa.Array
+            and not check_iterator_annotation(return_annotation)
+            and not check_tuple_annotation(return_annotation)
+        )
+    )
+    if is_iterator_array_agg:
+        return ArrowUDFType.GROUPED_AGG_ITER
+
     # pa.Array, ... -> Any
     is_array_agg = all(a == pa.Array for a in parameters_sig) and (
         return_annotation != pa.Array
@@ -249,6 +292,7 @@ def infer_eval_type(
     "ArrowScalarUDFType",
     "ArrowScalarIterUDFType",
     "ArrowGroupedAggUDFType",
+    "ArrowGroupedAggIterUDFType",
 ]:
     """
     Infers the evaluation type in :class:`pyspark.util.PythonEvalType` from
@@ -264,6 +308,7 @@ def infer_eval_type(
             "ArrowScalarUDFType",
             "ArrowScalarIterUDFType",
             "ArrowGroupedAggUDFType",
+            "ArrowGroupedAggIterUDFType",
         ]
     ] = None
     if kind == "pandas":
@@ -295,6 +340,7 @@ def infer_eval_type_for_udf(  # type: ignore[no-untyped-def]
         "ArrowScalarUDFType",
         "ArrowScalarIterUDFType",
         "ArrowGroupedAggUDFType",
+        "ArrowGroupedAggIterUDFType",
     ]
 ]:
     argspec = getfullargspec(f)

python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py

@@ -1059,6 +1059,90 @@ def my_grouped_agg_arrow_udf(x):
             ],
         )
 
+    def test_iterator_grouped_agg_single_column(self):
+        """
+        Test iterator API for grouped aggregation with single column.
+        """
+        import pyarrow as pa
+        from typing import Iterator
+
+        @arrow_udf("double")
+        def arrow_mean_iter(it: Iterator[pa.Array]) -> float:
+            sum_val = 0.0
+            cnt = 0
+            for v in it:
+                assert isinstance(v, pa.Array)
+                sum_val += pa.compute.sum(v).as_py()
+                cnt += len(v)
+            return sum_val / cnt if cnt > 0 else 0.0
+
+        df = self.spark.createDataFrame(
+            [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v")
+        )
+
+        result = df.groupby("id").agg(arrow_mean_iter(df["v"]).alias("mean")).sort("id")
+        expected = df.groupby("id").agg(sf.mean(df["v"]).alias("mean")).sort("id").collect()
+
+        self.assertEqual(expected, result.collect())
+
+    @unittest.skipIf(not have_numpy, numpy_requirement_message)
+    def test_iterator_grouped_agg_multiple_columns(self):
+        """
+        Test iterator API for grouped aggregation with multiple columns.
+        """
+        import pyarrow as pa
+        import numpy as np
+        from typing import Iterator, Tuple
+
+        @arrow_udf("double")
+        def arrow_weighted_mean_iter(it: Iterator[Tuple[pa.Array, pa.Array]]) -> float:
+            weighted_sum = 0.0
+            weight = 0.0
+            for v, w in it:
+                assert isinstance(v, pa.Array)
+                assert isinstance(w, pa.Array)
+                weighted_sum += np.dot(v, w)
+                weight += pa.compute.sum(w).as_py()
+            return weighted_sum / weight if weight > 0 else 0.0
+
+        df = self.spark.createDataFrame(
+            [(1, 1.0, 1.0), (1, 2.0, 2.0), (2, 3.0, 1.0), (2, 5.0, 2.0), (2, 10.0, 3.0)],
+            ("id", "v", "w"),
+        )
+
+        result = (
+            df.groupby("id")
+            .agg(arrow_weighted_mean_iter(df["v"], df["w"]).alias("wm"))
+            .sort("id")
+            .collect()
+        )
+
+        # Expected weighted means:
+        # Group 1: (1.0*1.0 + 2.0*2.0) / (1.0 + 2.0) = 5.0 / 3.0 = 1.6666666666666667
+        # Group 2: (3.0*1.0 + 5.0*2.0 + 10.0*3.0) / (1.0 + 2.0 + 3.0) = 43.0 / 6.0 = 7.166666666666667
+        expected = [(1, 5.0 / 3.0), (2, 43.0 / 6.0)]
+
+        self.assertEqual(len(result), len(expected))
+        for r, (exp_id, exp_wm) in zip(result, expected):
+            self.assertEqual(r["id"], exp_id)
+            self.assertAlmostEqual(r["wm"], exp_wm, places=5)
+
+    def test_iterator_grouped_agg_eval_type(self):
+        """
+        Test that the eval type is correctly inferred for iterator grouped agg UDFs.
+        """
+        import pyarrow as pa
+        from typing import Iterator
+
+        @arrow_udf("double")
+        def arrow_sum_iter(it: Iterator[pa.Array]) -> float:
+            total = 0.0
+            for v in it:
+                total += pa.compute.sum(v).as_py()
+            return total
+
+        self.assertEqual(arrow_sum_iter.evalType, PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF)
+
 
 class GroupedAggArrowUDFTests(GroupedAggArrowUDFTestsMixin, ReusedSQLTestCase):
     pass

python/pyspark/util.py

@@ -70,6 +70,7 @@
         ArrowScalarUDFType,
         ArrowScalarIterUDFType,
         ArrowGroupedAggUDFType,
+        ArrowGroupedAggIterUDFType,
         ArrowWindowAggUDFType,
     )
     from pyspark.sql._typing import (
@@ -660,6 +661,7 @@ class PythonEvalType:
     SQL_SCALAR_ARROW_ITER_UDF: "ArrowScalarIterUDFType" = 251
     SQL_GROUPED_AGG_ARROW_UDF: "ArrowGroupedAggUDFType" = 252
     SQL_WINDOW_AGG_ARROW_UDF: "ArrowWindowAggUDFType" = 253
+    SQL_GROUPED_AGG_ARROW_ITER_UDF: "ArrowGroupedAggIterUDFType" = 254
 
     SQL_TABLE_UDF: "SQLTableUDFType" = 300
     SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301