support date_add (facebookincubator#144)

liujiayi771 · Mar 7, 2023 · a50f208 · a50f208
1 parent b85f80c
commit a50f208
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 0 deletions.
diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -15,6 +15,7 @@
  */
 
 #include "velox/functions/lib/TimeUtils.h"
+#include "velox/functions/prestosql/DateTimeImpl.h"
 
 namespace facebook::velox::functions::sparksql {
 
@@ -36,4 +37,33 @@ struct YearFunction : public InitSessionTimezone<T> {
     result = getYear(getDateTime(date));
   }
 };
+
+template <typename T>
+struct DateAddFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Date>& result,
+      const arg_type<Date>& date,
+      const int32_t value) {
+    result = addToDate(date, DateTimeUnit::kDay, value);
+    return true;
+  }
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Date>& result,
+      const arg_type<Date>& date,
+      const int16_t value) {
+    result = addToDate(date, DateTimeUnit::kDay, (int32_t)value);
+    return true;
+  }
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Date>& result,
+      const arg_type<Date>& date,
+      const int8_t value) {
+    result = addToDate(date, DateTimeUnit::kDay, (int32_t)value);
+    return true;
+  }
+};
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
@@ -208,6 +208,9 @@ void registerFunctions(const std::string& prefix) {
       {prefix + "year_of_week"});
   registerFunction<YearOfWeekFunction, int32_t, TimestampWithTimezone>(
       {prefix + "year_of_week"});
+  registerFunction<DateAddFunction, Date, Date, int32_t>({"date_add"});
+  registerFunction<DateAddFunction, Date, Date, int16_t>({"date_add"});
+  registerFunction<DateAddFunction, Date, Date, int8_t>({"date_add"});
 }
 
 } // namespace sparksql

diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
@@ -29,6 +29,12 @@ class DateTimeFunctionsTest : public SparkFunctionBaseTest {
         {core::QueryConfig::kAdjustTimestampToTimezone, "true"},
     });
   }
+
+  Date parseDate(const std::string& dateStr) {
+    Date returnDate;
+    parseTo(dateStr, returnDate);
+    return returnDate;
+  }
 };
 
 TEST_F(DateTimeFunctionsTest, year) {
@@ -65,5 +71,48 @@ TEST_F(DateTimeFunctionsTest, yearDate) {
   EXPECT_EQ(1920, year(Date(-18262)));
 }
 
+TEST_F(DateTimeFunctionsTest, dateAdd) {
+  const auto dateAddInt32 = [&](std::optional<Date> date,
+                           std::optional<int32_t> value) {
+    return evaluateOnce<Date>(
+        "date_add(c0, c1)", date, value);
+  };
+  const auto dateAddInt16 = [&](std::optional<Date> date,
+                           std::optional<int16_t> value) {
+    return evaluateOnce<Date>(
+        "date_add(c0, c1)", date, value);
+  };
+  const auto dateAddInt8 = [&](std::optional<Date> date,
+                           std::optional<int8_t> value) {
+    return evaluateOnce<Date>(
+        "date_add(c0, c1)", date, value);
+  };
+
+  // Check null behaviors
+  EXPECT_EQ(std::nullopt, dateAddInt32(std::nullopt, 1));
+  EXPECT_EQ(std::nullopt, dateAddInt16(std::nullopt, 1));
+  EXPECT_EQ(std::nullopt, dateAddInt8(std::nullopt, 1));
+
+  // Simple tests
+  EXPECT_EQ(
+      parseDate("2019-03-01"), dateAddInt32(parseDate("2019-02-28"), 1));
+  EXPECT_EQ(
+      parseDate("2019-03-01"), dateAddInt16(parseDate("2019-02-28"), 1));
+  EXPECT_EQ(
+      parseDate("2019-03-01"), dateAddInt8(parseDate("2019-02-28"), 1));
+
+  // Account for the last day of a year-month
+  EXPECT_EQ(
+      parseDate("2020-02-29"), dateAddInt32(parseDate("2019-01-30"), 395));
+  EXPECT_EQ(
+      parseDate("2020-02-29"), dateAddInt16(parseDate("2019-01-30"), 395));
+
+  // Check for negative intervals
+  EXPECT_EQ(
+      parseDate("2019-02-28"), dateAddInt32(parseDate("2020-02-29"), -366));
+  EXPECT_EQ(
+      parseDate("2019-02-28"), dateAddInt16(parseDate("2020-02-29"), -366));
+}
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test