Skip to content

Commit

Permalink
Provide custom comparison functions in TimestampWithTimezone (#11025)
Browse files Browse the repository at this point in the history 
Summary:
Pull Request resolved: #11025

#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
  • Loading branch information
@facebook-github-bot
Kevin Wilfong authored and facebook-github-bot committed Sep 26, 2024
1 parent 50784a5 commit 9e1280a
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 26 deletions.
65 changes: 39 additions & 26 deletions velox/functions/prestosql/types/TimestampWithTimeZoneType.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<const CastOperator>& get() {
Expand Down Expand Up @@ -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() {
Expand All @@ -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";
}
Expand Down Expand Up @@ -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
67 changes: 67 additions & 0 deletions velox/functions/prestosql/types/tests/TimestampWithTimeZoneTypeTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand Down Expand Up @@ -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

0 comments on commit 9e1280a

Please sign in to comment.