Provide custom comparison functions in TimestampWithTimezone

Summary:
#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.

Differential Revision: D62906383

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<const CastOperator>& 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<const TimestampWithTimeZoneType>& 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<int64_t>()(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

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