remove casewhen if possible

dbtsai · dbtsai · commit dc8de5fec6ef · 2018-07-23T12:15:58.000-07:00
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala
@@ -414,6 +414,15 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
         // these branches can be pruned away
         val (h, t) = branches.span(_._1 != TrueLiteral)
         CaseWhen( h :+ t.head, None)
+
+      case CaseWhen(branches, Some(elseValue)) if {
+        // With previous rules, it's guaranteed that `branches.length >= 2`
+        val list = branches.map(_._2) :+ elseValue
+        list.tail.forall(list.head.semanticEquals)
+      } =>
+        // If all the values in the branches and elseValue are the same,
+        // `CaseWhen` condition can be removed.
+        elseValue
     }
   }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SimplifyConditionalSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SimplifyConditionalSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
+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.expressions.Literal.{FalseLiteral, TrueLiteral}
@@ -29,7 +31,8 @@ import org.apache.spark.sql.types.{IntegerType, NullType}
 class SimplifyConditionalSuite extends PlanTest with PredicateHelper {
 
   object Optimize extends RuleExecutor[LogicalPlan] {
-    val batches = Batch("SimplifyConditionals", FixedPoint(50), SimplifyConditionals) :: Nil
+    val batches = Batch("SimplifyConditionals", FixedPoint(50),
+      BooleanSimplification, ConstantFolding, SimplifyConditionals) :: Nil
   }
 
   protected def assertEquivalent(e1: Expression, e2: Expression): Unit = {
@@ -43,6 +46,12 @@ class SimplifyConditionalSuite extends PlanTest with PredicateHelper {
   private val unreachableBranch = (FalseLiteral, Literal(20))
   private val nullBranch = (Literal.create(null, NullType), Literal(30))
 
+  private val testRelation = LocalRelation('a.int, 'b.string, 'c.boolean)
+
+  val isNotNullCond = IsNotNull(UnresolvedAttribute(Seq("a")))
+  val isNullCond = IsNull(UnresolvedAttribute("b"))
+  val notCond = Not(UnresolvedAttribute("c"))
+
   test("simplify if") {
     assertEquivalent(
       If(TrueLiteral, Literal(10), Literal(20)),
@@ -100,4 +109,25 @@ class SimplifyConditionalSuite extends PlanTest with PredicateHelper {
         None),
       CaseWhen(normalBranch :: trueBranch :: Nil, None))
   }
+
+  test("remove entire CaseWhen if all the outputs are semantic equivalence") {
+    val originalQuery =
+      testRelation
+        .select(
+          CaseWhen((isNotNullCond, Subtract(Literal(3), Literal(2))) ::
+            (isNullCond, Literal(1)) ::
+            (notCond, Add(Literal(6), Literal(-5))) ::
+            Nil,
+            Add(Literal(2), Literal(-1))))
+        .analyze
+
+    val optimized = Optimize.execute(originalQuery.analyze).canonicalized
+    val correctAnswer =
+      testRelation
+        .select(Literal(1))
+        .analyze
+        .canonicalized
+
+    comparePlans(optimized, correctAnswer)
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -414,6 +414,15 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {`
`414`	`414`	`// these branches can be pruned away`
`415`	`415`	`val (h, t) = branches.span(_._1 != TrueLiteral)`
`416`	`416`	`CaseWhen( h :+ t.head, None)`
	`417`	`+`
	`418`	`+ case CaseWhen(branches, Some(elseValue)) if {`
	`419`	+ // With previous rules, it's guaranteed that `branches.length >= 2`
	`420`	`+ val list = branches.map(_._2) :+ elseValue`
	`421`	`+ list.tail.forall(list.head.semanticEquals)`
	`422`	`+ } =>`
	`423`	`+ // If all the values in the branches and elseValue are the same,`
	`424`	+ // `CaseWhen` condition can be removed.
	`425`	`+ elseValue`
`417`	`426`	`}`
`418`	`427`	`}`
`419`	`428`	`}`