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 444962585362 [SPARK-56312][PYTHON] Refactor SQL_COGROUPED_MAP_ARROW_UDF
444962585362 is described below
commit 444962585362c680116e4c3081ac2bf65d936403
Author: Yicong Huang <[email protected]>
AuthorDate: Fri May 8 18:25:55 2026 +0800
[SPARK-56312][PYTHON] Refactor SQL_COGROUPED_MAP_ARROW_UDF
### What changes were proposed in this pull request?
Refactor `SQL_COGROUPED_MAP_ARROW_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 comparison (`CogroupedMapArrowUDFTimeBench`, average of 5
runs):
```text
scenario udf before (ms) after (ms) diff
----------------------------------------------------------------------------
few_groups_sm identity_udf 12.36 9.91 -19.8%
few_groups_sm concat_udf 15.81 12.41 -21.5%
few_groups_sm left_semi_udf 70.58 67.06 -5.0%
few_groups_lg identity_udf 58.67 66.00 +12.5%
few_groups_lg concat_udf 87.44 84.23 -3.7%
few_groups_lg left_semi_udf 242.04 226.07 -6.6%
many_groups_sm identity_udf 393.54 323.73 -17.7%
many_groups_sm concat_udf 518.67 413.64 -20.2%
many_groups_sm left_semi_udf 1581.58 1489.98 -5.8%
many_groups_lg identity_udf 208.82 184.77 -11.5%
many_groups_lg concat_udf 291.11 257.13 -11.7%
many_groups_lg left_semi_udf 945.73 930.55 -1.6%
wide_values identity_udf 306.41 293.69 -4.2%
wide_values concat_udf 399.99 365.39 -8.7%
wide_values left_semi_udf 699.24 599.22 -14.3%
multi_key identity_udf 76.26 71.23 -6.6%
multi_key concat_udf 116.50 105.90 -9.1%
multi_key left_semi_udf 221.64 210.75 -4.9%
```
`few_groups_lg/identity_udf` +12.5% is a benchmark ordering artifact --
when run in isolation (54.25 -> 54.02 ms, -0.4%), no regression is observed.
The effect comes from prior scenarios polluting the Python process memory
state, which does not occur in production where each Spark task runs in a fresh
Python worker. 17 of 18 scenarios show improvement or no change.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55377 from Yicong-Huang/refactor/cogrouped-map-arrow-udf.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 484bb5b2cd718d27872235a30c1af69d5d4d34d7)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/sql/conversion.py | 14 +++++
python/pyspark/worker.py | 128 +++++++++++++++++++++------------------
2 files changed, 82 insertions(+), 60 deletions(-)
diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py
index 8fc4fa5cc0cc..a229386f3001 100644
--- a/python/pyspark/sql/conversion.py
+++ b/python/pyspark/sql/conversion.py
@@ -90,6 +90,20 @@ class ArrowBatchTransformer:
struct = batch.column(column_index)
return pa.RecordBatch.from_arrays(struct.flatten(),
schema=pa.schema(struct.type))
+ @classmethod
+ def select_columns(cls, batch: "pa.RecordBatch", column_indices:
list[int]) -> "pa.RecordBatch":
+ """
+ Select a subset of columns from a RecordBatch by index.
+
+ Used by: SQL_COGROUPED_MAP_ARROW_UDF handler in worker.py
+ """
+ import pyarrow as pa
+
+ return pa.RecordBatch.from_arrays(
+ [batch.columns[i] for i in column_indices],
+ [batch.schema.names[i] for i in column_indices],
+ )
+
@staticmethod
def wrap_struct(batch: "pa.RecordBatch") -> "pa.RecordBatch":
"""
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index b9f791ada1e6..4aa7bc68f3d5 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -75,7 +75,7 @@ from pyspark.sql.pandas.serializers import (
ArrowStreamGroupSerializer,
ArrowStreamPandasUDFSerializer,
ArrowStreamPandasUDTFSerializer,
- CogroupArrowUDFSerializer,
+ ArrowStreamCoGroupSerializer,
CogroupPandasUDFSerializer,
ApplyInPandasWithStateSerializer,
TransformWithStateInPandasSerializer,
@@ -525,37 +525,6 @@ def verify_pandas_result(result, return_type,
assign_cols_by_name, truncate_retu
)
-def wrap_cogrouped_map_arrow_udf(f, return_type, argspec, runner_conf):
- import pyarrow as pa
-
- if runner_conf.assign_cols_by_name:
- expected_cols_and_types = {
- col.name: to_arrow_type(col.dataType, timezone="UTC") for col in
return_type.fields
- }
- else:
- expected_cols_and_types = [
- (col.name, to_arrow_type(col.dataType, timezone="UTC")) for col in
return_type.fields
- ]
-
- def wrapped(left_key_table, left_value_table, right_key_table,
right_value_table):
- if len(argspec.args) == 2:
- result = f(left_value_table, right_value_table)
- elif len(argspec.args) == 3:
- key_table = left_key_table if left_key_table.num_rows > 0 else
right_key_table
- key = tuple(c[0] for c in key_table.columns)
- result = f(key, left_value_table, right_value_table)
-
- verify_return_type(result, pa.Table)
- verify_arrow_result(result, runner_conf.assign_cols_by_name,
expected_cols_and_types)
-
- return result.to_batches()
-
- return lambda kl, vl, kr, vr: (
- wrapped(kl, vl, kr, vr),
- to_arrow_type(return_type, timezone="UTC"),
- )
-
-
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
@@ -1097,7 +1066,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
return args_offsets, wrap_cogrouped_map_pandas_udf(func, return_type,
argspec, runner_conf)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF:
argspec = inspect.getfullargspec(chained_func) # signature was lost
when wrapping it
- return args_offsets, wrap_cogrouped_map_arrow_udf(func, return_type,
argspec, runner_conf)
+ return func, args_offsets, return_type, len(argspec.args)
elif eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF:
return wrap_grouped_agg_pandas_udf(
func, args_offsets, kwargs_offsets, return_type, runner_conf
@@ -2325,7 +2294,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
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)
+ ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
ser = CogroupPandasUDFSerializer(
timezone=runner_conf.timezone,
@@ -2999,6 +2968,71 @@ 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_COGROUPED_MAP_ARROW_UDF:
+ import pyarrow as pa
+
+ assert num_udfs == 1, "One COGROUPED_MAP_ARROW UDF expected here."
+ cogrouped_udf, arg_offsets, return_type, num_udf_args = udfs[0]
+
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+
+ # Pre-compute expected column names/types for strict result validation.
+ # Cogrouped map has a strict contract: missing, extra, or
type-mismatched
+ # columns must raise; no silent coercion.
+ if runner_conf.assign_cols_by_name:
+ expected_cols_and_types = {
+ col.name: to_arrow_type(col.dataType, timezone="UTC") for col
in return_type.fields
+ }
+ reorder_names = [col.name for col in return_type.fields]
+ else:
+ expected_cols_and_types = [
+ (col.name, to_arrow_type(col.dataType, timezone="UTC"))
+ for col in return_type.fields
+ ]
+ reorder_names = None
+
+ select_columns = ArrowBatchTransformer.select_columns
+ left_key_cols, left_val_cols = parsed_offsets[0]
+ right_key_cols, right_val_cols = parsed_offsets[1]
+
+ def table_from_batches(batches, cols):
+ return pa.Table.from_batches([select_columns(b, cols) for b in
batches])
+
+ def cogrouped_func(
+ split_index: int,
+ data: Iterator[Tuple[list[pa.RecordBatch], list[pa.RecordBatch]]],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply cogroupBy Arrow UDF."""
+ for left_batches, right_batches in data:
+ left_keys = table_from_batches(left_batches, left_key_cols)
+ left_values = table_from_batches(left_batches, left_val_cols)
+ right_keys = table_from_batches(right_batches, right_key_cols)
+ right_values = table_from_batches(right_batches,
right_val_cols)
+
+ if num_udf_args == 2:
+ result = cogrouped_udf(left_values, right_values)
+ else:
+ key_table = left_keys if left_keys.num_rows > 0 else
right_keys
+ key = tuple(c[0] for c in key_table.columns)
+ result = cogrouped_udf(key, left_values, right_values)
+
+ verify_return_type(result, pa.Table)
+ verify_arrow_result(
+ result, runner_conf.assign_cols_by_name,
expected_cols_and_types
+ )
+
+ for batch in result.to_batches():
+ if reorder_names is not None:
+ # Names and types already validated equal; pure
reorder, no cast.
+ batch = pa.RecordBatch.from_arrays(
+ [batch.column(name) for name in reorder_names],
+ names=reorder_names,
+ )
+ yield ArrowBatchTransformer.wrap_struct(batch)
+
+ # 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
@@ -3421,32 +3455,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
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_COGROUPED_MAP_ARROW_UDF:
- import pyarrow as pa
-
- # 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 batch_from_offset(batch, offsets):
- return pa.RecordBatch.from_arrays(
- arrays=[batch.columns[o] for o in offsets],
- names=[batch.schema.names[o] for o in offsets],
- )
-
- def table_from_batches(batches, offsets):
- return pa.Table.from_batches([batch_from_offset(batch, offsets)
for batch in batches])
-
- def mapper(a):
- df1_keys = table_from_batches(a[0], parsed_offsets[0][0])
- df1_vals = table_from_batches(a[0], parsed_offsets[0][1])
- df2_keys = table_from_batches(a[1], parsed_offsets[1][0])
- df2_vals = table_from_batches(a[1], 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]
