Add map_trim_values UDF (facebookincubator#15973)

Summary:
Pull Request resolved: facebookincubator#15973

Implements a new Velox-only function `map_trim_values(map(K, array(V)), n)` that trims the value arrays in a map to a specified maximum size. This is useful for optimizing memory usage and performance for large feature maps where value arrays may grow unbounded.

Behavior:
- If n >= 0: Each value array is trimmed to at most n elements
- If n < 0: The original map is returned unchanged
- Null elements in arrays are preserved in the output

The implementation follows the simple function pattern used by other map functions like `map_subset`.

Reviewed By: zacw7

Differential Revision: D90476691

fbshipit-source-id: 35ee713ca0dbf55857f189086311cea2919d3911

Author: abhinavmuk04

velox/docs/functions/presto/map.rst

@@ -192,6 +192,24 @@ 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_trim_values(map(K, array(V)), n) -> map(K, array(V))
+
+    Trims the value arrays in a map to a specified maximum size.
+    This function is useful for optimizing memory usage and performance for large feature maps
+    where the value arrays may grow unbounded.
+
+    Returns a map where each value array is trimmed to at most n elements.
+    If n is negative, returns the original map unchanged.
+    If n is 0, returns a map where all values are empty arrays.
+    If a value array has fewer than n elements, it is left unchanged.
+    Null elements in the arrays are preserved in the output. ::
+
+        SELECT map_trim_values(MAP(ARRAY['a', 'b'], ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5, 6, 7]]), 2); -- {a -> [1, 2], b -> [4, 5]}
+        SELECT map_trim_values(MAP(ARRAY['a'], ARRAY[ARRAY[1, 2]]), 5); -- {a -> [1, 2]}
+        SELECT map_trim_values(MAP(ARRAY['a'], ARRAY[ARRAY[1, NULL, 3]]), 2); -- {a -> [1, NULL]}
+        SELECT map_trim_values(MAP(ARRAY['a'], ARRAY[ARRAY[1, 2, 3]]), 0); -- {a -> []}
+        SELECT map_trim_values(MAP(ARRAY['a'], ARRAY[ARRAY[1, 2, 3]]), -1); -- {a -> [1, 2, 3]}
+
 .. 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.

velox/expression/fuzzer/ExpressionFuzzerTest.cpp

@@ -304,6 +304,7 @@ std::unordered_set<std::string> skipFunctionsSOT = {
     "map_keys_overlap", // Velox-only function, not available in Presto
     "map_append", // Velox-only function, not available in Presto
     "map_update", // Velox-only function, not available in Presto
+    "map_trim_values", // Velox-only function, not available in Presto
     "noisy_empty_approx_set_sfm", // non-deterministic because of privacy.
     // https://github.com/facebookincubator/velox/issues/11034
     "cast(real) -> varchar",

velox/functions/prestosql/MapTrimValues.h

@@ -0,0 +1,83 @@
+/*
+ * 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 <cstddef>
+#include <cstdint>
+
+#include "velox/expression/ComplexViewTypes.h"
+#include "velox/functions/Macros.h"
+
+namespace facebook::velox::functions {
+
+/// map_trim_values(map(K, array(V)), n) -> map(K, array(V))
+///
+/// Trims the value arrays in a map to a specified maximum size.
+/// This function is useful for optimizing memory usage and performance
+/// for large feature maps where the value arrays may grow unbounded.
+///
+/// Returns a map where each value array is trimmed to at most n elements.
+/// If n is negative, returns the original map unchanged.
+/// If n is 0, returns a map where all values are empty arrays.
+/// If a value array has fewer than n elements, it is left unchanged.
+/// Null map values are preserved in the output.
+/// Null elements within arrays are also preserved.
+template <typename TExec>
+struct MapTrimValuesFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(TExec);
+
+  void call(
+      out_type<Map<Generic<T1>, Array<Generic<T2>>>>& out,
+      const arg_type<Map<Generic<T1>, Array<Generic<T2>>>>& inputMap,
+      int64_t n) {
+    // If n is negative, preserve the original map as-is
+    if (n < 0) {
+      out.copy_from(inputMap);
+      return;
+    }
+
+    for (const auto& entry : inputMap) {
+      if (!entry.second.has_value()) {
+        auto& keyWriter = out.add_null();
+        keyWriter.copy_from(entry.first);
+      } else {
+        auto [keyWriter, valueWriter] = out.add_item();
+        keyWriter.copy_from(entry.first);
+
+        const auto& valueArray = entry.second.value();
+        const auto arraySize = static_cast<size_t>(valueArray.size());
+        const auto trimSize = std::min(static_cast<size_t>(n), arraySize);
+
+        size_t count = 0;
+        for (const auto& element : valueArray) {
+          if (count >= trimSize) {
+            break;
+          }
+          if (element.has_value()) {
+            auto& elementWriter = valueWriter.add_item();
+            elementWriter.copy_from(element.value());
+          } else {
+            valueWriter.add_null();
+          }
+          ++count;
+        }
+      }
+    }
+  }
+};
+
+} // namespace facebook::velox::functions

velox/functions/prestosql/registration/MapFunctionsRegistration.cpp

@@ -30,6 +30,7 @@
 #include "velox/functions/prestosql/MapTopN.h"
 #include "velox/functions/prestosql/MapTopNKeys.h"
 #include "velox/functions/prestosql/MapTopNValues.h"
+#include "velox/functions/prestosql/MapTrimValues.h"
 #include "velox/functions/prestosql/MapUpdate.h"
 #include "velox/functions/prestosql/MapValuesInRange.h"
 #include "velox/functions/prestosql/MultimapFromEntries.h"
@@ -355,6 +356,12 @@ void registerMapFunctions(const std::string& prefix) {
       Map<Generic<T1>, double>,
       double,
       double>({prefix + "map_values_in_range"});
+
+  registerFunction<
+      MapTrimValuesFunction,
+      Map<Generic<T1>, Array<Generic<T2>>>,
+      Map<Generic<T1>, Array<Generic<T2>>>,
+      int64_t>({prefix + "map_trim_values"});
 }
 
 void registerMapAllowingDuplicates(

velox/functions/prestosql/tests/CMakeLists.txt

@@ -93,6 +93,7 @@ add_executable(
   MapTopNTest.cpp
   MapTopNKeysTest.cpp
   MapTopNValuesTest.cpp
+  MapTrimValuesTest.cpp
   MapKeysAndValuesTest.cpp
   MapKeysByTopNValuesTest.cpp
   MapMatchTest.cpp