This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 85565f37fcd6 [SPARK-58577][PYTHON] Inline wrap_udf into the batched
UDF dispatch branch in worker.py
85565f37fcd6 is described below
commit 85565f37fcd620df0c705ed7e1d0b531309906c2
Author: Yicong Huang <[email protected]>
AuthorDate: Tue Aug 4 23:53:44 2026 -0700
[SPARK-58577][PYTHON] Inline wrap_udf into the batched UDF dispatch branch
in worker.py
### What changes were proposed in this pull request?
Inline the single-caller `wrap_udf` helper into the `SQL_BATCHED_UDF`
branch of `read_single_udf` in `worker.py` and delete its definition. The two
steps it did (`wrap_kwargs_support`, then a `needConversion()`/`toInternal`
wrap) are already open-coded inline in the neighboring dispatch branches, so
this just makes the outlier match the prevailing idiom. A short comment
preserves the semantic label.
### Why are the changes needed?
`wrap_udf` is a thin single-caller wrapper that diverges from the
inline-in-branch pattern of `read_single_udf`. Removing the indirection makes
the branch self-contained and consistent with its neighbors.
### Does this PR introduce _any_ user-facing change?
No. Pure internal refactor; the closure capture and the `needConversion()`
split are unchanged.
### How was this patch tested?
Existing `pyspark.sql.tests.test_udf UDFTests` passes.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #57775 from Yicong-Huang/inline-wrap-udf.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong Huang <[email protected]>
(cherry picked from commit 94ecd044ac05d2de513ebf74b1a135c647df7b26)
Signed-off-by: Yicong Huang <[email protected]>
---
python/pyspark/worker.py | 19 ++++++++-----------
1 file changed, 8 insertions(+), 11 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 5dc22650455c..3a5ae680e56c 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -371,16 +371,6 @@ def verify_iter_result_row_count(
verify_result_row_count(actual_rows, expected_rows())
-def wrap_udf(f, args_offsets, kwargs_offsets, return_type):
- func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets,
kwargs_offsets)
-
- if return_type.needConversion():
- toInternal = return_type.toInternal
- return args_kwargs_offsets, lambda *a: toInternal(func(*a))
- else:
- return args_kwargs_offsets, lambda *a: func(*a)
-
-
def _verify_column_schema(
actual_names: list, expected_names: list, *, assign_cols_by_name: bool
) -> None:
@@ -672,8 +662,15 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
return func, None, None, return_type
elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
return func, None, None, None
+ # Batched (plain Python) UDFs: (args_kwargs_offsets, eval func); apply
kwargs binding and
+ # convert each result to the internal representation only when the return
type requires it.
elif eval_type == PythonEvalType.SQL_BATCHED_UDF:
- return wrap_udf(func, args_offsets, kwargs_offsets, return_type)
+ func, args_kwargs_offsets = wrap_kwargs_support(func, args_offsets,
kwargs_offsets)
+ if return_type.needConversion():
+ toInternal = return_type.toInternal
+ return args_kwargs_offsets, lambda *a: toInternal(func(*a))
+ else:
+ return args_kwargs_offsets, lambda *a: func(*a)
else:
raise ValueError("Unknown eval type: {}".format(eval_type))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
