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/04/30 12:08:55 UTC
(spark) branch master updated: [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF
This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 1a4ed8f16635 [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF
1a4ed8f16635 is described below
commit 1a4ed8f16635c499bf4873e0240f3d24b1e5454e
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Thu Apr 30 20:08:42 2026 +0800
[SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF
### What changes were proposed in this pull request?
Refactor `SQL_GROUPED_MAP_PANDAS_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 (`GroupedMapPandasUDFTimeBench`):
```text
master: 174fc6014f vs PR: 41dc496616
Time (ms, lower = better)
scenario udf master PR Δ
-----------------------------------------------------------------
sm_grp_few_col identity_udf 424.0 423.3 -0.17%
sm_grp_few_col sort_udf 481.7 473.5 -1.70%
sm_grp_few_col key_identity_udf 427.6 384.3 -10.12%
sm_grp_many_col identity_udf 332.6 328.6 -1.21%
sm_grp_many_col sort_udf 343.4 341.9 -0.43%
sm_grp_many_col key_identity_udf 330.5 324.6 -1.77%
lg_grp_few_col identity_udf 242.2 236.1 -2.52%
lg_grp_few_col sort_udf 359.7 357.6 -0.57%
lg_grp_few_col key_identity_udf 212.3 216.1+ 1.77%
lg_grp_many_col identity_udf 492.1 517.9+ 5.24%
lg_grp_many_col sort_udf 598.7 613.2+ 2.42%
lg_grp_many_col key_identity_udf 479.2 488.7+ 1.97%
mixed_types identity_udf 422.9 440.2+ 4.08%
mixed_types sort_udf 449.4 456.9+ 1.65%
mixed_types key_identity_udf 398.1 383.2 -3.73%
-----------------------------------------------------------------
SUM 5994.5 5986.1 -0.14%
```
Aggregate essentially flat (-0.14%); per-scenario variation within run-to-run noise.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55495 from Yicong-Huang/SPARK-56477.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
python/pyspark/worker.py | 142 +++++++++++++++++++++++++----------------------
1 file changed, 76 insertions(+), 66 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 4bdb1f5de967..32ea54171184 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -652,33 +652,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_udf(f, return_type, argspec, runner_conf):
- def wrapped(key_series, value_series):
- import pandas as pd
-
- value_df = pd.concat(value_series, axis=1)
-
- if len(argspec.args) == 1:
- result = f(value_df)
- 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, value_df)
-
- verify_pandas_result(
- result, return_type, runner_conf.assign_cols_by_name, truncate_return_schema=False
- )
-
- yield 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_map_pandas_iter_udf(f, return_type, argspec, runner_conf):
def wrapped(key_series, value_batches):
import pandas as pd
@@ -1150,8 +1123,8 @@ def read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index):
elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
return func, None, None, None
elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF:
- argspec = inspect.getfullargspec(chained_func) # signature was lost when wrapping it
- return args_offsets, wrap_grouped_map_pandas_udf(func, return_type, argspec, runner_conf)
+ 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(
@@ -2404,6 +2377,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_ARROW_UDF,
PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
@@ -2424,10 +2398,7 @@ 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_UDF
- or eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
- ):
+ elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF:
ser = GroupPandasUDFSerializer(
timezone=runner_conf.timezone,
safecheck=runner_conf.safecheck,
@@ -2954,6 +2925,77 @@ 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_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One GROUPED_MAP_PANDAS 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 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 (non-iterator variant).
+
+ The explicit ``del`` calls below keep peakmem bounded across
+ groups. Without them, generator locals from the previous
+ iteration stay bound on the frame until each statement in
+ the next iteration rebinds its slot, so the input-side
+ DataFrames overlap with the next group's allocations and
+ the working set grows unbounded on wide-column, large-group
+ inputs. ``del result`` runs on resume from yield, before
+ ``data.__next__()`` is asked for the next group.
+ """
+ for group in data:
+ all_batches = list(group)
+ if all_batches:
+ table = pa.Table.from_batches(all_batches).combine_chunks()
+ else:
+ table = pa.table({})
+ all_series = ArrowBatchTransformer.to_pandas(
+ table,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+ value_df = pd.concat([all_series[o] for o in value_offsets], axis=1)
+
+ if num_udf_args == 1:
+ result = grouped_udf(value_df)
+ else:
+ key = tuple(all_series[o].iloc[0] for o in key_offsets)
+ result = grouped_udf(key, value_df)
+
+ del all_batches, table, all_series, value_df
+
+ verify_pandas_result(
+ result,
+ return_type,
+ runner_conf.assign_cols_by_name,
+ truncate_return_schema=False,
+ )
+
+ yield PandasToArrowConversion.convert(
+ [result],
+ 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,
+ )
+ del result
+
+ # 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
@@ -3198,39 +3240,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_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)
-
- key_offsets = parsed_offsets[0][0]
- value_offsets = parsed_offsets[0][1]
-
- def mapper(batch_iter):
- # Collect all Arrow batches and merge at Arrow level
- all_batches = list(batch_iter)
- if all_batches:
- table = pa.Table.from_batches(all_batches).combine_chunks()
- else:
- table = pa.table({})
- # Convert to pandas once for the entire group
- all_series = ArrowBatchTransformer.to_pandas(
- table,
- timezone=ser._timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- )
- key_series = [all_series[o] for o in key_offsets]
- value_series = [all_series[o] for o in value_offsets]
- yield from f(key_series, value_series)
-
- elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF:
+ 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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org