From c0fb0d36e1a30505ae1520debb5f437e7a8dc991 Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Tue, 17 Sep 2024 16:14:43 -0700 Subject: [PATCH] Add tests for Presto comparison functions for custom types with custom comparison functions (#11019) Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/11019 https://github.com/facebookincubator/velox/pull/11015 added support for custom types to provide custom comparison functions, motivated by the need to support TimestampWithTimezone's comparison semantics. I validated the UDFs provided in Comparisons.h on top of this change. It does not look like any further changes are needed in those files, but I added explicit tests for TimestampWithTimezone in the various comparison UDFs. I also added tests for the UDFs that support Generic arguments, to ensure that when the Type underlying the Generic provides custom comparison functions, that is respected. I considered adding checks to the UDFs that support templated arguments to ensure the type does not provide custom comparison functions, as these are not supported in any of them (e.g. between, <, >, and the SIMD versions of comparison functions). However, someone would have to explicitly register these functions for the custom type, and I don't think it's unreasonable to expect the user to test these functions work as expected when doing so, I don't think the per-batch check is worth the (admittedly small) overhead. Differential Revision: D62893143 --- .../prestosql/tests/ComparisonsTest.cpp | 130 ++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/velox/functions/prestosql/tests/ComparisonsTest.cpp b/velox/functions/prestosql/tests/ComparisonsTest.cpp index f05075121d05..5c56a4a93358 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,134 @@ 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")}, + 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")}, + 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})); + test::assertEqualVectors( + eval("<>"), + makeFlatVector( + {false, false, false, true, true, true, true, true, true})); + test::assertEqualVectors( + eval("<"), + makeFlatVector( + {false, false, false, true, true, true, false, false, false})); + test::assertEqualVectors( + eval(">"), + makeFlatVector( + {false, false, false, false, false, false, true, true, true})); + test::assertEqualVectors( + eval("<="), + makeFlatVector( + {true, true, true, true, true, true, false, false, false})); + test::assertEqualVectors( + eval(">="), + makeFlatVector( + {true, true, true, false, false, false, true, true, true})); + test::assertEqualVectors( + eval("is distinct from"), + makeFlatVector( + {false, false, false, true, true, 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")}, + 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"), + }, + 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"), + }, + TIMESTAMP_WITH_TIME_ZONE()), + }); + + test::assertEqualVectors( + evaluate>( + fmt::format("c0 between c1 and c2"), betweenInput), + makeFlatVector({true, true, true, 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 {