feat: Implement map_keys_by_top_n_values Function in Velox

velox/docs/functions/presto/map.rst

@@ -114,6 +114,15 @@ Map Functions
         SELECT map_top_n(map(ARRAY['a', 'b', 'c'], ARRAY[2, 3, 1]), 2) --- {'b' -> 3, 'a' -> 2}
         SELECT map_top_n(map(ARRAY['a', 'b', 'c'], ARRAY[NULL, 3, NULL]), 2) --- {'b' -> 3, 'c' -> NULL}
 
+.. function:: map_keys_by_top_n_values(map(K,V), n) -> array(K)
+
+    Returns an array of the top N keys from a map. Keeps only the top N elements by value. Keys are used to break ties with the max key being chosen. Both keys and values should be orderable.
+
+    ``n`` must be a non-negative BIGINT value.::
+
+        SELECT map_keys_by_top_n_values(map(ARRAY['a', 'b', 'c'], ARRAY[2, 3, 1]), 2) --- ['b', 'a']
+        SELECT map_keys_by_top_n_values(map(ARRAY['a', 'b', 'c'], ARRAY[NULL, 3, NULL]), 2) --- ['b', 'c']
+
 .. function:: map_top_n_keys(map(K,V), n) -> array(K)
 
     Constructs an array of the top N keys. Keys should be orderable.

velox/functions/prestosql/MapKeysByTopNValues.h

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <algorithm>
+#include "velox/functions/prestosql/MapTopN.h"
+
+namespace facebook::velox::functions {
+
+template <typename TExec>
+struct MapKeysByTopNValuesFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(TExec);
+  void call(
+      out_type<Array<Orderable<T1>>>& out,
+      const arg_type<Map<Orderable<T1>, Orderable<T2>>>& inputMap,
+      int64_t n) {
+    VELOX_USER_CHECK_GE(n, 0, "n must be greater than or equal to 0");
+
+    if (n == 0) {
+      return;
+    }
+
+    using It = typename arg_type<Map<Orderable<T1>, Orderable<T2>>>::Iterator;
+    // utilize comparator from MapTopNFunction to sort the input map.
+    using Compare = typename MapTopNFunction<TExec>::template Compare<It>;
+    const Compare comparator;
+
+    std::priority_queue<It, std::vector<It>, Compare> topEntries;
+
+    for (auto it = inputMap.begin(); it != inputMap.end(); ++it) {
+      if (topEntries.size() < n) {
+        topEntries.push(it);
+      } else if (comparator(it, topEntries.top())) {
+        topEntries.pop();
+        topEntries.push(it);
+      }
+    }
+
+    std::vector<It> result;
+    result.reserve(topEntries.size());
+    while (!topEntries.empty()) {
+      result.push_back(topEntries.top());
+      topEntries.pop();
+    }
+
+    // Output the results in descending order.
+    for (auto it = result.crbegin(); it != result.crend(); ++it) {
+      out.push_back((*it)->first);
+    }
+  }
+};
+
+} // namespace facebook::velox::functions

velox/functions/prestosql/registration/MapFunctionsRegistration.cpp

@@ -19,6 +19,7 @@
 #include "velox/functions/lib/MapConcat.h"
 #include "velox/functions/prestosql/Map.h"
 #include "velox/functions/prestosql/MapFunctions.h"
+#include "velox/functions/prestosql/MapKeysByTopNValues.h"
 #include "velox/functions/prestosql/MapNormalize.h"
 #include "velox/functions/prestosql/MapRemoveNullValues.h"
 #include "velox/functions/prestosql/MapSubset.h"
@@ -121,6 +122,12 @@ void registerMapFunctions(const std::string& prefix) {
       Map<Orderable<T1>, Orderable<T2>>,
       int64_t>({prefix + "map_top_n_keys"});
 
+  registerFunction<
+      MapKeysByTopNValuesFunction,
+      Array<Orderable<T1>>,
+      Map<Orderable<T1>, Orderable<T2>>,
+      int64_t>({prefix + "map_keys_by_top_n_values"});
+
   registerMapSubset(prefix);
 
   registerMapRemoveNullValues(prefix);

velox/functions/prestosql/tests/CMakeLists.txt

@@ -76,6 +76,7 @@ add_executable(
   MapTopNTest.cpp
   MapTopNKeysTest.cpp
   MapKeysAndValuesTest.cpp
+  MapKeysByTopNValuesTest.cpp
   MapMatchTest.cpp
   MapTest.cpp
   MapZipWithTest.cpp

velox/functions/prestosql/tests/MapKeysByTopNValuesTest.cpp

@@ -0,0 +1,163 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#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;
+
+namespace facebook::velox::functions {
+namespace {
+
+class MapKeysByTopNValuesTest : public test::FunctionBaseTest {};
+
+TEST_F(MapKeysByTopNValuesTest, emptyMap) {
+  RowVectorPtr input = makeRowVector({
+      makeMapVectorFromJson<int32_t, int64_t>({
+          "{}",
+      }),
+  });
+
+  assertEqualVectors(
+      evaluate("map_keys_by_top_n_values(c0, 3)", input),
+      makeArrayVectorFromJson<int32_t>({
+          "[]",
+      }));
+}
+
+TEST_F(MapKeysByTopNValuesTest, basic) {
+  auto data = makeRowVector({
+      makeMapVectorFromJson<int32_t, int64_t>({
+          "{1:3, 2:5, 3:1, 4:4, 5:2}",
+          "{1:3, 2:5, 3:null, 4:4, 5:2}",
+          "{1:null, 2:null, 3:1, 4:4, 5:null}",
+          "{1:10, 2:7, 3:11, 5:4}",
+          "{1:10, 2:10, 3:10, 4:10, 5:10}",
+          "{1:10, 2:7, 3:0}",
+          "{1:null, 2:10}",
+          "{}",
+          "{1:null, 2:null, 3:null}",
+      }),
+  });
+
+  auto result = evaluate("map_keys_by_top_n_values(c0, 3)", data);
+
+  auto expected = makeArrayVectorFromJson<int32_t>({
+      "[2, 4, 1]",
+      "[2, 4, 1]",
+      "[4, 3, 5]",
+      "[3, 1, 2]",
+      "[5, 4, 3]",
+      "[1, 2, 3]",
+      "[2, 1]",
+      "[]",
+      "[3, 2, 1]",
+  });
+
+  assertEqualVectors(expected, result);
+
+  // n = 0. Expect empty maps.
+  result = evaluate("map_keys_by_top_n_values(c0, 0)", data);
+
+  expected = makeArrayVectorFromJson<int32_t>({
+      "[]",
+      "[]",
+      "[]",
+      "[]",
+      "[]",
+      "[]",
+      "[]",
+      "[]",
+      "[]",
+  });
+
+  assertEqualVectors(expected, result);
+
+  // n is negative. Expect an error.
+  VELOX_ASSERT_THROW(
+      evaluate("map_keys_by_top_n_values(c0, -1)", data),
+      "n must be greater than or equal to 0");
+}
+
+TEST_F(MapKeysByTopNValuesTest, complexKeys) {
+  RowVectorPtr input =
+      makeRowVector({makeMapVectorFromJson<std::string, int64_t>(
+          {R"( {"x":1, "y":2} )",
+           R"( {"x":1, "x2":-2} )",
+           R"( {"ac":1, "cc":3, "dd": 4} )"})});
+
+  assertEqualVectors(
+      evaluate("map_keys_by_top_n_values(c0, 1)", input),
+      makeArrayVectorFromJson<std::string>({
+          "[\"y\"]",
+          "[\"x\"]",
+          "[\"dd\"]",
+      }));
+}
+
+TEST_F(MapKeysByTopNValuesTest, complexKeysWithLargeK) {
+  RowVectorPtr input =
+      makeRowVector({makeMapVectorFromJson<std::string, int64_t>(
+          {R"( {"x":1, "y":2} )",
+           R"( {"x":1, "x2":-2} )",
+           R"( {"ac":1, "cc":3, "dd": 4} )"})});
+
+  assertEqualVectors(
+      evaluate("map_keys_by_top_n_values(c0, 706100916841560005)", input),
+      makeArrayVectorFromJson<std::string>({
+          "[\"y\", \"x\"]",
+          "[\"x\", \"x2\"]",
+          "[\"dd\", \"cc\", \"ac\"]",
+      }));
+}
+
+TEST_F(MapKeysByTopNValuesTest, 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_keys_by_top_n_values(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(2, 2), pack(4, 4), pack(1, 1)});
+  testMapTopNKeys(
+      {pack(5, 1), pack(4, 2), pack(3, 3), pack(2, 4), pack(1, 5)},
+      {3, 5, 1, 4, 2},
+      {pack(4, 2), pack(2, 4), pack(5, 1)});
+  testMapTopNKeys(
+      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5)},
+      {1, 2, 3, 4, 5},
+      {pack(2, 5), pack(4, 4), pack(1, 3)});
+  testMapTopNKeys(
+      {pack(3, 5), pack(5, 4), pack(4, 2), pack(2, 1)},
+      {3, 3, 3, 3},
+      {pack(5, 4), pack(4, 2), pack(3, 5)});
+  testMapTopNKeys({}, {}, {});
+}
+} // namespace
+} // namespace facebook::velox::functions