(spark) branch branch-4.3 updated: [SPARK-58784][PYTHON] Merge the batched-UDF mapper into an explicit SQL_BATCHED_UDF branch in read_udfs

[yi...@apache.org]

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 f78dcff73c00 [SPARK-58784][PYTHON] Merge the batched-UDF mapper into an explicit SQL_BATCHED_UDF branch in read_udfs
f78dcff73c00 is described below

commit f78dcff73c00c8bb56d14609d4ecfff34df1b665
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Mon Aug 17 18:28:53 2026 +0000

    [SPARK-58784][PYTHON] Merge the batched-UDF mapper into an explicit SQL_BATCHED_UDF branch in read_udfs
    
    ### What changes were proposed in this pull request?
    
    The catch-all `else` at the end of `read_udfs` (`python/pyspark/worker.py`) is only reached by `SQL_BATCHED_UDF`: it is the sole eval type whose `read_single_udf` return is an `(arg_offsets, eval_func)` pair, and every other eval type returns from its own branch earlier. This PR makes that branch explicit (`elif eval_type == PythonEvalType.SQL_BATCHED_UDF:`, with a final `else: raise ValueError`), drops the intermediate `mapper` closure in favor of defining `func` directly with the sa [...]
    
    ### Why are the changes needed?
    
    The `else` hid that it serves one eval type, and the per-row `len(result) == 1` check re-evaluated a constant. An explicit branch plus the hoisted constant reads clearer and matches the rest of `read_udfs`.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Existing `pyspark.sql.tests.test_udf` (single-, multi-, complex-return, and nested-UDF projections). Behavior-preserving: `len(result) == num_udfs` for every row.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    No.
    
    Closes #58015 from Yicong-Huang/merge-mapper-batched-udf.
    
    Authored-by: Yicong Huang <17...@users.noreply.github.com>
    Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
    (cherry picked from commit 95d93a0116ac5e4a7229a87dfa7b7c85638e3f90)
    Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
 python/pyspark/worker.py | 32 ++++++++++++++++++--------------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 9ac8abbfa485..d75e1c5fac32 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -4102,22 +4102,26 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
         # profiling is not supported for UDF
         return func, None, ser, ser
 
-    else:
-
-        def mapper(a):
-            result = tuple(f(*[a[o] for o in arg_offsets]) for arg_offsets, f in udfs)
-            # In the special case of a single UDF this will return a single result rather
-            # than a tuple of results; this is the format that the JVM side expects.
-            if len(result) == 1:
-                return result[0]
-            else:
-                return result
+    elif eval_type == PythonEvalType.SQL_BATCHED_UDF:
+        # Plain Python (pickle) UDFs, the only eval type reaching this branch. read_single_udf
+        # prepared each UDF as an (arg_offsets, eval_func) pair. Apply every one to each input
+        # row: a single UDF yields its bare result, multiple UDFs yield a tuple of results,
+        # which is the shape the JVM side expects. num_udfs is fixed, so the single-result
+        # case is handled once here rather than by unwrapping a one-element tuple per row.
+        def func(split_index: int, data: Iterator[Any]) -> Iterator[Any]:
+            if num_udfs == 1:
+                arg_offsets, f = udfs[0]
+                return (f(*[row[offset] for offset in arg_offsets]) for row in data)
+            return (
+                tuple(f(*[row[offset] for offset in arg_offsets]) for arg_offsets, f in udfs)
+                for row in data
+            )
 
-    def func(_, it):
-        return map(mapper, it)
+        # profiling is not supported for UDF
+        return func, None, ser, ser
 
-    # profiling is not supported for UDF
-    return func, None, ser, ser
+    else:
+        raise ValueError("Unknown eval type: {}".format(eval_type))
 
 
 def invoke_udf(message_receiver: SparkMessageReceiver, outfile: BinaryIO):


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org