HyukjinKwon opened a new pull request, #57899:
URL: https://github.com/apache/spark/pull/57899
### What changes were proposed in this pull request?
This is a follow-up to SPARK-27052, which let a plain (row-at-a-time) Python
UDF be used inside a higher-order function's lambda by lifting it out of the
lambda and applying it once to the whole array. That work handled only
`SQL_BATCHED_UDF` / `SQL_ARROW_BATCHED_UDF`, lifting them to a single
element-wise eval type `SQL_ARROW_ELEMENTWISE_UDF` (=102).
This PR extends the same rewrite to the vectorized scalar UDFs, so that e.g.
```python
@pandas_udf("int")
def plus_one(s: pd.Series) -> pd.Series:
return s + 1
df.select(F.transform("values", lambda x: plus_one(x)))
```
works, instead of failing with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`. All four vectorized
scalar flavors are covered: scalar pandas (`pandas_udf`), scalar Arrow
(`arrow_udf`), and their iterator variants.
Each flavor lifts to its own dedicated element-wise eval type so the Python
worker keeps that flavor's native batching contract, mirroring the existing 102:
- `SQL_SCALAR_PANDAS_ELEMENTWISE_UDF` (=103)
- `SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF` (=104)
- `SQL_SCALAR_ARROW_ELEMENTWISE_UDF` (=105)
- `SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF` (=106)
`PythonUDF.liftedElementwiseEvalType` maps each base eval type to its lifted
counterpart, and `isElementwiseRewritableUDF` now accepts the vectorized types
(still excluding zero-argument, named-argument, and UDT shapes, which the lift
cannot preserve). In the Python worker, each argument arrives as `array<T>`
aligned with the iterated array; the worker flattens each list column to its
element column, runs the vectorized function once over that flat column (a
`pandas.Series` / `pyarrow.Array`, or an iterator of them), then re-nests the
flat result to `array<R>` using the input's offsets. For the iterator variants,
the JVM joins UDF output to input positionally by row, so the worker streams
the flattened elements through the user's iterator and regroups the streamed
results back into arrays via a FIFO of per-row shapes; output batch boundaries
need not match input ones.
### Why are the changes needed?
SPARK-27052 supported only plain Python UDFs inside HOF lambdas. Vectorized
scalar UDFs (pandas / Arrow) are the more performant and more common form, and
users expect them to compose with higher-order functions the same way.
### Does this PR introduce _any_ user-facing change?
Yes. A scalar pandas or Arrow UDF (and their iterator variants) can now be
used inside a higher-order function's lambda, e.g.
`df.select(F.transform("values", lambda x: pandas_plus_one(x)))`. Before this
change such a query failed at analysis with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`; now it evaluates. No
existing successful query changes behavior. This extends an unreleased feature
(SPARK-27052), so it is not a change relative to any released version.
### How was this patch tested?
- `pyspark.sql.tests.test_udf_in_higher_order_function` (classic and Connect
parity): new tests for all four vectorized flavors across `transform` /
`filter`, covering null arrays, null elements, empty arrays, and multi-batch
iterator regrouping.
- `ExtractPythonUDFFromLambdaSuite`: the previous "a pandas UDF inside a
lambda still fails" test is replaced by one asserting each vectorized flavor is
lifted to its element-wise eval type; the rewritable-predicate test now covers
all rewritable eval types and the base -> lifted mapping.
- `PythonUDFSuite` and `BatchEvalPythonExecSuite` pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
This pull request and its description were written by Isaac.
--
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]
