HyukjinKwon opened a new pull request, #57940:
URL: https://github.com/apache/spark/pull/57940
> **[DO-NOT-MERGE]** Draft for review/CI only. No JIRA yet.
### What changes were proposed in this pull request?
This is a follow-up to SPARK-27052 (plain Python UDFs in higher-order
function lambdas) and SPARK-58695 (vectorized scalar UDFs there).
`ExtractPythonUDFFromLambda` lifts a scalar Python UDF out of a HOF's lambda
and applies it once to the whole array outside the lambda. This PR extends that
rewrite to two shapes it previously rejected:
**1. A UDF inside a *nested* lambda**, e.g.
```python
df.select(F.transform("matrix", lambda row: F.transform(row, lambda x:
plus_one(x))))
```
The rule now rewrites a whole nest of higher-order functions root-first
(`apply` fires only at a nest root — a HOF that iterates real columns — and
`rewriteNest` rewrites the entire nest in one action, so a nested lambda's UDF
is never momentarily left as a free-variable `PythonUDF`). Each inner lambda's
UDF is lifted onto that lambda's variable and then re-lifted onto the enclosing
array, one array level per enclosing lambda: `plus_one` becomes a depth-`N`
element-wise UDF. The Python worker flattens `N` list levels down to the
leaves, runs the function once, and re-nests `N` levels. The per-UDF nesting
depth travels to the worker via a new `PythonUDF.elementwiseNestingDepth` field
carried in `evalConf` (`elementwise_nesting`). Works at any depth, for all four
vectorized flavors, and for a nested `aggregate`.
**2. A UDF inside `aggregate` / `reduce`'s `merge` lambda that reads only
the element**, e.g.
```python
df.select(F.aggregate("values", F.lit(0), lambda acc, x: acc + plus_one(x)))
```
`plus_one(x)` is element-independent of the fold, so it is precomputed over
the whole array and the merge reads it positionally, exactly like `transform`.
A UDF that reads the **accumulator**, or one in the `finish` lambda, remains
rejected — the fold is sequential and such a UDF cannot be precomputed.
`CheckAnalysis` and the rule continue to share one predicate
(`PythonUDF.canRewritePythonUDFInLambda`), so analysis accepts exactly what the
rule rewrites.
**Known limitation (rejected at analysis, not miscompiled):** a UDF in a
nested lambda that *captures an enclosing lambda's variable*, e.g.
`transform(m, row -> transform(row, x -> f(x, size(row))))`. Lifting it would
make the captured value a nested-lambda argument of the lifted UDF, and
`HigherOrderFunction.canonicalized` mishandles such a nested lambda (it
renumbers free lambda-variable references inconsistently, leaking a stray
attribute into the canonical `references`), which then breaks
`ExtractPythonUDFs`'s `ExpressionSet.filter(references.subsetOf(inputSet))` so
the UDF is left in the lambda and hits codegen. That is a pre-existing Catalyst
canonicalization issue this feature merely exposes; it is guarded precisely by
`PythonUDF.noUDFCapturesEnclosingLambdaVariable` and left as a separate
follow-up.
### Why are the changes needed?
`transform(col, lambda row: transform(row, lambda x: my_udf(x)))` over
nested arrays, and `aggregate(col, 0, lambda acc, x: acc + my_udf(x))`, are
natural things to write. The previous work supported only single, non-fold
lambdas; these are the remaining common shapes.
### Does this PR introduce _any_ user-facing change?
Yes. Queries that previously failed analysis with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]` now run:
```python
>>> df.select(F.transform("matrix", lambda row: F.transform(row, lambda x:
plus_one(x)))).show()
>>> df.select(F.aggregate("values", F.lit(0), lambda acc, x: acc +
plus_one(x))).show()
```
This extends unreleased features (SPARK-27052 / SPARK-58695), so it is not a
change relative to any released version. No existing successful query changes
behavior.
### How was this patch tested?
- `pyspark.sql.tests.test_udf_in_higher_order_function` (classic and Connect
parity): nested lambdas across `transform` / `filter`, three levels of nesting,
all four vectorized flavors, null/empty/multi-batch inputs; a
nondeterministic-inner-argument rejection; the enclosing-capture rejection;
`aggregate` over the element (plain and vectorized) plus accumulator/`finish`
rejections.
- `ExtractPythonUDFFromLambdaSuite`: plan-shape assertions for the aggregate
lift and rejection, the nested (depth-2) lift, the enclosing-capture rejection,
and the element-wise re-lift predicate.
- Classic 62/62, Connect parity 61 passed + 1 skipped (JVM-plan test),
`ExtractPythonUDFFromLambdaSuite` 26/26 locally.
### 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]
