[SPARK-12054] [SQL] Consider nullability of expression in codegen

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

@@ -69,10 +69,17 @@ case class BoundReference(ordinal: Int, dataType: DataType, nullable: Boolean)
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val javaType = ctx.javaType(dataType)
     val value = ctx.getValue(ctx.INPUT_ROW, dataType, ordinal.toString)
-    s"""
-      boolean ${ev.isNull} = ${ctx.INPUT_ROW}.isNullAt($ordinal);
-      $javaType ${ev.value} = ${ev.isNull} ? ${ctx.defaultValue(dataType)} : ($value);
-    """
+    if (nullable) {
+      s"""
+        boolean ${ev.isNull} = ${ctx.INPUT_ROW}.isNullAt($ordinal);
+        $javaType ${ev.value} = ${ev.isNull} ? ${ctx.defaultValue(dataType)} : ($value);
+      """
+    } else {
+      ev.isNull = "false"
+      s"""
+        $javaType ${ev.value} = $value;
+      """
+    }
   }
 }
 
@@ -92,7 +99,7 @@ object BindReferences extends Logging {
             sys.error(s"Couldn't find $a in ${input.mkString("[", ",", "]")}")
           }
         } else {
-          BoundReference(ordinal, a.dataType, a.nullable)
+          BoundReference(ordinal, a.dataType, input(ordinal).nullable)
         }
       }
     }.asInstanceOf[A] // Kind of a hack, but safe.  TODO: Tighten return type when possible.

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

@@ -87,18 +87,22 @@ object Cast {
   private def resolvableNullability(from: Boolean, to: Boolean) = !from || to
 
   private def forceNullable(from: DataType, to: DataType) = (from, to) match {
-    case (StringType, _: NumericType) => true
-    case (StringType, TimestampType) => true
-    case (DoubleType, TimestampType) => true
-    case (FloatType, TimestampType) => true
-    case (StringType, DateType) => true
-    case (_: NumericType, DateType) => true
-    case (BooleanType, DateType) => true
-    case (DateType, _: NumericType) => true
-    case (DateType, BooleanType) => true
-    case (DoubleType, _: DecimalType) => true
-    case (FloatType, _: DecimalType) => true
-    case (_, DecimalType.Fixed(_, _)) => true // TODO: not all upcasts here can really give null
+    case (NullType, _) => true
+    case (_, _) if from == to => false
+
+    case (StringType, BinaryType) => false
+    case (StringType, _) => true
+    case (_, StringType) => false
+
+    case (FloatType | DoubleType, TimestampType) => true
+    case (TimestampType, DateType) => false
+    case (_, DateType) => true
+    case (DateType, TimestampType) => false
+    case (DateType, _) => true
+    case (_, CalendarIntervalType) => true
+
+    case (_, _: DecimalType) => true  // overflow
+    case (_: FractionalType, _: IntegralType) => true  // NaN, infinity
     case _ => false
   }
 }

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

@@ -340,14 +340,21 @@ abstract class UnaryExpression extends Expression {
       ev: GeneratedExpressionCode,
       f: String => String): String = {
     val eval = child.gen(ctx)
-    val resultCode = f(eval.value)
-    eval.code + s"""
-      boolean ${ev.isNull} = ${eval.isNull};
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-      if (!${ev.isNull}) {
-        $resultCode
-      }
-    """
+    if (nullable) {
+      eval.code + s"""
+        boolean ${ev.isNull} = ${eval.isNull};
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        if (!${eval.isNull}) {
+          ${f(eval.value)}
+        }
+      """
+    } else {
+      ev.isNull = "false"
+      eval.code + s"""
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        ${f(eval.value)}
+      """
+    }
   }
 }
 
@@ -424,19 +431,30 @@ abstract class BinaryExpression extends Expression {
     val eval1 = left.gen(ctx)
     val eval2 = right.gen(ctx)
     val resultCode = f(eval1.value, eval2.value)
-    s"""
-      ${eval1.code}
-      boolean ${ev.isNull} = ${eval1.isNull};
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-      if (!${ev.isNull}) {
-        ${eval2.code}
-        if (!${eval2.isNull}) {
-          $resultCode
-        } else {
-          ${ev.isNull} = true;
+    if (nullable) {
+      s"""
+        ${eval1.code}
+        boolean ${ev.isNull} = ${eval1.isNull};
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        if (!${ev.isNull}) {
+          ${eval2.code}
+          if (!${eval2.isNull}) {
+            $resultCode
+          } else {
+            ${ev.isNull} = true;
+          }
         }
-      }
-    """
+      """
+
+    } else {
+      ev.isNull = "false"
+      s"""
+        ${eval1.code}
+        ${eval2.code}
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        $resultCode
+      """
+    }
   }
 }
 
@@ -548,20 +566,31 @@ abstract class TernaryExpression extends Expression {
     f: (String, String, String) => String): String = {
     val evals = children.map(_.gen(ctx))
     val resultCode = f(evals(0).value, evals(1).value, evals(2).value)
-    s"""
-      ${evals(0).code}
-      boolean ${ev.isNull} = true;
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-      if (!${evals(0).isNull}) {
-        ${evals(1).code}
-        if (!${evals(1).isNull}) {
-          ${evals(2).code}
-          if (!${evals(2).isNull}) {
-            ${ev.isNull} = false;  // resultCode could change nullability
-            $resultCode
+    if (nullable) {
+      s"""
+        ${evals(0).code}
+        boolean ${ev.isNull} = true;
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        if (!${evals(0).isNull}) {
+          ${evals(1).code}
+          if (!${evals(1).isNull}) {
+            ${evals(2).code}
+            if (!${evals(2).isNull}) {
+              ${ev.isNull} = false;  // resultCode could change nullability
+              $resultCode
+            }
           }
         }
-      }
-    """
+      """
+    } else {
+      ev.isNull = "false"
+      s"""
+        ${evals(0).code}
+        ${evals(1).code}
+        ${evals(2).code}
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        $resultCode
+      """
+    }
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala

@@ -53,7 +53,7 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
 
   override def children: Seq[Expression] = Seq(child)
 
-  override def nullable: Boolean = false
+  override def nullable: Boolean = true
 
   override def dataType: DataType = DoubleType
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala

@@ -20,7 +20,6 @@ package org.apache.spark.sql.catalyst.expressions.aggregate
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 /**
@@ -42,7 +41,7 @@ case class Corr(
 
   override def children: Seq[Expression] = Seq(left, right)
 
-  override def nullable: Boolean = false
+  override def nullable: Boolean = true
 
   override def dataType: DataType = DoubleType
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala

@@ -31,23 +31,32 @@ case class Count(children: Seq[Expression]) extends DeclarativeAggregate {
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.size)(AnyDataType)
 
-  private lazy val count = AttributeReference("count", LongType)()
+  private lazy val count = AttributeReference("count", LongType, nullable = false)()
 
   override lazy val aggBufferAttributes = count :: Nil
 
   override lazy val initialValues = Seq(
     /* count = */ Literal(0L)
   )
 
-  override lazy val updateExpressions = Seq(
-    /* count = */ If(children.map(IsNull).reduce(Or), count, count + 1L)
-  )
+  override lazy val updateExpressions = {
+    val nullableChildren = children.filter(_.nullable)
+    if (nullableChildren.isEmpty) {
+      Seq(
+        /* count = */ count + 1L
+      )
+    } else {
+      Seq(
+        /* count = */ If(nullableChildren.map(IsNull).reduce(Or), count, count + 1L)
+      )
+    }
+  }
 
   override lazy val mergeExpressions = Seq(
     /* count = */ count.left + count.right
   )
 
-  override lazy val evaluateExpression = Cast(count, LongType)
+  override lazy val evaluateExpression = count
 
   override def defaultResult: Option[Literal] = Option(Literal(0L))
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala

@@ -40,8 +40,6 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
   private lazy val resultType = child.dataType match {
     case DecimalType.Fixed(precision, scale) =>
       DecimalType.bounded(precision + 10, scale)
-    // TODO: Remove this line once we remove the NullType from inputTypes.
-    case NullType => IntegerType
     case _ => child.dataType
   }
 
@@ -57,18 +55,26 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
     /* sum = */ Literal.create(null, sumDataType)
   )
 
-  override lazy val updateExpressions: Seq[Expression] = Seq(
-    /* sum = */
-    Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType)), sum))
-  )
+  override lazy val updateExpressions: Seq[Expression] = {
+    if (child.nullable) {
+      Seq(
+        /* sum = */
+        Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType)), sum))
+      )
+    } else {
+      Seq(
+        /* sum = */
+        Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType))
+      )
+    }
+  }
 
   override lazy val mergeExpressions: Seq[Expression] = {
-    val add = Add(Coalesce(Seq(sum.left, zero)), Cast(sum.right, sumDataType))
     Seq(
       /* sum = */
-      Coalesce(Seq(add, sum.left))
+      Coalesce(Seq(Add(Coalesce(Seq(sum.left, zero)), sum.right), sum.left))
     )
   }
 
-  override lazy val evaluateExpression: Expression = Cast(sum, resultType)
+  override lazy val evaluateExpression: Expression = sum
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala

@@ -32,14 +32,23 @@ trait CodegenFallback extends Expression {
 
     ctx.references += this
     val objectTerm = ctx.freshName("obj")
-    s"""
-      /* expression: ${this.toCommentSafeString} */
-      java.lang.Object $objectTerm = expressions[${ctx.references.size - 1}].eval(${ctx.INPUT_ROW});
-      boolean ${ev.isNull} = $objectTerm == null;
-      ${ctx.javaType(this.dataType)} ${ev.value} = ${ctx.defaultValue(this.dataType)};
-      if (!${ev.isNull}) {
-        ${ev.value} = (${ctx.boxedType(this.dataType)}) $objectTerm;
-      }
-    """
+    if (nullable) {
+      s"""
+        /* expression: ${this.toCommentSafeString} */
+        Object $objectTerm = expressions[${ctx.references.size - 1}].eval(${ctx.INPUT_ROW});
+        boolean ${ev.isNull} = $objectTerm == null;
+        ${ctx.javaType(this.dataType)} ${ev.value} = ${ctx.defaultValue(this.dataType)};
+        if (!${ev.isNull}) {
+          ${ev.value} = (${ctx.boxedType(this.dataType)}) $objectTerm;
+        }
+      """
+    } else {
+      ev.isNull = "false"
+      s"""
+        /* expression: ${this.toCommentSafeString} */
+        Object $objectTerm = expressions[${ctx.references.size - 1}].eval(${ctx.INPUT_ROW});
+        ${ctx.javaType(this.dataType)} ${ev.value} = (${ctx.boxedType(this.dataType)}) $objectTerm;
+      """
+    }
   }
 }