You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/10/13 14:33:03 UTC
(spark) branch branch-4.x updated: [SPARK-60208][PYTHON] Migrate grouped aggregate and window pandas UDF eval types to pandas eval type handlers
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new dfa929d2f0e2 [SPARK-60208][PYTHON] Migrate grouped aggregate and window pandas UDF eval types to pandas eval type handlers
dfa929d2f0e2 is described below
commit dfa929d2f0e25a398f5ddb16e069bfaddede32ac
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Tue Oct 13 14:32:17 2026 +0000
[SPARK-60208][PYTHON] Migrate grouped aggregate and window pandas UDF eval types to pandas eval type handlers
### What changes were proposed in this pull request?
Continues the eval-type-handler migration (umbrella SPARK-59415), following SPARK-60100. Moves three pandas aggregate/window eval types off the `worker.py::read_udfs` if/elif chain onto handler classes in `python/pyspark/eval_handlers/_pandas.py`:
- `PandasGroupedAggUDFHandler` for `SQL_GROUPED_AGG_PANDAS_UDF`.
- `PandasGroupedAggIterUDFHandler` for `SQL_GROUPED_AGG_PANDAS_ITER_UDF`, keeping the drain of unread group batches and the existing `prefers_large_types=False` on output.
- `PandasWindowAggUDFHandler` for `SQL_WINDOW_AGG_PANDAS_UDF`, keeping the unbounded (repeat) and bounded (per-row `iloc` slice) frame paths.
Each is the pandas counterpart of the matching Arrow handler and calls `ArrowToPandasConversion` / `PandasToArrowConversion` inline. The corresponding branches, serializer-selection entries and now-unused imports are removed from `worker.py`.
```
EvalTypeHandler
GroupedEvalTypeHandler
PandasGroupedAggUDFHandler (SQL_GROUPED_AGG_PANDAS_UDF, new)
PandasGroupedAggIterUDFHandler (SQL_GROUPED_AGG_PANDAS_ITER_UDF, new)
PandasWindowAggUDFHandler (SQL_WINDOW_AGG_PANDAS_UDF, new)
```
### Why are the changes needed?
Part of SPARK-59415: replace the `read_udfs` if/elif chain with a per-eval-type handler registry so each eval type's logic lives in one directly-testable class.
### Does this PR introduce _any_ user-facing change?
No. Internal refactor; UDF API, wire format and behavior are unchanged.
### How was this patch tested?
- New unit tests in `python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py` (registration, multiple UDFs, kwargs, empty groups, draining, large-var-types handling, bounded/unbounded/mixed frames, invalid bound type).
- Existing suites pass: `pyspark.sql.tests.pandas.test_pandas_udf_grouped_agg`, `test_pandas_udf_window`.
- ASV `bench_eval_type`, master vs. this PR, one representative run:
```text
bench scenario udf before after diff
GroupedAggPandasUDF few_groups_sm sum_udf 17.1+-0.2ms 16.7+-0.1ms -2%
GroupedAggPandasUDF many_groups_sm sum_udf 589+-4ms 582+-0.8ms -1%
GroupedAggPandasUDF many_groups_lg mean_multi_udf 190+-0.8ms 186+-1ms -2%
GroupedAggPandasIter many_groups_sm sum_udf 573+-8ms 563+-5ms -2%
GroupedAggPandasIter wide_cols mean_multi_udf 164+-1ms 162+-2ms -1%
WindowAggPandasUDF few_groups_lg sum_udf 41.9+-0.4ms 41.8+-0.4ms 0%
WindowAggPandasUDF many_groups_sm sum_udf 668+-2ms 667+-2ms 0%
WindowAggPandasUDF many_groups_lg sum_udf 226+-1ms 229+-2ms +1%
```
All other scenarios are also within +-2%.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
This pull request and its description were written by Isaac.
Closes #59419 from Yicong-Huang/pandas-agg-window.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 8ac4e179476f550ad0e1a54620c01a106f9e87b0)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/eval_handlers/_pandas.py | 199 ++++++++++++++++++++-
.../tests/test_pandas_eval_type_handlers.py | 184 ++++++++++++++++++-
python/pyspark/worker.py | 193 +-------------------
3 files changed, 379 insertions(+), 197 deletions(-)
diff --git a/python/pyspark/eval_handlers/_pandas.py b/python/pyspark/eval_handlers/_pandas.py
index 4a3e0870e795..c4fa685568c0 100644
--- a/python/pyspark/eval_handlers/_pandas.py
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -31,7 +31,7 @@ from __future__ import annotations
from collections.abc import Iterator
from typing import TYPE_CHECKING, Any
-from pyspark.errors import PySparkTypeError, PySparkValueError
+from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.eval_handlers._base import (
BatchEvalTypeHandler,
CoGroupedEvalTypeHandler,
@@ -521,3 +521,200 @@ class PandasGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
int_to_decimal_coercion_enabled=self._runner_conf.int_to_decimal_coercion_enabled,
)
del result
+
+
+class PandasGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_GROUPED_AGG_PANDAS_UDF: each UDF reduces its input columns (pandas Series)
+ over the whole group to a single scalar; emit one row per group with one column
+ per UDF, converted back to Arrow. The pandas counterpart of
+ ArrowGroupedAggUDFHandler."""
+
+ eval_type = PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
+ ) -> None:
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ self._return_schema = StructType(
+ [StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
+ )
+
+ def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]":
+ import pandas as pd
+ import pyarrow as pa
+
+ for group in data:
+ batch_list = list(group)
+ if not batch_list:
+ continue
+ table = pa.Table.from_batches(batch_list).combine_chunks()
+ all_series = ArrowToPandasConversion.to_pandas(
+ table,
+ timezone=self._runner_conf.timezone,
+ prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
+ )
+ results = [
+ udf_func(
+ *[all_series[o] for o in args_offsets],
+ **{k: all_series[v] for k, v in kwargs_offsets.items()},
+ )
+ for udf_func, args_offsets, kwargs_offsets, _ in self._udfs
+ ]
+ yield PandasToArrowConversion.from_pandas(
+ [pd.Series([r]) for r in results],
+ self._return_schema,
+ timezone=self._runner_conf.timezone,
+ safecheck=self._runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=self._runner_conf.use_large_var_types,
+ assign_cols_by_name=self._runner_conf.assign_cols_by_name,
+ int_to_decimal_coercion_enabled=self._runner_conf.int_to_decimal_coercion_enabled,
+ )
+
+
+class PandasGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_GROUPED_AGG_PANDAS_ITER_UDF: the single UDF receives each group as an
+ iterator of its input columns (one pandas Series, or a tuple of Series, per input
+ batch) and returns one scalar; emit one row per group. The pandas counterpart of
+ ArrowGroupedAggIterUDFHandler."""
+
+ eval_type = PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
+ ) -> None:
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ assert len(udfs) == 1, "One GROUPED_AGG_PANDAS_ITER UDF expected here."
+ self._udf_func, self._args_offsets, _, return_type = udfs[0]
+ self._return_schema = StructType([StructField("_0", return_type)])
+
+ def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]":
+ import pandas as pd
+
+ args_offsets = self._args_offsets
+
+ def extract_series(batch: "pa.RecordBatch") -> Any:
+ # Convert one RecordBatch to a pandas Series per column, then select args:
+ # - pd.Series for a single column
+ # - tuple[pd.Series, ...] for multiple columns
+ all_series = ArrowToPandasConversion.to_pandas(
+ batch,
+ timezone=self._runner_conf.timezone,
+ prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
+ )
+ series = tuple(all_series[o] for o in args_offsets)
+ return series[0] if len(series) == 1 else series
+
+ for group in data:
+ series_iter = map(extract_series, group)
+ result = self._udf_func(series_iter)
+ # Drain remaining batches to maintain stream position
+ for _ in series_iter:
+ pass
+ yield PandasToArrowConversion.from_pandas(
+ [pd.Series([result])],
+ self._return_schema,
+ timezone=self._runner_conf.timezone,
+ safecheck=self._runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=False,
+ assign_cols_by_name=self._runner_conf.assign_cols_by_name,
+ int_to_decimal_coercion_enabled=self._runner_conf.int_to_decimal_coercion_enabled,
+ )
+
+
+class PandasWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_WINDOW_AGG_PANDAS_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 Series per row -- emitted as one column per UDF,
+ converted back to Arrow. The pandas counterpart of ArrowWindowAggUDFHandler."""
+
+ eval_type = PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
+ ) -> None:
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ self._window_bound_types = runner_conf.window_bound_types
+ self._return_schema = StructType(
+ [StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
+ )
+
+ def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]":
+ import pandas as pd
+ import pyarrow as pa
+
+ for group in data:
+ batch_list = list(group)
+ if not batch_list:
+ continue
+ table = pa.Table.from_batches(batch_list).combine_chunks()
+ all_series = ArrowToPandasConversion.to_pandas(
+ table,
+ timezone=self._runner_conf.timezone,
+ prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
+ )
+ num_rows = table.num_rows
+
+ result_series = []
+ 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(
+ *[all_series[o] for o in args_offsets],
+ **{k: all_series[v] for k, v in kwargs_offsets.items()},
+ )
+ # Repeat the scalar result to match the window (group) length.
+ result_series.append(pd.Series([result]).repeat(num_rows))
+ elif bound_type == "bounded":
+ # args_offsets[0] and args_offsets[1] are begin_index and end_index.
+ assert len(args_offsets) >= 2, len(args_offsets)
+ # Index operation is faster on np.ndarray, so we turn the
+ # index series into np arrays here for performance.
+ begin_array = all_series[args_offsets[0]].values
+ end_array = all_series[args_offsets[1]].values
+ series = [all_series[o] for o in args_offsets[2:]]
+ kw_series = {k: all_series[v] for k, v in kwargs_offsets.items()}
+ results = []
+ for i in range(num_rows):
+ # Note: Creating a slice from a series for each window is
+ # actually pretty expensive. However, there
+ # is no easy way to reduce cost here.
+ # Note: s.iloc[i : j] is about 30% faster than s[i: j], with
+ # the caveat that the created slices shares the same
+ # memory with s. Therefore, user are not allowed to
+ # change the value of input series inside the window
+ # function. It is rare that user needs to modify the
+ # input series in the window function, and therefore,
+ # it is a reasonable restriction.
+ # Note: Calling reset_index on the slices will increase the
+ # cost of creating slices by about 100%. Therefore, for
+ # performance reasons we don't do it here.
+ slices = [s.iloc[begin_array[i] : end_array[i]] for s in series]
+ kw_slices = {
+ k: s.iloc[begin_array[i] : end_array[i]] for k, s in kw_series.items()
+ }
+ results.append(udf_func(*slices, **kw_slices))
+ result_series.append(pd.Series(results))
+ else:
+ raise PySparkRuntimeError(
+ errorClass="INVALID_WINDOW_BOUND_TYPE",
+ messageParameters={"window_bound_type": bound_type},
+ )
+
+ yield PandasToArrowConversion.from_pandas(
+ result_series,
+ self._return_schema,
+ timezone=self._runner_conf.timezone,
+ safecheck=self._runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=self._runner_conf.use_large_var_types,
+ assign_cols_by_name=self._runner_conf.assign_cols_by_name,
+ int_to_decimal_coercion_enabled=self._runner_conf.int_to_decimal_coercion_enabled,
+ )
diff --git a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
index 69ff759ca00c..76687aee943d 100644
--- a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
@@ -50,11 +50,14 @@ if have_pandas and have_pyarrow:
from pyspark.eval_handlers._pandas import (
PandasCoGroupedMapUDFHandler,
+ PandasGroupedAggIterUDFHandler,
+ PandasGroupedAggUDFHandler,
PandasGroupedMapIterUDFHandler,
PandasGroupedMapUDFHandler,
PandasMapUDFHandler,
PandasScalarIterUDFHandler,
PandasScalarUDFHandler,
+ PandasWindowAggUDFHandler,
)
from pyspark.sql.conversion import ArrowBatchTransformer
@@ -100,6 +103,26 @@ class PandasEvalTypeHandlerTestCase(PySparkBaseTestCase):
"""Flatten the ``v`` field of each output batch's wrapping struct column."""
return [v for b in batches for v in b.column(0).field("v").to_pylist()]
+ @staticmethod
+ def _udf(func, return_type=None, args=(0,), kwargs=None):
+ """A pandas UDF tuple ``(func, args_offsets, kwargs_offsets, return_type)``."""
+ return (func, list(args), dict(kwargs or {}), return_type or LongType())
+
+ @staticmethod
+ def _agg_handler(handler_cls, *udfs, runner_conf=None):
+ """Build a grouped-agg/window-agg handler from one or more ``_udf`` tuples."""
+ return handler_cls(
+ udfs=list(udfs), runner_conf=runner_conf or RunnerConf({}), eval_conf=None
+ )
+
+ @staticmethod
+ def _columns(batches):
+ """Each output column's values, concatenated across batches."""
+ return [
+ [v for b in batches for v in b.column(i).to_pylist()]
+ for i in range(batches[0].num_columns)
+ ]
+
@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
class PandasEvalTypeHandlerRegistrationTests(PySparkBaseTestCase):
@@ -112,17 +135,15 @@ class PandasEvalTypeHandlerRegistrationTests(PySparkBaseTestCase):
(PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, PandasGroupedMapUDFHandler),
(PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF, PandasGroupedMapIterUDFHandler),
(PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, PandasCoGroupedMapUDFHandler),
+ (PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, PandasGroupedAggUDFHandler),
+ (PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, PandasGroupedAggIterUDFHandler),
+ (PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF, PandasWindowAggUDFHandler),
):
self.assertIs(get_eval_type_handler(eval_type), handler_cls)
@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
class PandasScalarUDFHandlerTests(PandasEvalTypeHandlerTestCase):
- @staticmethod
- def _udf(func, return_type=None, args=(0,), kwargs=None):
- """A scalar pandas UDF tuple ``(func, args_offsets, kwargs_offsets, return_type)``."""
- return (func, list(args), dict(kwargs or {}), return_type or LongType())
-
@staticmethod
def _handler(*udfs, runner_conf=None):
"""Build a PandasScalarUDFHandler from one or more ``_udf`` tuples."""
@@ -519,6 +540,159 @@ class PandasCoGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
list(handler.run(0, iter([cogroup])))
+@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
+class PandasGroupedAggUDFHandlerTests(PandasEvalTypeHandlerTestCase):
+ def test_one_row_per_group_one_column_per_udf(self):
+ handler = self._agg_handler(
+ PandasGroupedAggUDFHandler,
+ self._udf(lambda s: s.sum(), args=(1,)),
+ self._udf(lambda a, b: (a * b).max(), args=(0, 1)),
+ )
+ groups = [
+ iter([self._batch(a=[1, 1], b=[1, 2]), self._batch(a=[1], b=[3])]),
+ iter([self._batch(a=[2], b=[5])]),
+ ]
+ out = list(handler.run(0, iter(groups)))
+ self.assertEqual([b.num_rows for b in out], [1, 1])
+ self.assertEqual(self._columns(out), [[6, 5], [3, 10]])
+
+ def test_passes_kwargs(self):
+ handler = self._agg_handler(
+ PandasGroupedAggUDFHandler, self._udf(lambda *, x: x.sum(), args=(), kwargs={"x": 0})
+ )
+ out = list(handler.run(0, iter([iter([self._batch(a=[1, 2, 3])])])))
+ self.assertEqual(self._columns(out), [[6]])
+
+ def test_skips_empty_group(self):
+ handler = self._agg_handler(PandasGroupedAggUDFHandler, self._udf(lambda s: s.sum()))
+ out = list(handler.run(0, iter([iter([]), iter([self._batch(a=[4])])])))
+ self.assertEqual(self._columns(out), [[4]])
+
+ def test_honors_large_var_types(self):
+ conf = RunnerConf({"spark.sql.execution.arrow.useLargeVarTypes": "true"})
+ handler = self._agg_handler(
+ PandasGroupedAggUDFHandler,
+ self._udf(lambda s: str(s.sum()), return_type=StringType()),
+ runner_conf=conf,
+ )
+ out = list(handler.run(0, iter([iter([self._batch(a=[1, 2])])])))
+ self.assertEqual(out[0].column(0).type, pa.large_string())
+
+
+@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
+class PandasGroupedAggIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
+ def test_single_arg_receives_series_per_batch(self):
+ def total(series_iter):
+ return sum(s.sum() for s in series_iter)
+
+ handler = self._agg_handler(PandasGroupedAggIterUDFHandler, self._udf(total, args=(1,)))
+ groups = [
+ iter([self._batch(a=[1, 1], b=[1, 2]), self._batch(a=[1], b=[3])]),
+ iter([self._batch(a=[2], b=[5])]),
+ ]
+ out = list(handler.run(0, iter(groups)))
+ self.assertEqual(self._columns(out), [[6, 5]])
+
+ def test_multiple_args_receive_tuple_of_series(self):
+ def dot(pairs):
+ return sum((a * b).sum() for a, b in pairs)
+
+ handler = self._agg_handler(PandasGroupedAggIterUDFHandler, self._udf(dot, args=(0, 1)))
+ group = iter([self._batch(a=[1, 2], b=[3, 4]), self._batch(a=[5], b=[6])])
+ out = list(handler.run(0, iter([group])))
+ self.assertEqual(self._columns(out), [[41]])
+
+ def test_drains_unconsumed_group_batches(self):
+ # A UDF that reads only the first batch must not leave the rest of its group
+ # in the stream.
+ handler = self._agg_handler(
+ PandasGroupedAggIterUDFHandler, self._udf(lambda it: next(it).sum())
+ )
+ group1 = iter([self._batch(a=[1]), self._batch(a=[2])])
+ group2 = iter([self._batch(a=[3])])
+ out = list(handler.run(0, iter([group1, group2])))
+ self.assertEqual(self._columns(out), [[1, 3]])
+ self.assertIsNone(next(group1, None))
+
+ def test_ignores_large_var_types(self):
+ # This eval type always emits non-large Arrow types, regardless of the conf.
+ conf = RunnerConf({"spark.sql.execution.arrow.useLargeVarTypes": "true"})
+ handler = self._agg_handler(
+ PandasGroupedAggIterUDFHandler,
+ self._udf(lambda it: str(sum(s.sum() for s in it)), return_type=StringType()),
+ runner_conf=conf,
+ )
+ out = list(handler.run(0, iter([iter([self._batch(a=[1, 2])])])))
+ self.assertEqual(out[0].column(0).type, pa.string())
+
+
+@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
+class PandasWindowAggUDFHandlerTests(PandasEvalTypeHandlerTestCase):
+ @staticmethod
+ def _conf(*bound_types):
+ return RunnerConf({"window_bound_types": ",".join(bound_types)})
+
+ def test_unbounded_repeats_result_per_row(self):
+ handler = self._agg_handler(
+ PandasWindowAggUDFHandler,
+ self._udf(lambda s: s.sum()),
+ runner_conf=self._conf("unbounded"),
+ )
+ groups = [iter([self._batch(a=[1, 2]), self._batch(a=[3])]), iter([self._batch(a=[10])])]
+ out = list(handler.run(0, iter(groups)))
+ self.assertEqual([b.num_rows for b in out], [3, 1])
+ self.assertEqual(self._columns(out), [[6, 6, 6, 10]])
+
+ def test_bounded_slices_per_row(self):
+ # args_offsets[0] / [1] are the per-row frame begin / end indexes.
+ handler = self._agg_handler(
+ PandasWindowAggUDFHandler,
+ self._udf(lambda s: s.sum(), args=(0, 1, 2)),
+ runner_conf=self._conf("bounded"),
+ )
+ group = iter([self._batch(begin=[0, 0, 1], end=[1, 2, 3], v=[1, 2, 3])])
+ out = list(handler.run(0, iter([group])))
+ self.assertEqual(self._columns(out), [[1, 3, 5]])
+
+ def test_bounded_slices_kwargs(self):
+ handler = self._agg_handler(
+ PandasWindowAggUDFHandler,
+ self._udf(lambda *, x: x.max(), args=(0, 1), kwargs={"x": 2}),
+ runner_conf=self._conf("bounded"),
+ )
+ group = iter([self._batch(begin=[0, 1], end=[2, 2], v=[4, 7])])
+ out = list(handler.run(0, iter([group])))
+ self.assertEqual(self._columns(out), [[7, 7]])
+
+ def test_mixed_bound_types(self):
+ handler = self._agg_handler(
+ PandasWindowAggUDFHandler,
+ self._udf(lambda s: s.sum(), args=(2,)),
+ self._udf(lambda s: s.sum(), args=(0, 1, 2)),
+ runner_conf=self._conf("unbounded", "bounded"),
+ )
+ group = iter([self._batch(begin=[0, 1], end=[1, 2], v=[3, 4])])
+ out = list(handler.run(0, iter([group])))
+ self.assertEqual(self._columns(out), [[7, 7], [3, 4]])
+
+ def test_skips_empty_group(self):
+ handler = self._agg_handler(
+ PandasWindowAggUDFHandler,
+ self._udf(lambda s: s.sum()),
+ runner_conf=self._conf("unbounded"),
+ )
+ out = list(handler.run(0, iter([iter([]), iter([self._batch(a=[5])])])))
+ self.assertEqual(self._columns(out), [[5]])
+
+ def test_rejects_invalid_bound_type(self):
+ handler = self._agg_handler(
+ PandasWindowAggUDFHandler, self._udf(lambda s: s.sum()), runner_conf=self._conf("bogus")
+ )
+ with self.assertRaises(PySparkRuntimeError) as ctx:
+ list(handler.run(0, iter([iter([self._batch(a=[1])])])))
+ self.assertEqual(ctx.exception.getCondition(), "INVALID_WINDOW_BOUND_TYPE")
+
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 840d3b23e1bd..b39a43ba3023 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -44,8 +44,6 @@ if TYPE_CHECKING:
import pandas as pd
import pyarrow as pa
- from pyspark.sql.pandas._typing import GroupedBatch
-
from pyspark import _NoValue, shuffle
from pyspark.accumulators import (
SpecialAccumulatorIds,
@@ -88,10 +86,7 @@ from pyspark.sql.conversion import (
PandasToArrowConversion,
)
from pyspark.sql.functions import SkipRestOfInputTableException
-from pyspark.sql.pandas.serializers import (
- ArrowStreamGroupSerializer,
- ArrowStreamSerializer,
-)
+from pyspark.sql.pandas.serializers import ArrowStreamSerializer
from pyspark.sql.pandas.types import to_arrow_type
from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
@@ -1117,9 +1112,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
- PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
- PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
@@ -1127,14 +1119,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF,
):
# NOTE: if timezone is set here, that implies respectSessionTimeZone is True
- if eval_type in (
- PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
- PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
- PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
- ):
- ser = ArrowStreamGroupSerializer(write_start_stream=True)
- else:
- ser = ArrowStreamSerializer(write_start_stream=True)
+ ser = ArrowStreamSerializer(write_start_stream=True)
else:
batch_size = int(os.environ.get("PYTHON_UDF_BATCH_SIZE", "100"))
ser = BatchedSerializer(CPickleSerializer(), batch_size)
@@ -1146,180 +1131,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
num_udfs = len(udfs)
- if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF:
- import pandas as pd
- import pyarrow as pa
-
- col_names = ["_%d" % i for i in range(len(udfs))]
- output_schema = StructType(
- [StructField(name, rt) for name, (_, _, _, rt) in zip(col_names, udfs)]
- )
-
- 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
- table = pa.Table.from_batches(batch_list).combine_chunks()
- all_series = ArrowToPandasConversion.to_pandas(
- table,
- timezone=runner_conf.timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- )
- results = [
- udf_func(
- *[all_series[o] for o in args_offsets],
- **{k: all_series[v] for k, v in kwargs_offsets.items()},
- )
- for udf_func, args_offsets, kwargs_offsets, _ in udfs
- ]
- result_series = [pd.Series([r]) for r in results]
- yield PandasToArrowConversion.from_pandas(
- result_series,
- output_schema,
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- arrow_cast=True,
- prefers_large_types=runner_conf.use_large_var_types,
- assign_cols_by_name=runner_conf.assign_cols_by_name,
- int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
-
- return grouped_func, ser
-
- if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF:
- import pandas as pd
- import pyarrow as pa
-
- assert num_udfs == 1, "One GROUPED_AGG_PANDAS_ITER UDF expected here."
- udf_func, args_offsets, _, return_type = udfs[0]
-
- output_schema = StructType([StructField("_0", return_type)])
-
- def extract_series(
- batch: "pa.RecordBatch",
- ) -> Union["pd.Series", tuple["pd.Series", ...]]:
- # Convert one RecordBatch to a pandas Series per column, then select args:
- # - pd.Series for a single column
- # - tuple[pd.Series, ...] for multiple columns
- all_series = ArrowToPandasConversion.to_pandas(
- batch,
- timezone=runner_conf.timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- )
- series = tuple(all_series[o] for o in args_offsets)
- return series[0] if len(series) == 1 else series
-
- def grouped_func(
- split_index: int, data: Iterator["GroupedBatch"]
- ) -> Iterator[pa.RecordBatch]:
- for group in data:
- series_iter = map(extract_series, group)
- result = udf_func(series_iter)
- # Drain remaining batches to maintain stream position
- for _ in series_iter:
- pass
- yield PandasToArrowConversion.from_pandas(
- [pd.Series([result])],
- output_schema,
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- arrow_cast=True,
- prefers_large_types=False,
- assign_cols_by_name=runner_conf.assign_cols_by_name,
- int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
-
- return grouped_func, ser
-
- if eval_type == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF:
- import pandas as pd
- import pyarrow as pa
-
- window_bound_types = runner_conf.window_bound_types
-
- col_names = ["_%d" % i for i in range(len(udfs))]
- output_schema = StructType(
- [StructField(name, rt) for name, (_, _, _, rt) in zip(col_names, udfs)]
- )
-
- 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
- table = pa.Table.from_batches(batch_list).combine_chunks()
- all_series = ArrowToPandasConversion.to_pandas(
- table,
- timezone=runner_conf.timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- )
- num_rows = table.num_rows
-
- result_series = []
- for udf_index, (udf_func, args_offsets, kwargs_offsets, _) in enumerate(udfs):
- bound_type = window_bound_types[udf_index]
- if bound_type == "unbounded":
- result = udf_func(
- *[all_series[o] for o in args_offsets],
- **{k: all_series[v] for k, v in kwargs_offsets.items()},
- )
- # Repeat the scalar result to match the window (group) length.
- result_series.append(pd.Series([result]).repeat(num_rows))
- elif bound_type == "bounded":
- # args_offsets[0] and args_offsets[1] are begin_index and end_index.
- assert len(args_offsets) >= 2, len(args_offsets)
- # Index operation is faster on np.ndarray, so we turn the
- # index series into np arrays here for performance.
- begin_array = all_series[args_offsets[0]].values
- end_array = all_series[args_offsets[1]].values
- series = [all_series[o] for o in args_offsets[2:]]
- kw_series = {k: all_series[v] for k, v in kwargs_offsets.items()}
- results = []
- for i in range(num_rows):
- # Note: Creating a slice from a series for each window is
- # actually pretty expensive. However, there
- # is no easy way to reduce cost here.
- # Note: s.iloc[i : j] is about 30% faster than s[i: j], with
- # the caveat that the created slices shares the same
- # memory with s. Therefore, user are not allowed to
- # change the value of input series inside the window
- # function. It is rare that user needs to modify the
- # input series in the window function, and therefore,
- # it is a reasonable restriction.
- # Note: Calling reset_index on the slices will increase the
- # cost of creating slices by about 100%. Therefore, for
- # performance reasons we don't do it here.
- slices = [s.iloc[begin_array[i] : end_array[i]] for s in series]
- kw_slices = {
- k: s.iloc[begin_array[i] : end_array[i]]
- for k, s in kw_series.items()
- }
- results.append(udf_func(*slices, **kw_slices))
- result_series.append(pd.Series(results))
- else:
- raise PySparkRuntimeError(
- errorClass="INVALID_WINDOW_BOUND_TYPE",
- messageParameters={"window_bound_type": bound_type},
- )
-
- yield PandasToArrowConversion.from_pandas(
- result_series,
- output_schema,
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- arrow_cast=True,
- prefers_large_types=runner_conf.use_large_var_types,
- assign_cols_by_name=runner_conf.assign_cols_by_name,
- int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
-
- return grouped_func, ser
-
if (
eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF
and not runner_conf.use_legacy_pandas_udf_conversion
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org