Add day(interval) Presto function (facebookincubator#9054)

Summary: Pull Request resolved: facebookincubator#9054 Reviewed By: kagamiori Differential Revision: D54830164
mbasmanova · Mar 13, 2024 · 6d2989f · 6d2989f
1 parent de54d1e
commit 6d2989f
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 2 deletions.
diff --git a/velox/docs/functions/presto/datetime.rst b/velox/docs/functions/presto/datetime.rst
@@ -252,6 +252,8 @@ arbitrary large timestamps.
 
     Returns the day of the month from ``x``.
 
+    The supported types for ``x`` are DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, INTERVAL DAY TO SECOND.
+
 .. function:: day_of_month(x) -> bigint
 
     This is an alias for :func:`day`.

diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -330,6 +330,17 @@ struct DayFunction : public InitSessionTimezone<T>,
   }
 };
 
+template <typename T>
+struct DayFromIntervalFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      int64_t& result,
+      const arg_type<IntervalDayTime>& interval) {
+    result = interval / kMillisInDay;
+  }
+};
+
 template <typename T>
 struct LastDayOfMonthFunction : public InitSessionTimezone<T>,
                                 public TimestampWithTimezoneSupport<T> {

diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp
@@ -59,6 +59,8 @@ void registerSimpleFunctions(const std::string& prefix) {
       {prefix + "day", prefix + "day_of_month"});
   registerFunction<DayFunction, int64_t, Date>(
       {prefix + "day", prefix + "day_of_month"});
+  registerFunction<DayFromIntervalFunction, int64_t, IntervalDayTime>(
+      {prefix + "day", prefix + "day_of_month"});
   registerFunction<DateMinusInterval, Date, Date, IntervalDayTime>(
       {prefix + "minus"});
   registerFunction<DateMinusInterval, Date, Date, IntervalYearMonth>(

diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
@@ -268,8 +268,18 @@ TEST_F(DateTimeFunctionsTest, parseDatetimeSignatures) {
 }
 
 TEST_F(DateTimeFunctionsTest, dayOfXxxSignatures) {
-  for (const auto& name :
-       {"day_of_year", "doy", "day_of_month", "day_of_week", "dow"}) {
+  for (const auto& name : {"day", "day_of_month"}) {
+    SCOPED_TRACE(name);
+    auto signatures = getSignatureStrings(name);
+    ASSERT_EQ(4, signatures.size());
+
+    ASSERT_EQ(1, signatures.count("(timestamp with time zone) -> bigint"));
+    ASSERT_EQ(1, signatures.count("(date) -> bigint"));
+    ASSERT_EQ(1, signatures.count("(timestamp) -> bigint"));
+    ASSERT_EQ(1, signatures.count("(interval day to second) -> bigint"));
+  }
+
+  for (const auto& name : {"day_of_year", "doy", "day_of_week", "dow"}) {
     SCOPED_TRACE(name);
     auto signatures = getSignatureStrings(name);
     ASSERT_EQ(3, signatures.size());
@@ -804,6 +814,24 @@ TEST_F(DateTimeFunctionsTest, dayOfMonthDate) {
   EXPECT_EQ(2, day(-18262));
 }
 
+TEST_F(DateTimeFunctionsTest, dayOfMonthInterval) {
+  const auto day = [&](int64_t millis) {
+    auto result = evaluateOnce<int64_t, int64_t>(
+        "day_of_month(c0)", {millis}, {INTERVAL_DAY_TIME()});
+
+    auto result2 = evaluateOnce<int64_t, int64_t>(
+        "day(c0)", {millis}, {INTERVAL_DAY_TIME()});
+
+    EXPECT_EQ(result, result2);
+    return result;
+  };
+
+  EXPECT_EQ(1, day(kMillisInDay));
+  EXPECT_EQ(1, day(kMillisInDay + kMillisInHour));
+  EXPECT_EQ(10, day(10 * kMillisInDay + 7 * kMillisInHour));
+  EXPECT_EQ(-10, day(-10 * kMillisInDay - 7 * kMillisInHour));
+}
+
 TEST_F(DateTimeFunctionsTest, plusMinusDateIntervalYearMonth) {
   const auto makeInput = [&](const std::string& date, int32_t interval) {
     return makeRowVector({