Add month Spark function (#7114)

Summary: The function returns the month value 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#L838). 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#L172). Pull Request resolved: #7114 Reviewed By: Yuhta Differential Revision: D51107423 Pulled By: mbasmanova fbshipit-source-id: 03c7bf9718cbe88ce1e255bd6ea9d6d1bad080f8
facebookincubator · Nov 8, 2023 · b3e37fb · b3e37fb
1 parent 2b5c86d
commit b3e37fb
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 0 deletions.
diff --git a/velox/docs/functions/spark/datetime.rst b/velox/docs/functions/spark/datetime.rst
@@ -91,6 +91,12 @@ These functions support TIMESTAMP and DATE input types.
 
         SELECT quarter('2009-07-30'); -- 3
 
+.. spark:function:: month(date) -> integer
+
+    Returns the month of ``date``. ::
+
+        SELECT month('2009-07-30'); -- 7
+       
 .. spark:function:: to_unix_timestamp(string) -> integer
 
     Alias for ``unix_timestamp(string) -> integer``.

diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -373,6 +373,16 @@ struct AddMonthsFunction {
   }
 };
 
+template <typename T>
+struct MonthFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getMonth(getDateTime(date));
+  }
+};
+
+
 template <typename T>
 struct QuarterFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);

diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
@@ -267,6 +267,8 @@ void registerFunctions(const std::string& prefix) {
 
   registerFunction<QuarterFunction, int32_t, Date>({prefix + "quarter"});
 
+  registerFunction<MonthFunction, int32_t, Date>({prefix + "month"});
+
   // 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
@@ -432,6 +432,19 @@ TEST_F(DateTimeFunctionsTest, addMonths) {
       fmt::format("Integer overflow in add_months(2023-07-10, {})", kMax));
 }
 
+TEST_F(DateTimeFunctionsTest, monthDate) {
+  const auto month = [&](const std::string& dateString) {
+    return evaluateOnce<int32_t, int32_t>(
+        "month(c0)", {parseDate(dateString)}, {DATE()});
+  };
+
+  EXPECT_EQ(4, month("2015-04-08"));
+  EXPECT_EQ(11, month("2013-11-08"));
+  EXPECT_EQ(1, month("1987-01-08"));
+  EXPECT_EQ(8, month("1954-08-08"));
+}
+
+
 TEST_F(DateTimeFunctionsTest, quarterDate) {
   const auto quarter = [&](const std::string& dateString) {
     return evaluateOnce<int32_t, int32_t>(