Add date_diff Spark function (5993)

Author: rui-mo

velox/docs/functions/spark/datetime.rst

@@ -13,6 +13,13 @@ These functions support TIMESTAMP and DATE input types.
     If num_days is a negative value then these amount of days will be
     deducted from start_date.
 
+.. spark:function:: date_diff(endDate, startDate) -> integer
+    Returns the number of days from startDate to endDate. 
+    ``datediff`` is` an alias for ``date_diff``::
+
+    SELECT date_diff('2009-07-31', '2009-07-30'); -- 1
+    SELECT date_diff('2009-07-30', '2009-07-31'); -- -1
+
 .. spark:function:: date_sub(start_date, num_days) -> date
 
     Returns the date that is num_days before start_date. According to the inputs,

velox/functions/sparksql/DateTimeFunctions.h

@@ -319,4 +319,20 @@ struct DayOfWeekFunction : public InitSessionTimezone<T>,
   }
 };
 
+template <typename T>
+struct DateDiffFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE bool call(
+      int32_t& result,
+      const arg_type<Date>& endDate,
+      const arg_type<Date>& startDate) {
+    int64_t value = diffDate(DateTimeUnit::kDay, startDate, endDate);
+    if (value != (int32_t)value) {
+      VELOX_UNSUPPORTED("integer overflow");
+    }
+    result = (int32_t)value;
+    return true;
+  }
+};
 } // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/Register.cpp

@@ -233,7 +233,10 @@ void registerFunctions(const std::string& prefix) {
       Varchar>({prefix + "unix_timestamp", prefix + "to_unix_timestamp"});
   registerFunction<MakeDateFunction, Date, int32_t, int32_t, int32_t>(
       {prefix + "make_date"});
-
+  registerFunction<DateDiffFunction, int32_t, Date, Date>(
+      {prefix + "date_diff"});
+  registerFunction<DateDiffFunction, int32_t, Date, Date>(
+      {prefix + "datediff"});
   registerFunction<LastDayFunction, Date, Date>({prefix + "last_day"});
 
   registerFunction<DateAddFunction, Date, Date, int32_t>({prefix + "date_add"});

velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp

@@ -379,5 +379,41 @@ TEST_F(DateTimeFunctionsTest, dayofWeekTs) {
   }
 }
 
+TEST_F(DateTimeFunctionsTest, dateDiffDate) {
+  const auto dateDiff = [&](std::optional<int32_t> endDate,
+                            std::optional<int32_t> startDate) {
+    return evaluateOnce<int32_t, int32_t>(
+        "date_diff(c0, c1)", {endDate, startDate}, {DATE(), DATE()});
+  };
+
+  const auto dateDiffAlias = [&](std::optional<int32_t> endDate,
+                                 std::optional<int32_t> startDate) {
+    return evaluateOnce<int32_t, int32_t>(
+        "datediff(c0, c1)", {endDate, startDate}, {DATE(), DATE()});
+  };
+  // Check null behaviors
+  EXPECT_EQ(std::nullopt, dateDiff(1, std::nullopt));
+  EXPECT_EQ(std::nullopt, dateDiff(std::nullopt, 0));
+
+  // Simple tests
+  EXPECT_EQ(-1, dateDiff(parseDate("2019-02-28"), parseDate("2019-03-01")));
+  EXPECT_EQ(-358, dateDiff(parseDate("2019-02-28"), parseDate("2020-02-21")));
+  EXPECT_EQ(0, dateDiff(parseDate("1994-04-20"), parseDate("1994-04-20")));
+
+  // Account for the last day of a year-month
+  EXPECT_EQ(395, dateDiff(parseDate("2020-02-29"), parseDate("2019-01-30")));
+
+  // Check Large date
+  EXPECT_EQ(
+      -737790, dateDiff(parseDate("2020-02-29"), parseDate("4040-02-29")));
+
+  // Negative year
+  EXPECT_EQ(
+      2147474628,
+      dateDiff(parseDate("-5877641-06-23"), parseDate("1994-09-12")));
+
+  EXPECT_EQ(0, dateDiffAlias(parseDate("1994-04-20"), parseDate("1994-04-20")));
+}
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test