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 ea911d076840 [SPARK-27052][SQL][PYTHON] Support Python UDFs inside
higher-order function lambdas
ea911d076840 is described below
commit ea911d076840c8577381da422e5801df44c467fe
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Mon Aug 10 20:10:39 2026 +0900
[SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function
lambdas
### 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 fails today.** A `PythonUDF` runs in a separate operator that
`ExtractPythonUDFs` pulls
out of the enclosing one. A lambda's `NamedLambdaVariable`s only exist
while the function iterates,
so that operator cannot see them: the UDF can neither stay in the lambda
nor be lifted out. So
`CheckAnalysis` rejected it.
**How it works.** A new optimizer rule `ExtractPythonUDFFromLambda` applies
the UDF once to the
*whole array*, outside every lambda, and has the lambda read the result
positionally:
```
-- before (rejected)
transform(values, x -> plus_one(x))
-- after (the PythonUDF is outside every lambda)
transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s ->
s.u0)
```
`plus_one_over_array` is the same function re-typed `array<T> => array<R>`
and run with a new eval
type `SQL_ARROW_ELEMENTWISE_UDF`. Since the JVM cannot rewrap a pickled
function, the
array-at-a-time behaviour lives in the Python worker: it flattens each list
column once, calls the
function over all elements of the batch, and re-nests by the input's
offsets. One row in, one row
out; one Python round trip per batch. Every lifted argument is a
single-level `array<T>` aligned
with the iterated array (an outer column or constant is repeated into one),
so the worker flattens
all arguments the same way. The rule runs before `ExtractPythonUDFs` and is
non-excludable.
**What works.** All ten single-lambda functions: `transform`, `filter`,
`exists`, `forall`,
`zip_with`, `array_sort`, and the four map functions (desugared to
`map_keys`/`map_values` arrays,
rebuilt with `map_from_arrays`). Around the UDF, ordinary JVM work all
composes: arithmetic,
`when`, casts, the element index, multiple UDFs, and nested calls `f(g(x))`.
```python
df.select(F.transform("vals", lambda x: F.when(is_pos(x), x *
2).otherwise(0)))
df.select(F.filter("vals", lambda x: is_even(x)))
df.select(F.transform_values("m", lambda k, v: to_upper(v))) # map family
```
`array_sort` has two forms: a UDF used as a per-element sort key is
precomputed per element; a
pairwise comparator `(a, b) -> udf(a, b)`, which has no per-element key, is
precomputed over the
cross product of pairs (O(n^2) calls, so the key form is preferred where
possible).
To drive the rewrite generically, two marker traits are added to the
Catalyst higher-order
functions: `ResultTypeFromArgument` (result type is the input's - `filter`,
`map_filter`,
`array_sort`) and `ResultTypeFromFunction` (result type follows the lambda
- `transform`,
`transform_keys`/`transform_values`, `zip_with`, `map_zip_with`,
`aggregate`, `exists`/`forall`).
The first also hoists the shared `dataType = argument.dataType`. The rule
reads these instead of
special-casing each class, so it stays generic and a new function of a
familiar shape needs no
change.
**What does not work** (still rejected with the same error):
```python
# nested lambda: the inner array is the outer lambda's variable, not a real
column
df.select(F.transform("arr", lambda i: F.transform(i, lambda x: f(x))))
# rejected
# aggregate/reduce: the fold is sequential, so the UDF sees earlier steps,
not array elements
df.select(F.aggregate("vals", F.lit(0), lambda acc, x: acc + plus_one(x)))
# rejected
# pandas/arrow UDF: TBD after this PR
```
A UDF in a nested function's *argument* is fine, since it does not read the
inner lambda's variable:
```python
df.select(F.transform("arr", lambda x: F.transform(udf(x), lambda y: y)))
# works
```
`CheckAnalysis` and the rule share one predicate, so analysis accepts
exactly what the rule
rewrites. Turn the rewrite off with
`spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled=false`.
This builds on design notes from the
[elementwise-udf](https://github.com/HyukjinKwon/elementwise-udf) prototype
(Apache-2.0).
### Why are the changes needed?
`transform(col, lambda x: my_udf(x))` is a natural thing to write and a
long-standing gap
(SPARK-27052). The alternative, `explode` + regroup, is far slower and OOMs
on long arrays.
### 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]
```
The shapes that cannot be rewritten keep failing with the same error
condition. 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` (34 tests): all ten
functions, composition/index/multiple/nested calls, null and empty
inputs, ragged `zip_with`,
long arrays, integration with joins/caching/`groupBy`, and the rejected
cases above. Results are
asserted against the equivalent native expression wherever one exists.
- New plan-shape suite `ExtractPythonUDFFromLambdaSuite` (15 tests): no
`PythonUDF` remains in a
lambda, duplicate calls evaluated once, the rule is inert without a UDF
and not excludable, and
nested/aggregate UDFs are rejected.
- `PythonUDFSuite`, the `execution.python` suites and
`DataFrameFunctionsSuite` pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
Closes #57804 from HyukjinKwon/SPARK-27052-python-udf-in-hof.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
.../src/main/resources/error/error-conditions.json | 3 +-
.../org/apache/spark/api/python/PythonRunner.scala | 4 +
dev/sparktestsupport/modules.py | 2 +
.../docs/source/tutorial/sql/arrow_python_udf.rst | 11 +
python/pyspark/sql/_typing.pyi | 1 +
.../test_parity_udf_in_higher_order_function.py | 37 ++
.../sql/tests/test_udf_in_higher_order_function.py | 688 +++++++++++++++++++++
python/pyspark/util.py | 2 +
python/pyspark/worker.py | 133 ++++
.../sql/catalyst/analysis/CheckAnalysis.scala | 14 +-
.../spark/sql/catalyst/expressions/PythonUDF.scala | 104 ++++
.../expressions/higherOrderFunctions.scala | 52 +-
.../org/apache/spark/sql/internal/SQLConf.scala | 20 +
.../spark/sql/execution/SparkOptimizer.scala | 4 +
.../sql/execution/python/ArrowEvalPythonExec.scala | 1 +
.../sql/execution/python/ArrowPythonRunner.scala | 5 +-
.../python/ExtractPythonUDFFromLambda.scala | 619 ++++++++++++++++++
.../sql/execution/python/ExtractPythonUDFs.scala | 10 +-
.../python/BatchEvalPythonExecSuite.scala | 9 +-
.../python/ExtractPythonUDFFromLambdaSuite.scala | 379 ++++++++++++
.../sql/execution/python/PythonUDFSuite.scala | 20 +-
21 files changed, 2092 insertions(+), 26 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index 046302512298..7ba37615c364 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -8727,7 +8727,8 @@
},
"LAMBDA_FUNCTION_WITH_PYTHON_UDF" : {
"message" : [
- "Lambda function with Python UDF <funcName> in a higher order
function."
+ "Cannot evaluate the Python UDF <funcName> inside the lambda of a
higher-order function.",
+ "This placement is not supported. Rewrite the query so the UDF is
applied outside the lambda."
]
},
"LAMBDA_FUNCTION_WITH_SQL_UDF" : {
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 20424e4f37df..19a482850311 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -51,6 +51,9 @@ private[spark] object PythonEvalType {
val SQL_BATCHED_UDF = 100
val SQL_ARROW_BATCHED_UDF = 101
+ // A scalar Python UDF applied element-wise over the elements of an array
column, used to
+ // support Python UDFs inside higher-order function lambdas. See
ExtractPythonUDFFromLambda.
+ val SQL_ARROW_ELEMENTWISE_UDF = 102
val SQL_SCALAR_PANDAS_UDF = 200
val SQL_GROUPED_MAP_PANDAS_UDF = 201
@@ -86,6 +89,7 @@ private[spark] object PythonEvalType {
case NON_UDF => "NON_UDF"
case SQL_BATCHED_UDF => "SQL_BATCHED_UDF"
case SQL_ARROW_BATCHED_UDF => "SQL_ARROW_BATCHED_UDF"
+ case SQL_ARROW_ELEMENTWISE_UDF => "SQL_ARROW_ELEMENTWISE_UDF"
case SQL_SCALAR_PANDAS_UDF => "SQL_SCALAR_PANDAS_UDF"
case SQL_GROUPED_MAP_PANDAS_UDF => "SQL_GROUPED_MAP_PANDAS_UDF"
case SQL_GROUPED_AGG_PANDAS_UDF => "SQL_GROUPED_AGG_PANDAS_UDF"
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index b0c4d6e85444..2faeea33b6f6 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -642,6 +642,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.test_geometrytype",
"pyspark.sql.tests.test_udf",
"pyspark.sql.tests.test_udf_combinations",
+ "pyspark.sql.tests.test_udf_in_higher_order_function",
"pyspark.sql.tests.test_udf_profiler",
"pyspark.sql.tests.test_udf_transpile_hypothesis",
"pyspark.sql.tests.test_udf_transpile_parity",
@@ -1226,6 +1227,7 @@ pyspark_connect = Module(
"pyspark.sql.tests.connect.test_parity_column",
"pyspark.sql.tests.connect.test_parity_readwriter",
"pyspark.sql.tests.connect.test_parity_udf",
+ "pyspark.sql.tests.connect.test_parity_udf_in_higher_order_function",
"pyspark.sql.tests.connect.test_parity_udf_combinations",
"pyspark.sql.tests.connect.test_parity_udf_profiler",
"pyspark.sql.tests.connect.test_parity_unified_udf",
diff --git a/python/docs/source/tutorial/sql/arrow_python_udf.rst
b/python/docs/source/tutorial/sql/arrow_python_udf.rst
index d1430844a705..0a0509f05003 100644
--- a/python/docs/source/tutorial/sql/arrow_python_udf.rst
+++ b/python/docs/source/tutorial/sql/arrow_python_udf.rst
@@ -479,6 +479,17 @@ 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.
+
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
using Arrow's safe casting, which raises an error on overflow or precision
loss.
diff --git a/python/pyspark/sql/_typing.pyi b/python/pyspark/sql/_typing.pyi
index 94e3ccf77093..95bf27662ff1 100644
--- a/python/pyspark/sql/_typing.pyi
+++ b/python/pyspark/sql/_typing.pyi
@@ -64,6 +64,7 @@ RowLike = TypeVar("RowLike", List[Any], Tuple[Any, ...],
pyspark.sql.types.Row)
SQLBatchedUDFType = Literal[100]
SQLArrowBatchedUDFType = Literal[101]
+SQLArrowElementwiseUDFType = Literal[102]
SQLTableUDFType = Literal[300]
SQLArrowTableUDFType = Literal[301]
SQLArrowUDTFType = Literal[302]
diff --git
a/python/pyspark/sql/tests/connect/test_parity_udf_in_higher_order_function.py
b/python/pyspark/sql/tests/connect/test_parity_udf_in_higher_order_function.py
new file mode 100644
index 000000000000..67228b5c60fb
--- /dev/null
+++
b/python/pyspark/sql/tests/connect/test_parity_udf_in_higher_order_function.py
@@ -0,0 +1,37 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+
+from pyspark.sql.tests.test_udf_in_higher_order_function import (
+ UDFInHigherOrderFunctionTestsMixin,
+)
+from pyspark.testing.connectutils import ReusedConnectTestCase
+
+
+class UDFInHigherOrderFunctionParityTests(
+ UDFInHigherOrderFunctionTestsMixin, ReusedConnectTestCase
+):
+ @unittest.skip("Asserts on the JVM optimized plan via _jdf, unavailable in
Spark Connect.")
+ def test_lambda_without_udf_is_unchanged(self):
+ super().test_lambda_without_udf_is_unchanged()
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
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
new file mode 100644
index 000000000000..c027c4779fea
--- /dev/null
+++ b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py
@@ -0,0 +1,688 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+
+from pyspark.errors import AnalysisException
+from pyspark.sql import functions as sf
+from pyspark.sql.functions import udf
+from pyspark.sql.types import ArrayType, DoubleType, IntegerType, StringType
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import (
+ assertDataFrameEqual,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
[email protected](
+ not have_pandas or not have_pyarrow, pandas_requirement_message or
pyarrow_requirement_message
+)
+class UDFInHigherOrderFunctionTestsMixin:
+ """Tests for scalar Python UDFs used inside higher-order function lambdas
(SPARK-27052).
+
+ ``ExtractPythonUDFFromLambda`` rewrites such a plan so the UDF is applied
to the whole array
+ outside the lambda. Each test asserts the *result*, comparing against the
equivalent native
+ expression wherever one exists, so that a rewrite that runs but computes
the wrong thing
+ fails rather than passing quietly.
+ """
+
+ def test_transform(self):
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],), ([10],)],
"values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one(x)).alias("r")),
+ df.select(sf.transform("values", lambda x: x + 1).alias("r")),
+ )
+
+ def test_transform_null_array_and_null_elements(self):
+ # A null array must stay null, and a null *element* must reach the UDF
as None.
+ df = self.spark.createDataFrame([([1, None, 3],), (None,), ([],)],
"values array<int>")
+ # Null-aware so the UDF itself can observe the null element.
+ f = udf(lambda x: -1 if x is None else x * 2, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: f(x)).alias("r")),
+ [([2, -1, 6],), (None,), ([],)],
+ )
+
+ def test_transform_udf_returning_null(self):
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ f = udf(lambda x: None if x == 2 else x, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: f(x)).alias("r")),
+ [([1, None, 3],)],
+ )
+
+ def test_transform_with_index(self):
+ df = self.spark.createDataFrame([([10, 20, 30],), ([],)], "values
array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ # The index parameter must still work once the element is read from
the carrier struct.
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x, i: plus_one(x) +
i).alias("r")),
+ df.select(sf.transform("values", lambda x, i: (x + 1) +
i).alias("r")),
+ )
+
+ def test_composition_around_udf_result(self):
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ # Arithmetic, `when` and casts around the UDF result are ordinary JVM
work.
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda x: plus_one(x) * 2).alias("mul"),
+ sf.transform(
+ "values", lambda x: sf.when(plus_one(x) > 2,
sf.lit(1)).otherwise(sf.lit(0))
+ ).alias("cond"),
+ sf.transform("values", lambda x:
plus_one(x).cast("string")).alias("cast"),
+ ),
+ df.select(
+ sf.transform("values", lambda x: (x + 1) * 2).alias("mul"),
+ sf.transform(
+ "values", lambda x: sf.when((x + 1) > 2,
sf.lit(1)).otherwise(sf.lit(0))
+ ).alias("cond"),
+ sf.transform("values", lambda x: (x +
1).cast("string")).alias("cast"),
+ ),
+ )
+
+ def test_udf_argument_is_expression_over_element(self):
+ # `udf(x * 2)`: the argument is itself an expression over the element.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one(x *
2)).alias("r")),
+ df.select(sf.transform("values", lambda x: x * 2 + 1).alias("r")),
+ )
+
+ def test_multiple_udfs_in_one_lambda(self):
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ times_ten = udf(lambda x: x * 10, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one(x) +
times_ten(x)).alias("r")),
+ df.select(sf.transform("values", lambda x: (x + 1) + (x *
10)).alias("r")),
+ )
+
+ def test_nested_udfs(self):
+ # `f(g(x))`: both are lifted, and compose as array UDFs outside the
lambda.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ times_ten = udf(lambda x: x * 10, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
times_ten(plus_one(x))).alias("r")),
+ df.select(sf.transform("values", lambda x: (x + 1) *
10).alias("r")),
+ )
+
+ def test_nested_udf_inside_composite_argument(self):
+ # SPARK-27052: `f(g(x) + 1)` / `f(-g(x))`. The inner call is buried
inside a composite
+ # argument, not a direct child. The inner result must be substituted
before lifting `f`,
+ # or a raw `g` over the lambda variable would be left inside a lambda
and mis-extracted
+ # (SPARK-48706). The result must match evaluating the composition
element-wise.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ times_ten = udf(lambda x: x * 10, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: times_ten(plus_one(x) +
1)).alias("r")),
+ df.select(sf.transform("values", lambda x: ((x + 1) + 1) *
10).alias("r")),
+ )
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
times_ten(-plus_one(x))).alias("r")),
+ df.select(sf.transform("values", lambda x: (-(x + 1)) *
10).alias("r")),
+ )
+
+ def test_udf_with_outer_column_argument(self):
+ # A non-element argument must be broadcast to every element of its row.
+ df = self.spark.createDataFrame([([1, 2], 100), ([3], 200)], "values
array<int>, base int")
+ add = udf(lambda x, b: x + b, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: add(x,
sf.col("base"))).alias("r")),
+ [([101, 102],), ([203],)],
+ )
+
+ def test_udf_with_constant_argument_only(self):
+ # SPARK-27052: `transform(arr, x -> udf(lit(10)))` does not read the
element, but the UDF
+ # must still take the lambda's call domain: once per element, and zero
times for an empty
+ # or null array (where the lambda never runs), rather than once per
row.
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],), (None,)],
"values array<int>")
+ const = udf(lambda v: v * 2, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
const(sf.lit(10))).alias("r")),
+ [([20, 20, 20],), ([],), (None,)],
+ )
+
+ def test_filter(self):
+ # `filter`'s result is built from the input elements, not the lambda's
value.
+ df = self.spark.createDataFrame([([1, 2, 3, 4],), ([],), (None,)],
"values array<int>")
+ is_even = udf(lambda x: x % 2 == 0, "boolean")
+
+ assertDataFrameEqual(
+ df.select(sf.filter("values", lambda x: is_even(x)).alias("r")),
+ df.select(sf.filter("values", lambda x: (x % 2) == 0).alias("r")),
+ )
+
+ def test_exists_and_forall(self):
+ df = self.spark.createDataFrame([([1, 2, 3],), ([2, 4],), ([],)],
"values array<int>")
+ is_even = udf(lambda x: x % 2 == 0, "boolean")
+
+ assertDataFrameEqual(
+ df.select(
+ sf.exists("values", lambda x: is_even(x)).alias("e"),
+ sf.forall("values", lambda x: is_even(x)).alias("f"),
+ ),
+ df.select(
+ sf.exists("values", lambda x: (x % 2) == 0).alias("e"),
+ sf.forall("values", lambda x: (x % 2) == 0).alias("f"),
+ ),
+ )
+
+ def test_zip_with(self):
+ # Two arrays at once. `arrays_zip` pads the shorter side with nulls,
which is what
+ # `zip_with` does itself, so differing lengths must agree with the
native version.
+ df = self.spark.createDataFrame(
+ [([1, 2], [10, 20]), ([1, 2, 3], [10]), ([], []), (None, [1]),
([1], None)],
+ "l array<int>, r array<int>",
+ )
+ add = udf(lambda a, b: (0 if a is None else a) + (0 if b is None else
b), IntegerType())
+
+ # Compare against the equivalent native expression with the same null
handling.
+ assertDataFrameEqual(
+ df.select(sf.zip_with("l", "r", lambda a, b: add(a,
b)).alias("r")),
+ df.select(
+ sf.zip_with(
+ "l",
+ "r",
+ lambda a, b: sf.coalesce(a, sf.lit(0)) + sf.coalesce(b,
sf.lit(0)),
+ ).alias("r")
+ ),
+ )
+
+ def test_zip_with_udf_on_one_side_only(self):
+ df = self.spark.createDataFrame([([1, 2], [10, 20])], "l array<int>, r
array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.zip_with("l", "r", lambda a, b: plus_one(a) +
b).alias("r")),
+ [([12, 23],)],
+ )
+
+ def test_array_sort_with_udf_key(self):
+ # When the UDF applies per element, it is precomputed as a sort key
that the JVM comparator
+ # compares (a UDF taking both elements is instead precomputed over the
pairs; see the
+ # pairwise test). This must actually reorder, not be a no-op.
+ df = self.spark.createDataFrame([([3, 1, 2],), ([],), (None,)],
"values array<int>")
+ negate = udf(lambda x: -x, IntegerType())
+
+ # Sorting by -x gives descending order.
+ assertDataFrameEqual(
+ df.select(
+ sf.array_sort(
+ "values",
+ lambda a, b: sf.when(negate(a) < negate(b), sf.lit(-1))
+ .when(negate(a) > negate(b), sf.lit(1))
+ .otherwise(sf.lit(0)),
+ ).alias("r")
+ ),
+ [([3, 2, 1],), ([],), (None,)],
+ )
+
+ def test_array_sort_pairwise_comparator(self):
+ # One UDF call receiving both elements has no per-element key, so the
UDF is precomputed
+ # over every ordered pair and the comparator indexes that matrix.
Assert an actual
+ # reordering, not merely that the query runs.
+ df = self.spark.createDataFrame(
+ [([3, 1, 2],), ([],), (None,), ([5],), ([2, 2, 1],)], "values
array<int>"
+ )
+ cmp_udf = udf(lambda a, b: (a > b) - (a < b), IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.array_sort("values", lambda a, b: cmp_udf(a,
b)).alias("r")),
+ [([1, 2, 3],), ([],), (None,), ([5],), ([1, 2, 2],)],
+ )
+
+ def test_array_sort_pairwise_comparator_descending(self):
+ # Reversing the comparator must reverse the order, which a no-op
rewrite would not do.
+ df = self.spark.createDataFrame([([3, 1, 2],)], "values array<int>")
+ cmp_desc = udf(lambda a, b: (b > a) - (b < a), IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.array_sort("values", lambda a, b: cmp_desc(a,
b)).alias("r")),
+ [([3, 2, 1],)],
+ )
+
+ def test_transform_keys_and_values(self):
+ df = self.spark.createDataFrame(
+ [
+ ({"a": 1, "b": 2},),
+ ],
+ "m map<string,int>",
+ )
+ upper = udf(lambda s: s.upper(), StringType())
+ plus_one = udf(lambda v: v + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform_keys("m", lambda k, v:
upper(k)).alias("r")),
+ [({"A": 1, "B": 2},)],
+ )
+ assertDataFrameEqual(
+ df.select(sf.transform_values("m", lambda k, v:
plus_one(v)).alias("r")),
+ [({"a": 2, "b": 3},)],
+ )
+ # The lambda may read both the key and the value.
+ assertDataFrameEqual(
+ df.select(sf.transform_values("m", lambda k, v: plus_one(v) +
sf.length(k)).alias("r")),
+ [({"a": 3, "b": 4},)],
+ )
+
+ def test_map_filter(self):
+ df = self.spark.createDataFrame([({"a": 1, "b": 2, "c": 3},)], "m
map<string,int>")
+ is_odd = udf(lambda v: v % 2 == 1, "boolean")
+
+ assertDataFrameEqual(
+ df.select(sf.map_filter("m", lambda k, v: is_odd(v)).alias("r")),
+ [({"a": 1, "c": 3},)],
+ )
+
+ def test_map_zip_with(self):
+ # The visited key set is the union of both maps' keys; a key missing
from one side gives
+ # null on that side, matching map_zip_with's own semantics.
+ df = self.spark.createDataFrame(
+ [
+ (
+ {"a": 1, "b": 2},
+ {"b": 20, "c": 30},
+ )
+ ],
+ "l map<string,int>, r map<string,int>",
+ )
+ combine = udf(
+ lambda a, b: (0 if a is None else a) * 100 + (0 if b is None else
b), IntegerType()
+ )
+
+ assertDataFrameEqual(
+ df.select(sf.map_zip_with("l", "r", lambda k, v1, v2: combine(v1,
v2)).alias("r")),
+ [({"a": 100, "b": 220, "c": 30},)],
+ )
+
+ def test_transform_values_result_type_equals_key_type(self):
+ # SPARK-27052: transform_values on map<string,string> whose lambda
also returns string.
+ # The rewrite must replace the values, not the keys (dispatch is by
function, not type).
+ df = self.spark.createDataFrame([({"a": "x", "b": "y"},)], "m
map<string,string>")
+ tag = udf(lambda v: v + "!", StringType())
+ assertDataFrameEqual(
+ df.select(sf.transform_values("m", lambda k, v:
tag(v)).alias("r")),
+ [({"a": "x!", "b": "y!"},)],
+ )
+
+ def test_nondeterministic_udf_calls_are_distinct(self):
+ # SPARK-27052: two calls to a nondeterministic UDF in one lambda must
stay distinct, not be
+ # collapsed into one shared value. `rand_add` returns x plus a
per-call random draw, so
+ # f(x) + f(x) equals 2x only if the two calls were (wrongly)
deduplicated.
+ import random
+
+ df = self.spark.createDataFrame([([10, 20, 30],)], "values array<int>")
+ rand_add = udf(
+ lambda x: x + random.randint(1, 1_000_000), IntegerType()
+ ).asNondeterministic()
+ row = df.select(
+ sf.transform("values", lambda x: rand_add(x) -
rand_add(x)).alias("r")
+ ).collect()[0]
+ # If the two calls were deduplicated, every element would be exactly 0.
+ self.assertTrue(any(v != 0 for v in row["r"]), row["r"])
+
+ def test_fused_transforms_with_different_lengths(self):
+ # SPARK-27052: two transforms over arrays of different per-row lengths
and null layouts get
+ # fused by ExtractPythonUDFs into one element-wise batch. Each UDF
must be re-nested by its
+ # own array's shape, not a shared one.
+ df = self.spark.createDataFrame(
+ [([1, 2, 3], [10]), ([], None), (None, [7, 8])],
+ "a array<int>, b array<int>",
+ )
+ f = udf(lambda v: v + 1, IntegerType())
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("a", lambda x: f(x)).alias("ra"),
+ sf.transform("b", lambda x: f(x)).alias("rb"),
+ ),
+ [([2, 3, 4], [11]), ([], None), (None, [8, 9])],
+ )
+
+ def test_element_and_return_types(self):
+ df = self.spark.createDataFrame([(["a", "bb"],)], "values
array<string>")
+ upper_len = udf(lambda s: len(s), IntegerType())
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
upper_len(x)).alias("r")),
+ [([1, 2],)],
+ )
+
+ df2 = self.spark.createDataFrame([([1.5, 2.5],)], "values
array<double>")
+ half = udf(lambda v: v / 2, DoubleType())
+ assertDataFrameEqual(
+ df2.select(sf.transform("values", lambda x: half(x)).alias("r")),
+ [([0.75, 1.25],)],
+ )
+
+ # A UDF returning a non-atomic type, so the re-nesting handles nested
lists.
+ df3 = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ repeat = udf(lambda v: [v, v], ArrayType(IntegerType()))
+ assertDataFrameEqual(
+ df3.select(sf.transform("values", lambda x: repeat(x)).alias("r")),
+ [([[1, 1], [2, 2]],)],
+ )
+
+ def test_long_arrays_and_many_rows(self):
+ # Exercises batching: the wrapper evaluates all elements of a batch in
one pass.
+ df = self.spark.range(0, 200).select(
+ sf.transform(
+ sf.sequence(sf.lit(1), sf.lit(20)), lambda x: (x +
sf.col("id")).cast("int")
+ ).alias("values")
+ )
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one(x)).alias("r")),
+ df.select(sf.transform("values", lambda x: x + 1).alias("r")),
+ )
+
+ def test_all_null_rows(self):
+ df = self.spark.createDataFrame([(None,), (None,)], "values
array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one(x)).alias("r")),
+ [(None,), (None,)],
+ )
+
+ def test_empty_dataframe(self):
+ df = self.spark.createDataFrame([], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ self.assertEqual(
+ 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>>")
+ 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))
+
+ 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
+ # column, so this is rewritable even though a higher-order function is
also present.
+ df = self.spark.createDataFrame([([[1, 2], [3]],)], "values
array<array<int>>")
+ total = udf(lambda a: sum(a), IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda inner:
total(inner)).alias("r")),
+ [([3, 3],)],
+ )
+
+ def test_rewritable_higher_order_function_inside_outer_lambda(self):
+ # A rewritable inner HOF over a *real* column, sitting inside an outer
HOF's lambda:
+ # `transform(arr2, i -> array_max(transform(arr, x -> f(x))) + i)`.
Analysis accepts it
+ # (the inner `transform` iterates the real column `arr`, not the outer
lambda variable),
+ # the inner UDF is lifted out of the inner lambda, and the resulting
element-wise UDF stays
+ # inside the outer lambda for `ExtractPythonUDFs` to extract per row.
The result must match
+ # the native computation for a deterministic UDF.
+ df = self.spark.createDataFrame([([1, 2, 3], [10, 20])], "arr
array<int>, arr2 array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(
+ sf.transform(
+ "arr2",
+ lambda i: sf.array_max(sf.transform("arr", lambda x:
plus_one(x))) + i,
+ ).alias("r")
+ ),
+ # array_max([2, 3, 4]) = 4; 4 + 10 = 14, 4 + 20 = 24
+ [([14, 24],)],
+ )
+
+ def test_sql_string_syntax(self):
+ # The rewrite is on the analyzed plan, so it must fire for the SQL
string syntax too, not
+ # only the DataFrame API.
+ self.spark.udf.register("py_plus_one", udf(lambda x: x + 1,
IntegerType()))
+ with self.temp_view("t"):
+ self.spark.createDataFrame([([1, 2, 3],)], "values
array<int>").createOrReplaceTempView(
+ "t"
+ )
+ assertDataFrameEqual(
+ self.spark.sql("SELECT transform(values, x -> py_plus_one(x))
AS r FROM t"),
+ [([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>")
+ 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))
+
+ 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
+ # cannot express it; it must keep failing analysis rather than crash
the worker at runtime.
+ df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ const = udf(lambda: 7, IntegerType())
+
+ with self.assertRaises(AnalysisException) as ctx:
+ df.select(sf.transform("values", lambda x: const())).collect()
+ self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+
+ def test_nondeterministic_iterated_argument_still_fails(self):
+ # The rewrite references the iterated argument several times; a
nondeterministic one (e.g.
+ # shuffle) would be evaluated independently per reference and misalign
the results, so it
+ # must fail analysis instead of being rewritten.
+ df = self.spark.createDataFrame([([1, 2, 3, 4],)], "values array<int>")
+ is_even = udf(lambda x: x % 2 == 0, "boolean")
+
+ with self.assertRaises(AnalysisException) as ctx:
+ df.select(sf.filter(sf.shuffle("values"), lambda x:
is_even(x))).collect()
+ self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+
+ def test_decimal_timestamp_and_struct_element_types(self):
+ # Arrow conversion edge cases beyond int/double/string: decimal,
timestamp, and a struct
+ # return type. Each must round-trip through the element-wise wrapper
correctly.
+ import datetime
+ from decimal import Decimal
+ from pyspark.sql.types import (
+ DecimalType,
+ StructField,
+ StructType,
+ TimestampType,
+ )
+
+ dec_df = self.spark.createDataFrame(
+ [([Decimal("1.50"), Decimal("2.25")],)], "values
array<decimal(5,2)>"
+ )
+ add_half = udf(lambda v: v + Decimal("0.50"), DecimalType(5, 2))
+ assertDataFrameEqual(
+ dec_df.select(sf.transform("values", lambda x:
add_half(x)).alias("r")),
+ [([Decimal("2.00"), Decimal("2.75")],)],
+ )
+
+ ts_df = self.spark.createDataFrame(
+ [([datetime.datetime(2020, 1, 1, 0, 0, 0)],)], "values
array<timestamp>"
+ )
+ add_day = udf(lambda t: t + datetime.timedelta(days=1),
TimestampType())
+ assertDataFrameEqual(
+ ts_df.select(sf.transform("values", lambda x:
add_day(x)).alias("r")),
+ [([datetime.datetime(2020, 1, 2, 0, 0, 0)],)],
+ )
+
+ struct_type = StructType([StructField("a", IntegerType()),
StructField("b", IntegerType())])
+ int_df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ to_struct = udf(lambda v: (v, v * 10), struct_type)
+ assertDataFrameEqual(
+ int_df.select(sf.transform("values", lambda x:
to_struct(x)).alias("r")),
+ [([(1, 10), (2, 20)],)],
+ )
+
+ def test_integration_with_joins_grouping_and_caching(self):
+ left = self.spark.createDataFrame(
+ [
+ (1, [1, 2]),
+ (
+ 2,
+ [3],
+ ),
+ ],
+ "k int, values array<int>",
+ )
+ right = self.spark.createDataFrame([(1,), (2,)], "k int")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ joined = left.join(right, "k").select(
+ "k", sf.transform("values", lambda x: plus_one(x)).alias("r")
+ )
+ assertDataFrameEqual(joined, [(1, [2, 3]), (2, [4])])
+
+ cached = left.select(sf.transform("values", lambda x:
plus_one(x)).alias("r")).cache()
+ try:
+ assertDataFrameEqual(cached, [([2, 3],), ([4],)])
+ finally:
+ cached.unpersist()
+
+ grouped = left.groupBy().agg(
+ sf.sum(
+ sf.aggregate(
+ sf.transform("values", lambda x: plus_one(x)), sf.lit(0),
lambda a, x: a + x
+ )
+ ).alias("s")
+ )
+ # (1+1)+(2+1) + (3+1) = 9
+ assertDataFrameEqual(grouped, [(9,)])
+
+ def test_mixed_with_plain_python_udf(self):
+ df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ size_udf = udf(lambda a: len(a), IntegerType())
+
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda x: plus_one(x)).alias("r"),
+ size_udf("values").alias("n"),
+ ),
+ [([2, 3], 2)],
+ )
+
+ def test_lambda_without_udf_is_unchanged(self):
+ # The rewrite must be inert for plans that contain no Python UDF in a
lambda.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ native = df.select(sf.transform("values", lambda x: x + 1).alias("r"))
+ self.assertNotIn("pythonUDF",
native._jdf.queryExecution().optimizedPlan().toString())
+ assertDataFrameEqual(native, [([2, 3, 4],)])
+
+ def test_disabled_by_conf(self):
+ df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ with
self.sql_conf({"spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled":
False}):
+ with self.assertRaises(AnalysisException) as ctx:
+ 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.
+ 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)),
+ ]
+ for agg in aggregates:
+ with self.assertRaises(AnalysisException) as ctx:
+ df.select(agg).collect()
+ self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF",
str(ctx.exception))
+
+ def test_pandas_udf_in_lambda_still_fails(self):
+ # A vectorized (pandas) UDF in a lambda is not supported.
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+
+ df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+
+ @pandas_udf(IntegerType())
+ def plus_one_pandas(s: pd.Series) -> pd.Series:
+ return s + 1
+
+ with self.assertRaises(AnalysisException) as ctx:
+ df.select(sf.transform("values", lambda x:
plus_one_pandas(x))).collect()
+ self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+
+ def test_non_arrow_udf_is_also_supported(self):
+ # A UDF created with useArrow=False is still rewritable; the generated
array wrapper
+ # is Arrow-based regardless of how the user's UDF was declared.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType(), useArrow=False)
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one(x)).alias("r")),
+ [([2, 3, 4],)],
+ )
+
+ def test_string_return_type(self):
+ df = self.spark.createDataFrame([([1, 2],), (None,)], "values
array<int>")
+ to_str = udf(lambda x: f"v{x}", StringType())
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: to_str(x)).alias("r")),
+ [(["v1", "v2"],), (None,)],
+ )
+
+
+class UDFInHigherOrderFunctionTests(UDFInHigherOrderFunctionTestsMixin,
ReusedSQLTestCase):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 858dc1013e94..640728725ee3 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -90,6 +90,7 @@ if typing.TYPE_CHECKING:
)
from pyspark.sql._typing import (
SQLArrowBatchedUDFType,
+ SQLArrowElementwiseUDFType,
SQLArrowTableUDFType,
SQLBatchedUDFType,
SQLTableUDFType,
@@ -662,6 +663,7 @@ class PythonEvalType:
SQL_BATCHED_UDF: "SQLBatchedUDFType" = 100
SQL_ARROW_BATCHED_UDF: "SQLArrowBatchedUDFType" = 101
+ SQL_ARROW_ELEMENTWISE_UDF: "SQLArrowElementwiseUDFType" = 102
SQL_SCALAR_PANDAS_UDF: "PandasScalarUDFType" = 200
SQL_GROUPED_MAP_PANDAS_UDF: "PandasGroupedMapUDFType" = 201
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 9ac8abbfa485..c579c719c6e2 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -623,6 +623,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
# Scalar, aggregation and window UDFs: (func, args_offsets,
kwargs_offsets, return_type).
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
@@ -1821,6 +1822,7 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
@@ -2989,6 +2991,137 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# profiling is not supported for UDF
return func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF:
+ # This path exchanges data with the JVM over Arrow, so PyArrow is
required. Fail with a
+ # clear message rather than a bare ImportError from `import pyarrow`
below.
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+
+ 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>.
+
+ # UDF preparation
+ udf_infos = []
+ for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in
udfs:
+ 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
+ udf_infos.append(
+ (
+ wrapped_func,
+ args_kwargs_offsets,
+ to_arrow_type(
+ element_return_type,
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
+ ),
+ LocalDataToArrowConversion._create_converter(
+ element_return_type,
+ none_on_identity=True,
+
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ ),
+ )
+ )
+ 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:
+ return [udf_func(*row) for row in rows]
+ from concurrent.futures import ThreadPoolExecutor
+
+ 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.
+ 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])
+ )
+ # Stream the argument tuples rather than materializing a
batch-sized list.
+ rows = zip(*[columns[o] for o in offsets_meta])
+ results = _evaluate_elementwise_udf(wrapped_func, rows)
+ verify_result_row_count(len(results), total_elements)
+
+ # Convert results and re-nest to array<R> using that UDF's
offsets.
+ converted = (
+ [result_conv(r) for r in results] if result_conv is
not None else results
+ )
+ try:
+ flat_arr = pa.array(converted, type=arrow_element_type)
+ # Broader than the SQL_ARROW_BATCHED_UDF path above (which
catches only
+ # ArrowInvalid): the element-wise wrapper commonly returns
list/struct-typed
+ # elements, whose type mismatches surface as
ArrowTypeError, so both are caught
+ # before falling back to an explicit cast.
+ except (pa.lib.ArrowInvalid, pa.lib.ArrowTypeError):
+ flat_arr = pa.array(converted).cast(
+ target_type=arrow_element_type,
safe=runner_conf.safecheck
+ )
+ output_arrays.append(
+ list_cls.from_arrays(offsets_arr, flat_arr,
mask=null_mask)
+ )
+
+ yield pa.RecordBatch.from_arrays(output_arrays, col_names)
+
+ # profiling is not supported for UDF
+ return func, None, ser, ser
+
if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF:
import pandas as pd
import pyarrow as pa
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 f4f9a0296a06..66e2e78aef27 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
@@ -501,10 +501,20 @@ trait CheckAnalysis extends LookupCatalog with
QueryErrorsBase with PlanToString
hof.invalidFormat(checkRes)
}
+ // A Python UDF cannot be evaluated inside a lambda, because it
needs a separate
+ // physical operator that cannot see the lambda's variables. For
some shapes
+ // `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.
case hof: HigherOrderFunction
if hof.resolved && hof.functions
- .exists(_.exists(_.isInstanceOf[PythonUDF])) =>
- val u =
hof.functions.flatMap(_.find(_.isInstanceOf[PythonUDF])).head
+ .exists(_.exists(_.isInstanceOf[PythonUDF])) &&
+ !(conf.pythonUDFInHigherOrderFunctionEnabled &&
+ PythonUDF.canRewritePythonUDFInLambda(hof)) =>
+ // Name the offending UDF: the first one of an unsupported eval
type if any (e.g. a
+ // pandas UDF), otherwise the first Python UDF in the lambdas.
+ val udfs = hof.functions.flatMap(_.collect { case u: PythonUDF =>
u })
+ val u =
udfs.find(!PythonUDF.isElementwiseRewritableUDF(_)).getOrElse(udfs.head)
hof.failAnalysis(
errorClass =
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF",
messageParameters = Map("funcName" -> toSQLExpr(u)))
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 48da53ccf7c2..ec7e7832bb65 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
@@ -36,6 +36,10 @@ object PythonUDF {
private[this] val SCALAR_TYPES = Set(
PythonEvalType.SQL_BATCHED_UDF,
PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ // Element-wise UDFs are row-shaped from the plan's point of view: one
array column in, one
+ // array column out per row. They are extracted by `ExtractPythonUDFs`
like any other scalar
+ // UDF; only the Python worker treats them element-wise.
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
@@ -46,6 +50,106 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] &&
SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order
function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside
the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF
that takes one value per
+ * call. Otherwise-eligible shapes are excluded because the rewrite cannot
preserve them:
+ * - 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 UDF whose argument or return type involves a UDT: the lift forces
the 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.
+ * All keep the previous behavior (an analysis error) rather than being
rewritten.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts
are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
+ case udf: PythonUDF =>
+ (udf.evalType == PythonEvalType.SQL_BATCHED_UDF ||
+ udf.evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) &&
+ udf.children.nonEmpty &&
+ !udf.children.exists(_.isInstanceOf[NamedArgumentExpression]) &&
+ !containsUDT(udf.dataType) &&
+ !udf.children.exists(c => containsUDT(c.dataType))
+ case _ => false
+ }
+
+ /**
+ * 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 vectorized (scalar pandas / arrow) UDF, which is not supported;
+ * - 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.)
+ */
+ 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.
+ 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).
+ 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)
+ }
+
+ /**
+ * 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.
+ *
+ * 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
+ * marked today, but requiring it here keeps "analysis accepts exactly what
the rule rewrites"
+ * structural: a future HOF missing both traits is rejected at analysis
rather than slipping
+ * through and leaving the UDF inside the lambda at runtime.
+ */
+ private def isRewritableShape(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 &&
+ hof.functions.head.isInstanceOf[LambdaFunction] &&
+ hof.functions.head.asInstanceOf[LambdaFunction].arguments
+ .forall(_.isInstanceOf[NamedLambdaVariable]) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] ||
hof.isInstanceOf[ResultTypeFromFunction]) &&
+ hof.arguments.exists { a =>
+ a.dataType.isInstanceOf[ArrayType] || a.dataType.isInstanceOf[MapType]
+ }
+
+ /**
+ * Whether `e` references a [[NamedLambdaVariable]] that it does not itself
bind, i.e. one bound
+ * by an enclosing lambda. Such an expression cannot be evaluated outside
that lambda.
+ *
+ * Shared with `ExtractPythonUDFFromLambda` so the rule can re-check the
nested-lambda guard on
+ * its own, rather than relying only on `CheckAnalysis` having already
rejected such plans.
+ */
+ def hasFreeLambdaVariable(e: Expression): Boolean = {
+ def check(expr: Expression, bound: Set[ExprId]): Boolean = expr match {
+ case LambdaFunction(function, arguments, _) =>
+ check(function, bound ++ arguments.map(_.exprId))
+ case v: NamedLambdaVariable => !bound.contains(v.exprId)
+ case other => other.children.exists(check(_, bound))
+ }
+ check(e, Set.empty)
+ }
+
def isWindowPandasUDF(e: PythonFuncExpression): Boolean = {
// This is currently only `PythonUDAF` (which means
SQL_GROUPED_AGG_PANDAS_UDF or
// SQL_GROUPED_AGG_ARROW_UDF), but we might
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 3877de93a834..6c6b96f4f590 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
@@ -393,6 +393,29 @@ trait MapBasedSimpleHigherOrderFunction extends
SimpleHigherOrderFunction {
override def argumentType: AbstractDataType = MapType
}
+/**
+ * A higher-order function whose result type is its argument's type, because
it returns a subset or
+ * reordering of the input rather than the lambda's values. Members: `filter`,
`map_filter`,
+ * `array_sort` (e.g. `filter(array<int>, ...) => array<int>`). Provides the
shared `dataType`.
+ *
+ * The counterpart is [[ResultTypeFromFunction]]; the split is a real property
of the expression,
+ * not of any evaluation: `filter` keeps its input's type, `transform`'s type
follows the lambda.
+ */
+trait ResultTypeFromArgument extends SimpleHigherOrderFunction {
+ override def dataType: DataType = argument.dataType
+}
+
+/**
+ * A higher-order function whose result type follows its lambda, not its
argument. Members:
+ * `transform`, `transform_keys`, `transform_values`, `zip_with`,
`map_zip_with`, `aggregate`, and
+ * the predicates `exists` / `forall` (whose boolean lambda gives a boolean
result).
+ *
+ * Each member computes its own `dataType` (array of the element type, a
re-keyed/re-valued map, the
+ * fold result, ...), so this is a marker with no shared implementation - the
counterpart of
+ * [[ResultTypeFromArgument]].
+ */
+trait ResultTypeFromFunction extends HigherOrderFunction
+
/**
* Transform elements in an array using the transform function. This is
similar to
* a `map` in functional programming.
@@ -418,7 +441,7 @@ trait MapBasedSimpleHigherOrderFunction extends
SimpleHigherOrderFunction {
case class ArrayTransform(
argument: Expression,
function: Expression)
- extends ArrayBasedSimpleHigherOrderFunction {
+ extends ArrayBasedSimpleHigherOrderFunction with ResultTypeFromFunction {
override def dataType: ArrayType = ArrayType(function.dataType,
function.nullable)
@@ -549,7 +572,7 @@ case class ArraySort(
argument: Expression,
function: Expression,
allowNullComparisonResult: Boolean)
- extends ArrayBasedSimpleHigherOrderFunction with CodegenFallback {
+ extends ArrayBasedSimpleHigherOrderFunction with CodegenFallback with
ResultTypeFromArgument {
def this(argument: Expression, function: Expression) = {
this(
@@ -563,7 +586,6 @@ case class ArraySort(
@transient lazy val elementType: DataType =
argument.dataType.asInstanceOf[ArrayType].elementType
- override def dataType: ArrayType = argument.dataType.asInstanceOf[ArrayType]
override def checkInputDataTypes(): TypeCheckResult = {
checkArgumentDataTypes() match {
case TypeCheckResult.TypeCheckSuccess =>
@@ -680,7 +702,7 @@ object ArraySort {
case class MapFilter(
argument: Expression,
function: Expression)
- extends MapBasedSimpleHigherOrderFunction with CodegenFallback {
+ extends MapBasedSimpleHigherOrderFunction with CodegenFallback with
ResultTypeFromArgument {
@transient lazy val (keyVar, valueVar) = {
val args = function.asInstanceOf[LambdaFunction].arguments
@@ -710,8 +732,6 @@ case class MapFilter(
ArrayBasedMapData(retKeys.toArray, retValues.toArray)
}
- override def dataType: DataType = argument.dataType
-
override def functionType: AbstractDataType = BooleanType
override def nodeName: String = "map_filter"
@@ -750,9 +770,7 @@ case class MapFilter(
case class ArrayFilter(
argument: Expression,
function: Expression)
- extends ArrayBasedSimpleHigherOrderFunction {
-
- override def dataType: DataType = argument.dataType
+ extends ArrayBasedSimpleHigherOrderFunction with ResultTypeFromArgument {
override def functionType: AbstractDataType = BooleanType
@@ -897,7 +915,7 @@ case class ArrayExists(
argument: Expression,
function: Expression,
followThreeValuedLogic: Boolean)
- extends ArrayBasedSimpleHigherOrderFunction with Predicate {
+ extends ArrayBasedSimpleHigherOrderFunction with Predicate with
ResultTypeFromFunction {
def this(argument: Expression, function: Expression) = {
this(
@@ -1035,7 +1053,7 @@ object ArrayExists {
case class ArrayForAll(
argument: Expression,
function: Expression)
- extends ArrayBasedSimpleHigherOrderFunction with Predicate {
+ extends ArrayBasedSimpleHigherOrderFunction with Predicate with
ResultTypeFromFunction {
override def nullable: Boolean =
super.nullable || function.nullable
@@ -1165,7 +1183,7 @@ case class ArrayAggregate(
zero: Expression,
merge: Expression,
finish: Expression)
- extends HigherOrderFunction with QuaternaryLike[Expression] {
+ extends HigherOrderFunction with QuaternaryLike[Expression] with
ResultTypeFromFunction {
def this(argument: Expression, zero: Expression, merge: Expression) = {
this(argument, zero, merge, LambdaFunction.identity)
@@ -1382,7 +1400,7 @@ case class ArrayAggregate(
case class TransformKeys(
argument: Expression,
function: Expression)
- extends MapBasedSimpleHigherOrderFunction with CodegenFallback {
+ extends MapBasedSimpleHigherOrderFunction with CodegenFallback with
ResultTypeFromFunction {
@transient lazy val MapType(keyType, valueType, valueContainsNull) =
argument.dataType
@@ -1448,7 +1466,7 @@ case class TransformKeys(
case class TransformValues(
argument: Expression,
function: Expression)
- extends MapBasedSimpleHigherOrderFunction with CodegenFallback {
+ extends MapBasedSimpleHigherOrderFunction with CodegenFallback with
ResultTypeFromFunction {
@transient lazy val MapType(keyType, valueType, valueContainsNull) =
argument.dataType
@@ -1514,7 +1532,8 @@ case class TransformValues(
since = "3.0.0",
group = "lambda_funcs")
case class MapZipWith(left: Expression, right: Expression, function:
Expression)
- extends HigherOrderFunction with CodegenFallback with
TernaryLike[Expression] {
+ extends HigherOrderFunction with CodegenFallback with TernaryLike[Expression]
+ with ResultTypeFromFunction {
def functionForEval: Expression = functionsForEval.head
@@ -1752,7 +1771,8 @@ case class MapZipWith(left: Expression, right:
Expression, function: Expression)
group = "lambda_funcs")
// scalastyle:on line.size.limit
case class ZipWith(left: Expression, right: Expression, function: Expression)
- extends HigherOrderFunction with CodegenFallback with
TernaryLike[Expression] {
+ extends HigherOrderFunction with CodegenFallback with TernaryLike[Expression]
+ with ResultTypeFromFunction {
def functionForEval: Expression = functionsForEval.head
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 5f80aaf9623a..8d4856a9c6ce 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -5387,6 +5387,23 @@ object SQLConf {
" must be more than one.")
.createOptional
+ val PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED =
+ buildConf("spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled")
+ .doc("When true, a scalar Python UDF may be used inside the lambda of a
higher-order " +
+ "function such as `transform` or `filter`. The UDF is not evaluated
inside the lambda; " +
+ "it is applied to the whole array outside the lambda and the lambda
reads the " +
+ "precomputed result. This uses an Arrow-based execution path, so
PyArrow is required " +
+ "even for a UDF created with useArrow=False. Because the UDF is
precomputed over the " +
+ "whole array, it runs once per element regardless of any
short-circuiting (`exists`, " +
+ "`when`) in the lambda. In particular, an `array_sort` comparator that
calls the UDF on " +
+ "both elements, `(a, b) -> udf(a, b)`, precomputes it over every pair,
costing O(n^2) " +
+ "calls and memory for an array of length n; avoid it for large arrays.
When false, such " +
+ "queries fail with UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF
as before.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(true)
+
val PYTHON_UDF_ARROW_FALLBACK_ON_UDT =
buildConf("spark.sql.execution.pythonUDF.arrow.legacy.fallbackOnUDT")
.internal()
@@ -9286,6 +9303,9 @@ class SQLConf extends Serializable with Logging with
SqlApiConf {
def pythonUDFArrowFallbackOnUDT: Boolean =
getConf(PYTHON_UDF_ARROW_FALLBACK_ON_UDT)
+ def pythonUDFInHigherOrderFunctionEnabled: Boolean =
+ getConf(PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED)
+
def pysparkPlotMaxRows: Int = getConf(PYSPARK_PLOT_MAX_ROWS)
def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
index 9dd7bef549fa..2e9520742b0a 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
@@ -90,6 +90,8 @@ class SparkOptimizer(
ExtractPythonUDFFromAggregate,
// This must be executed after `ExtractPythonUDFFromAggregate` and
before `ExtractPythonUDFs`.
ExtractGroupingPythonUDFFromAggregate,
+ // `ExtractPythonUDFs` first lifts Python UDFs out of higher-order
function lambdas
+ // (via `ExtractPythonUDFFromLambda`) and then extracts them as ordinary
top-level UDFs.
ExtractPythonUDFs,
ExtractPythonUDTFs,
// The eval-python node may be between Project/Filter and the scan node,
which breaks
@@ -116,6 +118,8 @@ class SparkOptimizer(
ExtractPythonUDFFromJoinCondition.ruleName,
ExtractPythonUDFFromAggregate.ruleName,
ExtractGroupingPythonUDFFromAggregate.ruleName,
+ // Non-excludable: a plan with a Python UDF in a higher-order function
lambda only works
+ // because `ExtractPythonUDFs` lifts it out (via
`ExtractPythonUDFFromLambda`).
ExtractPythonUDFs.ruleName,
GroupBasedRowLevelOperationScanPlanning.ruleName,
V2ScanRelationPushDown.ruleName,
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 4a09898a58c7..7c7ed9c50ab6 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
@@ -162,6 +162,7 @@ case class ArrowEvalPythonExec(
private def supportedPythonEvalTypes: Array[Int] =
Array(
PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
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 75b8465e2607..4c1c472cb10d 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
@@ -138,7 +138,10 @@ class ArrowPythonWithNamedArgumentRunner(
override protected def runnerConf: Map[String, String] = super.runnerConf ++
pythonRunnerConf
override protected def evalConf: Map[String, String] = {
- if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
+ // 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.
+ if (evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF ||
+ evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
super.evalConf ++ Map(
"input_type" -> schema.json
)
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
new file mode 100644
index 000000000000..637ec259c651
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala
@@ -0,0 +1,619 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they
can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls
out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF
can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF
once to the *whole
+ * array*, outside every lambda, and has the lambda read the result
positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s
-> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> =>
array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour
lives in the Python
+ * worker: it flattens each list column once, calls the function over all
elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round
trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the
iterated array (an
+ * element-independent value is repeated into one with a native `transform`),
so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an
ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))`
all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions:
`transform`,
+ * `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.
+ *
+ * `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 vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): 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.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * 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
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean =
function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right:
NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * 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.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is
left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] ||
hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes
both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x
n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two
elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n)
for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar:
NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a
lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the
UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n),
Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with
the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform`
runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda,
Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body,
Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position and the shared flat result array so the
comparator can read
+ // its pair's precomputed cell, sort, then drop them again. `flatCells`
must be built here, in
+ // the sort's *argument*, not inside the comparator: `ArraySort`
re-evaluates the whole
+ // comparator body on every comparison, and `ExtractPythonUDFs` hoists
only the `PythonUDF`
+ // node - the surrounding
`arrays_zip`/`transform`/`flatten`/`array_repeat` that build the cells
+ // would otherwise be rebuilt O(n^2) per comparison (O(n^3 log n)
overall). In interpreted
+ // evaluation `array_repeat` stores n references to the one computed
`flatCells` array, so the
+ // carry is O(n^2); a later copy of the carrier into the Unsafe format
would materialize each
+ // reference into O(n^3) bytes. Either way the whole pairwise path is
already O(n^2) in Python
+ // calls, so it is only intended for small arrays (see the config doc).
+ // Also carry the row width `n` so the comparator does not re-evaluate
`Size(argument)` on every
+ // comparison (negligible for a column, but `argument` may be a computed
expression). Like the
+ // cells, it is repeated into the carrier once and read as a struct field.
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val cellsField = "cells"
+ val sizeField = "n"
+ val indexed = ArraysZip(
+ Seq(
+ argument,
+ ArrayTransform(argument, LambdaFunction(posIdx, Seq(posElem, posIdx))),
+ ArrayRepeat(flatCells, n),
+ ArrayRepeat(n, n)),
+ Seq(
+ Literal(s"${carrierElementPrefix}0"),
+ Literal(carrierIndexField),
+ Literal(cellsField),
+ Literal(sizeField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. The
cells and `n` live in
+ // struct fields carried by every element, so the comparator only does
field reads plus an
+ // `element_at`, all O(1). `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1,
Some(carrierIndexField))
+ val cells = GetStructField(cmpLeft, 2, Some(cellsField))
+ val width = GetStructField(cmpLeft, 3, Some(sizeField))
+ val comparison = ElementAt(
+ cells,
+ Add(Add(Multiply(idxOf(cmpLeft), width), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)),
allowNull), 0)
+ }
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so
everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The
lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is
rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped
afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's
type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key
sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly
its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type:
`transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever
pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on
`map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) =>
MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) =>
MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType,
keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None,
failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to
the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one
exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array,
indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is
special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption,
Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +:
built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps
its own class (via
+ // `withNewChildren`, children being arguments then functions); a
desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must
survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier,
newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so
project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides
to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * 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.
+ */
+ 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)
+ case _ => false
+ })
+
+ /**
+ * Whether `body` holds a rewritable UDF belonging to *this* lambda. A
nested function's lambda is
+ * skipped (its UDF reads that lambda's variable), but its *arguments* are
not: in
+ * `transform(arr, x -> transform(udf(x), y -> y))`, `udf(x)` is in the
inner argument and lifts
+ * onto `arr`.
+ */
+ private def hasDirectRewritableUDF(body: Expression): Boolean = body match {
+ case e if PythonUDF.isElementwiseRewritableUDF(e) => true
+ case hof: HigherOrderFunction =>
hof.arguments.exists(hasDirectRewritableUDF)
+ case e => e.children.exists(hasDirectRewritableUDF)
+ }
+
+ /** The pieces produced by [[buildCarrier]]. */
+ private case class Carrier(
+ carrier: Expression,
+ body: Expression,
+ boundVar: NamedLambdaVariable,
+ extraBoundVars: Seq[NamedLambdaVariable])
+
+ /**
+ * Builds the carrier array and the rewritten lambda body.
+ *
+ * The carrier is `arrays_zip` of the original arrays, one array per lifted
UDF, and - when the
+ * lambda declares an index parameter - an index array. The rewritten body
reads each of those
+ * through a struct field of the lambda variable bound to the carrier.
+ *
+ * `alsoBind` names further lambda variables that should read the same
carrier; it exists for
+ * `array_sort`'s comparator, whose two parameters are both elements of the
same array.
+ */
+ private def buildCarrier(
+ arguments: Seq[Expression],
+ function: Expression,
+ elementVars: Seq[NamedLambdaVariable],
+ indexVar: Option[NamedLambdaVariable],
+ alsoBind: Seq[NamedLambdaVariable] = Nil): Carrier = {
+ val LambdaFunction(body, _, _) = function
+ val lambdaExprIds =
+ (elementVars ++ indexVar.toSeq ++ alsoBind).map(_.exprId).toSet
+
+ // Collect the UDF calls to lift. Innermost first, so that a nested call
like `f(g(x))` has
+ // `g` lifted before `f`, letting `f`'s array UDF consume `g`'s array
result.
+ val liftableUDFs = collectLiftableUDFs(body)
+
+ // With more than one argument the arrays may be ragged (`zip_with` /
`map_zip_with` pad with
+ // nulls), so flattening them independently would misalign the elements.
Projecting each out of
+ // one common `arrays_zip` pads them to the same per-row length, which the
positional rewrite
+ // requires.
+ val alignedArguments =
+ if (arguments.length > 1) {
+ val names = arguments.indices.map(i => s"$carrierElementPrefix$i")
+ val zipped = ArraysZip(arguments, names.map(Literal(_)))
+ arguments.indices.map(i => unwrapCarrier(zipped, i))
+ } else {
+ arguments
+ }
+
+ // An index array, when the lambda asked for the element index.
+ val indexArray = indexVar.map { _ =>
+ val head = alignedArguments.head
+ val headType = head.dataType.asInstanceOf[ArrayType]
+ val v = NamedLambdaVariable("x", headType.elementType,
headType.containsNull)
+ val i = NamedLambdaVariable("i", IntegerType, nullable = false)
+ ArrayTransform(head, LambdaFunction(i, Seq(v, i)))
+ }
+
+ // Maps each element/index variable to the array it stands for, so a UDF
argument written in
+ // terms of the variables can be rewritten as an expression over whole
arrays. For a
+ // comparator, `alsoBind`'s variables denote the same array as the element
variable.
+ val arrayOfVar: Map[ExprId, Expression] =
+ elementVars.map(_.exprId).zip(alignedArguments).toMap ++
+ indexVar.map(_.exprId -> indexArray.get).toMap ++
+ alsoBind.map(_.exprId -> alignedArguments.head).toMap
+
+ // One lifted array UDF per distinct call. `arrayResults` maps each
original call (by `liftKey`)
+ // to the array holding its per-element results, so a nested call
`f(g(x))` and the carrier
+ // lookups can find it. Deterministic calls that lift to the *same*
function over the *same*
+ // array arguments share one lifted UDF: a key-form comparator's `udf(a)`
and `udf(b)`
+ // canonicalize differently (`a` != `b`) but both read the whole array, so
without this the
+ // Python function would run 2n times instead of n. Nondeterministic calls
stay distinct (their
+ // signature is the lifted node itself, carrying a distinct `resultId`),
matching `liftKey`.
+ var arrayResults = Map.empty[Expression, Expression]
+ val distinctLifted = scala.collection.mutable.ArrayBuffer.empty[PythonUDF]
+ val ordinalBySignature =
scala.collection.mutable.HashMap.empty[Expression, Int]
+ 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)
+ }
+ 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
+ // containsNull = true.
+ ArrayType(udf.dataType, containsNull = true),
+ arrayArgs,
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ udf.udfDeterministic)
+ val signature: Expression = if (udf.udfDeterministic)
lifted.canonicalized else lifted
+ val ordinal = ordinalBySignature.getOrElseUpdate(signature, {
+ val o = distinctLifted.length
+ distinctLifted += lifted
+ o
+ })
+ arrayResults += (liftKey(udf) -> distinctLifted(ordinal))
+ udfFieldByKey += (liftKey(udf) -> ordinal)
+ }
+ val liftedArrays = distinctLifted.toSeq
+
+ // The carrier: the original arrays first, then one field per lifted UDF,
then the index.
+ val carrierFields = alignedArguments ++ liftedArrays ++ indexArray.toSeq
+ val carrierNames =
+ arguments.indices.map(i => s"$carrierElementPrefix$i") ++
+ liftedArrays.indices.map(i => s"$carrierUDFFieldPrefix$i") ++
+ indexArray.map(_ => carrierIndexField).toSeq
+ val carrier = ArraysZip(carrierFields, carrierNames.map(Literal(_)))
+
+ val structType = carrier.dataType.asInstanceOf[ArrayType].elementType
+ val boundVar = NamedLambdaVariable("s", structType, nullable = false)
+ val extraBoundVars = alsoBind.map(v =>
+ NamedLambdaVariable(v.name, structType, nullable = false))
+
+ // Which struct field each lambda variable reads. For a comparator,
`alsoBind`'s variable reads
+ // the same ordinals but through its own bound variable.
+ val fieldOfVar: Map[ExprId, Int] =
+ elementVars.map(_.exprId).zipWithIndex.toMap ++
+ indexVar.map(_.exprId -> (carrierFields.length - 1)).toMap
+ val extraVarOf: Map[ExprId, NamedLambdaVariable] =
+ alsoBind.map(_.exprId).zip(extraBoundVars).toMap
+
+ // Rewrite the body. This must be top-down: a UDF call is matched by its
canonicalized form,
+ // and rewriting its arguments first (a variable becoming a struct field
read) would change
+ // that form so the call no longer matches and would be left inside the
lambda. Replacing the
+ // call outright also stops the traversal descending into arguments that
no longer exist.
+ def readerFor(v: NamedLambdaVariable, udfOrdinal: Option[Int]): Expression
= {
+ val base = extraVarOf.getOrElse(v.exprId, boundVar)
+ udfOrdinal match {
+ case Some(u) =>
+ GetStructField(base, arguments.length + u,
Some(s"$carrierUDFFieldPrefix$u"))
+ case None =>
+ val ordinal = fieldOfVar(v.exprId)
+ GetStructField(base, ordinal, Some(carrierNames(ordinal)))
+ }
+ }
+
+ val rewrittenBody = body.transformDown {
+ case udf: PythonUDF if udfFieldByKey.contains(liftKey(udf)) =>
+ val ordinal = udfFieldByKey(liftKey(udf))
+ // A UDF over a comparator's right-hand element must read that
element's key, so the
+ // struct field is read through whichever bound variable the call's
own arguments used.
+ val side = udf.collectFirst {
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => v
+ }
+ side match {
+ case Some(v) => readerFor(v, Some(ordinal))
+ case None =>
+ GetStructField(boundVar, arguments.length + ordinal,
+ Some(s"$carrierUDFFieldPrefix$ordinal"))
+ }
+ case v: NamedLambdaVariable if fieldOfVar.contains(v.exprId) =>
readerFor(v, None)
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) =>
+ // A comparator's right-hand element itself, read through its own
bound variable.
+ GetStructField(extraVarOf(v.exprId), 0, Some(carrierNames.head))
+ }
+
+ Carrier(carrier, rewrittenBody, boundVar, extraBoundVars)
+ }
+
+ /**
+ * Collects the Python UDF calls directly in `body` that must be lifted,
innermost first.
+ *
+ * Every rewritable UDF directly in the lambda body is lifted, even one
whose arguments do not
+ * read the lambda variable (`transform(arr, _ -> f(lit(10)))`). It cannot
stay inside the lambda,
+ * and lifting it - `overArray` repeats a constant/outer-column argument
into an aligned array -
+ * gives it the lambda's own call domain: once per element, and zero times
for a null or empty
+ * array. Leaving it to [[ExtractPythonUDFs]] would instead call it once per
row, including rows
+ * whose array is null or empty where the lambda never runs.
+ */
+ private def collectLiftableUDFs(body: Expression): Seq[PythonUDF] = {
+ val collected = Seq.newBuilder[PythonUDF]
+ def visit(e: Expression): Unit = {
+ // A nested higher-order function's lambda is not ours to rewrite, but
its arguments are
+ // evaluated outside that lambda and so belong to this body. See
`hasDirectRewritableUDF`.
+ val children = e match {
+ case hof: HigherOrderFunction => hof.arguments
+ case other => other.children
+ }
+ // Children first, so nested calls come out innermost-first.
+ children.foreach(visit)
+ e match {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ collected += udf
+ case _ =>
+ }
+ }
+ visit(body)
+ // Deduplicate identical calls so the same UDF is evaluated once per
array. Nondeterministic
+ // calls are kept distinct (see `liftKey`).
+ val seen = scala.collection.mutable.LinkedHashMap.empty[Expression,
PythonUDF]
+ collected.result().foreach(udf => seen.getOrElseUpdate(liftKey(udf), udf))
+ seen.values.toSeq
+ }
+
+ private def readsLambdaVariable(e: Expression, lambdaExprIds: Set[ExprId]):
Boolean =
+ e.exists {
+ case v: NamedLambdaVariable => lambdaExprIds.contains(v.exprId)
+ case _ => false
+ }
+
+ /**
+ * The key that decides whether two UDF calls are "the same call" for
lifting. A deterministic
+ * call is deduplicated by canonical form, so an identical call is evaluated
once. A
+ * nondeterministic call must stay distinct - `transform(arr, x -> f(x) +
f(x))` calls `f` twice
+ * and each call may return a different value - so it is keyed by its own
`resultId`-bearing node
+ * (`canonicalized` erases `resultId`, which would collapse the two calls
into one).
+ */
+ private def liftKey(udf: PythonUDF): Expression =
+ if (udf.udfDeterministic) udf.canonicalized else udf
+
+ /**
+ * Turns a UDF argument expression, written in terms of single elements,
into the equivalent
+ * `array<T>` aligned with the iterated array, so the worker can flatten
every argument uniformly.
+ *
+ * - a lambda variable becomes the array it stands for;
+ * - an already-lifted UDF call (a nested call `f(g(x))`) becomes its array
result;
+ * - an expression independent of the lambda is repeated into an aligned
array with a native
+ * `transform`, rather than passed through as a scalar to broadcast;
+ * - anything else is an expression over the elements, computed for every
element by a native
+ * `transform` that stays inside the JVM.
+ *
+ * An already-lifted UDF call nested inside a composite argument (`f(g(x) +
1)`, `f(-g(x))`) is
+ * handled by the last case too: each such call is replaced by a synthetic
variable standing for
+ * that call's aligned array result, which is then zipped in like any
element array. Leaving the
+ * raw `g` inside the generated `transform` lambda would put a `PythonUDF`
back inside a lambda -
+ * `ExtractPythonUDFs` would then extract a `g` whose child is a
`NamedLambdaVariable` (the
+ * SPARK-48706 failure mode), so this substitution is for correctness, not
just efficiency.
+ */
+ private def overArray(
+ child: Expression,
+ firstArgument: Expression,
+ arrayOfVar: Map[ExprId, Expression],
+ lambdaExprIds: Set[ExprId],
+ arrayResults: Map[Expression, Expression]): Expression = child match {
+ case v: NamedLambdaVariable if arrayOfVar.contains(v.exprId) =>
arrayOfVar(v.exprId)
+ case udf: PythonUDF if arrayResults.contains(liftKey(udf)) =>
+ arrayResults(liftKey(udf))
+ case e =>
+ // Replace each already-lifted nested UDF call with a synthetic variable
standing for that
+ // call's aligned array result, then fold those arrays into the
variable-to-array map so the
+ // logic below treats them exactly like element variables. This keeps a
lifted UDF buried in
+ // a composite argument (`f(g(x) + 1)`) from being left as a raw UDF
inside a lambda.
+ val nestedVars =
scala.collection.mutable.LinkedHashMap.empty[Expression, NamedLambdaVariable]
+ val expr = e.transformUp {
+ case u: PythonUDF if arrayResults.contains(liftKey(u)) =>
+ nestedVars.getOrElseUpdate(liftKey(u), {
+ val arrType =
arrayResults(liftKey(u)).dataType.asInstanceOf[ArrayType]
+ NamedLambdaVariable("g", arrType.elementType, arrType.containsNull)
+ })
+ }
+ val fullArrayOf = arrayOfVar ++
+ nestedVars.map { case (key, v) => v.exprId -> arrayResults(key) }
+
+ if (!readsLambdaVariable(expr, fullArrayOf.keySet)) {
+ // Independent of the element (an outer column or constant): repeat it
into an array aligned
+ // with the iterated array, so every UDF argument is a single-level
array the worker
+ // flattens the same way. `transform(arr, _ -> e)` keeps the value
constant, matching shape.
+ val arrType = firstArgument.dataType.asInstanceOf[ArrayType]
+ val v = NamedLambdaVariable("x", arrType.elementType,
arrType.containsNull)
+ ArrayTransform(firstArgument, LambdaFunction(expr, Seq(v)))
+ } else {
+ // An expression over the element(s) and/or nested results, e.g.
`udf(x * 2)` or
+ // `f(g(x) + 1)`. Compute it for every element with a native
`transform`, which stays inside
+ // the JVM. Multi-argument shapes need the values side by side, so zip
them first.
+ val arrays = fullArrayOf.values.toSeq.distinct
+ if (arrays.length == 1) {
+ val arr = arrays.head
+ val elemType = arr.dataType.asInstanceOf[ArrayType]
+ val v = NamedLambdaVariable("x", elemType.elementType,
elemType.containsNull)
+ val substituted = expr.transformUp {
+ case old: NamedLambdaVariable if fullArrayOf.contains(old.exprId)
=> v
+ }
+ ArrayTransform(arr, LambdaFunction(substituted, Seq(v)))
+ } else {
+ // Zip every array the expression may read, then project the fields
it needs.
+ val names = arrays.indices.map(i => s"$carrierElementPrefix$i")
+ val zipped = ArraysZip(arrays, names.map(Literal(_)))
+ val structType = zipped.dataType.asInstanceOf[ArrayType].elementType
+ val v = NamedLambdaVariable("z", structType, nullable = false)
+ val ordinalOf = fullArrayOf.map { case (id, arr) => id ->
arrays.indexOf(arr) }
+ val substituted = expr.transformUp {
+ case old: NamedLambdaVariable if ordinalOf.contains(old.exprId) =>
+ GetStructField(v, ordinalOf(old.exprId),
Some(names(ordinalOf(old.exprId))))
+ }
+ ArrayTransform(zipped, LambdaFunction(substituted, Seq(v)))
+ }
+ }
+ }
+
+ /**
+ * Projects field `ordinal` back out of a carrier array. Used where the
result is built from the
+ * input elements rather than from the lambda's return value: `filter`,
`array_sort` and the map
+ * family.
+ */
+ private def unwrapCarrier(carrierArray: Expression, ordinal: Int):
Expression = {
+ val structType = carrierArray.dataType.asInstanceOf[ArrayType].elementType
+ val v = NamedLambdaVariable("s", structType, nullable = false)
+ ArrayTransform(
+ carrierArray,
+ LambdaFunction(
+ GetStructField(v, ordinal, Some(s"$carrierElementPrefix$ordinal")),
Seq(v)))
+ }
+
+ private val carrierElementPrefix = "c"
+ private val carrierUDFFieldPrefix = "u"
+ private val carrierIndexField = "idx"
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
index af9b41e93a02..7a9ec702fea8 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
@@ -264,7 +264,14 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with
Logging {
expressions.flatMap(collectEvaluableUDFs)
}
- def apply(plan: LogicalPlan): LogicalPlan = plan match {
+ def apply(plan: LogicalPlan): LogicalPlan = applyInternal(
+ // Lift Python UDFs out of higher-order function lambdas first, turning
them into ordinary
+ // top-level UDFs this rule can then extract. Running it here rather than
as a separate batch
+ // entry makes the ordering structural: `CheckAnalysis` accepts a
lambda-UDF plan only because
+ // this rewrite will run, so the two must not be separable by rule
reordering or exclusion.
+ ExtractPythonUDFFromLambda(plan))
+
+ private def applyInternal(plan: LogicalPlan): LogicalPlan = plan match {
// SPARK-26293: A subquery will be rewritten into join later, and will go
through this rule
// eventually. Here we skip subquery, as Python UDF only needs to be
extracted once.
case s: Subquery if s.correlated => plan
@@ -342,6 +349,7 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with
Logging {
case PythonEvalType.SQL_SCALAR_PANDAS_UDF
| PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
| PythonEvalType.SQL_ARROW_BATCHED_UDF
+ | PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF
| PythonEvalType.SQL_SCALAR_ARROW_UDF
| PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF =>
ArrowEvalPython(validUdfs, resultAttrs, child, evalType)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
index 42da9f060b85..d0228f9341c1 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
@@ -27,7 +27,7 @@ import org.apache.spark.sql.connector.catalog.CatalogManager
import org.apache.spark.sql.execution.{FilterExec, InputAdapter,
WholeStageCodegenExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
-import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.test.{ExamplePointUDT, SharedSparkSession}
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
@@ -367,3 +367,10 @@ class MyDummyScalarPandasUDF extends
UserDefinedPythonFunction(
dataType = BooleanType,
pythonEvalType = PythonEvalType.SQL_SCALAR_PANDAS_UDF,
udfDeterministic = true)
+
+class MyDummyUDTPythonUDF extends UserDefinedPythonFunction(
+ name = "dummyUDTUDF",
+ func = new DummyUDF,
+ dataType = new ExamplePointUDT,
+ pythonEvalType = PythonEvalType.SQL_BATCHED_UDF,
+ udfDeterministic = true)
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
new file mode 100644
index 000000000000..b764b5c9cf9d
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala
@@ -0,0 +1,379 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.{AnalysisException, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{LambdaFunction, Literal,
NamedArgumentExpression,
+ PythonUDF}
+import org.apache.spark.sql.catalyst.plans.logical.ArrowEvalPython
+import org.apache.spark.sql.functions.{array_sort, col, forall, lit,
map_filter, map_zip_with,
+ transform, transform_keys, transform_values, when, zip_with}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.{ExamplePointUDT, SharedSparkSession}
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, StringType}
+
+/**
+ * Plan-shape tests for [[ExtractPythonUDFFromLambda]].
+ *
+ * These assert the structure the rewrite produces - that no `PythonUDF` is
left inside a
+ * `LambdaFunction`, that the lifted UDF is an element-wise UDF over an array,
and that plans
+ * without a UDF in a lambda are untouched. End-to-end result correctness is
covered by
+ * `pyspark.sql.tests.test_udf_in_higher_order_function`, which needs a real
Python worker.
+ */
+class ExtractPythonUDFFromLambdaSuite extends QueryTest with
SharedSparkSession {
+ import testImplicits._
+
+ private val pythonUDF = new MyDummyPythonUDF
+ private val scalarPandasUDF = new MyDummyScalarPandasUDF
+ // Used where a UDF call must receive two arguments, e.g. a pairwise
comparator.
+ private val pythonUDF2 = new MyDummyPythonUDF
+ private val nondeterministicUDF = new MyDummyNondeterministicPythonUDF
+ private val udtUDF = new MyDummyUDTPythonUDF
+
+ private def arrayDF = Seq(Seq(1, 2, 3)).toDF("values")
+
+ /** All `PythonUDF`s that remain inside a lambda in the optimized plan. */
+ private def udfsInsideLambda(df: org.apache.spark.sql.DataFrame):
Seq[PythonUDF] = {
+ df.queryExecution.optimizedPlan.expressions.flatMap { e =>
+ e.collect { case l: LambdaFunction => l }.flatMap { l =>
+ l.collect { case u: PythonUDF => u }
+ }
+ }
+ }
+
+ private def liftedUDFs(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] =
{
+ df.queryExecution.optimizedPlan.collect {
+ case a: ArrowEvalPython => a.udfs
+ }.flatten
+ }
+
+ test("transform: the UDF is lifted out of the lambda as an element-wise
array UDF") {
+ val df = arrayDF.select(transform(col("values"), x =>
pythonUDF(x)).as("r"))
+
+ // The whole point of the rewrite: nothing Python-shaped is left inside a
lambda.
+ assert(udfsInsideLambda(df).isEmpty)
+
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ assert(lifted.head.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)
+ // The lifted UDF takes the array and returns an array of the original
return type.
+ assert(lifted.head.dataType == ArrayType(pythonUDF.dataType, containsNull
= true))
+ assert(lifted.head.children.size == 1)
+ assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType])
+ }
+
+ test("filter/exists/forall: the UDF is lifted out of the lambda") {
+ val exprs = Seq(
+ org.apache.spark.sql.functions.filter(col("values"), x => pythonUDF(x)),
+ org.apache.spark.sql.functions.exists(col("values"), x => pythonUDF(x)),
+ org.apache.spark.sql.functions.forall(col("values"), x => pythonUDF(x)))
+
+ exprs.foreach { e =>
+ val df = arrayDF.select(e.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF still inside a lambda for $e")
+ assert(liftedUDFs(df).forall(_.evalType ==
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF))
+ }
+ }
+
+ test("aggregate: a UDF anywhere in the fold is rejected") {
+ // 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" ->
+ 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)))
+ cases.foreach { case (name, expr) =>
+ val e = intercept[AnalysisException] {
+ arrayDF.select(expr.as("r")).collect()
+ }
+ assert(e.getCondition ==
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF",
+ s"expected aggregate with a UDF in $name to be rejected")
+ }
+ }
+
+ test("several UDFs and nested UDFs in one lambda are all lifted") {
+ val several = arrayDF.select(
+ transform(col("values"), x => pythonUDF(x) || pythonUDF(x +
lit(1))).as("r"))
+ assert(udfsInsideLambda(several).isEmpty)
+ // Two distinct calls, so two lifted array UDFs.
+ assert(liftedUDFs(several).size == 2)
+
+ // The same deterministic call twice must be evaluated once.
+ val duplicated = arrayDF.select(
+ transform(col("values"), x => pythonUDF(x) || pythonUDF(x)).as("r"))
+ assert(udfsInsideLambda(duplicated).isEmpty)
+ assert(liftedUDFs(duplicated).size == 1)
+ }
+
+ test("a UDF nested inside a composite argument is lifted, not left inside a
lambda") {
+ // SPARK-27052: `f(g(x) + 1)` / `f(-g(x))`. The inner call `g(x)` is
already lifted; its
+ // occurrence buried inside the composite argument must be substituted
too, or a raw `g` over a
+ // lambda variable would be left inside the generated transform for
`ExtractPythonUDFs` to
+ // re-extract (the SPARK-48706 failure mode). Both nesting shapes must
leave no UDF in a lambda.
+ val plusOne = arrayDF.select(
+ transform(col("values"), x => pythonUDF(pythonUDF(x).cast("int") +
lit(1))).as("r"))
+ assert(udfsInsideLambda(plusOne).isEmpty)
+ // Two distinct calls (inner `g(x)`, outer `f(g(x) + 1)`), so two lifted
array UDFs.
+ assert(liftedUDFs(plusOne).size == 2)
+
+ val negated = arrayDF.select(
+ transform(col("values"), x =>
pythonUDF(-pythonUDF(x).cast("int"))).as("r"))
+ assert(udfsInsideLambda(negated).isEmpty)
+ assert(liftedUDFs(negated).size == 2)
+ }
+
+ test("identical nondeterministic calls are lifted distinctly, not
deduplicated") {
+ // Deduplicating `f(x)` with `f(x)` would collapse two independent draws
into one; a
+ // nondeterministic UDF must keep each call, so both are lifted.
+ val df = arrayDF.select(
+ transform(col("values"), x => nondeterministicUDF(x) ||
nondeterministicUDF(x)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ assert(liftedUDFs(df).size == 2)
+ }
+
+ test("transform_values whose result type equals the key type replaces
values, not keys") {
+ // SPARK-27052: dispatch is by concrete function, not result type. For
map<string, string>,
+ // transform_values' lambda also returns string, which must not be treated
as new keys.
+ val maps = Seq(Map("a" -> "x")).toDF("m")
+ val df = maps.select(transform_values(col("m"), (k, v) =>
pythonUDF(v).cast("string")).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ // Keys are preserved: the map still has the original key type and no
new-key projection.
+ val mapType =
df.queryExecution.analyzed.schema.head.dataType.asInstanceOf[MapType]
+ 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.
+ 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")
+ }
+
+ test("a UDF on the outer element of a nested array is lifted") {
+ // Here the UDF applies to the outer array's element, which is a real
column (the element just
+ // happens to be an array), so it is lifted onto that column like any
other single-level array.
+ val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
+ .select(transform(col("values"), inner => pythonUDF(inner)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+
+ test("a rewritable higher-order function inside another lambda leaves no UDF
in any lambda") {
+ // `transform(arr2, i -> array_max(transform(arr, x -> f(x))) + i)`: the
inner `transform`
+ // iterates the real column `arr` (not the outer lambda variable `i`), so
it is rewritable and
+ // its UDF is lifted out of the inner lambda. The lifted element-wise UDF
reads only `arr`, so
+ // `ExtractPythonUDFs` then hoists it out of the outer lambda too
(evaluated once per row). The
+ // end state must leave no `PythonUDF` inside any lambda.
+ val df = Seq((Seq(1, 2, 3), Seq(10, 20))).toDF("arr", "arr2")
+ .select(transform(col("arr2"), i =>
+ org.apache.spark.sql.functions.array_max(
+ transform(col("arr"), x => pythonUDF(x).cast("int"))) + i).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ assert(liftedUDFs(df).size == 1)
+ }
+
+ test("a lambda with no Python UDF is left unchanged") {
+ val df = arrayDF.select(transform(col("values"), x => x + lit(1)).as("r"))
+ val analyzed = df.queryExecution.analyzed
+ val optimized = df.queryExecution.optimizedPlan
+ // The rule must be inert: no eval-python node is introduced.
+ assert(optimized.collect { case a: ArrowEvalPython => a }.isEmpty)
+ assert(!optimized.toString.contains("pythonUDF"))
+ assert(analyzed.expressions.flatMap(_.collect { case u: PythonUDF => u
}).isEmpty)
+ }
+
+ test("every mapping higher-order function is rewritten") {
+ // One assertion per function, so that a shape regressing to "UDF left
inside a lambda" is
+ // caught here rather than only by the end-to-end Python suite.
`aggregate` / `reduce` are not
+ // here: they fold rather than map, so a UDF in them is rejected (see the
aggregate test above).
+ val arrays = Seq((Seq(1, 2), Seq(3, 4))).toDF("l", "r")
+ val maps = Seq((Map("a" -> 1), Map("a" -> 2))).toDF("l", "r")
+
+ val arrayCases = Seq(
+ "transform" -> transform(col("l"), x => pythonUDF(x)),
+ "transform with index" -> transform(col("l"), (x, i) => pythonUDF(x) ||
i > 0),
+ "filter" -> org.apache.spark.sql.functions.filter(col("l"), x =>
pythonUDF(x)),
+ "exists" -> org.apache.spark.sql.functions.exists(col("l"), x =>
pythonUDF(x)),
+ "forall" -> forall(col("l"), x => pythonUDF(x)),
+ "zip_with" -> zip_with(col("l"), col("r"), (a, b) => pythonUDF(a) ||
pythonUDF(b)),
+ "array_sort" -> array_sort(col("l"),
+ (a, b) => when(pythonUDF(a) === pythonUDF(b),
lit(0)).otherwise(lit(1))))
+ arrayCases.foreach { case (name, expr) =>
+ val df = arrays.select(expr.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for
$name")
+ }
+
+ val mapCases = Seq(
+ "transform_keys" -> transform_keys(col("l"), (k, v) =>
pythonUDF(k).cast("string")),
+ "transform_values" -> transform_values(col("l"), (k, v) => pythonUDF(v)),
+ "map_filter" -> map_filter(col("l"), (k, v) => pythonUDF(v)),
+ "map_zip_with" -> map_zip_with(col("l"), col("r"), (k, a, b) =>
pythonUDF(a) || pythonUDF(b)))
+ mapCases.foreach { case (name, expr) =>
+ val df = maps.select(expr.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for
$name")
+ }
+ }
+
+ test("a pairwise array_sort comparator is lifted over the cross product") {
+ // One UDF call receiving both elements has no per-element key, so the UDF
is precomputed over
+ // every ordered pair instead. The call must end up outside every lambda
like any other.
+ val df = arrayDF.select(
+ array_sort(col("values"), (a, b) => pythonUDF2(a,
b).cast("int")).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ // The UDF takes both pair sides, each a flat array of all n*n pairs.
+ assert(lifted.head.children.size == 2)
+ }
+
+ test("a key-form array_sort comparator lifts one UDF per element, not per
comparator side") {
+ // `(a, b) -> udf(a) < udf(b)`: `udf(a)` and `udf(b)` canonicalize
differently (distinct
+ // variable exprIds) but both lift to the same function over the whole
array, so they must be
+ // deduplicated into one lifted UDF - otherwise the Python function runs
2n times instead of n.
+ val df = arrayDF.select(
+ array_sort(col("values"),
+ (a, b) => when(pythonUDF(a) === pythonUDF(b),
lit(0)).otherwise(lit(1))).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ assert(liftedUDFs(df).size == 1)
+ }
+
+ test("a pandas UDF inside a lambda still fails analysis") {
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x =>
scalarPandasUDF(x))).collect()
+ }
+ assert(e.getCondition ==
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("a UDF with a UDT type inside a lambda still fails analysis") {
+ // Lifting forces the Arrow element-wise eval type, which has no UDT
fallback, so a UDF whose
+ // argument or return type involves a UDT is not rewritable and keeps the
previous analysis
+ // error rather than failing at runtime.
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x => udtUDF(x))).collect()
+ }
+ 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.
+ val plain = PythonUDF("f", null, IntegerType, Seq(Literal(1)),
+ PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
+ assert(PythonUDF.isElementwiseRewritableUDF(plain))
+
+ val zeroArg = plain.copy(children = Seq.empty)
+ assert(!PythonUDF.isElementwiseRewritableUDF(zeroArg))
+
+ val named = plain.copy(children = Seq(NamedArgumentExpression("k",
Literal(1))))
+ assert(!PythonUDF.isElementwiseRewritableUDF(named))
+
+ val udtReturn = plain.copy(dataType = new ExamplePointUDT)
+ assert(!PythonUDF.isElementwiseRewritableUDF(udtReturn))
+
+ val udtArg = plain.copy(children = Seq(Literal.create(null, new
ExamplePointUDT)))
+ assert(!PythonUDF.isElementwiseRewritableUDF(udtArg))
+ }
+
+ test("a zero-argument UDF inside a lambda still fails analysis") {
+ // `transform(arr, x -> f())` has no argument to carry the iterated shape,
so it is not
+ // rewritable and must keep failing analysis rather than crash the Python
worker at runtime.
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), _ => pythonUDF())).collect()
+ }
+ assert(e.getCondition ==
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("a nondeterministic iterated argument is rejected") {
+ // The rewrite references the iterated argument several times and
nondeterministic expressions
+ // are not subexpression-eliminated, so a nondeterministic argument like
`shuffle(arr)` would be
+ // evaluated independently per reference and misalign the results. It must
fail analysis.
+ val e = intercept[AnalysisException] {
+ arrayDF.select(
+ org.apache.spark.sql.functions.filter(
+ org.apache.spark.sql.functions.shuffle(col("values")),
+ x => pythonUDF(x))).collect()
+ }
+ assert(e.getCondition ==
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("the rewrite can be disabled by conf, restoring the previous error") {
+ withSQLConf(SQLConf.PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED.key ->
"false") {
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x => pythonUDF(x))).collect()
+ }
+ assert(e.getCondition ==
"UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+ }
+
+ test("the rewrite cannot be disabled via excludedRules") {
+ // The lambda rewrite is driven by `ExtractPythonUDFs` (not a standalone
batch rule), and that
+ // rule is non-excludable, so a plan that only works because of the
rewrite cannot be broken by
+ // excludedRules. Excluding either name must leave the UDF lifted out of
the lambda.
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ Seq(ExtractPythonUDFFromLambda.ruleName,
ExtractPythonUDFs.ruleName).mkString(",")) {
+ val df = arrayDF.select(transform(col("values"), x =>
pythonUDF(x)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+ }
+
+ test("a UDF over only constants is still lifted per element") {
+ // SPARK-27052: `transform(arr, x -> udf(lit(10)))` does not read the
element, but it is still
+ // lifted per element (`overArray` repeats the constant into an aligned
array) so it keeps the
+ // lambda's call domain - once per element, zero times for a null/empty
array - rather than
+ // being left to ExtractPythonUDFs, which would call it once per row.
+ val df = arrayDF.select(transform(col("values"), _ =>
pythonUDF(lit(10))).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+
+ test("a UDF argument that is an expression over the element is lifted") {
+ val df = arrayDF.select(transform(col("values"), x => pythonUDF(x *
lit(2))).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ // The argument became an array-valued expression over the whole array.
+ assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType])
+ }
+
+ test("an outer column argument is repeated into an aligned array for the
lifted UDF") {
+ val df = Seq((Seq(1, 2), 10)).toDF("values", "base")
+ .select(transform(col("values"), x => pythonUDF(x, col("base"))).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ // Both arguments are single-level arrays aligned with the iterated array:
the element argument
+ // is the array itself, and the outer column is repeated into an aligned
array so the worker
+ // flattens every argument uniformly.
+ lifted.head.children.foreach { c =>
+ assert(c.dataType.isInstanceOf[ArrayType])
+ assert(c.dataType.asInstanceOf[ArrayType].elementType == IntegerType)
+ }
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
index dfff8ed9e975..88e1e7a9122b 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
@@ -18,7 +18,7 @@
package org.apache.spark.sql.execution.python
import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, Row}
-import org.apache.spark.sql.functions.{array, avg, col, count, transform}
+import org.apache.spark.sql.functions.{aggregate, array, avg, col, count, lit,
transform}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.LongType
@@ -117,16 +117,28 @@ class PythonUDFSuite extends SharedSparkSession {
assert(df.agg(pandasTestUDF(df("id"))).schema.fieldNames.exists(_.startsWith(udfName)))
}
- test("SPARK-48706: Negative test case for Python UDF in higher order
functions") {
+ test("SPARK-27052: Python UDF in a higher order function lambda") {
assume(shouldTestPythonUDFs)
+ // SPARK-48706 originally rejected this. `ExtractPythonUDFFromLambda` now
rewrites it so the
+ // UDF is applied to the whole array outside the lambda, so it runs and
returns a result.
+ checkAnswer(
+ spark.range(1).select(transform(array("id"), x => pythonTestUDF(x))),
+ Row(Seq(0)))
+ }
+
+ test("SPARK-27052: Negative test case for Python UDF in higher order
functions") {
+ assume(shouldTestPythonUDFs)
+ // A UDF reading `aggregate`'s accumulator is sequential, so there is no
array to precompute
+ // over and the rewrite does not apply. This must still fail rather than
give a wrong answer.
checkError(
exception = intercept[AnalysisException] {
- spark.range(1).select(transform(array("id"), x =>
pythonTestUDF(x))).collect()
+ spark.range(1).select(
+ aggregate(array("id"), lit(0L), (acc, x) => pythonTestUDF(acc) +
x)).collect()
},
condition = "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF",
parameters = Map("funcName" -> "\"pyUDF(namedlambdavariable())\""),
context = ExpectedContext(
- "transform", s".*${this.getClass.getSimpleName}.*"))
+ "aggregate", s".*${this.getClass.getSimpleName}.*"))
}
test("SPARK-48666: Python UDF execution against partitioned column") {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
