You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "HyukjinKwon (via GitHub)" <gi...@apache.org> on 2026/08/06 00:44:38 UTC
[PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
HyukjinKwon opened a new pull request, #57804:
URL: https://github.com/apache/spark/pull/57804
### What changes were proposed in this pull request?
This PR lets a scalar Python UDF be used inside the lambda of a higher-order function, so that
```python
df.select(F.transform("values", lambda x: plus_one(x)))
```
works with a plain `pyspark.sql.functions.udf` instead of failing with
`[UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]`.
**Why it does not work today.** A `PythonUDF` is evaluated by a separate physical operator
(`ArrowEvalPython`), which `ExtractPythonUDFs` pulls out of the enclosing operator. A lambda's
`NamedLambdaVariable`s only exist while the higher-order function is iterating, so an extracted
operator cannot see them: the UDF can neither stay inside the lambda nor be lifted out by the
existing extraction rule. `CheckAnalysis` therefore rejected it outright.
**The rewrite.** Do not evaluate the UDF per element inside the lambda. Apply it once to the
*whole array*, outside every lambda, and let the lambda read the precomputed result positionally.
A new optimizer rule, `ExtractPythonUDFFromLambda`, does this:
```
-- before (rejected)
transform(values, x -> plus_one(x))
-- after (legal: the PythonUDF is outside every LambdaFunction)
transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
```
It runs in the `Extract Python UDFs` batch *before* `ExtractPythonUDFs`, which then extracts the
lifted UDF as an ordinary top-level `PythonUDF`, so no new physical operator is needed. It is
registered in `nonExcludableRules`, since a plan that only works because of this rewrite must not
be silently broken by `spark.sql.optimizer.excludedRules`.
**Where the array-at-a-time behaviour lives.** `_wrap_function` pickles `(func, returnType)` as
opaque bytes, so the JVM can never rewrap the user's function; the wrapper has to be worker-side.
A new eval type, `SQL_ARROW_ELEMENTWISE_UDF`, does it with Arrow: the worker flattens the incoming
list column, calls the user function once over the concatenated elements of the **whole batch**,
then re-nests the results using the input's offsets plus an explicit validity mask.
This is deliberately stronger than a per-row Python loop. The number of Python-level calls is
proportional to the number of elements, but the Arrow boundary is crossed once per batch rather
than once per row, which the design notes measured as the dominant cost. It also keeps one row in
and one row out: no `explode`, no shuffle.
Once the UDF result is an ordinary column, everything the lambda does around it is ordinary JVM
work, so arithmetic, `when`, casts, the element index, several UDFs in one lambda, nested UDF
calls (`f(g(x))`) and UDF arguments that are expressions over the element all follow without
special cases.
**Supported:** `transform`, `filter`, `exists`, `forall`, and `aggregate` when the UDF applies to
the element.
**Still rejected** (`CheckAnalysis` and the rule share one predicate, so analysis accepts exactly
what the optimizer rewrites):
- a UDF reading `aggregate`'s accumulator, which is sequential (step *n* depends on *n-1*);
- a UDF in `aggregate`'s `finish` — it runs once on the final accumulator, and lifting it out
would call it on the null a fold over a null array produces, where native Spark does not
evaluate `finish` at all;
- a pandas UDF, which receives a `Series` rather than one value per call;
- a nested higher-order function, e.g. `transform(arr, i -> transform(i, x -> udf(x)))`: the inner
array only exists while the outer function iterates, so a UDF lifted onto it would still sit
inside the outer lambda. (Note the reference prototype does not support this shape either — its
"nested" support is nested *UDF calls*, which this PR does support. Extending to nested
higher-order functions is possible by flattening and re-nesting one level deeper, and is left
as follow-up.)
- the map family (`transform_keys`, `transform_values`, `map_filter`, `map_zip_with`) and
`zip_with`, left as follow-up.
The rewrite can be turned off with
`spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled=false`, restoring the previous error.
This builds on the design notes derived from the
[elementwise-udf](https://github.com/HyukjinKwon/elementwise-udf) prototype (Apache-2.0), which
established which rewrites are correct and which are not worth doing.
### Why are the changes needed?
`transform(col, lambda x: my_udf(x))` is a natural thing to write and is a long-standing gap
(SPARK-27052). Today users must either avoid Python UDFs in lambdas entirely or fall back to
`explode` + regroup, which the prototype measured at 2-40x slower and which OOMs on long arrays
because carrying the source array alongside `posexplode` duplicates it per element.
### Does this PR introduce _any_ user-facing change?
Yes. A query that previously failed analysis now runs:
```python
>>> plus_one = udf(lambda x: x + 1, "int")
>>> df.select(F.transform("values", lambda x: plus_one(x))).show()
# before: AnalysisException [UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF]
# after: [2, 3, 4]
```
Shapes that cannot be rewritten keep failing with the same error condition as before. No existing
successful query changes behaviour.
### How was this patch tested?
New end-to-end suite `pyspark.sql.tests.test_udf_in_higher_order_function` (29 tests, all
passing), which asserts results against the equivalent native expression wherever one exists, so a
rewrite that runs but computes the wrong thing fails rather than passing quietly. It covers each
supported higher-order function; composition around the UDF result; the index parameter; several
and nested UDF calls; UDF arguments that are expressions over the element; broadcast outer-column
and constant-only arguments (the constant case must still yield one result per element); element
and return types including string, double and `array<int>`; null arrays, null elements, UDFs
returning null, empty arrays, all-null rows and an empty frame; long arrays over many rows to
exercise batching; integration with joins, caching and `groupBy`; mixing with an ordinary Python
UDF; that a lambda with no Python UDF is left untouched; and the negative cases above.
New plan-shape suite `ExtractPythonUDFFromLambdaSuite` (13 tests, all passing) asserts that no
`PythonUDF` remains inside a `LambdaFunction`, that the lifted UDF is an element-wise UDF over an
array, that duplicate calls are evaluated once, that the rule is inert without a UDF, and that it
is not excludable.
`PythonUDFSuite`'s `SPARK-48706` negative test asserted the old behaviour for exactly the case now
supported; it is updated to assert the result and a still-unsupported shape. `PythonUDFSuite`
(16 tests), the `execution.python` suites, `DataFrameFunctionsSuite`, `pyspark.sql.tests.test_udf`
and `arrow.test_arrow_python_udf` pass.
The Arrow flatten/re-nest algorithm was also validated standalone against null arrays, null
elements, empty arrays, sliced arrays, large lists and broadcast arguments before being wired in.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5234776854
Thanks @dongjoon-hyun for the thorough review. Addressed all points in 95ec08d:
- **`f(g(x) + 1)` / `f(-g(x))`** (correctness): nested lifted UDFs inside a composite argument are now substituted before lifting, so no raw UDF is left inside a lambda (avoiding the SPARK-48706 mode).
- **Pairwise `array_sort` comparator** (perf): `flatCells` moved into the sort argument (carried via `array_repeat`), so each comparison is O(1) instead of rebuilding the cells.
- **Predicate alignment**: `isRewritableShape` now requires a result-type marker trait; `isElementwiseRewritableUDF` rejects named-argument and UDT UDFs.
- **Key-form comparator dedup**: Python now runs n times, not 2n.
- Removed the unused `lambdaExprIds` param, stray blank lines, and stale test comment.
- **Connect parity test**: added `pyspark.sql.tests.connect.test_parity_udf_in_higher_order_function` and registered it in `modules.py`, per your suggestion.
Verified: `ExtractPythonUDFFromLambdaSuite` (21 tests) and the end-to-end `pyspark.sql.tests.test_udf_in_higher_order_function` suite both pass locally.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5212706429
cc @cloud-fan @zhengruifeng @Yicong-Huang
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3745890637
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
+ flatCells,
+ Add(Add(Multiply(idxOf(cmpLeft), n), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0)
+ }
+
+
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type: `transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array, indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption, Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps its own class (via
+ // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
+ * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ */
+ private def liftableHof(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 && (hof.functions.head match {
+ case LambdaFunction(body, args, _) =>
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
+ 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 += (liftKey(udf) -> 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 udfFieldByKey = liftableUDFs.map(liftKey).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 udfFieldByKey.contains(liftKey(udf)) =>
+ val ordinal = udfFieldByKey(liftKey(udf))
+ // A UDF over a comparator's right-hand element must read that element's key, so the
+ // struct field is read through whichever bound variable the call's own arguments used.
+ val side = udf.collectFirst {
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => v
+ }
+ side match {
+ case Some(v) => readerFor(v, Some(ordinal))
+ case None =>
+ GetStructField(boundVar, arguments.length + ordinal,
+ Some(s"$carrierUDFFieldPrefix$ordinal"))
+ }
+ case v: NamedLambdaVariable if fieldOfVar.contains(v.exprId) => readerFor(v, None)
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) =>
+ // A comparator's right-hand element itself, read through its own bound variable.
+ GetStructField(extraVarOf(v.exprId), 0, Some(carrierNames.head))
+ }
+
+ Carrier(carrier, rewrittenBody, boundVar, extraBoundVars)
+ }
+
+ /**
+ * Collects the Python UDF calls directly in `body` that must be lifted, innermost first.
+ *
+ * Every rewritable UDF directly in the lambda body is lifted, even one whose arguments do not
+ * read the lambda variable (`transform(arr, _ -> f(lit(10)))`). It cannot stay inside the lambda,
+ * and lifting it - `overArray` repeats a constant/outer-column argument into an aligned array -
+ * gives it the lambda's own call domain: once per element, and zero times for a null or empty
+ * array. Leaving it to [[ExtractPythonUDFs]] would instead call it once per row, including rows
+ * whose array is null or empty where the lambda never runs.
+ */
+ private def collectLiftableUDFs(
+ body: Expression,
+ 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) =>
+ collected += udf
+ case _ =>
+ }
+ }
+ visit(body)
+ // Deduplicate identical calls so the same UDF is evaluated once per array. Nondeterministic
+ // calls are kept distinct (see `liftKey`).
+ val seen = scala.collection.mutable.LinkedHashMap.empty[Expression, PythonUDF]
+ collected.result().foreach(udf => seen.getOrElseUpdate(liftKey(udf), udf))
+ seen.values.toSeq
+ }
+
+ private def readsLambdaVariable(e: Expression, lambdaExprIds: Set[ExprId]): Boolean =
+ e.exists {
+ case v: NamedLambdaVariable => lambdaExprIds.contains(v.exprId)
+ case _ => false
+ }
+
+ /**
+ * The key that decides whether two UDF calls are "the same call" for lifting. A deterministic
+ * call is deduplicated by canonical form, so an identical call is evaluated once. A
+ * nondeterministic call must stay distinct - `transform(arr, x -> f(x) + f(x))` calls `f` twice
+ * and each call may return a different value - so it is keyed by its own `resultId`-bearing node
+ * (`canonicalized` erases `resultId`, which would collapse the two calls into one).
+ */
+ private def liftKey(udf: PythonUDF): Expression =
+ if (udf.udfDeterministic) udf.canonicalized else udf
+
+ /**
+ * Turns a UDF argument expression, written in terms of single elements, into the equivalent
+ * `array<T>` aligned with the iterated array, so the worker can flatten every argument uniformly.
+ *
+ * - a lambda variable becomes the array it stands for;
+ * - an already-lifted UDF call (a nested call `f(g(x))`) becomes its array result;
+ * - an expression independent of the lambda is repeated into an aligned array with a native
+ * `transform`, rather than passed through as a scalar to broadcast;
+ * - anything else is an expression over the elements, computed for every element by a native
+ * `transform` that stays inside the JVM.
+ */
+ private def overArray(
+ child: Expression,
+ firstArgument: Expression,
+ arrayOfVar: Map[ExprId, Expression],
+ lambdaExprIds: Set[ExprId],
+ arrayResults: Map[Expression, Expression]): Expression = child match {
+ case v: NamedLambdaVariable if arrayOfVar.contains(v.exprId) => arrayOfVar(v.exprId)
+ case udf: PythonUDF if arrayResults.contains(liftKey(udf)) =>
+ arrayResults(liftKey(udf))
+ case e if !readsLambdaVariable(e, lambdaExprIds) =>
+ // Independent of the element (an outer column or constant): repeat it into an array aligned
+ // with the iterated array, so every UDF argument is a single-level array the worker flattens
+ // the same way. `transform(arr, _ -> e)` keeps the value constant while matching the shape.
+ val arrType = firstArgument.dataType.asInstanceOf[ArrayType]
+ val v = NamedLambdaVariable("x", arrType.elementType, arrType.containsNull)
+ ArrayTransform(firstArgument, LambdaFunction(e, Seq(v)))
+ case e =>
Review Comment:
Fixed in 95ec08d. `overArray` now substitutes an already-lifted UDF nested inside a composite argument (`f(g(x) + 1)`, `f(-g(x))`), not only a whole-child match: each such call is replaced with a synthetic variable standing for that call's aligned array result, which is zipped in like an element array. This keeps a raw `g` over a `NamedLambdaVariable` from being left inside the generated transform (the SPARK-48706 mode). Added plan-shape ("a UDF nested inside a composite argument is lifted, not left inside a lambda") and end-to-end (`test_nested_udf_inside_composite_argument`, covering both `f(g(x) + 1)` and `f(-g(x))`) regression tests.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,73 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
+ case udf: PythonUDF =>
+ udf.evalType == PythonEvalType.SQL_BATCHED_UDF ||
+ udf.evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF
+ case _ => false
+ }
+
+ /**
+ * Whether every Python UDF in `hof`'s lambdas can be lifted out by `ExtractPythonUDFFromLambda`.
+ * Used by `CheckAnalysis` to decide whether to reject the plan. Three shapes cannot be rewritten:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so it sees earlier steps' outputs;
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+ def canRewritePythonUDFInLambda(hof: HigherOrderFunction): Boolean = {
Review Comment:
Fixed in 95ec08d. `isRewritableShape` now also requires one of the `ResultTypeFromArgument` / `ResultTypeFromFunction` marker traits, so `CheckAnalysis` accepts exactly what the rule dispatches on. A future single-lambda HOF missing both traits is now rejected at analysis rather than leaving the UDF in the lambda at runtime.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
Review Comment:
Good catch - fixed in 95ec08d. `flatCells` is now built in the sort's *argument* and carried into every element via `array_repeat` (n references to the one computed array, so O(n^2) memory), and the comparator reads it with a `GetStructField` + `element_at`, both O(1). The `arrays_zip`/`transform`/`flatten`/`array_repeat` that build the cells no longer sit inside the comparator body, so `ArraySort` no longer rebuilds them per comparison.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5239349388
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/ea911d076840c8577381da422e5801df44c467fe
- merged into branch-4.x https://github.com/apache/spark/commit/7cb3ba5fc140927c8fce238b5d52aa4f8d7f9820
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3744954547
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
Review Comment:
I don't think this is O(1) per comparison yet. `ExtractPythonUDFs` replaces only the `PythonUDF` node with an attribute; the wrappers around it stay inside this comparator lambda: `ArraysZip(lefts, rights, attr)`, the `ArrayTransform` running the JVM part of the comparator body, and the `Flatten` / `ArrayRepeat` building `lefts` / `rights`. All of those are re-evaluated on every comparison, i.e. O(n^2) JVM work per comparison and roughly O(n^3 log n) per sort on top of the O(n^2) Python calls.
Worth benchmarking with a moderately sized array; if confirmed, the JVM-side body needs to be hoisted out of the comparator too (e.g. precompute `flatCells` alongside the carrier so the comparator only does the `element_at` lookup).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
+ flatCells,
+ Add(Add(Multiply(idxOf(cmpLeft), n), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0)
+ }
+
+
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type: `transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array, indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption, Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps its own class (via
+ // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
+ * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ */
+ private def liftableHof(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 && (hof.functions.head match {
+ case LambdaFunction(body, args, _) =>
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
+ 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 += (liftKey(udf) -> 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 udfFieldByKey = liftableUDFs.map(liftKey).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 udfFieldByKey.contains(liftKey(udf)) =>
+ val ordinal = udfFieldByKey(liftKey(udf))
+ // A UDF over a comparator's right-hand element must read that element's key, so the
+ // struct field is read through whichever bound variable the call's own arguments used.
+ val side = udf.collectFirst {
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => v
+ }
+ side match {
+ case Some(v) => readerFor(v, Some(ordinal))
+ case None =>
+ GetStructField(boundVar, arguments.length + ordinal,
+ Some(s"$carrierUDFFieldPrefix$ordinal"))
+ }
+ case v: NamedLambdaVariable if fieldOfVar.contains(v.exprId) => readerFor(v, None)
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) =>
+ // A comparator's right-hand element itself, read through its own bound variable.
+ GetStructField(extraVarOf(v.exprId), 0, Some(carrierNames.head))
+ }
+
+ Carrier(carrier, rewrittenBody, boundVar, extraBoundVars)
+ }
+
+ /**
+ * Collects the Python UDF calls directly in `body` that must be lifted, innermost first.
+ *
+ * Every rewritable UDF directly in the lambda body is lifted, even one whose arguments do not
+ * read the lambda variable (`transform(arr, _ -> f(lit(10)))`). It cannot stay inside the lambda,
+ * and lifting it - `overArray` repeats a constant/outer-column argument into an aligned array -
+ * gives it the lambda's own call domain: once per element, and zero times for a null or empty
+ * array. Leaving it to [[ExtractPythonUDFs]] would instead call it once per row, including rows
+ * whose array is null or empty where the lambda never runs.
+ */
+ private def collectLiftableUDFs(
+ body: Expression,
+ 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) =>
+ collected += udf
+ case _ =>
+ }
+ }
+ visit(body)
+ // Deduplicate identical calls so the same UDF is evaluated once per array. Nondeterministic
+ // calls are kept distinct (see `liftKey`).
+ val seen = scala.collection.mutable.LinkedHashMap.empty[Expression, PythonUDF]
+ collected.result().foreach(udf => seen.getOrElseUpdate(liftKey(udf), udf))
+ seen.values.toSeq
+ }
+
+ private def readsLambdaVariable(e: Expression, lambdaExprIds: Set[ExprId]): Boolean =
+ e.exists {
+ case v: NamedLambdaVariable => lambdaExprIds.contains(v.exprId)
+ case _ => false
+ }
+
+ /**
+ * The key that decides whether two UDF calls are "the same call" for lifting. A deterministic
+ * call is deduplicated by canonical form, so an identical call is evaluated once. A
+ * nondeterministic call must stay distinct - `transform(arr, x -> f(x) + f(x))` calls `f` twice
+ * and each call may return a different value - so it is keyed by its own `resultId`-bearing node
+ * (`canonicalized` erases `resultId`, which would collapse the two calls into one).
+ */
+ private def liftKey(udf: PythonUDF): Expression =
+ if (udf.udfDeterministic) udf.canonicalized else udf
+
+ /**
+ * Turns a UDF argument expression, written in terms of single elements, into the equivalent
+ * `array<T>` aligned with the iterated array, so the worker can flatten every argument uniformly.
+ *
+ * - a lambda variable becomes the array it stands for;
+ * - an already-lifted UDF call (a nested call `f(g(x))`) becomes its array result;
+ * - an expression independent of the lambda is repeated into an aligned array with a native
+ * `transform`, rather than passed through as a scalar to broadcast;
+ * - anything else is an expression over the elements, computed for every element by a native
+ * `transform` that stays inside the JVM.
+ */
+ private def overArray(
+ child: Expression,
+ firstArgument: Expression,
+ arrayOfVar: Map[ExprId, Expression],
+ lambdaExprIds: Set[ExprId],
+ arrayResults: Map[Expression, Expression]): Expression = child match {
+ case v: NamedLambdaVariable if arrayOfVar.contains(v.exprId) => arrayOfVar(v.exprId)
+ case udf: PythonUDF if arrayResults.contains(liftKey(udf)) =>
+ arrayResults(liftKey(udf))
+ case e if !readsLambdaVariable(e, lambdaExprIds) =>
+ // Independent of the element (an outer column or constant): repeat it into an array aligned
+ // with the iterated array, so every UDF argument is a single-level array the worker flattens
+ // the same way. `transform(arr, _ -> e)` keeps the value constant while matching the shape.
+ val arrType = firstArgument.dataType.asInstanceOf[ArrayType]
+ val v = NamedLambdaVariable("x", arrType.elementType, arrType.containsNull)
+ ArrayTransform(firstArgument, LambdaFunction(e, Seq(v)))
+ case e =>
Review Comment:
A lifted UDF nested inside a composite argument is left inside the generated lambda. This branch substitutes only `NamedLambdaVariable`s, so for `f(g(x) + 1)` the child `g(x) + 1` becomes `ArrayTransform(arr, v -> g(v) + 1)` with the raw `g` (SQL_BATCHED_UDF) still inside a lambda. `g` was already lifted into `arrayResults`, but its occurrence inside the composite expression is never replaced (only a whole-child match is, in the case above).
The optimized plan then violates this rule's own invariant, and `ExtractPythonUDFs` will happily extract that `g(v)` whose child is a `NamedLambdaVariable`, reintroducing exactly the failure mode SPARK-48706 guarded against (wrong results or a runtime error). `CheckAnalysis` accepts this shape via the shared predicate, so nothing rejects it.
Repro to check:
```python
df.select(sf.transform("values", lambda x: times_ten(plus_one(x) + 1)))
```
Either substitute lifted-UDF occurrences here as well (zip their result arrays into the transform and read them as struct fields), or reject UDFs nested in composite arguments in the shared predicate so it stays an analysis error. Either way please add a regression test for `f(g(x) + 1)` (and e.g. `f(-g(x))`).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
+ flatCells,
+ Add(Add(Multiply(idxOf(cmpLeft), n), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0)
+ }
+
+
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type: `transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array, indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption, Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps its own class (via
+ // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
+ * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ */
+ private def liftableHof(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 && (hof.functions.head match {
+ case LambdaFunction(body, args, _) =>
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
+ 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
Review Comment:
Minor: in the key-form comparator `(a, b) -> when(udf(a) < udf(b), ...)`, `udf(a)` and `udf(b)` canonicalize differently (different variable exprIds) but both map to the same array, so two identical lifted UDFs are produced and the Python function runs 2n times instead of n. Deduplicating by the post-substitution array arguments would halve that.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala:
##########
@@ -0,0 +1,279 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.{AnalysisException, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{LambdaFunction, PythonUDF}
+import org.apache.spark.sql.catalyst.plans.logical.ArrowEvalPython
+import org.apache.spark.sql.functions.{array_sort, col, forall, lit, map_filter, map_zip_with,
+ transform, transform_keys, transform_values, when, zip_with}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, StringType}
+
+/**
+ * Plan-shape tests for [[ExtractPythonUDFFromLambda]].
+ *
+ * These assert the structure the rewrite produces - that no `PythonUDF` is left inside a
+ * `LambdaFunction`, that the lifted UDF is an element-wise UDF over an array, and that plans
+ * without a UDF in a lambda are untouched. End-to-end result correctness is covered by
+ * `pyspark.sql.tests.test_udf_in_higher_order_function`, which needs a real Python worker.
+ */
+class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession {
+ import testImplicits._
+
+ private val pythonUDF = new MyDummyPythonUDF
+ private val scalarPandasUDF = new MyDummyScalarPandasUDF
+ // Used where a UDF call must receive two arguments, e.g. a pairwise comparator.
+ private val pythonUDF2 = new MyDummyPythonUDF
+ private val nondeterministicUDF = new MyDummyNondeterministicPythonUDF
+
+ private def arrayDF = Seq(Seq(1, 2, 3)).toDF("values")
+
+ /** All `PythonUDF`s that remain inside a lambda in the optimized plan. */
+ private def udfsInsideLambda(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] = {
+ df.queryExecution.optimizedPlan.expressions.flatMap { e =>
+ e.collect { case l: LambdaFunction => l }.flatMap { l =>
+ l.collect { case u: PythonUDF => u }
+ }
+ }
+ }
+
+ private def liftedUDFs(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] = {
+ df.queryExecution.optimizedPlan.collect {
+ case a: ArrowEvalPython => a.udfs
+ }.flatten
+ }
+
+ test("transform: the UDF is lifted out of the lambda as an element-wise array UDF") {
+ val df = arrayDF.select(transform(col("values"), x => pythonUDF(x)).as("r"))
+
+ // The whole point of the rewrite: nothing Python-shaped is left inside a lambda.
+ assert(udfsInsideLambda(df).isEmpty)
+
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ assert(lifted.head.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)
+ // The lifted UDF takes the array and returns an array of the original return type.
+ assert(lifted.head.dataType == ArrayType(pythonUDF.dataType, containsNull = true))
+ assert(lifted.head.children.size == 1)
+ assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType])
+ }
+
+ test("filter/exists/forall: the UDF is lifted out of the lambda") {
+ val exprs = Seq(
+ org.apache.spark.sql.functions.filter(col("values"), x => pythonUDF(x)),
+ org.apache.spark.sql.functions.exists(col("values"), x => pythonUDF(x)),
+ org.apache.spark.sql.functions.forall(col("values"), x => pythonUDF(x)))
+
+ exprs.foreach { e =>
+ val df = arrayDF.select(e.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF still inside a lambda for $e")
+ assert(liftedUDFs(df).forall(_.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF))
+ }
+ }
+
+ test("aggregate: a UDF anywhere in the fold is rejected") {
+ // The fold is sequential, so no array can precompute the values the UDF sees. A UDF in `merge`
+ // or in `finish` therefore has no rewrite and analysis must fail.
+ val cases = Seq(
+ "merge on element" ->
+ org.apache.spark.sql.functions.aggregate(
+ col("values"), lit(false), (acc, x) => acc || pythonUDF(x)),
+ "merge on accumulator" ->
+ org.apache.spark.sql.functions.aggregate(
+ col("values"), lit(false), (acc, x) => pythonUDF(acc) || x.cast("boolean")),
+ "finish" ->
+ org.apache.spark.sql.functions.aggregate(
+ col("values"), lit(0), (acc, x) => acc + x, acc => pythonUDF(acc)))
+ cases.foreach { case (name, expr) =>
+ val e = intercept[AnalysisException] {
+ arrayDF.select(expr.as("r")).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF",
+ s"expected aggregate with a UDF in $name to be rejected")
+ }
+ }
+
+ test("several UDFs and nested UDFs in one lambda are all lifted") {
+ val several = arrayDF.select(
+ transform(col("values"), x => pythonUDF(x) || pythonUDF(x + lit(1))).as("r"))
+ assert(udfsInsideLambda(several).isEmpty)
+ // Two distinct calls, so two lifted array UDFs.
+ assert(liftedUDFs(several).size == 2)
+
+ // The same deterministic call twice must be evaluated once.
+ val duplicated = arrayDF.select(
+ transform(col("values"), x => pythonUDF(x) || pythonUDF(x)).as("r"))
+ assert(udfsInsideLambda(duplicated).isEmpty)
+ assert(liftedUDFs(duplicated).size == 1)
+ }
+
+ test("identical nondeterministic calls are lifted distinctly, not deduplicated") {
+ // Deduplicating `f(x)` with `f(x)` would collapse two independent draws into one; a
+ // nondeterministic UDF must keep each call, so both are lifted.
+ val df = arrayDF.select(
+ transform(col("values"), x => nondeterministicUDF(x) || nondeterministicUDF(x)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ assert(liftedUDFs(df).size == 2)
+ }
+
+ test("transform_values whose result type equals the key type replaces values, not keys") {
+ // SPARK-27052: dispatch is by concrete function, not result type. For map<string, string>,
+ // transform_values' lambda also returns string, which must not be treated as new keys.
+ val maps = Seq(Map("a" -> "x")).toDF("m")
+ val df = maps.select(transform_values(col("m"), (k, v) => pythonUDF(v).cast("string")).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ // Keys are preserved: the map still has the original key type and no new-key projection.
+ val mapType = df.queryExecution.analyzed.schema.head.dataType.asInstanceOf[MapType]
+ assert(mapType.keyType == StringType)
+ }
+
+ test("a UDF inside a nested higher-order function's lambda is rejected") {
+ // The inner array is the outer lambda's variable, not a real column, so the UDF cannot be
+ // lifted onto an array outside the lambda. This must fail analysis rather than be rewritten.
+ val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
+ val e = intercept[AnalysisException] {
+ df.select(transform(col("values"), inner =>
+ transform(inner, x => pythonUDF(x))).as("r")).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("a UDF on the outer element of a nested array is lifted") {
+ // Here the UDF applies to the outer array's element, which is a real column (the element just
+ // happens to be an array), so it is lifted onto that column like any other single-level array.
+ val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
+ .select(transform(col("values"), inner => pythonUDF(inner)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+
+ test("a lambda with no Python UDF is left unchanged") {
+ val df = arrayDF.select(transform(col("values"), x => x + lit(1)).as("r"))
+ val analyzed = df.queryExecution.analyzed
+ val optimized = df.queryExecution.optimizedPlan
+ // The rule must be inert: no eval-python node is introduced.
+ assert(optimized.collect { case a: ArrowEvalPython => a }.isEmpty)
+ assert(!optimized.toString.contains("pythonUDF"))
+ assert(analyzed.expressions.flatMap(_.collect { case u: PythonUDF => u }).isEmpty)
+ }
+
+ test("every mapping higher-order function is rewritten") {
+ // One assertion per function, so that a shape regressing to "UDF left inside a lambda" is
+ // caught here rather than only by the end-to-end Python suite. `aggregate` / `reduce` are not
+ // here: they fold rather than map, so a UDF in them is rejected (see the aggregate test above).
+ val arrays = Seq((Seq(1, 2), Seq(3, 4))).toDF("l", "r")
+ val maps = Seq((Map("a" -> 1), Map("a" -> 2))).toDF("l", "r")
+
+ val arrayCases = Seq(
+ "transform" -> transform(col("l"), x => pythonUDF(x)),
+ "transform with index" -> transform(col("l"), (x, i) => pythonUDF(x) || i > 0),
+ "filter" -> org.apache.spark.sql.functions.filter(col("l"), x => pythonUDF(x)),
+ "exists" -> org.apache.spark.sql.functions.exists(col("l"), x => pythonUDF(x)),
+ "forall" -> forall(col("l"), x => pythonUDF(x)),
+ "zip_with" -> zip_with(col("l"), col("r"), (a, b) => pythonUDF(a) || pythonUDF(b)),
+ "array_sort" -> array_sort(col("l"),
+ (a, b) => when(pythonUDF(a) === pythonUDF(b), lit(0)).otherwise(lit(1))))
+ arrayCases.foreach { case (name, expr) =>
+ val df = arrays.select(expr.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for $name")
+ }
+
+ val mapCases = Seq(
+ "transform_keys" -> transform_keys(col("l"), (k, v) => pythonUDF(k).cast("string")),
+ "transform_values" -> transform_values(col("l"), (k, v) => pythonUDF(v)),
+ "map_filter" -> map_filter(col("l"), (k, v) => pythonUDF(v)),
+ "map_zip_with" -> map_zip_with(col("l"), col("r"), (k, a, b) => pythonUDF(a) || pythonUDF(b)))
+ mapCases.foreach { case (name, expr) =>
+ val df = maps.select(expr.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for $name")
+ }
+ }
+
+ test("a pairwise array_sort comparator is lifted over the cross product") {
+ // One UDF call receiving both elements has no per-element key, so the UDF is precomputed over
+ // every ordered pair instead. The call must end up outside every lambda like any other.
+ val df = arrayDF.select(
+ array_sort(col("values"), (a, b) => pythonUDF2(a, b).cast("int")).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ // The UDF takes both pair sides, each a flat array of all n*n pairs.
+ assert(lifted.head.children.size == 2)
+ }
+
+ test("a pandas UDF inside a lambda still fails analysis") {
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x => scalarPandasUDF(x))).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("the rewrite can be disabled by conf, restoring the previous error") {
+ withSQLConf(SQLConf.PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED.key -> "false") {
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x => pythonUDF(x))).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+ }
+
+ test("the rewrite cannot be disabled via excludedRules") {
+ // The lambda rewrite is driven by `ExtractPythonUDFs` (not a standalone batch rule), and that
+ // rule is non-excludable, so a plan that only works because of the rewrite cannot be broken by
+ // excludedRules. Excluding either name must leave the UDF lifted out of the lambda.
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ Seq(ExtractPythonUDFFromLambda.ruleName, ExtractPythonUDFs.ruleName).mkString(",")) {
+ val df = arrayDF.select(transform(col("values"), x => pythonUDF(x)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+ }
+
+ test("a UDF over only constants is still lifted per element") {
+ // SPARK-27052: `transform(arr, x -> udf(lit(10)))` does not read the element, so it is
Review Comment:
Stale after 742727c: the comment still says the constant-only UDF "is left to ExtractPythonUDFs", but it is now lifted per element, as the test title and `test_udf_with_constant_argument_only` assert. Please update the comment.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,73 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
+ case udf: PythonUDF =>
+ udf.evalType == PythonEvalType.SQL_BATCHED_UDF ||
+ udf.evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF
+ case _ => false
+ }
+
+ /**
+ * Whether every Python UDF in `hof`'s lambdas can be lifted out by `ExtractPythonUDFFromLambda`.
+ * Used by `CheckAnalysis` to decide whether to reject the plan. Three shapes cannot be rewritten:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so it sees earlier steps' outputs;
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+ def canRewritePythonUDFInLambda(hof: HigherOrderFunction): Boolean = {
Review Comment:
This predicate is not exactly the one the rule matches: the rewrite additionally requires the `ResultTypeFromArgument` / `ResultTypeFromFunction` marker traits, which this check does not know about. All built-in single-lambda HOFs are marked today, but a future one without either trait would pass analysis here and never be rewritten, leaving the UDF inside the lambda until runtime. Including the trait check here would make "analysis accepts exactly what the rule rewrites" structural rather than by convention.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,73 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
Review Comment:
Two edge shapes this predicate admits that the rewrite does not cover (or test):
1. A scalar UDF called with named arguments: the `NamedArgumentExpression` child ends up buried inside the generated `ArrayTransform`, so the kwargs mapping the runner derives from direct children is lost for the lifted UDF.
2. A UDF whose argument or return type involves a UDT: previously an analysis error, now routed to the Arrow elementwise path where it fails at runtime (the UDT fallback in `correctEvalType` only handles 101 -> 100).
Consider rejecting these here (keeping the old analysis error) or adding coverage.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
+ flatCells,
+ Add(Add(Multiply(idxOf(cmpLeft), n), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0)
+ }
+
+
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type: `transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array, indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption, Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps its own class (via
+ // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
+ * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ */
+ private def liftableHof(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 && (hof.functions.head match {
+ case LambdaFunction(body, args, _) =>
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
+ 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 += (liftKey(udf) -> 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 udfFieldByKey = liftableUDFs.map(liftKey).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 udfFieldByKey.contains(liftKey(udf)) =>
+ val ordinal = udfFieldByKey(liftKey(udf))
+ // A UDF over a comparator's right-hand element must read that element's key, so the
+ // struct field is read through whichever bound variable the call's own arguments used.
+ val side = udf.collectFirst {
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => v
+ }
+ side match {
+ case Some(v) => readerFor(v, Some(ordinal))
+ case None =>
+ GetStructField(boundVar, arguments.length + ordinal,
+ Some(s"$carrierUDFFieldPrefix$ordinal"))
+ }
+ case v: NamedLambdaVariable if fieldOfVar.contains(v.exprId) => readerFor(v, None)
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) =>
+ // A comparator's right-hand element itself, read through its own bound variable.
+ GetStructField(extraVarOf(v.exprId), 0, Some(carrierNames.head))
+ }
+
+ Carrier(carrier, rewrittenBody, boundVar, extraBoundVars)
+ }
+
+ /**
+ * Collects the Python UDF calls directly in `body` that must be lifted, innermost first.
+ *
+ * Every rewritable UDF directly in the lambda body is lifted, even one whose arguments do not
+ * read the lambda variable (`transform(arr, _ -> f(lit(10)))`). It cannot stay inside the lambda,
+ * and lifting it - `overArray` repeats a constant/outer-column argument into an aligned array -
+ * gives it the lambda's own call domain: once per element, and zero times for a null or empty
+ * array. Leaving it to [[ExtractPythonUDFs]] would instead call it once per row, including rows
+ * whose array is null or empty where the lambda never runs.
+ */
+ private def collectLiftableUDFs(
+ body: Expression,
+ lambdaExprIds: Set[ExprId]): Seq[PythonUDF] = {
Review Comment:
`lambdaExprIds` is unused in this method. (Also two stray double blank lines: after `rewritePairwiseComparator` and after `hasDirectRewritableUDF`.)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3748866853
##########
python/pyspark/worker.py:
##########
@@ -2989,6 +2991,137 @@ 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:
Review Comment:
Could we create the `ThreadPoolExecutor` once per `func` iterator and reuse it across UDFs and record batches? `_evaluate_elementwise_udf` is called once per UDF per batch, so the current placement repeatedly creates and tears down worker threads on this hot path when Arrow concurrency is enabled.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon closed pull request #57804: [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas
URL: https://github.com/apache/spark/pull/57804
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
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,
+)
+
+
+@unittest.skipIf(
+ 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3746565370
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,93 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call. Two otherwise-eligible shapes are excluded because the rewrite cannot preserve them:
+ * - a call with named arguments: its `NamedArgumentExpression` children would be buried inside
+ * the generated `ArrayTransform`, losing the kwargs mapping the runner derives from the
+ * direct children;
+ * - a UDF whose argument or return type involves a UDT: the lift forces the Arrow element-wise
+ * eval type, which has no UDT fallback (unlike `correctEvalType`'s Arrow -> pickle path), so
+ * it would fail at runtime instead of at analysis.
+ * Both keep the previous behavior (an analysis error) rather than being rewritten.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
+ case udf: PythonUDF =>
+ (udf.evalType == PythonEvalType.SQL_BATCHED_UDF ||
+ udf.evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) &&
+ !udf.children.exists(_.isInstanceOf[NamedArgumentExpression]) &&
+ !containsUDT(udf.dataType) &&
+ !udf.children.exists(c => containsUDT(c.dataType))
+ case _ => false
+ }
+
+ /**
+ * Whether every Python UDF in `hof`'s lambdas can be lifted out by `ExtractPythonUDFFromLambda`.
+ * Used by `CheckAnalysis` to decide whether to reject the plan. Three shapes cannot be rewritten:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so it sees earlier steps' outputs;
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+ def canRewritePythonUDFInLambda(hof: HigherOrderFunction): Boolean = {
Review Comment:
A nondeterministic *iterated argument* can produce silently wrong results, because the rewrite duplicates the argument expression: the carrier's `c0` and each lifted UDF's argument both re-reference it (`arrays_zip(shuffle(arr) AS c0, f_arr(shuffle(arr)) AS u0)`), the map desugar evaluates `MapKeys(m)` / `MapValues(m)` separately and re-references `keys` in the rebuild, and the pairwise path references `argument` several times. Nondeterministic expressions are excluded from subexpression elimination, so each copy evaluates independently.
Concretely, `filter(shuffle(arr), x -> is_even(x))` evaluates two independent shuffles - the predicate results are misaligned with the elements, so the wrong elements are kept - and `transform_values(<nondeterministic map expr>, ...)` can pair keys with values from a different evaluation. (This is distinct from the nondeterministic *UDF call* dedup already handled by `liftKey`.)
Suggest requiring `hof.arguments.forall(_.deterministic)` here, so the shape falls back to the previous analysis error - consistent with the PR's "analysis accepts exactly what the rule rewrites" principle.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,93 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call. Two otherwise-eligible shapes are excluded because the rewrite cannot preserve them:
+ * - a call with named arguments: its `NamedArgumentExpression` children would be buried inside
+ * the generated `ArrayTransform`, losing the kwargs mapping the runner derives from the
+ * direct children;
+ * - a UDF whose argument or return type involves a UDT: the lift forces the Arrow element-wise
+ * eval type, which has no UDT fallback (unlike `correctEvalType`'s Arrow -> pickle path), so
+ * it would fail at runtime instead of at analysis.
+ * Both keep the previous behavior (an analysis error) rather than being rewritten.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
Review Comment:
A zero-argument UDF slips through this predicate and crashes the Python worker at runtime. For `transform(arr, x -> zero_udf())` every check here is vacuously true on empty `children`, so analysis accepts it; `buildCarrier` then builds `arrayArgs = udf.children.map(...)` as an **empty** Seq, and the lifted element-wise UDF reaches the worker with no array argument. The worker indexes `input_batch.column(offsets_meta[0])` unconditionally - its comment says "The rewrite always passes at least one array argument", which does not hold for this shape - so the query dies with an `IndexError` from the worker instead of the previous clean `AnalysisException`.
Minimal fix: add `udf.children.nonEmpty` here so the shape keeps failing analysis as before (or alternatively inject the aligned shape array like the constant-argument case does). Either way please add a `transform(arr, x -> zero_udf())` test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,612 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position and the shared flat result array so the comparator can read
+ // its pair's precomputed cell, sort, then drop them again. `flatCells` must be built here, in
+ // the sort's *argument*, not inside the comparator: `ArraySort` re-evaluates the whole
+ // comparator body on every comparison, and `ExtractPythonUDFs` hoists only the `PythonUDF`
+ // node - the surrounding `arrays_zip`/`transform`/`flatten`/`array_repeat` that build the cells
+ // would otherwise be rebuilt O(n^2) per comparison (O(n^3 log n) overall). In interpreted
+ // evaluation `array_repeat` stores n references to the one computed `flatCells` array, so the
+ // carry is O(n^2); a later copy of the carrier into the Unsafe format would materialize each
+ // reference into O(n^3) bytes. Either way the whole pairwise path is already O(n^2) in Python
+ // calls, so it is only intended for small arrays (see the config doc).
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val cellsField = "cells"
+ val indexed = ArraysZip(
+ Seq(
+ argument,
+ ArrayTransform(argument, LambdaFunction(posIdx, Seq(posElem, posIdx))),
+ ArrayRepeat(flatCells, n)),
+ Seq(
+ Literal(s"${carrierElementPrefix}0"),
+ Literal(carrierIndexField),
+ Literal(cellsField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. The cells live in a
+ // struct field carried by every element, so the comparator only does a field read plus an
+ // `element_at`, both O(1). `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val cells = GetStructField(cmpLeft, 2, Some(cellsField))
+ val comparison = ElementAt(
Review Comment:
Minor, residual from the O(1)-comparator fix: `n` here is `Size(argument)`, embedded in the comparator body, so `argument` is re-evaluated on every comparison. Negligible for a column reference, but if `argument` is a computed expression (e.g. a `transform`), each comparison pays its full evaluation cost again. Carrying `n` in the carrier struct (like `idx`) - or deriving it as `Size(cells)`-free constant per row - would remove the last per-comparison re-evaluation.
##########
python/pyspark/worker.py:
##########
@@ -2989,6 +2991,133 @@ 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 renest_spec(shape):
+ """Offsets, list class and null mask that re-nest a flat result list by ``shape``.
+
+ Null rows stay null and consume no offsets. ``ListArray`` uses int32 offsets and
+ ``LargeListArray`` int64, so the input's list width is preserved.
+ """
+ lengths = pc.list_value_length(shape).to_pylist()
+ offsets = []
+ running = 0
+ for n in lengths:
+ offsets.append(running)
+ if n is not None:
+ running += n
+ offsets.append(running)
+ is_large = pa.types.is_large_list(shape.type)
+ list_cls = pa.LargeListArray if is_large else pa.ListArray
+ offsets_arr = pa.array(offsets, type=pa.int64() if is_large else pa.int32())
+ null_mask = pa.array([n is None for n in lengths], type=pa.bool_())
+ total_elements = running
+ return list_cls, offsets_arr, null_mask, total_elements
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ # Flatten each list column once to its element list, converting to Python.
+ columns = []
+ for col, conv in zip(input_batch.itercolumns(), arrow_to_py_converters):
+ values = ArrowTableToRowsConversion._to_pylist(col.flatten())
+ if conv is not None:
+ values = [conv(v) for v in values]
+ columns.append(values)
+
+ # Each UDF is re-nested by *its own* first argument's shape. ExtractPythonUDFs can
+ # fuse UDFs over differently shaped arrays into one batch, so a single shared shape
+ # would misalign every UDF but the first. The rewrite always passes at least one
+ # array argument, so `offsets_meta` is non-empty.
+ output_arrays = []
+ for wrapped_func, offsets_meta, arrow_element_type, result_conv in udf_infos:
+ list_cls, offsets_arr, null_mask, total_elements = renest_spec(
+ input_batch.column(offsets_meta[0])
+ )
+ # Stream the argument tuples rather than materializing a batch-sized list.
+ rows = zip(*[columns[o] for o in offsets_meta])
+ results = _evaluate_elementwise_udf(wrapped_func, rows)
+ verify_result_row_count(len(results), total_elements)
+
+ # Convert results and re-nest to array<R> using that UDF's offsets.
+ converted = (
+ [result_conv(r) for r in results] if result_conv is not None else results
+ )
+ try:
+ flat_arr = pa.array(converted, type=arrow_element_type)
+ except (pa.lib.ArrowInvalid, pa.lib.ArrowTypeError):
Review Comment:
Nit: this catches `(pa.lib.ArrowInvalid, pa.lib.ArrowTypeError)`, while the non-legacy `SQL_ARROW_BATCHED_UDF` path this mirrors catches only `pa.lib.ArrowInvalid`. If the broader catch is deliberate (it does look more robust for list-typed results), consider aligning the batched path or noting the difference in a comment, so the two paths don't drift apart silently.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3736199081
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -8657,7 +8657,8 @@
},
"LAMBDA_FUNCTION_WITH_PYTHON_UDF" : {
"message" : [
- "Lambda function with Python UDF <funcName> in a higher order function."
+ "Cannot evaluate the Python UDF <funcName> inside the lambda of a higher-order function.",
+ "This placement is not supported: the UDF reads a value that only exists while the lambda iterates (a fold accumulator in `aggregate`/`reduce`, or an array bound by an enclosing lambda in a nested higher-order function), so it cannot be applied to the whole column at once. Rewrite the query so the UDF is applied outside the lambda."
Review Comment:
Fixed in 742727c. The message no longer asserts a specific cause (it now reads "This placement is not supported. Rewrite the query so the UDF is applied outside the lambda."), and CheckAnalysis now names the offending UDF - the first of an unsupported eval type if any, else the first Python UDF in the lambdas.
##########
python/pyspark/worker.py:
##########
@@ -2989,6 +2991,128 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record
# profiling is not supported for UDF
return func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF:
+ # This path exchanges data with the JVM over Arrow, so PyArrow is required. Fail with a
+ # clear message rather than a bare ImportError from `import pyarrow` below.
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+
+ require_minimum_pyarrow_version()
+
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ # Element-wise UDFs back higher-order lambdas like transform(arr, x -> udf(x)).
+ # ExtractPythonUDFFromLambda rewrites them so the UDF receives *all* array elements
+ # at once (as ``array<T>``) rather than per-element. Flatten once, evaluate once over
+ # the batch, then re-nest with input offsets. Example: array<int> -> udf -> array<int>.
+
+ # UDF preparation
+ udf_infos = []
+ for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in udfs:
+ wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
+ udf_func, udf_args_offsets, udf_kwargs_offsets
+ )
+ # UDF returns one value per element; return type was pickled, unchanged.
+ # This is per-element, so element type equals the declared return type.
+ element_return_type = udf_return_type
+ udf_infos.append(
+ (
+ wrapped_func,
+ args_kwargs_offsets,
+ to_arrow_type(
+ element_return_type,
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
+ ),
+ LocalDataToArrowConversion._create_converter(
+ element_return_type,
+ none_on_identity=True,
+ int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ ),
+ )
+ )
+ col_names = [f"_{i}" for i in range(len(udfs))]
+
+ # Input: every argument arrives as ``array<T>`` aligned with the iterated array.
+ # Flatten once per column; convert elements with the array's element type.
+ input_fields = list(eval_conf.input_type)
+ arrow_to_py_converters = [
+ ArrowTableToRowsConversion._create_converter(
+ f.dataType.elementType,
+ none_on_identity=True,
+ binary_as_bytes=runner_conf.binary_as_bytes,
+ )
+ for f in input_fields
+ ]
+
+ @fail_on_stopiteration
+ def _evaluate_elementwise_udf(udf_func, rows):
+ if runner_conf.arrow_concurrency_level <= 0:
+ return [udf_func(*row) for row in rows]
+ from concurrent.futures import ThreadPoolExecutor
+
+ with ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool:
+ return list(pool.map(lambda row: udf_func(*row), rows))
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ # Flatten all array columns to element lists and convert to Python.
+ columns = []
+ for col, conv in zip(input_batch.itercolumns(), arrow_to_py_converters):
+ values = ArrowTableToRowsConversion._to_pylist(col.flatten())
+ if conv is not None:
+ values = [conv(v) for v in values]
+ columns.append(values)
+
+ # Extract shape (lengths per row, null mask) from first column.
+ shape = input_batch.column(0)
+ lengths = pc.list_value_length(shape).to_pylist()
+ total_elements = sum(n for n in lengths if n is not None)
+
+ # Build offsets to re-nest flat list back to array<R> with input shape.
+ # Null rows stay null; they consume no offsets. Preserve int32/int64 width.
+ offsets = []
+ running = 0
+ for n in lengths:
+ offsets.append(running)
+ if n is not None:
+ running += n
+ offsets.append(running)
+ is_large = pa.types.is_large_list(shape.type)
+ list_cls = pa.LargeListArray if is_large else pa.ListArray
+ offsets_arr = pa.array(offsets, type=pa.int64() if is_large else pa.int32())
+ null_mask = pa.array([n is None for n in lengths], type=pa.bool_())
+
+ # Evaluate all UDFs once over the flattened batch.
+ output_arrays = []
+ for wrapped_func, offsets_meta, arrow_element_type, result_conv in udf_infos:
+ rows = (
+ [() for _ in range(total_elements)]
+ if not offsets_meta
+ else list(zip(*[columns[o] for o in offsets_meta]))
Review Comment:
Fixed in 742727c. The worker now streams `zip(*[columns[o] for o in offsets_meta])` directly instead of materializing a batch-sized list.
##########
python/pyspark/sql/tests/test_udf_in_higher_order_function.py:
##########
@@ -0,0 +1,514 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+
+from pyspark.errors import AnalysisException
+from pyspark.sql import functions as sf
+from pyspark.sql.functions import udf
+from pyspark.sql.types import ArrayType, DoubleType, IntegerType, StringType
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import (
+ assertDataFrameEqual,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message
+)
+class UDFInHigherOrderFunctionTestsMixin:
+ """Tests for scalar Python UDFs used inside higher-order function lambdas (SPARK-27052).
+
+ ``ExtractPythonUDFFromLambda`` rewrites such a plan so the UDF is applied to the whole array
+ outside the lambda. Each test asserts the *result*, comparing against the equivalent native
+ expression wherever one exists, so that a rewrite that runs but computes the wrong thing
+ fails rather than passing quietly.
+ """
+
+ def test_transform(self):
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],), ([10],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")),
+ df.select(sf.transform("values", lambda x: x + 1).alias("r")),
+ )
+
+ def test_transform_null_array_and_null_elements(self):
+ # A null array must stay null, and a null *element* must reach the UDF as None.
+ df = self.spark.createDataFrame(
+ [([1, None, 3],), (None,), ([],)], "values array<int>"
+ )
+ # Null-aware so the UDF itself can observe the null element.
+ f = udf(lambda x: -1 if x is None else x * 2, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: f(x)).alias("r")),
+ [([2, -1, 6],), (None,), ([],)],
+ )
+
+ def test_transform_udf_returning_null(self):
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ f = udf(lambda x: None if x == 2 else x, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: f(x)).alias("r")),
+ [([1, None, 3],)],
+ )
+
+ def test_transform_with_index(self):
+ df = self.spark.createDataFrame([([10, 20, 30],), ([],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ # The index parameter must still work once the element is read from the carrier struct.
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x, i: plus_one(x) + i).alias("r")),
+ df.select(sf.transform("values", lambda x, i: (x + 1) + i).alias("r")),
+ )
+
+ def test_composition_around_udf_result(self):
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ # Arithmetic, `when` and casts around the UDF result are ordinary JVM work.
+ assertDataFrameEqual(
+ df.select(
+ sf.transform("values", lambda x: plus_one(x) * 2).alias("mul"),
+ sf.transform(
+ "values", lambda x: sf.when(plus_one(x) > 2, sf.lit(1)).otherwise(sf.lit(0))
+ ).alias("cond"),
+ sf.transform("values", lambda x: plus_one(x).cast("string")).alias("cast"),
+ ),
+ df.select(
+ sf.transform("values", lambda x: (x + 1) * 2).alias("mul"),
+ sf.transform(
+ "values", lambda x: sf.when((x + 1) > 2, sf.lit(1)).otherwise(sf.lit(0))
+ ).alias("cond"),
+ sf.transform("values", lambda x: (x + 1).cast("string")).alias("cast"),
+ ),
+ )
+
+ def test_udf_argument_is_expression_over_element(self):
+ # `udf(x * 2)`: the argument is itself an expression over the element.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one(x * 2)).alias("r")),
+ df.select(sf.transform("values", lambda x: x * 2 + 1).alias("r")),
+ )
+
+ def test_multiple_udfs_in_one_lambda(self):
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ times_ten = udf(lambda x: x * 10, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one(x) + times_ten(x)).alias("r")),
+ df.select(sf.transform("values", lambda x: (x + 1) + (x * 10)).alias("r")),
+ )
+
+ def test_nested_udfs(self):
+ # `f(g(x))`: both are lifted, and compose as array UDFs outside the lambda.
+ df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+ times_ten = udf(lambda x: x * 10, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: times_ten(plus_one(x))).alias("r")),
+ df.select(sf.transform("values", lambda x: (x + 1) * 10).alias("r")),
+ )
+
+ def test_udf_with_outer_column_argument(self):
+ # A non-element argument must be broadcast to every element of its row.
+ df = self.spark.createDataFrame(
+ [([1, 2], 100), ([3], 200)], "values array<int>, base int"
+ )
+ add = udf(lambda x, b: x + b, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: add(x, sf.col("base"))).alias("r")),
+ [([101, 102],), ([203],)],
+ )
+
+ def test_udf_with_constant_argument_only(self):
+ # SPARK-27052: `transform(arr, x -> udf(lit(10)))` must still yield one result per
+ # element rather than a single value.
+ df = self.spark.createDataFrame([([1, 2, 3],), ([],)], "values array<int>")
+ const = udf(lambda v: v * 2, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: const(sf.lit(10))).alias("r")),
+ [([20, 20, 20],), ([],)],
+ )
+
+ def test_filter(self):
+ # `filter`'s result is built from the input elements, not the lambda's value.
+ df = self.spark.createDataFrame([([1, 2, 3, 4],), ([],), (None,)], "values array<int>")
+ is_even = udf(lambda x: x % 2 == 0, "boolean")
+
+ assertDataFrameEqual(
+ df.select(sf.filter("values", lambda x: is_even(x)).alias("r")),
+ df.select(sf.filter("values", lambda x: (x % 2) == 0).alias("r")),
+ )
+
+ def test_exists_and_forall(self):
+ df = self.spark.createDataFrame([([1, 2, 3],), ([2, 4],), ([],)], "values array<int>")
+ is_even = udf(lambda x: x % 2 == 0, "boolean")
+
+ assertDataFrameEqual(
+ df.select(
+ sf.exists("values", lambda x: is_even(x)).alias("e"),
+ sf.forall("values", lambda x: is_even(x)).alias("f"),
+ ),
+ df.select(
+ sf.exists("values", lambda x: (x % 2) == 0).alias("e"),
+ sf.forall("values", lambda x: (x % 2) == 0).alias("f"),
+ ),
+ )
+
+ def test_zip_with(self):
+ # Two arrays at once. `arrays_zip` pads the shorter side with nulls, which is what
+ # `zip_with` does itself, so differing lengths must agree with the native version.
+ df = self.spark.createDataFrame(
+ [([1, 2], [10, 20]), ([1, 2, 3], [10]), ([], []), (None, [1]), ([1], None)],
+ "l array<int>, r array<int>",
+ )
+ add = udf(lambda a, b: (0 if a is None else a) + (0 if b is None else b), IntegerType())
+
+ # Compare against the equivalent native expression with the same null handling.
+ assertDataFrameEqual(
+ df.select(sf.zip_with("l", "r", lambda a, b: add(a, b)).alias("r")),
+ df.select(
+ sf.zip_with(
+ "l",
+ "r",
+ lambda a, b: sf.coalesce(a, sf.lit(0)) + sf.coalesce(b, sf.lit(0)),
+ ).alias("r")
+ ),
+ )
+
+ def test_zip_with_udf_on_one_side_only(self):
+ df = self.spark.createDataFrame([([1, 2], [10, 20])], "l array<int>, r array<int>")
+ plus_one = udf(lambda x: x + 1, IntegerType())
+
+ assertDataFrameEqual(
+ df.select(sf.zip_with("l", "r", lambda a, b: plus_one(a) + b).alias("r")),
+ [([12, 23],)],
+ )
+
+ def test_array_sort_with_udf_key(self):
+ # A comparator cannot be evaluated pairwise, but the UDF applied per element is a sort key
Review Comment:
Fixed in 742727c. Reworded so it no longer makes the absolute claim; it now notes the per-element key path and points to the pairwise test for the both-elements case.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -36,6 +36,10 @@ object PythonUDF {
private[this] val SCALAR_TYPES = Set(
PythonEvalType.SQL_BATCHED_UDF,
PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ // Element-wise UDFs are row-shaped from the plan's point of view: one array column in, one
Review Comment:
Fixed in 742727c. The comment now describes it as row-shaped from the plan's point of view: one array column in, one array column out per row.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,547 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * never names a concrete class: it reads arguments, lambdas and parameter roles off the
Review Comment:
Fixed in 742727c. Reworded to say the path names a concrete class only where a shape cannot be inferred otherwise (ArraySort's comparator, and the result-type traits telling ArrayFilter from ArrayTransform).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon closed pull request #57804: [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas
URL: https://github.com/apache/spark/pull/57804
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3747574863
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,93 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call. Two otherwise-eligible shapes are excluded because the rewrite cannot preserve them:
+ * - a call with named arguments: its `NamedArgumentExpression` children would be buried inside
+ * the generated `ArrayTransform`, losing the kwargs mapping the runner derives from the
+ * direct children;
+ * - a UDF whose argument or return type involves a UDT: the lift forces the Arrow element-wise
+ * eval type, which has no UDT fallback (unlike `correctEvalType`'s Arrow -> pickle path), so
+ * it would fail at runtime instead of at analysis.
+ * Both keep the previous behavior (an analysis error) rather than being rewritten.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
Review Comment:
Fixed in a657885. Added `udf.children.nonEmpty` to `isElementwiseRewritableUDF`, so a zero-argument call keeps failing analysis with `UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF` instead of reaching the worker with no array argument. Added `transform(arr, x -> zero_udf())` coverage - a predicate unit test, a Scala analysis test, and a Python end-to-end test (`test_zero_argument_udf_still_fails`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,93 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call. Two otherwise-eligible shapes are excluded because the rewrite cannot preserve them:
+ * - a call with named arguments: its `NamedArgumentExpression` children would be buried inside
+ * the generated `ArrayTransform`, losing the kwargs mapping the runner derives from the
+ * direct children;
+ * - a UDF whose argument or return type involves a UDT: the lift forces the Arrow element-wise
+ * eval type, which has no UDT fallback (unlike `correctEvalType`'s Arrow -> pickle path), so
+ * it would fail at runtime instead of at analysis.
+ * Both keep the previous behavior (an analysis error) rather than being rewritten.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
+ case udf: PythonUDF =>
+ (udf.evalType == PythonEvalType.SQL_BATCHED_UDF ||
+ udf.evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) &&
+ !udf.children.exists(_.isInstanceOf[NamedArgumentExpression]) &&
+ !containsUDT(udf.dataType) &&
+ !udf.children.exists(c => containsUDT(c.dataType))
+ case _ => false
+ }
+
+ /**
+ * Whether every Python UDF in `hof`'s lambdas can be lifted out by `ExtractPythonUDFFromLambda`.
+ * Used by `CheckAnalysis` to decide whether to reject the plan. Three shapes cannot be rewritten:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so it sees earlier steps' outputs;
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+ def canRewritePythonUDFInLambda(hof: HigherOrderFunction): Boolean = {
Review Comment:
Fixed in a657885. `canRewritePythonUDFInLambda` now also requires `hof.arguments.forall(_.deterministic)`, so a nondeterministic iterated argument falls back to the previous analysis error rather than being rewritten into misaligned copies. Kept it distinct from the nondeterministic-UDF-*call* handling in `liftKey`. Added a Scala test (`filter(shuffle(values), ...)`) and a Python end-to-end test (`test_nondeterministic_iterated_argument_still_fails`).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,612 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position and the shared flat result array so the comparator can read
+ // its pair's precomputed cell, sort, then drop them again. `flatCells` must be built here, in
+ // the sort's *argument*, not inside the comparator: `ArraySort` re-evaluates the whole
+ // comparator body on every comparison, and `ExtractPythonUDFs` hoists only the `PythonUDF`
+ // node - the surrounding `arrays_zip`/`transform`/`flatten`/`array_repeat` that build the cells
+ // would otherwise be rebuilt O(n^2) per comparison (O(n^3 log n) overall). In interpreted
+ // evaluation `array_repeat` stores n references to the one computed `flatCells` array, so the
+ // carry is O(n^2); a later copy of the carrier into the Unsafe format would materialize each
+ // reference into O(n^3) bytes. Either way the whole pairwise path is already O(n^2) in Python
+ // calls, so it is only intended for small arrays (see the config doc).
+ val posElem = NamedLambdaVariable("x", elementType, containsNull)
+ val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false)
+ val cellsField = "cells"
+ val indexed = ArraysZip(
+ Seq(
+ argument,
+ ArrayTransform(argument, LambdaFunction(posIdx, Seq(posElem, posIdx))),
+ ArrayRepeat(flatCells, n)),
+ Seq(
+ Literal(s"${carrierElementPrefix}0"),
+ Literal(carrierIndexField),
+ Literal(cellsField)))
+ val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. The cells live in a
+ // struct field carried by every element, so the comparator only does a field read plus an
+ // `element_at`, both O(1). `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val cells = GetStructField(cmpLeft, 2, Some(cellsField))
+ val comparison = ElementAt(
Review Comment:
Fixed in a657885. `n` is now carried in the carrier struct (an `array_repeat(n, n)` field, alongside the cells), and the comparator reads it as `GetStructField(..., "n")`, so `Size(argument)` is evaluated once per row instead of on every comparison - the last per-comparison re-evaluation is gone even when `argument` is a computed expression.
##########
python/pyspark/worker.py:
##########
@@ -2989,6 +2991,133 @@ 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 renest_spec(shape):
+ """Offsets, list class and null mask that re-nest a flat result list by ``shape``.
+
+ Null rows stay null and consume no offsets. ``ListArray`` uses int32 offsets and
+ ``LargeListArray`` int64, so the input's list width is preserved.
+ """
+ lengths = pc.list_value_length(shape).to_pylist()
+ offsets = []
+ running = 0
+ for n in lengths:
+ offsets.append(running)
+ if n is not None:
+ running += n
+ offsets.append(running)
+ is_large = pa.types.is_large_list(shape.type)
+ list_cls = pa.LargeListArray if is_large else pa.ListArray
+ offsets_arr = pa.array(offsets, type=pa.int64() if is_large else pa.int32())
+ null_mask = pa.array([n is None for n in lengths], type=pa.bool_())
+ total_elements = running
+ return list_cls, offsets_arr, null_mask, total_elements
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ # Flatten each list column once to its element list, converting to Python.
+ columns = []
+ for col, conv in zip(input_batch.itercolumns(), arrow_to_py_converters):
+ values = ArrowTableToRowsConversion._to_pylist(col.flatten())
+ if conv is not None:
+ values = [conv(v) for v in values]
+ columns.append(values)
+
+ # Each UDF is re-nested by *its own* first argument's shape. ExtractPythonUDFs can
+ # fuse UDFs over differently shaped arrays into one batch, so a single shared shape
+ # would misalign every UDF but the first. The rewrite always passes at least one
+ # array argument, so `offsets_meta` is non-empty.
+ output_arrays = []
+ for wrapped_func, offsets_meta, arrow_element_type, result_conv in udf_infos:
+ list_cls, offsets_arr, null_mask, total_elements = renest_spec(
+ input_batch.column(offsets_meta[0])
+ )
+ # Stream the argument tuples rather than materializing a batch-sized list.
+ rows = zip(*[columns[o] for o in offsets_meta])
+ results = _evaluate_elementwise_udf(wrapped_func, rows)
+ verify_result_row_count(len(results), total_elements)
+
+ # Convert results and re-nest to array<R> using that UDF's offsets.
+ converted = (
+ [result_conv(r) for r in results] if result_conv is not None else results
+ )
+ try:
+ flat_arr = pa.array(converted, type=arrow_element_type)
+ except (pa.lib.ArrowInvalid, pa.lib.ArrowTypeError):
Review Comment:
Addressed in a657885 by documenting the difference. The broader catch is deliberate: the element-wise wrapper commonly returns list/struct-typed elements, whose type mismatches surface as `ArrowTypeError` rather than `ArrowInvalid`. Added a comment at the element-wise site noting it intentionally catches a superset of the `SQL_ARROW_BATCHED_UDF` path so the two do not drift silently; left the batched path unchanged to avoid altering its established behavior.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3736197775
##########
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:
Fixed in 742727c. Deduplication now keys deterministic calls by canonical form and nondeterministic calls by identity (`liftKey`), so `transform(arr, x -> f(x) + f(x))` keeps two distinct lifted calls when `f` is nondeterministic. Added `test_nondeterministic_udf_calls_are_distinct` (Python) and "identical nondeterministic calls are lifted distinctly" (Scala).
##########
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:
Fixed in 742727c. Each fused element-wise UDF is now re-nested by its own first argument's shape (`renest_spec(input_batch.column(offsets_meta[0]))`) rather than a shared column-0 shape. Added `test_fused_transforms_with_different_lengths` covering two transforms over different lengths and null layouts.
##########
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:
Fixed in 742727c. `ExtractPythonUDFs.apply` now invokes `ExtractPythonUDFFromLambda` itself as a prerequisite, so the dependency is structural rather than relying on batch adjacency. The standalone batch entry is removed; only `ExtractPythonUDFs` remains in the batch and in `nonExcludableRules`.
##########
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:
Fixed in 742727c. The map rebuild now pattern-matches the concrete function (`TransformKeys`/`TransformValues`/`MapFilter`) instead of comparing result type to key type. Added a `transform_values` on `map<string,string>` regression test in both suites.
##########
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:
Fixed in 742727c. `collectLiftableUDFs` now lifts every rewritable UDF directly in the lambda body, including one whose arguments do not read the lambda variable; `overArray` repeats it into an aligned array. `test_udf_with_constant_argument_only` now also asserts the zero-call domain on empty and null arrays.
##########
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:
Fixed in 742727c. The comparator now indexes the flat n*n result array directly (`element_at(flatCells, a.idx*n + b.idx + 1)`); the per-row `matrix` `ArrayTransform` inside the comparator is gone, so `ExtractPythonUDFs` computes the results once and each comparison is O(1).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5235285604
Thanks @dongjoon-hyun for the careful pass and for confirming the dedup / worker re-nesting / ragged-zip semantics. Addressed the suggestions in 8d5e880:
1. **Eager-evaluation docs** - Added a note to the Arrow Python UDF tutorial ('Notes' section, next to the existing non-short-circuiting `WHERE` note): a UDF inside a higher-order function lambda is precomputed over the whole array, so it runs once per element (and over every pair for a two-argument `array_sort` comparator) regardless of `exists`/`when` short-circuiting, and a UDF that can fail must handle skipped inputs itself. Also expanded the config doc to say the same.
2. **Pairwise `array_sort` cost** - Documented the O(n^2) call/memory cost in the config doc and advise avoiding it for large arrays. I also corrected the code comment: the shared-reference carry is O(n^2) only in interpreted evaluation; a copy into the Unsafe format materializes each reference, degrading to O(n^3), as you noted. I stopped short of a hard runtime cap: it needs a new config + an injected error, and since these queries are opt-in and already O(n^2) in Python calls by construction, a documented warning seemed the more proportionate choice for this PR - happy to add a cap in a follow-up if you would prefer one.
3. **Defense-in-depth** - `liftableHof` now also calls `PythonUDF.hasFreeLambdaVariable(hof)` (made non-private and shared), so the rule re-checks the nested-lambda guard itself instead of trusting that `CheckAnalysis` ran. `transform(arr, i -> transform(i, x -> f(x)))` is now rejected by the rule independently.
4. **Test gaps** - Added: the SQL string syntax path (`spark.sql("SELECT transform(v, x -> pyUdf(x)) ...")`); an end-to-end kwargs call (`udf(x, y=lit(1))`) still failing analysis; a rewritable HOF inside another HOF's lambda over a real column (`transform(arr2, i -> array_max(transform(arr, x -> f(x))) + i)`, both as a Scala plan-shape test and an e2e result test); and decimal / timestamp / struct element types.
5. **`transform_keys` null-key error class** - Agreed this is expected for a new feature; the rebuild goes through `MapFromArrays` so null/duplicate-key dedup-policy behavior matches native, and the only difference is the error class text on a null key. Leaving as-is; noting it here for the record.
Verified locally: `ExtractPythonUDFFromLambdaSuite` (22 tests) and the end-to-end `pyspark.sql.tests.test_udf_in_higher_order_function` suite both pass.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3748922268
##########
python/pyspark/worker.py:
##########
@@ -2989,6 +2991,137 @@ 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:
Review Comment:
This is actually matched with the existent implementation, and it is off by default so let's probably leave it like this for now.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5239336607
I am working on Scalar Arrow UDF and Pandas UDF with this in followup works so merging it first. Let me know if there are more to fix.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5235781791
Thank you for keeping this PR up-to-date, @HyukjinKwon !
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5201153177
let me drop this for now
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5235039620
I went through the whole diff. Overall this is a well-designed and thoroughly tested PR: the shared predicate (`PythonUDF.isElementwiseRewritableUDF` / `canRewritePythonUDFInLambda`) keeps "what analysis accepts" aligned with "what the rule rewrites", the marker traits make the rewrite generic for future HOFs, and the SPARK-48706 failure mode (a raw `PythonUDF` left inside a generated lambda) is explicitly guarded and regression-tested. A few points below.
### Verified while reviewing
- Dedup semantics: deterministic duplicate calls collapse to one lifted UDF while nondeterministic calls stay distinct via `liftKey` (`resultId`-bearing node), and the key-form comparator's `udf(a)`/`udf(b)` correctly merge after lifting — all pinned by tests.
- Worker re-nesting is consistent with `pa.ListArray.flatten()` (null entries consume no offsets), and fused UDFs re-nest by their own first argument's shape; `verify_result_row_count` defends against misalignment.
- Ragged `zip_with`/`map_zip_with`: the common `arrays_zip` alignment and the `ArrayUnion` + `element_at(..., failOnError=false)` desugar match the native null-padding / key-union semantics.
- The conf's `.version("4.4.0")` matches the next feature release branch.
### Suggestions
1. **Document the eager-evaluation semantics.** After the rewrite, the UDF runs over *every* element regardless of `exists` short-circuiting, `when` branches, or boolean short-circuit around it. If the UDF raises on inputs a user expects to be skipped conditionally, the query now fails. Not a regression (these queries used to fail analysis), but the conf doc only partially covers it — a note in the SQL migration guide and/or the PySpark UDF docs would help (this PR has no doc changes).
2. **Consider a size guard for the pairwise `array_sort` comparator.** `(a, b) -> udf(a, b)` costs n² Python calls plus an n²-cell array per row; n = 10k means 100M calls per row, which will effectively hang or OOM an executor. Also, the "O(n²) carry via shared references" claim holds for `GenericArrayData`, but if the carrier is ever serialized to the Unsafe format the per-element copies degrade to O(n³). A cap with a clear error (or at least a warning in the conf doc) seems worthwhile.
3. **Defense-in-depth in the rule itself.** `liftableHof` doesn't check for free lambda variables — the nested-lambda guard lives only in `CheckAnalysis`. If the rule ever sees a plan that didn't go through analysis (internally constructed, or a shape reintroduced by a future rule), it would mis-rewrite `transform(arr, i -> transform(i, x -> f(x)))` and resurrect SPARK-48706. Reusing `hasFreeLambdaVariable` in `liftableHof` is cheap.
4. **Test gaps worth covering:**
- The SQL string syntax path, e.g. `spark.sql("SELECT transform(v, x -> pyUdf(x)) ...")`.
- An end-to-end test that a kwargs call (`udf(x, k=1)`) still fails analysis (only the predicate unit test exists).
- A rewritable HOF sitting *inside another HOF's lambda* while iterating a real column, e.g. `transform(arr2, i -> array_max(transform(arr, x -> f(x))) + i)` — analysis accepts it, the lifted element-wise UDF stays inside the outer lambda, and `ExtractPythonUDFs` then extracts it per-row. Deterministic UDFs give the same result, but the path is unusual enough to deserve a test.
- Decimal/timestamp/struct element types (Arrow conversion edge cases); current e2e coverage is int/double/string/array.
5. Minor: `transform_keys` rebuilds through `MapFromArrays`, so null-key / duplicate-key handling matches native dedup-policy behavior, but the error class on a null key may differ slightly from `TransformKeys`' own — fine for a new feature, just noting it.
None of these block the PR in my view — #1 and #2 are the substantive ones, the rest are nice-to-haves.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3745891476
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
+ flatCells,
+ Add(Add(Multiply(idxOf(cmpLeft), n), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0)
+ }
+
+
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type: `transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array, indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption, Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps its own class (via
+ // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
+ * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ */
+ private def liftableHof(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 && (hof.functions.head match {
+ case LambdaFunction(body, args, _) =>
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
+ 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
Review Comment:
Fixed in 95ec08d. Lifted UDFs are now deduplicated by their post-substitution array arguments (the canonical form of the lifted node), so a key-form comparator's `udf(a)` and `udf(b)` - which both read the whole array - share one lifted UDF and Python runs n times instead of 2n. Nondeterministic calls stay distinct (keyed by the node itself). Added "a key-form array_sort comparator lifts one UDF per element, not per comparator side".
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala:
##########
@@ -0,0 +1,556 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType}
+
+/**
+ * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated.
+ *
+ * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's
+ * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in
+ * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole
+ * array*, outside every lambda, and has the lambda read the result positionally:
+ *
+ * {{{
+ * -- before (rejected)
+ * transform(values, x -> plus_one(x))
+ *
+ * -- after (the PythonUDF is outside every lambda)
+ * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0)
+ * }}}
+ *
+ * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with
+ * [[PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF]]. The array-at-a-time behaviour lives in the Python
+ * worker: it flattens each list column once, calls the function over all elements of the batch, and
+ * re-nests by the input's offsets - one row in, one row out, one Python round trip per batch.
+ *
+ * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an
+ * element-independent value is repeated into one with a native `transform`), so the worker flattens
+ * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic,
+ * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work.
+ *
+ * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`,
+ * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to
+ * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a
+ * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs.
+ *
+ * `CheckAnalysis` still rejects what this rule does not handle:
+ * - a UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`: the inner array
+ * `i` is not a real column. (A UDF in a nested *argument*, `transform(arr, x ->
+ * transform(udf(x), y -> y))`, is fine - `udf(x)` lifts onto `arr`.)
+ * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps'
+ * outputs, not array elements.
+ * - a vectorized (scalar pandas / arrow) UDF, which is not supported.
+ */
+object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] {
+
+ def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.pythonUDFInHigherOrderFunctionEnabled) {
+ plan
+ } else {
+ // A single bottom-up pass lifts every liftable UDF: `transformExpressionsUpWithPruning`
+ // visits the innermost higher-order function first, and each rewrite lifts all of that
+ // lambda's UDFs at once. A UDF inside a *nested* function's lambda is not liftable at all
+ // (its argument is the outer lambda's variable, which is not a real column) and is rejected
+ // by `CheckAnalysis`, so no repeated fixed-point pass is needed.
+ plan.transformUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) {
+ case p =>
+ p.transformExpressionsUpWithPruning(
+ _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION))(rewrite)
+ }
+ }
+ }
+
+ /**
+ * Whether one UDF call in an `array_sort` comparator takes both elements, e.g.
+ * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross
+ * product of pairs rather than per element.
+ */
+ private def comparatorTakesBothElements(function: Expression): Boolean = function match {
+ case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) =>
+ body.exists {
+ case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) =>
+ def reads(id: ExprId) = udf.exists {
+ case v: NamedLambdaVariable => v.exprId == id
+ case _ => false
+ }
+ reads(left.exprId) && reads(right.exprId)
+ case _ => false
+ }
+ case _ => false
+ }
+
+ /**
+ * Rewrites one higher-order function whose lambda holds a rewritable Python UDF. The generic path
+ * reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and rebuilds
+ * with `withNewChildren` (children are `arguments` then `functions`), naming a concrete class
+ * only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the result-type
+ * traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` comparator, whose
+ * single call takes both elements, needs its own path.
+ */
+ private val rewrite: PartialFunction[Expression, Expression] = {
+ case sort @ ArraySort(_, function, _)
+ if liftableHof(sort) && comparatorTakesBothElements(function) =>
+ rewritePairwiseComparator(sort)
+
+ // Every result-typed higher-order function is handled; anything else is left alone.
+ case hof: HigherOrderFunction
+ if liftableHof(hof) &&
+ (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) =>
+ rewriteMapping(hof)
+ }
+
+ /**
+ * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is
+ * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with
+ * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions,
+ * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key.
+ */
+ private def rewritePairwiseComparator(sort: ArraySort): Expression = {
+ val ArraySort(argument, function, allowNull) = sort
+ val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) =
+ function
+ val arrayType = argument.dataType.asInstanceOf[ArrayType]
+ val elementType = arrayType.elementType
+ val containsNull = arrayType.containsNull
+ val n = Size(argument)
+
+ // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could
+ // capture the UDF; the one lambda here holds only the repeat, never the UDF.
+ val repeatVar = NamedLambdaVariable("a", elementType, containsNull)
+ val lefts = Flatten(
+ ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar))))
+ val rights = Flatten(ArrayRepeat(argument, n))
+
+ // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the
+ // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the
+ // comparator body (cast, `when`, arithmetic) once per pair in the JVM.
+ val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar))
+ val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None)
+ val flatCells = ArrayTransform(
+ pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar)))
+
+ // Carry each element's position so the comparator can read its pair's precomputed cell, sort,
+ // then drop the positions again.
+ 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
+
+ // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. `flatCells`'s Python UDF
+ // reads no comparator variable, so `ExtractPythonUDFs` hoists it to a computed-once column;
+ // indexing it here is O(1) per comparison. Building an n-row matrix inside the comparator
+ // instead would rebuild and slice it on every comparison. `element_at` is 1-based.
+ val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false)
+ val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false)
+ def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField))
+ val comparison = ElementAt(
+ flatCells,
+ Add(Add(Multiply(idxOf(cmpLeft), n), idxOf(cmpRight)), Literal(1)),
+ None,
+ failOnError = false)
+
+ unwrapCarrier(
+ ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0)
+ }
+
+
+
+ /**
+ * The generic rewrite for a mapping higher-order function.
+ *
+ * A map-valued argument is first desugared to its key and value arrays, so everything below works
+ * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then
+ * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a
+ * carrier that the single new lambda parameter reads.
+ */
+ private def rewriteMapping(hof: HigherOrderFunction): Expression = {
+ val lambda = hof.functions.head.asInstanceOf[LambdaFunction]
+ // The result is the input elements (so the carrier is unwrapped afterwards) rather than the
+ // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type.
+ val isFromElements = hof.isInstanceOf[ResultTypeFromArgument]
+
+ // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map
+ // up per key, which yields null for a key missing from one side - exactly its own semantics.
+ val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType])
+ val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) =
+ if (!mapValued) {
+ (hof.arguments, identity)
+ } else if (hof.arguments.length == 1) {
+ val map = hof.arguments.head
+ val keys = MapKeys(map)
+ val values = MapValues(map)
+ // Rebuild by the concrete function, not the result type: `transform_keys` replaces the
+ // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying
+ // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose
+ // lambda result type equals the key type.
+ val rebuild: Expression => Expression = hof match {
+ case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values)
+ case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues)
+ case _: MapFilter => (kept: Expression) =>
+ MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1))
+ }
+ (Seq(keys, values), rebuild)
+ } else {
+ val Seq(left, right) = hof.arguments
+ val keys = ArrayUnion(MapKeys(left), MapKeys(right))
+ val keyType = keys.dataType.asInstanceOf[ArrayType]
+ def valuesFor(map: Expression): Expression = {
+ val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull)
+ ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k)))
+ }
+ (Seq(keys, valuesFor(left), valuesFor(right)),
+ (newValues: Expression) => MapFromArrays(keys, newValues))
+ }
+
+ // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a
+ // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a
+ // comparator whose two parameters are two elements of the *same* array, indistinguishable from
+ // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here.
+ val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable])
+ val (elementVars, indexVar, alsoBind) =
+ if (hof.isInstanceOf[ArraySort]) {
+ (Seq(params.head), None, Seq(params.last))
+ } else {
+ (params.take(arrays.length), params.drop(arrays.length).headOption, Nil)
+ }
+
+ val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind)
+ val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars)
+
+ // Rebuild the node over the single carrier. A single-array function keeps its own class (via
+ // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array
+ // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so
+ // both key and value sides can be projected out.
+ val keepsOwnNode = hof.arguments.length == 1 && !mapValued
+ val iterated =
+ if (keepsOwnNode) {
+ hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression]
+ } else if (isFromElements) {
+ ArrayFilter(built.carrier, newLambda)
+ } else {
+ ArrayTransform(built.carrier, newLambda)
+ }
+
+ // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the
+ // carrier; for a map `rebuildResult` knows which of the key/value sides to keep.
+ if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0))
+ else rebuildResult(iterated)
+ }
+
+ /**
+ * True if `hof`'s single lambda holds a UDF belonging to *this* lambda (not a nested function's
+ * lambda). A UDF in a nested lambda is rejected by `CheckAnalysis`, so it is never matched here.
+ */
+ private def liftableHof(hof: HigherOrderFunction): Boolean =
+ hof.functions.length == 1 && (hof.functions.head match {
+ case LambdaFunction(body, args, _) =>
+ hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable])
+ 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 += (liftKey(udf) -> 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 udfFieldByKey = liftableUDFs.map(liftKey).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 udfFieldByKey.contains(liftKey(udf)) =>
+ val ordinal = udfFieldByKey(liftKey(udf))
+ // A UDF over a comparator's right-hand element must read that element's key, so the
+ // struct field is read through whichever bound variable the call's own arguments used.
+ val side = udf.collectFirst {
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => v
+ }
+ side match {
+ case Some(v) => readerFor(v, Some(ordinal))
+ case None =>
+ GetStructField(boundVar, arguments.length + ordinal,
+ Some(s"$carrierUDFFieldPrefix$ordinal"))
+ }
+ case v: NamedLambdaVariable if fieldOfVar.contains(v.exprId) => readerFor(v, None)
+ case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) =>
+ // A comparator's right-hand element itself, read through its own bound variable.
+ GetStructField(extraVarOf(v.exprId), 0, Some(carrierNames.head))
+ }
+
+ Carrier(carrier, rewrittenBody, boundVar, extraBoundVars)
+ }
+
+ /**
+ * Collects the Python UDF calls directly in `body` that must be lifted, innermost first.
+ *
+ * Every rewritable UDF directly in the lambda body is lifted, even one whose arguments do not
+ * read the lambda variable (`transform(arr, _ -> f(lit(10)))`). It cannot stay inside the lambda,
+ * and lifting it - `overArray` repeats a constant/outer-column argument into an aligned array -
+ * gives it the lambda's own call domain: once per element, and zero times for a null or empty
+ * array. Leaving it to [[ExtractPythonUDFs]] would instead call it once per row, including rows
+ * whose array is null or empty where the lambda never runs.
+ */
+ private def collectLiftableUDFs(
+ body: Expression,
+ lambdaExprIds: Set[ExprId]): Seq[PythonUDF] = {
Review Comment:
Fixed in 95ec08d. `collectLiftableUDFs` no longer takes `lambdaExprIds` (it was unused), and both stray double blank lines are removed.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -46,6 +50,73 @@ object PythonUDF {
e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType)
}
+ /**
+ * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by
+ * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda.
+ *
+ * Only the row-at-a-time eval types qualify, since the rule lifts a UDF that takes one value per
+ * call.
+ *
+ * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the
+ * optimizer rule can rewrite.
+ */
+ def isElementwiseRewritableUDF(e: Expression): Boolean = e match {
Review Comment:
Fixed in 95ec08d by rejecting both here. `isElementwiseRewritableUDF` now returns false when any child is a `NamedArgumentExpression` or when the argument/return type involves a UDT, so both keep the previous analysis error instead of failing at runtime. Added coverage: "a UDF with a UDT type inside a lambda still fails analysis" (end-to-end via analysis) and "the rewritable predicate excludes named-argument and UDT UDFs" (direct predicate unit test for named-arg, UDT return, and UDT argument).
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala:
##########
@@ -0,0 +1,279 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.{AnalysisException, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{LambdaFunction, PythonUDF}
+import org.apache.spark.sql.catalyst.plans.logical.ArrowEvalPython
+import org.apache.spark.sql.functions.{array_sort, col, forall, lit, map_filter, map_zip_with,
+ transform, transform_keys, transform_values, when, zip_with}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, StringType}
+
+/**
+ * Plan-shape tests for [[ExtractPythonUDFFromLambda]].
+ *
+ * These assert the structure the rewrite produces - that no `PythonUDF` is left inside a
+ * `LambdaFunction`, that the lifted UDF is an element-wise UDF over an array, and that plans
+ * without a UDF in a lambda are untouched. End-to-end result correctness is covered by
+ * `pyspark.sql.tests.test_udf_in_higher_order_function`, which needs a real Python worker.
+ */
+class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession {
+ import testImplicits._
+
+ private val pythonUDF = new MyDummyPythonUDF
+ private val scalarPandasUDF = new MyDummyScalarPandasUDF
+ // Used where a UDF call must receive two arguments, e.g. a pairwise comparator.
+ private val pythonUDF2 = new MyDummyPythonUDF
+ private val nondeterministicUDF = new MyDummyNondeterministicPythonUDF
+
+ private def arrayDF = Seq(Seq(1, 2, 3)).toDF("values")
+
+ /** All `PythonUDF`s that remain inside a lambda in the optimized plan. */
+ private def udfsInsideLambda(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] = {
+ df.queryExecution.optimizedPlan.expressions.flatMap { e =>
+ e.collect { case l: LambdaFunction => l }.flatMap { l =>
+ l.collect { case u: PythonUDF => u }
+ }
+ }
+ }
+
+ private def liftedUDFs(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] = {
+ df.queryExecution.optimizedPlan.collect {
+ case a: ArrowEvalPython => a.udfs
+ }.flatten
+ }
+
+ test("transform: the UDF is lifted out of the lambda as an element-wise array UDF") {
+ val df = arrayDF.select(transform(col("values"), x => pythonUDF(x)).as("r"))
+
+ // The whole point of the rewrite: nothing Python-shaped is left inside a lambda.
+ assert(udfsInsideLambda(df).isEmpty)
+
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ assert(lifted.head.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)
+ // The lifted UDF takes the array and returns an array of the original return type.
+ assert(lifted.head.dataType == ArrayType(pythonUDF.dataType, containsNull = true))
+ assert(lifted.head.children.size == 1)
+ assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType])
+ }
+
+ test("filter/exists/forall: the UDF is lifted out of the lambda") {
+ val exprs = Seq(
+ org.apache.spark.sql.functions.filter(col("values"), x => pythonUDF(x)),
+ org.apache.spark.sql.functions.exists(col("values"), x => pythonUDF(x)),
+ org.apache.spark.sql.functions.forall(col("values"), x => pythonUDF(x)))
+
+ exprs.foreach { e =>
+ val df = arrayDF.select(e.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF still inside a lambda for $e")
+ assert(liftedUDFs(df).forall(_.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF))
+ }
+ }
+
+ test("aggregate: a UDF anywhere in the fold is rejected") {
+ // The fold is sequential, so no array can precompute the values the UDF sees. A UDF in `merge`
+ // or in `finish` therefore has no rewrite and analysis must fail.
+ val cases = Seq(
+ "merge on element" ->
+ org.apache.spark.sql.functions.aggregate(
+ col("values"), lit(false), (acc, x) => acc || pythonUDF(x)),
+ "merge on accumulator" ->
+ org.apache.spark.sql.functions.aggregate(
+ col("values"), lit(false), (acc, x) => pythonUDF(acc) || x.cast("boolean")),
+ "finish" ->
+ org.apache.spark.sql.functions.aggregate(
+ col("values"), lit(0), (acc, x) => acc + x, acc => pythonUDF(acc)))
+ cases.foreach { case (name, expr) =>
+ val e = intercept[AnalysisException] {
+ arrayDF.select(expr.as("r")).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF",
+ s"expected aggregate with a UDF in $name to be rejected")
+ }
+ }
+
+ test("several UDFs and nested UDFs in one lambda are all lifted") {
+ val several = arrayDF.select(
+ transform(col("values"), x => pythonUDF(x) || pythonUDF(x + lit(1))).as("r"))
+ assert(udfsInsideLambda(several).isEmpty)
+ // Two distinct calls, so two lifted array UDFs.
+ assert(liftedUDFs(several).size == 2)
+
+ // The same deterministic call twice must be evaluated once.
+ val duplicated = arrayDF.select(
+ transform(col("values"), x => pythonUDF(x) || pythonUDF(x)).as("r"))
+ assert(udfsInsideLambda(duplicated).isEmpty)
+ assert(liftedUDFs(duplicated).size == 1)
+ }
+
+ test("identical nondeterministic calls are lifted distinctly, not deduplicated") {
+ // Deduplicating `f(x)` with `f(x)` would collapse two independent draws into one; a
+ // nondeterministic UDF must keep each call, so both are lifted.
+ val df = arrayDF.select(
+ transform(col("values"), x => nondeterministicUDF(x) || nondeterministicUDF(x)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ assert(liftedUDFs(df).size == 2)
+ }
+
+ test("transform_values whose result type equals the key type replaces values, not keys") {
+ // SPARK-27052: dispatch is by concrete function, not result type. For map<string, string>,
+ // transform_values' lambda also returns string, which must not be treated as new keys.
+ val maps = Seq(Map("a" -> "x")).toDF("m")
+ val df = maps.select(transform_values(col("m"), (k, v) => pythonUDF(v).cast("string")).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ // Keys are preserved: the map still has the original key type and no new-key projection.
+ val mapType = df.queryExecution.analyzed.schema.head.dataType.asInstanceOf[MapType]
+ assert(mapType.keyType == StringType)
+ }
+
+ test("a UDF inside a nested higher-order function's lambda is rejected") {
+ // The inner array is the outer lambda's variable, not a real column, so the UDF cannot be
+ // lifted onto an array outside the lambda. This must fail analysis rather than be rewritten.
+ val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
+ val e = intercept[AnalysisException] {
+ df.select(transform(col("values"), inner =>
+ transform(inner, x => pythonUDF(x))).as("r")).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("a UDF on the outer element of a nested array is lifted") {
+ // Here the UDF applies to the outer array's element, which is a real column (the element just
+ // happens to be an array), so it is lifted onto that column like any other single-level array.
+ val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values")
+ .select(transform(col("values"), inner => pythonUDF(inner)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+
+ test("a lambda with no Python UDF is left unchanged") {
+ val df = arrayDF.select(transform(col("values"), x => x + lit(1)).as("r"))
+ val analyzed = df.queryExecution.analyzed
+ val optimized = df.queryExecution.optimizedPlan
+ // The rule must be inert: no eval-python node is introduced.
+ assert(optimized.collect { case a: ArrowEvalPython => a }.isEmpty)
+ assert(!optimized.toString.contains("pythonUDF"))
+ assert(analyzed.expressions.flatMap(_.collect { case u: PythonUDF => u }).isEmpty)
+ }
+
+ test("every mapping higher-order function is rewritten") {
+ // One assertion per function, so that a shape regressing to "UDF left inside a lambda" is
+ // caught here rather than only by the end-to-end Python suite. `aggregate` / `reduce` are not
+ // here: they fold rather than map, so a UDF in them is rejected (see the aggregate test above).
+ val arrays = Seq((Seq(1, 2), Seq(3, 4))).toDF("l", "r")
+ val maps = Seq((Map("a" -> 1), Map("a" -> 2))).toDF("l", "r")
+
+ val arrayCases = Seq(
+ "transform" -> transform(col("l"), x => pythonUDF(x)),
+ "transform with index" -> transform(col("l"), (x, i) => pythonUDF(x) || i > 0),
+ "filter" -> org.apache.spark.sql.functions.filter(col("l"), x => pythonUDF(x)),
+ "exists" -> org.apache.spark.sql.functions.exists(col("l"), x => pythonUDF(x)),
+ "forall" -> forall(col("l"), x => pythonUDF(x)),
+ "zip_with" -> zip_with(col("l"), col("r"), (a, b) => pythonUDF(a) || pythonUDF(b)),
+ "array_sort" -> array_sort(col("l"),
+ (a, b) => when(pythonUDF(a) === pythonUDF(b), lit(0)).otherwise(lit(1))))
+ arrayCases.foreach { case (name, expr) =>
+ val df = arrays.select(expr.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for $name")
+ }
+
+ val mapCases = Seq(
+ "transform_keys" -> transform_keys(col("l"), (k, v) => pythonUDF(k).cast("string")),
+ "transform_values" -> transform_values(col("l"), (k, v) => pythonUDF(v)),
+ "map_filter" -> map_filter(col("l"), (k, v) => pythonUDF(v)),
+ "map_zip_with" -> map_zip_with(col("l"), col("r"), (k, a, b) => pythonUDF(a) || pythonUDF(b)))
+ mapCases.foreach { case (name, expr) =>
+ val df = maps.select(expr.as("r"))
+ assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for $name")
+ }
+ }
+
+ test("a pairwise array_sort comparator is lifted over the cross product") {
+ // One UDF call receiving both elements has no per-element key, so the UDF is precomputed over
+ // every ordered pair instead. The call must end up outside every lambda like any other.
+ val df = arrayDF.select(
+ array_sort(col("values"), (a, b) => pythonUDF2(a, b).cast("int")).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ val lifted = liftedUDFs(df)
+ assert(lifted.size == 1)
+ // The UDF takes both pair sides, each a flat array of all n*n pairs.
+ assert(lifted.head.children.size == 2)
+ }
+
+ test("a pandas UDF inside a lambda still fails analysis") {
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x => scalarPandasUDF(x))).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+
+ test("the rewrite can be disabled by conf, restoring the previous error") {
+ withSQLConf(SQLConf.PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED.key -> "false") {
+ val e = intercept[AnalysisException] {
+ arrayDF.select(transform(col("values"), x => pythonUDF(x))).collect()
+ }
+ assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF")
+ }
+ }
+
+ test("the rewrite cannot be disabled via excludedRules") {
+ // The lambda rewrite is driven by `ExtractPythonUDFs` (not a standalone batch rule), and that
+ // rule is non-excludable, so a plan that only works because of the rewrite cannot be broken by
+ // excludedRules. Excluding either name must leave the UDF lifted out of the lambda.
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ Seq(ExtractPythonUDFFromLambda.ruleName, ExtractPythonUDFs.ruleName).mkString(",")) {
+ val df = arrayDF.select(transform(col("values"), x => pythonUDF(x)).as("r"))
+ assert(udfsInsideLambda(df).isEmpty)
+ }
+ }
+
+ test("a UDF over only constants is still lifted per element") {
+ // SPARK-27052: `transform(arr, x -> udf(lit(10)))` does not read the element, so it is
Review Comment:
Fixed in 95ec08d. Updated the comment to say the constant-only UDF is now lifted per element (keeping the lambda's zero-or-N call domain) rather than left to `ExtractPythonUDFs`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-27052][SQL][PYTHON] Support Python UDFs inside higher-order function lambdas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #57804:
URL: https://github.com/apache/spark/pull/57804#discussion_r3736202185
##########
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:
Leaving this open for discussion rather than restricting. Spark already evaluates a Python UDF unconditionally when it sits under a JVM conditional: `ExtractPythonUDFs` pulls `when(cond, pyudf(x))` out into a separate operator evaluated before the enclosing expression, so a top-level guarded Python UDF is already eager today. Rejecting the in-lambda `when`/`if` case would make this feature stricter than plain Spark and would break the currently-passing `test_composition_around_udf_result`.
The genuinely-new over-evaluation is the short-circuiting shapes: `exists`/`forall` (stop early) and the pairwise `array_sort` (n^2 pairs, some never compared). Do you want those two specifically rejected when they carry a UDF that could throw, while keeping `when`/`if` composition consistent with existing Python-UDF semantics? Happy to implement whichever boundary you prefer.
##########
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:
For now I've documented the PyArrow requirement rather than adding a second transport: the config doc states the Arrow-based path needs PyArrow even for `useArrow=False`, and the worker calls `require_minimum_pyarrow_version()` up front so it fails with a clear `PACKAGE_NOT_INSTALLED` message instead of a bare ImportError.
A full non-Arrow (pickle) elementwise path - a new `SQL_BATCHED_ELEMENTWISE_UDF` eval type routed to `BatchEvalPython` with a mirrored flatten/re-nest worker branch - is doable but sizeable. Would you prefer that in this PR, or as a follow-up? Leaving this open for your call.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org