This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang 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 bf66715088c2 [SPARK-56648][4.X][PYTHON] Refactor SQL_SCALAR_PANDAS_UDF
bf66715088c2 is described below
commit bf66715088c2d08b4fbefbb753ed65f975db506d
Author: Yicong Huang <[email protected]>
AuthorDate: Tue Jun 9 22:48:06 2026 +0000
[SPARK-56648][4.X][PYTHON] Refactor SQL_SCALAR_PANDAS_UDF
### What changes were proposed in this pull request?
Backport of https://github.com/apache/spark/pull/55613 (commit 1f5d0e32531)
to `branch-4.x`. Clean cherry-pick, no code changes.
Refactor `SQL_SCALAR_PANDAS_UDF` to use `ArrowStreamSerializer` as pure
I/O, moving Arrow-to-Pandas and Pandas-to-Arrow conversion logic from
`ArrowStreamPandasUDFSerializer` into `read_udfs()` in `worker.py`.
Specifically:
- Remove the dedicated `wrap_scalar_pandas_udf` wrapper.
- Route `SQL_SCALAR_PANDAS_UDF` through
`ArrowStreamSerializer(write_start_stream=True)`.
- In `read_udfs()`, add a self-contained handler that:
- Converts each Arrow `RecordBatch` to pandas Series via
`ArrowBatchTransformer.to_pandas()` (with `struct_in_pandas="dict"`,
`df_for_struct=True`, `ndarray_as_list=False`).
- Invokes each UDF column-wise on the pandas inputs and validates the
return type (must be array-like) and row count (must match input).
- Enforces the existing rule that struct return types must be
`pandas.DataFrame`.
- Converts results back to an Arrow `RecordBatch` via
`PandasToArrowConversion.convert()`.
### Why are the changes needed?
Part of [SPARK-55388](https://issues.apache.org/jira/browse/SPARK-55388).
Backporting to `branch-4.x` keeps the eval-type processing paths consistent
with master and unblocks the backport of follow-up refactors (e.g. SPARK-56758)
that build on this change.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. No behavior change. See
https://github.com/apache/spark/pull/55613 for test details and ASV benchmark
results (latency within +-5% noise, peak memory flat).
### Was this patch authored or co-authored using generative AI tooling?
No (clean cherry-pick of the original commit).
This pull request and its description were written by Isaac.
Closes #56415 from Yicong-Huang/SPARK-56648-backport-4.x.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 122 ++++++++++++++++++++++++++++++++---------------
1 file changed, 83 insertions(+), 39 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index ad87e0570c13..076ef55dd21d 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -402,42 +402,6 @@ def wrap_udf(f, args_offsets, kwargs_offsets, return_type):
return args_kwargs_offsets, lambda *a: func(*a)
-def wrap_scalar_pandas_udf(f, args_offsets, kwargs_offsets, return_type,
runner_conf):
- func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets,
kwargs_offsets)
-
- def verify_result_type(result):
- if not hasattr(result, "__len__"):
- pd_type = "pandas.DataFrame" if isinstance(return_type,
StructType) else "pandas.Series"
- raise PySparkTypeError(
- errorClass="UDF_RETURN_TYPE",
- messageParameters={
- "expected": pd_type,
- "actual": type(result).__name__,
- },
- )
- return result
-
- def verify_result_length(result, length):
- if len(result) != length:
- raise PySparkRuntimeError(
- errorClass="SCHEMA_MISMATCH_FOR_PANDAS_UDF",
- messageParameters={
- "udf_type": "pandas_udf",
- "expected": str(length),
- "actual": str(len(result)),
- },
- )
- return result
-
- return (
- args_kwargs_offsets,
- lambda *a: (
- verify_result_length(verify_result_type(func(*a)), len(a[0])),
- return_type,
- ),
- )
-
-
def wrap_pandas_batch_iter_udf(f, return_type, runner_conf):
iter_type_label = "pandas.DataFrame" if isinstance(return_type,
StructType) else "pandas.Series"
@@ -998,7 +962,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
# the last returnType will be the return type of UDF
if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF:
- return wrap_scalar_pandas_udf(func, args_offsets, kwargs_offsets,
return_type, runner_conf)
+ 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:
@@ -2323,14 +2287,14 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_SCALAR_ARROW_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_UDF,
):
ser = ArrowStreamSerializer(write_start_stream=True)
else:
# Scalar Pandas UDF handles struct type arguments as pandas
DataFrames instead of
# pandas Series. See SPARK-27240.
df_for_struct = (
- eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF
- or eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
+ eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
or eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
)
@@ -3199,6 +3163,86 @@ 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_SCALAR_PANDAS_UDF:
+ import pandas as pd
+ import pyarrow as pa
+
+ # --- UDF preparation ---
+ udf_infos = []
+ for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in
udfs:
+ wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
+ udf_func, udf_args_offsets, udf_kwargs_offsets
+ )
+ udf_infos.append((wrapped_func, args_kwargs_offsets,
udf_return_type))
+ return_schema = StructType(
+ [StructField(f"_{i}", info[2]) for i, info in enumerate(udf_infos)]
+ )
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ num_rows = input_batch.num_rows
+
+ # --- Input: Arrow -> pandas Series (struct columns become
DataFrames) ---
+ pandas_columns = ArrowBatchTransformer.to_pandas(
+ input_batch,
+ timezone=runner_conf.timezone,
+ struct_in_pandas="dict",
+ ndarray_as_list=False,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ df_for_struct=True,
+ )
+
+ # --- Process: evaluate each UDF column-wise on pandas Series
---
+ results = []
+ for udf_func, offsets, udf_return_type in udf_infos:
+ result = udf_func(*[pandas_columns[o] for o in offsets])
+ if not hasattr(result, "__len__"):
+ pd_type = (
+ "pandas.DataFrame"
+ if isinstance(udf_return_type, StructType)
+ else "pandas.Series"
+ )
+ raise PySparkTypeError(
+ errorClass="UDF_RETURN_TYPE",
+ messageParameters={
+ "expected": pd_type,
+ "actual": type(result).__name__,
+ },
+ )
+ if len(result) != num_rows:
+ raise PySparkRuntimeError(
+ errorClass="SCHEMA_MISMATCH_FOR_PANDAS_UDF",
+ messageParameters={
+ "udf_type": "pandas_udf",
+ "expected": str(num_rows),
+ "actual": str(len(result)),
+ },
+ )
+ # struct_in_pandas="dict": UDF must return DataFrame for
struct types
+ if isinstance(udf_return_type, StructType) and not
isinstance(
+ result, pd.DataFrame
+ ):
+ raise PySparkValueError(
+ "Invalid return type. Please make sure that the
UDF returns a "
+ "pandas.DataFrame when the specified return type
is StructType."
+ )
+ results.append(result)
+
+ # --- Output: pandas -> Arrow ---
+ yield PandasToArrowConversion.convert(
+ results,
+ return_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,
+ )
+
+ # profiling is not supported for UDF
+ return func, None, ser, ser
+
is_scalar_iter = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
is_map_pandas_iter = eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
