Add quarter Spark function (#7111)

Summary: The function returns the quarter for the given number of days since 1970-01-01. It takes a DateType as input and returns an Integer. You can find the implementation of this function in Spark [here](https://github.com/apache/spark/blob/b0791b513da3f0671417b9fbcd3a0caddbb45318/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala#L822). The unit test for this Spark function can be found [here](https://github.com/apache/spark/blob/b0791b513da3f0671417b9fbcd3a0caddbb45318/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala#L149). Pull Request resolved: #7111 Reviewed By: Yuhta Differential Revision: D51107854 Pulled By: mbasmanova fbshipit-source-id: 874d3b824ece71a93754e8635a182c8a5e956cef
facebookincubator · Nov 8, 2023 · 2b5c86d · 2b5c86d
1 parent 4659a3d
commit 2b5c86d
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 0 deletions.
diff --git a/velox/docs/functions/spark/datetime.rst b/velox/docs/functions/spark/datetime.rst
@@ -85,6 +85,12 @@ These functions support TIMESTAMP and DATE input types.
     ``day`` need to be from 1 to 31, and matches the number of days in each month.
     days of ``year-month-day - 1970-01-01`` need to be in the range of INTEGER type.
 
+.. spark:function:: quarter(date) -> integer
+
+    Returns the quarter of ``date``. The value ranges from ``1`` to ``4``. ::
+
+        SELECT quarter('2009-07-30'); -- 3
+
 .. spark:function:: to_unix_timestamp(string) -> integer
 
     Alias for ``unix_timestamp(string) -> integer``.

diff --git a/velox/functions/lib/TimeUtils.h b/velox/functions/lib/TimeUtils.h
@@ -82,6 +82,10 @@ FOLLY_ALWAYS_INLINE int getDay(const std::tm& time) {
   return time.tm_mday;
 }
 
+FOLLY_ALWAYS_INLINE int32_t getQuarter(const std::tm& time) {
+  return time.tm_mon / 3 + 1;
+}
+
 template <typename T>
 struct InitSessionTimezone {
   VELOX_DEFINE_FUNCTION_TYPES(T);

diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -372,4 +372,14 @@ struct AddMonthsFunction {
     result = daysSinceEpoch;
   }
 };
+
+template <typename T>
+struct QuarterFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getQuarter(getDateTime(date));
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
@@ -265,6 +265,8 @@ void registerFunctions(const std::string& prefix) {
   registerFunction<DayOfWeekFunction, int32_t, Date>(
       {prefix + "dow", prefix + "dayofweek"});
 
+  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
@@ -431,5 +431,18 @@ TEST_F(DateTimeFunctionsTest, addMonths) {
       addMonths("2023-07-10", kMax),
       fmt::format("Integer overflow in add_months(2023-07-10, {})", kMax));
 }
+
+TEST_F(DateTimeFunctionsTest, quarterDate) {
+  const auto quarter = [&](const std::string& dateString) {
+    return evaluateOnce<int32_t, int32_t>(
+        "quarter(c0)", {parseDate(dateString)}, {DATE()});
+  };
+
+  EXPECT_EQ(2, quarter("2015-04-08"));
+  EXPECT_EQ(4, quarter("2013-11-08"));
+  EXPECT_EQ(1, quarter("1987-01-08"));
+  EXPECT_EQ(3, quarter("1954-08-08"));
+}
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test