Add dfs transformation function in XContentMapValues

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Change priority for scheduling reroute during timeout([#16445](https://github.com/opensearch-project/OpenSearch/pull/16445))
 - Renaming the node role search to warm ([#17573](https://github.com/opensearch-project/OpenSearch/pull/17573))
 - Introduce a new search node role to hold search only shards ([#17620](https://github.com/opensearch-project/OpenSearch/pull/17620))
+- Add dfs transformation function in XContentMapValues ([#17612](https://github.com/opensearch-project/OpenSearch/pull/17612))
 
 ### Dependencies
 - Bump `ch.qos.logback:logback-core` from 1.5.16 to 1.5.17 ([#17609](https://github.com/opensearch-project/OpenSearch/pull/17609))

server/src/main/java/org/opensearch/common/xcontent/support/XContentMapValues.java

@@ -43,9 +43,11 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.Strings;
 
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -60,6 +62,8 @@
  */
 public class XContentMapValues {
 
+    private static final String TRANSFORMER_TRIE_LEAF_KEY = "$transformer";
+
     /**
      * Extracts raw values (string, int, and so on) based on the path provided returning all of them
      * as a single list.
@@ -621,4 +625,149 @@ public static String[] nodeStringArrayValue(Object node) {
             return Strings.splitStringByCommaToArray(node.toString());
         }
     }
+
+    /**
+     * Performs a depth first traversal of a map and applies a transformation for each field matched along the way. For
+     * duplicated paths with transformers (i.e. "test.nested" and "test.nested.field"), only the transformer for
+     * the shorter path is applied.
+     *
+     * @param source Source map to perform transformation on
+     * @param transformers Map from path to transformer to apply to each path. Each transformer is a function that takes
+     *                    the current value and returns a transformed value
+     * @param inPlace If true, modify the source map directly; if false, create a copy
+     * @return Map with transformations applied
+     */
+    public static Map<String, Object> transform(
+        Map<String, Object> source,
+        Map<String, Function<Object, Object>> transformers,
+        boolean inPlace
+    ) {
+        return transform(transformers, inPlace).apply(source);
+    }
+
+    /**
+     * Returns function that performs a depth first traversal of a map and applies a transformation for each field
+     * matched along the way. For duplicated paths with transformers (i.e. "test.nested" and "test.nested.field"), only
+     * the transformer for the shorter path is applied.
+     *
+     * @param transformers Map from path to transformer to apply to each path. Each transformer is a function that takes
+     *                     the current value and returns a transformed value
+     * @param inPlace If true, modify the source map directly; if false, create a copy
+     * @return Function that takes a map and returns a transformed version of the map
+     */
+    public static Function<Map<String, Object>, Map<String, Object>> transform(
+        Map<String, Function<Object, Object>> transformers,
+        boolean inPlace
+    ) {
+        Map<String, Object> transformerTrie = buildTransformerTrie(transformers);
+        return source -> {
+            Deque<TransformContext> stack = new ArrayDeque<>();
+            Map<String, Object> result = inPlace ? source : new HashMap<>(source);
+            stack.push(new TransformContext(result, transformerTrie));
+
+            processStack(stack, inPlace);
+            return result;
+        };
+    }
+
+    @SuppressWarnings("unchecked")
+    private static Map<String, Object> buildTransformerTrie(Map<String, Function<Object, Object>> transformers) {
+        Map<String, Object> trie = new HashMap<>();
+        for (Map.Entry<String, Function<Object, Object>> entry : transformers.entrySet()) {
+            String[] pathElements = entry.getKey().split("\\.");
+            Map<String, Object> subTrie = trie;
+            for (String pathElement : pathElements) {
+                subTrie = (Map<String, Object>) subTrie.computeIfAbsent(pathElement, k -> new HashMap<>());
+            }
+            subTrie.put(TRANSFORMER_TRIE_LEAF_KEY, entry.getValue());
+        }
+        return trie;
+    }
+
+    private static void processStack(Deque<TransformContext> stack, boolean inPlace) {
+        while (!stack.isEmpty()) {
+            TransformContext ctx = stack.pop();
+            processMap(ctx.map, ctx.trie, stack, inPlace);
+        }
+    }
+
+    private static void processMap(
+        Map<String, Object> currentMap,
+        Map<String, Object> currentTrie,
+        Deque<TransformContext> stack,
+        boolean inPlace
+    ) {
+        for (Map.Entry<String, Object> entry : currentMap.entrySet()) {
+            processEntry(entry, currentTrie, stack, inPlace);
+        }
+    }
+
+    private static void processEntry(
+        Map.Entry<String, Object> entry,
+        Map<String, Object> currentTrie,
+        Deque<TransformContext> stack,
+        boolean inPlace
+    ) {
+        String key = entry.getKey();
+        Object value = entry.getValue();
+
+        Object subTrieObj = currentTrie.get(key);
+        if (subTrieObj instanceof Map == false) {
+            return;
+        }
+        Map<String, Object> subTrie = nodeMapValue(subTrieObj, "transform");
+
+        // Apply transformation if available
+        Function<Object, Object> transformer = (Function<Object, Object>) subTrie.get(TRANSFORMER_TRIE_LEAF_KEY);
+        if (transformer != null) {
+            entry.setValue(transformer.apply(value));
+            return;
+        }
+
+        // Process nested structures
+        if (value instanceof Map) {
+            Map<String, Object> subMap = nodeMapValue(value, "transform");
+            if (inPlace == false) {
+                subMap = new HashMap<>(subMap);
+                entry.setValue(subMap);
+            }
+            stack.push(new TransformContext(subMap, subTrie));
+        } else if (value instanceof List<?> list) {
+            List<Object> subList = (List<Object>) list;
+            if (inPlace == false) {
+                subList = new ArrayList<>(list);
+                entry.setValue(subList);
+            }
+            processList(subList, subTrie, stack, inPlace);
+        }
+    }
+
+    private static void processList(
+        List<Object> list,
+        Map<String, Object> transformerTrie,
+        Deque<TransformContext> stack,
+        boolean inPlace
+    ) {
+        for (int i = list.size() - 1; i >= 0; i--) {
+            Object value = list.get(i);
+            if (value instanceof Map) {
+                Map<String, Object> subMap = nodeMapValue(value, "transform");
+                if (inPlace == false) {
+                    subMap = new HashMap<>(subMap);
+                    list.set(i, subMap);
+                }
+                stack.push(new TransformContext(subMap, transformerTrie));
+            }
+        }
+    }
+
+    private static class TransformContext {
+        Map<String, Object> map;
+        Map<String, Object> trie;
+
+        TransformContext(Map<String, Object> map, Map<String, Object> trie) {
+            this.map = map;
+            this.trie = trie;
+        }
+    }
 }

server/src/test/java/org/opensearch/common/xcontent/support/XContentMapValuesTests.java

@@ -32,6 +32,7 @@
 
 package org.opensearch.common.xcontent.support;
 
+import org.opensearch.common.collect.MapBuilder;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.common.xcontent.XContentType;
@@ -48,9 +49,12 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.IntStream;
 
 import static java.util.Collections.emptySet;
 import static java.util.Collections.singleton;
@@ -629,6 +633,127 @@ public void testPrefix() {
         assertEquals(expected, filtered);
     }
 
+    public void testTransformFlat() {
+        Map<String, Object> mapToTransform = Map.of(
+            "test1",
+            "value_before",
+            "test2",
+            List.of("value_before", "value_before", "value_before")
+        );
+
+        Map<String, Function<Object, Object>> transformers = Map.of("test1", v -> "value_after", "test2", v -> "value_after");
+
+        Map<String, Object> expected = Map.of("test1", "value_after", "test2", "value_after");
+
+        Map<String, Object> transformedMapped = XContentMapValues.transform(mapToTransform, transformers, false);
+        assertEquals(expected, transformedMapped);
+    }
+
+    public void testTransformNested() {
+        Map<String, Object> mapToTransform = MapBuilder.<String, Object>newMapBuilder()
+            .put("test1", "value_before")
+            .put("test2", Map.of("nest2", "value_before"))
+            .put("test3", List.of(Map.of("nest3", "value_before"), Map.of("nest3", "value_before"), Map.of("nest3", "value_before")))
+            .put(
+                "test4",
+                List.of(
+                    Map.of(
+                        "nest4",
+                        List.of(Map.of("nest5", "value_before"), Map.of("nest5", "value_before"), Map.of("nest5", "value_before")),
+                        "test5",
+                        "no_change"
+                    ),
+                    Map.of(
+                        "nest4",
+                        List.of(
+                            Map.of("nest5", "value_before"),
+                            Map.of("nest5", "value_before"),
+                            Map.of("nest5", "value_before"),
+                            Map.of("nest5", "value_before")
+                        ),
+                        "test5",
+                        "no_change"
+                    ),
+                    Map.of("nest4", List.of(Map.of("nest5", "value_before"), Map.of("nest5", "value_before")), "test5", "no_change")
+                )
+            )
+            .put("test6", null)
+            .immutableMap();
+
+        Iterator<String> test3Stream = IntStream.rangeClosed(1, 3).mapToObj(i -> "value_after" + i).toList().iterator();
+        Iterator<String> test4Stream = IntStream.rangeClosed(1, 9).mapToObj(i -> "value_after" + i).toList().iterator();
+        Map<String, Function<Object, Object>> transformers = Map.of(
+            "test1",
+            v -> "value_after",
+            "test2.nest2",
+            v -> "value_after",
+            "test3.nest3",
+            v -> test3Stream.next(),
+            "test4.nest4.nest5",
+            v -> test4Stream.next(),
+            "test6",
+            v -> v == null ? v : "value_after"
+        );
+
+        Map<String, Object> expected = MapBuilder.<String, Object>newMapBuilder()
+            .put("test1", "value_after")
+            .put("test2", Map.of("nest2", "value_after"))
+            .put("test3", List.of(Map.of("nest3", "value_after1"), Map.of("nest3", "value_after2"), Map.of("nest3", "value_after3")))
+            .put(
+                "test4",
+                List.of(
+                    Map.of(
+                        "nest4",
+                        List.of(Map.of("nest5", "value_after1"), Map.of("nest5", "value_after2"), Map.of("nest5", "value_after3")),
+                        "test5",
+                        "no_change"
+                    ),
+                    Map.of(
+                        "nest4",
+                        List.of(
+                            Map.of("nest5", "value_after4"),
+                            Map.of("nest5", "value_after5"),
+                            Map.of("nest5", "value_after6"),
+                            Map.of("nest5", "value_after7")
+                        ),
+                        "test5",
+                        "no_change"
+                    ),
+                    Map.of("nest4", List.of(Map.of("nest5", "value_after8"), Map.of("nest5", "value_after9")), "test5", "no_change")
+                )
+            )
+            .put("test6", null)
+            .immutableMap();
+
+        Map<String, Object> transformedMapped = XContentMapValues.transform(mapToTransform, transformers, false);
+        assertEquals(expected, transformedMapped);
+    }
+
+    public void testTransformInPlace() {
+        Map<String, Object> mapToTransform = MapBuilder.<String, Object>newMapBuilder().put("test1", "value_before").map();
+        Map<String, Function<Object, Object>> transformers = Map.of("test1", v -> "value_after");
+        Map<String, Object> expected = MapBuilder.<String, Object>newMapBuilder().put("test1", "value_after").immutableMap();
+
+        Map<String, Object> transformedMapped = XContentMapValues.transform(mapToTransform, transformers, true);
+        assertEquals(expected, transformedMapped);
+    }
+
+    public void testTransformSharedPaths() {
+        Map<String, Object> mapToTransform = MapBuilder.<String, Object>newMapBuilder()
+            .put("test", "value_before")
+            .put("test.nested", "nested_value_before")
+            .map();
+        Map<String, Function<Object, Object>> transformers = Map.of("test", v -> "value_after", "test.nested", v -> "nested_value_after");
+
+        Map<String, Object> expected = MapBuilder.<String, Object>newMapBuilder()
+            .put("test", "value_after")
+            .put("test.nested", "nested_value_before")
+            .immutableMap();
+
+        Map<String, Object> transformedMap = XContentMapValues.transform(mapToTransform, transformers, true);
+        assertEquals(expected, transformedMap);
+    }
+
     private static Map<String, Object> toMap(Builder test, XContentType xContentType, boolean humanReadable) throws IOException {
         ToXContentObject toXContent = (builder, params) -> test.apply(builder);
         return convertToMap(toXContent(toXContent, xContentType, humanReadable), true, xContentType).v2();