custom comparisons array min max

kevinwilfong · Oct 1, 2024 · 80bd686 · 80bd686
1 parent 66c23ef
commit 80bd686
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 0 deletions.
diff --git a/velox/functions/prestosql/tests/ArrayMaxTest.cpp b/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,62 @@ 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.
+    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);
+}
+
 // Test documented example.
 TEST_F(ArrayMaxTest, docs) {
   {

diff --git a/velox/functions/prestosql/tests/ArrayMinTest.cpp b/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,59 @@ 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.
+    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);
+}