Support dayofmonth function in spark sql (7115)

oap-project · Oct 29, 2023 · 14af30f · 14af30f
1 parent 20a7a8c
commit 14af30f
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 3 deletions.
diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -427,4 +427,27 @@ struct MonthFunction : public InitSessionTimezone<T>,
   }
 };
 
+template <typename T>
+struct DayFunction : public InitSessionTimezone<T>,
+                     public TimestampWithTimezoneSupport<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<Timestamp>& timestamp) {
+    result = getDateTime(timestamp, this->timeZone_).tm_mday;
+  }
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getDateTime(date).tm_mday;
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
+    auto timestamp = this->toTimestamp(timestampWithTimezone);
+    result = getDateTime(timestamp, nullptr).tm_mday;
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
@@ -275,9 +275,9 @@ void registerFunctions(const std::string& prefix) {
   registerFunction<DateAddFunction, Date, Date, int32_t>({prefix + "date_add"});
   registerFunction<DateSubFunction, Date, Date, int32_t>({prefix + "date_sub"});
 
-  registerFunction<DayFunction, int64_t, Timestamp>(
+  registerFunction<DayFunction, int32_t, Timestamp>(
       {prefix + "day", prefix + "dayofmonth"});
-  registerFunction<DayFunction, int64_t, Date>(
+  registerFunction<DayFunction, int32_t, Date>(
       {prefix + "day", prefix + "dayofmonth"});
   registerFunction<DayOfYearFunction, int64_t, Timestamp>(
       {prefix + "doy", prefix + "dayofyear"});

diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
@@ -301,13 +301,36 @@ TEST_F(DateTimeFunctionsTest, dayOfYear) {
 
 TEST_F(DateTimeFunctionsTest, dayOfMonth) {
   const auto day = [&](std::optional<int32_t> date) {
-    return evaluateOnce<int64_t, int32_t>("dayofmonth(c0)", {date}, {DATE()});
+    return evaluateOnce<int32_t, int32_t>("dayofmonth(c0)", {date}, {DATE()});
   };
   EXPECT_EQ(std::nullopt, day(std::nullopt));
   EXPECT_EQ(30, day(parseDate("2009-07-30")));
   EXPECT_EQ(23, day(parseDate("2023-08-23")));
 }
 
+TEST_F(DateTimeFunctionsTest, dayOfMonthTimestamp) {
+  const auto day = [&](std::optional<Timestamp> date) {
+    return evaluateOnce<int32_t>("dayofmonth(c0)", date);
+  };
+  EXPECT_EQ(std::nullopt, day(std::nullopt));
+  EXPECT_EQ(1, day(Timestamp(0, 0)));
+  EXPECT_EQ(31, day(Timestamp(-1, 9000)));
+  EXPECT_EQ(30, day(Timestamp(1632989700, 0)));
+  EXPECT_EQ(1, day(Timestamp(1633076100, 0)));
+  EXPECT_EQ(6, day(Timestamp(1633508100, 0)));
+  EXPECT_EQ(31, day(Timestamp(1635668100, 0)));
+
+  setQueryTimeZone("Pacific/Apia");
+
+  EXPECT_EQ(std::nullopt, day(std::nullopt));
+  EXPECT_EQ(31, day(Timestamp(0, 0)));
+  EXPECT_EQ(31, day(Timestamp(-1, 9000)));
+  EXPECT_EQ(30, day(Timestamp(1632989700, 0)));
+  EXPECT_EQ(1, day(Timestamp(1633076100, 0)));
+  EXPECT_EQ(6, day(Timestamp(1633508100, 0)));
+  EXPECT_EQ(31, day(Timestamp(1635668100, 0)));
+}
+
 TEST_F(DateTimeFunctionsTest, dayOfWeekDate) {
   const auto dayOfWeek = [&](std::optional<int32_t> date,
                              const std::string& func) {