Yicong-Huang opened a new pull request, #58015:
URL: https://github.com/apache/spark/pull/58015
### What changes were proposed in this pull request?
`read_udfs` in `python/pyspark/worker.py` ends with a catch-all `else`
branch that builds a `mapper` closure, wraps it in `func` via `map(mapper,
it)`, and unwraps a single-element result tuple on every row. That branch is
only ever reached by `SQL_BATCHED_UDF`: it is the only eval type whose
`read_single_udf` return is a two-element `(arg_offsets, eval_func)` pair (the
shape the mapper unpacks), and every other eval type returns from its own
explicit branch earlier in the function.
This PR consolidates that logic:
- Replace the catch-all `else` with an explicit `elif eval_type ==
PythonEvalType.SQL_BATCHED_UDF:` branch, and add a final `else: raise
ValueError("Unknown eval type: ...")` mirroring the style of `read_single_udf`.
- Drop the intermediate `mapper` closure and define `func` directly, with
the same `(split_index: int, data: Iterator[Any]) -> Iterator[Any]` signature
the other branches in `read_udfs` use.
- Hoist the single-result handling out of the per-row path. The old code
built a tuple and checked `len(result) == 1` for every row, but `len(result)`
is always `num_udfs`, a build-time constant. A single UDF now yields its bare
result and multiple UDFs yield a tuple of results, decided once.
Pure internal refactor of the batched (plain Python / pickle) UDF worker
path; no change to the value shape sent to the JVM.
### Why are the changes needed?
The `else` branch obscured that it serves exactly one eval type, and the
per-row `len(result) == 1` check re-evaluated a constant on the hot path.
Making the branch explicit and hoisting the constant makes the batched-UDF path
easier to read and consistent with the rest of `read_udfs`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing batched-UDF coverage (`pyspark.sql.tests.test_udf`) exercises
single-UDF, multi-UDF, complex-return, and nested-UDF projections. The refactor
is behavior-preserving: for every row `len(result) == num_udfs`, so hoisting
the single-result check to build time is exact.
### Was this patch authored or co-authored using generative AI tooling?
No.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
