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") + } }