fix

Author: wangyum

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

@@ -536,6 +536,12 @@ abstract class UnaryExpression extends Expression {
   }
 }
 
+
+object UnaryExpression {
+  def unapply(e: UnaryExpression): Option[Expression] = Some(e.child)
+}
+
+
 /**
  * An expression with two inputs and one output. The output is by default evaluated to null
  * if any input is evaluated to null.

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

@@ -99,7 +99,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
         LikeSimplification,
         BooleanSimplification,
         SimplifyConditionals,
-        PushCastAndFoldableIntoBranches,
+        PushFoldableIntoBranches,
         RemoveDispensableExpressions,
         SimplifyBinaryComparison,
         ReplaceNullWithFalseInPredicate,

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

@@ -529,9 +529,9 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
 
 
 /**
- * Push the cast and foldable expression into (if / case) branches.
+ * Push the foldable expression into (if / case) branches.
  */
-object PushCastAndFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
+object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
 
   // To be conservative here: it's only a guaranteed win if all but at most only one branch
   // end up being not foldable.
@@ -542,41 +542,41 @@ object PushCastAndFoldableIntoBranches extends Rule[LogicalPlan] with PredicateH
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case q: LogicalPlan => q transformExpressionsUp {
-      case Cast(i @ If(_, trueValue, falseValue), dataType, timeZoneId)
-          if atMostOneUnfoldable(Seq(trueValue, falseValue)) =>
+      case u @ UnaryExpression(i @ If(_, trueValue, falseValue))
+          if !u.isInstanceOf[Alias] && atMostOneUnfoldable(Seq(trueValue, falseValue)) =>
         i.copy(
-          trueValue = Cast(trueValue, dataType, timeZoneId),
-          falseValue = Cast(falseValue, dataType, timeZoneId))
+          trueValue = u.withNewChildren(Array(trueValue)),
+          falseValue = u.withNewChildren(Array(falseValue)))
 
-      case Cast(c @ CaseWhen(branches, elseValue), dataType, timeZoneId)
-          if atMostOneUnfoldable(branches.map(_._2) ++ elseValue) =>
+      case u @ UnaryExpression(c @ CaseWhen(branches, elseValue))
+          if !u.isInstanceOf[Alias] && atMostOneUnfoldable(branches.map(_._2) ++ elseValue) =>
         c.copy(
-          branches.map(e => e.copy(_2 = Cast(e._2, dataType, timeZoneId))),
-          elseValue.map(e => Cast(e, dataType, timeZoneId)))
+          branches.map(e => e.copy(_2 = u.withNewChildren(Array(e._2)))),
+          elseValue.map(e => u.withNewChildren(Array(e))))
 
       case b @ BinaryExpression(i @ If(_, trueValue, falseValue), right)
           if right.foldable && atMostOneUnfoldable(Seq(trueValue, falseValue)) =>
         i.copy(
-          trueValue = b.makeCopy(Array(trueValue, right)),
-          falseValue = b.makeCopy(Array(falseValue, right)))
+          trueValue = b.withNewChildren(Array(trueValue, right)),
+          falseValue = b.withNewChildren(Array(falseValue, right)))
 
       case b @ BinaryExpression(left, i @ If(_, trueValue, falseValue))
           if left.foldable && atMostOneUnfoldable(Seq(trueValue, falseValue)) =>
         i.copy(
-          trueValue = b.makeCopy(Array(left, trueValue)),
-          falseValue = b.makeCopy(Array(left, falseValue)))
+          trueValue = b.withNewChildren(Array(left, trueValue)),
+          falseValue = b.withNewChildren(Array(left, falseValue)))
 
       case b @ BinaryExpression(c @ CaseWhen(branches, elseValue), right)
           if right.foldable && atMostOneUnfoldable(branches.map(_._2) ++ elseValue) =>
         c.copy(
-          branches.map(e => e.copy(_2 = b.makeCopy(Array(e._2, right)))),
-          elseValue.map(e => b.makeCopy(Array(e, right))))
+          branches.map(e => e.copy(_2 = b.withNewChildren(Array(e._2, right)))),
+          elseValue.map(e => b.withNewChildren(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.map(e => b.makeCopy(Array(left, e))))
+          branches.map(e => e.copy(_2 = b.withNewChildren(Array(left, e._2)))),
+          elseValue.map(e => b.withNewChildren(Array(left, e))))
     }
   }
 }

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

@@ -30,15 +30,12 @@ import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.types.{BooleanType, IntegerType, StringType}
 
 
-class PushCastAndFoldableIntoBranchesSuite
+class PushFoldableIntoBranchesSuite
   extends PlanTest with ExpressionEvalHelper with PredicateHelper {
 
   object Optimize extends RuleExecutor[LogicalPlan] {
-    val batches = Batch("PushCastAndFoldableIntoBranches", FixedPoint(50),
-      BooleanSimplification,
-      ConstantFolding,
-      SimplifyConditionals,
-      PushCastAndFoldableIntoBranches) :: Nil
+    val batches = Batch("PushFoldableIntoBranchesSuite", FixedPoint(50),
+      BooleanSimplification, ConstantFolding, SimplifyConditionals, PushFoldableIntoBranches) :: Nil
   }
 
   private val relation = LocalRelation('a.int, 'b.int, 'c.boolean)
@@ -226,16 +223,14 @@ class PushCastAndFoldableIntoBranchesSuite
     assertEquivalent(EqualTo(Literal(4), caseWhen), FalseLiteral)
   }
 
-  test("Push down cast through If") {
+  test("Push down cast through If/CaseWhen") {
     assertEquivalent(If(a, Literal(2), Literal(3)).cast(StringType),
       If(a, Literal("2"), Literal("3")))
     assertEquivalent(If(a, b, Literal(3)).cast(StringType),
       If(a, b.cast(StringType), Literal("3")))
     assertEquivalent(If(a, b, b + 1).cast(StringType),
       If(a, b, b + 1).cast(StringType))
-  }
 
-  test("Push down cast through CaseWhen") {
     assertEquivalent(
       CaseWhen(Seq((a, Literal(1))), Some(Literal(3))).cast(StringType),
       CaseWhen(Seq((a, Literal("1"))), Some(Literal("3"))))
@@ -247,6 +242,13 @@ class PushCastAndFoldableIntoBranchesSuite
       CaseWhen(Seq((a, b)), Some(b + 1)).cast(StringType))
   }
 
+  test("Push down abs through If/CaseWhen") {
+    assertEquivalent(Abs(If(a, Literal(-2), Literal(-3))), If(a, Literal(2), Literal(3)))
+    assertEquivalent(
+      Abs(CaseWhen(Seq((a, Literal(-1))), Some(Literal(-3)))),
+      CaseWhen(Seq((a, Literal(1))), Some(Literal(3))))
+  }
+
   test("Push down cast with binary expression through If/CaseWhen") {
     assertEquivalent(EqualTo(If(a, Literal(2), Literal(3)).cast(StringType), Literal("4")),
       FalseLiteral)