holdenk opened a new pull request, #58189:
URL: https://github.com/apache/spark/pull/58189
### What changes were proposed in this pull request?
The initial transpiled Python UDF contains a number of `_udf_param_N`
placeholders that are replaced by the bound argument, so an argument the body
reads N times is written out -- and evaluated -- N times. Instead we avoid N
time evaluation by:
- `UserDefinedPythonFunction.builder` marks every non-foldable substituted
argument with a new field on `TranspiledUDFParameter(child, index, id)`,
allocating one `ExprId` per parameter per call. The id, not structural
equality, is what ties a parameter's copies together: an argument whose seed
was still unresolved at call time (`expr("rand()")`, SQL text) is reseeded per
copy by `ResolveRandomSeed`, so the copies are no longer structurally equal.
- A new `PreEvaluateTranspiledUDFInputs`, which `ConvertToCatalyst` runs on
each plan node once that node's options have been substituted, turns those
marks into one aliased column (`_udf_input_N`) in a `Project` below the
operator, with every use reading that column. The plan surgery mirrors
`RewriteWithExpression.applyInternal`: pick the child that can compute the
column, add the `Project` there, and project the extra columns away again above.
- Which copies share a column follows what the interpreted UDF does with its
own argument columns. `EvalPythonEvaluatorFactory` reuses an input column for
any argument `semanticEquals` to an existing one, and `semanticEquals` is false
unless both sides are deterministic, so `f(a + 1, a + 1)` is reduced down to
one column and `f(rand(1), rand(1))` remains two. Deterministic params are
keyed on arguments; a nondeterministic one keys on the marker id.
- cheap arguments are in-lined instead of projected
An input stays inline, evaluated per use, where a column cannot help or
would not be safe:
| Shape | Why |
|---|---|
| a cheap argument -- a column, a literal, a struct field or fixed array
position | reading it twice costs nothing, and a literal is better folded at
each use site |
| an aggregate / window / generator argument | cannot live in a `Project`
(`PlanHelper.specialExpressionsInUnsupportedOperator`) |
| a lambda variable, or an argument spanning a join's children | no single
child of the operator can compute it |
| an argument carrying an `OuterReference` | it belongs to the enclosing
query |
| anything in a `GROUP BY` aggregation that no aggregate function encloses |
it has to match a grouping expression, and rewriting one side of that match
makes the `Aggregate` invalid |
| a nondeterministic argument where the operator does not emit one row per
input row (a join, `Generate`, `Expand`, a lateral join) | one draw would be
reused across every row the operator emits -- correlated, not redrawn per
output row |
| a `Command` | it has no output to widen or restore, so the schema
bookkeeping cannot check itself |
A call used **as a predicate** keeps the interpreted UDF unless its option
needs no column at all. `PushPredicateThroughNonJoin` inlines a column that is
not `Expression.expensive` -- which arithmetic is not -- back into the
predicate it pushes down, at every use site, which would both restore the
double evaluation and, with the input back inside the body's branches, let a
query that raises under interpreted Python return rows instead. So
`df.where(f(col("x")))` still avoids Python entirely, while
`df.where(f(col("a") / col("b")))` runs interpreted. A `Filter` condition, an
inner join condition and an inner lateral join condition are treated this way;
**any other join type keeps transpiling**, because
`ExtractPythonUDFFromJoinCondition` rejects a scalar Python UDF spanning both
sides of a non-inner join (`UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE`) --
there is no fallback to prefer when the fallback does not compile.
Other optimizer rules can put a *deterministic* input back resulting in
duplicate evaluation. However, the values never change -- `a / b` is `a / b`
however often you compute it -- and it's up to the other optimizer rules to
determine if inlining is a win or not.
A *nondeterministic* input's column is a guarantee: pushdown requires
`fields.forall(_.deterministic)` and `CollapseProject` never inlines a
nondeterministic producer. So the single `rand()` draw per row holds
unconditionally, while the single evaluation of `a / b` is a best effort.
Closing the deterministic case would need the hoisted column to report
`Expression.expensive` -- a new Catalyst node whose only purpose is to refuse
the pushdown -- and even that leaves `CollapseProject`'s force-inline branch,
which ignores cost. Not attempted here; the shape is pinned by a test instead.
### Why are the changes needed?
Multiple evaluations of the same param can be expensive.
### Does this PR introduce _any_ user-facing change?
Faster transpiled code
### How was this patch tested?
New and updated tests:
- `PreEvaluateTranspiledUDFInputsSuite` pins the plan shapes: one column for
a repeated input, a column for a single-use input, sharing across parameters
and across calls, reseeded copies sharing one draw, `f(rand(1), rand(1))`
keeping two, cheap inputs left alone, the nested-call case getting one column
per level, output schema preserved where an operator inherits its child's, and
each of the refusals in the table above.
- `ConvertToCatalystSuite` covers substitution leaving the markers for the
plan-level rewrite, and the predicate rules (a `Filter` and a join condition
keeping the interpreted UDF when a column is needed, transpiling when every
input is cheap, and a nested call deciding on its own inputs).
- `TranspiledUDFParameterSuite` covers what the builder marks and the ids it
hands out.
- `test_udf_transpile_unit` adds end-to-end coverage: inputs evaluated once,
the predicate rule (with a Python-free control in `select`), eagerness matching
interpreted Python, and the `GROUP BY` and subquery shapes.
### Was this patch authored or co-authored using generative AI tooling?
Yes claude 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: [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]
