[SPARK-9398] [SQL] Datetime cleanup

JIRA: https://issues.apache.org/jira/browse/SPARK-9398

Author: Yijie Shen <henry.yijieshen@gmail.com>

Closes #7725 from yjshen/date_null_check and squashes the following commits:

b4eade1 [Yijie Shen] inline daysToMonthEnd
d09acc1 [Yijie Shen] implement getLastDayOfMonth to avoid repeated evaluation
d857ec3 [Yijie Shen] add null check in DateExpressionSuite

3 files changed, 47 insertions, 43 deletions

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeFunctions.scala
index c37afc13f2..efecb771f2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeFunctions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeFunctions.scala
@@ -74,9 +74,7 @@ case class Hour(child: Expression) extends UnaryExpression with ImplicitCastInpu
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getHours($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getHours($c)")
   }
 }
 
@@ -92,9 +90,7 @@ case class Minute(child: Expression) extends UnaryExpression with ImplicitCastIn
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getMinutes($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getMinutes($c)")
   }
 }
 
@@ -110,9 +106,7 @@ case class Second(child: Expression) extends UnaryExpression with ImplicitCastIn
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getSeconds($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getSeconds($c)")
   }
 }
 
@@ -128,9 +122,7 @@ case class DayOfYear(child: Expression) extends UnaryExpression with ImplicitCas
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getDayInYear($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getDayInYear($c)")
   }
 }
 
@@ -147,9 +139,7 @@ case class Year(child: Expression) extends UnaryExpression with ImplicitCastInpu
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, c =>
-      s"""$dtu.getYear($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getYear($c)")
   }
 }
 
@@ -165,9 +155,7 @@ case class Quarter(child: Expression) extends UnaryExpression with ImplicitCastI
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getQuarter($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getQuarter($c)")
   }
 }
 
@@ -183,9 +171,7 @@ case class Month(child: Expression) extends UnaryExpression with ImplicitCastInp
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getMonth($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getMonth($c)")
   }
 }
 
@@ -201,9 +187,7 @@ case class DayOfMonth(child: Expression) extends UnaryExpression with ImplicitCa
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (c) =>
-      s"""$dtu.getDayOfMonth($c)"""
-    )
+    defineCodeGen(ctx, ev, c => s"$dtu.getDayOfMonth($c)")
   }
 }
 
@@ -226,7 +210,7 @@ case class WeekOfYear(child: Expression) extends UnaryExpression with ImplicitCa
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    nullSafeCodeGen(ctx, ev, (time) => {
+    nullSafeCodeGen(ctx, ev, time => {
       val cal = classOf[Calendar].getName
       val c = ctx.freshName("cal")
       ctx.addMutableState(cal, c,
@@ -250,8 +234,6 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
 
   override def inputTypes: Seq[AbstractDataType] = Seq(TimestampType, StringType)
 
-  override def prettyName: String = "date_format"
-
   override protected def nullSafeEval(timestamp: Any, format: Any): Any = {
     val sdf = new SimpleDateFormat(format.toString)
     UTF8String.fromString(sdf.format(new Date(timestamp.asInstanceOf[Long] / 1000)))
@@ -264,6 +246,8 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
           .format(new java.sql.Date($timestamp / 1000)))"""
     })
   }
+
+  override def prettyName: String = "date_format"
 }
 
 /**
@@ -277,15 +261,12 @@ case class LastDay(startDate: Expression) extends UnaryExpression with ImplicitC
   override def dataType: DataType = DateType
 
   override def nullSafeEval(date: Any): Any = {
-    val days = date.asInstanceOf[Int]
-    DateTimeUtils.getLastDayOfMonth(days)
+    DateTimeUtils.getLastDayOfMonth(date.asInstanceOf[Int])
   }
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (sd) => {
-      s"$dtu.getLastDayOfMonth($sd)"
-    })
+    defineCodeGen(ctx, ev, sd => s"$dtu.getLastDayOfMonth($sd)")
   }
 
   override def prettyName: String = "last_day"
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index 8b0b80c26d..93966a503c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -601,22 +601,43 @@ object DateTimeUtils {
   }
 
   /**
-   * number of days in a non-leap year.
-   */
-  private[this] val daysInNormalYear = Array(31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
-
-  /**
    * Returns last day of the month for the given date. The date is expressed in days
    * since 1.1.1970.
    */
   def getLastDayOfMonth(date: Int): Int = {
-    val dayOfMonth = getDayOfMonth(date)
-    val month = getMonth(date)
-    if (month == 2 && isLeapYear(getYear(date))) {
-      date + daysInNormalYear(month - 1) + 1 - dayOfMonth
+    var (year, dayInYear) = getYearAndDayInYear(date)
+    if (isLeapYear(year)) {
+      if (dayInYear > 31 && dayInYear <= 60) {
+        return date + (60 - dayInYear)
+      } else if (dayInYear > 60) {
+        dayInYear = dayInYear - 1
+      }
+    }
+    val lastDayOfMonthInYear = if (dayInYear <= 31) {
+      31
+    } else if (dayInYear <= 59) {
+      59
+    } else if (dayInYear <= 90) {
+      90
+    } else if (dayInYear <= 120) {
+      120
+    } else if (dayInYear <= 151) {
+      151
+    } else if (dayInYear <= 181) {
+      181
+    } else if (dayInYear <= 212) {
+      212
+    } else if (dayInYear <= 243) {
+      243
+    } else if (dayInYear <= 273) {
+      273
+    } else if (dayInYear <= 304) {
+      304
+    } else if (dayInYear <= 334) {
+      334
     } else {
-      date + daysInNormalYear(month - 1) - dayOfMonth
+      365
     }
+    date + (lastDayOfMonthInYear - dayInYear)
   }
-
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala
index 30c5769424..aca8d6eb35 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala
@@ -106,6 +106,7 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
         }
       }
     }
+    checkEvaluation(DayOfYear(Literal.create(null, DateType)), null)
   }
 
   test("Year") {
@@ -274,6 +275,7 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(LastDay(Literal(Date.valueOf("2015-12-05"))), Date.valueOf("2015-12-31"))
     checkEvaluation(LastDay(Literal(Date.valueOf("2016-01-06"))), Date.valueOf("2016-01-31"))
     checkEvaluation(LastDay(Literal(Date.valueOf("2016-02-07"))), Date.valueOf("2016-02-29"))
+    checkEvaluation(LastDay(Literal.create(null, DateType)), null)
   }
 
   test("next_day") {