[SPARK-8186][SPARK-8187][SPARK-8194][SPARK-8198][SQL] functions: date_add, date_sub, add_months, months_between, time-interval calculation

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

@@ -181,14 +181,18 @@ object FunctionRegistry {
     expression[Upper]("upper"),
 
     // datetime functions
+    expression[AddMonths]("add_months"),
     expression[CurrentDate]("current_date"),
     expression[CurrentTimestamp]("current_timestamp"),
+    expression[DateAdd]("date_add"),
     expression[DateFormatClass]("date_format"),
+    expression[DateSub]("date_sub"),
     expression[DayOfMonth]("day"),
     expression[DayOfYear]("dayofyear"),
     expression[DayOfMonth]("dayofmonth"),
     expression[Hour]("hour"),
     expression[Month]("month"),
+    expression[MonthsBetween]("months_between"),
     expression[Minute]("minute"),
     expression[Quarter]("quarter"),
     expression[Second]("second"),

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

@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.types._
-import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.unsafe.types.{Interval, UTF8String}
 
 /**
  * Returns the current date at the start of query evaluation.
@@ -62,6 +62,53 @@ case class CurrentTimestamp() extends LeafExpression with CodegenFallback {
   }
 }
 
+/**
+ * Adds a number of days to startdate.
+ */
+case class DateAdd(startDate: Expression, days: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  override def left: Expression = startDate
+  override def right: Expression = days
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(DateType, IntegerType)
+
+  override def dataType: DataType = DateType
+
+  override def nullSafeEval(start: Any, d: Any): Any = {
+    start.asInstanceOf[Int] + d.asInstanceOf[Int]
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    nullSafeCodeGen(ctx, ev, (sd, d) => {
+      s"""${ev.primitive} = $sd + $d;"""
+    })
+  }
+}
+
+/**
+ * Subtracts a number of days to startdate.
+ */
+case class DateSub(startDate: Expression, days: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+  override def left: Expression = startDate
+  override def right: Expression = days
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(DateType, IntegerType)
+
+  override def dataType: DataType = DateType
+
+  override def nullSafeEval(start: Any, d: Any): Any = {
+    start.asInstanceOf[Int] - d.asInstanceOf[Int]
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    nullSafeCodeGen(ctx, ev, (sd, d) => {
+      s"""${ev.primitive} = $sd - $d;"""
+    })
+  }
+}
+
 case class Hour(child: Expression) extends UnaryExpression with ImplicitCastInputTypes {
 
   override def inputTypes: Seq[AbstractDataType] = Seq(TimestampType)
@@ -258,3 +305,125 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
     })
   }
 }
+
+/**
+ * Time Adds Interval.
+ */
+case class TimeAdd(left: Expression, right: Expression)
+  extends BinaryExpression with ExpectsInputTypes {
+
+  override def toString: String = s"$left + $right"
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(DateType, TimestampType), IntervalType)
+
+  override def dataType: DataType = TimestampType
+
+  override def nullSafeEval(start: Any, interval: Any): Any = {
+    val itvl = interval.asInstanceOf[Interval]
+    left.dataType match {
+      case DateType =>
+        DateTimeUtils.dateAddFullInterval(
+          start.asInstanceOf[Int], itvl.months, itvl.microseconds)
+      case TimestampType =>
+        DateTimeUtils.timestampAddFullInterval(
+          start.asInstanceOf[Long], itvl.months, itvl.microseconds)
+    }
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+    left.dataType match {
+      case DateType =>
+        defineCodeGen(ctx, ev, (sd, i) => {
+          s"""$dtu.dateAddFullInterval($sd, $i.months, $i.microseconds)"""
+        })
+      case TimestampType => // TimestampType
+        defineCodeGen(ctx, ev, (sd, i) => {
+          s"""$dtu.timestampAddFullInterval($sd, $i.months, $i.microseconds)"""
+        })
+    }
+  }
+}
+
+/**
+ * Time Subtracts Interval.
+ */
+case class TimeSub(left: Expression, right: Expression)
+  extends BinaryExpression with ExpectsInputTypes {
+
+  override def toString: String = s"$left - $right"
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(DateType, TimestampType), IntervalType)
+
+  override def dataType: DataType = TimestampType
+
+  override def nullSafeEval(start: Any, interval: Any): Any = {
+    val itvl = interval.asInstanceOf[Interval]
+    left.dataType match {
+      case DateType =>
+        DateTimeUtils.dateAddFullInterval(
+          start.asInstanceOf[Int], 0 - itvl.months, 0 - itvl.microseconds)
+      case TimestampType =>
+        DateTimeUtils.timestampAddFullInterval(
+          start.asInstanceOf[Long], 0 - itvl.months, 0 - itvl.microseconds)
+    }
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+    left.dataType match {
+      case DateType =>
+        defineCodeGen(ctx, ev, (sd, i) => {
+          s"""$dtu.dateAddFullInterval($sd, 0 - $i.months, 0 - $i.microseconds)"""
+        })
+      case TimestampType => // TimestampType
+        defineCodeGen(ctx, ev, (sd, i) => {
+          s"""$dtu.timestampAddFullInterval($sd, 0 - $i.months, 0 - $i.microseconds)"""
+        })
+    }
+  }
+}
+
+/**
+ * Returns the date that is num_months after start_date.
+ */
+case class AddMonths(left: Expression, right: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(DateType, IntegerType)
+
+  override def dataType: DataType = DateType
+
+  override def nullSafeEval(start: Any, months: Any): Any = {
+    DateTimeUtils.dateAddYearMonthInterval(start.asInstanceOf[Int], months.asInstanceOf[Int])
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+    defineCodeGen(ctx, ev, (sd, m) => {
+      s"""$dtu.dateAddYearMonthInterval($sd, $m)"""
+    })
+  }
+}
+
+/**
+ * Returns number of months between dates date1 and date2.
+ */
+case class MonthsBetween(left: Expression, right: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(TimestampType, TimestampType)
+
+  override def dataType: DataType = DoubleType
+
+  override def nullSafeEval(t1: Any, t2: Any): Any = {
+    DateTimeUtils.monthsBetween(t1.asInstanceOf[Long], t2.asInstanceOf[Long])
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+    defineCodeGen(ctx, ev, (l, r) => {
+      s"""$dtu.monthsBetween($l, $r)"""
+    })
+  }
+}

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

@@ -573,4 +573,109 @@ object DateTimeUtils {
       dayInYear - 334
     }
   }
+
+  /**
+   * Returns the date value for the first day of the given month.
+   * The month is expressed in months since 1.1.1970, starting from 0.
+   */
+  def getDaysFromMonths(months: Int): Int = {
+    val yearSinceEpoch = if (months < 0) months / 12 - 1 else months / 12
+    val monthInYearFromZero = months - yearSinceEpoch * 12
+    val daysFromYears = getDaysFromYears(yearSinceEpoch)
+    val febDays = if (isLeapYear(1970 + yearSinceEpoch)) 29 else 28
+    val daysForMonths = Seq(31, febDays, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
+    daysForMonths.slice(0, monthInYearFromZero).sum + daysFromYears
+  }
+
+  /**
+   * Returns the date value for January 1 of the given year.
+   * The year is expressed in years since 1.1.1970, starting from 0.
+   */
+  def getDaysFromYears(years: Int): Int = {
+    val remainYear = (years % 400 + 400) % 400
+    val cycles = (years - remainYear) / 400
+    val numLeaps = if (remainYear < 130) {
+      remainYear / 4
+    } else if (remainYear < 230) {
+      remainYear / 4 - 1
+    } else if (remainYear < 330) {
+      remainYear / 4 - 2
+    } else {
+      remainYear / 4 - 3
+    }
+    cycles * (365 * 400 + 397) + remainYear * 365 + numLeaps
+  }
+
+  /**
+   * Add date and year-month interval.
+   * Returns a date value, expressed in days since 1.1.1970.
+   */
+  def dateAddYearMonthInterval(days: Int, months: Int): Int = {
+    val currentMonth = (getYear(days) - 1970) * 12 + getMonth(days) - 1 + months
+    val currentMonthInYear = if (currentMonth < 0) {
+      ((currentMonth % 12) + 12) % 12
+    } else {
+      currentMonth % 12
+    }
+    val currentYear = if (currentMonth < 0) (currentMonth / 12) - 1 else currentMonth / 12
+    val febDays = if (isLeapYear(1970 + currentYear)) 29 else 28
+    val daysForMonths = Seq(31, febDays, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
+    getDaysFromMonths(currentMonth) + math.min(
+      getDayOfMonth(days), daysForMonths(currentMonthInYear)) - 1
+  }
+
+  /**
+   * Add date and full interval.
+   * Returns a timestamp value, expressed in microseconds since 1.1.1970 00:00:00.
+   */
+  def dateAddFullInterval(days: Int, months: Int, microseconds: Long): Long = {
+    daysToMillis(dateAddYearMonthInterval(days, months)) * 1000L + microseconds
+  }
+
+  /**
+   * Add timestamp and full interval.
+   * Returns a timestamp value, expressed in microseconds since 1.1.1970 00:00:00.
+   */
+  def timestampAddFullInterval(micros: Long, months: Int, microseconds: Long): Long = {
+    val days = millisToDays(micros / 1000L)
+    dateAddFullInterval(days, months, microseconds) + micros - daysToMillis(days) * 1000L
+  }
+
+  /**
+   * Returns the last dayInMonth in the month it belongs to. The date is expressed
+   * in days since 1.1.1970. the return value starts from 1.
+   */
+  def getLastDayInMonthOfMonth(date: Int): Int = {
+    val month = getMonth(date)
+
+    val febDay = if (isLeapYear(getYear(date))) 29 else 28
+    val days = Seq(31, febDay, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
+    days(month - 1)
+  }
+
+  /**
+   * Returns number of months between time1 and time2. time1 and time2 are expressed in
+   * microseconds since 1.1.1970
+   */
+  def monthsBetween(time1: Long, time2: Long): Double = {
+    val millis1 = time1.asInstanceOf[Long] / 1000L
+    val millis2 = time2.asInstanceOf[Long] / 1000L
+    val date1 = millisToDays(millis1)
+    val date2 = millisToDays(millis2)
+    val dayInMonth1 = getDayOfMonth(date1)
+    val dayInMonth2 = getDayOfMonth(date2)
+    val lastDayMonth1 = getLastDayInMonthOfMonth(date1)
+    val lastDayMonth2 = getLastDayInMonthOfMonth(date2)
+    val months1 = getYear(date1) * 12 + getMonth(date1) - 1
+    val months2 = getYear(date2) * 12 + getMonth(date2) - 1
+    val timeInDay1 = time1 - daysToMillis(date1) * 1000L
+    val timeInDay2 = time2 - daysToMillis(date2) * 1000L
+    if (dayInMonth1 == dayInMonth2 || (lastDayMonth1 == dayInMonth1 &&
+      lastDayMonth2 == dayInMonth2)) {
+      (months1 - months2).toDouble
+    } else {
+      val timesBetween = (timeInDay1 - timeInDay2).toDouble / (MILLIS_PER_DAY * 1000)
+      (months1 - months2).toDouble + (dayInMonth1 - dayInMonth2 + timesBetween) / 31.0
+    }
+  }
 }

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

@@ -22,7 +22,9 @@ import java.text.SimpleDateFormat
 import java.util.Calendar
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.types.{StringType, TimestampType, DateType}
+import org.apache.spark.unsafe.types.Interval
 
 class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
 
@@ -246,4 +248,57 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("date_add") {
+    checkEvaluation(
+      DateAdd(Literal(Date.valueOf("2016-02-28")), Literal(1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2016-02-29")))
+    checkEvaluation(
+      DateAdd(Literal(Date.valueOf("2016-02-28")), Literal(-365)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2015-02-28")))
+  }
+
+  test("date_sub") {
+    checkEvaluation(
+      DateSub(Literal(Date.valueOf("2015-01-01")), Literal(1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2014-12-31")))
+    checkEvaluation(
+      DateSub(Literal(Date.valueOf("2015-01-01")), Literal(-1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2015-01-02")))
+  }
+
+  test("time_add") {
+    checkEvaluation(
+      TimeAdd(Literal(Date.valueOf("2016-01-29")), Literal(new Interval(1, 0))),
+      DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2016-02-29 00:00:00")))
+    checkEvaluation(
+      TimeAdd(Literal(Date.valueOf("2016-01-31")), Literal(new Interval(1, 2000000.toLong))),
+      DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2016-02-29 00:00:02")))
+  }
+
+  test("time_sub") {
+    checkEvaluation(
+      TimeSub(Literal(Timestamp.valueOf("2016-03-31 10:00:00")), Literal(new Interval(1, 0))),
+      DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2016-02-29 10:00:00")))
+    checkEvaluation(
+      TimeSub(
+        Literal(Timestamp.valueOf("2016-03-30 00:00:01")),
+        Literal(new Interval(1, 2000000.toLong))),
+      DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2016-02-28 23:59:59")))
+  }
+
+  test("add_months") {
+    checkEvaluation(AddMonths(Literal(Date.valueOf(
+      "2015-01-30")), Literal(1)), DateTimeUtils.fromJavaDate(Date.valueOf("2015-02-28")))
+    checkEvaluation(AddMonths(Literal(Date.valueOf(
+      "2016-03-30")), Literal(-1)), DateTimeUtils.fromJavaDate(Date.valueOf("2016-02-29")))
+  }
+
+  test("months_between") {
+    checkEvaluation(MonthsBetween(Literal(Timestamp.valueOf(
+      "2015-01-30 11:52:00")), Literal(Timestamp.valueOf("2015-01-30 11:50:00"))), 0.0)
+    checkEvaluation(MonthsBetween(Literal(Timestamp.valueOf(
+      "2015-01-31 00:00:00")), Literal(Timestamp.valueOf("2015-03-31 22:00:00"))), -2.0)
+    checkEvaluation(MonthsBetween(Literal(Timestamp.valueOf(
+      "2015-03-31 22:00:00")), Literal(Timestamp.valueOf("2015-02-28 00:00:00"))), 1.0)
+  }
 }