HyukjinKwon opened a new pull request, #57804:
URL: https://github.com/apache/spark/pull/57804
### What changes were proposed in this pull request?
This PR lets a scalar Python UDF be used inside the lambda of a higher-order
function, so that
```python
df.select(F.transform("values", lambda x: plus_one(x)))
```
works with a plain `pyspark.sql.functions.udf` instead of failing with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`.
**Why it does not work today.** A `PythonUDF` is evaluated by a separate
physical operator
(`ArrowEvalPython`), which `ExtractPythonUDFs` pulls out of the enclosing
operator. A lambda's
`NamedLambdaVariable`s only exist while the higher-order function is
iterating, so an extracted
operator cannot see them: the UDF can neither stay inside the lambda nor be
lifted out by the
existing extraction rule. `CheckAnalysis` therefore rejected it outright.
**The rewrite.** Do not evaluate the UDF per element inside the lambda.
Apply it once to the
*whole array*, outside every lambda, and let the lambda read the precomputed
result positionally.
A new optimizer rule, `ExtractPythonUDFFromLambda`, does this:
```
-- before (rejected)
transform(values, x -> plus_one(x))
-- after (legal: the PythonUDF is outside every LambdaFunction)
transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s ->
s.u0)
```
It runs in the `Extract Python UDFs` batch *before* `ExtractPythonUDFs`,
which then extracts the
lifted UDF as an ordinary top-level `PythonUDF`, so no new physical operator
is needed. It is
registered in `nonExcludableRules`, since a plan that only works because of
this rewrite must not
be silently broken by `spark.sql.optimizer.excludedRules`.
**Where the array-at-a-time behaviour lives.** `_wrap_function` pickles
`(func, returnType)` as
opaque bytes, so the JVM can never rewrap the user's function; the wrapper
has to be worker-side.
A new eval type, `SQL_ARROW_ELEMENTWISE_UDF`, does it with Arrow: the worker
flattens the incoming
list column, calls the user function once over the concatenated elements of
the **whole batch**,
then re-nests the results using the input's offsets plus an explicit
validity mask.
This is deliberately stronger than a per-row Python loop. The number of
Python-level calls is
proportional to the number of elements, but the Arrow boundary is crossed
once per batch rather
than once per row, which the design notes measured as the dominant cost. It
also keeps one row in
and one row out: no `explode`, no shuffle.
Once the UDF result is an ordinary column, everything the lambda does around
it is ordinary JVM
work, so arithmetic, `when`, casts, the element index, several UDFs in one
lambda, nested UDF
calls (`f(g(x))`) and UDF arguments that are expressions over the element
all follow without
special cases.
**Supported:** `transform`, `filter`, `exists`, `forall`, and `aggregate`
when the UDF applies to
the element.
**Still rejected** (`CheckAnalysis` and the rule share one predicate, so
analysis accepts exactly
what the optimizer rewrites):
- a UDF reading `aggregate`'s accumulator, which is sequential (step *n*
depends on *n-1*);
- a UDF in `aggregate`'s `finish` — it runs once on the final accumulator,
and lifting it out
would call it on the null a fold over a null array produces, where native
Spark does not
evaluate `finish` at all;
- a pandas UDF, which receives a `Series` rather than one value per call;
- a nested higher-order function, e.g. `transform(arr, i -> transform(i, x
-> udf(x)))`: the inner
array only exists while the outer function iterates, so a UDF lifted onto
it would still sit
inside the outer lambda. (Note the reference prototype does not support
this shape either — its
"nested" support is nested *UDF calls*, which this PR does support.
Extending to nested
higher-order functions is possible by flattening and re-nesting one level
deeper, and is left
as follow-up.)
- the map family (`transform_keys`, `transform_values`, `map_filter`,
`map_zip_with`) and
`zip_with`, left as follow-up.
The rewrite can be turned off with
`spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled=false`,
restoring the previous error.
This builds on the design notes derived from the
[elementwise-udf](https://github.com/HyukjinKwon/elementwise-udf) prototype
(Apache-2.0), which
established which rewrites are correct and which are not worth doing.
### Why are the changes needed?
`transform(col, lambda x: my_udf(x))` is a natural thing to write and is a
long-standing gap
(SPARK-27052). Today users must either avoid Python UDFs in lambdas entirely
or fall back to
`explode` + regroup, which the prototype measured at 2-40x slower and which
OOMs on long arrays
because carrying the source array alongside `posexplode` duplicates it per
element.
### Does this PR introduce _any_ user-facing change?
Yes. A query that previously failed analysis now runs:
```python
>>> plus_one = udf(lambda x: x + 1, "int")
>>> df.select(F.transform("values", lambda x: plus_one(x))).show()
# before: AnalysisException
[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]
# after: [2, 3, 4]
```
Shapes that cannot be rewritten keep failing with the same error condition
as before. No existing
successful query changes behaviour.
### How was this patch tested?
New end-to-end suite `pyspark.sql.tests.test_udf_in_higher_order_function`
(29 tests, all
passing), which asserts results against the equivalent native expression
wherever one exists, so a
rewrite that runs but computes the wrong thing fails rather than passing
quietly. It covers each
supported higher-order function; composition around the UDF result; the
index parameter; several
and nested UDF calls; UDF arguments that are expressions over the element;
broadcast outer-column
and constant-only arguments (the constant case must still yield one result
per element); element
and return types including string, double and `array<int>`; null arrays,
null elements, UDFs
returning null, empty arrays, all-null rows and an empty frame; long arrays
over many rows to
exercise batching; integration with joins, caching and `groupBy`; mixing
with an ordinary Python
UDF; that a lambda with no Python UDF is left untouched; and the negative
cases above.
New plan-shape suite `ExtractPythonUDFFromLambdaSuite` (13 tests, all
passing) asserts that no
`PythonUDF` remains inside a `LambdaFunction`, that the lifted UDF is an
element-wise UDF over an
array, that duplicate calls are evaluated once, that the rule is inert
without a UDF, and that it
is not excludable.
`PythonUDFSuite`'s `SPARK-48706` negative test asserted the old behaviour
for exactly the case now
supported; it is updated to assert the result and a still-unsupported shape.
`PythonUDFSuite`
(16 tests), the `execution.python` suites, `DataFrameFunctionsSuite`,
`pyspark.sql.tests.test_udf`
and `arrow.test_arrow_python_udf` pass.
The Arrow flatten/re-nest algorithm was also validated standalone against
null arrays, null
elements, empty arrays, sliced arrays, large lists and broadcast arguments
before being wired in.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
--
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]
