From 6fd0a892bdee5ad6eabf7438b8bae42fada9930a Mon Sep 17 00:00:00 2001
From: Kevin Wilfong <kevinwilfong@meta.com>
Date: Thu, 10 Oct 2024 16:58:33 -0700
Subject: [PATCH] Add tests for array_union and array_remove with
 TimestampWithTimezone (#11222)

Summary:

With https://github.com/facebookincubator/velox/pull/11136 array_union and
array_remove just work with TimestampWithTimezone.

Adding unit tests to verify this and ensure it doesn't break in the future.

Differential Revision: D64196032
---
 .../prestosql/tests/ArrayRemoveTest.cpp       | 51 ++++++++++++
 .../prestosql/tests/ArrayUnionTest.cpp        | 77 +++++++++++++++++++
 2 files changed, 128 insertions(+)

diff --git a/velox/functions/prestosql/tests/ArrayRemoveTest.cpp b/velox/functions/prestosql/tests/ArrayRemoveTest.cpp
index 20920bdc83d7..f5cb6b851e6d 100644
--- a/velox/functions/prestosql/tests/ArrayRemoveTest.cpp
+++ b/velox/functions/prestosql/tests/ArrayRemoveTest.cpp
@@ -15,6 +15,7 @@
  */
 #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;
@@ -207,4 +208,54 @@ TEST_F(ArrayRemoveTest, arrayWithNull) {
   testExpression(
       "array_remove(c0, c1)", {arrayVector, elementVector}, expected);
 }
+
+TEST_F(ArrayRemoveTest, timestampWithTimeZone) {
+  const auto arrayVector = makeArrayVector(
+      {0, 6, 13, 16},
+      makeNullableFlatVector<int64_t>(
+          {pack(1, 1),
+           std::nullopt,
+           pack(2, 2),
+           pack(3, 3),
+           std::nullopt,
+           pack(4, 4),
+           pack(3, 5),
+           pack(4, 6),
+           pack(5, 7),
+           std::nullopt,
+           pack(3, 8),
+           pack(4, 9),
+           pack(5, 10),
+           pack(7, 11),
+           pack(8, 12),
+           pack(9, 13),
+           pack(10, 14),
+           pack(20, 15),
+           pack(30, 16)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  const auto elementVector = makeFlatVector<int64_t>(
+      {pack(3, 3), pack(3, 3), pack(33, 1), pack(30, 1)},
+      TIMESTAMP_WITH_TIME_ZONE());
+  const auto expected = makeArrayVector(
+      {0, 5, 10, 13},
+      makeNullableFlatVector<int64_t>(
+          {pack(1, 1),
+           std::nullopt,
+           pack(2, 2),
+           std::nullopt,
+           pack(4, 4),
+           pack(4, 6),
+           pack(5, 7),
+           std::nullopt,
+           pack(4, 9),
+           pack(5, 10),
+           pack(7, 11),
+           pack(8, 12),
+           pack(9, 13),
+           pack(10, 14),
+           pack(20, 15)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  testExpression(
+      "array_remove(c0, c1)", {arrayVector, elementVector}, expected);
+}
 } // namespace
diff --git a/velox/functions/prestosql/tests/ArrayUnionTest.cpp b/velox/functions/prestosql/tests/ArrayUnionTest.cpp
index 9173f3725f54..3648d5543c0a 100644
--- a/velox/functions/prestosql/tests/ArrayUnionTest.cpp
+++ b/velox/functions/prestosql/tests/ArrayUnionTest.cpp
@@ -15,6 +15,7 @@
  */
 #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;
@@ -168,4 +169,80 @@ TEST_F(ArrayUnionTest, floatingPointType) {
   floatArrayTest<float>();
   floatArrayTest<double>();
 }
+
+TEST_F(ArrayUnionTest, timestampWithTimeZone) {
+  const auto array1 = makeArrayVector(
+      {0, 4, 7, 10},
+      makeFlatVector<int64_t>(
+          {pack(1, 1),
+           pack(2, 2),
+           pack(3, 3),
+           pack(4, 4),
+           pack(3, 5),
+           pack(4, 6),
+           pack(5, 7),
+           pack(7, 8),
+           pack(8, 9),
+           pack(9, 10),
+           pack(10, 11),
+           pack(20, 12),
+           pack(30, 13)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  const auto array2 = makeArrayVector(
+      {0, 3, 6, 6},
+      makeFlatVector<int64_t>(
+          {pack(2, 10),
+           pack(4, 11),
+           pack(5, 12),
+           pack(3, 13),
+           pack(4, 14),
+           pack(5, 15),
+           pack(40, 16),
+           pack(50, 17)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+
+  VectorPtr expected = makeArrayVector(
+      {0, 5, 8, 11},
+      makeFlatVector<int64_t>(
+          {pack(1, 1),
+           pack(2, 2),
+           pack(3, 3),
+           pack(4, 4),
+           pack(5, 12),
+           pack(3, 5),
+           pack(4, 6),
+           pack(5, 7),
+           pack(7, 8),
+           pack(8, 9),
+           pack(9, 10),
+           pack(10, 11),
+           pack(20, 12),
+           pack(30, 13),
+           pack(40, 16),
+           pack(50, 17)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  testExpression("array_union(c0, c1)", {array1, array2}, expected);
+
+  expected = makeArrayVector(
+      {0, 5, 8, 11},
+      makeFlatVector<int64_t>(
+          {pack(2, 10),
+           pack(4, 11),
+           pack(5, 12),
+           pack(1, 1),
+           pack(3, 3),
+           pack(3, 13),
+           pack(4, 14),
+           pack(5, 15),
+           pack(7, 8),
+           pack(8, 9),
+           pack(9, 10),
+           pack(40, 16),
+           pack(50, 17),
+           pack(10, 11),
+           pack(20, 12),
+           pack(30, 13)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  testExpression("array_union(c0, c1)", {array2, array1}, expected);
+}
 } // namespace