Track conditionally evaluated expressions to resolve as subexpressions for cases they are already being evaluated

Author: Kimahriman

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala

@@ -22,6 +22,7 @@ import scala.collection.mutable
 import org.apache.spark.TaskContext
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
 import org.apache.spark.sql.catalyst.expressions.objects.LambdaVariable
+import org.apache.spark.sql.internal.SQLConf
 
 /**
  * This class is used to compute equality of (sub)expression trees. Expressions can be added
@@ -43,6 +44,9 @@ class EquivalentExpressions {
 
   // For each expression, the set of equivalent expressions.
   private val equivalenceMap = mutable.HashMap.empty[Expr, mutable.ArrayBuffer[Expression]]
+  // Maintain optionally evaluated expressions that we can resolve as well
+  private val conditionalEquivalenceMap =
+    mutable.HashMap.empty[Expr, mutable.ArrayBuffer[Expression]]
 
   /**
    * Adds each expression to this data structure, grouping them with existing equivalent
@@ -65,6 +69,24 @@ class EquivalentExpressions {
     }
   }
 
+  private def addConditionalExpr(commonExprs: mutable.Set[Expr])(expr: Expression): Boolean = {
+    // If we've extract it as a common expression, we don't want to double count it as conditional
+    val inCommonExprs = commonExprs.contains(Expr(expr))
+    if (expr.deterministic && !inCommonExprs) {
+      val e: Expr = Expr(expr)
+      val f = conditionalEquivalenceMap.get(e)
+      if (f.isDefined) {
+        f.get += expr
+        true
+      } else {
+        conditionalEquivalenceMap.put(e, mutable.ArrayBuffer(expr))
+        equivalenceMap.contains(e)
+      }
+    } else {
+      inCommonExprs
+    }
+  }
+
   private def addExprToSet(expr: Expression, set: mutable.Set[Expr]): Boolean = {
     if (expr.deterministic) {
       val e = Expr(expr)
@@ -91,9 +113,9 @@ class EquivalentExpressions {
    * For example, if `((a + b) + c)` and `(a + b)` are common expressions, we only add
    * `((a + b) + c)`.
    */
-  private def addCommonExprs(
+  private def findCommonExprs(
       exprs: Seq[Expression],
-      addFunc: Expression => Boolean = addExpr): Unit = {
+      addFunc: Expression => Boolean = addExpr): mutable.Set[Expr] = {
     val exprSetForAll = mutable.Set[Expr]()
     addExprTree(exprs.head, addExprToSet(_, exprSetForAll))
 
@@ -105,13 +127,11 @@ class EquivalentExpressions {
 
     // Not all expressions in the set should be added. We should filter out the related
     // children nodes.
-    val commonExprSet = candidateExprs.filter { candidateExpr =>
+    candidateExprs.filter { candidateExpr =>
       candidateExprs.forall { expr =>
         expr == candidateExpr || expr.e.find(_.semanticEquals(candidateExpr.e)).isEmpty
       }
     }
-
-    commonExprSet.foreach(expr => addExprTree(expr.e, addFunc))
   }
 
   // There are some special expressions that we should not recurse into all of its children.
@@ -134,24 +154,18 @@ class EquivalentExpressions {
 
   // For some special expressions we cannot just recurse into all of its children, but we can
   // recursively add the common expressions shared between all of its children.
-  private def commonChildrenToRecurse(expr: Expression): Seq[Seq[Expression]] = expr match {
-    case i: If => Seq(Seq(i.trueValue, i.falseValue))
-    case c: CaseWhen =>
-      // We look at subexpressions in conditions and values of `CaseWhen` separately. It is
-      // because a subexpression in conditions will be run no matter which condition is matched
-      // if it is shared among conditions, but it doesn't need to be shared in values. Similarly,
-      // a subexpression among values doesn't need to be in conditions because no matter which
-      // condition is true, it will be evaluated.
-      val conditions = c.branches.tail.map(_._1)
-      // For an expression to be in all branch values of a CaseWhen statement, it must also be in
-      // the elseValue.
-      val values = if (c.elseValue.nonEmpty) {
-        c.branches.map(_._2) ++ c.elseValue
-      } else {
-        Nil
-      }
-      Seq(conditions, values)
-    case c: Coalesce => Seq(c.children.tail)
+  private def commonChildrenToRecurse(expr: Expression): Seq[Expression] = expr match {
+    case i: If => Seq(i.trueValue, i.falseValue)
+    case c: CaseWhen if c.elseValue.nonEmpty => c.branches.map(_._2) ++ c.elseValue
+    case _ => Nil
+  }
+
+  // Finds expressions that are conditionally evaluated, so that if they are definitely evaluated
+  // elsewhere, we can create a subexpression to optimize the conditional case.
+  private def conditionallyEvaluatedChildren(expr: Expression): Seq[Expression] = expr match {
+    case i: If => Seq(i.trueValue, i.falseValue)
+    case c: CaseWhen => c.branches.tail.map(_._1) ++ c.branches.map(_._2) ++ c.elseValue
+    case c: Coalesce => c.children.tail
     case _ => Nil
   }
 
@@ -172,7 +186,19 @@ class EquivalentExpressions {
 
     if (!skip && !addFunc(expr)) {
       childrenToRecurse(expr).foreach(addExprTree(_, addFunc))
-      commonChildrenToRecurse(expr).filter(_.nonEmpty).foreach(addCommonExprs(_, addFunc))
+
+      val commonChildrenCandidates = commonChildrenToRecurse(expr)
+      val commonChildrenExprs = if (commonChildrenCandidates.nonEmpty) {
+        findCommonExprs(commonChildrenCandidates)
+      } else {
+        mutable.Set.empty[Expr]
+      }
+      commonChildrenExprs.foreach(e => addExprTree(e.e, addFunc))
+
+      if (SQLConf.get.subexpressionEliminationConditionalsEnabled) {
+        val conditionallyEvaluatedExprs = conditionallyEvaluatedChildren(expr)
+        conditionallyEvaluatedExprs.foreach(addExprTree(_, addConditionalExpr(commonChildrenExprs)))
+      }
     }
   }
 
@@ -189,7 +215,9 @@ class EquivalentExpressions {
    * times.
    */
   def getAllEquivalentExprs(repeatTimes: Int = 0): Seq[Seq[Expression]] = {
-    equivalenceMap.values.map(_.toSeq).filter(_.size > repeatTimes).toSeq
+    equivalenceMap.values.map(_.toSeq)
+      .map(exprs => exprs ++ conditionalEquivalenceMap.getOrElse(Expr(exprs.head), Seq.empty))
+      .filter(_.size > repeatTimes).toSeq
       .sortBy(_.head)(new ExpressionContainmentOrdering)
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -626,6 +626,14 @@ object SQLConf {
       .checkValue(_ >= 0, "The maximum must not be negative")
       .createWithDefault(100)
 
+  val SUBEXPRESSION_ELIMINATION_CONDITIONALS_ENABLED =
+    buildConf("spark.sql.subexpressionElimination.conditionals.enabled")
+      .internal()
+      .doc("When true, common conditional subexpressions will be eliminated.")
+      .version("3.2.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val CASE_SENSITIVE = buildConf("spark.sql.caseSensitive")
     .internal()
     .doc("Whether the query analyzer should be case sensitive or not. " +
@@ -3599,6 +3607,9 @@ class SQLConf extends Serializable with Logging {
   def subexpressionEliminationCacheMaxEntries: Int =
     getConf(SUBEXPRESSION_ELIMINATION_CACHE_MAX_ENTRIES)
 
+  def subexpressionEliminationConditionalsEnabled: Boolean =
+    getConf(SUBEXPRESSION_ELIMINATION_CONDITIONALS_ENABLED)
+
   def autoBroadcastJoinThreshold: Long = getConf(AUTO_BROADCASTJOIN_THRESHOLD)
 
   def limitScaleUpFactor: Int = getConf(LIMIT_SCALE_UP_FACTOR)

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala

@@ -193,18 +193,6 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
   test("Children of conditional expressions: CaseWhen") {
     val add1 = Add(Literal(1), Literal(2))
     val add2 = Add(Literal(2), Literal(3))
-    val conditions1 = (GreaterThan(add2, Literal(3)), add1) ::
-      (GreaterThan(add2, Literal(4)), add1) ::
-      (GreaterThan(add2, Literal(5)), add1) :: Nil
-
-    val caseWhenExpr1 = CaseWhen(conditions1, None)
-    val equivalence1 = new EquivalentExpressions
-    equivalence1.addExprTree(caseWhenExpr1)
-
-    // `add2` is repeatedly in all conditions.
-    assert(equivalence1.getAllEquivalentExprs().count(_.size == 2) == 1)
-    assert(equivalence1.getAllEquivalentExprs().filter(_.size == 2).head == Seq(add2, add2))
-
     val conditions2 = (GreaterThan(add1, Literal(3)), add1) ::
       (GreaterThan(add2, Literal(4)), add1) ::
       (GreaterThan(add2, Literal(5)), add1) :: Nil
@@ -229,30 +217,31 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
   }
 
   test("Children of conditional expressions: Coalesce") {
-    val add1 = Add(Literal(1), Literal(2))
-    val add2 = Add(Literal(2), Literal(3))
-    val conditions1 = GreaterThan(add2, Literal(3)) ::
-      GreaterThan(add2, Literal(4)) ::
-      GreaterThan(add2, Literal(5)) :: Nil
-
-    val coalesceExpr1 = Coalesce(conditions1)
-    val equivalence1 = new EquivalentExpressions
-    equivalence1.addExprTree(coalesceExpr1)
-
-    // `add2` is repeatedly in all conditions.
-    assert(equivalence1.getAllEquivalentExprs().count(_.size == 2) == 1)
-    assert(equivalence1.getAllEquivalentExprs().filter(_.size == 2).head == Seq(add2, add2))
-
-    // Negative case. `add1` and `add2` both are not used in all branches.
-    val conditions2 = GreaterThan(add1, Literal(3)) ::
-      GreaterThan(add2, Literal(4)) ::
-      GreaterThan(add2, Literal(5)) :: Nil
-
-    val coalesceExpr2 = Coalesce(conditions2)
-    val equivalence2 = new EquivalentExpressions
-    equivalence2.addExprTree(coalesceExpr2)
-
-    assert(equivalence2.getAllEquivalentExprs().count(_.size == 2) == 0)
+    withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_CONDITIONALS_ENABLED.key -> "true") {
+      val add1 = Add(Literal(1), Literal(2))
+      val add2 = Add(Literal(2), Literal(3))
+      val conditions1 = GreaterThan(add2, Literal(3)) ::
+        GreaterThan(add2, Literal(4)) :: Nil
+
+      val coalesceExpr1 = Coalesce(conditions1)
+      val equivalence1 = new EquivalentExpressions
+      equivalence1.addExprTree(coalesceExpr1)
+
+      // `add2` is repeatedly in all conditions.
+      assert(equivalence1.getAllEquivalentExprs(1).size == 1)
+      assert(equivalence1.getAllEquivalentExprs(1).head == Seq(add2, add2))
+
+      // Negative case. `add1` and `add2` both are not used in all branches.
+      val conditions2 = GreaterThan(add1, Literal(3)) ::
+        GreaterThan(add2, Literal(4)) ::
+        GreaterThan(add2, Literal(5)) :: Nil
+
+      val coalesceExpr2 = Coalesce(conditions2)
+      val equivalence2 = new EquivalentExpressions
+      equivalence2.addExprTree(coalesceExpr2)
+
+      assert(equivalence2.getAllEquivalentExprs(1).size == 0)
+    }
   }
 
   test("SPARK-34723: Correct parameter type for subexpression elimination under whole-stage") {
@@ -359,9 +348,10 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
     + "elseValue") {
     val add1 = Add(Literal(1), Literal(2))
     val add2 = Add(Literal(2), Literal(3))
-    val conditions = (GreaterThan(add1, Literal(3)), add1) ::
-      (GreaterThan(add2, Literal(4)), add1) ::
-      (GreaterThan(add2, Literal(5)), add1) :: Nil
+    val add3 = Add(Literal(3), Literal(4))
+    val conditions = (GreaterThan(add2, Literal(3)), add1) ::
+      (GreaterThan(add3, Literal(4)), add1) ::
+      (GreaterThan(add3, Literal(5)), add1) :: Nil
 
     val caseWhenExpr = CaseWhen(conditions, None)
     val equivalence = new EquivalentExpressions
@@ -371,6 +361,44 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
     assert(equivalence.getAllEquivalentExprs().count(_.size == 2) == 0)
   }
 
+  test("SPARK-35564: Subexpressions should be extracted from conditional values if that value "
+    + "will always be evaluated elsewhere") {
+    withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_CONDITIONALS_ENABLED.key -> "true") {
+      val add1 = Add(Literal(1), Literal(2))
+      val add2 = Add(Literal(2), Literal(3))
+
+      val conditions1 = (GreaterThan(add1, Literal(3)), add1) :: Nil
+      val caseWhenExpr1 = CaseWhen(conditions1, None)
+      val equivalence1 = new EquivalentExpressions
+      equivalence1.addExprTree(caseWhenExpr1)
+
+      // `add1` is evaluated once in the first condition, and optionally in the first value
+      assert(equivalence1.getAllEquivalentExprs(1).size == 1)
+
+      val ifExpr = If(GreaterThan(add1, Literal(3)), add1, add2)
+      val equivalence2 = new EquivalentExpressions
+      equivalence2.addExprTree(ifExpr)
+
+      // `add1` is evaluated once in the condition, and optionally in the true value
+      assert(equivalence2.getAllEquivalentExprs(1).size == 1)
+    }
+  }
+
+  test("SPARK-35564: Don't double count conditional expressions if present in all branches") {
+    withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_CONDITIONALS_ENABLED.key -> "true") {
+      val add1 = Add(Literal(1), Literal(2))
+      val add2 = Add(Literal(2), Literal(3))
+      val add3 = Add(add2, Literal(4))
+
+      val caseWhenExpr1 = CaseWhen((GreaterThan(add1, Literal(3)), add3) :: Nil, add2)
+      val equivalence1 = new EquivalentExpressions
+      equivalence1.addExprTree(caseWhenExpr1)
+
+      // `add2` will only be evaluated once so don't create a subexpression
+      assert(equivalence1.getAllEquivalentExprs(1).size == 0)
+    }
+  }
+
   test("SPARK-35439: sort exprs with ExpressionContainmentOrdering") {
     val exprOrdering = new ExpressionContainmentOrdering