diff --git a/velox/docs/functions/spark/datetime.rst b/velox/docs/functions/spark/datetime.rst index 2f5bbf9afff3..2a2c906950be 100644 --- 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 index 54bcf8901d5f..7f51092c63b1 100644 --- a/velox/functions/sparksql/DateTimeFunctions.h +++ b/velox/functions/sparksql/DateTimeFunctions.h @@ -373,6 +373,16 @@ struct AddMonthsFunction { } }; +template +struct MonthFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type& date) { + result = getMonth(getDateTime(date)); + } +}; + + template struct QuarterFunction { VELOX_DEFINE_FUNCTION_TYPES(T); diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp index 16298d2d7fd3..23f1dda00b73 100644 --- a/velox/functions/sparksql/Register.cpp +++ b/velox/functions/sparksql/Register.cpp @@ -267,6 +267,8 @@ void registerFunctions(const std::string& prefix) { registerFunction({prefix + "quarter"}); + registerFunction({prefix + "month"}); + // Register bloom filter function registerFunction( {prefix + "might_contain"}); diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp index 48410856b46f..ecdc7132e290 100644 --- 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( + "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(