[SPARK-59704][PYTHON] Migrate Arrow aggregate and window UDF eval types to the EvalTypeHandler pipeline

python/pyspark/eval_handlers/_arrow.py

@@ -31,6 +31,7 @@
 from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
+from pyspark.errors import PySparkRuntimeError
 from pyspark.eval_handlers._base import (
     BatchEvalTypeHandler,
     CoGroupedEvalTypeHandler,
@@ -250,6 +251,163 @@ def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.Rec
                 yield ArrowBatchTransformer.wrap_struct(batch)
 
 
+def _concat_group_batches(batch_list: list["pa.RecordBatch"]) -> "pa.RecordBatch":
+    """Concatenate a group's RecordBatches into a single one, with a fallback for
+    pyarrow before 19.0.0 (which lacks ``pa.concat_batches``). Remove the fallback
+    once support for those versions is dropped."""
+    import pyarrow as pa
+
+    if hasattr(pa, "concat_batches"):
+        return pa.concat_batches(batch_list)
+    return pa.RecordBatch.from_struct_array(
+        pa.concat_arrays([b.to_struct_array() for b in batch_list])
+    )
+
+
+class ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+    """SQL_GROUPED_AGG_ARROW_UDF: each UDF reduces its input columns over the whole
+    group to a single scalar; emit one row per group with one column per UDF,
+    coerced to the declared schema."""
+
+    eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF
+
+    def __init__(
+        self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
+    ) -> None:
+        require_minimum_pyarrow_version()
+        super().__init__(udfs, runner_conf, eval_conf)
+        self._col_names = ["_%d" % i for i in range(len(udfs))]
+        self._return_schema = to_arrow_schema(
+            StructType([StructField(n, rt) for n, (_, _, _, rt) in zip(self._col_names, udfs)]),
+            timezone="UTC",
+            prefers_large_types=runner_conf.use_large_var_types,
+        )
+
+    def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]:
+        import pyarrow as pa
+
+        for group in data:
+            batch_list = list(group)
+            if not batch_list:
+                continue
+            concatenated = _concat_group_batches(batch_list)
+            results = [
+                udf_func(
+                    *[concatenated.column(o) for o in args_offsets],
+                    **{k: concatenated.column(v) for k, v in kwargs_offsets.items()},
+                )
+                for udf_func, args_offsets, kwargs_offsets, _ in self._udfs
+            ]
+            result_arrays = [pa.array([r]) for r in results]
+            batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names)
+            yield ArrowBatchTransformer.enforce_schema(batch, self._return_schema)
+
+
+class ArrowGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+    """SQL_GROUPED_AGG_ARROW_ITER_UDF: the single UDF receives each group as an
+    iterator of its input columns and returns one scalar; emit one row per group."""
+
+    eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF
+
+    def __init__(
+        self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
+    ) -> None:
+        require_minimum_pyarrow_version()
+        super().__init__(udfs, runner_conf, eval_conf)
+        assert len(udfs) == 1, "One GROUPED_AGG_ARROW_ITER UDF expected here."
+        self._udf_func, self._args_offsets, _, return_type = udfs[0]
+        self._return_schema = to_arrow_schema(
+            StructType([StructField("_0", return_type)]),
+            timezone="UTC",
+            prefers_large_types=runner_conf.use_large_var_types,
+        )
+
+    def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]:
+        import pyarrow as pa
+
+        args_offsets = self._args_offsets
+
+        def extract_args(batch: pa.RecordBatch) -> Any:
+            args = tuple(batch.column(o) for o in args_offsets)
+            return args[0] if len(args) == 1 else args
+
+        for group in data:
+            batch_iter = map(extract_args, group)
+            result = self._udf_func(batch_iter)
+            # Drain remaining batches to maintain stream position
+            for _ in batch_iter:
+                pass
+            batch = pa.RecordBatch.from_arrays([pa.array([result])], ["_0"])
+            yield ArrowBatchTransformer.enforce_schema(batch, self._return_schema)
+
+
+class ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+    """SQL_WINDOW_AGG_ARROW_UDF: each UDF produces one value per input row over its
+    window frame -- an unbounded frame is computed once and repeated for every row,
+    a bounded frame slices the input columns per row -- emitted as one column per
+    UDF, coerced to the declared schema."""
+
+    eval_type = PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF
+
+    def __init__(
+        self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
+    ) -> None:
+        require_minimum_pyarrow_version()
+        super().__init__(udfs, runner_conf, eval_conf)
+        self._window_bound_types = runner_conf.window_bound_types
+        self._col_names = ["_%d" % i for i in range(len(udfs))]
+        self._return_schema = to_arrow_schema(
+            StructType([StructField(n, rt) for n, (_, _, _, rt) in zip(self._col_names, udfs)]),
+            timezone="UTC",
+            prefers_large_types=runner_conf.use_large_var_types,
+        )
+
+    def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]:
+        import pyarrow as pa
+
+        for group in data:
+            batch_list = list(group)
+            if not batch_list:
+                continue
+            concatenated = _concat_group_batches(batch_list)
+            num_rows = concatenated.num_rows
+
+            result_arrays = []
+            for udf_index, (udf_func, args_offsets, kwargs_offsets, _) in enumerate(self._udfs):
+                bound_type = self._window_bound_types[udf_index]
+                if bound_type == "unbounded":
+                    result = udf_func(
+                        *[concatenated.column(o) for o in args_offsets],
+                        **{k: concatenated.column(v) for k, v in kwargs_offsets.items()},
+                    )
+                    result_arrays.append(pa.repeat(result, num_rows))
+                elif bound_type == "bounded":
+                    begin_col = concatenated.column(args_offsets[0])
+                    end_col = concatenated.column(args_offsets[1])
+                    results = []
+                    for i in range(num_rows):
+                        offset = begin_col[i].as_py()
+                        length = end_col[i].as_py() - offset
+                        slices = [
+                            concatenated.column(o).slice(offset=offset, length=length)
+                            for o in args_offsets[2:]
+                        ]
+                        kw_slices = {
+                            k: concatenated.column(v).slice(offset=offset, length=length)
+                            for k, v in kwargs_offsets.items()
+                        }
+                        results.append(udf_func(*slices, **kw_slices))
+                    result_arrays.append(pa.array(results))
+                else:
+                    raise PySparkRuntimeError(
+                        errorClass="INVALID_WINDOW_BOUND_TYPE",
+                        messageParameters={"window_bound_type": bound_type},
+                    )
+
+            batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names)
+            yield ArrowBatchTransformer.enforce_schema(batch, self._return_schema)
+
+
 class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
     """SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input
     RecordBatch stream and yields a RecordBatch stream, exchanged as flattened

python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py

@@ -36,11 +36,14 @@
 
     from pyspark.eval_handlers._arrow import (
         ArrowCoGroupedMapUDFHandler,
+        ArrowGroupedAggIterUDFHandler,
+        ArrowGroupedAggUDFHandler,
         ArrowGroupedMapIterUDFHandler,
         ArrowGroupedMapUDFHandler,
         ArrowMapUDFHandler,
         ArrowScalarIterUDFHandler,
         ArrowScalarUDFHandler,
+        ArrowWindowAggUDFHandler,
     )
     from pyspark.sql.conversion import ArrowBatchTransformer
 
@@ -52,6 +55,8 @@ class _RunnerConf:
     use_large_var_types = False
     assign_cols_by_name = True
     map_in_batch_legacy_accept_any_iterable = False
+    # One bound type per window UDF; the bounded test overrides this per instance.
+    window_bound_types = ["unbounded"]
 
 
 def _batch(**columns):
@@ -120,14 +125,17 @@ def _grouped_handler(handler_cls, udf, arg_offsets, num_udf_args):
 @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
 class ArrowEvalTypeHandlerRegistrationTests(unittest.TestCase):
     def test_arrow_eval_types_are_registered(self):
-        # Every migrated Arrow map/iter eval type dispatches to its handler by lookup.
+        # Every migrated Arrow eval type dispatches to its handler by lookup.
         for eval_type, handler_cls in (
             (PythonEvalType.SQL_SCALAR_ARROW_UDF, ArrowScalarUDFHandler),
             (PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, ArrowScalarIterUDFHandler),
             (PythonEvalType.SQL_MAP_ARROW_ITER_UDF, ArrowMapUDFHandler),
             (PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF, ArrowGroupedMapUDFHandler),
             (PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF, ArrowGroupedMapIterUDFHandler),
             (PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, ArrowCoGroupedMapUDFHandler),
+            (PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, ArrowGroupedAggUDFHandler),
+            (PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, ArrowGroupedAggIterUDFHandler),
+            (PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, ArrowWindowAggUDFHandler),
         ):
             self.assertIs(get_eval_type_handler(eval_type), handler_cls)
 
@@ -255,6 +263,65 @@ def cogrouped_udf(key, left_values, right_values):
         self.assertEqual(out[0].column(0).field("v").to_pylist(), [35])
 
 
+@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
+class ArrowGroupedAggUDFHandlerTests(unittest.TestCase):
+    # Grouped-agg batches arrive un-wrapped (flat columns); the UDF reads columns by offset and
+    # returns one scalar, emitted as a single-row batch per group.
+    def test_reduces_group_to_one_row(self):
+        def sum_udf(col):
+            return sum(c.as_py() for c in col)
+
+        handler = ArrowGroupedAggUDFHandler(
+            udfs=[(sum_udf, [0], {}, LongType())], runner_conf=_RunnerConf(), eval_conf=None
+        )
+        out = list(handler.run(0, _one_group(_batch(v=[1, 2, 3]), _batch(v=[4]))))
+        self.assertEqual(out[0].column("_0").to_pylist(), [10])
+
+
+@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
+class ArrowGroupedAggIterUDFHandlerTests(unittest.TestCase):
+    # The UDF receives the group's input columns as an iterator and returns one scalar.
+    def test_reduces_group_to_one_row(self):
+        def sum_iter_udf(col_iter):
+            return sum(c.as_py() for col in col_iter for c in col)
+
+        handler = ArrowGroupedAggIterUDFHandler(
+            udfs=[(sum_iter_udf, [0], {}, LongType())], runner_conf=_RunnerConf(), eval_conf=None
+        )
+        out = list(handler.run(0, _one_group(_batch(v=[10]), _batch(v=[20]))))
+        self.assertEqual(out[0].column("_0").to_pylist(), [30])
+
+
+@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
+class ArrowWindowAggUDFHandlerTests(unittest.TestCase):
+    # One output value per input row over the UDF's window frame.
+    def test_unbounded_frame_repeats_one_value(self):
+        def sum_udf(col):
+            return sum(c.as_py() for c in col)
+
+        conf = _RunnerConf()
+        conf.window_bound_types = ["unbounded"]
+        handler = ArrowWindowAggUDFHandler(
+            udfs=[(sum_udf, [0], {}, LongType())], runner_conf=conf, eval_conf=None
+        )
+        out = list(handler.run(0, _one_group(_batch(v=[1, 2, 3]))))
+        self.assertEqual(out[0].column("_0").to_pylist(), [6, 6, 6])
+
+    def test_bounded_frame_slices_per_row(self):
+        # args_offsets = [begin_col, end_col, *value_cols]; each row's frame is ``[begin, end)``.
+        def sum_udf(col):
+            return sum(c.as_py() for c in col)
+
+        conf = _RunnerConf()
+        conf.window_bound_types = ["bounded"]
+        handler = ArrowWindowAggUDFHandler(
+            udfs=[(sum_udf, [0, 1, 2], {}, LongType())], runner_conf=conf, eval_conf=None
+        )
+        # Row 0 frame [0, 1) -> [10]; row 1 frame [0, 2) -> [10, 20].
+        out = list(handler.run(0, _one_group(_batch(begin=[0, 0], end=[1, 2], v=[10, 20]))))
+        self.assertEqual(out[0].column("_0").to_pylist(), [10, 30])
+
+
 @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
 class CoGroupedBatchTests(unittest.TestCase):
     def test_deserialized_co_group_is_a_pair_of_lists(self):