Add tests for Presto's map_subset UDF with TimestampWithTimezones (fa…

…cebookincubator#11242)

Summary:
Pull Request resolved: facebookincubator#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

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