Merge remote-tracking branch 'upstream/master' into SPARK-22856

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

@@ -67,10 +67,15 @@ object ExprValue {
 }
 
 // A literal evaluation of [[ExprCode]].
-case class LiteralValue(val value: String) extends ExprValue {
+class LiteralValue(val value: String) extends ExprValue {
   override def toString: String = value
 }
 
+object LiteralValue {
+  def apply(value: String): LiteralValue = new LiteralValue(value)
+  def unapply(literal: LiteralValue): Option[String] = Some(literal.value)
+}
+
 // A variable evaluation of [[ExprCode]].
 case class VariableValue(val variableName: String) extends ExprValue {
   override def toString: String = variableName
@@ -86,10 +91,12 @@ case class GlobalValue(val value: String) extends ExprValue {
   override def toString: String = value
 }
 
+case object TrueLiteral extends LiteralValue("true")
+case object FalseLiteral extends LiteralValue("false")
 
 object ExprCode {
-  def forNonNullValue(value: String): ExprCode = {
-    ExprCode(code = "", isNull = "false", value = value)
+  def forNonNullValue(value: ExprValue): ExprCode = {
+    ExprCode(code = "", isNull = FalseLiteral, value = value)
   }
 }
 

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

@@ -279,38 +279,44 @@ case class Literal (value: Any, dataType: DataType) extends LeafExpression {
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val javaType = ctx.javaType(dataType)
     if (value == null) {
-      ev.isNull = LiteralValue("true")
-      ev.copy(s"final $javaType ${ev.value} = ${ctx.defaultValue(dataType)};")
+      val defaultValueLiteral = ctx.defaultValue(javaType) match {
+        case "null" => s"(($javaType)null)"
+        case lit => lit
+      }
+      ExprCode(code = "", isNull = LiteralValue("true"), value = LiteralValue(defaultValueLiteral))
     } else {
-      ev.isNull = LiteralValue("false")
       dataType match {
         case BooleanType | IntegerType | DateType =>
-          ev.copy(code = "", value = LiteralValue(value.toString))
+          ExprCode.forNonNullValue(LiteralValue(value.toString))
         case FloatType =>
-          val v = value.asInstanceOf[Float]
-          if (v.isNaN || v.isInfinite) {
-            val boxedValue = ctx.addReferenceObj("boxedValue", v)
-            val code = s"final $javaType ${ev.value} = ($javaType) $boxedValue;"
-            ev.copy(code = code)
-          } else {
-            ev.copy(code = "", value = LiteralValue(s"${value}f"))
+          value.asInstanceOf[Float] match {
+            case v if v.isNaN =>
+              ExprCode.forNonNullValue(LiteralValue("Float.NaN"))
+            case Float.PositiveInfinity =>
+              ExprCode.forNonNullValue(LiteralValue("Float.POSITIVE_INFINITY"))
+            case Float.NegativeInfinity =>
+              ExprCode.forNonNullValue(LiteralValue("Float.NEGATIVE_INFINITY"))
+            case _ =>
+              ExprCode.forNonNullValue(LiteralValue(s"${value}F"))
           }
         case DoubleType =>
-          val v = value.asInstanceOf[Double]
-          if (v.isNaN || v.isInfinite) {
-            val boxedValue = ctx.addReferenceObj("boxedValue", v)
-            val code = s"final $javaType ${ev.value} = ($javaType) $boxedValue;"
-            ev.copy(code = code)
-          } else {
-            ev.copy(code = "", value = LiteralValue(s"${value}D"))
+          value.asInstanceOf[Double] match {
+            case v if v.isNaN =>
+              ExprCode.forNonNullValue(LiteralValue("Double.NaN"))
+            case Double.PositiveInfinity =>
+              ExprCode.forNonNullValue(LiteralValue("Double.POSITIVE_INFINITY"))
+            case Double.NegativeInfinity =>
+              ExprCode.forNonNullValue(LiteralValue("Double.NEGATIVE_INFINITY"))
+            case _ =>
+              ExprCode.forNonNullValue(LiteralValue(s"${value}D"))
           }
         case ByteType | ShortType =>
-          ev.copy(code = "", value = LiteralValue(s"($javaType)$value"))
+          ExprCode.forNonNullValue(LiteralValue(s"($javaType)$value"))
         case TimestampType | LongType =>
-          ev.copy(code = "", value = LiteralValue(s"${value}L"))
+          ExprCode.forNonNullValue(LiteralValue(s"${value}L"))
         case _ =>
-          ev.copy(code = "", value = GlobalValue(ctx.addReferenceObj("literal", value,
-            ctx.javaType(dataType))))
+          val constRef = ctx.addReferenceObj("literal", value, javaType)
+          ExprCode.forNonNullValue(GlobalValue(constRef))
       }
     }
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarBatchScan.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.execution
 
-import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode, LiteralValue, VariableValue}
 import org.apache.spark.sql.execution.metric.SQLMetrics
 import org.apache.spark.sql.types.DataType
@@ -50,7 +50,7 @@ private[sql] trait ColumnarBatchScan extends CodegenSupport {
       dataType: DataType,
       nullable: Boolean): ExprCode = {
     val javaType = ctx.javaType(dataType)
-    val value = ctx.getValue(columnVar, dataType, ordinal)
+    val value = ctx.getValueFromVector(columnVar, dataType, ordinal)
     val isNullVar = if (nullable) {
       VariableValue(ctx.freshName("isNull"))
     } else {

sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala

@@ -111,7 +111,7 @@ trait CodegenSupport extends SparkPlan {
 
   private def prepareRowVar(ctx: CodegenContext, row: String, colVars: Seq[ExprCode]): ExprCode = {
     if (row != null) {
-      ExprCode("", "false", row)
+      ExprCode("", LiteralValue("false"), VariableValue(row))
     } else {
       if (colVars.nonEmpty) {
         val colExprs = output.zipWithIndex.map { case (attr, i) =>
@@ -126,10 +126,10 @@ trait CodegenSupport extends SparkPlan {
           |$evaluateInputs
           |${ev.code.trim}
          """.stripMargin.trim
-        ExprCode(code, "false", ev.value)
+        ExprCode(code, LiteralValue("false"), ev.value)
       } else {
         // There is no columns
-        ExprCode("", "false", "unsafeRow")
+        ExprCode("", LiteralValue("false"), VariableValue("unsafeRow"))
       }
     }
   }
@@ -154,28 +154,7 @@ trait CodegenSupport extends SparkPlan {
         }
       }
 
-    val rowVar = if (row != null) {
-      ExprCode("", LiteralValue("false"), VariableValue(row))
-    } else {
-      if (outputVars.nonEmpty) {
-        val colExprs = output.zipWithIndex.map { case (attr, i) =>
-          BoundReference(i, attr.dataType, attr.nullable)
-        }
-        val evaluateInputs = evaluateVariables(outputVars)
-        // generate the code to create a UnsafeRow
-        ctx.INPUT_ROW = row
-        ctx.currentVars = outputVars
-        val ev = GenerateUnsafeProjection.createCode(ctx, colExprs, false)
-        val code = s"""
-          |$evaluateInputs
-          |${ev.code.trim}
-         """.stripMargin.trim
-        ExprCode(code, LiteralValue("false"), ev.value)
-      } else {
-        // There is no columns
-        ExprCode("", LiteralValue("false"), VariableValue("unsafeRow"))
-      }
-    }
+    val rowVar = prepareRowVar(ctx, row, outputVars)
 
     // Set up the `currentVars` in the codegen context, as we generate the code of `inputVars`
     // before calling `parent.doConsume`. We can't set up `INPUT_ROW`, because parent needs to
@@ -261,15 +240,15 @@ trait CodegenSupport extends SparkPlan {
       parameters += s"$paramType $paramName"
       val paramIsNull = if (!attributes(i).nullable) {
         // Use constant `false` without passing `isNull` for non-nullable variable.
-        "false"
+        LiteralValue("false")
       } else {
         val isNull = ctx.freshName(s"exprIsNull_$i")
         arguments += ev.isNull
         parameters += s"boolean $isNull"
-        isNull
+        VariableValue(isNull)
       }
 
-      paramVars += ExprCode("", paramIsNull, paramName)
+      paramVars += ExprCode("", paramIsNull, VariableValue(paramName))
     }
     (arguments, parameters, paramVars)
   }