From b3c18a6ecb6e13a1c4be87d11d3ed21cf56f4d84 Mon Sep 17 00:00:00 2001 From: Masha Basmanova Date: Wed, 12 Jun 2024 21:50:16 -0700 Subject: [PATCH] Add support for INTERVAL input to year, month, hour, minute, second, millisecond Presto functions Differential Revision: D58509062 --- velox/docs/functions/presto/datetime.rst | 17 +++-- velox/functions/prestosql/DateTimeFunctions.h | 66 +++++++++++++++++++ .../DateTimeFunctionsRegistration.cpp | 23 +++++++ .../prestosql/tests/DateTimeFunctionsTest.cpp | 34 ++++++++++ 4 files changed, 135 insertions(+), 5 deletions(-) diff --git a/velox/docs/functions/presto/datetime.rst b/velox/docs/functions/presto/datetime.rst index 480872c97a09..1fe23bf651bd 100644 --- a/velox/docs/functions/presto/datetime.rst +++ b/velox/docs/functions/presto/datetime.rst @@ -341,6 +341,8 @@ arbitrary large timestamps. .. function:: hour(x) -> bigint Returns the hour of the day from ``x``. The value ranges from 0 to 23. + Supported types for ``x`` are: DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, + INTERVAL DAY TO SECOND¶. .. function:: last_day_of_month(x) -> date @@ -348,15 +350,18 @@ arbitrary large timestamps. .. function:: millisecond(x) -> int64 - Returns the millisecond of the second from ``x``. + Returns the millisecond of the second from ``x``. Supported types for ``x`` are: + DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, INTERVAL DAY TO SECOND¶. .. function:: minute(x) -> bigint - Returns the minute of the hour from ``x``. + Returns the minute of the hour from ``x``. Supported types for ``x`` are: + DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, INTERVAL DAY TO SECOND¶. .. function:: month(x) -> bigint - Returns the month of the year from ``x``. + Returns the month of the year from ``x``. Supported types for ``x`` are: + DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, INTERVAL YEAR TO MONTH. .. function:: quarter(x) -> bigint @@ -364,7 +369,8 @@ arbitrary large timestamps. .. function:: second(x) -> bigint - Returns the second of the minute from ``x``. + Returns the second of the minute from ``x``. Supported types for ``x`` are: + DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, INTERVAL DAY TO SECOND¶. .. function:: timezone_hour(timestamp) -> bigint @@ -386,7 +392,8 @@ arbitrary large timestamps. .. function:: year(x) -> bigint - Returns the year from ``x``. + Returns the year from ``x``. Supported types for ``x`` are: + DATE, TIMESTAMP, TIMESTAMP WITH TIME ZONE, INTERVAL YEAR TO MONTH. .. function:: year_of_week(x) -> bigint diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h index 7be0e1866476..86c274970439 100644 --- a/velox/functions/prestosql/DateTimeFunctions.h +++ b/velox/functions/prestosql/DateTimeFunctions.h @@ -238,6 +238,17 @@ struct YearFunction : public InitSessionTimezone, } }; +template +struct YearFromIntervalFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& months) { + result = months / 12; + } +}; + template struct QuarterFunction : public InitSessionTimezone, public TimestampWithTimezoneSupport { @@ -292,6 +303,17 @@ struct MonthFunction : public InitSessionTimezone, } }; +template +struct MonthFromIntervalFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& months) { + result = months % 12; + } +}; + template struct DayFunction : public InitSessionTimezone, public TimestampWithTimezoneSupport { @@ -687,6 +709,17 @@ struct HourFunction : public InitSessionTimezone, } }; +template +struct HourFromIntervalFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& millis) { + result = (millis % kMillisInDay) / kMillisInHour; + } +}; + template struct MinuteFunction : public InitSessionTimezone, public TimestampWithTimezoneSupport { @@ -710,6 +743,17 @@ struct MinuteFunction : public InitSessionTimezone, } }; +template +struct MinuteFromIntervalFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& millis) { + result = (millis % kMillisInHour) / kMillisInMinute; + } +}; + template struct SecondFunction : public TimestampWithTimezoneSupport { VELOX_DEFINE_FUNCTION_TYPES(T); @@ -732,6 +776,17 @@ struct SecondFunction : public TimestampWithTimezoneSupport { } }; +template +struct SecondFromIntervalFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& millis) { + result = (millis % kMillisInMinute) / kMillisInSecond; + } +}; + template struct MillisecondFunction : public TimestampWithTimezoneSupport { VELOX_DEFINE_FUNCTION_TYPES(T); @@ -757,6 +812,17 @@ struct MillisecondFunction : public TimestampWithTimezoneSupport { } }; +template +struct MillisecondFromIntervalFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& millis) { + result = millis % kMillisecondsInSecond; + } +}; + namespace { inline std::optional fromDateTimeUnitString( const StringView& unitString, diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp index 8c4288da4256..376ee6572e97 100644 --- a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp @@ -83,10 +83,14 @@ void registerSimpleFunctions(const std::string& prefix) { registerFunction( {prefix + "timezone_minute"}); + registerFunction({prefix + "year"}); registerFunction({prefix + "year"}); registerFunction( {prefix + "year"}); + registerFunction( + {prefix + "year"}); + registerFunction( {prefix + "week", prefix + "week_of_year"}); registerFunction( @@ -97,16 +101,21 @@ void registerSimpleFunctions(const std::string& prefix) { registerFunction({prefix + "quarter"}); registerFunction( {prefix + "quarter"}); + registerFunction({prefix + "month"}); registerFunction({prefix + "month"}); registerFunction( {prefix + "month"}); + registerFunction( + {prefix + "month"}); + registerFunction( {prefix + "day", prefix + "day_of_month"}); registerFunction( {prefix + "day", prefix + "day_of_month"}); registerFunction( {prefix + "day", prefix + "day_of_month"}); + registerFunction( {prefix + "minus"}); registerFunction( @@ -147,30 +156,44 @@ void registerSimpleFunctions(const std::string& prefix) { {prefix + "yow", prefix + "year_of_week"}); registerFunction( {prefix + "yow", prefix + "year_of_week"}); + registerFunction({prefix + "hour"}); registerFunction({prefix + "hour"}); registerFunction( {prefix + "hour"}); + registerFunction( + {prefix + "hour"}); + registerFunction( {prefix + "last_day_of_month"}); registerFunction( {prefix + "last_day_of_month"}); registerFunction( {prefix + "last_day_of_month"}); + registerFunction({prefix + "minute"}); registerFunction({prefix + "minute"}); registerFunction( {prefix + "minute"}); + registerFunction( + {prefix + "minute"}); + registerFunction({prefix + "second"}); registerFunction({prefix + "second"}); registerFunction( {prefix + "second"}); + registerFunction( + {prefix + "second"}); + registerFunction( {prefix + "millisecond"}); registerFunction( {prefix + "millisecond"}); registerFunction( {prefix + "millisecond"}); + registerFunction( + {prefix + "millisecond"}); + registerFunction( {prefix + "date_trunc"}); registerFunction( diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp index a006c8f3ed27..9521a1d32170 100644 --- a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp +++ b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp @@ -1561,6 +1561,40 @@ TEST_F(DateTimeFunctionsTest, millisecondTimestampWithTimezone) { "millisecond(c0)", -980, "+05:30")); } +TEST_F(DateTimeFunctionsTest, extractFromIntervalDayTime) { + const auto millis = 5 * kMillisInDay + 7 * kMillisInHour + + 11 * kMillisInMinute + 13 * kMillisInSecond + 17; + + auto extract = [&](const std::string& unit, int64_t millis) { + return evaluateOnce( + fmt::format("{}(c0)", unit), + INTERVAL_DAY_TIME(), + std::optional(millis)) + .value(); + }; + + EXPECT_EQ(17, extract("millisecond", millis)); + EXPECT_EQ(13, extract("second", millis)); + EXPECT_EQ(11, extract("minute", millis)); + EXPECT_EQ(7, extract("hour", millis)); + EXPECT_EQ(5, extract("day", millis)); +} + +TEST_F(DateTimeFunctionsTest, extractFromIntervalYearMonth) { + const auto months = 3 * 12 + 4; + + auto extract = [&](const std::string& unit, int32_t months) { + return evaluateOnce( + fmt::format("{}(c0)", unit), + INTERVAL_YEAR_MONTH(), + std::optional(months)) + .value(); + }; + + EXPECT_EQ(3, extract("year", months)); + EXPECT_EQ(4, extract("month", months)); +} + TEST_F(DateTimeFunctionsTest, dateTrunc) { const auto dateTrunc = [&](const std::string& unit, std::optional timestamp) {