cloud-fan commented on code in PR #59749:
URL: https://github.com/apache/spark/pull/59749#discussion_r4360943401
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -244,16 +244,26 @@ trait HigherOrderFunction extends Expression with
ExpectsInputTypes {
override lazy val canonicalized: Expression = {
var currExprId = -1
- // Number the lambda variables of this higher-order function, but only
those not already
- // canonicalized. A canonical `NamedLambdaVariable` carries `value = null`
(see the rename
- // below); an original one carries an `AtomicReference`. When this HOF is
nested inside another,
- // the enclosing HOF's `canonicalized` renames every variable in the whole
subtree first, then
- // canonicalizes the children - which re-enters this method on the nested
HOF. Re-numbering the
- // already-renamed variables here would give a reference to an *enclosing*
lambda's variable a
- // different id than its binding, leaking it into `references`; skipping
them keeps a variable
- // and its references in agreement across nesting levels.
+ // Number the lambda variables bound in this higher-order function: the
arguments of its lambda
+ // functions, including those of nested higher-order functions. A free
variable, bound by an
+ // enclosing lambda function, keeps its exprId. Otherwise `filter(ys, y ->
y > x1)` and
+ // `filter(ys, y -> y > x2)` would be semantically equal whenever `x1` and
`x2` have the same
+ // data type and nullability.
+ //
+ // Only variables not already canonicalized are numbered. A canonical
`NamedLambdaVariable`
+ // carries `value = null` (see the rename below); an original one carries
an `AtomicReference`.
+ // When this HOF is nested inside another, the enclosing HOF's
`canonicalized` renames the
+ // variables bound in its whole subtree first, then canonicalizes the
children - which re-enters
+ // this method on the nested HOF. Its variables are then already numbered
uniquely across the
+ // nesting levels, and numbering them again here could make them collide
with the variables of
+ // the enclosing lambda functions.
+ val boundExprIds = functions.flatMap(_.collect {
+ case f: LambdaFunction => f.arguments.collect {
+ case l: NamedLambdaVariable if l.value != null => l.exprId
+ }
+ }.flatten).toSet
val argumentMap = functions.flatMap(_.collect {
- case l: NamedLambdaVariable if l.value != null =>
+ case l: NamedLambdaVariable if l.value != null &&
boundExprIds.contains(l.exprId) =>
Review Comment:
**Non-blocking (P2):** This guard leaves a captured variable untouched, so
it reaches the canonical form as the original `NamedLambdaVariable` with its
`value: AtomicReference` still attached. `NamedLambdaVariable` is a case class
with no `equals` override, and `AtomicReference` compares by identity. That
makes a captured variable's canonical identity depend on per-occurrence
evaluation state, not just its exprId.
That matters because the optimizer never sees the analyzed plan's shared
instances. `QueryExecution` optimizes
`cloneWithFreshStatefulExpressions(withCachedData)`, and since
`NamedLambdaVariable` is `stateful`, each occurrence becomes `copy(value = new
AtomicReference())`. Take this query on a Parquet table:
```sql
SELECT transform(xs, x -> exists(ys, y -> y > x) AND exists(ys, y -> y > x))
FROM t
```
On master, both `x` occurrences turn into the same value-null `none` node.
The two `exists` are then `semanticEquals`, and `BooleanSimplification` (`a And
b if a.semanticEquals(b) => a`) keeps only one. With this change, each `exists`
keeps its own `AtomicReference`, so they no longer compare equal and both run
for every element. The same loss hits `SimplifyConditionals` equal-branch
folding, `SimplifyBinaryComparison`, `ExpressionSet` dedupe, and the pre-lift
`seen` dedupe in `ExtractPythonUDFFromLambda`. It affects any pair of
structurally identical higher-order functions in a lambda body that capture the
same enclosing variable.
Query results stay correct. This is an optimization lost relative to master,
and the different-capture fix itself is right.
`CodeGenerator.withoutLambdaValues` already strips these values for the same
fresh-copy reason.
The new same-capture assertions in `HigherOrderFunctionsSuite` (`filter1`
and `nested1`) build both sides from the single `x1` instance, so they can't
catch this.
**Recommended change:** In HigherOrderFunction.canonicalized, keep renaming
the bound lambda variables exactly as this PR does, and also replace the
runtime value of every captured lambda variable (not bound in this function and
not yet canonical) with one shared non-null placeholder. A captured variable is
then identified only by its name, data type, nullability, and exprId. Add
higher-order function expression tests that compare same-capture functions
built from distinct copies of the captured variable, and a DataFrame
higher-order function test that checks the optimizer still folds same-capture
duplicates on a file source.
**Why this works:** The rename pass gets a second case after the
bound-variable case: a lambda variable with a non-null value whose exprId is
not in the bound set is copied with the shared placeholder as its value.
Occurrences of the same captured variable then compare equal whichever
AtomicReference each carried, which restores the merge-base equality for
same-capture functions, while different captured variables still differ by
exprId. The placeholder is non-null and renamed bound variables carry null, so
a captured variable can never equal a renamed bound variable, even when it is
named `none` and its exprId numerically matches a canonical id. The `value !=
null` re-entry marker keeps working: a nested function re-entered during the
enclosing function's canonicalization still skips the null-valued variables the
enclosing function renamed, and re-applying the placeholder to an already
normalized captured variable is idempotent.
**Scope:** Normalize the runtime value of captured lambda variables inside
higher-order-function canonicalization, and add same-capture regression
coverage at the expression and optimizer levels.
**Compatibility:** Higher-order functions that capture different lambda
variables must stay not semantically equal; that is the wrong-result fix this
PR delivers.
**Risks:** If the placeholder case runs before or instead of the
bound-variable rename, bound variables could be conflated with captured ones or
the nested-function skip could be defeated, so the bound-variable case must
take precedence in the rename.
**Constraints:** Canonical forms of higher-order functions that capture no
lambda variables, such as all top-level ones, must stay exactly as they are, so
plan reuse, cache matching, and grouping-expression matching are unaffected.
**Exclusions:**
- Making ordinary lambda-body expressions that read the same lambda variable
compare equal during optimization, such as `x + 1` versus `x + 1`, is a
pre-existing limitation of how a bare lambda variable compares, and this fix
leaves it unchanged.
**Success:** Two structurally identical higher-order functions that capture
the same enclosing lambda variable through distinct occurrences, as produced by
the optimizer's fresh stateful copy, are semantically equal and have equal
semantic hashes. Higher-order functions that capture different lambda variables
remain not semantically equal. A captured lambda variable never compares equal
to a renamed bound variable, including when it is named `none` and its exprId
equals a canonical id. Canonical forms of higher-order functions that capture
no lambda variables are unchanged.
**Verification:**
- **Regression:** Build two otherwise identical nested higher-order
functions whose captured-variable occurrences are distinct copies with fresh
runtime values, as the optimizer's fresh stateful copy produces, and assert
that they are semantically equal and have equal semantic hashes.
- **Behavior:** On a file-backed table, optimize a lambda body that combines
two same-capture higher-order functions, such as `exists(ys, y -> y > x) AND
exists(ys, y -> y > x)` inside `transform`, and assert that the optimized plan
keeps a single `exists`, while the four different-capture queries from this PR
still return their expected results.
- **Regression:** Assert that a captured lambda variable named `none`, whose
exprId equals the canonical id given to a bound variable of the same function,
does not make two different higher-order functions semantically equal.
- **Compatibility:** Confirm that the different-capture expression
assertions, the enclosing-function collision assertions, and the lifting of two
different-capture Python UDF calls as two UDFs still hold, and that canonical
forms of higher-order functions without captured variables are unchanged.
--
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]
