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 2f47c1ff0012 [SPARK-58276][PYTHON] Consolidate serializer selection
branches in worker.py
2f47c1ff0012 is described below
commit 2f47c1ff0012a3cab6d89f03827403193c06132b
Author: Yicong Huang <[email protected]>
AuthorDate: Fri Jul 24 03:43:31 2026 +0000
[SPARK-58276][PYTHON] Consolidate serializer selection branches in worker.py
### What changes were proposed in this pull request?
This PR consolidates redundant `eval_type` branches in
`python/pyspark/worker.py` where different branches produced identical results.
No logic changes:
- `read_single_udf`: grouped the return branches by the shape of the value
returned (scalar/agg/window, grouped-map/cogrouped-map,
with-state/transform-with-state, map-iterator, batched), collapsing 13 branches
into 6. The cogrouped-map branches computed
`len(inspect.getfullargspec(chained_func).args)`, identical to the grouped-map
`num_udf_args`, so they now share one branch.
- `read_udtf`: merged the two branches that both build
`ArrowStreamSerializer(write_start_stream=True)`; the per-eval-type comments
are preserved inline.
- `read_udfs`: merged the window-agg branch into the grouped-map/agg branch
(both build `ArrowStreamGroupSerializer`) and the two cogrouped-map branches
(both build `ArrowStreamCoGroupSerializer`), collapsing 5 branches into 3.
### Why are the changes needed?
The serializer selection logic had accumulated many one-line branches that
returned identical values, making the mapping from eval type to behavior harder
to read. Grouping by outcome makes the equivalence classes explicit.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. No behavior change.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #57446 from Yicong-Huang/simplify-worker-serializer-selection.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 110 +++++++++++++++++++++--------------------------
1 file changed, 49 insertions(+), 61 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index cd5b0fb73c32..2db77d73148e 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -641,61 +641,50 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
args_offsets, kwargs_offsets = udf_info.args, udf_info.kwargs
- # the last returnType will be the return type of UDF
- if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF:
- return func, args_offsets, kwargs_offsets, return_type
- elif eval_type == PythonEvalType.SQL_SCALAR_ARROW_UDF:
- return func, args_offsets, kwargs_offsets, return_type
- elif eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF:
- return func, args_offsets, kwargs_offsets, return_type
- elif eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF:
- return func, args_offsets, kwargs_offsets, return_type
- elif eval_type == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF:
+ # The last returnType will be the return type of UDF. Eval types are
grouped below by the
+ # shape of the value they return.
+
+ # Scalar, aggregation and window UDFs: (func, args_offsets,
kwargs_offsets, return_type).
+ if eval_type in (
+ PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
+ ):
return func, args_offsets, kwargs_offsets, return_type
- elif eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF:
- return func, None, None, return_type
- elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
- return func, None, None, None
+ # Grouped-map and cogrouped-map UDFs: (func, args_offsets, return_type,
num_udf_args).
elif eval_type in (
- PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
+ PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF,
+ PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
):
+ # signature was lost when wrapping it
num_udf_args = len(inspect.getfullargspec(chained_func).args)
return func, args_offsets, return_type, num_udf_args
+ # Grouped-map-with-state and transform-with-state UDFs: (func,
args_offsets, return_type).
elif eval_type in (
- PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
- PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF,
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_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_UDF_WITH_STATE:
- return func, args_offsets, return_type
- elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
- return func, args_offsets, return_type
- elif eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF:
return func, args_offsets, return_type
- elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF:
- return func, args_offsets, return_type
- elif eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF:
- return func, args_offsets, return_type
- elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
- argspec = inspect.getfullargspec(chained_func) # signature was lost
when wrapping it
- 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)
- elif eval_type in (
- 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 in (
- PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
- PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
- ):
- return func, args_offsets, kwargs_offsets, return_type
+ # Map iterator UDFs take no offsets.
+ elif eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF:
+ return func, None, None, return_type
+ elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
+ return func, None, None, None
elif eval_type == PythonEvalType.SQL_BATCHED_UDF:
return wrap_udf(func, args_offsets, kwargs_offsets, return_type)
else:
@@ -707,14 +696,15 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
# ensure the UDTF is valid. This function also prepares a mapper function for
applying
# the UDTF logic to input rows.
def read_udtf(pickleSer, udtf_info, eval_type, runner_conf, eval_conf):
- if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF:
+ if eval_type in (
# Pure Arrow stream I/O for both the legacy pandas conversion path and
the
# non-legacy path; the pandas (de)serialization for the legacy path
and the
# output struct wrapping are both handled in the func below.
- ser = ArrowStreamSerializer(write_start_stream=True)
- elif eval_type == PythonEvalType.SQL_ARROW_UDTF:
+ PythonEvalType.SQL_ARROW_TABLE_UDF,
# Pure Arrow stream I/O; table-arg flattening and output coercion
# are handled in the func below.
+ PythonEvalType.SQL_ARROW_UDTF,
+ ):
ser = ArrowStreamSerializer(write_start_stream=True)
else:
# Each row is a group so do not batch but send one by one.
@@ -1874,24 +1864,22 @@ def read_udfs(pickleSer, udf_info_list, 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_PANDAS_UDF,
- PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
- ):
- ser = ArrowStreamGroupSerializer(write_start_stream=True)
- elif eval_type in (
+ PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
):
ser = ArrowStreamGroupSerializer(write_start_stream=True)
- elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF:
- ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
- elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
+ elif eval_type in (
+ PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF,
+ PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
+ ):
ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
else:
ser = ArrowStreamSerializer(write_start_stream=True)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
