You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2026/08/12 22:59:57 UTC
(spark) branch master updated: [SPARK-58728][SQL][PYTHON] Support Python UDFs in nested higher-order function lambdas
This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 6d8c619fb82d [SPARK-58728][SQL][PYTHON] Support Python UDFs in nested higher-order function lambdas
6d8c619fb82d is described below
commit 6d8c619fb82dc5b5f7af41614035e61910dd89b1
Author: Hyukjin Kwon <gu...@apache.org>
AuthorDate: Thu Aug 13 07:59:40 2026 +0900
[SPARK-58728][SQL][PYTHON] Support Python UDFs in nested higher-order function lambdas
### What changes were proposed in this pull request?
Follow-up to SPARK-27052 / SPARK-58695, which let a scalar Python UDF (plain, or vectorized pandas/Arrow + iterator variants) be used inside a *single* higher-order function lambda by lifting it out of the lambda and applying it once to the whole array. This PR extends `ExtractPythonUDFFromLambda` to two more shapes, and fixes a latent Catalyst canonicalization bug that blocked one of them.
1. **Nested lambdas** -- `transform(matrix, lambda row: transform(row, lambda x: udf(x)))`. The rule rewrites a whole nest root-first (`rewriteNest`), lifting the UDF out one lambda level at a time; `udf` becomes a depth-`N` element-wise UDF whose per-UDF nesting depth (`PythonUDF.elementwiseNestingDepth`) travels to the worker via `evalConf`. The worker flattens `N` list levels down to the leaves, runs the function once, and re-nests. Works at any depth, for all four vectorized flavo [...]
2. **Named-argument UDF calls** -- `transform(vals, lambda x: udf(x, y=lit(1)))` on the non-iterator flavors. The lift keeps the `NamedArgumentExpression` as a direct child of the lifted UDF (only its value becomes an aligned array), so the runner still derives the kwargs mapping.
3. **Catalyst fix in `HigherOrderFunction.canonicalized`.** It renumbered *all* lambda-variable occurrences, including free references to an enclosing lambda's variable, so a nested HOF's recursive re-canonicalization renumbered an enclosing variable inconsistently and leaked a stray attribute into canonical `references`. That broke `ExtractPythonUDFs`' `ExpressionSet.filter(references.subsetOf(inputSet))` (the filter applies its predicate to `e.canonicalized`). The fix numbers only n [...]
`CheckAnalysis` and the rule share one predicate (`PythonUDF.canRewritePythonUDFInLambda`), so analysis accepts exactly what the rule rewrites.
#### What is supported after this PR
| Case | Before this PR | After this PR |
|---|---|---|
| Scalar UDF in a single lambda (`transform`, `filter`, `exists`, `forall`, `zip_with`, `array_sort`, `map_*`) | supported | supported |
| All 4 vectorized flavors (pandas / Arrow scalar + their iterator variants) and the plain UDF | supported | supported |
| **UDF in a nested lambda, at any depth** (`transform(m, r -> transform(r, x -> udf(x)))`) | rejected | **supported** |
| **Nested lambda that captures the enclosing variable** (`... x -> f(x, size(r))`) | rejected | **supported** |
| **UDF called with named arguments** (`udf(x, y=lit(1))`) on non-iterator flavors | rejected | **supported** |
#### Examples of the newly-supported cases
UDF in a nested lambda, at any depth:
```python
from pyspark.sql import functions as sf
sf.udf("int")
def add_one(x):
return x + 1
df = spark.createDataFrame([([[1, 2], [3]],)], ["matrix"])
df.select(
sf.transform("matrix", lambda row: sf.transform(row, lambda x: add_one(x)))
).show(truncate=False)
# +----------------+
# |transform(mat...)|
# +----------------+
# |[[2, 3], [4]] |
# +----------------+
```
Nested lambda that captures the enclosing variable (`row` is read inside the inner lambda):
```python
from pyspark.sql import functions as sf
sf.udf("int")
def add(x, n):
return x + n
df = spark.createDataFrame([([[1, 2], [3]],)], ["matrix"])
df.select(
sf.transform(
"matrix",
lambda row: sf.transform(row, lambda x: add(x, sf.size(row))),
)
).show(truncate=False)
# +----------------+
# |transform(mat...)|
# +----------------+
# |[[3, 4], [4]] |
# +----------------+
```
UDF called with named arguments (non-iterator flavors):
```python
from pyspark.sql import functions as sf
sf.udf("int")
def add(x, y):
return x + y
df = spark.createDataFrame([([1, 2, 3],)], ["vals"])
df.select(
sf.transform("vals", lambda x: add(x, y=sf.lit(10)))
).show(truncate=False)
# +----------------+
# |transform(vals...)|
# +----------------+
# |[11, 12, 13] |
# +----------------+
```
A vectorized flavor (`arrow_udf`) works the same way in a nested lambda:
```python
import pyarrow as pa
from pyspark.sql import functions as sf
from pyspark.sql.functions import arrow_udf
arrow_udf("int")
def add_one(x: pa.Array) -> pa.Array:
return pa.compute.add(x, 1)
df = spark.createDataFrame([([[1, 2], [3]],)], ["matrix"])
df.select(
sf.transform("matrix", lambda row: sf.transform(row, lambda x: add_one(x)))
).show(truncate=False)
# +----------------+
# |transform(mat...)|
# +----------------+
# |[[2, 3], [4]] |
# +----------------+
```
#### What is still rejected (unchanged)
| Case | Why |
|---|---|
| UDF inside `aggregate` / `reduce` | The fold is sequential -- each step depends on the previous accumulator -- so the UDF cannot be applied once to the whole array. |
| Nondeterministic iterated argument | Lifting changes evaluation order/count, so a nondeterministic per-element argument is unsafe. |
| Zero-argument UDF | Nothing to align to the array being iterated. |
| UDT argument / return type | Not supported by the element-wise lift. |
All rejections fail cleanly at analysis with `[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`.
### Why are the changes needed?
`transform(col, lambda row: transform(row, lambda x: my_udf(x)))` over nested arrays, and named-argument UDF calls, are natural things to write; the earlier work covered only single, flat lambdas. The canonicalization bug is a pre-existing correctness hazard for any nested HOF that references an enclosing lambda variable.
### Does this PR introduce _any_ user-facing change?
Yes. Queries that previously failed analysis with `[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]` now run (nested lambdas including enclosing capture; named-argument calls). This extends unreleased features, so it is not a change relative to any released version, and no existing successful query changes behavior. As with the existing feature, the lifted UDF runs over *every* element even where a lambda would short-circuit, so a UDF that can fail on some input should handle it [...]
### How was this patch tested?
- `pyspark.sql.tests.test_udf_in_higher_order_function` (classic + Connect parity): nested lambdas across `transform` / `filter`, three levels, all four vectorized flavors, enclosing-variable capture, null / empty / multi-batch inputs; named-argument calls; a nondeterministic nested-argument rejection; and an `aggregate` / `reduce` rejection.
- `ExtractPythonUDFFromLambdaSuite`: plan-shape assertions for each case, including the depth bookkeeping.
- Regression for the canonicalization fix: `CanonicalizeSuite`, `HigherOrderFunctionsSuite`, `SubexpressionEliminationSuite`.
- Local run: catalyst regression 65/65, `ExtractPythonUDFFromLambdaSuite` 27/27, Python classic + Connect parity 125 passed + 1 skipped.
### 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.
Closes #57940 from HyukjinKwon/SPARK-udf-in-hof-nested-and-aggregate.
Authored-by: Hyukjin Kwon <gu...@apache.org>
Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
---
.../docs/source/tutorial/sql/arrow_python_udf.rst | 27 +-
.../sql/tests/test_udf_in_higher_order_function.py | 176 ++++++++++--
python/pyspark/worker.py | 309 ++++++++++++++-------
.../sql/catalyst/analysis/CheckAnalysis.scala | 7 +
.../spark/sql/catalyst/expressions/PythonUDF.scala | 114 ++++++--
.../expressions/higherOrderFunctions.scala | 10 +-
.../sql/execution/python/ArrowEvalPythonExec.scala | 5 +-
.../sql/execution/python/ArrowPythonRunner.scala | 46 ++-
.../ColumnarArrowEvalPythonEvaluatorFactory.scala | 7 +-
.../python/ColumnarArrowPythonRunner.scala | 21 +-
.../python/ExtractPythonUDFFromLambda.scala | 105 ++++---
.../python/ExtractPythonUDFFromLambdaSuite.scala | 74 +++--
12 files changed, 639 insertions(+), 262 deletions(-)
diff --git a/python/docs/source/tutorial/sql/arrow_python_udf.rst b/python/docs/source/tutorial/sql/arrow_python_udf.rst
index fb1ff7515218..0ed9212edf11 100644
--- a/python/docs/source/tutorial/sql/arrow_python_udf.rst
+++ b/python/docs/source/tutorial/sql/arrow_python_udf.rst
@@ -479,21 +479,18 @@ SQL boolean expressions do not short-circuit: in ``WHERE cond AND udf(x)``, the
called on all rows regardless of ``cond``. If the function can fail on certain input values
(e.g., division by zero), handle those cases inside the function itself.
-The same holds for a Python UDF used inside a higher-order function's lambda, such as
-``transform(values, lambda x: udf(x))`` or ``array_sort(values, lambda a, b: udf(a, b))``.
-The UDF is not evaluated element by element inside the lambda; it is applied once to the whole
-array outside the lambda, and the lambda reads the precomputed result. As a result the UDF runs
-over *every* element (for ``array_sort`` with a two-argument comparator, over every pair of
-elements), even ones a lambda would otherwise skip, such as the elements after ``exists`` finds a
-match or the branch of a ``when`` that is not taken. If the function can fail on some input,
-handle those cases inside the function itself. This path is enabled by
-``spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled`` (default ``true``); set it to
-``false`` to reject such queries at analysis instead.
-
-This works for a plain Python UDF as well as for a vectorized scalar UDF - a scalar pandas UDF
-(``pandas_udf``), a scalar Arrow UDF (``arrow_udf``), and their iterator variants. A vectorized
-UDF still receives its native batch (a ``pandas.Series`` / ``pyarrow.Array``, or an iterator of
-them) over the flattened array elements, and one output value is produced per input element.
+The same holds for a Python UDF inside a higher-order function's lambda -- a plain, scalar pandas
+(``pandas_udf``), scalar Arrow (``arrow_udf``) or iterator UDF. It is not evaluated element by
+element inside the lambda; it is applied once, outside the lambda, and the lambda reads the
+precomputed result. So the UDF runs over *every* element, even ones a lambda would otherwise skip
+(the elements after ``exists`` matches, or the untaken branch of a ``when``) -- if it can fail on
+some input, handle that inside the function. This covers ``transform``, ``filter``, ``exists``,
+``forall``, ``zip_with``, ``array_sort`` and the map functions, as well as *nested* lambdas
+(``transform(matrix, lambda row: transform(row, lambda x: udf(x)))``, including ones that capture
+the enclosing variable). A UDF inside ``aggregate`` / ``reduce`` is not supported, because the fold
+is sequential and cannot be applied once to the whole array. Toggle with
+``spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled`` (default ``true``; set ``false`` to
+reject at analysis).
The Arrow data type of the returned ``pyarrow.Array`` should match the declared ``returnType``.
When there is a mismatch, Spark will attempt to convert the returned data to the expected type
diff --git a/python/pyspark/sql/tests/test_udf_in_higher_order_function.py b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
index 9d0beb024abe..ff0788cf4cf0 100644
--- a/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
+++ b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
@@ -423,23 +423,142 @@ class UDFInHigherOrderFunctionTestsMixin:
df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")).count(), 0
)
- def test_udf_inside_nested_higher_order_function_fails(self):
- # The inner array is the outer lambda's variable, not a real column, so a UDF in the inner
- # lambda cannot be lifted onto an array outside the lambda. This must fail analysis. (A UDF
- # over the outer element is fine; see test_udf_outside_inner_higher_order_function.)
- df = self.spark.createDataFrame([([[1, 2], [3]],)], "values array<array<int>>")
+ def test_udf_inside_nested_lambda(self):
+ # A UDF in a *nested* lambda is lifted onto the fully flattened leaves and the nested
+ # structure is rebuilt around the result: `transform(matrix, row -> transform(row, x ->
+ # f(x)))` applies `f` to every leaf. The UDF runs once over all leaves (a depth-2 lift), and
+ # the result is compared against the equivalent native expression.
+ df = self.spark.createDataFrame(
+ [([[1, 2], [3]],), ([[], [4, 5]],), (None,), ([None, [6]],)],
+ "values array<array<int>>",
+ )
plus_one = udf(lambda x: x + 1, IntegerType())
is_even = udf(lambda x: x % 2 == 0, "boolean")
- nested = [
- sf.transform("values", lambda inner: sf.transform(inner, lambda x: plus_one(x))),
- # A `filter` inside a `transform`, with the UDF in the inner predicate.
- sf.transform("values", lambda inner: sf.filter(inner, lambda x: is_even(x))),
- ]
- for expr in nested:
- with self.assertRaises(AnalysisException) as ctx:
- df.select(expr).collect()
- self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+ # transform inside transform.
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda row: sf.transform(row, lambda x: plus_one(x))).alias(
+ "r"
+ )
+ ),
+ df.select(
+ sf.transform("values", lambda row: sf.transform(row, lambda x: x + 1)).alias("r")
+ ),
+ )
+ # filter inside transform: the inner result length differs from the input, but the UDF still
+ # runs over every leaf before the (JVM) filtering.
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda row: sf.filter(row, lambda x: is_even(x))).alias("r")
+ ),
+ df.select(
+ sf.transform("values", lambda row: sf.filter(row, lambda x: x % 2 == 0)).alias("r")
+ ),
+ )
+
+ def test_udf_inside_nested_lambda_capturing_outer_variable(self):
+ # The inner UDF reads both the inner element and the *enclosing* lambda's variable
+ # (`sf.size(row)`, where `row` is the outer element), so its argument depends on two nesting
+ # levels; the lift aligns both onto the leaves. Compared against the equivalent native
+ # expression, over null outer rows, null inner arrays, and empty inner arrays.
+ df = self.spark.createDataFrame(
+ [([[1, 2], [3, 4]],), (None,), ([[]],), ([None, [5]],)],
+ "values array<array<int>>",
+ )
+ add = udf(lambda a, b: a + b, IntegerType())
+ assertDataFrameEqual(
+ df.select(
+ sf.transform(
+ "values", lambda row: sf.transform(row, lambda x: add(x, sf.size(row)))
+ ).alias("r")
+ ),
+ df.select(
+ sf.transform(
+ "values", lambda row: sf.transform(row, lambda x: x + sf.size(row))
+ ).alias("r")
+ ),
+ )
+
+ def test_udf_inside_three_level_nested_lambda(self):
+ # Three levels of nesting: `f` is lifted to a depth-3 element-wise UDF, so the worker
+ # flattens three array levels to the leaves and re-nests three levels.
+ df = self.spark.createDataFrame(
+ [([[[1, 2], [3]], [[4]]],), (None,), ([[[], None]],)],
+ "values array<array<array<int>>>",
+ )
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ assertDataFrameEqual(
+ df.select(
+ sf.transform(
+ "values",
+ lambda a: sf.transform(a, lambda b: sf.transform(b, lambda x: plus_one(x))),
+ ).alias("r")
+ ),
+ df.select(
+ sf.transform(
+ "values",
+ lambda a: sf.transform(a, lambda b: sf.transform(b, lambda x: x + 1)),
+ ).alias("r")
+ ),
+ )
+
+ def test_vectorized_udf_inside_nested_lambda(self):
+ # All four vectorized flavors lifted out of a nested lambda (depth 2): the worker flattens
+ # two array levels to the leaves, runs the native-batch function once, and re-nests two
+ # levels. Includes null outer rows, null inner arrays, and empty inner arrays.
+ from typing import Iterator
+ import pandas as pd
+ import pyarrow as pa
+ from pyspark.sql.functions import pandas_udf, arrow_udf
+
+ df = self.spark.createDataFrame(
+ [([[1, 2], [3]],), ([[], [4, 5]],), (None,), ([None, [6]],)],
+ "values array<array<int>>",
+ )
+
+ @pandas_udf(IntegerType())
+ def plus_one_pandas(s: pd.Series) -> pd.Series:
+ return s + 1
+
+ @arrow_udf(IntegerType())
+ def plus_one_arrow(a: pa.Array) -> pa.Array:
+ return pa.compute.add(a, 1)
+
+ @pandas_udf(IntegerType())
+ def plus_one_pandas_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield s + 1
+
+ @arrow_udf(IntegerType())
+ def plus_one_arrow_iter(it: Iterator[pa.Array]) -> Iterator[pa.Array]:
+ for a in it:
+ yield pa.compute.add(a, 1)
+
+ native = df.select(
+ sf.transform("values", lambda row: sf.transform(row, lambda x: x + 1)).alias("r")
+ )
+ for f in (plus_one_pandas, plus_one_arrow, plus_one_pandas_iter, plus_one_arrow_iter):
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda row: sf.transform(row, lambda x: f(x))).alias("r")
+ ),
+ native,
+ )
+
+ def test_nested_lambda_with_nondeterministic_inner_argument_fails(self):
+ # The inner iterated argument is nondeterministic, so the rewrite - which references it more
+ # than once - would evaluate it independently and misalign the results. It must keep failing
+ # analysis at the nest root, even though the UDF itself is liftable.
+ df = self.spark.createDataFrame([([[1, 2], [3]],)], "values array<array<int>>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ with self.assertRaises(AnalysisException) as ctx:
+ df.select(
+ sf.transform(
+ "values", lambda row: sf.transform(sf.shuffle(row), lambda x: plus_one(x))
+ )
+ ).collect()
+ self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
def test_udf_outside_inner_higher_order_function(self):
# The UDF applies to the outer array's element (itself an array), which *is* a real
@@ -486,15 +605,17 @@ class UDFInHigherOrderFunctionTestsMixin:
[([2, 3, 4],)],
)
- def test_kwargs_call_still_fails(self):
- # A UDF called with a keyword argument carries a NamedArgumentExpression, which the lift
- # cannot preserve, so it must keep failing analysis rather than being rewritten.
- df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ def test_kwargs_call(self):
+ # A UDF called with a keyword argument is lifted too: the NamedArgumentExpression stays a
+ # direct child of the lifted UDF (only its value becomes an aligned array), so the runner
+ # still derives the kwargs mapping. add(x, y=10) = x + 10.
+ df = self.spark.createDataFrame([([1, 2],), (None,)], "values array<int>")
add = udf(lambda x, y: x + y, IntegerType())
- with self.assertRaises(AnalysisException) as ctx:
- df.select(sf.transform("values", lambda x: add(x, y=sf.lit(1)))).collect()
- self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: add(x, y=sf.lit(10))).alias("r")),
+ [([11, 12],), (None,)],
+ )
def test_zero_argument_udf_still_fails(self):
# A zero-argument UDF has no argument to carry the iterated array's shape, so the rewrite
@@ -618,22 +739,17 @@ class UDFInHigherOrderFunctionTestsMixin:
df.select(sf.transform("values", lambda x: plus_one(x))).collect()
self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
- def test_udf_in_aggregate_still_fails(self):
- # The fold is sequential: the values the UDF sees are outputs of earlier steps rather than
- # elements of a collection (folding [1,2,3] with `acc*2 + x` calls it on 0, 1, 4), and
- # computing the final accumulator alternates Python and JVM work once per element, which no
- # single UDF call can do. Restating the fold as an iteration over the whole column would
- # work but is an operator-level change, so a UDF anywhere in `aggregate` / `reduce` -
- # `merge` or `finish` - must fail rather than mislead.
+ def test_udf_in_aggregate_fails(self):
+ # `aggregate` / `reduce` is a sequential fold: the values a UDF sees are outputs of earlier
+ # steps, not elements of a collection, so it cannot be applied once to the whole array. A
+ # UDF anywhere in `aggregate` / `reduce` - `merge` or `finish` - must fail analysis.
df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
plus_one = udf(lambda x: x + 1, IntegerType())
double_it = udf(lambda acc: acc * 2, IntegerType())
aggregates = [
- # A UDF in `merge`, on the element and on the accumulator.
sf.aggregate("values", sf.lit(0), lambda acc, x: acc + plus_one(x)),
sf.aggregate("values", sf.lit(0), lambda acc, x: plus_one(acc) + x),
- # A UDF in `finish`.
sf.aggregate("values", sf.lit(0), lambda acc, x: acc + x, lambda acc: double_it(acc)),
# `reduce` is an alias of `aggregate`, so it is rejected the same way.
sf.reduce("values", sf.lit(0), lambda acc, x: acc + plus_one(x)),
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 4e77c630777d..215ddf9a7fe0 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -224,6 +224,17 @@ class EvalConf(Conf):
return None
return _parse_datatype_json_string(input_type)
+ @property
+ def elementwise_nesting(self) -> Optional[list]:
+ # Per-UDF nesting depth (parallel to the UDF list) for the element-wise lift: how many
+ # ``array`` levels the worker flattens off each argument and re-nests onto the result. A UDF
+ # in a single lambda is depth 1; one lifted out of nested lambdas is deeper. Absent/empty
+ # means depth 1 for every UDF. See ExtractPythonUDFFromLambda.
+ raw = self.get("elementwise_nesting", None)
+ if raw is None or raw == "":
+ return None
+ return [int(x) for x in raw.split(",")]
+
@property
def table_arg_offsets(self) -> Optional[list[int]]:
offsets = self.get("table_arg_offsets", None)
@@ -1857,6 +1868,68 @@ def _elementwise_renest(flat_values, shape_lengths, is_large):
return list_cls.from_arrays(offsets_arr, flat_values, mask=null_mask)
+def _elementwise_leaf_type(data_type, depth):
+ """The element type ``depth`` ``ArrayType`` levels below ``data_type``.
+
+ A lifted UDF's argument arrives as ``array^depth<T>`` (one ``array`` level per enclosing higher-
+ order function lambda); this peels them off to the scalar leaf ``T`` the user function sees. See
+ ``ExtractPythonUDFFromLambda``.
+ """
+ for _ in range(depth):
+ data_type = data_type.elementType
+ return data_type
+
+
+def _elementwise_flatten_deep(col, depth):
+ """Flatten ``depth`` list levels off ``col``, keeping each level's shape for re-nesting.
+
+ Returns ``(leaf, shape_levels, is_large_levels)``: ``leaf`` is the fully flattened element
+ ``pa.Array`` (the leaves of the ``depth``-deep nesting), ``shape_levels[k]`` is the per-slot
+ length (``None`` for a null slot) at level ``k`` (0 = outermost), and ``is_large_levels[k]``
+ whether that level is a ``LargeListArray``. ``depth`` is 1 for a UDF in a single lambda and more
+ for one lifted out of nested lambdas. Shared by the element-wise worker paths. See
+ ``ExtractPythonUDFFromLambda``.
+ """
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ shape_levels = []
+ is_large_levels = []
+ cur = col
+ for _ in range(depth):
+ shape_levels.append(pc.list_value_length(cur).to_pylist())
+ is_large_levels.append(pa.types.is_large_list(cur.type))
+ cur = cur.flatten()
+ return cur, shape_levels, is_large_levels
+
+
+def _elementwise_flatten_leaf(col, depth):
+ """Flatten ``depth`` list levels off ``col`` to its leaf ``pa.Array``, without capturing shape.
+
+ A lifted UDF re-nests its result by the *first* argument's per-level shapes only, so the other
+ arguments need just their leaves. This skips the ``pc.list_value_length(...).to_pylist()`` and
+ ``is_large`` bookkeeping ``_elementwise_flatten_deep`` does for the first argument. See
+ ``ExtractPythonUDFFromLambda``.
+ """
+ cur = col
+ for _ in range(depth):
+ cur = cur.flatten()
+ return cur
+
+
+def _elementwise_renest_deep(flat_values, shape_levels, is_large_levels):
+ """Re-nest a flat leaf Array back through ``len(shape_levels)`` list levels, innermost first.
+
+ Inverse of ``_elementwise_flatten_deep``: rebuilds the ``array^depth<R>`` result from the flat
+ per-leaf results and the per-level shapes captured while flattening the input. For ``depth`` 1
+ this is a single ``_elementwise_renest``.
+ """
+ result = flat_values
+ for lengths, is_large in zip(reversed(shape_levels), reversed(is_large_levels)):
+ result = _elementwise_renest(result, lengths, is_large)
+ return result
+
+
def _elementwise_flatten_column(flat, element_type, is_pandas, runner_conf):
"""Adapt one already-flattened ``array<T>`` element column to the vectorized fn's input.
@@ -3099,33 +3172,50 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
require_minimum_pyarrow_version()
import pyarrow as pa
- import pyarrow.compute as pc
# Element-wise UDFs back higher-order lambdas like transform(arr, x -> udf(x)).
# ExtractPythonUDFFromLambda rewrites them so the UDF receives *all* array elements
- # at once (as ``array<T>``) rather than per-element. Flatten once, evaluate once over
- # the batch, then re-nest with input offsets. Example: array<int> -> udf -> array<int>.
+ # at once (as ``array<T>``) rather than per-element. Flatten each argument down to its
+ # leaves, evaluate once over the batch, then re-nest with the input offsets. A UDF lifted
+ # out of nested lambdas (e.g. transform(arr, i -> transform(i, x -> udf(x)))) flattens more
+ # than one ``array`` level - its per-UDF depth comes from ``elementwise_nesting``.
+ # Example: array<array<int>> -> udf(depth 2) -> array<array<int>>.
# UDF preparation
+ input_fields = list(eval_conf.input_type)
+ nesting = eval_conf.elementwise_nesting
udf_infos = []
- for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in udfs:
+ for udf_index, udf in enumerate(udfs):
+ udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type = udf
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
udf_func, udf_args_offsets, udf_kwargs_offsets
)
- # UDF returns one value per element; return type was pickled, unchanged.
- # This is per-element, so element type equals the declared return type.
- element_return_type = udf_return_type
+ depth = nesting[udf_index] if nesting is not None else 1
+ # Each argument arrives as ``array^depth<T>``; convert its leaves with the element type
+ # ``T`` reached by peeling ``depth`` array levels.
+ arg_converters = [
+ ArrowTableToRowsConversion._create_converter(
+ _elementwise_leaf_type(input_fields[o].dataType, depth),
+ none_on_identity=True,
+ binary_as_bytes=runner_conf.binary_as_bytes,
+ )
+ for o in args_kwargs_offsets
+ ]
udf_infos.append(
(
wrapped_func,
args_kwargs_offsets,
+ depth,
+ arg_converters,
+ # UDF returns one value per element; return type was pickled, unchanged. This is
+ # per-element, so element type equals the declared return type.
to_arrow_type(
- element_return_type,
+ udf_return_type,
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
),
LocalDataToArrowConversion._create_converter(
- element_return_type,
+ udf_return_type,
none_on_identity=True,
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
),
@@ -3133,18 +3223,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
)
col_names = [f"_{i}" for i in range(len(udfs))]
- # Input: every argument arrives as ``array<T>`` aligned with the iterated array.
- # Flatten once per column; convert elements with the array's element type.
- input_fields = list(eval_conf.input_type)
- arrow_to_py_converters = [
- ArrowTableToRowsConversion._create_converter(
- f.dataType.elementType,
- none_on_identity=True,
- binary_as_bytes=runner_conf.binary_as_bytes,
- )
- for f in input_fields
- ]
-
@fail_on_stopiteration
def _evaluate_elementwise_udf(udf_func, rows):
if runner_conf.arrow_concurrency_level <= 0:
@@ -3154,48 +3232,42 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
with ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool:
return list(pool.map(lambda row: udf_func(*row), rows))
- def renest_spec(shape):
- """Offsets, list class and null mask that re-nest a flat result list by ``shape``.
-
- Null rows stay null and consume no offsets. ``ListArray`` uses int32 offsets and
- ``LargeListArray`` int64, so the input's list width is preserved.
- """
- lengths = pc.list_value_length(shape).to_pylist()
- offsets = []
- running = 0
- for n in lengths:
- offsets.append(running)
- if n is not None:
- running += n
- offsets.append(running)
- is_large = pa.types.is_large_list(shape.type)
- list_cls = pa.LargeListArray if is_large else pa.ListArray
- offsets_arr = pa.array(offsets, type=pa.int64() if is_large else pa.int32())
- null_mask = pa.array([n is None for n in lengths], type=pa.bool_())
- total_elements = running
- return list_cls, offsets_arr, null_mask, total_elements
-
def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
for input_batch in data:
- # Flatten each list column once to its element list, converting to Python.
- columns = []
- for col, conv in zip(input_batch.itercolumns(), arrow_to_py_converters):
- values = ArrowTableToRowsConversion._to_pylist(col.flatten())
- if conv is not None:
- values = [conv(v) for v in values]
- columns.append(values)
-
# Each UDF is re-nested by *its own* first argument's shape. ExtractPythonUDFs can
- # fuse UDFs over differently shaped arrays into one batch, so a single shared shape
- # would misalign every UDF but the first. The rewrite always passes at least one
- # array argument, so `offsets_meta` is non-empty.
+ # fuse UDFs over differently shaped/nested arrays into one batch, so a single shared
+ # shape would misalign every UDF but the first. The rewrite always passes at least
+ # one array argument, so `offsets` is non-empty.
output_arrays = []
- for wrapped_func, offsets_meta, arrow_element_type, result_conv in udf_infos:
- list_cls, offsets_arr, null_mask, total_elements = renest_spec(
- input_batch.column(offsets_meta[0])
+ for info in udf_infos:
+ (
+ wrapped_func,
+ offsets,
+ depth,
+ arg_converters,
+ arrow_element_type,
+ result_conv,
+ ) = info
+ # Flatten each argument `depth` list levels to its leaves; the first argument's
+ # per-level shapes drive the re-nest.
+ leaf0, shape_levels, is_large_levels = _elementwise_flatten_deep(
+ input_batch.column(offsets[0]), depth
)
+ columns = []
+ for i, (o, conv) in enumerate(zip(offsets, arg_converters)):
+ leaf = (
+ leaf0
+ if i == 0
+ else _elementwise_flatten_leaf(input_batch.column(o), depth)
+ )
+ values = ArrowTableToRowsConversion._to_pylist(leaf)
+ if conv is not None:
+ values = [conv(v) for v in values]
+ columns.append(values)
+
+ total_elements = len(columns[0])
# Stream the argument tuples rather than materializing a batch-sized list.
- rows = zip(*[columns[o] for o in offsets_meta])
+ rows = zip(*columns)
results = _evaluate_elementwise_udf(wrapped_func, rows)
verify_result_row_count(len(results), total_elements)
@@ -3214,7 +3286,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
target_type=arrow_element_type, safe=runner_conf.safecheck
)
output_arrays.append(
- list_cls.from_arrays(offsets_arr, flat_arr, mask=null_mask)
+ _elementwise_renest_deep(flat_arr, shape_levels, is_large_levels)
)
yield pa.RecordBatch.from_arrays(output_arrays, col_names)
@@ -3231,7 +3303,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
require_minimum_pyarrow_version()
import pyarrow as pa
- import pyarrow.compute as pc
# A scalar pandas or Arrow UDF lifted out of a higher-order function's lambda by
# ExtractPythonUDFFromLambda. Each argument arrives as ``array<T>`` aligned with the
@@ -3241,8 +3312,11 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
# input's offsets - one row in, one row out, one Python round trip per batch.
is_pandas = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF
+ input_fields = list(eval_conf.input_type)
+ nesting = eval_conf.elementwise_nesting
udf_infos = []
- for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in udfs:
+ for udf_index, udf in enumerate(udfs):
+ udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type = udf
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
udf_func, udf_args_offsets, udf_kwargs_offsets
)
@@ -3251,37 +3325,59 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
arrow_element_type = to_arrow_type(
udf_return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types
)
+ depth = nesting[udf_index] if nesting is not None else 1
+ # Each argument arrives as ``array^depth<T>``; the vectorized function must see the leaf
+ # element type ``T`` reached by peeling ``depth`` array levels.
+ arg_leaf_types = [
+ _elementwise_leaf_type(input_fields[o].dataType, depth) for o in args_kwargs_offsets
+ ]
udf_infos.append(
- (wrapped_func, args_kwargs_offsets, udf_return_type, arrow_element_type)
+ (
+ wrapped_func,
+ args_kwargs_offsets,
+ udf_return_type,
+ arrow_element_type,
+ depth,
+ arg_leaf_types,
+ )
)
col_names = [f"_{i}" for i in range(len(udfs))]
if is_pandas:
import pandas as pd
- # Each argument is ``array<T>``; the vectorized function must see the element type ``T``.
- element_types = [f.dataType.elementType for f in eval_conf.input_type]
-
def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
for input_batch in data:
- # Flatten each list column to its element array once per batch and share it across
- # fused UDFs (the 102 path does the same), rather than re-flattening per UDF.
- flat_columns = [
- _elementwise_flatten_column(
- col.flatten(), element_types[i], is_pandas, runner_conf
+ output_arrays = []
+ for (
+ wrapped_func,
+ offsets,
+ return_type,
+ arrow_element_type,
+ depth,
+ arg_leaf_types,
+ ) in udf_infos:
+ # Flatten each argument `depth` list levels to its leaves and adapt to the
+ # vectorized fn's input. Different UDFs in one operator may iterate differently
+ # shaped or differently nested arrays, so each flattens and re-nests by its own
+ # argument (the first argument's per-level shapes drive the re-nest).
+ leaf0, shape_levels, is_large_levels = _elementwise_flatten_deep(
+ input_batch.column(offsets[0]), depth
)
- for i, col in enumerate(input_batch.itercolumns())
- ]
+ total_elements = len(leaf0)
+ flat_columns = [
+ _elementwise_flatten_column(
+ leaf0
+ if i == 0
+ else _elementwise_flatten_leaf(input_batch.column(o), depth),
+ t,
+ is_pandas,
+ runner_conf,
+ )
+ for i, (o, t) in enumerate(zip(offsets, arg_leaf_types))
+ ]
- output_arrays = []
- for wrapped_func, offsets, return_type, arrow_element_type in udf_infos:
- # Re-nest by this UDF's first argument's shape. Different UDFs in one operator
- # may iterate differently shaped arrays, so each re-nests by its own argument.
- shape = input_batch.column(offsets[0])
- shape_lengths = pc.list_value_length(shape).to_pylist()
- total_elements = sum(n for n in shape_lengths if n is not None)
-
- result = wrapped_func(*[flat_columns[o] for o in offsets])
+ result = wrapped_func(*flat_columns)
if is_pandas:
if not hasattr(result, "__len__"):
pd_type = (
@@ -3316,9 +3412,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
flat_arr = _elementwise_result_to_arrow(
result, return_type, arrow_element_type, is_pandas, runner_conf
)
- nested = _elementwise_renest(
- flat_arr, shape_lengths, pa.types.is_large_list(shape.type)
- )
+ nested = _elementwise_renest_deep(flat_arr, shape_levels, is_large_levels)
output_arrays.append(nested)
yield pa.RecordBatch.from_arrays(output_arrays, col_names)
@@ -3337,7 +3431,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
import collections
import pyarrow as pa
- import pyarrow.compute as pc
is_pandas = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF
if is_pandas:
@@ -3354,36 +3447,45 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
# output to input positionally by row (one ``array<R>`` per input ``array<T>`` row, in
# order), buffering a FIFO of the per-row element counts to re-group the streamed flat
# results back into arrays. Output batch boundaries need not match input ones.
- # Each argument is ``array<T>``; the pandas function must see each argument's own element
- # type ``T`` (arguments may differ, e.g. an outer column repeated into an aligned array).
- element_types = [f.dataType.elementType for f in eval_conf.input_type]
+ # Each argument arrives as ``array^depth<T>``; the vectorized function must see the leaf
+ # element type ``T`` (arguments may differ, e.g. an outer column repeated into an aligned
+ # array). ``depth`` > 1 for a UDF lifted out of nested lambdas.
+ input_fields = list(eval_conf.input_type)
+ nesting = eval_conf.elementwise_nesting
+ depth = nesting[0] if nesting is not None else 1
+ arg_leaf_types = [
+ _elementwise_leaf_type(input_fields[o].dataType, depth) for o in args_offsets
+ ]
arrow_element_type = to_arrow_type(
return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types
)
- is_large = None # set from the first input batch; the list width is uniform per column.
def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
- # FIFO of per-row element counts (None for a null array) awaiting their flat results,
- # and the flat elements produced so far but not yet enough to complete the head shapes.
+ # FIFO of per-input-batch nesting shapes awaiting their flat leaf results. Each entry
+ # (per-level shapes, per-level list width, leaf count) re-nests one input batch's worth
+ # of rows back to ``array^depth<R>`` once that many leaf results have streamed in.
pending_shapes: "collections.deque" = collections.deque()
num_input_elements = 0
def extract_flat(batch: pa.RecordBatch):
- nonlocal is_large, num_input_elements
- shape = batch.column(args_offsets[0])
- if is_large is None:
- is_large = pa.types.is_large_list(shape.type)
- pending_shapes.append(pc.list_value_length(shape).to_pylist())
- # Flatten each argument to its element column; the user function sees the flat
- # elements as a pandas Series / DataFrame (pandas) or a pa.Array (Arrow). Each
- # argument is converted with its own element type.
+ nonlocal num_input_elements
+ # Flatten each argument `depth` levels to its leaves; the first argument's per-level
+ # shapes re-nest this batch's rows. The user function sees the flat leaves as a
+ # pandas Series / DataFrame (pandas) or a pa.Array (Arrow), each with its leaf type.
+ leaf0, shape_levels, is_large_levels = _elementwise_flatten_deep(
+ batch.column(args_offsets[0]), depth
+ )
+ pending_shapes.append((shape_levels, is_large_levels, len(leaf0)))
+ num_input_elements += len(leaf0)
flat_cols = [
_elementwise_flatten_column(
- batch.column(o).flatten(), element_types[o], is_pandas, runner_conf
+ leaf0 if i == 0 else _elementwise_flatten_leaf(batch.column(o), depth),
+ arg_leaf_types[i],
+ is_pandas,
+ runner_conf,
)
- for o in args_offsets
+ for i, o in enumerate(args_offsets)
]
- num_input_elements += len(flat_cols[0])
return flat_cols[0] if len(flat_cols) == 1 else tuple(flat_cols)
flat_args_iter = map(extract_flat, data)
@@ -3421,8 +3523,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
def emit_ready():
nonlocal pending_chunks, pending_len
while pending_shapes:
- lengths = pending_shapes[0]
- needed = sum(n for n in lengths if n is not None)
+ shape_levels, is_large_levels, needed = pending_shapes[0]
if needed > pending_len:
break
pending_shapes.popleft()
@@ -3438,7 +3539,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
remainder = combined.slice(needed)
pending_chunks = [remainder] if len(remainder) else []
pending_len -= needed
- nested = _elementwise_renest(flat, lengths, bool(is_large))
+ nested = _elementwise_renest_deep(flat, shape_levels, is_large_levels)
yield pa.RecordBatch.from_arrays([nested], ["_0"])
def process_results():
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
index 66e2e78aef27..6687c163efad 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
@@ -506,9 +506,16 @@ trait CheckAnalysis extends LookupCatalog with QueryErrorsBase with PlanToString
// `ExtractPythonUDFFromLambda` rewrites the plan in the optimizer so the UDF is
// applied to the whole array outside the lambda instead; those are allowed through
// here. Everything else must still fail, or nothing downstream can evaluate it.
+ //
+ // Judge only at a *nest root* - a HOF that iterates real columns, not a free lambda
+ // variable. `canRewritePythonUDFInLambda` validates the whole nest below the root
+ // (a UDF in a nested lambda is lifted out one level at a time), and a root's
+ // `functions.exists` sees UDFs at every depth, so firing on inner HOFs too would
+ // double-report and reject nests the rule actually handles.
case hof: HigherOrderFunction
if hof.resolved && hof.functions
.exists(_.exists(_.isInstanceOf[PythonUDF])) &&
+ !PythonUDF.hasFreeLambdaVariable(hof) &&
!(conf.pythonUDFInHigherOrderFunctionEnabled &&
PythonUDF.canRewritePythonUDFInLambda(hof)) =>
// Name the offending UDF: the first one of an unsupported eval type if any (e.g. a
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
index e421d946a669..c7d9c49b640c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
@@ -69,9 +69,11 @@ object PythonUDF {
* - a zero-argument call, `f()`: the lift turns each argument into an aligned array, so with no
* argument there is no array to carry the iterated shape, and the element-wise UDF would
* reach the worker with no input column and crash there instead of failing analysis;
- * - a call with named arguments: its `NamedArgumentExpression` children would be buried inside
- * the generated `ArrayTransform`, losing the kwargs mapping the runner derives from the
- * direct children;
+ * - a call with named arguments on an *iterator* UDF (scalar pandas / Arrow iterator): iterator
+ * UDFs do not take keyword arguments (the worker ignores their kwargs offsets), so such a
+ * call is invalid regardless of the lift. Named arguments on the non-iterator flavors are
+ * supported: the lift keeps each `NamedArgumentExpression` as a direct child of the lifted
+ * UDF (only its value becomes an aligned array), so the runner still derives the kwargs map;
* - a UDF whose argument or return type involves a UDT: the lift forces an Arrow element-wise
* eval type, which has no UDT fallback (unlike `correctEvalType`'s Arrow -> pickle path), so
* it would fail at runtime instead of at analysis.
@@ -84,12 +86,26 @@ object PythonUDF {
case udf: PythonUDF =>
isElementwiseRewritableEvalType(udf.evalType) &&
udf.children.nonEmpty &&
- !udf.children.exists(_.isInstanceOf[NamedArgumentExpression]) &&
+ (supportsNamedArgumentsWhenLifted(udf.evalType) ||
+ !udf.children.exists(_.isInstanceOf[NamedArgumentExpression])) &&
!containsUDT(udf.dataType) &&
!udf.children.exists(c => containsUDT(c.dataType))
case _ => false
}
+ /**
+ * Whether a lifted UDF of this eval type can carry keyword arguments. Iterator UDFs (scalar
+ * pandas / Arrow iterator, and their lifted element-wise forms) cannot - the worker binds only
+ * positional arguments for them - so a named-argument call on an iterator UDF is not rewritable.
+ */
+ private def supportsNamedArgumentsWhenLifted(evalType: Int): Boolean = evalType match {
+ case PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF |
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF |
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF |
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF => false
+ case _ => true
+ }
+
private def isElementwiseRewritableEvalType(evalType: Int): Boolean = evalType match {
case PythonEvalType.SQL_BATCHED_UDF |
PythonEvalType.SQL_ARROW_BATCHED_UDF |
@@ -97,15 +113,21 @@ object PythonUDF {
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF |
PythonEvalType.SQL_SCALAR_ARROW_UDF |
PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF => true
- case _ => false
+ // The already-lifted element-wise types are rewritable again: a UDF inside a nested lambda is
+ // lifted once onto the inner lambda's variable (producing an element-wise UDF over that
+ // variable) and then re-lifted onto the enclosing array, incrementing its nesting depth. Users
+ // cannot create these eval types directly, so they only appear mid-rewrite - `CheckAnalysis`,
+ // which runs before this rule, never sees them.
+ case _ => PythonEvalType.isElementwiseUDF(evalType)
}
/**
* The eval type a rewritable UDF runs under once lifted out of the lambda. Each maps to the
* element-wise flavor that preserves its worker contract: the row-at-a-time types share the one
* pickle-based element-wise path, while each vectorized scalar type keeps its own pandas- vs.
- * Arrow-shaped batching and iterator behavior. `evalType` must satisfy
- * [[isElementwiseRewritableEvalType]].
+ * Arrow-shaped batching and iterator behavior. An already-lifted element-wise type maps to itself
+ * (re-lifting for a nested lambda keeps the flavor and only bumps the nesting depth). `evalType`
+ * must satisfy [[isElementwiseRewritableEvalType]].
*/
def liftedElementwiseEvalType(evalType: Int): Int = evalType match {
case PythonEvalType.SQL_BATCHED_UDF | PythonEvalType.SQL_ARROW_BATCHED_UDF =>
@@ -118,44 +140,71 @@ object PythonUDF {
PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF
case PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF =>
PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF
+ case elementwise if PythonEvalType.isElementwiseUDF(elementwise) => elementwise
case other =>
throw internalError(s"Not a rewritable elementwise UDF eval type: $other")
}
/**
* Whether every Python UDF in `hof`'s lambdas can be lifted out by `ExtractPythonUDFFromLambda`.
- * Used by `CheckAnalysis` to decide whether to reject the plan. These shapes cannot be rewritten:
- * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
- * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
- * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
- * - a UDF in `aggregate` / `reduce`: the fold is sequential, so it sees earlier steps' outputs;
- * - a *nondeterministic iterated argument*, `filter(shuffle(arr), x -> f(x))`: the rewrite
- * references that argument several times (the carrier's `c0`, each lifted UDF's argument, the
- * `map_keys`/`map_values` desugar, the pairwise `array_sort` path), and nondeterministic
- * expressions are not subexpression-eliminated, so the copies would evaluate independently
- * and disagree - keeping the results misaligned. (This is distinct from a nondeterministic
- * UDF *call*, which `ExtractPythonUDFFromLambda.liftKey` keeps distinct but well-defined.)
+ * Used by `CheckAnalysis` to decide whether to reject the plan; `hof` must be a *nest root* - one
+ * that iterates real columns, not a free lambda variable - because `CheckAnalysis` fires only at
+ * nest roots (see its guard) and this predicate validates the whole nest below the root.
+ *
+ * Both the row-at-a-time and the vectorized scalar eval types are liftable, in a single lambda or
+ * nested lambdas: an inner lambda's UDF is lifted onto its (enclosing-variable) argument and then
+ * re-lifted outward one array level at a time, so `transform(arr, i -> transform(i, x -> f(x)))`
+ * works (`f` lifts to a depth-2 element-wise UDF over `arr`). A UDF in a nested *argument*,
+ * `transform(arr, x -> transform(udf(x), y -> y))`, lifts onto `arr` the same way.
+ *
+ * These shapes still cannot be rewritten and are rejected:
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements, and cannot be applied once to the whole array (see
+ * [[isRewritableShape]]).
+ * - a *nondeterministic iterated argument*, `filter(shuffle(arr), x -> f(x))` (at the root or
+ * any nested level): the rewrite references that argument several times (the carrier's `c0`,
+ * each lifted UDF's argument, the `map_keys`/`map_values` desugar, the pairwise `array_sort`
+ * path), and nondeterministic expressions are not subexpression-eliminated, so the copies
+ * would evaluate independently and disagree - keeping the results misaligned. (This is
+ * distinct from a nondeterministic UDF *call*, which `ExtractPythonUDFFromLambda.liftKey`
+ * keeps distinct but well-defined.)
+ * - a HOF (at any level) whose shape the rewrite does not model (see [[isRewritableShape]]).
*/
def canRewritePythonUDFInLambda(hof: HigherOrderFunction): Boolean = {
- // Only row-at-a-time eval types are supported; a vectorized UDF is not rewritable regardless
- // of the enclosing function.
+ // Every Python UDF anywhere in the lambdas must be a rewritable flavor. `collect` is recursive,
+ // so this also covers UDFs in nested lambdas.
val allUDFsRewritable = hof.functions.forall { f =>
f.collect { case udf: PythonUDF => udf }.forall(isElementwiseRewritableUDF)
}
// Reading a free lambda variable means `hof` is itself nested in an enclosing lambda, so the
- // array it iterates is not a real column (the nested case above).
+ // array it iterates is not a real column. Such an inner HOF is validated as part of its
+ // enclosing root's nest by `everyHofInNestRewritable`, never on its own.
val iteratesRealColumns = !hasFreeLambdaVariable(hof)
- // The rewrite duplicates the iterated argument expression, so a nondeterministic one would be
- // evaluated more than once with diverging results.
- val deterministicArguments = hof.arguments.forall(_.deterministic)
- allUDFsRewritable && iteratesRealColumns && deterministicArguments && isRewritableShape(hof)
+ iteratesRealColumns && allUDFsRewritable && everyHofInNestRewritable(hof)
+ }
+
+ /**
+ * Whether `hof` and every higher-order function nested within its lambda bodies is a rewritable
+ * shape (see [[isRewritableShape]]) with deterministic arguments. This is the recursive core that
+ * supports UDFs in *nested* lambdas: every HOF on the path from the root down to a UDF must be
+ * rewritable, because the rule lifts the UDF out one lambda level at a time. A nested HOF's
+ * iterated argument is legitimately an enclosing lambda variable, which is why the free-variable
+ * check in [[canRewritePythonUDFInLambda]] applies only at the root, not here.
+ */
+ private def everyHofInNestRewritable(hof: HigherOrderFunction): Boolean = {
+ val nestedHofs = hof.functions.flatMap(_.collect { case h: HigherOrderFunction => h })
+ (hof +: nestedHofs).forall { h =>
+ isRewritableShape(h) && h.arguments.forall(_.deterministic)
+ }
}
/**
* The structural assumption the rewrite makes: one lambda with plain-variable parameters, over at
- * least one array- or map-valued argument. `aggregate` / `reduce` fail this - they have two
- * lambdas (`merge`, `finish`) - so a UDF in a fold is rejected. Checking the shape rather than
- * listing classes means a new function of a familiar shape needs no change here.
+ * least one array- or map-valued argument (`transform`, `filter`, the map family, ...).
+ * `aggregate` / `reduce` fail this - they have two lambdas (`merge`, `finish`) - so a UDF in a
+ * fold is rejected: the fold is sequential, so the UDF sees earlier steps' outputs, not array
+ * elements. Checking the shape rather than listing classes means a new function of a familiar
+ * shape needs no change here.
*
* The function must also carry one of the result-type marker traits the rewrite dispatches on
* ([[ResultTypeFromArgument]] or [[ResultTypeFromFunction]]). Every built-in single-lambda HOF is
@@ -280,7 +329,14 @@ case class PythonUDF(
children: Seq[Expression],
evalType: Int,
udfDeterministic: Boolean,
- resultId: ExprId = NamedExpression.newExprId)
+ resultId: ExprId = NamedExpression.newExprId,
+ // For an element-wise UDF lifted out of a higher-order function's lambda (see
+ // `ExtractPythonUDFFromLambda`), the number of `array` levels the Python worker flattens off
+ // each argument before invoking the function, and re-nests onto the result: 1 for a UDF in a
+ // single lambda, and one more for each enclosing lambda when the UDF is lifted out of a nested
+ // lambda (e.g. `transform(arr, i -> transform(i, x -> f(x)))` lifts `f` to depth 2). Ignored
+ // for every non-element-wise eval type, where it stays at its default of 1.
+ elementwiseNestingDepth: Int = 1)
extends Expression with PythonFuncExpression with Unevaluable {
lazy val resultAttribute: Attribute = AttributeReference(toPrettySQL(this), dataType, nullable)(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala
index 6c6b96f4f590..2f90d718294b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala
@@ -244,8 +244,16 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes {
override lazy val canonicalized: Expression = {
var currExprId = -1
+ // Number the lambda variables of this higher-order function, but only those not already
+ // canonicalized. A canonical `NamedLambdaVariable` carries `value = null` (see the rename
+ // below); an original one carries an `AtomicReference`. When this HOF is nested inside another,
+ // the enclosing HOF's `canonicalized` renames every variable in the whole subtree first, then
+ // canonicalizes the children - which re-enters this method on the nested HOF. Re-numbering the
+ // already-renamed variables here would give a reference to an *enclosing* lambda's variable a
+ // different id than its binding, leaking it into `references`; skipping them keeps a variable
+ // and its references in agreement across nesting levels.
val argumentMap = functions.flatMap(_.collect {
- case l: NamedLambdaVariable =>
+ case l: NamedLambdaVariable if l.value != null =>
currExprId += 1
l.exprId -> currExprId
}).toMap
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
index 5206b15877c5..0170d4354a6a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
@@ -216,7 +216,10 @@ class ArrowEvalPythonEvaluatorFactory(
pythonRunnerConf,
pythonMetrics,
jobArtifactUUID,
- sessionUUID) with BatchedPythonArrowInput
+ sessionUUID,
+ // Parallel to `funcs` (both come from `udfs` in order); tells the worker how many `array`
+ // levels each element-wise UDF flattens/re-nests (see `PythonUDF.elementwiseNestingDepth`).
+ udfs.map(_.elementwiseNestingDepth)) with BatchedPythonArrowInput
val columnarBatchIter = pyRunner.compute(batchIter, context.partitionId(), context)
columnarBatchIter.flatMap { batch =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
index d0383c32a889..22ff6f711819 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
@@ -130,27 +130,20 @@ class ArrowPythonWithNamedArgumentRunner(
pythonRunnerConf: Map[String, String],
pythonMetrics: Map[String, SQLMetric],
jobArtifactUUID: Option[String],
- sessionUUID: Option[String])
+ sessionUUID: Option[String],
+ // Per-UDF element-wise nesting depth, parallel to `funcs` (see
+ // `PythonUDF.elementwiseNestingDepth`). Empty (the default) means depth 1 for every UDF, so the
+ // many non-element-wise construction sites need not pass it.
+ elementwiseNestingDepths: Seq[Int] = Nil)
extends RowInputArrowPythonRunner(
funcs, evalType, argMetas.map(_.map(_.offset)), schema, timeZoneId, largeVarTypes,
pythonMetrics, jobArtifactUUID, sessionUUID) {
override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf
- override protected def evalConf: Map[String, String] = {
- // An element-wise UDF receives each argument as an `array<T>` column and flattens it, so like
- // the Arrow batched UDF it needs the input schema to convert the incoming Arrow types. This
- // covers the row-at-a-time lift (SQL_ARROW_ELEMENTWISE_UDF) as well as the vectorized lifts
- // (scalar pandas / Arrow, and their iterator variants), which flatten the same way.
- if (PythonEvalType.isElementwiseUDF(evalType) ||
- evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
- super.evalConf ++ Map(
- "input_type" -> schema.json
- )
- } else {
- super.evalConf
- }
- }
+ override protected def evalConf: Map[String, String] =
+ ArrowPythonRunner.elementwiseEvalConf(
+ super.evalConf, evalType, schema, elementwiseNestingDepths)
override protected def writeUDF(dataOut: DataOutputStream): Unit = {
PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas)
@@ -158,6 +151,29 @@ class ArrowPythonWithNamedArgumentRunner(
}
object ArrowPythonRunner {
+ /**
+ * Adds the `input_type` (and, for element-wise UDFs, the per-UDF `elementwise_nesting`) entries
+ * an Arrow runner sends to the worker via eval-conf. An element-wise UDF receives each argument
+ * as an `array<T>` column and flattens it, so like the Arrow batched UDF it needs the input
+ * schema to convert the incoming Arrow types; a lifted UDF from *nested* lambdas re-nests more
+ * than one level, so it also needs its per-UDF nesting depth. Shared by row and columnar runners.
+ */
+ def elementwiseEvalConf(
+ base: Map[String, String],
+ evalType: Int,
+ schema: StructType,
+ elementwiseNestingDepths: Seq[Int]): Map[String, String] = {
+ if (PythonEvalType.isElementwiseUDF(evalType)) {
+ base ++ Map(
+ "input_type" -> schema.json,
+ "elementwise_nesting" -> elementwiseNestingDepths.mkString(","))
+ } else if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
+ base ++ Map("input_type" -> schema.json)
+ } else {
+ base
+ }
+ }
+
/** Return Map with conf settings to be used in ArrowPythonRunner */
def getPythonRunnerConfMap(conf: SQLConf): Map[String, String] = {
val confMap = collection.mutable.Map.empty[String, String]
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
index bad609919a4b..3981602875ad 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
@@ -202,7 +202,7 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory(
pyFuncs, evalType, argMetas, udfInputSchema,
sessionLocalTimeZone, largeVarTypes, pythonRunnerConf,
pythonMetrics, jobArtifactUUID, sessionUUID,
- columnIndices)
+ columnIndices, udfs.map(_.elementwiseNestingDepth))
val resultIter = pyRunner.compute(
bufferedIter, context.partitionId(), context)
@@ -262,7 +262,7 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory(
pyFuncs, evalType, argMetas, udfInputSchema,
sessionLocalTimeZone, largeVarTypes, pythonRunnerConf,
pythonMetrics, jobArtifactUUID, sessionUUID,
- inputColumnIndices.get)
+ inputColumnIndices.get, udfs.map(_.elementwiseNestingDepth))
pyRunner.compute(
bufferedIter, context.partitionId(), context)
} else {
@@ -279,7 +279,8 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory(
val pyRunner = new ArrowPythonWithNamedArgumentRunner(
pyFuncs, evalType, argMetas, udfInputSchema,
sessionLocalTimeZone, largeVarTypes, pythonRunnerConf,
- pythonMetrics, jobArtifactUUID, sessionUUID
+ pythonMetrics, jobArtifactUUID, sessionUUID,
+ udfs.map(_.elementwiseNestingDepth)
) with BasicPythonArrowInput
pyRunner.compute(
batchIter, context.partitionId(), context)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
index 41a563a7302f..de19b43d32c1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
@@ -42,7 +42,10 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner(
override val pythonMetrics: Map[String, SQLMetric],
jobArtifactUUID: Option[String],
sessionUUID: Option[String],
- override protected val inputColumnIndices: Array[Int])
+ override protected val inputColumnIndices: Array[Int],
+ // Per-UDF element-wise nesting depth, parallel to `funcs` (see
+ // `PythonUDF.elementwiseNestingDepth`); empty means depth 1 for every UDF.
+ elementwiseNestingDepths: Seq[Int] = Nil)
extends BaseArrowPythonRunner[ColumnarBatch, ColumnarBatch](
funcs, evalType, argMetas.map(_.map(_.offset)), schema, timeZoneId, largeVarTypes,
pythonMetrics, jobArtifactUUID, sessionUUID)
@@ -51,20 +54,14 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner(
override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf
- // The input schema travels in evalConf (key "input_type"), exactly like
- // ArrowPythonWithNamedArgumentRunner. It must NOT be written into the UDF command section:
+ // The input schema (and per-UDF nesting depth) travel in evalConf, exactly like
+ // ArrowPythonWithNamedArgumentRunner. They must NOT be written into the UDF command section:
// the worker's WorkerInitInfo.from_stream reads that section as a UDF count followed by UDF
// entries, so an extra UTF string there desyncs the whole init-message parse -- the worker
// then blocks waiting for bytes past the end of the init message and the task hangs forever.
- override protected def evalConf: Map[String, String] = {
- if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
- super.evalConf ++ Map(
- "input_type" -> schema.json
- )
- } else {
- super.evalConf
- }
- }
+ override protected def evalConf: Map[String, String] =
+ ArrowPythonRunner.elementwiseEvalConf(
+ super.evalConf, evalType, schema, elementwiseNestingDepths)
override protected def writeUDF(dataOut: DataOutputStream): Unit = {
PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
index cab6c9895ca9..dfaa4732ca9d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
@@ -17,6 +17,7 @@
package org.apache.spark.sql.execution.python
+import org.apache.spark.api.python.PythonEvalType
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
@@ -58,13 +59,13 @@ import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
* `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
* `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
* per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ * A UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`, is handled too: the
+ * whole nest is rewritten root-first (see `apply` / `rewriteNest`), lifting the UDF out one lambda
+ * level at a time so it ends up applied to the fully flattened leaves.
*
* `CheckAnalysis` still rejects what this rule does not handle:
- * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
- * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
- * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
- * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
- * outputs, not array elements.
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential (the UDF sees earlier steps' outputs,
+ * not array elements), so it cannot be applied once to the whole array.
*/
object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
@@ -72,20 +73,37 @@ object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
plan
} else {
- // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
- // visits the innermost higher-order function first, and each rewrite lifts all of that
- // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
- // (its argument is the outer lambda's variable, which is not a real column) and is rejected
- // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ // Rewrite each expression tree top-down. The first (outermost) higher-order function that is
+ // a liftable *nest root* - it iterates real columns and its whole, possibly nested, nest is
+ // rewritable - has its entire nest rewritten in one action by `rewriteNest`. Handling the
+ // nest atomically means a UDF in a nested lambda is never momentarily left as a free-variable
+ // `PythonUDF` that downstream could not evaluate (SPARK-48706): the rule either rewrites a
+ // nest completely or leaves it untouched for `CheckAnalysis` to reject. The gate mirrors
+ // `CheckAnalysis` exactly, so analysis accepts precisely what this rewrites.
plan.transformUpWithPruning(
_.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
case p =>
- p.transformExpressionsUpWithPruning(
- _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ p.transformExpressionsDownWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case hof: HigherOrderFunction
+ if hof.functions.exists(_.exists(_.isInstanceOf[PythonUDF])) &&
+ PythonUDF.canRewritePythonUDFInLambda(hof) =>
+ rewriteNest(hof)
+ }
}
}
}
+ /**
+ * Rewrites a whole nest of higher-order functions rooted at `root`, already validated liftable by
+ * [[PythonUDF.canRewritePythonUDFInLambda]]. `transformUp` visits the innermost function first,
+ * so each inner lambda's UDFs are lifted onto that lambda's (enclosing-variable) argument -
+ * becoming element-wise UDFs in an argument position - and then the enclosing function re-lifts
+ * them onto its own array one level deeper (see `buildCarrier` / `overArray`). A single bottom-up
+ * pass therefore lifts every UDF out of every lambda in the nest, innermost first.
+ */
+ private def rewriteNest(root: HigherOrderFunction): Expression = root.transformUp(rewriteOne)
+
/**
* Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
* `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
@@ -106,14 +124,20 @@ object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
}
/**
- * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
- * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
- * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
- * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
- * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
- * single call takes both elements, needs its own path.
+ * Rewrites one higher-order function whose own lambda holds a rewritable Python UDF. The generic
+ * path reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and
+ * rebuilds with `withNewChildren` (children are `arguments` then `functions`), naming a concrete
+ * class only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the
+ * result-type traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort`
+ * comparator, whose single call takes both elements, needs its own path.
+ *
+ * Applied by `rewriteNest` to every function in a validated nest, innermost first. Unlike the
+ * nest-root gate in `apply`, `liftableHof` here does not re-check free lambda variables: within
+ * an already-validated nest an inner function iterating an enclosing variable is expected, and
+ * `rewriteNest` guarantees the enclosing function re-lifts whatever an inner step leaves in an
+ * argument position.
*/
- private val rewrite: PartialFunction[Expression, Expression] = {
+ private val rewriteOne: PartialFunction[Expression, Expression] = {
case sort @ ArraySort(_, function, _)
if liftableHof(sort) && comparatorTakesBothElements(function) =>
rewritePairwiseComparator(sort)
@@ -286,19 +310,18 @@ object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
}
/**
- * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
- * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ * True if `hof`'s single lambda holds a UDF to lift at *this* level - either directly in the body
+ * or in a nested function's argument (which `hasDirectRewritableUDF` reaches, but a nested
+ * function's own lambda is not this level's concern; `rewriteNest` handles that level itself).
+ *
+ * Does not re-check free lambda variables: `apply` enters `rewriteNest` only on a validated nest
+ * root, so within the nest an inner function iterating an enclosing variable is expected, and the
+ * enclosing function is guaranteed to re-lift whatever this level leaves in an argument position.
*/
private def liftableHof(hof: HigherOrderFunction): Boolean =
hof.functions.length == 1 && (hof.functions.head match {
case LambdaFunction(body, args, _) =>
- hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable]) &&
- // Defense in depth: `CheckAnalysis` already rejects a `hof` nested in an enclosing
- // lambda (its argument is a free lambda variable, not a real column), but re-check here
- // so the rule stays correct even for a plan that did not go through analysis - otherwise
- // it could mis-rewrite `transform(arr, i -> transform(i, x -> f(x)))` and resurrect
- // SPARK-48706 (a raw `PythonUDF` left inside a generated lambda).
- !PythonUDF.hasFreeLambdaVariable(hof)
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
case _ => false
})
@@ -388,23 +411,37 @@ object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
val udfFieldByKey = scala.collection.mutable.LinkedHashMap.empty[Expression, Int]
liftableUDFs.foreach { udf =>
// `overArray` turns each argument into an `array<T>` aligned with the iterated array, so the
- // worker flattens every one exactly once (no per-argument shape to track).
- val arrayArgs = udf.children.map { child =>
- overArray(child, alignedArguments.head, arrayOfVar, lambdaExprIds, arrayResults)
+ // worker flattens every one exactly once (no per-argument shape to track). A keyword argument
+ // keeps its `NamedArgumentExpression` wrapper as a direct child of the lifted UDF - only its
+ // value is lifted - so the runner still derives the kwargs mapping from the direct children.
+ val arrayArgs = udf.children.map {
+ case NamedArgumentExpression(key, value) =>
+ NamedArgumentExpression(
+ key, overArray(value, alignedArguments.head, arrayOfVar, lambdaExprIds, arrayResults))
+ case child =>
+ overArray(child, alignedArguments.head, arrayOfVar, lambdaExprIds, arrayResults)
}
+ // Each lift wraps the arguments in exactly one more `array` level. Lifting a base UDF gives
+ // depth 1; re-lifting an already-lifted element-wise UDF (a UDF from a *nested* lambda,
+ // lifted once onto the inner variable and now again onto the enclosing array) adds one more
+ // level, so the worker flattens `depth` levels down to the scalar element and re-nests them.
+ val newDepth =
+ if (PythonEvalType.isElementwiseUDF(udf.evalType)) udf.elementwiseNestingDepth + 1 else 1
val lifted = PythonUDF(
udf.name,
udf.func,
// The wrapper returns one element per input element, i.e. one array level on top of the
- // user function's scalar return. Elements may be null (the UDF can return null), hence
+ // UDF's previous return type. Elements may be null (the UDF can return null), hence
// containsNull = true.
ArrayType(udf.dataType, containsNull = true),
arrayArgs,
// Each rewritable flavor lifts to its own element-wise eval type so the worker keeps that
// flavor's batching contract (pickle row-at-a-time, pandas Series, Arrow Array, or an
- // iterator of batches). See `PythonUDF.liftedElementwiseEvalType`.
+ // iterator of batches); an already-lifted type maps to itself. See
+ // `PythonUDF.liftedElementwiseEvalType`.
PythonUDF.liftedElementwiseEvalType(udf.evalType),
- udf.udfDeterministic)
+ udf.udfDeterministic,
+ elementwiseNestingDepth = newDepth)
val signature: Expression = if (udf.udfDeterministic) lifted.canonicalized else lifted
val ordinal = ordinalBySignature.getOrElseUpdate(signature, {
val o = distinctLifted.length
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
index 6ec31e67ca69..09d982026006 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
@@ -98,12 +98,9 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession
// The fold is sequential, so no array can precompute the values the UDF sees. A UDF in `merge`
// or in `finish` therefore has no rewrite and analysis must fail.
val cases = Seq(
- "merge on element" ->
+ "merge" ->
org.apache.spark.sql.functions.aggregate(
col("values"), lit(false), (acc, x) => acc || pythonUDF(x)),
- "merge on accumulator" ->
- org.apache.spark.sql.functions.aggregate(
- col("values"), lit(false), (acc, x) => pythonUDF(acc) || x.cast("boolean")),
"finish" ->
org.apache.spark.sql.functions.aggregate(
col("values"), lit(0), (acc, x) => acc + x, acc => pythonUDF(acc)))
@@ -167,15 +164,36 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession
assert(mapType.keyType == StringType)
}
- test("a UDF inside a nested higher-order function's lambda is rejected") {
- // The inner array is the outer lambda's variable, not a real column, so the UDF cannot be
- // lifted onto an array outside the lambda. This must fail analysis rather than be rewritten.
+ test("a UDF inside a nested higher-order function's lambda is lifted, deepening the nesting") {
+ // `transform(matrix, row -> transform(row, x -> f(x)))`: the UDF is lifted onto the inner
+ // variable and then re-lifted onto the real `array<array<int>>` column, so it becomes a
+ // depth-2 element-wise UDF (flattening two array levels). Nothing is left inside a lambda.
val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
- val e = intercept[AnalysisException] {
- df.select(transform(col("values"), inner =>
- transform(inner, x => pythonUDF(x))).as("r")).collect()
- }
- assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ .select(transform(col("values"), inner =>
+ transform(inner, x => pythonUDF(x))).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ assert(lifted.head.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)
+ // Re-lifted once per enclosing lambda: depth 2, over an array<array<...>> argument.
+ assert(lifted.head.elementwiseNestingDepth == 2)
+ assert(lifted.head.dataType ==
+ ArrayType(ArrayType(pythonUDF.dataType, containsNull = true), containsNull = true))
+ }
+
+ test("a UDF in a nested lambda that captures an enclosing lambda variable is lifted") {
+ // `transform(m, row -> transform(row, x -> f(x, size(row))))`: the UDF reads the inner element
+ // and the *enclosing* variable `row`. The captured value is repeated into an aligned array, so
+ // the UDF still lifts to a depth-2 element-wise UDF over the real column and nothing is left in
+ // a lambda. Relies on the fixed `HigherOrderFunction.canonicalized`, which no longer leaks the
+ // captured variable into the lifted UDF's references.
+ val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
+ .select(transform(col("values"), row =>
+ transform(row, x => pythonUDF2(x, org.apache.spark.sql.functions.size(row)))).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ assert(lifted.head.elementwiseNestingDepth == 2)
}
test("a UDF on the outer element of a nested array is lifted") {
@@ -295,11 +313,11 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession
assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
}
- test("the rewritable predicate excludes zero-argument, named-argument and UDT UDFs") {
- // These shapes cannot be preserved by the lift (a zero-arg call has no array to carry the
- // iterated shape and would crash the worker; a `NamedArgumentExpression` child would be buried
- // inside the generated transform, losing kwargs; a UDT would hit the Arrow path with no
- // fallback), so the shared predicate must reject them, keeping the previous analysis error.
+ test("the rewritable predicate: eval types, zero-argument / iterator-kwarg / UDT UDFs") {
+ // A zero-arg call has no array to carry the iterated shape and would crash the worker; an
+ // iterator UDF takes no kwargs; a UDT would hit the Arrow path with no fallback - so the shared
+ // predicate rejects those, keeping the previous analysis error. A named argument on a
+ // non-iterator flavor is accepted (the lift keeps the NamedArgumentExpression).
val plain = PythonUDF("f", null, IntegerType, Seq(Literal(1)),
PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
assert(PythonUDF.isElementwiseRewritableUDF(plain))
@@ -321,11 +339,31 @@ class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession
assert(PythonUDF.liftedElementwiseEvalType(PythonEvalType.SQL_BATCHED_UDF) ==
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)
+ // An already-lifted element-wise UDF is rewritable again (nested lambdas re-lift it), and its
+ // lifted eval type is itself - only the nesting depth changes.
+ Seq(
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF).foreach { ew =>
+ assert(PythonUDF.isElementwiseRewritableUDF(plain.copy(evalType = ew)))
+ assert(PythonUDF.liftedElementwiseEvalType(ew) == ew)
+ }
+
val zeroArg = plain.copy(children = Seq.empty)
assert(!PythonUDF.isElementwiseRewritableUDF(zeroArg))
+ // A named argument is rewritable on the non-iterator flavors (the lift keeps the
+ // NamedArgumentExpression as a direct child), but not on an iterator UDF (no kwargs there).
val named = plain.copy(children = Seq(NamedArgumentExpression("k", Literal(1))))
- assert(!PythonUDF.isElementwiseRewritableUDF(named))
+ assert(PythonUDF.isElementwiseRewritableUDF(named))
+ assert(PythonUDF.isElementwiseRewritableUDF(
+ named.copy(evalType = PythonEvalType.SQL_SCALAR_PANDAS_UDF)))
+ assert(!PythonUDF.isElementwiseRewritableUDF(
+ named.copy(evalType = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF)))
+ assert(!PythonUDF.isElementwiseRewritableUDF(
+ named.copy(evalType = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF)))
val udtReturn = plain.copy(dataType = new ExamplePointUDT)
assert(!PythonUDF.isElementwiseRewritableUDF(udtReturn))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org