zhengruifeng commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3734762318
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala:
##########
@@ -89,6 +89,9 @@ class SparkOptimizer(
ExtractPythonUDFFromAggregate,
// This must be executed after `ExtractPythonUDFFromAggregate` and
before `ExtractPythonUDFs`.
ExtractGroupingPythonUDFFromAggregate,
+ // Lifts Python UDFs out of higher-order function lambdas. Must run
before
+ // `ExtractPythonUDFs`, which then extracts the lifted UDF as an
ordinary top-level UDF.
+ ExtractPythonUDFFromLambda,
Review Comment:
Make this prerequisite structural instead of relying on adjacency in a
`Once` batch. CheckAnalysis now accepts plans that become executable only after
this rewrite, so a future reorder before `ExtractPythonUDFs` can leave an
accepted UDF inside a lambda. Invoking the lambda rewrite from the extraction
rule, or combining both behind one non-excludable rule, would enforce the
dependency.
##########
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)
Review Comment:
Derive offsets and the null mask per UDF from that UDF's first argument
offset. `ExtractPythonUDFs` can fuse `transform(a, f)` and `transform(b, g)`
into one eval-type-102 batch even when `a` and `b` have different lengths.
Reusing column 0's shape makes the second UDF fail row-count validation or
re-nest results against the wrong offsets. Please add a two-array regression
test with different lengths and null layouts.
##########
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
+ * [[HigherOrderFunction]] API and rebuilds with `withNewChildren` (children
are `arguments` then
+ * `functions`). Only a pairwise `array_sort` comparator needs a separate
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 so the comparator can index the matrix,
sort by
+ // `matrix[a.idx][b.idx]`, then drop the positions again. `element_at` is
1-based.
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val indexed = ArraysZip(
+ Seq(argument, ArrayTransform(argument, LambdaFunction(posIdx,
Seq(posElem, posIdx)))),
+ Seq(Literal(s"${carrierElementPrefix}0"), Literal(carrierIndexField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // The matrix, as n rows of n taken from the flat results.
+ val rowElem = NamedLambdaVariable("x", elementType, containsNull)
+ val rowIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val matrix = ArrayTransform(
+ argument,
+ LambdaFunction(
+ Slice(flatCells, Add(Multiply(rowIdx, n), Literal(1)), n),
+ Seq(rowElem, rowIdx)))
+
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def indexOf(v: NamedLambdaVariable): Expression =
+ Add(GetStructField(v, 1, Some(carrierIndexField)), Literal(1))
+ val comparison = ElementAt(
+ ElementAt(matrix, indexOf(cmpLeft), None, failOnError = false),
Review Comment:
Index the lifted flat result arrays directly from the comparator. `matrix`
is an `ArrayTransform` expression inside this lambda, and `ArraySort` evaluates
the lambda for every comparison, so each comparison rebuilds and slices the
full n-by-n matrix. That turns the documented O(n^2) pairwise preparation into
O(n^2) work per sort comparison. Keep the JVM wrapper body in the comparator
and make its UDF reads direct flat-array lookups.
##########
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
+ * [[HigherOrderFunction]] API and rebuilds with `withNewChildren` (children
are `arguments` then
+ * `functions`). Only a pairwise `array_sort` comparator needs a separate
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 so the comparator can index the matrix,
sort by
+ // `matrix[a.idx][b.idx]`, then drop the positions again. `element_at` is
1-based.
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val indexed = ArraysZip(
+ Seq(argument, ArrayTransform(argument, LambdaFunction(posIdx,
Seq(posElem, posIdx)))),
+ Seq(Literal(s"${carrierElementPrefix}0"), Literal(carrierIndexField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // The matrix, as n rows of n taken from the flat results.
+ val rowElem = NamedLambdaVariable("x", elementType, containsNull)
+ val rowIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val matrix = ArrayTransform(
+ argument,
+ LambdaFunction(
+ Slice(flatCells, Add(Multiply(rowIdx, n), Literal(1)), n),
+ Seq(rowElem, rowIdx)))
+
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def indexOf(v: NamedLambdaVariable): Expression =
+ Add(GetStructField(v, 1, Some(carrierIndexField)), Literal(1))
+ val comparison = ElementAt(
+ ElementAt(matrix, indexOf(cmpLeft), None, failOnError = false),
+ indexOf(cmpRight),
+ 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)
+ // `map_filter` keeps whichever pairs survive; `transform_keys`
replaces the keys and
+ // `transform_values` the values, told apart by whether the result key
type is the lambda's.
+ val rebuild: Expression => Expression =
+ if (isFromElements) { (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ } else if (hof.dataType.asInstanceOf[MapType].keyType ==
lambda.dataType) {
Review Comment:
Distinguish `TransformKeys` from `TransformValues` by expression type, not
result type. For `map<string,string>`, a `transform_values` UDF also returns
`string`, so this condition rebuilds the UDF output as new keys and keeps the
old values. Pattern-match the concrete operation and add a
same-key-and-result-type regression test.
##########
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:
Describe this as array columns in and one array column out per row. A
multi-argument scalar UDF receives one aligned array per child, and the worker
flattens every input column.
##########
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
+ * [[HigherOrderFunction]] API and rebuilds with `withNewChildren` (children
are `arguments` then
+ * `functions`). Only a pairwise `array_sort` comparator needs a separate
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 so the comparator can index the matrix,
sort by
+ // `matrix[a.idx][b.idx]`, then drop the positions again. `element_at` is
1-based.
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val indexed = ArraysZip(
+ Seq(argument, ArrayTransform(argument, LambdaFunction(posIdx,
Seq(posElem, posIdx)))),
+ Seq(Literal(s"${carrierElementPrefix}0"), Literal(carrierIndexField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // The matrix, as n rows of n taken from the flat results.
+ val rowElem = NamedLambdaVariable("x", elementType, containsNull)
+ val rowIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val matrix = ArrayTransform(
+ argument,
+ LambdaFunction(
+ Slice(flatCells, Add(Multiply(rowIdx, n), Literal(1)), n),
+ Seq(rowElem, rowIdx)))
+
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def indexOf(v: NamedLambdaVariable): Expression =
+ Add(GetStructField(v, 1, Some(carrierIndexField)), Literal(1))
+ val comparison = ElementAt(
+ ElementAt(matrix, indexOf(cmpLeft), None, failOnError = false),
+ indexOf(cmpRight),
+ 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)
+ // `map_filter` keeps whichever pairs survive; `transform_keys`
replaces the keys and
+ // `transform_values` the values, told apart by whether the result key
type is the lambda's.
+ val rebuild: Expression => Expression =
+ if (isFromElements) { (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ } else if (hof.dataType.asInstanceOf[MapType].keyType ==
lambda.dataType) {
+ (newKeys: Expression) => MapFromArrays(newKeys, values)
+ } else {
+ (newValues: Expression) => MapFromArrays(keys, newValues)
+ }
+ (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])
+ 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, lambdaExprIds)
+
+ // 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
+
+ var arrayResults = Map.empty[Expression, Expression]
+ val liftedArrays = liftableUDFs.map { 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)
+ arrayResults += (udf.canonicalized -> lifted)
+ lifted
+ }
+
+ // 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
+ val udfFieldByCanonical =
liftableUDFs.map(_.canonicalized).zipWithIndex.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 udfFieldByCanonical.contains(udf.canonicalized) =>
+ val ordinal = udfFieldByCanonical(udf.canonicalized)
+ // 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 in `body` that must be lifted, innermost
first.
+ *
+ * Only calls that actually read the lambda's variables need lifting; a UDF
over constants or
+ * outer columns is already valid outside the lambda and is left to
[[ExtractPythonUDFs]].
+ */
+ private def collectLiftableUDFs(
+ body: Expression,
+ lambdaExprIds: Set[ExprId]): 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) &&
+ readsLambdaVariable(udf, lambdaExprIds) =>
+ collected += udf
+ case _ =>
+ }
+ }
+ visit(body)
+ // Deduplicate identical calls so the same UDF is evaluated once per array.
+ val seen = scala.collection.mutable.LinkedHashMap.empty[Expression,
PythonUDF]
+ collected.result().foreach(udf => seen.getOrElseUpdate(udf.canonicalized,
udf))
Review Comment:
Preserve nondeterministic call identity here. `ExtractPythonUDFs`
canonicalizes only deterministic UDFs, but this map collapses every matching
call. As a result, `transform(arr, x => f(x) + f(x))` reuses one value when `f`
is marked nondeterministic. Keep nondeterministic calls distinct by `resultId`,
or exclude them from this rewrite.
##########
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()
Review Comment:
Preserve a non-Arrow transport for UDFs created with `useArrow=False`. The
eligibility predicate accepts ordinary `SQL_BATCHED_UDF` calls, but the rewrite
changes all of them to eval type 102 and reaches this unconditional PyArrow
requirement. Base PySpark installations can run plain scalar UDFs without that
optional package, so add a non-Arrow elementwise eval path and cover it without
pandas/PyArrow installed.
##########
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:
Stream these rows instead of materializing another batch-sized list. The
flattened columns, output results, and converted results are already resident,
so this tuple list adds O(total elements x arity) memory on the long-array
workloads this feature targets. Pass `zip(...)` directly and use
`itertools.repeat((), total_elements)` for the zero-argument case.
##########
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
+ * [[HigherOrderFunction]] API and rebuilds with `withNewChildren` (children
are `arguments` then
+ * `functions`). Only a pairwise `array_sort` comparator needs a separate
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 so the comparator can index the matrix,
sort by
+ // `matrix[a.idx][b.idx]`, then drop the positions again. `element_at` is
1-based.
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val indexed = ArraysZip(
+ Seq(argument, ArrayTransform(argument, LambdaFunction(posIdx,
Seq(posElem, posIdx)))),
+ Seq(Literal(s"${carrierElementPrefix}0"), Literal(carrierIndexField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // The matrix, as n rows of n taken from the flat results.
+ val rowElem = NamedLambdaVariable("x", elementType, containsNull)
+ val rowIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val matrix = ArrayTransform(
+ argument,
+ LambdaFunction(
+ Slice(flatCells, Add(Multiply(rowIdx, n), Literal(1)), n),
+ Seq(rowElem, rowIdx)))
+
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def indexOf(v: NamedLambdaVariable): Expression =
+ Add(GetStructField(v, 1, Some(carrierIndexField)), Literal(1))
+ val comparison = ElementAt(
+ ElementAt(matrix, indexOf(cmpLeft), None, failOnError = false),
+ indexOf(cmpRight),
+ 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)
+ // `map_filter` keeps whichever pairs survive; `transform_keys`
replaces the keys and
+ // `transform_values` the values, told apart by whether the result key
type is the lambda's.
+ val rebuild: Expression => Expression =
+ if (isFromElements) { (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ } else if (hof.dataType.asInstanceOf[MapType].keyType ==
lambda.dataType) {
+ (newKeys: Expression) => MapFromArrays(newKeys, values)
+ } else {
+ (newValues: Expression) => MapFromArrays(keys, newValues)
+ }
+ (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])
+ 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, lambdaExprIds)
+
+ // 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
+
+ var arrayResults = Map.empty[Expression, Expression]
+ val liftedArrays = liftableUDFs.map { 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)
+ arrayResults += (udf.canonicalized -> lifted)
+ lifted
+ }
+
+ // 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
+ val udfFieldByCanonical =
liftableUDFs.map(_.canonicalized).zipWithIndex.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 udfFieldByCanonical.contains(udf.canonicalized) =>
+ val ordinal = udfFieldByCanonical(udf.canonicalized)
+ // 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 in `body` that must be lifted, innermost
first.
+ *
+ * Only calls that actually read the lambda's variables need lifting; a UDF
over constants or
Review Comment:
Lift direct UDF calls even when their arguments do not reference a lambda
variable. Leaving `transform(arr, _ => f(lit(10)))` to ordinary
`ExtractPythonUDFs` runs `f` once per input row, including null and empty
arrays, although the lambda executes zero times there. The existing `overArray`
branch already repeats an independent argument into an aligned array, which
preserves the correct zero-or-N call domain.
##########
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:
Use a cause-specific message here. A top-level pandas UDF is rejected
because its eval type is unsupported, and setting the new config to false
rejects an otherwise rewritable scalar UDF; neither case is a nested lambda or
fold accumulator. Mixed lambdas can also name the first rewritable UDF instead
of the offending call. Select the failed predicate and offending UDF before
constructing the error.
##########
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:
Narrow this to the result-type dispatch. The generic path explicitly checks
`ArraySort` to distinguish comparator parameters and constructs `ArrayFilter`
or `ArrayTransform`, so the whole rewrite is not class-independent.
##########
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
+ * [[HigherOrderFunction]] API and rebuilds with `withNewChildren` (children
are `arguments` then
+ * `functions`). Only a pairwise `array_sort` comparator needs a separate
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 so the comparator can index the matrix,
sort by
+ // `matrix[a.idx][b.idx]`, then drop the positions again. `element_at` is
1-based.
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val indexed = ArraysZip(
+ Seq(argument, ArrayTransform(argument, LambdaFunction(posIdx,
Seq(posElem, posIdx)))),
+ Seq(Literal(s"${carrierElementPrefix}0"), Literal(carrierIndexField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // The matrix, as n rows of n taken from the flat results.
+ val rowElem = NamedLambdaVariable("x", elementType, containsNull)
+ val rowIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val matrix = ArrayTransform(
+ argument,
+ LambdaFunction(
+ Slice(flatCells, Add(Multiply(rowIdx, n), Literal(1)), n),
+ Seq(rowElem, rowIdx)))
+
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def indexOf(v: NamedLambdaVariable): Expression =
+ Add(GetStructField(v, 1, Some(carrierIndexField)), Literal(1))
+ val comparison = ElementAt(
+ ElementAt(matrix, indexOf(cmpLeft), None, failOnError = false),
+ indexOf(cmpRight),
+ 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)
+ // `map_filter` keeps whichever pairs survive; `transform_keys`
replaces the keys and
+ // `transform_values` the values, told apart by whether the result key
type is the lambda's.
+ val rebuild: Expression => Expression =
+ if (isFromElements) { (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ } else if (hof.dataType.asInstanceOf[MapType].keyType ==
lambda.dataType) {
+ (newKeys: Expression) => MapFromArrays(newKeys, values)
+ } else {
+ (newValues: Expression) => MapFromArrays(keys, newValues)
+ }
+ (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])
+ 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, lambdaExprIds)
+
+ // 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
+
+ var arrayResults = Map.empty[Expression, Expression]
+ val liftedArrays = liftableUDFs.map { 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)
+ arrayResults += (udf.canonicalized -> lifted)
+ lifted
+ }
+
+ // 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
+ val udfFieldByCanonical =
liftableUDFs.map(_.canonicalized).zipWithIndex.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 udfFieldByCanonical.contains(udf.canonicalized) =>
+ val ordinal = udfFieldByCanonical(udf.canonicalized)
+ // 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 in `body` that must be lifted, innermost
first.
+ *
+ * Only calls that actually read the lambda's variables need lifting; a UDF
over constants or
+ * outer columns is already valid outside the lambda and is left to
[[ExtractPythonUDFs]].
+ */
+ private def collectLiftableUDFs(
+ body: Expression,
+ lambdaExprIds: Set[ExprId]): 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
Review Comment:
Reject calls whose original evaluation is conditional until the rewrite can
preserve their mask and schedule. This traversal hoists UDFs from unselected
`when` branches, while `exists`/`forall` skip a decisive suffix and
`array_sort` requests only selected pairs. A deterministic UDF that throws only
on one skipped value therefore changes a successful query into an error.
Focused tests should cover each skipped-call boundary.
##########
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:
State that Python cannot run on demand inside the JVM comparator. The
pairwise test below shows that a two-argument UDF is evaluated over the
Cartesian product, so the current absolute claim contradicts the supported path.
--
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]
