From 0403db4cabffc05b956e24b581c7bb36abdb8e1a Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Mon, 7 Oct 2024 12:14:50 -0700 Subject: [PATCH 1/2] Add tests for array_min/array_max with TimestampWithTimezone (#11141) Summary: With https://github.com/facebookincubator/velox/pull/11136 array_min and array_max just work with TimestampWithTimezone. Adding unit tests to verify this and ensure it doesn't break in the future. Differential Revision: D63708110 --- .../prestosql/tests/ArrayMaxTest.cpp | 61 +++++++++++++++++++ .../prestosql/tests/ArrayMinTest.cpp | 61 +++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/velox/functions/prestosql/tests/ArrayMaxTest.cpp b/velox/functions/prestosql/tests/ArrayMaxTest.cpp index 2d7bd98a8eac..6f2e15ad4dfa 100644 --- a/velox/functions/prestosql/tests/ArrayMaxTest.cpp +++ b/velox/functions/prestosql/tests/ArrayMaxTest.cpp @@ -16,6 +16,7 @@ #include #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" using namespace facebook::velox; using namespace facebook::velox::test; @@ -143,6 +144,66 @@ TEST_F(ArrayMaxTest, longVarcharNoNulls) { testArrayMax(input, expected); } +TEST_F(ArrayMaxTest, timestampWithTimezone) { + auto testArrayMax = [this]( + const std::vector>& inputArray, + std::optional expectedValue) { + // Test with primitive types. + auto input = makeRowVector({makeArrayVector( + {0}, makeNullableFlatVector(inputArray, TIMESTAMP_WITH_TIME_ZONE()))}); + VectorPtr expected = makeNullableFlatVector( + {expectedValue}, TIMESTAMP_WITH_TIME_ZONE()); + + auto result = evaluate("array_max(C0)", input); + assertEqualVectors(expected, result); + + // array_max does not support nulls inside complex types. To exclude these + // we exclude tests where the expected result is null. However, the result + // can also be null when the input is empty, so we add an exception for that + // case. + if (expectedValue.has_value() || inputArray.empty()) { + // Test wrapped in complex type. + input = makeRowVector({makeArrayVector( + {0}, + makeRowVector({makeNullableFlatVector( + inputArray, TIMESTAMP_WITH_TIME_ZONE())}))}); + expected = makeRowVector( + {expected}, [&](vector_size_t) { return inputArray.empty(); }); + + result = evaluate("array_max(C0)", input); + assertEqualVectors(expected, result); + } + }; + + testArrayMax( + {pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4), pack(4, 5)}, + pack(4, 5)); + testArrayMax( + {pack(4, 0), + pack(3, 1), + pack(2, 2), + pack(1, 3), + pack(0, 4), + pack(-1, 5), + pack(-2, 6)}, + pack(4, 0)); + testArrayMax( + {pack(-5, 3), pack(-4, 2), pack(-3, 1), pack(-2, 0), pack(-1, 4)}, + pack(-1, 4)); + testArrayMax( + {pack(101, 4), pack(102, 0), pack(103, 1), pack(104, 2), pack(105, 3)}, + pack(105, 3)); + testArrayMax({}, std::nullopt); + testArrayMax( + {pack(101, 4), pack(102, 0), pack(103, 1), pack(104, 2), std::nullopt}, + std::nullopt); + testArrayMax( + {std::nullopt, pack(-1, 4), pack(-2, 5), pack(-3, 1), pack(-4, 0)}, + std::nullopt); + testArrayMax({std::nullopt}, std::nullopt); + testArrayMax({pack(1, 0), pack(1, 1), pack(1, 2)}, pack(1, 0)); +} + // Test documented example. TEST_F(ArrayMaxTest, docs) { { diff --git a/velox/functions/prestosql/tests/ArrayMinTest.cpp b/velox/functions/prestosql/tests/ArrayMinTest.cpp index 4b63f924ddbd..6b5636dd875c 100644 --- a/velox/functions/prestosql/tests/ArrayMinTest.cpp +++ b/velox/functions/prestosql/tests/ArrayMinTest.cpp @@ -17,6 +17,7 @@ #include #include "velox/common/base/tests/GTestUtils.h" #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" using namespace facebook::velox; using namespace facebook::velox::test; @@ -284,3 +285,63 @@ TEST_F(ArrayMinTest, complexTypeElements) { }); assertEqualVectors(expected, result); } + +TEST_F(ArrayMinTest, timestampWithTimezone) { + auto testArrayMin = [this]( + const std::vector>& inputArray, + std::optional expectedValue) { + // Test with primitive types. + auto input = makeRowVector({makeArrayVector( + {0}, makeNullableFlatVector(inputArray, TIMESTAMP_WITH_TIME_ZONE()))}); + VectorPtr expected = makeNullableFlatVector( + {expectedValue}, TIMESTAMP_WITH_TIME_ZONE()); + + auto result = evaluate("array_min(C0)", input); + assertEqualVectors(expected, result); + + // array_min does not support nulls inside complex types. To exclude these + // we exclude tests where the expected result is null. However, the result + // can also be null when the input is empty, so we add an exception for that + // case. + if (expectedValue.has_value() || inputArray.empty()) { + // Test wrapped in complex type. + input = makeRowVector({makeArrayVector( + {0}, + makeRowVector({makeNullableFlatVector( + inputArray, TIMESTAMP_WITH_TIME_ZONE())}))}); + expected = makeRowVector( + {expected}, [&](vector_size_t) { return inputArray.empty(); }); + + result = evaluate("array_min(C0)", input); + assertEqualVectors(expected, result); + } + }; + + testArrayMin( + {pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4), pack(4, 5)}, + pack(-1, 0)); + testArrayMin( + {pack(4, 0), + pack(3, 1), + pack(2, 2), + pack(1, 3), + pack(0, 4), + pack(-1, 5), + pack(-2, 6)}, + pack(-2, 6)); + testArrayMin( + {pack(-5, 3), pack(-4, 2), pack(-3, 1), pack(-2, 0), pack(-1, 4)}, + pack(-5, 3)); + testArrayMin( + {pack(101, 4), pack(102, 0), pack(103, 1), pack(104, 2), pack(105, 3)}, + pack(101, 4)); + testArrayMin({}, std::nullopt); + testArrayMin( + {pack(101, 4), pack(102, 0), pack(103, 1), pack(104, 2), std::nullopt}, + std::nullopt); + testArrayMin( + {std::nullopt, pack(-1, 4), pack(-2, 5), pack(-3, 1), pack(-4, 0)}, + std::nullopt); + testArrayMin({std::nullopt}, std::nullopt); + testArrayMin({pack(1, 0), pack(1, 1), pack(1, 2)}, pack(1, 0)); +} From 8f06c64fd982205184be621ba1b22d309071a710 Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Mon, 7 Oct 2024 12:14:50 -0700 Subject: [PATCH 2/2] Add tests for array_sort with TimestampWithTImezone (#11142) 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. Reviewed By: bikramSingh91 Differential Revision: D63711984 --- .../prestosql/tests/ArraySortTest.cpp | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/velox/functions/prestosql/tests/ArraySortTest.cpp b/velox/functions/prestosql/tests/ArraySortTest.cpp index 574b3d591302..2c492e35b467 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,149 @@ TEST_F(ArraySortTest, failOnRowNullCompare) { } } +TEST_F(ArraySortTest, timestampWithTimezone) { + auto testArraySort = + [this]( + const std::vector>& inputArray, + const std::vector>& expectedAscArray, + const std::vector>& expectedDescArray) { + const auto input = makeRowVector({makeArrayVector( + {0}, + makeNullableFlatVector( + inputArray, TIMESTAMP_WITH_TIME_ZONE()))}); + const auto expectedAsc = makeArrayVector( + {0}, + makeNullableFlatVector( + expectedAscArray, TIMESTAMP_WITH_TIME_ZONE())); + const auto expectedDesc = makeArrayVector( + {0}, + makeNullableFlatVector( + expectedDescArray, TIMESTAMP_WITH_TIME_ZONE())); + + auto resultAsc = evaluate("array_sort(c0)", input); + assertEqualVectors(expectedAsc, resultAsc); + + auto resultDesc = evaluate("array_sort_desc(c0)", input); + assertEqualVectors(expectedDesc, resultDesc); + }; + + testArraySort( + {pack(2, 0), pack(1, 1), pack(0, 2)}, + {pack(0, 2), pack(1, 1), pack(2, 0)}, + {pack(2, 0), pack(1, 1), pack(0, 2)}); + testArraySort( + {pack(0, 0), pack(1, 1), pack(2, 2)}, + {pack(0, 0), pack(1, 1), pack(2, 2)}, + {pack(2, 2), pack(1, 1), pack(0, 0)}); + testArraySort( + {pack(0, 0), pack(0, 1), pack(0, 2)}, + {pack(0, 0), pack(0, 1), pack(0, 2)}, + {pack(0, 0), pack(0, 1), pack(0, 2)}); + testArraySort( + {pack(1, 0), pack(0, 1), pack(2, 2)}, + {pack(0, 1), pack(1, 0), pack(2, 2)}, + {pack(2, 2), pack(1, 0), pack(0, 1)}); + testArraySort( + {std::nullopt, pack(1, 0), pack(0, 1), pack(2, 2)}, + {pack(0, 1), pack(1, 0), pack(2, 2), std::nullopt}, + {pack(2, 2), pack(1, 0), pack(0, 1), std::nullopt}); + testArraySort( + {std::nullopt, std::nullopt, pack(1, 2), pack(0, 1), pack(2, 0)}, + {pack(0, 1), pack(1, 2), pack(2, 0), std::nullopt, std::nullopt}, + {pack(2, 0), pack(1, 2), pack(0, 1), std::nullopt, std::nullopt}); + testArraySort( + {std::nullopt, pack(1, 1), pack(0, 2), std::nullopt, pack(2, 0)}, + {pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt, std::nullopt}, + {pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt, std::nullopt}); + testArraySort( + {pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2), std::nullopt}, + {pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt, std::nullopt}, + {pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt, std::nullopt}); + testArraySort( + {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt}, + {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt}, + {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt}); +} + +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 testArraySort = + [this]( + const std::vector>& inputArray, + const std::vector>& expectedAscArray, + const std::vector>& expectedDescArray) { + const auto input = makeRowVector({makeArrayVector( + {0}, + makeNullableFlatVector( + inputArray, TIMESTAMP_WITH_TIME_ZONE()))}); + const auto expectedAsc = makeArrayVector( + {0}, + makeNullableFlatVector( + expectedAscArray, TIMESTAMP_WITH_TIME_ZONE())); + const auto expectedDesc = makeArrayVector( + {0}, + makeNullableFlatVector( + expectedDescArray, TIMESTAMP_WITH_TIME_ZONE())); + + auto resultAsc = evaluate("array_sort(c0, x -> timezone(x))", input); + assertEqualVectors(expectedAsc, resultAsc); + + auto resultDesc = + evaluate("array_sort_desc(c0, x -> timezone(x))", input); + assertEqualVectors(expectedDesc, resultDesc); + }; + + testArraySort( + {pack(2, 0), pack(1, 1), pack(0, 2)}, + {pack(2, 0), pack(1, 1), pack(0, 2)}, + {pack(0, 2), pack(1, 1), pack(2, 0)}); + testArraySort( + {pack(0, 0), pack(1, 1), pack(2, 2)}, + {pack(0, 0), pack(1, 1), pack(2, 2)}, + {pack(2, 2), pack(1, 1), pack(0, 0)}); + testArraySort( + {pack(0, 0), pack(0, 1), pack(0, 2)}, + {pack(0, 0), pack(0, 1), pack(0, 2)}, + {pack(0, 2), pack(0, 1), pack(0, 0)}); + testArraySort( + {pack(1, 0), pack(0, 1), pack(2, 2)}, + {pack(1, 0), pack(0, 1), pack(2, 2)}, + {pack(2, 2), pack(0, 1), pack(1, 0)}); + testArraySort( + {std::nullopt, pack(1, 0), pack(0, 1), pack(2, 2)}, + {pack(1, 0), pack(0, 1), pack(2, 2), std::nullopt}, + {pack(2, 2), pack(0, 1), pack(1, 0), std::nullopt}); + testArraySort( + {std::nullopt, std::nullopt, pack(1, 2), pack(0, 1), pack(2, 0)}, + {pack(2, 0), pack(0, 1), pack(1, 2), std::nullopt, std::nullopt}, + {pack(1, 2), pack(0, 1), pack(2, 0), std::nullopt, std::nullopt}); + testArraySort( + {std::nullopt, pack(1, 1), pack(0, 2), std::nullopt, pack(2, 0)}, + {pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt, std::nullopt}, + {pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt, std::nullopt}); + testArraySort( + {pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2), std::nullopt}, + {pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt, std::nullopt}, + {pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt, std::nullopt}); + testArraySort( + {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt}, + {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt}, + {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt}); +} + TEST_F(ArraySortTest, floatingPointExtremes) { testFloatingPoint(); testFloatingPoint();