From ee1283ada4de81f532a74797c021438e6f6129d8 Mon Sep 17 00:00:00 2001
From: rui-mo <rui.mo@intel.com>
Date: Wed, 18 Oct 2023 14:45:31 +0800
Subject: [PATCH] Add quarter function support in spark sql (7111)

---
 velox/functions/sparksql/DateTimeFunctions.h  | 28 ++++++++++++++
 velox/functions/sparksql/Register.cpp         |  3 ++
 .../sparksql/tests/DateTimeFunctionsTest.cpp  | 38 +++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
index eef388ebb0c2..2030416b40f9 100644
--- a/velox/functions/sparksql/DateTimeFunctions.h
+++ b/velox/functions/sparksql/DateTimeFunctions.h
@@ -372,4 +372,32 @@ struct AddMonthsFunction {
     result = daysSinceEpoch;
   }
 };
+
+template <typename T>
+struct QuarterFunction : public InitSessionTimezone<T>,
+                         public TimestampWithTimezoneSupport<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE int32_t getQuarter(const std::tm& time) {
+    return time.tm_mon / 3 + 1;
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<Timestamp>& timestamp) {
+    result = getQuarter(getDateTime(timestamp, this->timeZone_));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getQuarter(getDateTime(date));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
+    auto timestamp = this->toTimestamp(timestampWithTimezone);
+    result = getQuarter(getDateTime(timestamp, nullptr));
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
index 98d65ec8d0ae..a1ffe74fda29 100644
--- a/velox/functions/sparksql/Register.cpp
+++ b/velox/functions/sparksql/Register.cpp
@@ -289,6 +289,9 @@ void registerFunctions(const std::string& prefix) {
   registerFunction<DayOfWeekFunction, int32_t, Date>(
       {prefix + "dow", prefix + "dayofweek"});
 
+  registerFunction<QuarterFunction, int32_t, Timestamp>({prefix + "quarter"});
+  registerFunction<QuarterFunction, int32_t, Date>({prefix + "quarter"});
+
   // Register bloom filter function
   registerFunction<BloomFilterMightContainFunction, bool, Varbinary, int64_t>(
       {prefix + "might_contain"});
diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
index 7250bdc957e8..17f8df2e3f5d 100644
--- a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
+++ b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
@@ -431,5 +431,43 @@ TEST_F(DateTimeFunctionsTest, addMonths) {
       addMonths("2023-07-10", kMax),
       fmt::format("Integer overflow in add_months(2023-07-10, {})", kMax));
 }
+
+TEST_F(DateTimeFunctionsTest, quarter) {
+  const auto quarter = [&](std::optional<Timestamp> date) {
+    return evaluateOnce<int32_t>("quarter(c0)", date);
+  };
+  EXPECT_EQ(std::nullopt, quarter(std::nullopt));
+  EXPECT_EQ(1, quarter(Timestamp(0, 0)));
+  EXPECT_EQ(4, quarter(Timestamp(-1, 9000)));
+  EXPECT_EQ(4, quarter(Timestamp(4000000000, 0)));
+  EXPECT_EQ(4, quarter(Timestamp(4000000000, 123000000)));
+  EXPECT_EQ(2, quarter(Timestamp(990000000, 321000000)));
+  EXPECT_EQ(3, quarter(Timestamp(998423705, 321000000)));
+
+  setQueryTimeZone("Pacific/Apia");
+
+  EXPECT_EQ(std::nullopt, quarter(std::nullopt));
+  EXPECT_EQ(4, quarter(Timestamp(0, 0)));
+  EXPECT_EQ(4, quarter(Timestamp(-1, Timestamp::kMaxNanos)));
+  EXPECT_EQ(4, quarter(Timestamp(4000000000, 0)));
+  EXPECT_EQ(4, quarter(Timestamp(4000000000, 123000000)));
+  EXPECT_EQ(2, quarter(Timestamp(990000000, 321000000)));
+  EXPECT_EQ(3, quarter(Timestamp(998423705, 321000000)));
+}
+
+TEST_F(DateTimeFunctionsTest, quarterDate) {
+  const auto quarter = [&](std::optional<int32_t> date) {
+    return evaluateOnce<int32_t, int32_t>("quarter(c0)", {date}, {DATE()});
+  };
+  EXPECT_EQ(std::nullopt, quarter(std::nullopt));
+  EXPECT_EQ(1, quarter(0));
+  EXPECT_EQ(4, quarter(-1));
+  EXPECT_EQ(4, quarter(-40));
+  EXPECT_EQ(2, quarter(110));
+  EXPECT_EQ(3, quarter(200));
+  EXPECT_EQ(1, quarter(18262));
+  EXPECT_EQ(1, quarter(-18262));
+}
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test