address comments from davies

Author: adrian-wang

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

@@ -277,6 +277,8 @@ case class LastDay(startDate: Expression) extends UnaryExpression with ImplicitC
  */
 case class ToDate(child: Expression) extends UnaryExpression with ImplicitCastInputTypes {
 
+  // Implicit casting of spark will accept string in both date and timestamp format, as
+  // well as TimestampType.
   override def inputTypes: Seq[AbstractDataType] = Seq(DateType)
 
   override def dataType: DataType = DateType
@@ -361,37 +363,76 @@ case class Trunc(date: Expression, format: Expression)
   override def inputTypes: Seq[AbstractDataType] = Seq(DateType, StringType)
   override def dataType: DataType = DateType
 
-  override def nullSafeEval(d: Any, fmt: Any): Any = {
-    val minItem = DateTimeUtils.getFmt(fmt.asInstanceOf[UTF8String])
-    if (minItem == -1) {
-      // unknown format
-      null
+  lazy val constFmt = format.eval().asInstanceOf[UTF8String]
+
+  override def eval(input: InternalRow): Any = {
+    if (format.foldable) {
+      val minItem = DateTimeUtils.getFmt(constFmt)
+      if (minItem == -1) {
+        // unknown format
+        null
+      } else {
+        val d = date.eval(input)
+        if (d == null) {
+          null
+        } else {
+          DateTimeUtils.dateTrunc(d.asInstanceOf[Int], minItem)
+        }
+      }
     } else {
-      val days = d.asInstanceOf[Int]
-      if (minItem == Calendar.YEAR) {
-        days - DateTimeUtils.getDayInYear(days) + 1
+      val fmt = format.eval(input).asInstanceOf[UTF8String]
+      val d = date.eval(input)
+      if (d == null) {
+        null
       } else {
-        // trunc to MONTH
-        days - DateTimeUtils.getDayOfMonth(days) + 1
+        val minItem = DateTimeUtils.getFmt(fmt)
+        if (minItem == -1) {
+          // unknown format
+          null
+        } else {
+          DateTimeUtils.dateTrunc(d.asInstanceOf[Int], minItem)
+        }
       }
     }
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    nullSafeCodeGen(ctx, ev, (dateVal, fmt) => {
-      val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-      val form = ctx.freshName("form")
-      s"""
-        int $form = $dtu.getFmt($fmt);
-        if ($form == ${Calendar.YEAR}) {
-          ${ev.primitive} = $dateVal - $dtu.getDayInYear($dateVal) + 1;
-        } else if ($form == ${Calendar.MONTH}) {
-          ${ev.primitive} = $dateVal - $dtu.getDayInYear($dateVal) + 1;
-        } else {
-          ${ev.isNull} = true;
-        }
-      """
-    })
+    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+    if (date.foldable) {
+      val d = date.gen(ctx)
+      val minItem = DateTimeUtils.getFmt(constFmt)
+      if (d == null || minItem == -1) {
+        s"""
+          boolean ${ev.isNull} = true;
+          ${ctx.javaType(dataType)} ${ev.primitive} = ${ctx.defaultValue(dataType)};
+        """
+      } else {
+        s"""
+          ${d.code}
+          boolean ${ev.isNull} = ${d.isNull};
+          ${ctx.javaType(dataType)} ${ev.primitive} = ${ctx.defaultValue(dataType)};
+          if (!${ev.isNull}) {
+            if ($minItem == -1) {
+              ${ev.isNull} = true;
+            } else {
+              ${ev.primitive} = $dtu.dateTrunc(${d.primitive}, $minItem);
+            }
+          }
+        """
+      }
+    } else {
+      nullSafeCodeGen(ctx, ev, (dateVal, fmt) => {
+        val form = ctx.freshName("form")
+        s"""
+          int $form = $dtu.getFmt($fmt);
+          if ($form == -1) {
+            ${ev.isNull} = true;
+          } else {
+            ${ev.primitive} = $dtu.dateTrunc($dateVal, $form);
+          }
+        """
+      })
+    }
   }
 
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala

@@ -642,15 +642,33 @@ object DateTimeUtils {
   }
 
   /**
-   * Returns the truncate level, could be [[Calendar.MONTH]]/[[Calendar.YEAR]]/-1
+   * Returns the trunc date from original date and trunc level.
+   * Trunc level should be generated using `this.getFmt()`.
+   */
+  def dateTrunc(d: Int, minItem: Int): Int = {
+    if (minItem == 2) {
+      // trunc to year
+      d - DateTimeUtils.getDayInYear(d) + 1
+    } else {
+      // trunc to MONTH
+      d - DateTimeUtils.getDayOfMonth(d) + 1
+    }
+  }
+
+  /**
+   * Returns the truncate level, could be 1 for month, 2 for year, -1 for invalid/null
    * -1 means unsupported truncate level.
    */
   def getFmt(string: UTF8String): Int = {
-    val fmtString = string.toString.toUpperCase
-    fmtString match {
-      case "MON" | "MONTH" | "MM" => Calendar.MONTH
-      case "YEAR"| "YYYY" | "YY" => Calendar.YEAR
-      case _ => -1
+    if (string == null) {
+      -1
+    } else {
+      val fmtString = string.toString.toUpperCase
+      fmtString match {
+        case "MON" | "MONTH" | "MM" => 1
+        case "YEAR" | "YYYY" | "YY" => 2
+        case _ => -1
+      }
     }
   }
 }