From a34cad5f49025cce41e8b295ad99a5d24adadc91 Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Tue, 17 Sep 2024 15:55:25 -0700 Subject: [PATCH] custom compare timesteamp with timezone --- .../types/TimestampWithTimeZoneType.h | 65 +++++++++++-------- .../tests/TimestampWithTimeZoneTypeTest.cpp | 60 +++++++++++++++++ 2 files changed, 99 insertions(+), 26 deletions(-) diff --git a/velox/functions/prestosql/types/TimestampWithTimeZoneType.h b/velox/functions/prestosql/types/TimestampWithTimeZoneType.h index 78b56483e862..9e34092ed1d4 100644 --- a/velox/functions/prestosql/types/TimestampWithTimeZoneType.h +++ b/velox/functions/prestosql/types/TimestampWithTimeZoneType.h @@ -22,6 +22,31 @@ namespace facebook::velox { +using TimeZoneKey = int16_t; + +constexpr int32_t kTimezoneMask = 0xFFF; +constexpr int32_t kMillisShift = 12; + +inline int64_t unpackMillisUtc(int64_t dateTimeWithTimeZone) { + return dateTimeWithTimeZone >> kMillisShift; +} + +inline TimeZoneKey unpackZoneKeyId(int64_t dateTimeWithTimeZone) { + return dateTimeWithTimeZone & kTimezoneMask; +} + +inline int64_t pack(int64_t millisUtc, int16_t timeZoneKey) { + return (millisUtc << kMillisShift) | (timeZoneKey & kTimezoneMask); +} + +inline int64_t pack(const Timestamp& timestamp, int16_t timeZoneKey) { + return pack(timestamp.toMillis(), timeZoneKey); +} + +inline Timestamp unpackTimestampUtc(int64_t dateTimeWithTimeZone) { + return Timestamp::fromMillis(unpackMillisUtc(dateTimeWithTimeZone)); +} + class TimestampWithTimeZoneCastOperator : public exec::CastOperator { public: static const std::shared_ptr& get() { @@ -56,7 +81,7 @@ class TimestampWithTimeZoneCastOperator : public exec::CastOperator { /// Represents timestamp with time zone as a number of milliseconds since epoch /// and time zone ID. class TimestampWithTimeZoneType : public BigintType { - TimestampWithTimeZoneType() = default; + TimestampWithTimeZoneType() : BigintType(true) {} public: static const std::shared_ptr& get() { @@ -72,6 +97,19 @@ class TimestampWithTimeZoneType : public BigintType { return this == &other; } + int32_t compare(const int64_t* left, const int64_t* right) const override { + int64_t leftUnpacked = unpackMillisUtc(*left); + int64_t rightUnpacked = unpackMillisUtc(*right); + + return leftUnpacked < rightUnpacked ? -1 + : leftUnpacked == rightUnpacked ? 0 + : 1; + } + + uint64_t hash(const int64_t* value) const override { + return folly::hasher()(unpackMillisUtc(*value)); + } + const char* name() const override { return "TIMESTAMP WITH TIME ZONE"; } @@ -125,29 +163,4 @@ class TimestampWithTimeZoneTypeFactories : public CustomTypeFactories { void registerTimestampWithTimeZoneType(); -using TimeZoneKey = int16_t; - -constexpr int32_t kTimezoneMask = 0xFFF; -constexpr int32_t kMillisShift = 12; - -inline int64_t unpackMillisUtc(int64_t dateTimeWithTimeZone) { - return dateTimeWithTimeZone >> kMillisShift; -} - -inline TimeZoneKey unpackZoneKeyId(int64_t dateTimeWithTimeZone) { - return dateTimeWithTimeZone & kTimezoneMask; -} - -inline int64_t pack(int64_t millisUtc, int16_t timeZoneKey) { - return (millisUtc << kMillisShift) | (timeZoneKey & kTimezoneMask); -} - -inline int64_t pack(const Timestamp& timestamp, int16_t timeZoneKey) { - return pack(timestamp.toMillis(), timeZoneKey); -} - -inline Timestamp unpackTimestampUtc(int64_t dateTimeWithTimeZone) { - return Timestamp::fromMillis(unpackMillisUtc(dateTimeWithTimeZone)); -} - } // namespace facebook::velox diff --git a/velox/functions/prestosql/types/tests/TimestampWithTimeZoneTypeTest.cpp b/velox/functions/prestosql/types/tests/TimestampWithTimeZoneTypeTest.cpp index 8cfa5f578b6e..86268bf292ba 100644 --- a/velox/functions/prestosql/types/tests/TimestampWithTimeZoneTypeTest.cpp +++ b/velox/functions/prestosql/types/tests/TimestampWithTimeZoneTypeTest.cpp @@ -15,6 +15,7 @@ */ #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/functions/prestosql/types/tests/TypeTestBase.h" +#include "velox/type/tz/TimeZoneMap.h" namespace facebook::velox::test { @@ -65,4 +66,63 @@ TEST_F(TimestampWithTimeZoneTypeTest, pack) { } } +TEST_F(TimestampWithTimeZoneTypeTest, compare) { + auto compare = [](int32_t expected, + int64_t millis1, + const std::string& tz1, + int64_t millis2, + const std::string& tz2) { + int64_t left = pack(millis1, tz::getTimeZoneID(tz1)); + int64_t right = pack(millis2, tz::getTimeZoneID(tz2)); + + ASSERT_EQ(expected, TIMESTAMP_WITH_TIME_ZONE()->compare(&left, &right)); + }; + + compare(0, 1639426440000, "+01:00", 1639426440000, "+03:00"); + compare(0, 1639426440000, "+01:00", 1639426440000, "-14:00"); + compare(0, 1639426440000, "+03:00", 1639426440000, "-14:00"); + + compare(-1, 1549770072000, "+01:00", 1639426440000, "+03:00"); + compare(-1, 1549770072000, "+01:00", 1639426440000, "-14:00"); + compare(-1, 1549770072000, "+03:00", 1639426440000, "-14:00"); + + compare(1, 1639426440000, "+01:00", 1549770072000, "+03:00"); + compare(1, 1639426440000, "+01:00", 1549770072000, "-14:00"); + compare(1, 1639426440000, "+03:00", 1549770072000, "-14:00"); +} + +TEST_F(TimestampWithTimeZoneTypeTest, hash) { + auto expectHashesEq = [](int64_t millis1, + const std::string& tz1, + int64_t millis2, + const std::string& tz2) { + int64_t left = pack(millis1, tz::getTimeZoneID(tz1)); + int64_t right = pack(millis2, tz::getTimeZoneID(tz2)); + + ASSERT_EQ( + TIMESTAMP_WITH_TIME_ZONE()->hash(&left), + TIMESTAMP_WITH_TIME_ZONE()->hash(&right)); + }; + + auto expectHashesNeq = [](int64_t millis1, + const std::string& tz1, + int64_t millis2, + const std::string& tz2) { + int64_t left = pack(millis1, tz::getTimeZoneID(tz1)); + int64_t right = pack(millis2, tz::getTimeZoneID(tz2)); + + ASSERT_NE( + TIMESTAMP_WITH_TIME_ZONE()->hash(&left), + TIMESTAMP_WITH_TIME_ZONE()->hash(&right)); + }; + + expectHashesEq(1639426440000, "+01:00", 1639426440000, "+03:00"); + expectHashesEq(1639426440000, "+01:00", 1639426440000, "-14:00"); + expectHashesEq(1639426440000, "+03:00", 1639426440000, "-14:00"); + + expectHashesNeq(1549770072000, "+01:00", 1639426440000, "+03:00"); + expectHashesNeq(1549770072000, "+01:00", 1639426440000, "-14:00"); + expectHashesNeq(1549770072000, "+03:00", 1639426440000, "-14:00"); +} + } // namespace facebook::velox::test