HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3736199081
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -8657,7 +8657,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: the UDF reads a value that only
exists while the lambda iterates (a fold accumulator in `aggregate`/`reduce`,
or an array bound by an enclosing lambda in a nested higher-order function), so
it cannot be applied to the whole column at once. Rewrite the query so the UDF
is applied outside the lambda."
Review Comment:
Fixed in 742727c. The message no longer asserts a specific cause (it now
reads "This placement is not supported. Rewrite the query so the UDF is applied
outside the lambda."), and CheckAnalysis now names the offending UDF - the
first of an unsupported eval type if any, else the first Python UDF in the
lambdas.
##########
python/pyspark/worker.py:
##########
@@ -2989,6 +2991,128 @@ def func(split_index: int, data:
Iterator[pa.RecordBatch]) -> Iterator[pa.Record
# 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 func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ # Flatten all array columns to element lists and convert 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)
+
+ # Extract shape (lengths per row, null mask) from first column.
+ shape = input_batch.column(0)
+ lengths = pc.list_value_length(shape).to_pylist()
+ total_elements = sum(n for n in lengths if n is not None)
+
+ # Build offsets to re-nest flat list back to array<R> with
input shape.
+ # Null rows stay null; they consume no offsets. Preserve
int32/int64 width.
+ 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_())
+
+ # Evaluate all UDFs once over the flattened batch.
+ output_arrays = []
+ for wrapped_func, offsets_meta, arrow_element_type,
result_conv in udf_infos:
+ rows = (
+ [() for _ in range(total_elements)]
+ if not offsets_meta
+ else list(zip(*[columns[o] for o in offsets_meta]))
Review Comment:
Fixed in 742727c. The worker now streams `zip(*[columns[o] for o in
offsets_meta])` directly instead of materializing a batch-sized list.
##########
python/pyspark/sql/tests/test_udf_in_higher_order_function.py:
##########
@@ -0,0 +1,514 @@
+#
+# 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_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)))` must still yield
one result per
+ # element rather than a single value.
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],)], "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],), ([],)],
+ )
+
+ 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):
+ # A comparator cannot be evaluated pairwise, but the UDF applied per
element is a sort key
Review Comment:
Fixed in 742727c. Reworded so it no longer makes the absolute claim; it now
notes the per-element key path and points to the pairwise test for the
both-elements case.
##########
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
Review Comment:
Fixed in 742727c. The comment now describes it as row-shaped from the plan's
point of view: one array column in, one array column out per row.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,547 @@
+/*
+ * 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
+ * never names a concrete class: it reads arguments, lambdas and parameter
roles off the
Review Comment:
Fixed in 742727c. Reworded to say the path names a concrete class only where
a shape cannot be inferred otherwise (ArraySort's comparator, and the
result-type traits telling ArrayFilter from ArrayTransform).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
