You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "LuciferYang (via GitHub)" <gi...@apache.org> on 2026/08/17 09:15:59 UTC
[PR] [SPARK-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
LuciferYang opened a new pull request, #58045:
URL: https://github.com/apache/spark/pull/58045
### What changes were proposed in this pull request?
`RewriteWithExpression` inlines a `With` sitting in a conditional branch by blind substitution, without the reference-count and cheapness checks the main branch applies. When the common expression is nondeterministic, each reference then evaluates it separately and gets a different value, which is exactly what `With` exists to prevent.
This PR pre-evaluates such a common expression in a project instead, but only when it cannot raise. Pre-evaluating runs the expression for every row, including the rows whose branch is not taken, so an expression that raises would trade a wrong result for a spurious error.
That safety condition rules out more than it may seem. `randstr(-1, 0)` raises on its constant length, and `reflect('java.lang.Integer', 'parseInt', 'abc')` raises on constant arguments, so neither "the nondeterministic node is the root" nor "all children are foldable" is sufficient. The guard therefore admits `rand`/`randn`, whose only child is a seed the analyzer already requires to be foldable, and the nondeterministic leaves (`uuid`, `monotonically_increasing_id`, `spark_partition_id`, `input_file_name`, the input-file-block pair), which evaluate no arguments at all.
**This is a partial fix.** A nondeterministic expression wrapped in anything else, `rand() / col` for example, keeps the existing inlining and its existing wrong result. Covering it needs a way to evaluate a branch lazily, not a wider guard here — I measured what a wider guard does and it converts the wrong result into a `DIVIDE_BY_ZERO` on rows whose branch was never taken, which is not an improvement.
### Why are the changes needed?
The corrupt evaluation is a silent wrong result reachable from ordinary SQL with default configuration. `BETWEEN` references its input twice, and the parser produces `Between` by default (`spark.sql.legacy.duplicateBetweenInput` is false), so:
```sql
SELECT CASE WHEN id < 0 THEN false
ELSE monotonically_increasing_id() BETWEEN 3 AND 5 END AS r
FROM range(0, 10, 1, 1)
```
returns 6 `true` rows instead of 3. Each reference got its own counter, and the second one advanced only on the rows where the first predicate passed. The optimized plan shows the duplication:
```
Project [CASE WHEN (id < 0) THEN false
ELSE ((monotonically_increasing_id() >= 3) AND (monotonically_increasing_id() <= 5)) END]
```
`rand()` is wrong the same way. Over 200k rows, `CASE WHEN id < 0 THEN false ELSE rand(42) BETWEEN 0.4 AND 0.6 END` selects 36.0% of rows where the correct rate is 20.0%; written outside a conditional branch the same predicate gives 20.1%.
`NullIf` builds a `With` of the same shape, so it is affected wherever it appears in a conditional branch.
### Does this PR introduce _any_ user-facing change?
Yes, it fixes the wrong results above. A query whose conditional branch contains `rand`, `randn`, `uuid`, `monotonically_increasing_id`, `spark_partition_id` or `input_file_name` under a `With` now evaluates it once per row rather than once per reference.
One consequence is worth naming: those expressions are now evaluated for every row, including rows whose branch is not taken. For a generator this only advances the RNG or the counter, but it does mean `monotonically_increasing_id()` values and `rand()` draws shift relative to the previous behavior. That behavior was already wrong, so there is nothing to preserve, and the same expression outside a conditional branch has always been pre-evaluated this way.
### How was this patch tested?
New tests, plan-level and end-to-end, both written against the unfixed tree first and confirmed to fail there:
- `RewriteWithExpressionSuite`: one `comparePlans` test with five cases — `rand()` is pre-evaluated; a deterministic common expression in the same position is still inlined; `randstr(-1, 0)` stays inlined because it can raise; `rand() / a` stays inlined because it reads row data; a single reference stays inlined because it evaluates once either way.
- `ColumnExpressionSuite`: a `checkAnswer` test on the `monotonically_increasing_id()` shape above, with the non-branch form as a control.
On the unfixed tree the catalyst test reports the two `rand(1)` copies, and the end-to-end test reports 3 wrong rows out of 10.
Mutation-tested each half of the guard: flipping the fallback to admit everything, dropping the `rand`/`randn` arm, and dropping the reference-count check each fail at least one of the five cases.
Also verified with throwaway probes (deleted) that a `randstr(-1, 0)` and a `reflect(...)` in a never-taken branch behave the same before and after, that the fix produces the right shape under `Filter`, a join condition and `Aggregate`, and that a nested nondeterministic `With` terminates and is idempotent.
Regressions: `catalyst/testOnly org.apache.spark.sql.catalyst.optimizer.* NullExpressionsSuite ConditionalExpressionSuite RandomSuite` (1437 passed), `sql/testOnly org.apache.spark.sql.ColumnExpressionSuite org.apache.spark.sql.DataFrameFunctionsSuite` (315 passed).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude 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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5508819001
can you fix merge conflcts?
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5509596288
> can you fix merge conflcts?
done
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5602051167
Thank you @cloud-fan @dongjoon-hyun and @uros-b
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3805954200
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -98,6 +98,70 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether pre-evaluating this common expression in a conditional branch is safe, where inlining
+ * is otherwise preferred because the branch may not be evaluated at all.
+ *
+ * It is safe only when the expression cannot raise. Pre-evaluating happens for every row,
+ * including the rows whose branch is not taken, so an expression that raises would turn a wrong
+ * result into a spurious error. That rules out more than it may seem: `randstr(-1, 0)` raises on
+ * its constant length and `reflect(...)` raises on constant arguments, so neither a
+ * nondeterministic root nor foldable children are enough. The generators below produce a value
+ * from a seed or from the task context, evaluate no argument beyond a foldable seed, and so have
+ * nothing to raise on. Any other nondeterministic expression, say `rand() / col`, keeps the
+ * existing inlining and its existing wrong result.
+ *
+ * Only nondeterministic expressions are worth pre-evaluating here at all: inlining a
+ * deterministic one repeats the work but returns the same value, so the branch keeps deciding
+ * whether it runs.
+ */
+ private def canPreEvaluateInBranch(child: Expression): Boolean = {
+ val cannotRaise = child match {
+ // `rand`/`randn`: the seed is their only child and the analyzer rejects a non-foldable one
+ // with `SEED_EXPRESSION_IS_UNFOLDABLE`, so there is no row data left to raise on.
+ case _: NondeterministicUnaryRDG => true
+ // `uuid`, `monotonically_increasing_id`, `spark_partition_id`, `input_file_name` and the
+ // input-file-block pair: no children to evaluate at all.
+ case _: LeafExpression => true
+ case _ => false
+ }
+ !child.deterministic && cannotRaise
+ }
+
+ /**
+ * Adds `child` to the input plan that supplies its references as a pre-evaluated column, and
Review Comment:
**Nit:**
```suggestion
* Adds `child` as a pre-evaluated column to the input plan that supplies `child`'s references, and
```
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5405711425
> We should be careful when we change the expression evaluation order or invocation times. Do the existing expression properties sufficient to make safe judgement here? e.g. stateful, throwable, deterministic, etc.
They are not sufficient, and investigating that is what changed this PR. It no longer pre-evaluates anything, so there is no such judgement left to make.
On the properties: the bug set coincides with `stateful`. `spark_partition_id()` inlined is still correct, while `monotonically_increasing_id()` returns 6 `true` rows where 3 is right, so excluding stateful expressions from pre-evaluation would have left the PR empty. `throwable` is not available as a property either. I tried deciding it from the expression, and what that admits is too narrow to cover the bug: `rand() / col` and `randstr(3, 0)` kept the wrong result.
What the PR does now is make `With` evaluable and memoize each common expression for the row being evaluated. Nothing is relocated: a definition is evaluated where the child would have evaluated it, so an expression that raises raises on exactly the rows it did before, and a generator advances on exactly the rows that reach a reference. Invocation times change in one direction only, from once per reference to once. A reference's domain is in fact narrower than its branch's — `BETWEEN` builds an `And`, and the second reference is not read on a row whose first comparison is false, which pre-evaluation cannot express.
If this approach looks acceptable, I will file a follow-up for the two shapes that are still substituted: a common expression in a join condition that references both sides, and one whose definition cannot be put in a `Project`.
Could you take another look when you have time? Thanks.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang closed pull request #58045: [SPARK-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch
URL: https://github.com/apache/spark/pull/58045
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5586693462
Thanks, all five are real. Fixed in 1133af6 and 85fd1a4.
**P2, fallback initialization.** Right, and it lands on the one thing this PR is the first to depend on: a `With` is not a `CodegenFallback`, so when it falls back it relies on `CodegenFallback.generate` walking the subtree and registering every `Nondeterministic` for partition initialization. Both tests now build `Counter()` directly and let `proj.initialize(0)` reach it. Checked that this is load-bearing: narrowing the registration to direct children fails the nested case with `Nondeterministic expression ... should be initialized before eval`.
**P2, migration guide.** Agreed, the entry promised more than `canSubstitute` delivers. It now says what changed -- a nondeterministic common expression read more than once in a branch is evaluated once -- and names the exception: a definition that is cheap and reports itself deterministic is still substituted at every read. The `aes_encrypt` case you name is that exception; `canSubstitute`'s scaladoc records it, along with why the fix belongs on those expressions or in a purity predicate shared with `isSafeToDuplicate` rather than here. The PR description is narrowed the same way.
**P3, zero references.** Correct, `!multiplyReferenced.contains(id)` admits a definition with no references at all. Changed to `referenced at most once` in `canSubstitute` and `inlineDefsThatGainNothing`, and in the description.
**P3, PythonUDF agreement.** Fixed.
**P3, "the failure to want".** That sentence is gone: 1133af6 cut the comments in these files roughly in half (they had grown to about 50% of the added lines across five review rounds), and this paragraph went with it. Nothing to reword.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3862864751
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -162,7 +319,80 @@ case class CommonExpressionDef(child: Expression, id: CommonExpressionId = new C
* referenced, so that we can determine the data type and nullable of the reference node.
*/
case class CommonExpressionRef(id: CommonExpressionId, dataType: DataType, nullable: Boolean)
- extends LeafExpression with Unevaluable {
+ extends LeafExpression {
def this(exprDef: CommonExpressionDef) = this(exprDef.id, exprDef.dataType, exprDef.nullable)
+
+ /**
+ * The definition this reference names, and the cell holding its value for the current row. Both
+ * are wired by the enclosing [[With]] before it evaluates its child, and are left out of the case
+ * class parameters so that equality and canonicalization are unchanged -- and so that a rule
+ * comparing two references does not compare their cells.
+ */
+ private var definition: Expression = _
+ private var cell: CommonExpressionCell = _
+
+ private[expressions] def bindTo(exprDef: CommonExpressionDef): Unit = {
+ definition = exprDef.child
+ cell = exprDef.cell
+ }
+
override val nodePatterns: Seq[TreePattern] = Seq(COMMON_EXPR_REF)
+
+ // The cell is cleared by the enclosing `With` on every entry, so this reads mutable state.
+ override def stateful: Boolean = true
+
+ /**
+ * A copy must not carry this reference's binding: the copy belongs to a different `With`, which
+ * wires it to its own cell. `LeafLike` returns `this` here, which would hand two `With`s one
+ * reference object and let whichever wires last decide what both of them read --
+ * `NamedLambdaVariable` overrides this for the same reason.
+ */
+ override def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): CommonExpressionRef = copy()
+
+ override def eval(input: InternalRow): Any = {
+ if (cell == null) {
+ throw SparkException.internalError(
+ s"Cannot evaluate a common expression reference outside its With: $this")
+ }
+ cell.get(definition, input)
+ }
+
+ /**
+ * Computes the definition into the shared slots if this row has not done so yet, then reads them.
+ * The definition's code is emitted here rather than by the enclosing `With`, so it runs where the
+ * first reference is reached -- behind a short-circuiting operator or a nested conditional, if
+ * that is where the reference sits.
+ *
+ * A second reference emits the same code text again, which never runs because the flag is set.
+ * The text is generated once and cached on the slots, so every copy shares whatever mutable
+ * state the definition allocated, and a nested `With` does not grow its code by a factor per
+ * level. The definition's locals are declared inside each guard, whose blocks are siblings, so
+ * repeating the text declares nothing twice in one scope.
+ *
+ * Whether the isNull slot exists is decided by the definition, so it is read off the slot rather
+ * than off this reference's own `nullable`: taking it from both would let the two disagree, and
+ * either emit `false = <isNull>;`, which does not compile, or leave the slot holding the previous
+ * row's nullness.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val slots = ctx.getCommonExpr(id.id)
+ val defGen = slots.definitionGen(ctx)
+ val assignIsNull = if (slots.value.isNull == FalseLiteral) {
+ ""
+ } else {
+ s"${slots.value.isNull} = ${defGen.isNull};"
+ }
+ ev.copy(
+ code = code"""
+ |if (!${slots.computed}) {
+ | ${defGen.code}
Review Comment:
**Blocking:**
Please emit one reusable lazy-definition block per `With` scope instead of pasting `defGen.code` at every reference. Nested `nullif` leaves nested `With` nodes, so each level's two references duplicate the complete inner body (`T(n) = 2T(n-1) + O(1)`). Whole-stage codegen cannot use `Expression.reduceCodeSize` here because it has `currentVars` and no `INPUT_ROW`; a deep nested-`nullif` compile or source-size regression would catch this.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala:
##########
@@ -25,13 +25,32 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block._
*/
trait CodegenFallback extends Expression {
- protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ CodegenFallback.generate(this, ctx, ev)
+}
+
+object CodegenFallback {
+
+ /**
+ * Generates code that evaluates `e` by calling its `eval`, rather than by generating code for it.
+ * This is what [[CodegenFallback]] gives the expressions that mix it in, and is also used by an
+ * expression that can generate code in general but has to fall back for a particular shape of
+ * its own tree -- see `With.doGenCode`.
+ *
+ * Three places reason about which subtrees run interpretively, and all of them dispatch on the
+ * trait: `CollapseCodegenStages.supportCodegen`, which turns whole-stage codegen off;
+ * `EquivalentExpressions.childrenToRecurse`, which keeps a subtree that only `eval` reaches out
+ * of subexpression elimination; and `With.refUnderCodegenFallback`, which decides whether a
+ * memoized reference is reached that way. A caller that is not a `CodegenFallback` is invisible
+ * to all three, so it has to be named in each of them as well.
Review Comment:
**Nit:**
Please narrow this to the sites that cannot discover the caller structurally. `With` calls `generate` without mixing in `CodegenFallback`, but `CollapseCodegenStages` still finds the descendant fallback through its recursive `plan.expressions.exists(_.exists(...))`; a non-trait caller therefore does not necessarily have to be named at all three sites.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -430,6 +430,136 @@ class ColumnExpressionSuite extends SharedSparkSession {
checkAnswer(testData.filter($"a".between($"b", $"c")), expectAnswer)
}
+ // Runs `f` on each of the three evaluation paths, since a `With` left in a conditional branch is
+ // evaluated by all three and the memoization is implemented separately for interpretation and for
+ // codegen.
+ private def onEachEvalPath(f: => Unit): Unit = {
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN")(f)
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
+ }
+
+ test("SPARK-58902: BETWEEN on a nondeterministic input inside a conditional branch") {
+ onEachEvalPath {
+ // `BETWEEN` reads its input twice. Inlining the common expression into a branch gave each
+ // read its own value, so the two comparisons saw two different ids and 6 of the 10 rows came
+ // back true where 3 is correct. A single partition makes the id sequence 0, 1, 2, ...
+ val df = spark.range(0, 10, 1, 1)
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN id < 0 THEN false ELSE monotonically_increasing_id() BETWEEN 3 AND 5 END"),
+ (0 until 10).map(i => Row(i >= 3 && i <= 5)))
+ }
+ }
+
+ test("SPARK-58902: a definition in a branch is evaluated only on the rows that reach it") {
+ onEachEvalPath {
+ // Rows 0 to 4 take the first branch, so the five rows that reach the ELSE see ids 0 to 4 --
+ // the same values they would see if the branch were the whole expression.
+ val df = spark.range(0, 10, 1, 1)
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN id < 5 THEN NULL ELSE monotonically_increasing_id() BETWEEN 1 AND 2 END"),
+ (0 until 10).map { i =>
+ if (i < 5) Row(null) else Row(i - 5 >= 1 && i - 5 <= 2)
+ })
+ }
+ }
+
+ test("SPARK-58902: a branch condition that can raise is not evaluated on other rows") {
+ onEachEvalPath {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ // Nothing is relocated, so a condition is evaluated only where it always was. `a` runs
+ // -2, -1, 0, 1, 2, 3: the first two take the third branch, 0 the first, 1 and 2 the second,
+ // and 3 falls through. `rand` is in [0, 1), so a row reaching a BETWEEN is true whatever it
+ // draws -- the answers here do not depend on memoization. What this rules out is the
+ // alternative that was measured and rejected: hoisting the definition into a `Project` and
+ // guarding that column with the branch condition, which puts `6 / a` outside conditional
+ // evaluation and raises on the `a = 0` row under ANSI.
+ val df = spark.range(0, 6, 1, 1).selectExpr("cast(id as int) - 2 as a")
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN a = 0 THEN false " +
+ "WHEN 6 / a > 2 THEN rand(1) BETWEEN 0 AND 1 " +
+ "WHEN 6 / a < -2 THEN rand(2) BETWEEN 0 AND 1 " +
+ "ELSE false END"),
+ Seq(Row(true), Row(true), Row(false), Row(true), Row(true), Row(false)))
+ }
+ }
+ }
+
+ test("SPARK-58902: a nondeterministic input a branch cannot pre-evaluate is still read once") {
+ onEachEvalPath {
+ // Enough rows that the two copies an inlining implementation makes have to fall out of step:
+ // each copy owns its own generator seeded the same way, so they only differ once the first
+ // comparison has skipped the second copy on some row.
+ val df = spark.range(0, 20, 1, 1)
+ // `randstr(3, 0)` cannot be pre-evaluated into a project, because it raises on a negative
Review Comment:
**Nit:**
Please replace the negative-length rationale with the nondeterministic memoization behavior this test actually exercises. This invocation's foldable length is literal `3`, so it cannot raise for negative length; codegen also calls `lengthInteger()` while generating code rather than selectively by row.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,121 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with the definition each
+ * names. The list is found once, since the tree does not change between evaluations.
+ *
+ * Only `child` is scanned, which relies on a reference to one of these definitions never living
+ * inside another one of them: `With.apply` cannot build that, and `RewriteWithExpression`, which
Review Comment:
**Nit:**
Please narrow this sentence to the curried `With(commonExprs: _*)(replaced)` helper. The case class also exposes direct `apply(child, defs)` construction, so a definition can contain a reference to its own id; because `refsToBind` scans only `child`, that nested reference remains unbound and fails during evaluation.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3966357431
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -162,7 +354,85 @@ case class CommonExpressionDef(child: Expression, id: CommonExpressionId = new C
* referenced, so that we can determine the data type and nullable of the reference node.
*/
case class CommonExpressionRef(id: CommonExpressionId, dataType: DataType, nullable: Boolean)
- extends LeafExpression with Unevaluable {
+ extends LeafExpression {
def this(exprDef: CommonExpressionDef) = this(exprDef.id, exprDef.dataType, exprDef.nullable)
+
+ /**
+ * The definition this reference names, and the cell holding its value for the current row. Both
+ * are wired by the enclosing [[With]] before it evaluates its child, and are left out of the case
+ * class parameters so that equality and canonicalization are unchanged -- and so that a rule
+ * comparing two references does not compare their cells.
+ *
+ * `@transient` for the reason [[CommonExpressionCell]]'s own fields are: a binding lives only for
+ * the duration of one `With.eval`, which restores it on the way out, so serializing a task sees
+ * null. A serialization that happens while an evaluation is still on the stack would otherwise
+ * capture a live binding, and the deserialized reference would evaluate that definition rather
+ * than raise, dragging the definition's whole subtree along with it.
+ */
+ @transient private var definition: Expression = _
+ @transient private var cell: CommonExpressionCell = _
+
+ private[expressions] def bindTo(exprDef: CommonExpressionDef): Unit = {
+ definition = exprDef.child
+ cell = exprDef.cell
+ }
+
+ private[expressions] def boundDefinition: Expression = definition
+ private[expressions] def boundCell: CommonExpressionCell = cell
+
+ private[expressions] def bindTo(
+ newDefinition: Expression,
+ newCell: CommonExpressionCell): Unit = {
+ definition = newDefinition
+ cell = newCell
+ }
+
override val nodePatterns: Seq[TreePattern] = Seq(COMMON_EXPR_REF)
+
+ // The cell is cleared by the enclosing `With` on every entry, so this reads mutable state.
+ override def stateful: Boolean = true
+
+ /**
+ * A copy must not carry this reference's binding: the copy belongs to a different `With`, which
+ * wires it to its own cell. `LeafLike` returns `this` here, which would hand two `With`s one
+ * reference object and let whichever wires last decide what both of them read --
+ * `NamedLambdaVariable` overrides this for the same reason.
+ */
+ override def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): CommonExpressionRef = copy()
+
+ override def eval(input: InternalRow): Any = {
+ if (cell == null) {
+ throw SparkException.internalError(
+ s"Cannot evaluate a common expression reference outside its With: $this")
+ }
+ cell.get(definition, input)
+ }
+
+ /**
+ * Computes the definition into the shared slots if this row has not done so yet, then reads them.
+ * The code that computes it is emitted here rather than by the enclosing `With`, so it runs where
+ * the first reference is reached -- behind a short-circuiting operator or a nested conditional,
+ * if that is where the reference sits.
+ *
+ * A second reference emits the same code again, which never runs because the flag is set. What
+ * that code is depends on the definition: a call, where the definition can be put in a method, so
+ * that a definition that is or holds another `With` is not pasted once per reference at every
+ * level; the body itself otherwise, whose locals are declared inside each guard. No copy of one
+ * body encloses another -- for that, a definition would have to reference its own id, directly or
Review Comment:
**Nit (P3):** An acyclic sibling dependency can enclose another fill body without recursing: `withCommonExprs` registers every sibling slot before generating the child, so `d2 = Add(ref(d1), 1)` resolves `d1` and completes normally. Only self-reference or a dependency cycle re-enters a slot already being filled. Could we state that distinction here and narrow the earlier `refsToBind` claim that codegen refuses every definition-reference shape?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,163 @@ class CodegenContext extends Logging {
throw QueryExecutionErrors.lambdaVariableNotDefinedError(id))
}
+ /**
+ * The slots a `CommonExpressionRef` reads: the value and its nullness, plus the flag saying
+ * whether this row has computed them yet, and the definition to compute them from.
+ */
+ case class CommonExprSlots(
+ value: ExprCode,
+ computed: String,
+ definition: Expression) {
+
+ // Not constructor parameters: a mutable one takes part in `equals`/`hashCode`/`copy`, so a
+ // slot's hash would change as it is filled and a `copy` would carry another scope's code.
+ private var fillCode: Option[Block] = None
+ private var filling: Boolean = false
+
+ /**
+ * The code that computes the definition into the slots and sets `computed`, which a reference
+ * emits behind that flag. Cached, so every reference shares whatever mutable state the
+ * definition allocated, such as an RNG.
+ *
+ * The cache lives on the slot, so it lasts exactly as long as the scope: the code names the
+ * `INPUT_ROW` and `currentVars` in effect when it was generated. `GenerateOrdering` generates
+ * its key once per comparison side under a different row variable, so a slot shared between the
+ * sides would read the wrong row, or not compile where `Expression.reduceCodeSize` has hoisted
+ * the reference into a method taking one row.
+ *
+ * `filling` catches a definition that references its own id, which would otherwise re-enter and
+ * recurse, since `fillCode` is set only after `definition.genCode` returns.
+ *
+ * The body goes into a method where it can and is worth it -- a definition that is or holds
+ * another `With`, or a body past the split threshold -- so it is emitted once per scope rather
+ * than once per reference, which for nested `With`s would double per level. A method is only
+ * possible where the definition reads the input row rather than local variables, the condition
+ * `reduceCodeSize` splits under. That is not the same as whole-stage codegen being off: a
+ * whole-stage `Project` or `Filter` passes local variables, while
+ * `SortMergeJoinExec.createJoinKey` and the aggregate output paths generate against a row.
+ */
+ def fill: Block = {
+ if (fillCode.isEmpty) {
+ if (filling) {
+ throw SparkException.internalError(
+ "Cannot generate a common expression whose definition references it: " +
+ definition.toString)
+ }
+ filling = true
+ try {
+ fillCode = Some(build)
+ } finally {
+ filling = false
+ }
+ }
+ fillCode.get
+ }
+
+ private def build: Block = {
+ val defGen = definition.genCode(CodegenContext.this)
+ // Whether the isNull slot exists is decided by the definition, so it is read off the slot
+ // rather than off a reference's own `nullable`: taking it from both would let the two
+ // disagree, and either emit `false = <isNull>;`, which does not compile, or leave the slot
+ // holding the previous row's nullness.
+ val assignIsNull = if (value.isNull == FalseLiteral) {
+ ""
+ } else {
+ s"${value.isNull} = ${defGen.isNull};"
+ }
+ val body = code"""
+ |${defGen.code}
+ |$assignIsNull
+ |${value.value} = ${defGen.value};
+ |$computed = true;
+ """.stripMargin
+ // TODO(SPARK-59295): cover the local-variable case too, by passing the `currentVars` values a
+ // definition reads into the method as parameters, the way
+ // `subexpressionEliminationForWholeStageCodegen` does. It needs a decision first:
+ // `getLocalInputVariableValues` hoists an input variable that is not evaluated yet to
+ // before the call, which for a reference behind a branch means evaluating it on rows that
+ // never reach the reference.
+ val canPutInMethod = INPUT_ROW != null && currentVars == null
+ // A definition that is or holds another `With` is the shape whose code doubles per level,
+ // and what this is aimed at. It is not the only one -- a definition referencing a sibling
+ // definition of the same `With` doubles the same way, and codegen accepts that, since the
+ // sibling's slots are in scope while this definition is generated (`With.refsToBind` says
+ // why nothing builds that tree, and that evaluating one raises). What bounds those is not
+ // the length arm below: `body` is assembled after `definition.genCode` already ran
+ // `reduceCodeSize`, so the arm fires only in the band just under the threshold. It is
+ // `reduceCodeSize` itself, which hoists whichever node's code first passes the threshold as
+ // generation walks up, capping what one level contributes, so the code stays linear in the
+ // depth either way. The length arm just keeps the same body from being split once per
+ // reference, which leaves the methods small and the code as large.
+ val worthAMethod = definition.exists(_.isInstanceOf[With]) ||
Review Comment:
**Nit (P3):** `With` already sets `WITH_EXPRESSION` in `nodePatterns`, and those bits include descendants, so this recursive `exists` walk repeats work for every definition. Could we import the pattern and use `definition.containsPattern(WITH_EXPRESSION)` here? It preserves the same root-or-descendant test without revisiting the expression subtree.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +70,154 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with the definition each
+ * names, found once since the tree does not change between evaluations. Only `child` is scanned,
+ * which is why the checks below refuse a definition holding a reference of this scope.
+ */
+ @transient private lazy val refsToBind: IndexedSeq[(CommonExpressionRef, CommonExpressionDef)] = {
+ // Three shapes whose bindings cannot be kept straight, each of which this path used to answer
Review Comment:
**Nit (P3):** `this path used to answer for with` is not grammatical, so the outcome of these three shapes is hard to follow. Could this say that the path previously evaluated them incorrectly or overflowed the stack?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WithExpressionEvalSuite.scala:
##########
@@ -0,0 +1,372 @@
+/*
+ * 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.catalyst.expressions
+
+import org.apache.spark.{SparkException, SparkFunSuite}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodegenFallback, ExprCode, GenerateMutableProjection}
+import org.apache.spark.sql.catalyst.plans.SQLHelper
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType}
+
+/**
+ * Evaluation of [[With]] and the memoization it gives a [[CommonExpressionRef]]. The rewrite that
+ * decides which `With`s reach evaluation at all is covered by `RewriteWithExpressionSuite`.
+ */
+class WithExpressionEvalSuite extends SparkFunSuite with SQLHelper {
+
+ /**
+ * A stand-in for a stateful generator: every evaluation returns the next integer, so a second
+ * evaluation within one row is directly observable. Only used on the interpreted path.
+ */
+ private case class Counter() extends LeafExpression with Nondeterministic {
+ @transient private var n = 0
+ override def stateful: Boolean = true
+ override def dataType: DataType = IntegerType
+ override def nullable: Boolean = false
+ override protected def initializeInternal(partitionIndex: Int): Unit = {}
+ override protected def evalInternal(input: InternalRow): Any = { n += 1; n }
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ throw new UnsupportedOperationException
+ }
+
+ private def counter(): Counter = {
+ val c = Counter()
+ c.initialize(0)
+ c
+ }
+
+ /**
+ * A node that has to be evaluated interpretively even when its parent is generated, so that a
+ * reference below it takes the interpreted path out of generated code.
+ */
+ private case class Fallback(child: Expression) extends UnaryExpression with CodegenFallback {
+ override def dataType: DataType = child.dataType
+ override def eval(input: InternalRow): Any = child.eval(input)
+ override protected def withNewChildInternal(newChild: Expression): Fallback =
+ copy(child = newChild)
+ }
+
+ test("a definition is evaluated once per row however many references read it") {
+ // `ref + ref` is the shape `BETWEEN` produces. Memoized, both references read one value, so the
+ // sum is 2n on the nth row; inlined it would be n + (n + 1).
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ assert((1 to 4).map(_ => w.eval(InternalRow.empty)) == Seq(2, 4, 6, 8))
+ }
+
+ test("a definition is not evaluated on a row that reaches no reference") {
+ val c = counter()
+ // The branch is inside the `With`, so the `With` is entered on every row and does clear its
+ // cells. What it must not do is evaluate the definition before a reference is reached; putting
+ // the branch outside would only test that `If` does not evaluate the arm it did not take.
+ val w = With(c) { case Seq(ref) => If(Literal.TrueLiteral, Literal(-1), Add(ref, ref)) }
+ // -1 whatever the definition does, since the taken arm is a literal: the assertion below is the
+ // load-bearing one, and this only says the evaluation ran.
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(-1, -1, -1))
+ // The counter is still at 0, so the first row that does reach a reference sees 1. Eager filling
+ // would have consumed three values and this would read 8.
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty) == 2)
+ }
+
+ test("a reference behind a short-circuiting operator is not read when the left side is false") {
+ // Neither pre-evaluating the definition into a project nor guarding that column by the branch
+ // can express this: both evaluate on every row the branch is reached on, while `And`
+ // short-circuits before the reference. The `And` is inside the `With` so that the `With` is
+ // entered and only the reference is skipped.
+ val c = counter()
+ val w = With(c) { case Seq(ref) =>
+ And(Literal.FalseLiteral, GreaterThanOrEqual(ref, Literal(1)))
+ }
+ // False whatever the definition does, as in the case above: the counter read is what bites.
+ // This is not that case repeated -- there the arm is skipped by `If`, here by `And`'s short
+ // circuit, and the two are separate paths in both evaluation and codegen.
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(false, false, false))
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty) == 2)
+ }
+
+ test("a nested With memoizes each scope separately") {
+ val outer = counter()
+ val inner = counter()
+ // The outer reference is read twice, and one of its uses wraps an inner `With` whose own
+ // reference is also read twice, so each row is `o + (2i + o)` with `o` and `i` both n.
+ val w = With(outer) { case Seq(o) =>
+ Add(o, With(inner) { case Seq(i) => Add(Add(i, i), o) })
+ }
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(4, 8, 12))
+ }
+
+ test("a With rebuilt by a rule still memoizes") {
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ // Evaluate the original first. A rule rebuilds the definition, and therefore its cell, while
+ // handing the new `With` the original's reference objects, so an implementation that bound the
+ // references once would leave the rebuilt `With` reading this evaluation's cell.
+ assert(w.eval(InternalRow.empty) == 2)
+ // The rule has to hand back a node that is not `==` the one it replaces, or `transformUp` keeps
+ // the original -- `Counter()` is a case class with no parameters, so a fresh one compares equal
+ // to the old one. Wrapping it changes the tree while leaving the definition's value sequence,
+ // its data type and its nullability alone, so the references are carried over rather than
+ // rebuilt, which is the shape that needs the per-evaluation rebinding.
+ val rewritten = w.transformUp { case _: Counter => Add(counter(), Literal(0)) }
+ .asInstanceOf[With]
+ assert(rewritten ne w, "the rule has to have rebuilt something")
+ assert((1 to 3).map(_ => rewritten.eval(InternalRow.empty)) == Seq(2, 4, 6))
+ }
+
+ test("SPARK-58902: a copy of a With owns its references and its cells") {
+ // A deliberately non-stateful definition, which is what makes `CommonExpressionDef.stateful`
+ // load-bearing: `mapChildren` finds nothing changed below it, so without the override the copy
+ // gets the original definition object, cell included. (A `Rand` definition would hide that.)
+ // Likewise `CommonExpressionRef.stateful`, since `LeafLike.withNewChildrenInternal` returns
+ // `this` and the two `With`s would otherwise share references while owning two cells.
+ val w1 = With(Literal(1)) { case Seq(ref) => Add(ref, ref) }
+ val w2 = w1.freshCopyIfContainsStatefulExpression().asInstanceOf[With]
+ def refOf(e: Expression): CommonExpressionRef =
+ e.collect { case r: CommonExpressionRef => r }.head
+ assert(w1 ne w2, "a stateful With has to be copied")
+ assert(refOf(w1.child) ne refOf(w2.child), "the copy has to own its references")
+ assert(w1.defs.head.cell ne w2.defs.head.cell, "the copy has to own its cell")
+ }
+
+ test("SPARK-58902: two Withs over one set of references each read their own definition") {
+ // The shape a rule produces when it rewrites only the definition: the new `With` is built over
+ // the original's references, since a rebuilt reference compares equal to the one it replaces.
+ // Visibly different definitions show which cell the references actually read.
+ val w1 = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ val w2 = w1.withNewChildren(
+ IndexedSeq(w1.child, CommonExpressionDef(Literal(100), w1.defs.head.id))).asInstanceOf[With]
+
+ assert(w1.eval(InternalRow.empty) == 2, "w1 reads its own counter")
+ assert(w2.eval(InternalRow.empty) == 200, "w2 reads its own literal")
+ assert(w1.eval(InternalRow.empty) == 4, "w1 read the value w2 memoized")
+ }
+
+ test("SPARK-58902: a definition's code is generated once however many references read it") {
+ // Short and holding no `With`, so the body is emitted inline at each reference. It has to be
+ // the same text, generated once: otherwise each copy calls `addMutableState` again and the
+ // definition owns one counter per reference, which for references in mutually exclusive
+ // positions hands out one value on two rows.
+ val ctx = new CodegenContext
+ With(MonotonicallyIncreasingID()) { case Seq(ref) => Add(ref, ref) }.genCode(ctx)
+ val counters = ctx.inlinedMutableStates.count { case (_, name) => name.startsWith("count") }
+ assert(counters == 1, s"the definition was generated $counters times")
+ }
+
+ test("SPARK-58902: nesting does not multiply a definition's body when it can go in a method") {
+ // Each level reads its definition twice, so pasting the body doubles it per level;
+ // `CommonExprSlots.fill` putting it in a method makes it 2 at any depth. Not covered here: the
+ // shape where the input arrives as local variables, as a whole-stage `Project` or `Filter`
+ // passes it, where no method is possible and the count is 2, 4, 8, ... 256 at depths 1 to 8 --
+ // see SPARK-59295. A bare `CodegenContext` has `INPUT_ROW` set and `currentVars` null.
+ val marker = 1234567
+ def nested(depth: Int): Expression = {
+ val leaf: Expression = Add(BoundReference(0, IntegerType, nullable = false), Literal(marker))
+ (1 to depth).foldLeft(leaf) { (inner, _) =>
+ With(inner) { case Seq(ref) => Add(ref, ref) }
+ }
+ }
+ def markerCount(depth: Int): Int = {
+ val ctx = new CodegenContext
+ val source = nested(depth).genCode(ctx).code.toString + ctx.declareAddedFunctions()
+ marker.toString.r.findAllMatchIn(source).size
+ }
+ // Pin the threshold: below the innermost body's length the fill would go into a method of its
+ // own and the count would be 1, for a reason unrelated to nesting.
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1024") {
+ (1 to 6).foreach { depth =>
+ val count = markerCount(depth)
+ assert(count == 2, s"the innermost body was emitted $count times at depth $depth")
+ }
+ // The 2 is the innermost body pasted at its two references, at every depth; what stops the
+ // doubling from depth 2 on is that each enclosing definition is itself a `With` and so goes
+ // into a method, whose call sites carry no marker.
+ //
+ // The values double per level, which a deterministic leaf does whether a reference reads a
+ // slot or recomputes -- an arithmetic and compile check, not a second reading of the count.
+ // Four levels keeps the product inside Int, which ANSI `Add` would raise on past depth 10.
+ val proj = GenerateMutableProjection.generate(Seq(nested(4)))
+ assert(proj(InternalRow(1)).getInt(0) == (1 + marker) * 16)
+ }
+ }
+
+ test("SPARK-58902: a nested With sharing a reference object restores the outer binding") {
+ // Two `With`s over one reference object, the inner one nested in the outer one's child and
+ // redefining the id the outer one defines. Only a caller building the case class directly
+ // produces this -- `withNewChildrenInternal` shares references between a `With` and its
+ // replacement, which are siblings, and no rule nests a redefinition of a live id -- so this is
+ // hardening rather than a reachable wrong answer. Binding on entry alone is not enough: the
+ // inner `With` rebinds the shared reference to its own definition, so a read after the inner
+ // one returned would answer with the inner definition unless the outer binding is put back.
+ val id = new CommonExpressionId()
+ val outerDef = CommonExpressionDef(counter(), id)
+ val innerDef = CommonExpressionDef(Literal(10), id)
+ val ref = new CommonExpressionRef(outerDef)
+ val inner = With(Add(ref, ref), Seq(innerDef))
+ val outer = With(Add(Add(ref, inner), ref), Seq(outerDef))
+ // The outer definition counts, the inner one does not, so row n is n + (10 + 10) + n. A counter
+ // is what makes the repetition worth something: leaving the inner binding in place gives
+ // n + 20 + 10, and losing memoization of the outer definition gives n + 20 + (n + 1).
+ assert((1 to 3).map(_ => outer.eval(InternalRow.empty)) == Seq(22, 24, 26))
+
+ // The generated path never has to answer this: one id in two nested scopes is refused while
+ // generating, so it cannot quietly disagree with the values above.
+ val literalOuter = CommonExpressionDef(Literal(1), new CommonExpressionId())
+ val literalRef = new CommonExpressionRef(literalOuter)
+ val nested = With(
+ Add(literalRef, With(Add(literalRef, literalRef),
+ Seq(CommonExpressionDef(Literal(10), literalOuter.id)))),
+ Seq(literalOuter))
+ val generated = intercept[SparkException](GenerateMutableProjection.generate(Seq(nested)))
+ assert(generated.getMessage.contains("is already being generated"))
+ }
+
+ test("SPARK-58902: a definition that references its own id fails without recursing") {
+ // Only a caller building the case class directly can produce this, and it is caught rather than
+ // left to recurse: generating the definition re-enters the same slot, and evaluating it would
Review Comment:
**Nit (P3):** The generated half does re-enter the same slot, but the interpreted half of this first construction uses two distinct reference objects. Without the new validation, `refsToBind` would bind only the child reference and the reference inside `selfDef` would fail as unbound; only the shared-reference variant below re-enters one cell. Could we split those two failure mechanisms in this comment?
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5602038473
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/53b19220d3a2c2f2ff0ac54c017b234a948326a7
- merged into branch-4.x https://github.com/apache/spark/commit/93e2850ba31f500f174368cd190839d9e747c2bc
*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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5420184617
@cloud-fan done in [9ad8a39](https://github.com/apache/spark/pull/58045/commits/9ad8a390420cb9cbf7b16ce2eefe76d9dc4a9364), Could you take another look when you have time? Thanks.
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5364502418
We should be careful when we change the expression evaluation order or invocation times. Do the existing expression properties sufficient to make safe judgement here? e.g. stateful, throwable, deterministic, etc.
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3821187764
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -183,6 +215,15 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
// Use transformUp to handle nested With.
newExpr.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w @ With(_, defs) if defs.exists(shouldPreEvaluateInBranch(_, commonExprIdSet)) =>
Review Comment:
Thanks for the pointer on lazy memoization. I went and tried the closest thing reachable in the plan, and it does not hold up. Reporting the negative result here since it rules out a direction rather than just failing to reach one.
**The guard experiment.** The idea was to keep the eager pre-evaluation but wrap the column in `If(branch is reached, definition, default)`, so a stateful generator is only advanced on the rows that reach the branch. Once-when-taken, zero-times-otherwise, expressed in the plan.
It breaks on subexpression elimination. A guard in a plain `Project` is outside conditional evaluation, so to CSE it is an `alwaysEvaluatedInputs` of the `If`: a part repeated across guards gets hoisted to the top of the projection and evaluated up front (`spark.sql.subexpressionElimination.skipForShortcutExpr` is false by default). A condition that can raise then raises on rows the branch never reached:
```sql
SELECT CASE WHEN a = 0 THEN false
WHEN 6 / a > 2 THEN rand(1) BETWEEN 0 AND 1
WHEN 6 / a < -2 THEN rand(2) BETWEEN 0 AND 1
ELSE false END
FROM t -- t contains a = 0
```
gives `DIVIDE_BY_ZERO`, frame in `project_subExpr_0$`. Limiting how many guards a projection may hold does not help, since a single guard repeating `6 / a` internally does the same.
The only fix I found is to build guards out of nodes that cannot raise, and that leaves very little:
| condition | guard |
|---|---|
| `a > 0`, `a > b`, `a IS NULL`, `a IN (1,2,3)`, `s = 'x'` | kept |
| `a + b > 0`, `a * 2 > b`, `abs(a) > 0`, `upper(s) = 'X'`, `cast(a as long) > 0` | dropped |
Any arithmetic, function call or cast on a column loses it, and `WHEN a + b > 0` is an ordinary thing to write. It is also dropped when the guard reads the other side of a join, when the condition carries a subquery, when a preceding condition is nondeterministic, and for conditional expressions outside `If`/`CaseWhen`/`Coalesce`/`NaNvl`. All silently. A mechanism that only applies to plain column comparisons and fails invisibly elsewhere does not deliver the property you asked for, so I left it out of the PR rather than shipping a version of it.
**On excluding stateful generators.** I don't think that can be satisfied together with the fix. The generators whose extra evaluations are observable are exactly the ones that produce the wrong result: `rand`, `randn`, `uuid`, `monotonically_increasing_id`. What remains after excluding them — `spark_partition_id`, `input_file_name`, the input-file-block pair — is constant within a task, so evaluating it twice was never wrong to begin with. Excluding them removes the fix rather than narrowing it, and the end-to-end test in `ColumnExpressionSuite` would have to go with it.
For what the eager path does cost: Spark makes no commitment about which row draws which value from a seeded generator. The seed is `seed + partitionIndex`, so the sequence already moves with partitioning, splits and AQE, and `spark.sql.alwaysInlineCommonExpr` already makes it configuration-dependent. There is precedent for shipping that kind of shift as a migration note — since 4.3 nondeterministic filters are no longer pushed to DSv2, for the same class of reason. Happy to add a migration-guide entry here if you want one.
**Lazy memoization itself.** Filed as SPARK-58902 with the shapes it would fix. It is worth more than the guard was: the allowlist in this PR turns expressions down by shape rather than behavior, so `randstr(3, 0) BETWEEN 'a' AND 'b'` inside a branch still gets two independent draws today even though only `randstr(-1, 0)` raises. Memoization would fix that class in one go, and it would also close the pre-existing join TODO in this rule. It needs an evaluable expression with per-row caching plus codegen support, which `With` / `CommonExpressionDef` / `CommonExpressionRef` cannot be today, so it is a separate change rather than something to fold in here.
Your call on the scope of this PR.
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3797825399
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -183,6 +215,15 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
// Use transformUp to handle nested With.
newExpr.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w @ With(_, defs) if defs.exists(shouldPreEvaluateInBranch(_, commonExprIdSet)) =>
Review Comment:
**Non-blocking:**
Could we preserve branch-local lazy evaluation instead of moving the definition into an eager `Project`? Hoisting `rand(42)` consumes draws on untaken rows, so later taken rows receive different reproducible values; the `defs.exists` guard can also hoist an unsafe sibling such as `RandStr(-1, 0)`. Please investigate lazy per-row memoization inside the conditional branch—once when taken and zero times otherwise—and cover RNG advancement plus mixed definitions in tests.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -250,6 +251,79 @@ class RewriteWithExpressionSuite extends PlanTest {
)
}
+ test("SPARK-58818: nondeterministic common expression in a conditional branch") {
+ val a = testRelation.output.head
+ // The shape `input BETWEEN lower AND upper` builds: the reference appears twice.
Review Comment:
**Nit:**
```suggestion
// The shape built for `input BETWEEN lower AND upper` references the input twice.
```
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5365308253
> We should be careful when we change the expression evaluation order or invocation times. Do the existing expression properties sufficient to make safe judgement here? e.g. stateful, throwable, deterministic, etc.
Let me investigate this further.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3854922288
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala:
##########
@@ -25,13 +25,31 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block._
*/
trait CodegenFallback extends Expression {
- protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ CodegenFallback.generate(this, ctx, ev)
+}
+
+object CodegenFallback {
+
+ /**
+ * Generates code that evaluates `e` by calling its `eval`, rather than by generating code for it.
+ * This is what [[CodegenFallback]] gives the expressions that mix it in, and is also used by an
+ * expression that can generate code in general but has to fall back for a particular shape of
+ * its own tree -- see `With.doGenCode`.
+ *
+ * The two places that reason about which subtrees run interpretively both dispatch on the trait:
Review Comment:
**Nit:**
Please add `EquivalentExpressions.childrenToRecurse` to this maintenance list. It also dispatches on `CodegenFallback` to keep interpretive subtrees out of CSE, so the current `two places` claim is incomplete.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -226,15 +228,85 @@ class RewriteWithExpressionSuite extends PlanTest {
)
}
+ test("SPARK-58902: a With left in a conditional branch of an aggregate still converges") {
+ val Seq(a, b) = testRelation.output
+ // Not cheap and referenced twice, so it stays a memoizing `With` rather than being inlined.
+ val inBranch = With(a + b) { case Seq(ref) => ref * ref }
+ val plan = testRelation.groupBy(a)(max(Coalesce(Seq(a, inBranch))).as("col"))
+ // The `PhysicalAggregation` arm restructures the aggregate into a `Project` above it, and its
+ // guard is "the expressions contain a `With`", which a surviving one keeps true on every
+ // iteration of this fixed-point batch. Without the eq check in the rule this raises
+ // `Max iterations (5) reached for batch Rewrite With expression`, one `Project` per iteration.
+ val rewritten = Optimizer.execute(plan)
+ // Idempotent: running the batch again changes nothing.
+ comparePlans(Optimizer.execute(rewritten), rewritten)
+ assert(rewritten.collect { case p: Project => p }.size <= 1,
+ s"the rule stacked a Project per iteration:\n$rewritten")
+ }
+
+ test("SPARK-58902: a cheap or single-reference definition in a branch is still inlined") {
+ val Seq(a, b) = testRelation.output
+ // A bare attribute is cheap, so inlining it costs nothing and keeps the branch foldable.
Review Comment:
**Nit:**
`Attribute.foldable` is false, so inlining `a` does not keep this branch foldable.
```suggestion
// A bare attribute is cheap, so inlining it costs nothing.
```
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WithExpressionEvalSuite.scala:
##########
@@ -0,0 +1,194 @@
+/*
+ * 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.catalyst.expressions
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodegenFallback, ExprCode, GenerateMutableProjection}
+import org.apache.spark.sql.types.{DataType, IntegerType}
+
+/**
+ * Evaluation of [[With]] and the memoization it gives a [[CommonExpressionRef]]. The rewrite that
+ * decides which `With`s reach evaluation at all is covered by `RewriteWithExpressionSuite`.
+ */
+class WithExpressionEvalSuite extends SparkFunSuite {
+
+ /**
+ * A stand-in for a stateful generator: every evaluation returns the next integer, so a second
+ * evaluation within one row is directly observable. Only used on the interpreted path.
+ */
+ private case class Counter() extends LeafExpression with Nondeterministic {
+ @transient private var n = 0
+ override def stateful: Boolean = true
+ override def dataType: DataType = IntegerType
+ override def nullable: Boolean = false
+ override protected def initializeInternal(partitionIndex: Int): Unit = {}
+ override protected def evalInternal(input: InternalRow): Any = { n += 1; n }
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ throw new UnsupportedOperationException
+ }
+
+ private def counter(): Counter = {
+ val c = Counter()
+ c.initialize(0)
+ c
+ }
+
+ /**
+ * A node that has to be evaluated interpretively even when its parent is generated, so that a
+ * reference below it takes the interpreted path out of generated code.
+ */
+ private case class Fallback(child: Expression) extends UnaryExpression with CodegenFallback {
+ override def dataType: DataType = child.dataType
+ override def eval(input: InternalRow): Any = child.eval(input)
+ override protected def withNewChildInternal(newChild: Expression): Fallback =
+ copy(child = newChild)
+ }
+
+ test("a definition is evaluated once per row however many references read it") {
+ // `ref + ref` is the shape `BETWEEN` produces. Memoized, both references read one value, so the
+ // sum is 2n on the nth row; inlined it would be n + (n + 1).
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ assert((1 to 4).map(_ => w.eval(InternalRow.empty)) == Seq(2, 4, 6, 8))
+ }
+
+ test("a definition is not evaluated on a row that reaches no reference") {
+ val c = counter()
+ // The branch is inside the `With`, so the `With` is entered on every row and does clear its
+ // cells. What it must not do is evaluate the definition before a reference is reached; putting
+ // the branch outside would only test that `If` does not evaluate the arm it did not take.
+ val w = With(c) { case Seq(ref) => If(Literal.TrueLiteral, Literal(-1), Add(ref, ref)) }
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(-1, -1, -1))
+ // The counter is still at 0, so the first row that does reach a reference sees 1.
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty) == 2)
+ }
+
+ test("a reference behind a short-circuiting operator is not read when the left side is false") {
+ // Neither pre-evaluating the definition into a project nor guarding that column by the branch
+ // can express this: both evaluate on every row the branch is reached on, while `And` short
+ // circuits before the reference. The `And` is inside the `With` so that the `With` is entered
Review Comment:
**Nit:**
`short-circuit` is a compound verb here.
```suggestion
// can express this: both evaluate on every row the branch is reached on, while `And`
// short-circuits before the reference. The `And` stays inside the `With` so it is entered
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,121 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with the definition each
+ * names. The list is found once, since the tree does not change between evaluations.
+ *
+ * Only `child` is scanned, which relies on a reference to one of these definitions never living
+ * inside another one of them: `With.apply` cannot build that, and `RewriteWithExpression`, which
+ * does rewrite inside a `With`, only ever replaces a reference with its definition's child or
+ * with an attribute -- it never puts a reference inside a definition. Scanning `children` instead
+ * would look safer and be worse -- it would also bind a reference sitting inside the definition
+ * it names, and since `CommonExpressionCell.get` sets `computed` only after the nested evaluation
+ * returns, that turns today's loud "outside its With" error into a StackOverflowError. A nested
+ * `With` is not affected either way: `children` is `child +: defs`, so this scan already descends
+ * into an inner `With`'s own definitions.
+ */
+ @transient private lazy val refsToBind: Seq[(CommonExpressionRef, CommonExpressionDef)] = {
+ val idToDef = defs.map(d => d.id -> d).toMap
+ child.collect { case r: CommonExpressionRef if idToDef.contains(r.id) => (r, idToDef(r.id)) }
+ }
+
+ /**
+ * Binds this `With`'s references to its own cells, clears them, and evaluates the child. A
+ * reference reached by that evaluation computes its definition once and every later reference
+ * reads the value back, so a definition is evaluated where the child would have evaluated it,
+ * once, rather than once per reference. See [[CommonExpressionCell]].
+ *
+ * The binding is redone on every evaluation rather than once, because a reference can be reached
+ * from two `With`s. `withNewChildrenInternal` cannot hand the new `With` its own references: a
+ * rebuilt reference compares equal to the one it replaces, since the binding it carries is not
+ * part of its equality, so `transform` keeps the original. Binding once would then leave the
+ * `With` that bound last deciding what both of them read. Rebinding costs one pass over the
+ * references, two for a `BETWEEN`, and makes the `With` currently evaluating always the owner.
+ */
+ override def eval(input: InternalRow): Any = {
+ refsToBind.foreach { case (ref, exprDef) => ref.bindTo(exprDef) }
+ defs.foreach(_.cell.clear())
+ child.eval(input)
+ }
+
+ // The cells are cleared on entry, so this holds state for the duration of one evaluation.
+ override def stateful: Boolean = true
+
+ /**
+ * Whether one of this `With`'s references sits somewhere that will be evaluated interpretively
+ * even though this `With` is generated. Two shapes do that: a [[CodegenFallback]], which is
+ * evaluated by calling `eval` on it from the generated code, and a nested `With` that itself
+ * takes the fallback below -- `With` does not mix in `CodegenFallback`, so it has to be named
+ * here rather than matched as one. A reference reached that way needs its cell bound and
+ * cleared, which the generated code does not do: it clears the codegen flags.
+ *
+ * This is the same shape `EquivalentExpressions.childrenToRecurse` already refuses to look past,
+ * for the same reason.
+ *
+ * Each level memoizes, but `holdsMyRef` runs again at every nested `With` the scan passes, so a
+ * chain of them nested in each other's `child` costs on the order of the square of the depth.
+ * `nullif(a, nullif(b, c))` does produce such a chain -- only the memoized input becomes a
+ * definition, the rest stays in `child` -- but these chains are shallow in practice. Reading a
+ * nested `With`'s own `lazy val` from here also takes its monitor while holding this one; the
+ * edges only ever run from an ancestor to a proper descendant of an immutable tree, so the order
+ * is a strict partial one and cannot deadlock. `canonicalizationIdMap` below relies on the same.
+ */
+ @transient private lazy val refUnderCodegenFallback: Boolean = {
+ val ids = defs.map(_.id).toSet
+ def holdsMyRef(e: Expression): Boolean = e.exists {
+ case r: CommonExpressionRef => ids.contains(r.id)
+ case _ => false
+ }
+ child.exists {
+ case f: CodegenFallback => holdsMyRef(f)
+ case w: With if w.refUnderCodegenFallback => holdsMyRef(w)
+ case _ => false
+ }
+ }
+
+ /**
+ * Clears each definition's flag, then generates the child. The flags are cleared in the same
+ * block the child is generated into, so a reference cannot run against a flag left set by an
+ * earlier row: on a row that does not reach the branch holding this `With`, neither the clearing
+ * nor any reference runs.
+ *
+ * When a reference sits under a [[CodegenFallback]], or inside a nested `With` that itself falls
+ * back, the whole `With` is evaluated interpretively instead. Generating the child would leave
+ * that reference reading a cell nobody bound and nobody clears, and generating part of it is
+ * worse still: a definition reached from both sides would be computed once through the flags and
+ * once through the cell, holding two values for one row. [[eval]] binds and clears both, so
+ * handing it the whole subtree keeps one mechanism in play. `ctx.INPUT_ROW` is available on that
+ * path because `CollapseCodegenStages.supportCodegen` turns whole-stage codegen off for a plan
+ * whose expressions hold the offending `CodegenFallback` -- it is visible there, since a `With`
+ * in a conditional branch reaches execution inside `plan.expressions` like any other expression.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ if (refUnderCodegenFallback) {
+ return CodegenFallback.generate(this, ctx, ev)
+ }
+ ctx.withCommonExprs(defs) { slots =>
+ val clearFlags = slots.map(s => s"${s.computed} = false;").mkString("\n")
+ val childGen = child.genCode(ctx)
+ ev.copy(
+ code = code"""
+ |$clearFlags
+ |${childGen.code}
+ """.stripMargin,
+ isNull = childGen.isNull,
+ value = childGen.value)
+ }
+ }
+
override protected def withNewChildrenInternal(
newChildren: IndexedSeq[Expression]): Expression = {
val newDefs = newChildren.tail.map(_.asInstanceOf[CommonExpressionDef])
// If any `CommonExpressionDef` has been updated (data type or nullability), also update its
- // `CommonExpressionRef` in the `child`.
+ // `CommonExpressionRef` in the `child`. This cannot be used to hand the new `With` its own
+ // reference objects: a rebuilt reference is `==` the one it replaces, since the binding it
+ // carries is not part of its equality, so `transform` keeps the original. `eval` rebinds
+ // instead of relying on the references being unshared -- see [[refsToBind]].
Review Comment:
**Nit:**
Scaladoc links are not parsed in `//` comments, so this remains literal bracket markup.
```suggestion
// instead of relying on the references being unshared -- see `refsToBind`.
```
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3854925487
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -183,6 +215,15 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
// Use transformUp to handle nested With.
newExpr.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w @ With(_, defs) if defs.exists(shouldPreEvaluateInBranch(_, commonExprIdSet)) =>
Review Comment:
The new branch-local `With` memoization addresses this concern: definitions stay in the conditional branch and evaluate once on first reference. Thanks for reworking this.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5428585049
Thanks. The finding is real, and the comment sitting above that code said the opposite, which bothers me more than the code does.
I measured it before touching anything, and the numbers moved my view of what the fix is worth. Copies of a marked innermost body in the generated source, nested `With`s, depths 1 to 8, default `methodSplitThreshold` of 1024, leaf body 237 characters. The two rows differ in what the expression is generated against: an input row (`INPUT_ROW` set, `currentVars` null), or the child's output columns as local variables, which is what whole-stage codegen hands an expression.
```
against a row, before: 2 4 4 4 4 4 4 4 source 919 -> 7221
against a row, after: 2 2 2 2 2 2 2 2 source 919 -> 4457
against locals, both: 2 4 8 16 32 64 128 256 source 823 -> 163371
```
Generating against a row was never exponential. `Expression.genCode` runs `reduceCodeSize` on every expression, including a `With` serving as a definition, under the same `INPUT_ROW != null && currentVars == null` guard and the same threshold. Once a level's text crosses it, that level becomes a method and its parent's two references emit two calls rather than two bodies, so the growth resets. Hence the plateau at 4.
The exponential case is the one you named. Whole-stage codegen has `reduceCodeSize` disabled by the condition you quoted, and it is the path this fix cannot reach, because a method cannot see the `currentVars` locals a definition reads.
So what the fix buys is smaller than the comment claimed: the body once per scope instead of once per reference, and a bound that no longer grows with depth wherever the threshold sits. Raise `methodSplitThreshold` and the "before" plateau rises with it, while the "after" stays at 2. The body goes into `private void computeCommonExpr_N(InternalRow)` and each reference keeps a call behind its `computed` flag. The new test in `WithExpressionEvalSuite` holds the count at 2 for depths 1 to 6 with the threshold pinned, so the expectation does not ride on the default. It reads 4 at depth 2 without the change.
Covering whole-stage codegen means passing the `currentVars` locals a definition needs into the method as parameters. That is the `// TODO: support whole stage codegen` on `splitExpressionsWithCurrentInputs` (`CodeGenerator.scala:1051`), blocked by the same condition for the same reason. I would rather not fold it in here.
The whole-stage half therefore stays as it was, which I do not think is a regression. On master a nested `nullif` in a conditional branch gets inlined by `RewriteWithExpression`, so the plan is already 2^depth nodes before codegen sees it. This PR makes the plan linear, and the generated code has a size-based net already: `spark.sql.codegen.hugeMethodLimit` and the 64KB method limit demote oversized whole-stage output, and the projection it falls back to is where this fix applies.
If you would rather close that half here too, the cheap lever is `CollapseCodegenStages.supportCodegen` returning false for a `With` whose definition holds another one, which keeps it out of the fused stage. I left it out because whole-stage source at depth 2 is only 2091 characters, so the trade is giving up whole-stage codegen for an ordinary `nullif(nullif(a, b), c)`, and bounding by depth instead means picking a magic number. Happy to add it if you prefer that trade.
The three nits were all right and are applied. The `CodegenFallback` note now says `CollapseCodegenStages` finds a non-leaf descendant fallback structurally, so a non-trait caller only has to be named at the two sites that dispatch on the node in hand. `refsToBind` names the curried `With(commonExprs: _*)(replaced)` helper instead of claiming the case class constructor cannot build a definition that references one of these ids, and says what happens when a caller does: the reference stays unbound and evaluating it raises. And the `randstr` comment no longer explains a raise that a foldable length of `3` cannot produce.
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5324037877
Thanks for the review. I fixed the `defs.exists` problem you flagged, and I looked into the lazy-memoization suggestion. Reporting what I found on the second one, because I do not think I should decide its scope alone.
### The `defs.exists` hoisting is real, fixed in 01f629aa0d3
You were right, and my first attempt to reproduce it was wrong: putting the `With` *outside* the `CaseWhen` goes down the main rewrite, which hoists everything anyway. With the `With` inside the ELSE branch (the shape `Between` actually produces) and two definitions, `rand(1)` plus `randstr(-1, 0)`, the unsafe sibling did get pre-evaluated. On bb3296963f2:
```
Project [CASE WHEN (a#0 > 0) THEN true ELSE (((_common_expr_0#3 >= 0.4) AND (_common_expr_0#3 <= 0.6)) AND ((_common_expr_1#4 >= a) AND (_common_expr_1#4 <= z))) END AS col#2]
+- Project [a#0, b#1, rand(1) AS _common_expr_0#3, randstr(-1, 0, false) AS _common_expr_1#4]
+- LocalRelation <empty>, [a#0, b#1]
```
`randstr(-1, 0)` then raises `INVALID_PARAMETER_VALUE.LENGTH` on rows whose branch is never taken. The follow-up moves the decision to each definition, so `rand()` is pre-evaluated while the `randstr` sibling keeps its inlining:
```
Project [CASE WHEN (a#0 > 0) THEN true ELSE (((_common_expr_0#3 >= 0.4) AND (_common_expr_0#3 <= 0.6)) AND ((randstr(-1, 0, false) >= a) AND (randstr(-1, 0, false) <= z))) END AS col#2]
+- Project [a#0, b#1, rand(1) AS _common_expr_0#3]
+- LocalRelation <empty>, [a#0, b#1]
```
A test covers this shape. It also let me delete the recursion back into the main branch, so the two paths now read more alike. (The alias names above are index-based; with `spark.sql.useCommonExprIdForAlias` at its default of true you would see the common-expression id instead.)
### On branch-local lazy memoization: no mechanism exists today, and I would rather you decide the scope
I checked whether this could be built on the existing subexpression-elimination machinery. It cannot, in either execution mode, and the blocker is not that CSE avoids branches — it does look inside them — but that it never admits a nondeterministic expression and never evaluates lazily.
- **CSE does reach into conditional branches.** `childrenToRecurse` takes only `alwaysEvaluatedInputs` for a `ConditionalExpression` (EquivalentExpressions.scala:160), but `commonChildrenToRecurse` additionally recurses into `branchGroups` (line 169), and `updateCommonExprs` (lines 103-130) hoists whatever is common to *every* branch of a group — `If.branchGroups` is `Seq(Seq(trueValue, falseValue))` (conditionalExpressions.scala:72). So the machinery already reasons about branches; it admits expressions that are guaranteed to be evaluated whichever branch runs.
- **But it excludes nondeterministic expressions entirely.** Every insertion goes through `updateExprInMap`, which is gated on `if (expr.deterministic)` (EquivalentExpressions.scala:65). `rand`/`randn` never become common subexpressions. The nondeterministic leaves — `uuid`, `monotonically_increasing_id`, `spark_partition_id`, `input_file_name` — are excluded twice over, because `updateExprTree` also skips every `LeafExpression` regardless of determinism (line 198).
- **And where it does apply, it is eager, not lazy.** Interpreted evaluation has genuine per-row memoization: `SubExprEvaluationRuntime` invalidates its cache in `setInput` (SubExprEvaluationRuntime.scala:67-70), which every caller invokes once per row (`InterpretedPredicate.eval`, `InterpretedUnsafeProjection.apply`, `InterpretedMutableProjection.apply`, `InterpretedSafeProjection.apply`), and `ExpressionProxy.eval` loads through it (line 135). But proxies come only from that same gated map (`proxyExpressions`, line 88). Codegen emits one `subExpr` function per common expression and calls them all up front (CodeGenerator.scala:1322-1367, `subexprFunctionsCode` at GeneratePredicate.scala:42-63). `SubExprEliminationState.children` (line 81) orders dependencies, it does not defer them. Whole-stage codegen has its own path, `subexpressionEliminationForWholeStageCodegen` (line 1193), which lets the operator choose where to place the evaluation instead of emitting it all at once — bu
t it builds the same `EquivalentExpressions` (lines 1195-1196), so the same gate applies, and the placement is still a fixed point rather than per-reference memoization.
A runtime check, hand-duplicating the expression inside the branch so no `With` is involved at all:
```sql
SELECT CASE WHEN id < 0 THEN false
ELSE monotonically_increasing_id() >= 3
AND monotonically_increasing_id() <= 5 END
FROM range(0, 10, 1, 1)
```
returns 6 true rows instead of 3, with `spark.sql.subexpressionElimination.enabled` both true and false. To be precise about what this does and does not show: it demonstrates that no existing machinery memoizes a repeated nondeterministic expression in a branch, but because `monotonically_increasing_id` is a leaf it cannot isolate the determinism gate from the leaf skip.
Repeated evaluation is of course not a bug when the user writes the duplicates: `SELECT rand(), rand()` is meant to give two draws. The bug is specific to `With`, which promises the opposite — "The common expressions are guaranteed to be evaluated only once even if it's referenced more than once" (With.scala:26-27) — and inlining silently downgraded it to the repeated-call behavior in a place the user never wrote a duplicate, because `Between` put it there.
So lazy memoization needs a new evaluable expression with per-row caching plus codegen support, and `With`/`CommonExpressionDef`/`CommonExpressionRef` are all `Unevaluable` today (With.scala:30,154,165), which is why the rewrite has to eliminate them before execution.
That is a bigger change than this bug fix, and it would subsume the eager projection here rather than sit beside it. But I do not want to declare it out of scope unilaterally: if you would rather this PR wait for the real mechanism than ship a partial fix that trades RNG advancement on untaken rows, I am happy to do it that way. Which would you prefer?
For the record, the RNG advancement you named is real and the PR description calls it out: the pre-evaluated generator now advances on untaken rows, so `monotonically_increasing_id()` values and `rand()` draws shift relative to the current behavior. My argument for accepting it is that the current behavior is already wrong, and that outside a conditional branch the main rewrite already hoists a multi-referenced, non-cheap definition into a project without any determinism check. Where it cannot — e.g. the `childPlanIndex == -1` force-inline — the file already carries a TODO saying exactly this: "things can go wrong if the common expression is nondeterministic. We don't fix it for now to match the old buggy behavior when certain `RuntimeReplaceable` did not use the `With` expression" (RewriteWithExpression.scala:198-200). So inline-plus-nondeterministic is already documented upstream as known-buggy legacy behavior; this PR removes one class of it rather than introducing a new po
licy.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3910503167
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -98,6 +104,55 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether substituting this definition into its references is as good as evaluating it once.
+ * True when it is referenced once, since substituting then evaluates it once either way, or when
+ * it is cheap enough to evaluate repeatedly and deterministic, so the repeated evaluations agree.
+ *
+ * `CollapseProject.isCheap` answers what one evaluation costs, not whether a second one is
+ * allowed: it admits a `PythonUDF`, which may be nondeterministic. Determinism is checked here
+ * rather than read into the cost test, so that a reader of either call site can see which of the
+ * two questions is being asked.
+ */
+ private def canSubstitute(
+ child: Expression,
+ id: CommonExpressionId,
+ commonExprIdSet: Set[CommonExpressionId]): Boolean = {
+ !commonExprIdSet.contains(id) || (CollapseProject.isCheap(child) && child.deterministic)
+ }
+
+ /**
+ * `w` with every definition that gains nothing from being memoized inlined into its references:
+ * one cheap enough to evaluate twice, and one that is referenced once anyway. This is the test
+ * the main rewrite already applies before it hoists a definition into a project.
+ *
+ * Inlining matters beyond the per-row bookkeeping it saves. A `With` is not foldable, so it hides
+ * whatever it wraps from `ConstantFolding`, `PushFoldableIntoBranches`, `SimplifyConditionals`
+ * and `ReplaceNullWithFalseInPredicate`, all of which run in later batches. Dropping the `With`
+ * once nothing is left to memoize keeps `CASE WHEN c THEN nullif(1, 1) END` folding as it did
+ * before this rule learned to leave one behind.
+ */
+ private def inlineDefsThatGainNothing(
+ w: With,
+ commonExprIdSet: Set[CommonExpressionId]): Expression = {
+ val (toInline, toKeep) = w.defs.partition { d =>
+ canSubstitute(d.child, d.id, commonExprIdSet)
Review Comment:
**Blocking (P1):** `commonExprIdSet` describes the expression before this bottom-up rewrite, but rewriting an inner `With` can duplicate a reference owned by this outer `With`. With `spark.sql.optimizer.avoidCollapseUDFWithExpensiveExpr=false`, an outer nondeterministic Python UDF used once inside a cheap deterministic inner Python UDF becomes two outer references when that inner definition is substituted at two sites; this stale single-use decision then substitutes the nondeterministic definition twice, so the consumers can observe different values. Please base this classification on reference multiplicity in the rewritten `With` and add the nested regression.
**Recommended change:** Recompute reference multiplicity for each rewritten With immediately before classifying its definitions, and add a focused nested-UDF regression.
**Why this works:** Count CommonExpressionRef occurrences from the post-transform child and definitions so references introduced by inner substitution prevent an outer nondeterministic definition from being inlined.
**Scope:** RewriteWithExpression and its focused Catalyst optimizer suite.
**Compatibility:** No public API or data-format change; this only corrects optimizer behavior for nested With expressions.
**Risks:** The recount must respect With binding boundaries while including references introduced into the rewritten node. More definitions may remain memoized in nested shapes where the original count was stale.
**Constraints:** Preserve the existing cheap deterministic and true single-use substitutions when their post-rewrite counts permit them. Preserve branch-local laziness and optimizer idempotence.
**Success:** The nested configuration invokes the outer nondeterministic definition once per reached With entry, and the optimizer's existing substitution and convergence tests continue to pass.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,150 @@ class CodegenContext extends Logging {
throw QueryExecutionErrors.lambdaVariableNotDefinedError(id))
}
+ /**
+ * The slots a `CommonExpressionRef` reads: the value and its nullness, plus the flag saying
+ * whether this row has computed them yet, and the definition to compute them from.
+ */
+ case class CommonExprSlots(
+ value: ExprCode,
+ computed: String,
+ definition: Expression) {
+
+ // Not constructor parameters: every one of those takes part in a case class's `equals`,
+ // `hashCode`, `copy` and `toString`, so a mutable one would make a slot's hash change as it is
+ // filled, and would hand a `copy` the code generated for another scope.
+ private var fillCode: Option[Block] = None
+ private var filling: Boolean = false
+
+ /**
+ * The code that computes the definition into the slots and sets `computed`, which a reference
+ * emits behind that flag. Generated once and cached, so that every reference shares whatever
+ * mutable state the definition allocated, such as an RNG, rather than getting its own.
+ *
+ * A definition that references its own id would re-enter this while the first call is still
+ * generating it, since `fillCode` is set only after `definition.genCode` returns. `filling`
+ * turns that into the error `With.refsToBind` documents for the interpreted path, rather than a
+ * StackOverflowError from generating the same definition inside itself forever.
+ *
+ * The body goes into a method where it can and is worth it -- a definition that is or holds
+ * another `With`, or a body past the split threshold -- leaving a reference's own code a single
+ * call: such a definition would otherwise have its body pasted once per reference at every
+ * level. `Expression.reduceCodeSize` already keeps that from running away, by hoisting
+ * whichever node's code first passes the same threshold, so what a method here buys is the body
+ * once per scope rather than once per reference up to that threshold, and a bound that no
+ * longer grows with depth wherever the threshold sits. It is only possible where the definition
+ * reads the input row rather than local variables -- the condition `reduceCodeSize` splits
+ * under, and for the same reason. Whole-stage codegen generates the operators a `With` reaches
+ * with `currentVars` set, so neither applies there and the body is pasted per reference.
+ */
+ def fill: Block = {
+ if (fillCode.isEmpty) {
+ if (filling) {
+ throw SparkException.internalError(
+ "Cannot generate a common expression whose definition references it: " +
+ definition.toString)
+ }
+ filling = true
+ try {
+ fillCode = Some(build)
+ } finally {
+ filling = false
+ }
+ }
+ fillCode.get
+ }
+
+ private def build: Block = {
+ val defGen = definition.genCode(CodegenContext.this)
+ // Whether the isNull slot exists is decided by the definition, so it is read off the slot
+ // rather than off a reference's own `nullable`: taking it from both would let the two
+ // disagree, and either emit `false = <isNull>;`, which does not compile, or leave the slot
+ // holding the previous row's nullness.
+ val assignIsNull = if (value.isNull == FalseLiteral) {
+ ""
+ } else {
+ s"${value.isNull} = ${defGen.isNull};"
+ }
+ val body = code"""
+ |${defGen.code}
+ |$assignIsNull
+ |${value.value} = ${defGen.value};
+ |$computed = true;
+ """.stripMargin
+ val canPutInMethod = INPUT_ROW != null && currentVars == null
+ // A definition that is or holds another `With` is the shape whose code doubles per level,
+ // and what this is aimed at. It is not the only one -- a definition referencing a sibling
+ // definition of the same `With` doubles the same way, and codegen accepts that, since the
+ // sibling's slots are in scope while this definition is generated (`With.refsToBind` says
+ // why nothing builds that tree, and that evaluating one raises). What bounds those is not
+ // the length arm below: `body` is assembled after `definition.genCode` already ran
+ // `reduceCodeSize`, so the arm fires only in the band just under the threshold. It is
+ // `reduceCodeSize` itself, which hoists whichever node's code first passes the threshold as
+ // generation walks up, capping what one level contributes, so the code stays linear in the
+ // depth either way. The length arm just keeps the same body from being split once per
+ // reference, which leaves the methods small and the code as large.
+ val worthAMethod = definition.exists(_.isInstanceOf[With]) ||
+ body.length > SQLConf.get.methodSplitThreshold
+ if (canPutInMethod && worthAMethod) {
+ val funcName = freshName("computeCommonExpr")
+ val funcFullName = addNewFunction(funcName,
+ s"""
+ |private void $funcName(InternalRow $INPUT_ROW) {
+ | $body
+ |}
+ """.stripMargin)
+ code"$funcFullName($INPUT_ROW);"
+ } else {
+ body
+ }
+ }
+ }
+
+ /**
+ * Holding a map of the common expressions of the `With` expressions currently being generated,
+ * the same way [[currentLambdaVars]] holds the variables of the enclosing lambdas.
+ */
+ var currentCommonExprs: mutable.Map[Long, CommonExprSlots] = mutable.HashMap.empty
+
+ /**
+ * Allocates a value slot and a `computed` flag per definition, generates `f` with them in scope,
+ * then takes them out of scope again. A reference generated inside `f` reads the slots back by
+ * id and fills them the first time it is reached on a row -- the enclosing `With` only clears the
+ * flags.
+ */
+ def withCommonExprs(defs: Seq[CommonExpressionDef])(f: Seq[CommonExprSlots] => ExprCode)
+ : ExprCode = {
+ val slots = defs.map { d =>
Review Comment:
**Non-blocking (P2):** Please add interpreted and generated execution coverage for one `With` containing two surviving, distinguishable definitions, each referenced twice. Every current runtime case has one definition, and the optimizer's two-definition case inlines one before evaluation, so a slot-index or cleanup bug that aliases these per-definition states would not be observed. Exact values and once-per-definition advancement across rows would cover both allocation and clearing.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -162,7 +373,79 @@ case class CommonExpressionDef(child: Expression, id: CommonExpressionId = new C
* referenced, so that we can determine the data type and nullable of the reference node.
*/
case class CommonExpressionRef(id: CommonExpressionId, dataType: DataType, nullable: Boolean)
- extends LeafExpression with Unevaluable {
+ extends LeafExpression {
def this(exprDef: CommonExpressionDef) = this(exprDef.id, exprDef.dataType, exprDef.nullable)
+
+ /**
+ * The definition this reference names, and the cell holding its value for the current row. Both
+ * are wired by the enclosing [[With]] before it evaluates its child, and are left out of the case
+ * class parameters so that equality and canonicalization are unchanged -- and so that a rule
+ * comparing two references does not compare their cells.
+ */
+ private var definition: Expression = _
+ private var cell: CommonExpressionCell = _
+
+ private[expressions] def bindTo(exprDef: CommonExpressionDef): Unit = {
+ definition = exprDef.child
+ cell = exprDef.cell
+ }
+
+ private[expressions] def boundDefinition: Expression = definition
+ private[expressions] def boundCell: CommonExpressionCell = cell
+
+ private[expressions] def bindTo(
+ newDefinition: Expression,
+ newCell: CommonExpressionCell): Unit = {
+ definition = newDefinition
+ cell = newCell
+ }
+
override val nodePatterns: Seq[TreePattern] = Seq(COMMON_EXPR_REF)
+
+ // The cell is cleared by the enclosing `With` on every entry, so this reads mutable state.
+ override def stateful: Boolean = true
+
+ /**
+ * A copy must not carry this reference's binding: the copy belongs to a different `With`, which
+ * wires it to its own cell. `LeafLike` returns `this` here, which would hand two `With`s one
+ * reference object and let whichever wires last decide what both of them read --
+ * `NamedLambdaVariable` overrides this for the same reason.
+ */
+ override def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): CommonExpressionRef = copy()
+
+ override def eval(input: InternalRow): Any = {
+ if (cell == null) {
+ throw SparkException.internalError(
+ s"Cannot evaluate a common expression reference outside its With: $this")
+ }
+ cell.get(definition, input)
+ }
+
+ /**
+ * Computes the definition into the shared slots if this row has not done so yet, then reads them.
+ * The code that computes it is emitted here rather than by the enclosing `With`, so it runs where
+ * the first reference is reached -- behind a short-circuiting operator or a nested conditional,
+ * if that is where the reference sits.
+ *
+ * A second reference emits the same code again, which never runs because the flag is set. What
+ * that code is depends on the definition: a call, where the definition can be put in a method, so
+ * that a definition that is or holds another `With` is not pasted once per reference at every
+ * level; the body itself otherwise, whose locals are declared inside each guard. No copy of one
+ * body encloses another -- for that, a definition would have to reference its own id, directly or
+ * through a sibling, which recurses in `fill` before any Java exists -- so repeating it declares
+ * nothing twice in one scope. See `CommonExprSlots.fill`, which also says what bounds the code
+ * when a method is not possible.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val slots = ctx.getCommonExpr(id.id)
+ ev.copy(
+ code = code"""
+ |if (!${slots.computed}) {
Review Comment:
**Non-blocking (P2):** Could the generated path get a case that enters this `With` but takes a child branch containing no reference? The current no-reference checks use interpreted `Counter`, while the SQL skip case avoids entering the `With` altogether. A `GenerateMutableProjection` over a stateful definition with inputs `false, true, false, true` should return `-1, 0, -1, 1`; that would catch an eager fill here that consumes values on the false rows.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -181,15 +236,12 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
rewriteWithExprAndInputPlans(
_, inputPlans, commonExprsPerChild, commonExprIdSet, isNestedWith))
val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
- // Use transformUp to handle nested With.
+ // A `With` in a conditional branch cannot go into a project, which is always evaluated
+ // while the branch may not be. It stays where it is and memoizes its definition per row
Review Comment:
**Nit (P3):** This memoization is per `With` entry/evaluation, not per row: `With` clears the cell or generated flag whenever it is entered, and the same expression object can occur at two positions for one row. Please use `per-entry` here (and in the corresponding test/PR description wording) so the documented lifetime matches `CommonExpressionCell` and `With.eval`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -98,6 +104,55 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether substituting this definition into its references is as good as evaluating it once.
+ * True when it is referenced once, since substituting then evaluates it once either way, or when
+ * it is cheap enough to evaluate repeatedly and deterministic, so the repeated evaluations agree.
+ *
+ * `CollapseProject.isCheap` answers what one evaluation costs, not whether a second one is
+ * allowed: it admits a `PythonUDF`, which may be nondeterministic. Determinism is checked here
+ * rather than read into the cost test, so that a reader of either call site can see which of the
+ * two questions is being asked.
+ */
+ private def canSubstitute(
+ child: Expression,
+ id: CommonExpressionId,
+ commonExprIdSet: Set[CommonExpressionId]): Boolean = {
+ !commonExprIdSet.contains(id) || (CollapseProject.isCheap(child) && child.deterministic)
+ }
+
+ /**
+ * `w` with every definition that gains nothing from being memoized inlined into its references:
+ * one cheap enough to evaluate twice, and one that is referenced once anyway. This is the test
+ * the main rewrite already applies before it hoists a definition into a project.
+ *
+ * Inlining matters beyond the per-row bookkeeping it saves. A `With` is not foldable, so it hides
+ * whatever it wraps from `ConstantFolding`, `PushFoldableIntoBranches`, `SimplifyConditionals`
+ * and `ReplaceNullWithFalseInPredicate`, all of which run in later batches. Dropping the `With`
+ * once nothing is left to memoize keeps `CASE WHEN c THEN nullif(1, 1) END` folding as it did
+ * before this rule learned to leave one behind.
+ */
+ private def inlineDefsThatGainNothing(
+ w: With,
+ commonExprIdSet: Set[CommonExpressionId]): Expression = {
+ val (toInline, toKeep) = w.defs.partition { d =>
+ canSubstitute(d.child, d.id, commonExprIdSet)
+ }
+ if (toInline.isEmpty) {
+ w
+ } else {
+ val refToExpr = toInline.map(d => d.id -> d.child).toMap
+ val newChild = w.child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
+ // A ref of a definition kept here, or of an enclosing `With`, is left for its owner.
+ case ref: CommonExpressionRef if refToExpr.contains(ref.id) => refToExpr(ref.id)
+ }
+ // `copy` rather than `withNewChildren`, which requires the child count to be unchanged. The
+ // references of the kept definitions are carried over unbound; the new `With` binds them, and
+ // `w` is discarded here, so no two `With`s hold the same reference.
Review Comment:
**Nit (P3):** Discarding this parent does not guarantee unique reference objects. If the same `With` instance is reused at two parent positions, each transform can build a distinct copy while retaining the same unchanged references for `toKeep`. The new bind/restore logic makes that DAG-shaped reuse safe, so could this comment describe that mechanism instead of claiming two `With` expressions cannot hold the same reference?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,150 @@ class CodegenContext extends Logging {
throw QueryExecutionErrors.lambdaVariableNotDefinedError(id))
}
+ /**
+ * The slots a `CommonExpressionRef` reads: the value and its nullness, plus the flag saying
+ * whether this row has computed them yet, and the definition to compute them from.
+ */
+ case class CommonExprSlots(
+ value: ExprCode,
+ computed: String,
+ definition: Expression) {
+
+ // Not constructor parameters: every one of those takes part in a case class's `equals`,
+ // `hashCode`, `copy` and `toString`, so a mutable one would make a slot's hash change as it is
+ // filled, and would hand a `copy` the code generated for another scope.
+ private var fillCode: Option[Block] = None
+ private var filling: Boolean = false
+
+ /**
+ * The code that computes the definition into the slots and sets `computed`, which a reference
+ * emits behind that flag. Generated once and cached, so that every reference shares whatever
+ * mutable state the definition allocated, such as an RNG, rather than getting its own.
+ *
+ * A definition that references its own id would re-enter this while the first call is still
+ * generating it, since `fillCode` is set only after `definition.genCode` returns. `filling`
+ * turns that into the error `With.refsToBind` documents for the interpreted path, rather than a
+ * StackOverflowError from generating the same definition inside itself forever.
+ *
+ * The body goes into a method where it can and is worth it -- a definition that is or holds
+ * another `With`, or a body past the split threshold -- leaving a reference's own code a single
+ * call: such a definition would otherwise have its body pasted once per reference at every
+ * level. `Expression.reduceCodeSize` already keeps that from running away, by hoisting
+ * whichever node's code first passes the same threshold, so what a method here buys is the body
+ * once per scope rather than once per reference up to that threshold, and a bound that no
+ * longer grows with depth wherever the threshold sits. It is only possible where the definition
+ * reads the input row rather than local variables -- the condition `reduceCodeSize` splits
+ * under, and for the same reason. Whole-stage codegen generates the operators a `With` reaches
Review Comment:
**Nit (P3):** Whole-stage mode does not always imply `currentVars` is populated. For example, `SortMergeJoinExec.createJoinKey` generates keys with `INPUT_ROW` set and `currentVars` cleared, which satisfies the helper-method predicate below. Please describe the exact `INPUT_ROW != null && currentVars == null` boundary and identify local-variable whole-stage `Project` paths as the case where extraction is unavailable.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -430,6 +430,136 @@ class ColumnExpressionSuite extends SharedSparkSession {
checkAnswer(testData.filter($"a".between($"b", $"c")), expectAnswer)
}
+ // Runs `f` on each of the three evaluation paths, since a `With` left in a conditional branch is
+ // evaluated by all three and the memoization is implemented separately for interpretation and for
+ // codegen.
+ private def onEachEvalPath(f: => Unit): Unit = {
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN")(f)
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
+ }
+
+ test("SPARK-58902: BETWEEN on a nondeterministic input inside a conditional branch") {
+ onEachEvalPath {
+ // `BETWEEN` reads its input twice. Inlining the common expression into a branch gave each
+ // read its own value, so the two comparisons saw two different ids and 6 of the 10 rows came
+ // back true where 3 is correct. A single partition makes the id sequence 0, 1, 2, ...
+ val df = spark.range(0, 10, 1, 1)
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN id < 0 THEN false ELSE monotonically_increasing_id() BETWEEN 3 AND 5 END"),
+ (0 until 10).map(i => Row(i >= 3 && i <= 5)))
+ }
+ }
+
+ test("SPARK-58902: a definition in a branch is evaluated only on the rows that reach it") {
+ onEachEvalPath {
+ // Rows 0 to 4 take the first branch, so the five rows that reach the ELSE see ids 0 to 4 --
+ // the same values they would see if the branch were the whole expression.
+ val df = spark.range(0, 10, 1, 1)
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN id < 5 THEN NULL ELSE monotonically_increasing_id() BETWEEN 1 AND 2 END"),
+ (0 until 10).map { i =>
+ if (i < 5) Row(null) else Row(i - 5 >= 1 && i - 5 <= 2)
+ })
+ }
+ }
+
+ test("SPARK-58902: a branch condition that can raise is not evaluated on other rows") {
+ onEachEvalPath {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ // Nothing is relocated, so a condition is evaluated only where it always was. `a` runs
+ // -2, -1, 0, 1, 2, 3: the first two take the third branch, 0 the first, 1 and 2 the second,
+ // and 3 falls through. `rand` is in [0, 1), so a row reaching a BETWEEN is true whatever it
+ // draws -- the answers here do not depend on memoization. What this rules out is the
+ // alternative that was measured and rejected: hoisting the definition into a `Project` and
+ // guarding that column with the branch condition, which puts `6 / a` outside conditional
+ // evaluation and raises on the `a = 0` row under ANSI.
+ val df = spark.range(0, 6, 1, 1).selectExpr("cast(id as int) - 2 as a")
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN a = 0 THEN false " +
+ "WHEN 6 / a > 2 THEN rand(1) BETWEEN 0 AND 1 " +
+ "WHEN 6 / a < -2 THEN rand(2) BETWEEN 0 AND 1 " +
+ "ELSE false END"),
+ Seq(Row(true), Row(true), Row(false), Row(true), Row(true), Row(false)))
+ }
+ }
+ }
+
+ test("SPARK-58902: a nondeterministic input a branch cannot pre-evaluate is still read once") {
+ onEachEvalPath {
+ // Enough rows that the two copies an inlining implementation makes have to fall out of step:
+ // each copy owns its own generator seeded the same way, so they only differ once the first
+ // comparison has skipped the second copy on some row.
+ val df = spark.range(0, 20, 1, 1)
+ // `randstr` draws a new string per row, so inlining gives the two comparisons of the BETWEEN
+ // two draws, while memoizing gives them one and makes the answer inside a branch the same as
+ // outside one. Its value is a string rather than a primitive, so this also exercises a slot
+ // that cannot be an inlined field. `id` is selected alongside so that `checkAnswer`, which
+ // compares rows as a bag, is comparing per row rather than just counting trues.
+ val inBranch = df.selectExpr(
+ "id", "CASE WHEN id < 0 THEN false ELSE randstr(3, 0) BETWEEN 'a' AND 'zzzz' END")
+ checkAnswer(
+ inBranch, df.selectExpr("id", "randstr(3, 0) BETWEEN 'a' AND 'zzzz'").collect().toSeq)
+
+ // The same for a nondeterministic input wrapped in arithmetic. Comparing against the
+ // branch-free form under a fixed seed is what makes this bite: inlined, the second comparison
+ // draws again, so the row a given draw lands on shifts. The window also has to be one the
+ // first comparison can fail: `BETWEEN 0 AND 1` over `rand` is satisfied by every draw, which
+ // keeps the two copies in lockstep and makes the comparison hold either way. The assertion
+ // below is what keeps that true for these rows rather than merely intended.
+ val branchFree = df.selectExpr("id", "(rand(7) / 1.0) BETWEEN 0.3 AND 0.6").collect()
+ assert(branchFree.exists(!_.getBoolean(1)) && branchFree.exists(_.getBoolean(1)),
+ s"the window is not selective over these rows: ${branchFree.mkString(", ")}")
+ val inBranchRand = df.selectExpr(
+ "id", "CASE WHEN id < 0 THEN false ELSE (rand(7) / 1.0) BETWEEN 0.3 AND 0.6 END")
+ checkAnswer(inBranchRand, branchFree.toSeq)
+ }
+ }
+
+ test("SPARK-58902: a nested With agrees across the evaluation paths") {
+ // The outer `nullif`'s definition is the inner `nullif`, which is itself a `With`, and neither
+ // one is cheap or read once, so both survive inside the branch. That nesting is where each
+ // reference's copy of the definition code matters: the code is generated once per definition
+ // and the text reused, so a nested `With` does not grow its code by a factor per level.
Review Comment:
**Nit (P3):** This is true only when the generated body can be extracted into a helper method. In a whole-stage `Project`, `currentVars` prevents extraction and each reference pastes the cached inner Block, so the focused source-size test still measures `2, 4, 8, ...` growth. Could this comment distinguish generating the definition once for shared state from emitting it once, and scope the bounded-source claim to the row-based standalone 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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5599486418
Thanks for the approval. All four are right; fixed in 1d9ff29.
**Acyclic sibling references.** Agreed, and my comment was wrong rather than merely vague: `withCommonExprs` registers every sibling slot before generating the child, so `d2 = Add(ref(d1), 1)` resolves and completes. What codegen refuses is the two id-keyed shapes plus a definition that reaches its own id, around a cycle or directly, through `filling`. `refsToBind` now says that, and `doGenCode`'s reason for no body enclosing another is restated accordingly: a definition reading a sibling is generated inside that sibling's guard, not inside its own.
**Self-reference test.** Split. With distinct reference objects the interpreted failure was an unbound reference, since the scan over `child` binds only the object in `child`; the shared-object variant below is the one that re-entered a cell and overflowed.
**Malformed sentence.** Rewritten. It was mangled by the comment trim in 1133af6.
**Tree-pattern bit.** Good catch, and cleaner: `containsPattern(WITH_EXPRESSION)` is the same root-or-descendant test in a bit check. Also re-ran `ColumnExpressionSuite` and `WholeStageCodegenSuite` for this one, since it decides `worthAMethod`.
On the red `Build` check at the previous head: that is the `pull_request_target` waiter timing out, not a test failure -- it concluded at 16:48Z while the fork's `sql - other tests`, `sql - extended tests` and Docker integration shards were still running, and the fork run had zero failed jobs.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5578537643
friendly ping @cloud-fan cloud you take another look when you have time? Thanks ~
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3958110074
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -217,30 +209,127 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether substituting this definition into its references is as good as evaluating it once.
+ * True when it is referenced once, since substituting then evaluates it once either way, or when
Review Comment:
**Nit (P3):** `!multiplyReferenced.contains(id)` also accepts definitions with zero references, including expensive nondeterministic ones that are then dropped without evaluation. Please change `referenced once` to `referenced at most once` here and in the `inlineDefsThatGainNothing` Scaladoc so the documented boundary includes that case.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WithExpressionEvalSuite.scala:
##########
@@ -0,0 +1,391 @@
+/*
+ * 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.catalyst.expressions
+
+import org.apache.spark.{SparkException, SparkFunSuite}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodegenFallback, ExprCode, GenerateMutableProjection}
+import org.apache.spark.sql.catalyst.plans.SQLHelper
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType}
+
+/**
+ * Evaluation of [[With]] and the memoization it gives a [[CommonExpressionRef]]. The rewrite that
+ * decides which `With`s reach evaluation at all is covered by `RewriteWithExpressionSuite`.
+ */
+class WithExpressionEvalSuite extends SparkFunSuite with SQLHelper {
+
+ /**
+ * A stand-in for a stateful generator: every evaluation returns the next integer, so a second
+ * evaluation within one row is directly observable. Only used on the interpreted path.
+ */
+ private case class Counter() extends LeafExpression with Nondeterministic {
+ @transient private var n = 0
+ override def stateful: Boolean = true
+ override def dataType: DataType = IntegerType
+ override def nullable: Boolean = false
+ override protected def initializeInternal(partitionIndex: Int): Unit = {}
+ override protected def evalInternal(input: InternalRow): Any = { n += 1; n }
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ throw new UnsupportedOperationException
+ }
+
+ private def counter(): Counter = {
+ val c = Counter()
+ c.initialize(0)
+ c
+ }
+
+ /**
+ * A node that has to be evaluated interpretively even when its parent is generated, so that a
+ * reference below it takes the interpreted path out of generated code.
+ */
+ private case class Fallback(child: Expression) extends UnaryExpression with CodegenFallback {
+ override def dataType: DataType = child.dataType
+ override def eval(input: InternalRow): Any = child.eval(input)
+ override protected def withNewChildInternal(newChild: Expression): Fallback =
+ copy(child = newChild)
+ }
+
+ test("a definition is evaluated once per row however many references read it") {
+ // `ref + ref` is the shape `BETWEEN` produces. Memoized, both references read one value, so the
+ // sum is 2n on the nth row; inlined it would be n + (n + 1).
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ assert((1 to 4).map(_ => w.eval(InternalRow.empty)) == Seq(2, 4, 6, 8))
+ }
+
+ test("a definition is not evaluated on a row that reaches no reference") {
+ val c = counter()
+ // The branch is inside the `With`, so the `With` is entered on every row and does clear its
+ // cells. What it must not do is evaluate the definition before a reference is reached; putting
+ // the branch outside would only test that `If` does not evaluate the arm it did not take.
+ val w = With(c) { case Seq(ref) => If(Literal.TrueLiteral, Literal(-1), Add(ref, ref)) }
+ // -1 whatever the definition does, since the taken arm is a literal: the assertion below is the
+ // load-bearing one, and this only says the evaluation ran.
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(-1, -1, -1))
+ // The counter is still at 0, so the first row that does reach a reference sees 1. Eager filling
+ // would have consumed three values and this would read 8.
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty) == 2)
+ }
+
+ test("a reference behind a short-circuiting operator is not read when the left side is false") {
+ // Neither pre-evaluating the definition into a project nor guarding that column by the branch
+ // can express this: both evaluate on every row the branch is reached on, while `And`
+ // short-circuits before the reference. The `And` is inside the `With` so that the `With` is
+ // entered and only the reference is skipped.
+ val c = counter()
+ val w = With(c) { case Seq(ref) =>
+ And(Literal.FalseLiteral, GreaterThanOrEqual(ref, Literal(1)))
+ }
+ // False whatever the definition does, as in the case above: the counter read is what bites.
+ // This is not that case repeated -- there the arm is skipped by `If`, here by `And`'s short
+ // circuit, and the two are separate paths in both evaluation and codegen.
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(false, false, false))
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty) == 2)
+ }
+
+ test("a nested With memoizes each scope separately") {
+ val outer = counter()
+ val inner = counter()
+ // The outer reference is read twice, and one of its uses wraps an inner `With` whose own
+ // reference is also read twice, so each row is `o + (2i + o)` with `o` and `i` both n.
+ val w = With(outer) { case Seq(o) =>
+ Add(o, With(inner) { case Seq(i) => Add(Add(i, i), o) })
+ }
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(4, 8, 12))
+ }
+
+ test("a With rebuilt by a rule still memoizes") {
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ // Evaluate the original first. A rule rebuilds the definition, and therefore its cell, while
+ // handing the new `With` the original's reference objects, so an implementation that bound the
+ // references once would leave the rebuilt `With` reading this evaluation's cell.
+ assert(w.eval(InternalRow.empty) == 2)
+ // The rule has to hand back a node that is not `==` the one it replaces, or `transformUp` keeps
+ // the original -- `Counter()` is a case class with no parameters, so a fresh one compares equal
+ // to the old one. Wrapping it changes the tree while leaving the definition's value sequence,
+ // its data type and its nullability alone, so the references are carried over rather than
+ // rebuilt, which is the shape that needs the per-evaluation rebinding.
+ val rewritten = w.transformUp { case _: Counter => Add(counter(), Literal(0)) }
+ .asInstanceOf[With]
+ assert(rewritten ne w, "the rule has to have rebuilt something")
+ assert((1 to 3).map(_ => rewritten.eval(InternalRow.empty)) == Seq(2, 4, 6))
+ }
+
+ test("SPARK-58902: a copy of a With owns its references and its cells") {
+ // The definition is deliberately not stateful, which is what makes
+ // `CommonExpressionDef.stateful` load-bearing: `mapChildren` finds nothing changed below the
+ // definition, so without the override the copy is handed the original definition object, cell
+ // included. A `Rand` definition would hide that, since copying it changes the definition's
+ // child and forces a rebuild anyway.
+ // The references are stateful leaves, and `LeafLike.withNewChildrenInternal` returns `this`, so
+ // without the override on `CommonExpressionRef` the two `With`s would share one set of
+ // reference objects while owning two cells.
+ val w1 = With(Literal(1)) { case Seq(ref) => Add(ref, ref) }
+ val w2 = w1.freshCopyIfContainsStatefulExpression().asInstanceOf[With]
+ def refOf(e: Expression): CommonExpressionRef =
+ e.collect { case r: CommonExpressionRef => r }.head
+ assert(w1 ne w2, "a stateful With has to be copied")
+ assert(refOf(w1.child) ne refOf(w2.child), "the copy has to own its references")
+ assert(w1.defs.head.cell ne w2.defs.head.cell, "the copy has to own its cell")
+ }
+
+ test("SPARK-58902: two Withs over one set of references each read their own definition") {
+ // The shape a rule produces when it rewrites only the definition: `mapChildren` hands back the
+ // same `child` object with a new definition, so `withNewChildrenInternal` builds a `With` over
+ // the original's references -- and a rebuilt reference compares equal to the one it replaces,
+ // so `transform` cannot hand the new `With` references of its own. Giving the two `With`s
+ // visibly different definitions shows which cell the references actually read.
+ val w1 = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ val w2 = w1.withNewChildren(
+ IndexedSeq(w1.child, CommonExpressionDef(Literal(100), w1.defs.head.id))).asInstanceOf[With]
+
+ assert(w1.eval(InternalRow.empty) == 2, "w1 reads its own counter")
+ assert(w2.eval(InternalRow.empty) == 200, "w2 reads its own literal")
+ assert(w1.eval(InternalRow.empty) == 4, "w1 read the value w2 memoized")
+ }
+
+ test("SPARK-58902: a definition's code is generated once however many references read it") {
+ // This definition is short and holds no `With`, so it is emitted inline: every reference emits
+ // it inside its own `if (!computed)` guard, and the text appears once per reference -- but it
+ // has to be the same text, generated once, or each copy calls `addMutableState` again and the
+ // definition ends up owning one counter per reference. Two references in mutually exclusive
+ // positions would then draw from two counters sitting at the same position in their sequences,
+ // and hand out one value on two rows.
+ val ctx = new CodegenContext
+ With(MonotonicallyIncreasingID()) { case Seq(ref) => Add(ref, ref) }.genCode(ctx)
+ val counters = ctx.inlinedMutableStates.count { case (_, name) => name.startsWith("count") }
+ assert(counters == 1, s"the definition was generated $counters times")
+ }
+
+ test("SPARK-58902: nesting does not multiply a definition's body when it can go in a method") {
+ // Each level of nested `With`s reads its definition twice, so a reference that pastes the body
+ // doubles it per level. What already stopped that from running away is
+ // `Expression.reduceCodeSize`, which hoists into a method whichever node's code first passes
+ // `methodSplitThreshold` as generation walks up: measured with the 237-character leaf below and
+ // the default threshold, the count settled at 4 from depth 2 on rather than doubling -- where
+ // it settles is where the threshold catches the doubling, so it is a property of those two
+ // numbers and not a law. Emitting the body into a method here makes it 2 at any depth, so the
+ // bound no longer grows with depth wherever the threshold sits.
+ //
+ // What is not covered here is the shape where the input arrives as local variables, which is
+ // what a whole-stage `Project` or `Filter` hands an expression: `currentVars` is set, so
+ // neither this method nor `reduceCodeSize` applies and the body is pasted per reference -- 2,
+ // 4, 8, ... 256 at depths 1 to 8. A bare `CodegenContext` has `INPUT_ROW` set and `currentVars`
+ // null, which is what a method needs.
+ val marker = 1234567
+ def nested(depth: Int): Expression = {
+ val leaf: Expression = Add(BoundReference(0, IntegerType, nullable = false), Literal(marker))
+ (1 to depth).foldLeft(leaf) { (inner, _) =>
+ With(inner) { case Seq(ref) => Add(ref, ref) }
+ }
+ }
+ def markerCount(depth: Int): Int = {
+ val ctx = new CodegenContext
+ val source = nested(depth).genCode(ctx).code.toString + ctx.declareAddedFunctions()
+ marker.toString.r.findAllMatchIn(source).size
+ }
+ // The threshold decides whether the innermost body is inlined at all, so pin it rather than
+ // depend on the default: below the length of that body the fill would go into a method of its
+ // own and the count would be 1, for a reason that has nothing to do with nesting.
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1024") {
+ (1 to 6).foreach { depth =>
+ val count = markerCount(depth)
+ assert(count == 2, s"the innermost body was emitted $count times at depth $depth")
+ }
+ // The 2 is the same fact at every depth: the innermost body, pasted at the innermost `With`'s
+ // two references. What changes from depth 2 on is why it stops there -- each enclosing
+ // level's definition is itself a `With`, so that level's body goes into a method instead of
+ // being pasted, and a call site carries no marker.
+ //
+ // The values still come out right: each level doubles what the one below it produced. Four
+ // levels are enough to show that and keep the product inside Int, which ANSI `Add` would
+ // raise on rather than wrap past depth 10. A deterministic leaf doubles whether a reference
+ // reads a slot or recomputes, so this is an arithmetic and compile check, not a second
+ // reading of the count above.
+ val proj = GenerateMutableProjection.generate(Seq(nested(4)))
+ assert(proj(InternalRow(1)).getInt(0) == (1 + marker) * 16)
+ }
+ }
+
+ test("SPARK-58902: a nested With sharing a reference object restores the outer binding") {
+ // Two `With`s over one reference object, the inner one nested in the outer one's child and
+ // redefining the id the outer one defines. Only a caller building the case class directly
+ // produces this -- `withNewChildrenInternal` shares references between a `With` and its
+ // replacement, which are siblings, and no rule nests a redefinition of a live id -- so this is
+ // hardening rather than a reachable wrong answer. Binding on entry alone is not enough: the
+ // inner `With` rebinds the shared reference to its own definition, so a read after the inner
+ // one returned would answer with the inner definition unless the outer binding is put back.
+ val id = new CommonExpressionId()
+ val outerDef = CommonExpressionDef(counter(), id)
+ val innerDef = CommonExpressionDef(Literal(10), id)
+ val ref = new CommonExpressionRef(outerDef)
+ val inner = With(Add(ref, ref), Seq(innerDef))
+ val outer = With(Add(Add(ref, inner), ref), Seq(outerDef))
+ // The outer definition counts, the inner one does not, so row n is n + (10 + 10) + n. A counter
+ // is what makes the repetition worth something: leaving the inner binding in place gives
+ // n + 20 + 10, and losing memoization of the outer definition gives n + 20 + (n + 1).
+ assert((1 to 3).map(_ => outer.eval(InternalRow.empty)) == Seq(22, 24, 26))
+
+ // The generated path never has to answer this: one id in two nested scopes is refused while
+ // generating, so it cannot quietly disagree with the values above.
+ val literalOuter = CommonExpressionDef(Literal(1), new CommonExpressionId())
+ val literalRef = new CommonExpressionRef(literalOuter)
+ val nested = With(
+ Add(literalRef, With(Add(literalRef, literalRef),
+ Seq(CommonExpressionDef(Literal(10), literalOuter.id)))),
+ Seq(literalOuter))
+ val generated = intercept[SparkException](GenerateMutableProjection.generate(Seq(nested)))
+ assert(generated.getMessage.contains("is already being generated"))
+ }
+
+ test("SPARK-58902: a definition that references its own id fails without recursing") {
+ // Only a caller building the case class directly can produce this, and it is caught rather than
+ // left to recurse: generating the definition re-enters the same slot, and evaluating it would
+ // re-enter the same cell, whose `computed` is set only after the nested evaluation returns.
+ val id = new CommonExpressionId()
+ val proto = CommonExpressionDef(Literal(1), id)
+ val selfDef = CommonExpressionDef(Add(new CommonExpressionRef(proto), Literal(1)), id)
+ val w = With(new CommonExpressionRef(proto), Seq(selfDef))
+
+ val generated = intercept[SparkException] {
+ GenerateMutableProjection.generate(Seq(w))
+ }
+ assert(generated.getMessage.contains("references it"))
+ val interpreted = intercept[SparkException](w.eval(InternalRow.empty))
+ assert(interpreted.getMessage.contains("references an id its own With defines"))
+
+ // The same shape with one reference object in both places, which is the variant that recursed:
+ // the scan over `child` binds that object, so the definition's read of it found a cell still
+ // being filled. Measured before the check: StackOverflowError.
+ val shared = new CommonExpressionRef(proto)
+ val sharedSelfDef = CommonExpressionDef(Add(shared, Literal(1)), id)
+ val sharedFailure =
+ intercept[SparkException](With(shared, Seq(sharedSelfDef)).eval(InternalRow.empty))
+ assert(sharedFailure.getMessage.contains("references an id its own With defines"))
+ }
+
+ test("SPARK-58902: the generated path does not compute a definition on a row that skips it") {
+ // The interpreted no-reference tests use `Counter`, which has no codegen, and the end-to-end
+ // skip case never enters the `With` at all. Here the `With` is entered on every row and the
+ // branch inside it decides whether a reference is reached, so an eager fill would show up as
+ // ids consumed on the rows that take the other arm.
+ val id = MonotonicallyIncreasingID()
+ val w = With(id) { case Seq(ref) =>
+ If(BoundReference(0, BooleanType, nullable = false), ref, Literal(-1L))
+ }
+ val proj = GenerateMutableProjection.generate(Seq(w))
+ proj.initialize(0)
+ val got = Seq(false, true, false, true).map { reached =>
+ proj(InternalRow(reached)).getLong(0)
+ }
+ // Only the rows that reach the reference take an id, so the ids stay 0 and 1.
+ assert(got == Seq(-1L, 0L, -1L, 1L))
+ }
+
+ test("SPARK-58902: two surviving definitions in one With keep their own state") {
+ // Every other runtime case has a single definition, so a slot that aliased two definitions, or
+ // a clear that reset the wrong one, would go unnoticed. The two are scaled differently to make
+ // which slot answered visible in the value.
+ def build(): Expression = With(
+ MonotonicallyIncreasingID(),
+ Multiply(MonotonicallyIncreasingID(), Literal(100L))) { case Seq(a, b) =>
+ Add(Add(a, a), Add(b, b))
+ }
+ // Row n reads a = n and b = 100n once each, so the sum is 2n + 200n. Aliasing the two slots
+ // would give 4n or 400n, and clearing only one of them would drift after the first row.
+ val proj = GenerateMutableProjection.generate(Seq(build()))
+ proj.initialize(0)
+ assert((0 until 3).map(_ => proj(InternalRow.empty).getLong(0)) == Seq(0L, 202L, 404L))
+
+ val interpreted = build()
+ interpreted.foreach {
+ case n: Nondeterministic => n.initialize(0)
+ case _ =>
+ }
+ assert((0 until 3).map(_ => interpreted.eval(InternalRow.empty)) == Seq(0L, 202L, 404L))
+ }
+
+ test("SPARK-58902: eval refuses the two shapes whose bindings it could not keep straight") {
+ // Both need the case class constructor, and both were answered by `eval` before this check --
+ // with a wrong value, while `GenerateMutableProjection` raised. Refusing them keeps the two
+ // paths saying the same thing about the same tree.
+ val id = new CommonExpressionId()
+ val first = CommonExpressionDef(Literal(1), id)
+ val ref = new CommonExpressionRef(first)
+ // Two definitions with one id: `defs.map(d => d.id -> d).toMap` kept the last, so both
+ // references read 2 and the answer was 4 where the first definition says 2.
+ val duplicateIds = With(Add(ref, ref), Seq(first, CommonExpressionDef(Literal(2), id)))
+ val duplicate = intercept[SparkException](duplicateIds.eval(InternalRow.empty))
+ assert(duplicate.getMessage.contains("Duplicate common expression ids"))
+
+ // One definition object in two scopes: the inner `With` clears the cell the outer one had
+ // already filled, so the definition answered twice inside one entry -- 1, then 2 from the
+ // counter, giving 3 where the memoization contract says 1 + 0 + 1.
+ val counterDef = CommonExpressionDef(counter(), new CommonExpressionId())
+ val counterRef = new CommonExpressionRef(counterDef)
+ val shared = With(
+ Add(Add(counterRef, With(Literal(0), Seq(counterDef))), counterRef),
+ Seq(counterDef))
+ val sharedFailure = intercept[SparkException](shared.eval(InternalRow.empty))
+ assert(sharedFailure.getMessage.contains("shares a common expression definition object"))
+ }
+
+ test("SPARK-58902: a refused scope leaves no slot behind in the context") {
+ // The duplicate-id check throws partway through allocating, after earlier definitions of the
+ // same `With` are already in scope. If those were left there, a later reference to one of those
+ // ids in the same context would resolve an orphan and generate code from it instead of saying
+ // the id is not in scope.
+ val ctx = new CodegenContext
+ val id = new CommonExpressionId()
+ val w = With(Literal(0), Seq(CommonExpressionDef(Literal(1), id), CommonExpressionDef(
+ Literal(2), id)))
+ val refused = intercept[SparkException](w.genCode(ctx))
+ assert(refused.getMessage.contains("is already being generated"))
+ val gone = intercept[SparkException](ctx.getCommonExpr(id.id))
+ assert(gone.getMessage.contains("is not in scope"))
+ }
+
+ test("SPARK-58902: a reference under a CodegenFallback still memoizes") {
+ // The generated code clears the codegen flags, not the cells, so a reference reached through a
+ // `CodegenFallback`'s `eval` cannot use them: the whole `With` falls back to `eval`, which
+ // binds and clears the cells itself. Generating part of it would be worse than either -- a
+ // definition reached from both sides would hold one value in the slots and another in the cell.
+ val w = With(counter()) { case Seq(ref) => Add(Fallback(ref), Fallback(ref)) }
Review Comment:
**Non-blocking (P2):** Both fallback tests create `Counter` through `counter()`, which calls `initialize(0)` before projection generation. They therefore still pass if `CodegenFallback.generate` stops registering nested `Nondeterministic` descendants for partition initialization. Please construct `Counter()` directly in a generated fallback `With`, call `proj.initialize(0)`, and assert `Seq(2, 4, 6, 8)` so this covers initialization as well as memoization.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +75,222 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with the definition each
+ * names. The list is found once, since the tree does not change between evaluations.
+ *
+ * Only `child` is scanned for references, which relies on a reference to one of these definitions
+ * never living inside another one of them. The helper `With(commonExprs: _*)(replaced)` builds
+ * the references outside the definitions and cannot produce that, and `RewriteWithExpression`,
+ * which does rewrite inside a `With`, only ever replaces a reference with its definition's child
+ * or with an attribute -- it never puts a reference inside a definition. The case class
+ * constructor does take `child` and `defs` directly, so a caller can hand a definition a
+ * reference to any of these ids; that reference is then never bound, and evaluating it raises
+ * "Cannot evaluate a common expression reference outside its With", which is the failure to want.
Review Comment:
**Nit (P3):** `the failure to want` is not grammatical wording for an intentional error. Please use `the desired failure` or `the failure we want`.
##########
docs/sql-migration-guide.md:
##########
@@ -26,6 +26,7 @@ license: |
- Since Spark 4.4, the optimizer rule `SimplifyCaseConversionExpressions` no longer simplifies mixed case conversions `LOWER(UPPER(str))` to `LOWER(str)` or `UPPER(LOWER(str))` to `UPPER(str)`. This preserves standard Unicode case mapping semantics for non-ASCII characters such as Latin small letter dotless i (`ı` U+0131) and micro sign (`µ` U+00B5).
- Since Spark 4.4, the Oracle JDBC datasource maps Oracle DATE and TIMESTAMP to TimestampNTZType unconditionally, ignoring the JDBC read option `preferTimestampNTZ` (DATE only when the driver default `oracle.jdbc.mapDateToTimestamp=true` surfaces it as TIMESTAMP; otherwise DATE stays DateType). In 4.3 and earlier these were read as TimestampType, or as TimestampNTZType only when `preferTimestampNTZ=true`. This faithfully represents these zoneless Oracle types. Unaffected: Oracle DATE read with `oracle.jdbc.mapDateToTimestamp=false` (mapped to DateType, as before), and TIMESTAMP WITH TIME ZONE / TIMESTAMP WITH LOCAL TIME ZONE. Symmetrically, a TimestampNTZType column written to Oracle is now bound zoneless via `setObject` rather than a JVM-default-zone `java.sql.Timestamp`, which can change the stored value for a wall-clock that falls in the JVM default zone's DST gap. To restore the previous read and write behavior, set `spark.sql.legacy.oracle.timestampNTZMapping.enabled` to `true`
.
+- Since Spark 4.4, a common expression that a conditional branch reads more than once is evaluated once rather than once per read. `BETWEEN` and `NULLIF` read their input twice, so an expression such as `CASE WHEN c THEN monotonically_increasing_id() BETWEEN 3 AND 5 END` previously gave each read its own value and could return a different result; both reads now see one value. Queries whose results depended on the earlier behaviour change, and an optimized or physical plan may now show a `with` node for such a branch.
Review Comment:
**Non-blocking (P2):** This promises once-only evaluation for every multiply read common expression, but `canSubstitute` still inlines multiply referenced definitions that are cheap and report deterministic. For example, literal `aes_encrypt` without an IV is foldable and reports deterministic while each inlined copy can generate a different random IV. Please limit the entry to definitions retained for memoization, especially nondeterministic ones, or explicitly disclose the cheap-and-deterministic substitution exception.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -217,30 +209,127 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether substituting this definition into its references is as good as evaluating it once.
+ * True when it is referenced once, since substituting then evaluates it once either way, or when
+ * it is cheap enough to evaluate repeatedly and deterministic, so the repeated evaluations agree.
+ *
+ * `CollapseProject.isCheap` answers what one evaluation costs, not whether a second one is
+ * allowed: it admits a `PythonUDF`, which may be nondeterministic. Determinism is checked here
+ * rather than read into the cost test, so that a reader of either call site can see which of the
+ * two questions is being asked. A `PythonUDF` is still in the expression tree when this runs --
+ * `SparkOptimizer` extracts them in a batch after the one holding this rule -- so that case is
Review Comment:
**Nit (P3):** The singular `A PythonUDF` does not agree with `extracts them`. Please pluralize `PythonUDF` or change `them` to `it`.
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5439773368
Could you take another look when you have time? Thanks. @cloud-fan
--
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-58818][SQL] Pre-evaluate a nondeterministic common expression in a conditional branch [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on PR #58045:
URL: https://github.com/apache/spark/pull/58045#issuecomment-5315373457
Thank you @LuciferYang!
--
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-58818][SQL] Fix wrong results from evaluating a common expression once per reference in a conditional branch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3904659842
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,124 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with the definition each
+ * names. The list is found once, since the tree does not change between evaluations.
+ *
+ * Only `child` is scanned, which relies on a reference to one of these definitions never living
+ * inside another one of them. The helper `With(commonExprs: _*)(replaced)` builds the references
+ * outside the definitions and cannot produce that, and `RewriteWithExpression`, which does
+ * rewrite inside a `With`, only ever replaces a reference with its definition's child or with an
+ * attribute -- it never puts a reference inside a definition. The case class constructor does
+ * take `child` and `defs` directly, so a caller can hand a definition a reference to any of these
+ * ids; that reference is then never bound, and evaluating it raises "Cannot evaluate a common
+ * expression reference outside its With", which is the failure to want. Scanning `children`
+ * instead would look safer and be worse -- it would bind such a reference, and since
+ * `CommonExpressionCell.get` sets `computed` only after the nested evaluation returns, that loud
+ * error would become a StackOverflowError. A nested `With` is not affected either way: `children`
+ * is `child +: defs`, so this scan already descends into an inner `With`'s own definitions.
+ */
+ @transient private lazy val refsToBind: Seq[(CommonExpressionRef, CommonExpressionDef)] = {
+ val idToDef = defs.map(d => d.id -> d).toMap
+ child.collect { case r: CommonExpressionRef if idToDef.contains(r.id) => (r, idToDef(r.id)) }
+ }
+
+ /**
+ * Binds this `With`'s references to its own cells, clears them, and evaluates the child. A
+ * reference reached by that evaluation computes its definition once and every later reference
+ * reads the value back, so a definition is evaluated where the child would have evaluated it,
+ * once, rather than once per reference. See [[CommonExpressionCell]].
+ *
+ * The binding is redone on every evaluation rather than once, because a reference can be reached
+ * from two `With`s. `withNewChildrenInternal` cannot hand the new `With` its own references: a
+ * rebuilt reference compares equal to the one it replaces, since the binding it carries is not
+ * part of its equality, so `transform` keeps the original. Binding once would then leave the
+ * `With` that bound last deciding what both of them read. Rebinding costs one pass over the
+ * references, two for a `BETWEEN`, and makes the `With` currently evaluating always the owner.
+ */
+ override def eval(input: InternalRow): Any = {
+ refsToBind.foreach { case (ref, exprDef) => ref.bindTo(exprDef) }
Review Comment:
**Blocking (P1):** Rebinding on entry does not preserve the outer scope when a nested copied `With` shares reference objects. `CommonExpressionRef` has one mutable definition/cell pair, so the inner `bindTo` overwrites it and this method never restores the previous pair. With outer definition `1` and an inner definition with the same id returning `10`, `outer(Add(Add(ref, inner(ref + ref)), ref))` returns `31` instead of the lexically scoped `22`. Please save and restore each prior binding in `finally` (or use a binding stack), and cover this nested-copy case repeatedly.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,124 @@ case class With(child: Expression, defs: Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with the definition each
+ * names. The list is found once, since the tree does not change between evaluations.
+ *
+ * Only `child` is scanned, which relies on a reference to one of these definitions never living
+ * inside another one of them. The helper `With(commonExprs: _*)(replaced)` builds the references
+ * outside the definitions and cannot produce that, and `RewriteWithExpression`, which does
+ * rewrite inside a `With`, only ever replaces a reference with its definition's child or with an
+ * attribute -- it never puts a reference inside a definition. The case class constructor does
+ * take `child` and `defs` directly, so a caller can hand a definition a reference to any of these
+ * ids; that reference is then never bound, and evaluating it raises "Cannot evaluate a common
+ * expression reference outside its With", which is the failure to want. Scanning `children`
+ * instead would look safer and be worse -- it would bind such a reference, and since
+ * `CommonExpressionCell.get` sets `computed` only after the nested evaluation returns, that loud
+ * error would become a StackOverflowError. A nested `With` is not affected either way: `children`
+ * is `child +: defs`, so this scan already descends into an inner `With`'s own definitions.
+ */
+ @transient private lazy val refsToBind: Seq[(CommonExpressionRef, CommonExpressionDef)] = {
+ val idToDef = defs.map(d => d.id -> d).toMap
+ child.collect { case r: CommonExpressionRef if idToDef.contains(r.id) => (r, idToDef(r.id)) }
+ }
+
+ /**
+ * Binds this `With`'s references to its own cells, clears them, and evaluates the child. A
+ * reference reached by that evaluation computes its definition once and every later reference
+ * reads the value back, so a definition is evaluated where the child would have evaluated it,
+ * once, rather than once per reference. See [[CommonExpressionCell]].
+ *
+ * The binding is redone on every evaluation rather than once, because a reference can be reached
+ * from two `With`s. `withNewChildrenInternal` cannot hand the new `With` its own references: a
+ * rebuilt reference compares equal to the one it replaces, since the binding it carries is not
+ * part of its equality, so `transform` keeps the original. Binding once would then leave the
+ * `With` that bound last deciding what both of them read. Rebinding costs one pass over the
+ * references, two for a `BETWEEN`, and makes the `With` currently evaluating always the owner.
+ */
+ override def eval(input: InternalRow): Any = {
+ refsToBind.foreach { case (ref, exprDef) => ref.bindTo(exprDef) }
+ defs.foreach(_.cell.clear())
+ child.eval(input)
+ }
+
+ // The cells are cleared on entry, so this holds state for the duration of one evaluation.
+ override def stateful: Boolean = true
+
+ /**
+ * Whether one of this `With`'s references sits somewhere that will be evaluated interpretively
+ * even though this `With` is generated. Two shapes do that: a [[CodegenFallback]], which is
+ * evaluated by calling `eval` on it from the generated code, and a nested `With` that itself
+ * takes the fallback below -- `With` does not mix in `CodegenFallback`, so it has to be named
+ * here rather than matched as one. A reference reached that way needs its cell bound and
+ * cleared, which the generated code does not do: it clears the codegen flags.
+ *
+ * This is the same shape `EquivalentExpressions.childrenToRecurse` already refuses to look past,
+ * for the same reason.
+ *
+ * Each level memoizes, but `holdsMyRef` runs again at every nested `With` the scan passes, so a
+ * chain of them nested in each other's `child` costs on the order of the square of the depth.
+ * `nullif(a, nullif(b, c))` does produce such a chain -- only the memoized input becomes a
+ * definition, the rest stays in `child` -- but these chains are shallow in practice. Reading a
+ * nested `With`'s own `lazy val` from here also takes its monitor while holding this one; the
+ * edges only ever run from an ancestor to a proper descendant of an immutable tree, so the order
+ * is a strict partial one and cannot deadlock. `canonicalizationIdMap` below relies on the same.
+ */
+ @transient private lazy val refUnderCodegenFallback: Boolean = {
+ val ids = defs.map(_.id).toSet
+ def holdsMyRef(e: Expression): Boolean = e.exists {
+ case r: CommonExpressionRef => ids.contains(r.id)
+ case _ => false
+ }
+ child.exists {
+ case f: CodegenFallback => holdsMyRef(f)
+ case w: With if w.refUnderCodegenFallback => holdsMyRef(w)
+ case _ => false
+ }
+ }
+
+ /**
+ * Clears each definition's flag, then generates the child. The flags are cleared in the same
+ * block the child is generated into, so a reference cannot run against a flag left set by an
+ * earlier row: on a row that does not reach the branch holding this `With`, neither the clearing
+ * nor any reference runs.
+ *
+ * When a reference sits under a [[CodegenFallback]], or inside a nested `With` that itself falls
+ * back, the whole `With` is evaluated interpretively instead. Generating the child would leave
+ * that reference reading a cell nobody bound and nobody clears, and generating part of it is
+ * worse still: a definition reached from both sides would be computed once through the flags and
+ * once through the cell, holding two values for one row. [[eval]] binds and clears both, so
+ * handing it the whole subtree keeps one mechanism in play. `ctx.INPUT_ROW` is available on that
+ * path because `CollapseCodegenStages.supportCodegen` turns whole-stage codegen off for a plan
+ * whose expressions hold the offending `CodegenFallback` -- it is visible there, since a `With`
+ * in a conditional branch reaches execution inside `plan.expressions` like any other expression.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ if (refUnderCodegenFallback) {
+ return CodegenFallback.generate(this, ctx, ev)
Review Comment:
**Blocking (P1):** This fallback registers the same stateful `With` instance for every generated occurrence. `GenerateOrdering` generates the key once for row A and once for row B, so both emitted calls advance one stateful definition. For example, a `With(MonotonicallyIncreasingID())` whose reference is below a non-leaf `CodegenFallback` makes `compare(row, row)` observe consecutive values and return nonzero. Please give each fallback occurrence independent state (as the interpreted ordering does for its right side) and add a `GenerateOrdering` regression asserting reflexivity.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,126 @@ class CodegenContext extends Logging {
throw QueryExecutionErrors.lambdaVariableNotDefinedError(id))
}
+ /**
+ * The slots a `CommonExpressionRef` reads: the value and its nullness, plus the flag saying
+ * whether this row has computed them yet, and the definition to compute them from.
+ */
+ case class CommonExprSlots(
+ value: ExprCode,
+ computed: String,
+ definition: Expression,
+ private var fillCode: Option[Block] = None) {
+
+ /**
+ * The code that computes the definition into the slots and sets `computed`, which a reference
+ * emits behind that flag. Generated once and cached, so that every reference shares whatever
+ * mutable state the definition allocated, such as an RNG, rather than getting its own.
+ *
+ * The body goes into a method where it can and is worth it -- a definition that is or holds
+ * another `With`, or a body past the split threshold -- leaving a reference's own code a single
+ * call: such a definition would otherwise have its body pasted once per reference at every
+ * level. `Expression.reduceCodeSize` already keeps that from running away, by hoisting
+ * whichever node's code first passes the same threshold, so what a method here buys is the body
+ * once per scope rather than once per reference up to that threshold, and a bound that no
+ * longer grows with depth wherever the threshold sits. It is only possible where the definition
+ * reads the input row rather than local variables -- the condition `reduceCodeSize` splits
+ * under, and for the same reason. Whole-stage codegen generates the operators a `With` reaches
+ * with `currentVars` set, so neither applies there and the body is pasted per reference.
+ */
+ def fill: Block = {
+ if (fillCode.isEmpty) {
+ val defGen = definition.genCode(CodegenContext.this)
Review Comment:
**Non-blocking (P2):** `fillCode` is still empty while `definition.genCode` runs. If a directly constructed definition contains a `CommonExpressionRef` to its own id, that reference resolves this same slot and recursively re-enters `fill` until `StackOverflowError`; interpreted evaluation instead reaches the controlled unbound-reference Spark error documented in `With.scala`. Please publish an in-progress state or reject the cycle before recursing, and add a generated-evaluation regression that explicitly rules out `StackOverflowError`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -98,6 +104,55 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether substituting this definition into its references is as good as evaluating it once.
+ * True when it is referenced once, since substituting then evaluates it once either way, or when
+ * it is cheap enough to evaluate repeatedly and deterministic, so the repeated evaluations agree.
+ *
+ * `CollapseProject.isCheap` answers what one evaluation costs, not whether a second one is
+ * allowed: it admits a `PythonUDF`, which may be nondeterministic. Determinism is checked here
+ * rather than read into the cost test, so that a reader of either call site can see which of the
+ * two questions is being asked.
+ */
+ private def canSubstitute(
+ child: Expression,
+ id: CommonExpressionId,
+ commonExprIdSet: Set[CommonExpressionId]): Boolean = {
+ !commonExprIdSet.contains(id) || (CollapseProject.isCheap(child) && child.deterministic)
+ }
+
+ /**
+ * `w` with every definition that gains nothing from being memoized inlined into its references:
+ * one cheap enough to evaluate twice, and one that is referenced once anyway. This is the test
+ * the main rewrite already applies before it hoists a definition into a project.
+ *
+ * Inlining matters beyond the per-row bookkeeping it saves. A `With` is not foldable, so it hides
+ * whatever it wraps from `ConstantFolding`, `PushFoldableIntoBranches`, `SimplifyConditionals`
+ * and `ReplaceNullWithFalseInPredicate`, all of which run in later batches. Dropping the `With`
+ * once nothing is left to memoize keeps `CASE WHEN c THEN nullif(1, 1) END` folding as it did
+ * before this rule learned to leave one behind.
+ */
+ private def inlineDefsThatGainNothing(
+ w: With,
+ commonExprIdSet: Set[CommonExpressionId]): Expression = {
+ val (toInline, toKeep) = w.defs.partition { d =>
+ canSubstitute(d.child, d.id, commonExprIdSet)
+ }
+ if (toInline.isEmpty) {
+ w
+ } else {
+ val refToExpr = toInline.map(d => d.id -> d.child).toMap
+ val newChild = w.child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
+ // A ref of a definition kept here, or of an enclosing `With`, is left for its owner.
+ case ref: CommonExpressionRef if refToExpr.contains(ref.id) => refToExpr(ref.id)
+ }
+ // `copy` rather than `withNewChildren`, which requires the child count to be unchanged. The
+ // references of the kept definitions are carried over unbound; the new `With` binds them, and
+ // `w` is discarded here, so no two `With`s hold the same reference.
+ if (toKeep.isEmpty) newChild else w.copy(child = newChild, defs = toKeep)
Review Comment:
**Non-blocking (P2):** Could we add a focused test for this partial branch? The current branch-local cases each use one definition, so they exercise all-inline or all-keep results but not `copy(child = newChild, defs = toKeep)`. A two-definition `With` under `CaseWhen`/`Coalesce`, with one cheap deterministic definition and one expensive multiply referenced definition, should assert that only the expensive definition and its id remain, no removed-id reference survives, and a second optimizer run is identical.
--
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