Push the cast into (if / case) branches.

Author: wangyum

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,
-        PushFoldableIntoBranches,
+        PushCastAndFoldableIntoBranches,
         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 foldable expression into (if / case) branches.
+ * Push the cast and foldable expression into (if / case) branches.
  */
-object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
+object PushCastAndFoldableIntoBranches 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,6 +542,18 @@ object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
 
   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)) =>
+        i.copy(
+          trueValue = Cast(trueValue, dataType, timeZoneId),
+          falseValue = Cast(falseValue, dataType, timeZoneId))
+
+      case Cast(c @ CaseWhen(branches, elseValue), dataType, timeZoneId)
+          if 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)))
+
       case b @ BinaryExpression(i @ If(_, trueValue, falseValue), right)
           if right.foldable && atMostOneUnfoldable(Seq(trueValue, falseValue)) =>
         i.copy(

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

@@ -27,15 +27,18 @@ import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLite
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
-import org.apache.spark.sql.types.{BooleanType, IntegerType}
+import org.apache.spark.sql.types.{BooleanType, IntegerType, StringType}
 
 
 class PushFoldableIntoBranchesSuite
   extends PlanTest with ExpressionEvalHelper with PredicateHelper {
 
   object Optimize extends RuleExecutor[LogicalPlan] {
-    val batches = Batch("PushFoldableIntoBranches", FixedPoint(50),
-      BooleanSimplification, ConstantFolding, SimplifyConditionals, PushFoldableIntoBranches) :: Nil
+    val batches = Batch("PushCastAndFoldableIntoBranches", FixedPoint(50),
+      BooleanSimplification,
+      ConstantFolding,
+      SimplifyConditionals,
+      PushCastAndFoldableIntoBranches) :: Nil
   }
 
   private val relation = LocalRelation('a.int, 'b.int, 'c.boolean)
@@ -222,4 +225,36 @@ class PushFoldableIntoBranchesSuite
     assertEquivalent(EqualTo(Literal(4), ifExp), FalseLiteral)
     assertEquivalent(EqualTo(Literal(4), caseWhen), FalseLiteral)
   }
+
+  test("Push down cast through If") {
+    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"))))
+    assertEquivalent(
+      CaseWhen(Seq((a, Literal(1))), Some(b)).cast(StringType),
+      CaseWhen(Seq((a, Literal("1"))), Some(b.cast(StringType))))
+    assertEquivalent(
+      CaseWhen(Seq((a, b)), Some(b + 1)).cast(StringType),
+      CaseWhen(Seq((a, b)), Some(b + 1)).cast(StringType))
+  }
+
+  test("Push down cast with binary expression through If/CaseWhen") {
+    assertEquivalent(EqualTo(If(a, Literal(2), Literal(3)).cast(StringType), Literal("4")),
+      FalseLiteral)
+    assertEquivalent(
+      EqualTo(CaseWhen(Seq((a, Literal(1))), Some(Literal(3))).cast(StringType), Literal("4")),
+      FalseLiteral)
+    assertEquivalent(
+      EqualTo(CaseWhen(Seq((a, Literal(1)), (c, Literal(2))), None).cast(StringType), Literal("4")),
+      CaseWhen(Seq((a, FalseLiteral), (c, FalseLiteral)), None))
+  }
 }