Add tests for Presto's map_top_n/map_top_n_keys UDFs with TimestampWi…

…thTimezones (#11243)

Summary:
Pull Request resolved: #11243

Presto's map_top_n and map_top_n_keys implementations use 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 these UDFs
work 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: amitkdutta

Differential Revision: D64267032

fbshipit-source-id: b7fdd725cae5918c98a9cb0a6d053966991f00c5

velox/functions/prestosql/tests/MapTopNKeysTest.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;
 
@@ -79,5 +80,46 @@ TEST_F(MapTopNKeysTest, nIsNegative) {
       "n must be greater than or equal to 0");
 }
 
+TEST_F(MapTopNKeysTest, timestampWithTimeZone) {
+  auto testMapTopNKeys = [&](const std::vector<int64_t>& keys,
+                             const std::vector<int32_t>& values,
+                             const std::vector<int64_t>& expectedKeys) {
+    const auto map = makeMapVector(
+        {0},
+        makeFlatVector(keys, TIMESTAMP_WITH_TIME_ZONE()),
+        makeFlatVector(values));
+    const auto expected = makeArrayVector(
+        {0}, makeFlatVector(expectedKeys, TIMESTAMP_WITH_TIME_ZONE()));
+
+    const auto result = evaluate("map_top_n_keys(c0, 3)", makeRowVector({map}));
+
+    assertEqualVectors(expected, result);
+  };
+
+  testMapTopNKeys(
+      {pack(1, 1), pack(2, 2), pack(3, 3), pack(4, 4), pack(5, 5)},
+      {3, 5, 1, 4, 2},
+      {pack(5, 5), pack(4, 4), pack(3, 3)});
+  testMapTopNKeys(
+      {pack(5, 1), pack(4, 2), pack(3, 3), pack(2, 4), pack(1, 5)},
+      {3, 5, 1, 4, 2},
+      {pack(5, 1), pack(4, 2), pack(3, 3)});
+  testMapTopNKeys(
+      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5)},
+      {1, 2, 3, 4, 5},
+      {pack(5, 2), pack(4, 4), pack(3, 1)});
+  testMapTopNKeys(
+      {pack(3, 5), pack(5, 4), pack(4, 2), pack(2, 1)},
+      {1, 2, 4, 5},
+      {pack(5, 4), pack(4, 2), pack(3, 5)});
+  testMapTopNKeys(
+      {pack(10, 3), pack(7, 2), pack(0, 1)},
+      {1, 2, 3},
+      {pack(10, 3), pack(7, 2), pack(0, 1)});
+  testMapTopNKeys(
+      {pack(1, 10), pack(10, 1)}, {1, 2}, {pack(10, 1), pack(1, 10)});
+  testMapTopNKeys({}, {}, {});
+}
+
 } // namespace
 } // namespace facebook::velox::functions

velox/functions/prestosql/tests/MapTopNTest.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;
 
@@ -130,5 +131,85 @@ TEST_F(MapTopNTest, equalValues) {
   assertEqualVectors(expectedResults, result);
 }
 
+TEST_F(MapTopNTest, timestampWithTimeZone) {
+  auto testMapTopN =
+      [&](const std::vector<int32_t>& keys,
+          const std::vector<std::optional<int64_t>>& values,
+          const std::vector<int32_t>& expectedKeys,
+          const std::vector<std::optional<int64_t>>& expectedValues) {
+        const auto map = makeMapVector(
+            {0},
+            makeFlatVector(keys),
+            makeNullableFlatVector(values, TIMESTAMP_WITH_TIME_ZONE()));
+        const auto expectedMap = makeMapVector(
+            {0},
+            makeFlatVector(expectedKeys),
+            makeNullableFlatVector(expectedValues, TIMESTAMP_WITH_TIME_ZONE()));
+
+        const auto result = evaluate("map_top_n(c0, 3)", makeRowVector({map}));
+
+        assertEqualVectors(expectedMap, result);
+      };
+
+  testMapTopN(
+      {1, 2, 3, 4, 5},
+      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5)},
+      {2, 4, 1},
+      {pack(5, 2), pack(4, 4), pack(3, 1)});
+  testMapTopN(
+      {1, 2, 3, 4, 5},
+      {pack(3, 5), pack(5, 4), std::nullopt, pack(4, 2), pack(2, 1)},
+      {2, 4, 1},
+      {pack(5, 4), pack(4, 2), pack(3, 5)});
+  testMapTopN(
+      {1, 2, 3, 4, 5},
+      {std::nullopt, std::nullopt, pack(1, 1), pack(4, 4), std::nullopt},
+      {4, 3, 5},
+      {pack(4, 4), pack(1, 1), std::nullopt});
+  testMapTopN(
+      {1, 2, 3, 5},
+      {pack(10, 1), pack(7, 2), pack(11, 3), pack(4, 4)},
+      {3, 1, 2},
+      {pack(11, 3), pack(10, 1), pack(7, 2)});
+  testMapTopN(
+      {1, 2, 3},
+      {pack(10, 3), pack(7, 2), pack(0, 1)},
+      {1, 2, 3},
+      {pack(10, 3), pack(7, 2), pack(0, 1)});
+  testMapTopN(
+      {1, 2}, {std::nullopt, pack(10, 1)}, {2, 1}, {pack(10, 1), std::nullopt});
+  testMapTopN({}, {}, {}, {});
+  testMapTopN(
+      {1, 2, 3},
+      {std::nullopt, std::nullopt, std::nullopt},
+      {1, 2, 3},
+      {std::nullopt, std::nullopt, std::nullopt});
+  testMapTopN(
+      {6, 2, 3, 4, 5, 1},
+      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5), pack(3, 1)},
+      {2, 4, 6},
+      {pack(5, 2), pack(4, 4), pack(3, 1)});
+  testMapTopN(
+      {1, 2, 3, 4, 5, 6},
+      {pack(3, 6),
+       pack(5, 5),
+       std::nullopt,
+       pack(4, 3),
+       pack(2, 2),
+       pack(5, 5)},
+      {6, 2, 4},
+      {pack(5, 5), pack(5, 5), pack(4, 3)});
+  testMapTopN(
+      {5, 2, 3, 4, 1},
+      {std::nullopt, std::nullopt, pack(1, 1), pack(4, 2), std::nullopt},
+      {4, 3, 5},
+      {pack(4, 2), pack(1, 1), std::nullopt});
+  testMapTopN(
+      {1, 5, 3, 4, 2},
+      {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt},
+      {4, 3, 5},
+      {std::nullopt, std::nullopt, std::nullopt});
+}
+
 } // namespace
 } // namespace facebook::velox::functions