Add month function support in spark sql (7114)

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

Hide file tree

Showing 3 changed files with 65 additions and 0 deletions.
diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -400,4 +400,31 @@ struct QuarterFunction : public InitSessionTimezone<T>,
   }
 };
 
+template <typename T>
+struct MonthFunction : public InitSessionTimezone<T>,
+                       public TimestampWithTimezoneSupport<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE int32_t getMonth(const std::tm& time) {
+    return 1 + time.tm_mon;
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<Timestamp>& timestamp) {
+    result = getMonth(getDateTime(timestamp, this->timeZone_));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getMonth(getDateTime(date));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
+    auto timestamp = this->toTimestamp(timestampWithTimezone);
+    result = getMonth(getDateTime(timestamp, nullptr));
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
@@ -291,6 +291,8 @@ void registerFunctions(const std::string& prefix) {
 
   registerFunction<QuarterFunction, int32_t, Timestamp>({prefix + "quarter"});
   registerFunction<QuarterFunction, int32_t, Date>({prefix + "quarter"});
+  registerFunction<MonthFunction, int32_t, Timestamp>({prefix + "month"});
+  registerFunction<MonthFunction, int32_t, Date>({prefix + "month"});
 
   // Register bloom filter function
   registerFunction<BloomFilterMightContainFunction, bool, Varbinary, int64_t>(

diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
@@ -469,5 +469,41 @@ TEST_F(DateTimeFunctionsTest, quarterDate) {
   EXPECT_EQ(1, quarter(-18262));
 }
 
+TEST_F(DateTimeFunctionsTest, month) {
+  const auto month = [&](std::optional<Timestamp> date) {
+    return evaluateOnce<int32_t>("month(c0)", date);
+  };
+  EXPECT_EQ(std::nullopt, month(std::nullopt));
+  EXPECT_EQ(1, month(Timestamp(0, 0)));
+  EXPECT_EQ(12, month(Timestamp(-1, 9000)));
+  EXPECT_EQ(10, month(Timestamp(4000000000, 0)));
+  EXPECT_EQ(10, month(Timestamp(4000000000, 123000000)));
+  EXPECT_EQ(8, month(Timestamp(998474645, 321000000)));
+  EXPECT_EQ(8, month(Timestamp(998423705, 321000000)));
+
+  setQueryTimeZone("Pacific/Apia");
+
+  EXPECT_EQ(std::nullopt, month(std::nullopt));
+  EXPECT_EQ(12, month(Timestamp(0, 0)));
+  EXPECT_EQ(12, month(Timestamp(-1, Timestamp::kMaxNanos)));
+  EXPECT_EQ(10, month(Timestamp(4000000000, 0)));
+  EXPECT_EQ(10, month(Timestamp(4000000000, 123000000)));
+  EXPECT_EQ(8, month(Timestamp(998474645, 321000000)));
+  EXPECT_EQ(8, month(Timestamp(998423705, 321000000)));
+}
+
+TEST_F(DateTimeFunctionsTest, monthDate) {
+  const auto month = [&](std::optional<int32_t> date) {
+    return evaluateOnce<int32_t, int32_t>("month(c0)", {date}, {DATE()});
+  };
+  EXPECT_EQ(std::nullopt, month(std::nullopt));
+  EXPECT_EQ(1, month(0));
+  EXPECT_EQ(12, month(-1));
+  EXPECT_EQ(11, month(-40));
+  EXPECT_EQ(2, month(40));
+  EXPECT_EQ(1, month(18262));
+  EXPECT_EQ(1, month(-18262));
+}
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test