implement getLastDayOfMonth to avoid repeated evaluation

Author: yjshen

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"

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

@@ -601,22 +601,51 @@ object DateTimeUtils {
   }
 
   /**
-   * number of days in a non-leap year.
+   * Returns the number of days till the month end.
+   * if the `date` itself is the last day of a month, just return 0.
    */
-  private[this] val daysInNormalYear = Array(31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
+  private[this] def daysToMonthEnd(date: Int): Int = {
+    var (year, dayInYear) = getYearAndDayInYear(date)
+    if (isLeapYear(year)) {
+      if (dayInYear > 31 && dayInYear <= 60) {
+        return 60 - dayInYear
+      } else if (dayInYear > 60) {
+        dayInYear = dayInYear - 1
+      }
+    }
+    val lastDayOfMonth = 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 {
+      365
+    }
+    lastDayOfMonth - dayInYear
+  }
 
   /**
    * 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
-    } else {
-      date + daysInNormalYear(month - 1) - dayOfMonth
-    }
+    date + daysToMonthEnd(date)
   }
-
 }