function date_add , date_sub

Author: adrian-wang

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

@@ -183,7 +183,9 @@ object FunctionRegistry {
     // datetime functions
     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"),

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

@@ -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)

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

@@ -22,7 +22,8 @@ import java.text.SimpleDateFormat
 import java.util.Calendar
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.sql.types.{StringType, TimestampType, DateType}
+import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.types.{IntegerType, StringType, TimestampType, DateType}
 
 class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
 
@@ -246,4 +247,36 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("date_add") {
+    checkResult(
+      DateAdd(Literal(Date.valueOf("2016-02-28")), Literal(1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2016-02-29")))
+    checkResult(
+      DateAdd(Literal("2016-03-01"), Literal(-1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2016-02-29")))
+    checkResult(
+      DateAdd(Literal(Timestamp.valueOf("2016-03-01 23:59:59")), Literal(-2)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2016-02-28")))
+    checkResult(
+      DateAdd(Literal("2016-03-01 23:59:59"), Literal(-3)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2016-02-27")))
+    checkResult(DateAdd(Literal(null), Literal(-1)), null)
+  }
+
+  test("date_sub") {
+    checkResult(
+      DateSub(Literal("2015-01-01"), Literal(1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2014-12-31")))
+    checkResult(
+      DateSub(Literal(Date.valueOf("2015-01-01")), Literal(-1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2015-01-02")))
+    checkResult(
+      DateSub(Literal(Timestamp.valueOf("2015-01-01 01:00:00")), Literal(-1)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2015-01-02")))
+    checkResult(
+      DateSub(Literal("2015-01-01 01:00:00"), Literal(0)),
+      DateTimeUtils.fromJavaDate(Date.valueOf("2015-01-01")))
+    checkResult(
+      DateSub(Literal("2015-01-01"), Literal.create(null, IntegerType)), null)
+  }
 }

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

@@ -2293,6 +2293,20 @@ object functions {
   def date_format(dateColumnName: String, format: String): Column =
     date_format(Column(dateColumnName), format)
 
+  /**
+   * Extracts the year as an integer from a given date/timestamp/string.
+   * @group datetime_funcs
+   * @since 1.5.0
+   */
+  def date_add(startdate: Column, days: Column): Column = DateAdd(startdate.expr, days.expr)
+
+  /**
+   * Extracts the year as an integer from a given date/timestamp/string.
+   * @group datetime_funcs
+   * @since 1.5.0
+   */
+  def date_sub(startdate: Column, days: Column): Column = DateSub(startdate.expr, days.expr)
+
   /**
    * Extracts the year as an integer from a given date/timestamp/string.
    * @group datetime_funcs

sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala

@@ -184,4 +184,34 @@ class DateFunctionsSuite extends QueryTest {
       Row(15, 15, 15))
   }
 
+  test("function date_add") {
+    val df =
+      Seq((1, Date.valueOf("2015-06-01")), (3, Date.valueOf("2015-06-02"))).toDF("num", "day")
+    checkAnswer(
+      df.select(date_add(col("day"), col("num"))),
+      Seq(Row(Date.valueOf("2015-06-02")), Row(Date.valueOf("2015-06-05"))))
+    checkAnswer(
+      df.select(date_add(column("day"), lit(null))).limit(1), Row(null))
+
+    checkAnswer(df.selectExpr("DATE_ADD(null, num)"), Seq(Row(null), Row(null)))
+    checkAnswer(
+      df.selectExpr("""DATE_ADD(day, num)"""),
+      Seq(Row(Date.valueOf("2015-06-02")), Row(Date.valueOf("2015-06-05"))))
+  }
+
+  test("function date_sub") {
+    val df =
+      Seq((1, Date.valueOf("2015-06-01")), (3, Date.valueOf("2015-06-02"))).toDF("num", "day")
+    checkAnswer(
+      df.select(date_sub(col("day"), col("num"))),
+      Seq(Row(Date.valueOf("2015-05-31")), Row(Date.valueOf("2015-05-30"))))
+    checkAnswer(
+      df.select(date_sub(lit(null), column("num"))).limit(1), Row(null))
+
+    checkAnswer(df.selectExpr("""DATE_SUB(day, null)"""), Seq(Row(null), Row(null)))
+    checkAnswer(
+      df.selectExpr("""DATE_SUB(day, num)"""),
+      Seq(Row(Date.valueOf("2015-05-31")), Row(Date.valueOf("2015-05-30"))))
+  }
+
 }