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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
