[SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas

[HyukjinKwon]

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

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
NamedLambdaVariables 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 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. The Arrow
boundary is crossed once per batch rather than once per row, and it stays one row in / one row out:
no explode, no shuffle.

All twelve lambda-taking functions are supported: transform, filter, exists, forall,
zip_with, aggregate/reduce (both on merge's element and in finish), array_sort (as a
per-element sort key), transform_keys, transform_values, map_filter and map_zip_with. The
map family is desugared to the array case over map_keys/map_values and rebuilt with
map_from_arrays; map_zip_with visits the union of both key sets, looking each map up per key so
a key missing from one side yields null, matching its own semantics.

Nested higher-order functions are supported, e.g.
transform(arr, i -> transform(i, x -> f(x))). One pass cannot lift this, because the inner array
is the outer lambda's variable. The rewrite composes instead: the inner pass lifts f onto i, and
because the result then sits in the inner function's argument — which is evaluated outside its
lambda — the outer pass lifts it again onto arr, incrementing the flatten depth. Applied
repeatedly to a fixed point, so arbitrary nesting depth works (tested three deep, and with mixed
kinds such as a filter inside a transform).

The flatten depth cannot be inferred by the worker: transform(arr, i -> total(i)) and
transform(arr, i -> transform(i, x -> f(x))) hand it the identical Arrow type list<list<int32>>
but need one and two flattens respectively. It is therefore carried explicitly on
PythonUDF.elementwiseDepths and sent in the eval conf.

The rule is generic. Rather than a case per function, there is one rewrite that reads what it
needs off the HigherOrderFunction API — arguments/functions for the arity, withNewChildren
to rebuild (a higher-order function's children are always arguments ++ functions, verified for all
twelve), and positional matching of lambda parameters to collection arguments. Only two facts per
function cannot be derived, and both default safely for an unknown function: whether the result is
the lambda's value or the input elements, and whether the lambda is a comparator (which parameter
types cannot reveal, since (T, Int) is ambiguous). Adding a new higher-order function of a
familiar shape needs no change to the rule.

Still rejected (CheckAnalysis and the rule share one predicate, so analysis accepts exactly
what the optimizer rewrites):

• a UDF reading aggregate's accumulator. Not rewritable at any cost, which is stronger than
  "expensive". Every other shape precomputes over values that are elements of some collection; here
  they are not. Folding [1,2,3] with acc*2 + x calls the UDF on 0, 1, 4 — outputs of earlier
  steps, which do not exist until the fold runs. So there is no input set to precompute over, and
  even the cross-product trick that would make a pairwise comparator work does not apply.
  Supporting it would need either a physical operator that calls Python from inside higher-order
  function evaluation (which the extraction architecture does not allow) or unrolling the fold
  (which needs a static bound on array length).
• an array_sort comparator whose UDF receives both elements in one call. This one is
  implementable, by precomputing the whole n×n matrix of pairs and having the comparator index it —
  but that is O(n²) Python calls where sorting needs only O(n log n) comparisons, plus O(n²) memory.
  Returning a per-element sort key instead stays on the O(n) path and is supported, so the error
  points there. I prototyped the matrix rewrite and left it out rather than ship a pessimisation.
• a pandas UDF, which receives a Series rather than one value per call.

The rewrite can be turned off with
spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled=false, restoring the previous error.

This builds on design notes derived from the
elementwise-udf prototype (Apache-2.0), which
established which rewrites are correct and which are not worth doing. Note the Catalyst
implementation goes further: the prototype cannot support a higher-order function nested inside a
lambda, because it resolves each argument with F.col(...) outside the lambda.

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:

>>> 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]

The two 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 (41 tests), 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 all twelve
higher-order functions; nested higher-order functions (two and three deep, and mixed kinds);
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; ragged zip_with inputs; 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 (15 tests) asserts that no PythonUDF
remains inside a LambdaFunction for any of the twelve functions, that the lifted UDF is an
element-wise UDF over an array with the right flatten depth, 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. 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, two-level nesting and broadcast arguments before
being wired in.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 5)