From adbebfdad56b1c07a274cf89ede25732201c2360 Mon Sep 17 00:00:00 2001
From: Kevin Wilfong <kevinwilfong@meta.com>
Date: Fri, 11 Oct 2024 15:40:05 -0700
Subject: [PATCH] Add tests for Presto's map_subset UDF with
 TimestampWithTimezones (#11242)

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

Presto's map_subset implementation for Generic types uses the hash and equals
functions implemented for GenericViews.  For TimestampWithTImeZone this
recursively calls these functions on the CustomTypeWithCustomComparisonView the
GenericView wraps, which supports custom comparison operators. So this UDF works
as is with the recent changes earlier in this stack.

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

Reviewed By: xiaoxmeng

Differential Revision: D64262952

fbshipit-source-id: 8f98e5707e5b2fb90d2efadcf3c12d1e24f829ef
---
 .../prestosql/tests/MapSubsetTest.cpp         | 75 +++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/velox/functions/prestosql/tests/MapSubsetTest.cpp b/velox/functions/prestosql/tests/MapSubsetTest.cpp
index 8d120003d639..1645c204972d 100644
--- a/velox/functions/prestosql/tests/MapSubsetTest.cpp
+++ b/velox/functions/prestosql/tests/MapSubsetTest.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::test;
 
@@ -198,5 +199,79 @@ TEST_F(MapSubsetTest, floatNaNs) {
   testFloatNaNs<double>();
 }
 
+TEST_F(MapSubsetTest, timestampWithTimeZone) {
+  const auto keys = makeFlatVector<int64_t>(
+      {pack(1, 1),
+       pack(2, 2),
+       pack(3, 3),
+       pack(4, 4),
+       pack(5, 5),
+       pack(6, 6),
+       pack(1, 7),
+       pack(2, 8),
+       pack(4, 9),
+       pack(5, 10),
+       pack(2, 11),
+       pack(4, 12),
+       pack(6, 13)},
+      TIMESTAMP_WITH_TIME_ZONE());
+  const auto values = makeNullableFlatVector<int32_t>(
+      {10, 20, std::nullopt, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130});
+  const auto maps = makeMapVector({0, 6, 10, 10}, keys, values);
+
+  // Test map with TimestampWithTimeZone keys and constant second arg.
+  const auto constLookup = BaseVector::wrapInConstant(
+      4,
+      0,
+      makeArrayVector(
+          {0},
+          makeFlatVector<int64_t>(
+              {pack(1, 1), pack(3, 2), pack(5, 3)},
+              TIMESTAMP_WITH_TIME_ZONE())));
+  const auto expectedKeys = makeFlatVector<int64_t>(
+      {pack(1, 1), pack(3, 3), pack(5, 50), pack(1, 7), pack(5, 10)},
+      TIMESTAMP_WITH_TIME_ZONE());
+  const auto expectedValues =
+      makeNullableFlatVector<int32_t>({10, std::nullopt, 50, 70, 100});
+  const auto expected =
+      makeMapVector({0, 3, 5, 5}, expectedKeys, expectedValues);
+  auto result =
+      evaluate("map_subset(c0, c1)", makeRowVector({maps, constLookup}));
+
+  assertEqualVectors(expected, result);
+
+  // Test map with TimestampWithTimeZone keys and non-constant second arg.
+  const auto lookupKeys = makeFlatVector<int64_t>(
+      {pack(1, 1),
+       pack(3, 3),
+       pack(5, 5),
+       pack(1, 10),
+       pack(3, 12),
+       pack(5, 13),
+       pack(7, 14),
+       pack(3, 15),
+       pack(5, 16),
+       pack(1, 17),
+       pack(3, 18)},
+      TIMESTAMP_WITH_TIME_ZONE());
+  const auto lookup = makeArrayVector({0, 3, 7, 9}, lookupKeys);
+
+  result = evaluate("map_subset(c0, c1)", makeRowVector({maps, lookup}));
+  assertEqualVectors(expected, result);
+
+  // Test map with TimestampWithTimeZone wrapped in a complex type as keys.
+  const auto mapsWithRowKeys =
+      makeMapVector({0, 6, 10, 10}, makeRowVector({keys}), values);
+  const auto lookupWithRowKeys =
+      makeArrayVector({0, 3, 7, 9}, makeRowVector({lookupKeys}));
+  const auto expectedWithRowKeys = makeMapVector(
+      {0, 3, 5, 5}, makeRowVector({expectedKeys}), expectedValues);
+
+  result = evaluate(
+      "map_subset(c0, c1)",
+      makeRowVector({mapsWithRowKeys, lookupWithRowKeys}));
+  assertEqualVectors(expectedWithRowKeys, result);
+}
+
 } // namespace
 } // namespace facebook::velox::functions