You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2026/05/06 01:38:41 UTC
(spark) branch branch-4.x updated: [SPARK-56691][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_ITER_UDF
This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 01a6781aa3ff [SPARK-56691][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_ITER_UDF
01a6781aa3ff is described below
commit 01a6781aa3ffb9500f610a56e3def7f4c76c9c21
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Wed May 6 09:38:15 2026 +0800
[SPARK-56691][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_ITER_UDF
### What changes were proposed in this pull request?
Refactor `SQL_GROUPED_MAP_PANDAS_ITER_UDF` to be self-contained in `read_udfs()`.
### Why are the changes needed?
Part of SPARK-55388 (Refactor PythonEvalType processing logic). Making each eval type self-contained in `read_udfs()` improves readability and makes it easier to reason about the data flow for each eval type independently.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. No behavior change.
ASV benchmark (`GroupedMapPandasIterUDFTimeBench`, single run with `-a repeat=5`):
```text
master: 4b3f8c3796e vs PR: 29538fd7980
Time (ms, lower = better)
scenario udf master PR diff
sm_grp_few_col identity_udf 447.4 441.0 -1.43%
sm_grp_few_col sort_udf 499.5 498.8 -0.14%
sm_grp_few_col key_identity_udf 449.9 411.8 -8.46%
sm_grp_many_col identity_udf 358.3 375.5 +4.79%
sm_grp_many_col sort_udf 378.5 388.7 +2.70%
sm_grp_many_col key_identity_udf 371.3 341.1 -8.14%
lg_grp_few_col identity_udf 802.7 791.6 -1.39%
lg_grp_few_col sort_udf 993.7 949.8 -4.42%
lg_grp_few_col key_identity_udf 682.4 691.2 +1.30%
lg_grp_many_col identity_udf 928.7 911.1 -1.89%
lg_grp_many_col sort_udf 1010.4 963.1 -4.69%
lg_grp_many_col key_identity_udf 897.8 919.7 +2.44%
mixed_types identity_udf 446.2 431.3 -3.34%
mixed_types sort_udf 471.2 450.0 -4.50%
mixed_types key_identity_udf 399.8 383.4 -4.10%
SUM 9137.8 8948.1 -2.08%
```
Aggregate slightly improved (-2.08%); per-scenario variation within run-to-run noise.
Peakmem benchmark (`GroupedMapPandasIterUDFPeakmemBench`) was essentially flat (SUM -0.02%).
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55675 from Yicong-Huang/SPARK-56691.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
(cherry picked from commit 512605461ca9c0f84066d5945268a5011181cb1a)
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
python/pyspark/worker.py | 165 +++++++++++++++++++++++------------------------
1 file changed, 82 insertions(+), 83 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 32ea54171184..2e38838b1028 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -62,7 +62,6 @@ from pyspark.sql.pandas.serializers import (
ArrowStreamGroupSerializer,
ArrowStreamPandasUDFSerializer,
ArrowStreamPandasUDTFSerializer,
- GroupPandasUDFSerializer,
CogroupArrowUDFSerializer,
CogroupPandasUDFSerializer,
ApplyInPandasWithStateSerializer,
@@ -652,39 +651,6 @@ def verify_arrow_batch(batch, assign_cols_by_name, expected_cols_and_types):
verify_arrow_result(batch, assign_cols_by_name, expected_cols_and_types)
-def wrap_grouped_map_pandas_iter_udf(f, return_type, argspec, runner_conf):
- def wrapped(key_series, value_batches):
- import pandas as pd
-
- # value_batches is an Iterator[list[pd.Series]] (one list per batch)
- # Convert each list of Series into a DataFrame
- def dataframe_iter():
- for value_series in value_batches:
- yield pd.concat(value_series, axis=1)
-
- if len(argspec.args) == 1:
- result = f(dataframe_iter())
- elif len(argspec.args) == 2:
- # Extract key from pandas Series, preserving numpy types
- key = tuple(s.iloc[0] for s in key_series)
- result = f(key, dataframe_iter())
-
- def verify_element(df):
- verify_pandas_result(
- df, return_type, runner_conf.assign_cols_by_name, truncate_return_schema=False
- )
- return df
-
- yield from map(verify_element, result)
-
- def flatten_wrapper(k, v):
- # Return Iterator[[(df, spark_type)]] directly
- for df in wrapped(k, v):
- yield [(df, return_type)]
-
- return flatten_wrapper
-
-
def wrap_grouped_transform_with_state_pandas_udf(f, return_type, runner_conf):
def wrapped(stateful_processor_api_client, mode, key, value_series_gen):
result_iter = f(stateful_processor_api_client, mode, key, value_series_gen)
@@ -1122,14 +1088,12 @@ def read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index):
return args_offsets, wrap_pandas_batch_iter_udf(func, return_type, runner_conf)
elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
return func, None, None, None
- elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF:
+ elif eval_type in (
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
+ ):
num_udf_args = len(inspect.getfullargspec(chained_func).args)
return func, args_offsets, return_type, num_udf_args
- elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF:
- argspec = inspect.getfullargspec(chained_func) # signature was lost when wrapping it
- return args_offsets, wrap_grouped_map_pandas_iter_udf(
- func, return_type, argspec, runner_conf
- )
elif eval_type in (
PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
@@ -2378,6 +2342,7 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
# NOTE: if timezone is set here, that implies respectSessionTimeZone is True
if eval_type in (
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
@@ -2398,14 +2363,6 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
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_GROUPED_MAP_PANDAS_ITER_UDF:
- ser = GroupPandasUDFSerializer(
- 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 = CogroupArrowUDFSerializer(assign_cols_by_name=runner_conf.assign_cols_by_name)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
@@ -2996,6 +2953,82 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One GROUPED_MAP_PANDAS_ITER UDF expected here."
+ grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0]
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+ assert len(parsed_offsets) == 1, (
+ "Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF."
+ )
+
+ key_offsets = parsed_offsets[0][0]
+ value_offsets = parsed_offsets[0][1]
+ output_schema = StructType([StructField("_0", return_type)])
+
+ def grouped_func(
+ split_index: int,
+ data: Iterator[Iterator[pa.RecordBatch]],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply groupBy Pandas UDF (iterator variant).
+
+ The UDF receives an Iterator[pd.DataFrame] per group and
+ returns an Iterator[pd.DataFrame]. Input batches are
+ converted to pandas lazily so peakmem stays bounded by a
+ single batch rather than the whole group.
+ """
+ for group in data:
+ group_iter = iter(group)
+ # Read the first batch to extract grouping keys.
+ first_series = ArrowBatchTransformer.to_pandas(
+ next(group_iter),
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+
+ def dataframe_iter():
+ yield pd.concat([first_series[o] for o in value_offsets], axis=1)
+ for batch in group_iter:
+ series = ArrowBatchTransformer.to_pandas(
+ batch,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+ yield pd.concat([series[o] for o in value_offsets], axis=1)
+
+ if num_udf_args == 1:
+ result = grouped_udf(dataframe_iter())
+ else:
+ key = tuple(first_series[o].iloc[0] for o in key_offsets)
+ result = grouped_udf(key, dataframe_iter())
+
+ for df in result:
+ verify_pandas_result(
+ df,
+ return_type,
+ runner_conf.assign_cols_by_name,
+ truncate_return_schema=False,
+ )
+ yield PandasToArrowConversion.convert(
+ [df],
+ 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,
+ )
+
+ # Drain remaining input batches to maintain stream position.
+ for _ in group_iter:
+ pass
+
+ # profiling is not supported for UDF
+ return grouped_func, None, ser, ser
+
if (
eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF
and not runner_conf.use_legacy_pandas_udf_conversion
@@ -3240,41 +3273,7 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
# profiling is not supported for UDF
return func, None, ser, ser
- if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF:
- import pyarrow as pa
-
- # We assume there is only one UDF here because grouped map doesn't
- # support combining multiple UDFs.
- assert num_udfs == 1
-
- # See FlatMapGroupsInPandasExec for how arg_offsets are used to
- # distinguish between grouping attributes and data attributes
- arg_offsets, f = udfs[0]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
-
- def mapper(batch_iter):
- # Convert first Arrow batch to pandas to extract keys
- first_series = ArrowBatchTransformer.to_pandas(
- next(batch_iter),
- timezone=ser._timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- )
- key_series = [first_series[o] for o in parsed_offsets[0][0]]
-
- # Lazily convert remaining Arrow batches to pandas Series
- def value_series_gen():
- yield [first_series[o] for o in parsed_offsets[0][1]]
- for batch in batch_iter:
- series = ArrowBatchTransformer.to_pandas(
- batch,
- timezone=ser._timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- )
- yield [series[o] for o in parsed_offsets[0][1]]
-
- yield from f(key_series, value_series_gen())
-
- elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
+ if eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
# We assume there is only one UDF here because grouped map doesn't
# support combining multiple UDFs.
assert num_udfs == 1
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org