address feedback

dbtsai · dbtsai · commit bf0b2d91a89e · 2018-07-23T19:04:40.000-07:00
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala
@@ -64,6 +64,8 @@ case class PrintToStderr(child: Expression) extends UnaryExpression {
   """)
 case class AssertTrue(child: Expression) extends UnaryExpression with ImplicitCastInputTypes {
 
+  override lazy val deterministic: Boolean = false
+
   override def nullable: Boolean = true
 
   override def inputTypes: Seq[DataType] = Seq(BooleanType)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala
@@ -1627,6 +1627,8 @@ case class InitializeJavaBean(beanInstance: Expression, setters: Map[String, Exp
 case class AssertNotNull(child: Expression, walkedTypePath: Seq[String] = Nil)
   extends UnaryExpression with NonSQLExpression {
 
+  override lazy val deterministic: Boolean = false
+
   override def dataType: DataType = child.dataType
   override def foldable: Boolean = false
   override def nullable: Boolean = false
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
@@ -390,7 +390,8 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
       case If(TrueLiteral, trueValue, _) => trueValue
       case If(FalseLiteral, _, falseValue) => falseValue
       case If(Literal(null, _), _, falseValue) => falseValue
-      case If(_, trueValue, falseValue) if trueValue.semanticEquals(falseValue) => trueValue
+      case If(cond, trueValue, falseValue)
+        if cond.deterministic && trueValue.semanticEquals(falseValue) => trueValue
 
       case e @ CaseWhen(branches, elseValue) if branches.exists(x => falseOrNullLiteral(x._1)) =>
         // If there are branches that are always false, remove them.
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
@@ -18,6 +18,7 @@
 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}
@@ -30,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 = {
@@ -44,6 +46,8 @@ 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)
+
   test("simplify if") {
     assertEquivalent(
       If(TrueLiteral, Literal(10), Literal(20)),
@@ -64,6 +68,34 @@ class SimplifyConditionalSuite extends PlanTest with PredicateHelper {
         Subtract(Literal(10), Literal(1)),
         Add(Literal(6), Literal(3))),
       Literal(9))
+
+    // For non-deterministic condition, we don't remove the `If` statement.
+    assertEquivalent(
+      If(GreaterThan(Rand(0), Literal(0.5)),
+        Subtract(Literal(10), Literal(1)),
+        Add(Literal(6), Literal(3))),
+      If(GreaterThan(Rand(0), Literal(0.5)),
+        Literal(9),
+        Literal(9)))
+
+    // For non-deterministic condition, we don't remove the `If` statement.
+    val originalQuery =
+      testRelation
+        .select(If(AssertTrue(IsNull(UnresolvedAttribute("a"))),
+          Subtract(Literal(10), Literal(1)),
+          Add(Literal(6), Literal(3)))).analyze
+
+    val optimized = Optimize.execute(originalQuery.analyze).canonicalized
+
+    val correctAnswer =
+      testRelation
+        .select(If(AssertTrue(IsNull(UnresolvedAttribute("a"))),
+          Literal(9),
+          Literal(9)))
+        .analyze
+        .canonicalized
+
+    comparePlans(optimized, correctAnswer)
   }
 
   test("remove unreachable branches") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
@@ -393,7 +393,7 @@ private[sql] trait SQLTestUtilsBase
   }
 
   /**
-   * Returns full path to the given file in the resouce folder
+   * Returns full path to the given file in the resource folder
    */
   protected def testFile(fileName: String): String = {
     Thread.currentThread().getContextClassLoader.getResource(fileName).toString

Original file line number	Diff line number	Diff line change
`@@ -393,7 +393,7 @@ private[sql] trait SQLTestUtilsBase`
`393`	`393`	`}`
`394`	`394`
`395`	`395`	`/**`
`396`		`- * Returns full path to the given file in the resouce folder`
	`396`	`+ * Returns full path to the given file in the resource folder`
`397`	`397`	`*/`
`398`	`398`	`protected def testFile(fileName: String): String = {`
`399`	`399`	`Thread.currentThread().getContextClassLoader.getResource(fileName).toString`