Removes NullTypes.

Author: rxin

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala

@@ -36,6 +36,7 @@ object HiveTypeCoercion {
   val typeCoercionRules =
     PropagateTypes ::
       InConversion ::
+      RemoveNullTypes ::
       WidenTypes ::
       PromoteStrings ::
       DecimalPrecision ::
@@ -147,6 +148,47 @@ object HiveTypeCoercion {
     }
   }
 
+  /**
+   * Removes [[NullType]] (from null literals in SQL) from expressions by adding an explicit cast
+   * into the type the expression supports. This rule is here to avoid handling [[NullType]] in
+   * every expression implementations.
+   *
+   * This works by looking up the expected input types for expressions, and cast [[NullType]]
+   * into some other specific data type that the expression expects. For example, consider [[Add]],
+   * an expression that supports any numeric types.
+   *
+   * When applying this rule on `Add(NullType, NullType)`, the expression will be converted to
+   * `Add(DoubleType, DoubleType)`.
+   */
+  object RemoveNullTypes extends Rule[LogicalPlan] {
+    def apply(plan: LogicalPlan): LogicalPlan = plan transform {
+      case q: LogicalPlan => q transformExpressions {
+        // Skip nodes who's children have not been resolved yet.
+        case e if !e.childrenResolved => e
+
+        // For binary operators whose input types are NullType, cast them to some specific type.
+        case b@BinaryOperator(left, right)
+          if left.dataType == NullType && right.dataType == NullType &&
+            !b.inputType.acceptsType(NullType) =>
+          // If both inputs are null type (from null literals), cast the null type into some
+          // specific type the expression expects, so expressions don't need to handle NullType
+          val newLeft = Cast(left, b.inputType.defaultConcreteType)
+          val newRight = Cast(right, b.inputType.defaultConcreteType)
+          b.makeCopy(Array(newLeft, newRight))
+
+        case e: ExpectsInputTypes if e.inputTypes.nonEmpty =>
+          val children: Seq[Expression] = e.children.zip(e.inputTypes).map { case (in, expected) =>
+            if (in.dataType == NullType && !expected.acceptsType(NullType)) {
+              Cast(in, expected.defaultConcreteType)
+            } else {
+              in
+            }
+          }
+          e.withNewChildren(children)
+      }
+    }
+  }
+
   /**
    * Widens numeric types and converts strings to numbers when appropriate.
    *
@@ -220,14 +262,6 @@ object HiveTypeCoercion {
         // Skip nodes who's children have not been resolved yet.
         case e if !e.childrenResolved => e
 
-        case b @ BinaryOperator(left, right)
-          if left.dataType == NullType &&  right.dataType == NullType =>
-          // If both inputs are null type (from null literals), cast the null type into some
-          // specific type the expression expects, so expressions don't need to handle NullType
-          val newLeft = Cast(left, b.inputType.defaultConcreteType)
-          val newRight = Cast(right, b.inputType.defaultConcreteType)
-          b.makeCopy(Array(newLeft, newRight))
-
         case b @ BinaryOperator(left, right) if left.dataType != right.dataType =>
           findTightestCommonTypeOfTwo(left.dataType, right.dataType).map { commonType =>
             if (b.inputType.acceptsType(commonType)) {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercionSuite.scala

@@ -194,10 +194,28 @@ class HiveTypeCoercionSuite extends PlanTest {
       Project(Seq(Alias(transformed, "a")()), testRelation))
   }
 
-  test("null literals handling for binary operators") {
-    ruleTest(HiveTypeCoercion.WidenTypes,
-      Add(Literal.create(null, NullType), Literal.create(null, NullType)),
-      Add(
+  test("RemoveNullTypes for expressions that define ExpectsInputTypes") {
+    import HiveTypeCoercionSuite._
+
+    ruleTest(HiveTypeCoercion.RemoveNullTypes,
+      AnyTypeUnaryExpression(Literal.create(null, NullType)),
+      AnyTypeUnaryExpression(Literal.create(null, NullType)))
+
+    ruleTest(HiveTypeCoercion.RemoveNullTypes,
+      NumericTypeUnaryExpression(Literal.create(null, NullType)),
+      NumericTypeUnaryExpression(Cast(Literal.create(null, NullType), DoubleType)))
+  }
+
+  test("RemoveNullTypes for binary operators") {
+    import HiveTypeCoercionSuite._
+
+    ruleTest(HiveTypeCoercion.RemoveNullTypes,
+      AnyTypeBinaryOperator(Literal.create(null, NullType), Literal.create(null, NullType)),
+      AnyTypeBinaryOperator(Literal.create(null, NullType), Literal.create(null, NullType)))
+
+    ruleTest(HiveTypeCoercion.RemoveNullTypes,
+      NumericTypeBinaryOperator(Literal.create(null, NullType), Literal.create(null, NullType)),
+      NumericTypeBinaryOperator(
         Cast(Literal.create(null, NullType), DoubleType),
         Cast(Literal.create(null, NullType), DoubleType)))
   }
@@ -310,3 +328,33 @@ class HiveTypeCoercionSuite extends PlanTest {
     )
   }
 }
+
+
+object HiveTypeCoercionSuite {
+
+  case class AnyTypeUnaryExpression(child: Expression)
+    extends UnaryExpression with ExpectsInputTypes {
+    override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+    override def dataType: DataType = NullType
+  }
+
+  case class NumericTypeUnaryExpression(child: Expression)
+    extends UnaryExpression with ExpectsInputTypes {
+    override def inputTypes: Seq[AbstractDataType] = Seq(NumericType)
+    override def dataType: DataType = NullType
+  }
+
+  case class AnyTypeBinaryOperator(left: Expression, right: Expression)
+    extends BinaryOperator with ExpectsInputTypes {
+    override def dataType: DataType = NullType
+    override def inputType: AbstractDataType = AnyDataType
+    override def symbol: String = "anytype"
+  }
+
+  case class NumericTypeBinaryOperator(left: Expression, right: Expression)
+    extends BinaryOperator with ExpectsInputTypes {
+    override def dataType: DataType = NullType
+    override def inputType: AbstractDataType = NumericType
+    override def symbol: String = "numerictype"
+  }
+}