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 9147749923af [SPARK-57676][PYTHON] Refactor
SQL_GROUPED_AGG_PANDAS_ITER_UDF
9147749923af is described below
commit 9147749923af8ec71eb8b3c0bd9acc2e8edb0921
Author: Yicong Huang <[email protected]>
AuthorDate: Thu Jun 25 23:34:00 2026 +0000
[SPARK-57676][PYTHON] Refactor SQL_GROUPED_AGG_PANDAS_ITER_UDF
### What changes were proposed in this pull request?
Refactor `SQL_GROUPED_AGG_PANDAS_ITER_UDF` to use
`ArrowStreamGroupSerializer` as a pure I/O layer, moving all processing logic
into a dedicated function in `read_udfs()` in `worker.py`. This mirrors the
Arrow path done in SPARK-56123 for `SQL_GROUPED_AGG_ARROW_ITER_UDF`. The
now-unused `ArrowStreamAggPandasUDFSerializer` and
`wrap_grouped_agg_pandas_iter_udf` are removed from `worker.py` (the serializer
class itself is removed in a follow-up, SPARK-57680).
### Why are the changes needed?
Part of SPARK-55388. `SQL_GROUPED_AGG_PANDAS_ITER_UDF` was the last
consumer of `ArrowStreamAggPandasUDFSerializer`; this migration unblocks
deleting that serializer.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. No behavior change.
ASV `GroupedAggPandasIterUDFTimeBench` (`-a repeat=3`, `--python=same`),
one representative run per side; conclusion consistent across multiple runs.
```text
scenario udf before after diff
few_groups_sm sum 40.6+-0.2ms 39.4+-0.1ms -3%
few_groups_sm mean_multi 46.3+-0.2ms 45.6+-0.3ms -2%
few_groups_lg sum 67.2+-0.5ms 65.6+-0.1ms -2%
few_groups_lg mean_multi 74.5+-0.4ms 74.6+-2ms 0%
many_groups_sm sum 1.53+-0s 1.49+-0.01s -3%
many_groups_sm mean_multi 1.77+-0.01s 1.72+-0s -3%
many_groups_lg sum 421+-2ms 411+-2ms -2%
many_groups_lg mean_multi 489+-2ms 474+-1ms -3%
wide_cols sum 403+-0.7ms 393+-0.7ms -2%
wide_cols mean_multi 423+-0.2ms 414+-1ms -2%
```
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #56754 from Yicong-Huang/SPARK-57676.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit d52093ae249241acc7e7a400e17b1b4311706948)
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 103 ++++++++++++++++++++++-------------------------
1 file changed, 48 insertions(+), 55 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 5ceebca15338..1b19409ec562 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -83,7 +83,6 @@ from pyspark.sql.pandas.serializers import (
TransformWithStateInPandasInitStateSerializer,
TransformWithStateInPySparkRowSerializer,
TransformWithStateInPySparkRowInitStateSerializer,
- ArrowStreamAggPandasUDFSerializer,
ArrowStreamUDTFSerializer,
)
from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type
@@ -645,24 +644,6 @@ def wrap_grouped_map_pandas_udf_with_state(f, return_type,
runner_conf):
return lambda k, v, s: [(wrapped(k, v, s), return_type)]
-def wrap_grouped_agg_pandas_iter_udf(f, args_offsets, kwargs_offsets,
return_type, runner_conf):
- func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets,
kwargs_offsets)
-
- def wrapped(series_iter):
- import pandas as pd
-
- # series_iter: Iterator[pd.Series] (single column) or
- # Iterator[Tuple[pd.Series, ...]] (multiple columns)
- # This has already been adapted by the mapper function in read_udfs
- result = func(series_iter)
- return pd.Series([result])
-
- return (
- args_kwargs_offsets,
- lambda *a: (wrapped(*a), return_type),
- )
-
-
def wrap_kwargs_support(f, args_offsets, kwargs_offsets):
if len(kwargs_offsets):
keys = list(kwargs_offsets.keys())
@@ -857,12 +838,9 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
):
return func, args_offsets, kwargs_offsets, return_type
- elif eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF:
- return wrap_grouped_agg_pandas_iter_udf(
- func, args_offsets, kwargs_offsets, return_type, runner_conf
- )
elif eval_type in (
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
@@ -2059,6 +2037,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
):
ser = ArrowStreamGroupSerializer(write_start_stream=True)
elif eval_type in (
@@ -2066,14 +2045,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
):
ser = ArrowStreamGroupSerializer(write_start_stream=True)
- elif eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF:
- ser = ArrowStreamAggPandasUDFSerializer(
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- assign_cols_by_name=runner_conf.assign_cols_by_name,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
-
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF:
ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
@@ -2389,6 +2360,52 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ 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 = ArrowBatchTransformer.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.convert(
+ [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,
+ )
+
+ # profiling is not supported for UDF
+ return grouped_func, None, ser, ser
+
if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF:
import pyarrow as pa
@@ -3622,30 +3639,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
return f(keys, vals, state)
- elif eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF:
- # We assume there is only one UDF here because grouped agg doesn't
- # support combining multiple UDFs.
- assert num_udfs == 1
-
- arg_offsets, f = udfs[0]
-
- # Convert to iterator of pandas Series:
- # - Iterator[pd.Series] for single column
- # - Iterator[Tuple[pd.Series, ...]] for multiple columns
- def mapper(batch_iter):
- # batch_iter is Iterator[Tuple[pd.Series, ...]] where each tuple
represents one batch
- # Convert to Iterator[pd.Series] or Iterator[Tuple[pd.Series,
...]] based on arg_offsets
- if len(arg_offsets) == 1:
- # Single column: Iterator[Tuple[pd.Series, ...]] ->
Iterator[pd.Series]
- series_iter = (batch_series[arg_offsets[0]] for batch_series
in batch_iter)
- else:
- # Multiple columns: Iterator[Tuple[pd.Series, ...]] ->
- # Iterator[Tuple[pd.Series, ...]]
- series_iter = (
- tuple(batch_series[o] for o in arg_offsets) for
batch_series in batch_iter
- )
- return f(series_iter)
-
else:
def mapper(a):
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
