From 11aabadcaf8073467c4b6b824ca1cf117061745c Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 17 Sep 2026 04:58:38 -0400 Subject: [PATCH] [SPARK-59603][SQL] Memoize a subexpression repeated inside a conditional branch Subexpression elimination cannot reach a subexpression that is repeated inside a single branch of an `if` or a `case when`, and the reason is structural rather than an oversight: it evaluates its candidates before the projection, so `EquivalentExpressions` only collects what is always evaluated (`ConditionalExpression.alwaysEvaluatedInputs`) plus what every branch of a group shares (`branchGroups`, an intersection). A subexpression repeated inside one branch body and nowhere else belongs to no group, so each occurrence is evaluated for every row that reaches that branch. Since SPARK-58818 a `With` is evaluable and memoizes per evaluation, which is what this gap needs: the definition sits inside the branch, so nothing is computed for a row that takes another one, and the references read the one value. `MemoizeCommonExpressionsInBranches` rewrites the tallest repeated subexpression of each conditionally evaluated child of an `if` / `case when` into a `With`. It runs after the simplification rules, so what it memoizes is what survives them, and it leaves the `With` in the branch -- the shape `RewriteWithExpression`, which runs right after `FinishAnalysis`, keeps anyway, so that rule does not have to run again. Over `select case when id < 5 then udf(id) + udf(id) else 0 end from range(0, 10, 1, 1)` with a counting UDF, the definition is evaluated 5 times with the rule and 10 times without: once for each of the five rows that reach the branch, rather than once per reference. The rule is behind `spark.sql.optimizer.memoizeCommonExpressionsInBranches.enabled`, default false, so no plan changes unless it is turned on. A candidate is skipped when reading it back would not be cheaper than recomputing it (`CollapseProject.isCheap`), when it holds a common-expression reference or a lambda variable -- neither can be evaluated where the definition would sit -- or when it holds a subquery expression. An aggregate, window or generator expression rules out the whole branch body rather than just a candidate holding one: the planner takes those out of the tree they stand in, and a reference left behind would be evaluated where its `With`, and so its definition, is no longer above it. A branch that already holds a `With` is left alone too: `RewriteWithExpression` defers a nested `With` to a pass that no longer runs by this point. Being stateful is deliberately not a reason to refuse. A `ScalaUDF` is stateful because its encoder reuses an `UnsafeRow`, and an expensive UDF repeated in a branch is what this rule is for; the definition is evaluated once and read back within the same row, which is what a `With` left in a branch by `nullif(udf(x), 0)` already does. What would be unsafe is a value that changes per evaluation, and those are nondeterministic -- `EquivalentExpressions` never records one. --- .../MemoizeCommonExpressionsInBranches.scala | 117 +++++++++++++++++ .../sql/catalyst/optimizer/Optimizer.scala | 4 + .../apache/spark/sql/internal/SQLConf.scala | 15 +++ ...oizeCommonExpressionsInBranchesSuite.scala | 120 ++++++++++++++++++ .../spark/sql/ColumnExpressionSuite.scala | 27 ++++ 5 files changed, 283 insertions(+) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranches.scala create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranchesSuite.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranches.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranches.scala new file mode 100644 index 0000000000000..99ad2381c4989 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranches.scala @@ -0,0 +1,117 @@ +/* + * 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.optimizer + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.expressions.objects.LambdaVariable +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{CASE_WHEN, IF} +import org.apache.spark.sql.internal.SQLConf + +/** + * Rewrites a subexpression that occurs more than once inside a single branch of an `if` or a + * `case when` into a [[With]], so that the rows reaching that branch evaluate it once. + * + * This is the one place subexpression elimination structurally cannot reach. It evaluates its + * candidates before the projection, so `EquivalentExpressions` only collects expressions that are + * always evaluated (`ConditionalExpression.alwaysEvaluatedInputs`) plus those shared by every + * branch of a group (`branchGroups`), which is an intersection: a subexpression repeated inside one + * branch body and nowhere else is in no group and is never eliminated. A `With` covers it because + * it memoizes per evaluation rather than ahead of it -- nothing is computed for a row that does not + * reach the branch. + * + * The rule runs after the simplification rules, so what it memoizes is what survives them, and it + * leaves the `With` in the branch: `RewriteWithExpression` has already run by then, and a + * definition inside a conditional branch is one that rule keeps anyway. + */ +object MemoizeCommonExpressionsInBranches extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.MEMOIZE_COMMON_EXPRESSIONS_IN_BRANCHES)) { + plan + } else { + plan.transformWithPruning(_.containsAnyPattern(IF, CASE_WHEN)) { + case p => p.transformExpressionsUpWithPruning(_.containsAnyPattern(IF, CASE_WHEN)) { + case i: If => i.copy(trueValue = memoize(i.trueValue), falseValue = memoize(i.falseValue)) + case c: CaseWhen => + // The first condition is always evaluated, so it is subexpression elimination's to + // take; every other condition and every value is reached only for some rows. + val branches = c.branches.zipWithIndex.map { case ((cond, value), i) => + (if (i == 0) cond else memoize(cond), memoize(value)) + } + c.copy(branches = branches, elseValue = c.elseValue.map(memoize)) + } + } + } + } + + private def memoize(body: Expression): Expression = { + if (skipBody(body)) { + body + } else { + val equivalence = new EquivalentExpressions + equivalence.addExprTree(body) + // `getCommonSubexpressions` is ordered by height, so the last one that qualifies is the + // tallest: memoizing it subsumes every repeated subtree inside it. + equivalence.getCommonSubexpressions.reverse.find(worthMemoizing) match { + case Some(common) => + With(common) { case Seq(ref) => + body.transformDown { case e if e.semanticEquals(common) => ref } + } + case None => body + } + } + } + + private def skipBody(body: Expression): Boolean = body.exists { + // A body that already holds a `With` is left alone: `RewriteWithExpression` defers a nested + // `With` to its next pass and that rule is behind us, so a nested definition would never be + // looked at again. + case _: With => true + // An aggregate, window or generator expression anywhere in the body rules the whole body out, + // not just candidates holding one. The planner takes those expressions out of the tree they + // stand in -- `PhysicalAggregation` gives an aggregate its own operator -- and a reference left + // behind would be evaluated with its `With`, and so its definition, no longer above it. + case _: AggregateExpression | _: WindowExpression | _: Generator => true + case _ => false + } + + private def worthMemoizing(candidate: Expression): Boolean = { + // Reading the value back has to cost less than computing it again. + !CollapseProject.isCheap(candidate) && + !candidate.exists { + // A reference belongs to the `With` that binds it and a lambda variable to its loop; + // neither can be evaluated where the definition would sit. An aggregate, window or + // generator expression has to stay where the planner looks for it, and a subquery + // expression carries a plan that later rules still rewrite. + case _: CommonExpressionRef | _: CommonExpressionDef => true + case _: NamedLambdaVariable | _: LambdaVariable => true + case _: AggregateExpression | _: WindowExpression | _: Generator => true + case _: PlanExpression[_] => true + case _ => false + } + // `stateful` is deliberately not a reason to refuse. A `ScalaUDF` is stateful because its + // encoder reuses an `UnsafeRow`, and it is exactly what this rule exists for; the definition is + // evaluated once and read back within the same row, which is what `RewriteWithExpression` + // already does with a `With` a `nullif(udf(x), 0)` leaves in a branch. What would be unsafe is + // an expression whose value changes per evaluation, and those are nondeterministic -- + // `EquivalentExpressions` never records one. + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index f1484ff02e154..758e17cfebccf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -301,6 +301,10 @@ abstract class Optimizer(catalogManager: CatalogManager) RemoveNoopOperators), // This batch must be executed after the `RewriteSubquery` batch, which creates joins. Batch("NormalizeFloatingNumbers", Once, NormalizeFloatingNumbers), + // Runs this late so that what gets memoized is what survives the simplification rules. The + // `With` it creates stays in the branch, which is the shape `RewriteWithExpression` -- long + // since run, right after `FinishAnalysis` -- keeps anyway. + Batch("Memoize common expressions in branches", Once, MemoizeCommonExpressionsInBranches), Batch("ReplaceUpdateFieldsExpression", Once, ReplaceUpdateFieldsExpression))) // remove any batches with no rules. this may happen when subclasses do not add optional rules. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 2ec67312f0218..90b1186e9694b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1670,6 +1670,18 @@ object SQLConf { .booleanConf .createWithDefault(false) + val MEMOIZE_COMMON_EXPRESSIONS_IN_BRANCHES = + buildConf("spark.sql.optimizer.memoizeCommonExpressionsInBranches.enabled") + .internal() + .doc("When true, a subexpression that occurs more than once inside a single branch of an " + + "`if` or `case when` is rewritten into a `With`, so that the rows reaching that branch " + + "evaluate it once. Subexpression elimination cannot cover this: it evaluates its " + + "candidates before the projection, so it only considers expressions that are always " + + "evaluated, plus those shared by every branch of a group.") + .version("4.4.0") + .booleanConf + .createWithDefault(false) + val SUBEXPRESSION_ELIMINATION_FILTER_EXEC_ENABLED = buildConf("spark.sql.subexpressionElimination.filterExec.enabled") .internal() @@ -9284,6 +9296,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def subexpressionEliminationSkipForShotcutExpr: Boolean = getConf(SUBEXPRESSION_ELIMINATION_SKIP_FOR_SHORTCUT_EXPR) + def memoizeCommonExpressionsInBranches: Boolean = + getConf(MEMOIZE_COMMON_EXPRESSIONS_IN_BRANCHES) + def subexpressionEliminationFilterExecEnabled: Boolean = getConf(SUBEXPRESSION_ELIMINATION_FILTER_EXEC_ENABLED) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranchesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranchesSuite.scala new file mode 100644 index 0000000000000..5b241daef85ab --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MemoizeCommonExpressionsInBranchesSuite.scala @@ -0,0 +1,120 @@ +/* + * 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.optimizer + +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.dsl.plans._ +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.PlanTest +import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.internal.SQLConf + +class MemoizeCommonExpressionsInBranchesSuite extends PlanTest { + + private object Optimize extends RuleExecutor[LogicalPlan] { + val batches = + Batch("Memoize common expressions in branches", Once, + MemoizeCommonExpressionsInBranches) :: Nil + } + + private val relation = LocalRelation($"a".int, $"b".int) + // Resolved attributes, so that the expressions built here are the ones the analyzed plan holds. + private val Seq(a, b) = relation.output + private val common = Multiply(a, b) + + private def optimize(plan: LogicalPlan, enabled: Boolean = true): LogicalPlan = { + withSQLConf(SQLConf.MEMOIZE_COMMON_EXPRESSIONS_IN_BRANCHES.key -> enabled.toString) { + Optimize.execute(plan) + } + } + + private def withExprs(plan: LogicalPlan): Seq[With] = + plan.expressions.flatMap(_.collect { case w: With => w }) + + test("a subexpression repeated in one branch body is memoized") { + val plan = relation + .select(CaseWhen(Seq((GreaterThan(a, Literal(0)), Add(common, common))), Literal(0)).as("r")) + .analyze + val optimized = optimize(plan) + val memoized = withExprs(optimized) + assert(memoized.length == 1, s"expected one With: $optimized") + val definitions = memoized.head.defs.map(_.child) + assert(definitions == Seq(common), s"memoized the wrong expression: $definitions") + val references = memoized.head.child.collect { case r: CommonExpressionRef => r } + assert(references.length == 2, s"expected two references: ${memoized.head.child}") + // The `With` has to sit inside the branch. Above the `CaseWhen` it would be evaluated for every + // row, which is what subexpression elimination already does and what this cannot do. + val branchValues = optimized.expressions.flatMap(_.collect { + case c: CaseWhen => c.branches.map(_._2) + }.flatten) + assert(branchValues.forall(_.isInstanceOf[With]), s"the With left the branch: $optimized") + } + + test("the plan is untouched while the config is off") { + val plan = relation + .select(CaseWhen(Seq((GreaterThan(a, Literal(0)), Add(common, common))), Literal(0)).as("r")) + .analyze + comparePlans(optimize(plan, enabled = false), plan) + } + + test("a cheap subexpression is left alone") { + // Reading back a memoized value costs a field read and a flag check, so a foldable expression + // -- or an attribute, or anything else `CollapseProject.isCheap` accepts -- gains nothing. + val cheap = Add(Literal(1), Literal(2)) + val plan = relation + .select(CaseWhen(Seq((GreaterThan(a, Literal(0)), Add(cheap, cheap))), Literal(0)).as("r")) + .analyze + assert(withExprs(optimize(plan)).isEmpty, "a cheap expression was memoized") + } + + test("the first condition of a case when is left to subexpression elimination") { + // It is evaluated for every row that reaches the conditional, so it is already covered by + // `ConditionalExpression.alwaysEvaluatedInputs`, where elimination costs no flag check. + val plan = relation + .select(CaseWhen( + Seq((GreaterThan(Add(common, common), Literal(0)), Literal(1))), Literal(0)).as("r")) + .analyze + assert(withExprs(optimize(plan)).isEmpty, "the always-evaluated condition was memoized") + } + + test("a branch body holding an aggregate expression is left alone") { + // `PhysicalAggregation` gives each aggregate expression its own operator, so a `With` wrapped + // around one would no longer be above the reference when that reference is evaluated. The + // candidate here (`a * b`) is fine on its own; what rules the body out is the aggregate above + // it. + val aggregate = sum(common) + val plan = relation + .groupBy(b)(CaseWhen( + Seq((GreaterThan(b, Literal(0)), Add(aggregate, aggregate))), Literal(0)).as("r")) + assert(withExprs(optimize(plan)).isEmpty, "a body holding an aggregate was memoized") + } + + test("a branch that already holds a With is left alone") { + // `RewriteWithExpression` defers a nested `With` to its next pass, and that rule has already + // run, so a definition put inside one would never be looked at again. + val existing = With(common) { case Seq(ref) => Add(ref, ref) } + val plan = relation + .select(CaseWhen( + Seq((GreaterThan(a, Literal(0)), Add(existing, Multiply(a, a)))), Literal(0)).as("r")) + .analyze + val memoized = withExprs(optimize(plan)) + assert(memoized.length == 1, s"expected the existing With and nothing more: $memoized") + assert(memoized.head.fastEquals(existing), s"the branch was rewritten: ${memoized.head}") + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala index c79fd0878ce25..c878c605f7c3d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala @@ -3414,4 +3414,31 @@ class ColumnExpressionSuite extends SharedSparkSession { } } } + + test("SPARK-59603: a subexpression repeated in one branch is evaluated once for a row") { + val counter = spark.sparkContext.longAccumulator + val counted = udf((x: Long) => { counter.add(1); x * 2 }) + def evaluations(memoize: Boolean): (Long, Boolean) = { + withSQLConf(SQLConf.MEMOIZE_COMMON_EXPRESSIONS_IN_BRANCHES.key -> memoize.toString) { + counter.reset() + // Five of the ten rows take the branch, and only those rows evaluate the UDF at all: the + // memoized definition stays inside the branch, where a `Project` column could not. + // `collect` rather than `checkAnswer`, which runs the plan more than once and would count + // the evaluations of each run. + val df = spark.range(0, 10, 1, 1) + .select(when($"id" < 5, counted($"id") + counted($"id")).otherwise(0L).as("r")) + val kept = df.queryExecution.optimizedPlan.expressions.exists(_.exists { + case _: With => true + case _ => false + }) + val rows = df.collect() + assert(rows.map(_.getLong(0)).toSeq == Seq(0L, 4L, 8L, 12L, 16L, 0L, 0L, 0L, 0L, 0L)) + (counter.value, kept) + } + } + val memoized = evaluations(memoize = true) + val inlined = evaluations(memoize = false) + assert(memoized == (5L, true) && inlined == (10L, false), + s"expected (5, true) and (10, false), got $memoized and $inlined") + } }