From 5bedca0f38a57b87f28cb235ac85f0bee3ed4967 Mon Sep 17 00:00:00 2001
From: Kevin Wilfong <kevinwilfong@meta.com>
Date: Fri, 11 Oct 2024 11:28:14 -0700
Subject: [PATCH] Add tests constructing maps with duplicate custom
 TimestampWithTimezones (#11228)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/11228

When constructing a map in Presto and checking for duplicate keys, it already uses
the key Vector's equalValuesAt function which calls the custom comparison operator
provided by the custom type (if that applies).  So this just works automatically for
types with custom comparison.

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

Reviewed By: xiaoxmeng

Differential Revision: D64214359

fbshipit-source-id: 15261dd7cf32693a199eba9b7efc43f47c39177b
---
 velox/functions/prestosql/tests/MapTest.cpp | 69 +++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/velox/functions/prestosql/tests/MapTest.cpp b/velox/functions/prestosql/tests/MapTest.cpp
index 4113b4c5d166..218f9e5877c9 100644
--- a/velox/functions/prestosql/tests/MapTest.cpp
+++ b/velox/functions/prestosql/tests/MapTest.cpp
@@ -18,6 +18,7 @@
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/functions/prestosql/registration/RegistrationFunctions.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;
@@ -497,4 +498,72 @@ TEST_F(MapTest, unknownType) {
       "map key cannot be null");
 }
 
+TEST_F(MapTest, timestampWithTimeZone) {
+  functions::prestosql::registerMapAllowingDuplicates("map2");
+
+  auto values = makeArrayVector<int32_t>({{1, 2, 3, 4, 5, 6}});
+  auto checkDuplicate = [&](const VectorPtr& keys, std::string expectedError) {
+    VELOX_ASSERT_THROW(
+        evaluate("map(c0, c1)", makeRowVector({keys, values})), expectedError);
+
+    ASSERT_NO_THROW(
+        evaluate("try(map(c0, c1))", makeRowVector({keys, values})));
+
+    // Trying the map version with allowing duplicates.
+    ASSERT_NO_THROW(evaluate("map2(c0, c1)", makeRowVector({keys, values})));
+  };
+
+  // Check for duplicate keys with identical timestamps.
+  const auto keysIdenticalTimestamps = makeArrayVector(
+      {0},
+      makeFlatVector<int64_t>(
+          {pack(1, 1),
+           pack(2, 2),
+           pack(3, 3),
+           pack(4, 4),
+           pack(5, 5),
+           pack(3, 3)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  checkDuplicate(
+      keysIdenticalTimestamps, "Duplicate map keys (12291) are not allowed");
+
+  // Check for duplicate keys with the same timestamps in different time zones.
+  const auto keysDifferentTimeZones = makeArrayVector(
+      {0},
+      makeFlatVector<int64_t>(
+          {pack(1, 1),
+           pack(2, 2),
+           pack(3, 3),
+           pack(4, 4),
+           pack(5, 5),
+           pack(3, 6)},
+          TIMESTAMP_WITH_TIME_ZONE()));
+  checkDuplicate(
+      keysDifferentTimeZones, "Duplicate map keys (12294) are not allowed");
+
+  // Check for duplicate keys when the keys vector is a constant.
+  VectorPtr keysConstant =
+      BaseVector::wrapInConstant(1, 0, keysDifferentTimeZones);
+  checkDuplicate(keysConstant, "Duplicate map keys (12294) are not allowed");
+
+  // Check for duplicate keys when the keys vector wrapped in a dictionary.
+  VectorPtr keysInDictionary =
+      wrapInDictionary(makeIndices(1, folly::identity), keysDifferentTimeZones);
+  checkDuplicate(
+      keysInDictionary, "Duplicate map keys (12294) are not allowed");
+
+  // Check for duplicate keys nested inside a complex key.
+  VectorPtr arrayOfRows = makeArrayVector(
+      {0},
+      makeRowVector({makeFlatVector<int64_t>(
+          {pack(1, 1),
+           pack(2, 2),
+           pack(3, 3),
+           pack(4, 4),
+           pack(5, 5),
+           pack(3, 6)},
+          TIMESTAMP_WITH_TIME_ZONE())}));
+  checkDuplicate(arrayOfRows, "Duplicate map keys ({12294}) are not allowed");
+}
+
 } // namespace