Improve SimplifyConditionals

Author: wangyum

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

@@ -560,15 +560,13 @@ object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
           if right.foldable && atMostOneUnfoldable(branches.map(_._2) ++ elseValue) =>
         c.copy(
           branches.map(e => e.copy(_2 = b.makeCopy(Array(e._2, right)))),
-          elseValue.orElse(Some(Literal.create(null, right.dataType)))
-            .map(e => b.makeCopy(Array(e, right))))
+          elseValue.map(e => b.makeCopy(Array(e, right))))
 
       case b @ BinaryExpression(left, c @ CaseWhen(branches, elseValue))
           if left.foldable && atMostOneUnfoldable(branches.map(_._2) ++ elseValue) =>
         c.copy(
           branches.map(e => e.copy(_2 = b.makeCopy(Array(left, e._2)))),
-          elseValue.orElse(Some(Literal.create(null, left.dataType)))
-            .map(e => b.makeCopy(Array(left, e))))
+          elseValue.map(e => b.makeCopy(Array(left, e))))
     }
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushFoldableIntoBranchesSuite.scala

@@ -124,17 +124,15 @@ class PushFoldableIntoBranchesSuite
       CaseWhen(Seq((a, FalseLiteral), (c, FalseLiteral)), Some(TrueLiteral)))
     assertEquivalent(
       EqualTo(CaseWhen(Seq((a, Literal(1)), (c, Literal(2))), None), Literal(4)),
-      CaseWhen(Seq((a, FalseLiteral), (c, FalseLiteral)), nullBoolean))
-    assertEquivalent(
-      EqualTo(CaseWhen(Seq((a, nullInt), (c, nullInt)), None), Literal(4)), nullBoolean)
+      CaseWhen(Seq((a, FalseLiteral), (c, FalseLiteral)), None))
 
     assertEquivalent(
       And(EqualTo(caseWhen, Literal(5)), EqualTo(caseWhen, Literal(6))),
       FalseLiteral)
 
     // Push down at most one branch is not foldable expressions.
     assertEquivalent(EqualTo(CaseWhen(Seq((a, b), (c, Literal(1))), None), Literal(1)),
-      CaseWhen(Seq((a, EqualTo(b, Literal(1))), (c, TrueLiteral)), nullBoolean))
+      CaseWhen(Seq((a, EqualTo(b, Literal(1))), (c, TrueLiteral)), None))
     assertEquivalent(EqualTo(CaseWhen(Seq((a, b), (c, b + 1)), None), Literal(1)),
       EqualTo(CaseWhen(Seq((a, b), (c, b + 1)), None), Literal(1)))
     assertEquivalent(EqualTo(CaseWhen(Seq((a, b)), None), Literal(1)),
@@ -225,7 +223,5 @@ class PushFoldableIntoBranchesSuite
     assertEquivalent(EqualTo(Literal(4), ifExp), FalseLiteral)
     assertEquivalent(EqualTo(Literal(4), If(a, nullInt, nullInt)), nullBoolean)
     assertEquivalent(EqualTo(Literal(4), caseWhen), FalseLiteral)
-    assertEquivalent(EqualTo(Literal(4), CaseWhen(Seq((a, nullInt), (c, nullInt)), None)),
-      nullBoolean)
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceNullWithFalseInPredicateSuite.scala

@@ -21,7 +21,7 @@ import org.apache.spark.sql.AnalysisException
 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.{And, ArrayExists, ArrayFilter, ArrayTransform, CaseWhen, Expression, GreaterThan, If, LambdaFunction, Literal, MapFilter, NamedExpression, Or, UnresolvedNamedLambdaVariable}
+import org.apache.spark.sql.catalyst.expressions.{And, ArrayExists, ArrayFilter, ArrayTransform, CaseWhen, Expression, GreaterThan, If, LambdaFunction, LessThanOrEqual, Literal, MapFilter, NamedExpression, Or, UnresolvedNamedLambdaVariable}
 import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral}
 import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest}
 import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, LocalRelation, LogicalPlan, UpdateTable}
@@ -236,12 +236,13 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
       Literal(2) === nestedCaseWhen,
       TrueLiteral,
       FalseLiteral)
-    val branches = Seq((UnresolvedAttribute("i") > Literal(10)) -> branchValue)
-    val condition = CaseWhen(branches)
-    testFilter(originalCond = condition, expectedCond = condition)
-    testJoin(originalCond = condition, expectedCond = condition)
-    testDelete(originalCond = condition, expectedCond = condition)
-    testUpdate(originalCond = condition, expectedCond = condition)
+    val condition = CaseWhen(Seq((UnresolvedAttribute("i") > Literal(10)) -> branchValue))
+    val expectedCond =
+      CaseWhen(Seq((UnresolvedAttribute("i") > Literal(10)) -> (Literal(2) === nestedCaseWhen)))
+    testFilter(originalCond = condition, expectedCond = expectedCond)
+    testJoin(originalCond = condition, expectedCond = expectedCond)
+    testDelete(originalCond = condition, expectedCond = expectedCond)
+    testUpdate(originalCond = condition, expectedCond = expectedCond)
   }
 
   test("inability to replace null in non-boolean branches of If inside another If") {
@@ -252,10 +253,14 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
         Literal(3)),
       TrueLiteral,
       FalseLiteral)
-    testFilter(originalCond = condition, expectedCond = condition)
-    testJoin(originalCond = condition, expectedCond = condition)
-    testDelete(originalCond = condition, expectedCond = condition)
-    testUpdate(originalCond = condition, expectedCond = condition)
+    val expectedCond = Literal(5) > If(
+      UnresolvedAttribute("i") === Literal(15),
+      Literal(null, IntegerType),
+      Literal(3))
+    testFilter(originalCond = condition, expectedCond = expectedCond)
+    testJoin(originalCond = condition, expectedCond = expectedCond)
+    testDelete(originalCond = condition, expectedCond = expectedCond)
+    testUpdate(originalCond = condition, expectedCond = expectedCond)
   }
 
   test("replace null in If used as a join condition") {
@@ -405,9 +410,9 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     val lambda1 = LambdaFunction(
       function = If(cond, Literal(null, BooleanType), TrueLiteral),
       arguments = lambdaArgs)
-    // the optimized lambda body is: if(arg > 0, false, true)
+    // the optimized lambda body is: if(arg > 0, false, true) => arg <= 0
     val lambda2 = LambdaFunction(
-      function = If(cond, FalseLiteral, TrueLiteral),
+      function = LessThanOrEqual(condArg, Literal(0)),
       arguments = lambdaArgs)
     testProjection(
       originalExpr = createExpr(argument, lambda1) as 'x,

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SimplifyConditionalSuite.scala

@@ -79,16 +79,6 @@ class SimplifyConditionalSuite extends PlanTest with ExpressionEvalHelper with P
         Literal(9)))
   }
 
-  test("remove unnecessary if when the outputs are boolean type") {
-    assertEquivalent(
-      If(IsNotNull(UnresolvedAttribute("a")), TrueLiteral, FalseLiteral),
-      IsNotNull(UnresolvedAttribute("a")))
-
-    assertEquivalent(
-      If(IsNotNull(UnresolvedAttribute("a")), FalseLiteral, TrueLiteral),
-      IsNull(UnresolvedAttribute("a")))
-  }
-
   test("remove unreachable branches") {
     // i.e. removing branches whose conditions are always false
     assertEquivalent(
@@ -209,4 +199,20 @@ class SimplifyConditionalSuite extends PlanTest with ExpressionEvalHelper with P
         If(Factorial(5) > 100L, b, nullLiteral).eval(EmptyRow))
     }
   }
+
+  test("remove unnecessary if when the outputs are boolean type") {
+    assertEquivalent(
+      If(IsNotNull(UnresolvedAttribute("a")), TrueLiteral, FalseLiteral),
+      IsNotNull(UnresolvedAttribute("a")))
+    assertEquivalent(
+      If(IsNotNull(UnresolvedAttribute("a")), FalseLiteral, TrueLiteral),
+      IsNull(UnresolvedAttribute("a")))
+
+    assertEquivalent(
+      If(GreaterThan(Rand(0), UnresolvedAttribute("a")), TrueLiteral, FalseLiteral),
+      If(GreaterThan(Rand(0), UnresolvedAttribute("a")), TrueLiteral, FalseLiteral))
+    assertEquivalent(
+      If(GreaterThan(Rand(0), UnresolvedAttribute("a")), FalseLiteral, TrueLiteral),
+      If(GreaterThan(Rand(0), UnresolvedAttribute("a")), FalseLiteral, TrueLiteral))
+  }
 }