custom compare timesteamp with timezone

kevinwilfong · Sep 17, 2024 · a34cad5 · a34cad5
1 parent 88043f0
commit a34cad5
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 26 deletions.
diff --git 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<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
diff --git 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