Add tests for array_min/array_max with TimestampWithTimezone (#11141)

Summary:
Pull Request resolved: #11141

With #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.

Reviewed By: kgpai

Differential Revision: D63708110

fbshipit-source-id: 4dfd162c0914a5a2ec6304f6aaedd505a3921913

velox/functions/prestosql/tests/ArrayMaxTest.cpp

@@ -16,6 +16,7 @@
 
 #include <optional>
 #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<std::optional<int64_t>>& inputArray,
+                          std::optional<int64_t> expectedValue) {
+    // Test with primitive types.
+    auto input = makeRowVector({makeArrayVector(
+        {0}, makeNullableFlatVector(inputArray, TIMESTAMP_WITH_TIME_ZONE()))});
+    VectorPtr expected = makeNullableFlatVector<int64_t>(
+        {expectedValue}, TIMESTAMP_WITH_TIME_ZONE());
+
+    auto result = evaluate<BaseVector>("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<BaseVector>("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) {
   {

velox/functions/prestosql/tests/ArrayMinTest.cpp

@@ -17,6 +17,7 @@
 #include <optional>
 #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<std::optional<int64_t>>& inputArray,
+                          std::optional<int64_t> expectedValue) {
+    // Test with primitive types.
+    auto input = makeRowVector({makeArrayVector(
+        {0}, makeNullableFlatVector(inputArray, TIMESTAMP_WITH_TIME_ZONE()))});
+    VectorPtr expected = makeNullableFlatVector<int64_t>(
+        {expectedValue}, TIMESTAMP_WITH_TIME_ZONE());
+
+    auto result = evaluate<BaseVector>("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<BaseVector>("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));
+}