This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang 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 65e3e0b96a27 [SPARK-56718][PYTHON] Refactor
SQL_COGROUPED_MAP_PANDAS_UDF
65e3e0b96a27 is described below
commit 65e3e0b96a27aca5e95a0b17c2286ad9229d1f1f
Author: Yicong Huang <[email protected]>
AuthorDate: Fri Jun 12 19:02:51 2026 +0000
[SPARK-56718][PYTHON] Refactor SQL_COGROUPED_MAP_PANDAS_UDF
### What changes were proposed in this pull request?
Refactor `SQL_COGROUPED_MAP_PANDAS_UDF` to be self-contained in
`read_udfs()`, moving logic from `CogroupPandasUDFSerializer` into `worker.py`,
using `ArrowStreamCoGroupSerializer` as pure I/O.
### 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 (`CogroupedMapPandasUDFTimeBench`), 3 runs each, averaged:
```text
master: ebfdf8b6b4 vs PR: 87d96ebe40
Time (ms, lower = better)
scenario udf master PR diff
---------------------------------------------------------------------
few_groups_sm identity_udf 160.3 159.0 -0.81%
few_groups_sm concat_udf 181.0 181.0 0.00%
few_groups_sm left_semi_udf 182.7 181.7 -0.55%
few_groups_sm key_identity_udf 160.3 162.0 +1.04%
few_groups_lg identity_udf 413.3 411.3 -0.48%
few_groups_lg concat_udf 545.7 537.3 -1.54%
few_groups_lg left_semi_udf 508.3 500.7 -1.51%
few_groups_lg key_identity_udf 412.3 412.3 0.00%
many_groups_sm identity_udf 1356.7 1356.7 0.00%
many_groups_sm concat_udf 1480.0 1473.3 -0.45%
many_groups_sm left_semi_udf 1496.7 1503.3 +0.44%
many_groups_sm key_identity_udf 1370.0 1363.3 -0.49%
many_groups_lg identity_udf 765.3 760.7 -0.61%
many_groups_lg concat_udf 922.7 908.7 -1.52%
many_groups_lg left_semi_udf 887.0 886.0 -0.11%
many_groups_lg key_identity_udf 764.7 767.0 +0.30%
wide_values identity_udf 1073.3 1073.3 0.00%
wide_values concat_udf 1276.7 1266.7 -0.78%
wide_values left_semi_udf 1130.0 1130.0 0.00%
wide_values key_identity_udf 1080.0 1073.3 -0.62%
multi_key identity_udf 425.3 423.7 -0.39%
multi_key concat_udf 473.7 473.0 -0.14%
multi_key left_semi_udf 463.0 451.3 -2.52%
multi_key key_identity_udf 430.3 428.3 -0.46%
---------------------------------------------------------------------
SUM 17995.0 17834.0 -0.90%
```
Aggregate essentially flat (-0.90%); per-scenario variation within
run-to-run noise.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55674 from Yicong-Huang/refactor/cogrouped-map-pandas-udf.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 118 ++++++++++++++++++++++++++++-------------------
1 file changed, 71 insertions(+), 47 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 2dfb4a96bf07..e28bbecf5e02 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -79,7 +79,6 @@ from pyspark.sql.pandas.serializers import (
ArrowStreamPandasUDFSerializer,
ArrowStreamPandasUDTFSerializer,
ArrowStreamCoGroupSerializer,
- CogroupPandasUDFSerializer,
ApplyInPandasWithStateSerializer,
TransformWithStateInPandasSerializer,
TransformWithStateInPandasInitStateSerializer,
@@ -493,28 +492,6 @@ def verify_pandas_result(
_verify_column_schema(actual_names, field_names,
assign_cols_by_name=by_name)
-def wrap_cogrouped_map_pandas_udf(f, return_type, argspec, runner_conf):
- def wrapped(left_key_series, left_value_series, right_key_series,
right_value_series):
- import pandas as pd
-
- left_df = pd.concat(left_value_series, axis=1)
- right_df = pd.concat(right_value_series, axis=1)
-
- if len(argspec.args) == 2:
- result = f(left_df, right_df)
- elif len(argspec.args) == 3:
- key_series = left_key_series if not left_df.empty else
right_key_series
- key = tuple(s[0] for s in key_series)
- result = f(key, left_df, right_df)
- verify_pandas_result(
- result, return_type, runner_conf.assign_cols_by_name,
truncate_return_schema=False
- )
-
- return result
-
- return lambda kl, vl, kr, vr: [(wrapped(kl, vl, kr, vr), return_type)]
-
-
def verify_arrow_result(
result: Union["pa.Table", "pa.RecordBatch"],
assign_cols_by_name: bool,
@@ -924,7 +901,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
argspec = inspect.getfullargspec(chained_func) # signature was lost
when wrapping it
- return args_offsets, wrap_cogrouped_map_pandas_udf(func, return_type,
argspec, runner_conf)
+ return func, args_offsets, return_type, len(argspec.args)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF:
argspec = inspect.getfullargspec(chained_func) # signature was lost
when wrapping it
return func, args_offsets, return_type, len(argspec.args)
@@ -2152,14 +2129,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF:
ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
- ser = CogroupPandasUDFSerializer(
- 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,
- arrow_cast=True,
- )
+ ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE:
ser = ApplyInPandasWithStateSerializer(
timezone=runner_conf.timezone,
@@ -3056,6 +3026,75 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# profiling is not supported for UDF
return func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One COGROUPED_MAP_PANDAS UDF expected here."
+ cogrouped_udf, arg_offsets, return_type, num_udf_args = udfs[0]
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+
+ left_key_offsets, left_value_offsets = parsed_offsets[0]
+ right_key_offsets, right_value_offsets = parsed_offsets[1]
+ output_schema = StructType([StructField("_0", return_type)])
+
+ def cogrouped_func(
+ split_index: int,
+ data: Iterator[Tuple[list[pa.RecordBatch], list[pa.RecordBatch]]],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply cogroupBy Pandas UDF."""
+ for left_batches, right_batches in data:
+ left_table = pa.Table.from_batches(left_batches)
+ right_table = pa.Table.from_batches(right_batches)
+ left_series = ArrowBatchTransformer.to_pandas(
+ left_table,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+ right_series = ArrowBatchTransformer.to_pandas(
+ right_table,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+ left_df = pd.concat([left_series[o] for o in
left_value_offsets], axis=1)
+ right_df = pd.concat([right_series[o] for o in
right_value_offsets], axis=1)
+
+ if num_udf_args == 2:
+ result = cogrouped_udf(left_df, right_df)
+ else:
+ key_series = (
+ [left_series[o] for o in left_key_offsets]
+ if not left_df.empty
+ else [right_series[o] for o in right_key_offsets]
+ )
+ key = tuple(s.iloc[0] for s in key_series)
+ result = cogrouped_udf(key, left_df, right_df)
+
+ del left_batches, right_batches, left_table, right_table
+ del left_series, right_series, left_df, right_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 cogrouped_func, None, ser, ser
+
if (
eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF
and not runner_conf.use_legacy_pandas_udf_conversion
@@ -3535,21 +3574,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
return f(keys, vals, state)
- elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
- # We assume there is only one UDF here because cogrouped map doesn't
- # support combining multiple UDFs.
- assert num_udfs == 1
- arg_offsets, f = udfs[0]
-
- parsed_offsets = extract_key_value_indexes(arg_offsets)
-
- def mapper(a):
- df1_keys = [a[0][o] for o in parsed_offsets[0][0]]
- df1_vals = [a[0][o] for o in parsed_offsets[0][1]]
- df2_keys = [a[1][o] for o in parsed_offsets[1][0]]
- df2_vals = [a[1][o] for o in parsed_offsets[1][1]]
- return f(df1_keys, df1_vals, df2_keys, df2_vals)
-
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.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
