Fix bugs and merge into SimplifyConditionals

sunchao · sunchao · commit e4d929cae26b · 2020-08-28T10:51:08.000-07:00
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
@@ -103,7 +103,6 @@ abstract class Optimizer(catalogManager: CatalogManager)
         ReplaceNullWithFalseInPredicate,
         PruneFilters,
         SimplifyCasts,
-        SimplifyIf,
         SimplifyCaseConversionExpressions,
         RewriteCorrelatedScalarSubquery,
         EliminateSerialization,
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
@@ -19,7 +19,6 @@ package org.apache.spark.sql.catalyst.optimizer
 
 import scala.collection.immutable.HashSet
 import scala.collection.mutable.{ArrayBuffer, Stack}
-
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral}
@@ -463,6 +462,8 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
       case If(Literal(null, _), _, falseValue) => falseValue
       case If(cond, trueValue, falseValue)
         if cond.deterministic && trueValue.semanticEquals(falseValue) => trueValue
+      case If(p, l @ Literal(null, _), FalseLiteral) if !p.nullable => And(p, l)
+      case If(p, l @ Literal(null, _), TrueLiteral) if !p.nullable => Or(Not(p), l)
 
       case e @ CaseWhen(branches, elseValue) if branches.exists(x => falseOrNullLiteral(x._1)) =>
         // If there are branches that are always false, remove them.
@@ -716,18 +717,6 @@ object SimplifyCasts extends Rule[LogicalPlan] {
   }
 }
 
-/**
- * Simplify if clauses of pattern `if(p, null, true|false)`, by replacing them with
- * AND or OR clauses which are simpler and can better be pushed down
- */
-object SimplifyIf extends Rule[LogicalPlan] {
-  val nullLiteral = Literal(null, BooleanType)
-  def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
-    case If(p, Literal(null, _), FalseLiteral) => And(p, nullLiteral)
-    case If(p, Literal(null, _), TrueLiteral) => Or(p, nullLiteral)
-  }
-}
-
 /**
  * Removes nodes that are not necessary.
  */
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
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.types.{BooleanType, IntegerType}
 
 
-class SimplifyConditionalSuite extends PlanTest with PredicateHelper {
+class SimplifyConditionalSuite extends PlanTest with ExpressionEvalHelper with PredicateHelper {
 
   object Optimize extends RuleExecutor[LogicalPlan] {
     val batches = Batch("SimplifyConditionals", FixedPoint(50),
@@ -165,4 +165,23 @@ class SimplifyConditionalSuite extends PlanTest with PredicateHelper {
         Literal(1))
     )
   }
+
+  test("simplify if when then clause is null and else clause is boolean") {
+    val p = IsNull('a)
+    val nullLiteral = Literal(null, BooleanType)
+    assertEquivalent(If(p, nullLiteral, FalseLiteral), And(p, nullLiteral))
+    assertEquivalent(If(p, nullLiteral, TrueLiteral), Or(IsNotNull('a), nullLiteral))
+
+    // the rule should not apply to nullable predicate
+    Seq(TrueLiteral, FalseLiteral).foreach { b =>
+      assertEquivalent(If(GreaterThan('a, 42), nullLiteral, b),
+        If(GreaterThan('a, 42), nullLiteral, b))
+    }
+
+    // check evaluation also
+    Seq(TrueLiteral, FalseLiteral).foreach { p =>
+        checkEvaluation(If(p, nullLiteral, FalseLiteral), And(p, nullLiteral).eval(EmptyRow))
+        checkEvaluation(If(p, nullLiteral, TrueLiteral), Or(Not(p), nullLiteral).eval(EmptyRow))
+    }
+ }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SimplifyIfSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SimplifyIfSuite.scala
