From 9e1280a7223fa324b8307d379605bcfc89f8447a Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Wed, 25 Sep 2024 17:44:42 -0700 Subject: [PATCH] Provide custom comparison functions in TimestampWithTimezone (#11025) Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/11025 https://github.com/facebookincubator/velox/pull/11021 introduced the ability for custom types to provide their own compare and hash semantics. This change updates TimestampWithTimezone to take advantage of this new feature. It is represented as a 64-bit integer, but only the top 42 bits of the value, which represent the millis since epoch in UTC, should be used for the purposes of comparison and hashing. The bottom 12 bits which represent the timezone should be ignored. Reviewed By: pansatadru Differential Revision: D62906383 fbshipit-source-id: 93c5c4e9a1f8bf716d3b804f3bd7019d3b2b0ad3 --- .../types/TimestampWithTimeZoneType.h | 65 +++++++++++------- .../tests/TimestampWithTimeZoneTypeTest.cpp | 67 +++++++++++++++++++ 2 files changed, 106 insertions(+), 26 deletions(-) diff --git a/velox/functions/prestosql/types/TimestampWithTimeZoneType.h b/velox/functions/prestosql/types/TimestampWithTimeZoneType.h index 78b56483e862..95d711232c77 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, TimeZoneKey timeZoneKey) { + return (millisUtc << kMillisShift) | (timeZoneKey & kTimezoneMask); +} + +inline int64_t pack(const Timestamp& timestamp, TimeZoneKey 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..65a85fb0d2c3 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,70 @@ 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(0, -1639426440000, "+01:00", -1639426440000, "+03: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", -1539426440000, "+03:00"); + compare(-1, -1639426440000, "+01: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"); + compare(1, 1639426440000, "+01:00", -1639426440000, "+03:00"); + compare(1, -1539426440000, "+01:00", -1639426440000, "-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"); + 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"); + expectHashesNeq(-1639426440000, "+03:00", 1639426440000, "-14:00"); +} + } // namespace facebook::velox::test