diff --git a/velox/functions/prestosql/tests/ComparisonsTest.cpp b/velox/functions/prestosql/tests/ComparisonsTest.cpp index f05075121d05..38098b5628dd 100644 --- a/velox/functions/prestosql/tests/ComparisonsTest.cpp +++ b/velox/functions/prestosql/tests/ComparisonsTest.cpp @@ -19,6 +19,8 @@ #include "velox/functions/Udf.h" #include "velox/functions/lib/RegistrationHelpers.h" #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" +#include "velox/type/tests/utils/CustomTypesForTesting.h" +#include "velox/type/tz/TimeZoneMap.h" using namespace facebook::velox; @@ -738,6 +740,259 @@ TEST_F(ComparisonsTest, nanComparison) { testNaN("", input, false); } +TEST_F(ComparisonsTest, TimestampWithTimezone) { + auto makeTimestampWithTimezone = [](int64_t millis, const std::string& tz) { + return pack(millis, tz::getTimeZoneID(tz)); + }; + + auto lhs = makeFlatVector( + {makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(1549770072000, "+01:00"), + makeTimestampWithTimezone(1549770072000, "+01:00"), + makeTimestampWithTimezone(1549770072000, "+03:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(-1639426440000, "+01:00"), + makeTimestampWithTimezone(-1539426440000, "+03:00"), + makeTimestampWithTimezone(-1639426440000, "-14:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(-1639426440000, "-14:00")}, + TIMESTAMP_WITH_TIME_ZONE()); + + auto rhs = makeFlatVector( + {makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(1639426440000, "-14:00"), + makeTimestampWithTimezone(1639426440000, "-14:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(1639426440000, "-14:00"), + makeTimestampWithTimezone(1639426440000, "-14:00"), + makeTimestampWithTimezone(1549770072000, "+03:00"), + makeTimestampWithTimezone(1549770072000, "-14:00"), + makeTimestampWithTimezone(1549770072000, "-14:00"), + makeTimestampWithTimezone(-1639426440000, "+01:00"), + makeTimestampWithTimezone(-1639426440000, "-14:00"), + makeTimestampWithTimezone(-1539426440000, "+03:00"), + makeTimestampWithTimezone(-1639426440000, "+03:00"), + makeTimestampWithTimezone(1639426440000, "+01:00")}, + TIMESTAMP_WITH_TIME_ZONE()); + + auto input = makeRowVector({lhs, rhs}); + + auto eval = [&](const std::string& expr) { + return evaluate>(fmt::format("c0 {} c1", expr), input); + }; + + test::assertEqualVectors( + eval("="), + makeFlatVector( + {true, + true, + true, + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false})); + test::assertEqualVectors( + eval("<>"), + makeFlatVector( + {false, + false, + false, + true, + true, + true, + true, + true, + true, + false, + true, + true, + true, + true})); + test::assertEqualVectors( + eval("<"), + makeFlatVector( + {false, + false, + false, + true, + true, + true, + false, + false, + false, + false, + false, + true, + false, + true})); + test::assertEqualVectors( + eval(">"), + makeFlatVector( + {false, + false, + false, + false, + false, + false, + true, + true, + true, + false, + true, + false, + true, + false})); + test::assertEqualVectors( + eval("<="), + makeFlatVector( + {true, + true, + true, + true, + true, + true, + false, + false, + false, + true, + false, + true, + false, + true})); + test::assertEqualVectors( + eval(">="), + makeFlatVector( + {true, + true, + true, + false, + false, + false, + true, + true, + true, + true, + true, + false, + true, + false})); + test::assertEqualVectors( + eval("is distinct from"), + makeFlatVector( + {false, + false, + false, + true, + true, + true, + true, + true, + true, + false, + true, + true, + true, + true})); + + auto betweenInput = makeRowVector({ + makeFlatVector( + {makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(1539426440000, "+03:00"), + makeTimestampWithTimezone(1739426440000, "+01:00"), + makeTimestampWithTimezone(-1639426440000, "+01:00"), + makeTimestampWithTimezone(-1639426440000, "+03:00"), + makeTimestampWithTimezone(1639426440000, "-14:00"), + makeTimestampWithTimezone(-1739426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "-14:00")}, + TIMESTAMP_WITH_TIME_ZONE()), + makeFlatVector( + {makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(1539426440000, "+01:00"), + makeTimestampWithTimezone(1539426440000, "+03:00"), + makeTimestampWithTimezone(1539426440000, "+03:00"), + makeTimestampWithTimezone(1539426440000, "+03:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1539426440000, "+03:00"), + makeTimestampWithTimezone(-1739426440000, "+01:00"), + makeTimestampWithTimezone(-1739426440000, "+01:00"), + makeTimestampWithTimezone(-1639426440000, "+01:00"), + makeTimestampWithTimezone(-1639426440000, "+01:00"), + makeTimestampWithTimezone(-1639426440000, "+01:00")}, + TIMESTAMP_WITH_TIME_ZONE()), + makeFlatVector( + {makeTimestampWithTimezone(1739426440000, "+01:00"), + makeTimestampWithTimezone(1739426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+03:00"), + makeTimestampWithTimezone(1739426440000, "+03:00"), + makeTimestampWithTimezone(1739426440000, "+03:00"), + makeTimestampWithTimezone(1739426440000, "+01:00"), + makeTimestampWithTimezone(1639426440000, "+01:00"), + makeTimestampWithTimezone(-1539426440000, "+03:00"), + makeTimestampWithTimezone(1539426440000, "+03:00"), + makeTimestampWithTimezone(1739426440000, "+01:00"), + makeTimestampWithTimezone(-1539426440000, "+03:00"), + makeTimestampWithTimezone(-1539426440000, "+03:00")}, + TIMESTAMP_WITH_TIME_ZONE()), + }); + + test::assertEqualVectors( + evaluate>( + fmt::format("c0 between c1 and c2"), betweenInput), + makeFlatVector( + {true, + true, + true, + true, + true, + true, + false, + false, + true, + true, + true, + false, + false})); +} + +TEST_F(ComparisonsTest, CustomComparisonWithGenerics) { + // Tests that functions that support signatures with generics handle custom + // comparison correctly. + auto input = makeRowVector({ + makeFlatVector( + {0, 1}, test::BIGINT_TYPE_WITH_CUSTOM_COMPARISON()), + makeFlatVector( + {256, 258}, test::BIGINT_TYPE_WITH_CUSTOM_COMPARISON()), + }); + + test::assertEqualVectors( + evaluate>("c0 = c1", input), + makeFlatVector({true, false})); + test::assertEqualVectors( + evaluate>("c0 <> c1", input), + makeFlatVector({false, true})); + test::assertEqualVectors( + evaluate>("c0 is distinct from c1", input), + makeFlatVector({false, true})); +} + namespace { template struct ComparisonTypeOp {