diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h index cf68fa924215..b437a3eeca62 100644 --- a/velox/functions/prestosql/DateTimeFunctions.h +++ b/velox/functions/prestosql/DateTimeFunctions.h @@ -1902,6 +1902,38 @@ struct AtTimezoneFunction : public TimestampWithTimezoneSupport { } }; +template +struct AtTimezoneTimeFunction : public TimestampWithTimezoneSupport { + VELOX_DEFINE_FUNCTION_TYPES(T); + + std::optional targetTimezoneID_; + + FOLLY_ALWAYS_INLINE void initialize( + const std::vector& /*inputTypes*/, + const core::QueryConfig& config, + const arg_type* /*timeWithTz*/, + const arg_type* timezone) { + if (timezone) { + targetTimezoneID_ = tz::getTimeZoneID( + std::string_view(timezone->data(), timezone->size())); + } + } + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& timeWithTz, + const arg_type& timezone) { + const auto inputMs = unpackMillisUtc(*timeWithTz); + const auto targetTimezoneID = targetTimezoneID_.has_value() + ? targetTimezoneID_.value() + : tz::getTimeZoneID(std::string_view(timezone.data(), timezone.size())); + + // Similar to timestamp version - only timezone ID changes, not the time + // value. + result = pack(inputMs, targetTimezoneID); + } +}; + template struct ToMillisecondFunction { VELOX_DEFINE_FUNCTION_TYPES(TExec); diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp index 870aa8f7b9d0..3d3b3973d227 100644 --- a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp @@ -17,6 +17,7 @@ #include "velox/functions/Registerer.h" #include "velox/functions/prestosql/DateTimeFunctions.h" #include "velox/functions/prestosql/types/TimeWithTimezoneRegistration.h" +#include "velox/functions/prestosql/types/TimeWithTimezoneType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneRegistration.h" namespace facebook::velox::functions { @@ -318,6 +319,12 @@ void registerSimpleFunctions(const std::string& prefix) { TimestampWithTimezone, Varchar>({prefix + "at_timezone"}); + registerFunction< + AtTimezoneTimeFunction, + TimeWithTimezone, + TimeWithTimezone, + Varchar>({prefix + "at_timezone"}); + registerFunction( {prefix + "to_milliseconds"}); diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp index 3305ef3f2354..2b23a9fb5095 100644 --- a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp +++ b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp @@ -20,6 +20,7 @@ #include "velox/common/base/tests/GTestUtils.h" #include "velox/external/tzdb/zoned_time.h" #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" +#include "velox/functions/prestosql/types/TimeWithTimezoneType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/type/tz/TimeZoneMap.h" @@ -6135,6 +6136,123 @@ TEST_F(DateTimeFunctionsTest, atTimezoneTest) { EXPECT_EQ(at_timezone(std::nullopt, "Pacific/Fiji"), std::nullopt); } +TEST_F(DateTimeFunctionsTest, atTimezoneTimeTest) { + const auto at_timezone_time = [&](std::optional timeWithTimezone, + std::optional targetTimezone) { + return evaluateOnce( + "at_timezone(c0, c1)", + {TIME_WITH_TIME_ZONE(), VARCHAR()}, + timeWithTimezone, + targetTimezone); + }; + + const int64_t millisTenOClockWarsawWinter = 9 * 60 * 60 * 1000; + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw")), + "UTC"), + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("UTC"))); + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw")), + "+00:45"), + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("+00:45"))); + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw")), + "America/New_York"), + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("America/New_York"))); + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw")), + "Europe/Berlin"), + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Berlin"))); + + const int64_t millisTenOClockUTC = 10 * 60 * 60 * 1000; + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockUTC, tz::getTimeZoneID("UTC")), "UTC"), + pack(millisTenOClockUTC, tz::getTimeZoneID("UTC"))); + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw")), + "Europe/Warsaw"), + pack(millisTenOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw"))); + + const int64_t millisTwoOClockWarsawWinter = 1 * 60 * 60 * 1000; + const int64_t millisTwentyOClockNewYork = 20 * 60 * 60 * 1000; + EXPECT_EQ( + at_timezone_time( + pack(millisTwoOClockWarsawWinter, tz::getTimeZoneID("Europe/Warsaw")), + "America/New_York"), + pack(millisTwentyOClockNewYork, tz::getTimeZoneID("America/New_York"))); + + const int64_t millisTwentyTwoOClockNewYork = 22 * 60 * 60 * 1000; + const int64_t millisFourOClockWarsaw = 4 * 60 * 60 * 1000; + EXPECT_EQ( + at_timezone_time( + pack( + millisTwentyTwoOClockNewYork, + tz::getTimeZoneID("America/New_York")), + "Europe/Warsaw"), + pack(millisFourOClockWarsaw, tz::getTimeZoneID("Europe/Warsaw"))); + + const int64_t millisMidnight = 0; + const int64_t millisTwentyThreeOClock = 23 * 60 * 60 * 1000; + EXPECT_EQ( + at_timezone_time( + pack(millisMidnight, tz::getTimeZoneID("+14:00")), "+13:00"), + pack(millisTwentyThreeOClock, tz::getTimeZoneID("+13:00"))); + + const int64_t millisTwentyOClock = 20 * 60 * 60 * 1000; + EXPECT_EQ( + at_timezone_time( + pack(millisMidnight, tz::getTimeZoneID("+14:00")), "-14:00"), + pack(millisTwentyOClock, tz::getTimeZoneID("-14:00"))); + + const int64_t millisMaxTime = + 23 * 60 * 60 * 1000 + 59 * 60 * 1000 + 59 * 1000 + 999; + const int64_t millisTwentyTwoFiftyNine = + 22 * 60 * 60 * 1000 + 59 * 60 * 1000 + 59 * 1000 + 999; + EXPECT_EQ( + at_timezone_time( + pack(millisMaxTime, tz::getTimeZoneID("+14:00")), "+13:00"), + pack(millisTwentyTwoFiftyNine, tz::getTimeZoneID("+13:00"))); + + const int64_t millisNineteenFiftyNine = + 19 * 60 * 60 * 1000 + 59 * 60 * 1000 + 59 * 1000 + 999; + EXPECT_EQ( + at_timezone_time( + pack(millisMaxTime, tz::getTimeZoneID("+14:00")), "-14:00"), + pack(millisNineteenFiftyNine, tz::getTimeZoneID("-14:00"))); + + const int64_t millisTenOClockKathmandu = 10 * 60 * 60 * 1000; + const int64_t millisFourFifteenUTC = 4 * 60 * 60 * 1000 + 15 * 60 * 1000; + EXPECT_EQ( + at_timezone_time( + pack(millisFourFifteenUTC, tz::getTimeZoneID("Asia/Kathmandu")), + "UTC"), + pack(millisFourFifteenUTC, tz::getTimeZoneID("UTC"))); + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockKathmandu, tz::getTimeZoneID("Asia/Kabul")), + "Asia/Kabul"), + pack(millisTenOClockKathmandu, tz::getTimeZoneID("Asia/Kabul"))); + + EXPECT_EQ( + at_timezone_time( + pack(millisTenOClockUTC, tz::getTimeZoneID("UTC")), std::nullopt), + std::nullopt); + + EXPECT_EQ(at_timezone_time(std::nullopt, "UTC"), std::nullopt); +} + TEST_F(DateTimeFunctionsTest, toMilliseconds) { EXPECT_EQ( 123,