From 2bcdf6d8988978629d6a777fc069d158a41f5d2c Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Tue, 1 Oct 2024 12:02:45 -0700 Subject: [PATCH] Add tests for array_sort with TimestampWithTImezone Summary: With https://github.com/facebookincubator/velox/pull/11136 array_sort just works with TimestampWithTimezone. Adding unit tests to verify this and ensure it doesn't break in the future. Differential Revision: D63711984 --- .../prestosql/tests/ArraySortTest.cpp | 111 ++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/velox/functions/prestosql/tests/ArraySortTest.cpp b/velox/functions/prestosql/tests/ArraySortTest.cpp index 574b3d5913024..f668304aafade 100644 --- a/velox/functions/prestosql/tests/ArraySortTest.cpp +++ b/velox/functions/prestosql/tests/ArraySortTest.cpp @@ -14,7 +14,10 @@ * limitations under the License. */ #include "velox/common/base/tests/GTestUtils.h" +#include "velox/functions/Macros.h" +#include "velox/functions/Registerer.h" #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include #include @@ -722,6 +725,114 @@ TEST_F(ArraySortTest, failOnRowNullCompare) { } } +TEST_F(ArraySortTest, timestampWithTimezone) { + auto data = makeRowVector({makeArrayVector( + {0, 3, 6, 9, 12, 16, 21, 26, 31}, + makeNullableFlatVector( + {pack(2, 0), pack(1, 1), pack(0, 2), pack(0, 0), pack(1, 1), + pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(1, 0), + pack(0, 1), pack(2, 2), std::nullopt, pack(1, 0), pack(0, 1), + pack(2, 2), std::nullopt, std::nullopt, pack(1, 2), pack(0, 1), + pack(2, 0), std::nullopt, pack(1, 1), pack(0, 2), std::nullopt, + pack(2, 0), pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2), + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE()))}); + + auto expectedAsc = makeArrayVector( + {0, 3, 6, 9, 12, 16, 21, 26, 31}, + makeNullableFlatVector( + {pack(0, 2), pack(1, 1), pack(2, 0), pack(0, 0), pack(1, 1), + pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(0, 1), + pack(1, 0), pack(2, 2), pack(0, 1), pack(1, 0), pack(2, 2), + std::nullopt, pack(0, 1), pack(1, 2), pack(2, 0), std::nullopt, + std::nullopt, pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt, + std::nullopt, pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())); + + auto expectedDesc = makeArrayVector( + {0, 3, 6, 9, 12, 16, 21, 26, 31}, + makeNullableFlatVector( + {pack(2, 0), pack(1, 1), pack(0, 2), pack(2, 2), pack(1, 1), + pack(0, 0), pack(0, 0), pack(0, 1), pack(0, 2), pack(2, 2), + pack(1, 0), pack(0, 1), pack(2, 2), pack(1, 0), pack(0, 1), + std::nullopt, pack(2, 0), pack(1, 2), pack(0, 1), std::nullopt, + std::nullopt, pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt, + std::nullopt, pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())); + + auto resultAsc = evaluate("array_sort(c0)", data); + assertEqualVectors(expectedAsc, resultAsc); + + auto resultDesc = evaluate("array_sort_desc(c0)", data); + assertEqualVectors(expectedDesc, resultDesc); +} + +template +struct TimeZoneFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + int64_t& result, + const arg_type& ts) { + result = unpackZoneKeyId(*ts); + } +}; + +TEST_F(ArraySortTest, timestampWithTimezoneWithLambda) { + registerFunction( + {"timezone"}); + + auto data = makeRowVector({makeArrayVector( + {0, 3, 6, 9, 12, 16, 21, 26, 31}, + makeNullableFlatVector( + {pack(2, 0), pack(1, 1), pack(0, 2), pack(0, 0), pack(1, 1), + pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(1, 0), + pack(0, 1), pack(2, 2), std::nullopt, pack(1, 0), pack(0, 1), + pack(2, 2), std::nullopt, std::nullopt, pack(1, 2), pack(0, 1), + pack(2, 0), std::nullopt, pack(1, 1), pack(0, 2), std::nullopt, + pack(2, 0), pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2), + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE()))}); + + auto expectedAsc = makeArrayVector( + {0, 3, 6, 9, 12, 16, 21, 26, 31}, + makeNullableFlatVector( + {pack(2, 0), pack(1, 1), pack(0, 2), pack(0, 0), pack(1, 1), + pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(1, 0), + pack(0, 1), pack(2, 2), pack(1, 0), pack(0, 1), pack(2, 2), + std::nullopt, pack(2, 0), pack(0, 1), pack(1, 2), std::nullopt, + std::nullopt, pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt, + std::nullopt, pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())); + + auto expectedDesc = makeArrayVector( + {0, 3, 6, 9, 12, 16, 21, 26, 31}, + makeNullableFlatVector( + {pack(0, 2), pack(1, 1), pack(2, 0), pack(2, 2), pack(1, 1), + pack(0, 0), pack(0, 2), pack(0, 1), pack(0, 0), pack(2, 2), + pack(0, 1), pack(1, 0), pack(2, 2), pack(0, 1), pack(1, 0), + std::nullopt, pack(1, 2), pack(0, 1), pack(2, 0), std::nullopt, + std::nullopt, pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt, + std::nullopt, pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())); + + auto resultAsc = evaluate("array_sort(c0, x -> timezone(x))", data); + assertEqualVectors(expectedAsc, resultAsc); + + auto resultDesc = evaluate("array_sort_desc(c0, x -> timezone(x))", data); + assertEqualVectors(expectedDesc, resultDesc); +} + TEST_F(ArraySortTest, floatingPointExtremes) { testFloatingPoint(); testFloatingPoint();