Revert the renaming in SubDate and make the returned type as input type (date/timestamp)

Author: HyukjinKwon

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

@@ -324,7 +324,7 @@ object FunctionRegistry {
     expression[DateDiff]("datediff"),
     expression[AddDays]("date_add"),
     expression[DateFormatClass]("date_format"),
-    expression[SubDays]("date_sub"),
+    expression[DateSub]("date_sub"),
     expression[DayOfMonth]("day"),
     expression[DayOfYear]("dayofyear"),
     expression[DayOfMonth]("dayofmonth"),
@@ -342,7 +342,7 @@ object FunctionRegistry {
     expression[ToDate]("to_date"),
     expression[ToUnixTimestamp]("to_unix_timestamp"),
     expression[ToUTCTimestamp]("to_utc_timestamp"),
-    expression[TruncateTimestamp]("trunc"),
+    expression[TruncInstant]("trunc"),
     expression[UnixTimestamp]("unix_timestamp"),
     expression[WeekOfYear]("weekofyear"),
     expression[Year]("year"),

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

@@ -127,7 +127,7 @@ case class AddDays(instant: Expression, days: Expression) extends AddDaysBase(in
 @ExpressionDescription(
   usage = "_FUNC_(instant, num_days) - Returns the date/timestamp that is num_days before instant.",
   extended = "> SELECT _FUNC_('2016-07-30', 1);\n '2016-07-29'")
-case class SubDays(instant: Expression, days: Expression) extends AddDaysBase(instant, days) {
+case class DateSub(instant: Expression, days: Expression) extends AddDaysBase(instant, days) {
 
   override def signModifier: Int = -1
 
@@ -935,18 +935,19 @@ case class ToDate(child: Expression) extends UnaryExpression with ImplicitCastIn
  */
 // scalastyle:off line.size.limit
 @ExpressionDescription(
-  usage = "_FUNC_(timestamp, fmt) - Returns returns timestamp with the time portion truncated to the unit specified by the format model fmt.",
+  usage = "_FUNC_(instant, fmt) - Returns returns date/timestamp with the time portion truncated to the unit specified by the format model fmt.",
   extended = "> SELECT _FUNC_('2009-02-12', 'MM')\n '2009-02-01 00:00:00'\n> SELECT _FUNC_('2015-10-27', 'YEAR');\n '2015-01-01 00:00:00'")
 // scalastyle:on line.size.limit
-case class TruncateTimestamp(timestamp: Expression, format: Expression)
+case class TruncInstant(instant: Expression, format: Expression)
   extends BinaryExpression with ImplicitCastInputTypes {
 
-  override def left: Expression = timestamp
+  override def left: Expression = instant
   override def right: Expression = format
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(TimestampType, StringType)
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(DateType, TimestampType), StringType)
 
-  override def dataType: DataType = TimestampType
+  override def dataType: DataType = instant.dataType
 
   override def nullable: Boolean = true
 
@@ -965,11 +966,12 @@ case class TruncateTimestamp(timestamp: Expression, format: Expression)
       // unknown format
       null
     } else {
-      val ts = timestamp.eval(input)
-      if (ts == null) {
-        null
-      } else {
-        DateTimeUtils.truncateTimestamp(ts.asInstanceOf[Long], level)
+      (instant.dataType, instant.eval(input)) match {
+        case (_: DateType, date: Int) =>
+          DateTimeUtils.truncateInstant(date, level)
+        case (_: TimestampType, timestamp: Long) =>
+          DateTimeUtils.truncateInstant(timestamp, level)
+        case (_, null) => null
       }
     }
   }
@@ -983,13 +985,13 @@ case class TruncateTimestamp(timestamp: Expression, format: Expression)
           boolean ${ev.isNull} = true;
           ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};""")
       } else {
-        val ts = timestamp.genCode(ctx)
+        val ist = instant.genCode(ctx)
         ev.copy(code = s"""
-          ${ts.code}
-          boolean ${ev.isNull} = ${ts.isNull};
+          ${ist.code}
+          boolean ${ev.isNull} = ${ist.isNull};
           ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
           if (!${ev.isNull}) {
-            ${ev.value} = $dtu.truncateTimestamp(${ts.value}, $truncLevel);
+            ${ev.value} = $dtu.truncateInstant(${ist.value}, $truncLevel);
           }""")
       }
     } else {
@@ -1000,7 +1002,7 @@ case class TruncateTimestamp(timestamp: Expression, format: Expression)
           if ($form == -1) {
             ${ev.isNull} = true;
           } else {
-            ${ev.value} = $dtu.truncateTimestamp($dateVal, $form);
+            ${ev.value} = $dtu.truncateInstant($dateVal, $form);
           }
         """
       })

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

@@ -839,9 +839,9 @@ object DateTimeUtils {
    * Returns the trunc timestamp from original timestamp and trunc level.
    * Trunc level should be generated using `parseTruncLevel()`, should only be 1 - 6.
    */
-  def truncateTimestamp(ts: SQLTimestamp, level: Int): SQLTimestamp = {
+  def truncateInstant(ts: SQLTimestamp, level: Int): SQLTimestamp = {
     if (level == TRUNC_TO_YEAR || level == TRUNC_TO_MONTH) {
-      daysToMillis(truncateDate(millisToDays(ts / 1000L), level)) * 1000L
+      daysToMillis(truncateInstant(millisToDays(ts / 1000L), level)) * 1000L
     } else {
       truncateTime(ts, level)
     }
@@ -851,7 +851,7 @@ object DateTimeUtils {
    * Returns the trunc date from original date and trunc level.
    * Trunc level should be generated using `parseTruncLevel()`, should only be 1 or 2.
    */
-  private def truncateDate(d: SQLDate, level: Int): SQLDate = {
+  def truncateInstant(d: SQLDate, level: Int): SQLDate = {
     if (level == TRUNC_TO_YEAR) {
       d - DateTimeUtils.getDayInYear(d) + 1
     } else if (level == TRUNC_TO_MONTH) {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala

@@ -234,22 +234,22 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
 
   test("date_sub") {
     checkEvaluation(
-      SubDays(Literal(Date.valueOf("2015-01-01")), Literal(1)),
+      DateSub(Literal(Date.valueOf("2015-01-01")), Literal(1)),
       DateTimeUtils.fromJavaDate(Date.valueOf("2014-12-31")))
     checkEvaluation(
-      SubDays(Literal(Date.valueOf("2015-01-01")), Literal(-1)),
+      DateSub(Literal(Date.valueOf("2015-01-01")), Literal(-1)),
       DateTimeUtils.fromJavaDate(Date.valueOf("2015-01-02")))
-    checkEvaluation(SubDays(Literal.create(null, DateType), Literal(1)), null)
-    checkEvaluation(SubDays(Literal(Date.valueOf("2016-02-28")), Literal.create(null, IntegerType)),
+    checkEvaluation(DateSub(Literal.create(null, DateType), Literal(1)), null)
+    checkEvaluation(DateSub(Literal(Date.valueOf("2016-02-28")), Literal.create(null, IntegerType)),
       null)
-    checkEvaluation(SubDays(Literal.create(null, DateType), Literal.create(null, IntegerType)),
+    checkEvaluation(DateSub(Literal.create(null, DateType), Literal.create(null, IntegerType)),
       null)
     checkEvaluation(
-      SubDays(Literal(Date.valueOf("2016-02-28")), positiveIntLit), -15909)
+      DateSub(Literal(Date.valueOf("2016-02-28")), positiveIntLit), -15909)
     checkEvaluation(
-      SubDays(Literal(Date.valueOf("2016-02-28")), negativeIntLit), 49628)
-    checkConsistencyBetweenInterpretedAndCodegen(SubDays, DateType, IntegerType)
-    checkEvaluation(SubDays(Literal(Timestamp.valueOf("2015-01-10 12:00:00")), Literal(1)),
+      DateSub(Literal(Date.valueOf("2016-02-28")), negativeIntLit), 49628)
+    checkConsistencyBetweenInterpretedAndCodegen(DateSub, DateType, IntegerType)
+    checkEvaluation(DateSub(Literal(Timestamp.valueOf("2015-01-10 12:00:00")), Literal(1)),
       DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2015-01-09 12:00:00")))
   }
 
@@ -395,13 +395,35 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkConsistencyBetweenInterpretedAndCodegen(ToDate, DateType)
   }
 
-  test("function trunc") {
+  test("function trunc - date") {
+    def testTrunc(input: Date, fmt: String, expected: Date): Unit = {
+      checkEvaluation(TruncInstant(Literal.create(input, DateType), Literal.create(fmt, StringType)),
+        expected)
+      checkEvaluation(
+        TruncInstant(Literal.create(input, DateType), NonFoldableLiteral.create(fmt, StringType)),
+        expected)
+    }
+
+    val date = Date.valueOf("2015-07-22")
+    Seq("yyyy", "YYYY", "year", "YEAR", "yy", "YY").foreach { fmt =>
+      testTrunc(date, fmt, Date.valueOf("2015-01-01"))
+    }
+    Seq("month", "MONTH", "mon", "MON", "mm", "MM").foreach { fmt =>
+      testTrunc(date, fmt, Date.valueOf("2015-07-01"))
+    }
+
+    testTrunc(date, null, null)
+    testTrunc(null, "MON", null)
+    testTrunc(null, null, null)
+  }
+
+  test("function trunc - timestamp") {
     def testTrunc(input: Timestamp, fmt: String, expected: Timestamp): Unit = {
       checkEvaluation(
-        TruncateTimestamp(Literal.create(input, TimestampType), Literal.create(fmt, StringType)),
+        TruncInstant(Literal.create(input, TimestampType), Literal.create(fmt, StringType)),
         expected)
       checkEvaluation(
-        TruncateTimestamp(
+        TruncInstant(
           Literal.create(input, TimestampType), NonFoldableLiteral.create(fmt, StringType)),
         expected)
     }

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -2381,7 +2381,7 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def date_sub(start: Column, days: Int): Column = withExpr { SubDays(start.expr, Literal(days)) }
+  def date_sub(start: Column, days: Int): Column = withExpr { DateSub(start.expr, Literal(days)) }
 
   /**
    * Returns the number of days from `start` to `end`.
@@ -2561,7 +2561,7 @@ object functions {
    * @since 1.5.0
    */
   def trunc(date: Column, format: String): Column = withExpr {
-    TruncateTimestamp(date.expr, Literal(format))
+    TruncInstant(date.expr, Literal(format))
   }
 
   /**