Rehydrate vectors at the deepest available object in the JSON path (#131663)

This change ensures that when restoring vectors removed from the source, we locate the deepest existing object along the JSON path to rehydrate them accurately.

Author: jimczi

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.core.TimeValue;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -574,4 +575,118 @@ public static String[] nodeStringArrayValue(Object node) {
             return Strings.splitStringByCommaToArray(node.toString());
         }
     }
+
+    public static void insertValue(String path, Map<?, ?> map, Object newValue) {
+        insertValue(path, map, newValue, true);
+    }
+
+    /**
+     * <p>
+     * Insert or replace the path's value in the map with the provided new value. The map will be modified in-place.
+     * If the complete path does not exist in the map, it will be added to the deepest (sub-)map possible.
+     * </p>
+     * <p>
+     * For example, given the map:
+     * </p>
+     * <pre>
+     * {
+     *   "path1": {
+     *     "path2": {
+     *       "key1": "value1"
+     *     }
+     *   }
+     * }
+     * </pre>
+     * <p>
+     * And the caller wanted to insert {@code "path1.path2.path3.key2": "value2"}, the method would emit the modified map:
+     * </p>
+     * <pre>
+     * {
+     *   "path1": {
+     *     "path2": {
+     *       "key1": "value1",
+     *       "path3.key2": "value2"
+     *     }
+     *   }
+     * }
+     * </pre>
+     *
+     * @param path the value's path in the map.
+     * @param map the map to search and modify in-place.
+     * @param newValue the new value to assign to the path.
+     * @param failOnMultiMap whether the insertion should fail with an {@link IllegalArgumentException}
+     *                       if the path can be resolved in multiple ways.
+     *
+     * @throws IllegalArgumentException If either the path cannot be fully traversed.
+     */
+    public static void insertValue(String path, Map<?, ?> map, Object newValue, boolean failOnMultiMap) {
+        String[] pathElements = path.split("\\.");
+        if (pathElements.length == 0) {
+            return;
+        }
+
+        List<SuffixMap> suffixMaps = extractSuffixMaps(pathElements, 0, map);
+        if (suffixMaps.isEmpty()) {
+            // This should never happen. Throw in case it does for some reason.
+            throw new IllegalStateException("extractSuffixMaps returned an empty suffix map list");
+        } else if (suffixMaps.size() > 1 && failOnMultiMap) {
+            throw new IllegalArgumentException(
+                "Path [" + path + "] could be inserted in " + suffixMaps.size() + " distinct ways, it is ambiguous which one to use"
+            );
+        }
+        SuffixMap suffixMap = suffixMaps.getFirst();
+        suffixMap.map().put(suffixMap.suffix(), newValue);
+    }
+
+    record SuffixMap(String suffix, Map<String, Object> map) {}
+
+    private static List<SuffixMap> extractSuffixMaps(String[] pathElements, int index, Object currentValue) {
+        if (currentValue instanceof List<?> valueList) {
+            List<SuffixMap> suffixMaps = new ArrayList<>(valueList.size());
+            for (Object o : valueList) {
+                suffixMaps.addAll(extractSuffixMaps(pathElements, index, o));
+            }
+
+            return suffixMaps;
+        } else if (currentValue instanceof Map<?, ?>) {
+            @SuppressWarnings("unchecked")
+            Map<String, Object> map = (Map<String, Object>) currentValue;
+            List<SuffixMap> suffixMaps = new ArrayList<>(map.size());
+
+            String key = pathElements[index];
+            while (index < pathElements.length) {
+                if (map.containsKey(key)) {
+                    if (index + 1 == pathElements.length) {
+                        // We found the complete path
+                        suffixMaps.add(new SuffixMap(key, map));
+                    } else {
+                        // We've matched that path partially, keep traversing to try to match it fully
+                        suffixMaps.addAll(extractSuffixMaps(pathElements, index + 1, map.get(key)));
+                    }
+                }
+
+                if (++index < pathElements.length) {
+                    key += "." + pathElements[index];
+                }
+            }
+
+            if (suffixMaps.isEmpty()) {
+                // We checked for all remaining elements in the path, and they do not exist. This means we found a leaf map that we should
+                // add the value to.
+                suffixMaps.add(new SuffixMap(key, map));
+            }
+
+            return suffixMaps;
+        } else {
+            throw new IllegalArgumentException(
+                "Path ["
+                    + String.join(".", Arrays.copyOfRange(pathElements, 0, index))
+                    + "] has value ["
+                    + currentValue
+                    + "] of type ["
+                    + currentValue.getClass().getSimpleName()
+                    + "], which cannot be traversed into further"
+            );
+        }
+    }
 }

server/src/main/java/org/elasticsearch/index/mapper/SourceLoader.java

@@ -476,7 +476,7 @@ private static void applyPatches(String rootPath, Map<String, Object> map, List<
         for (SyntheticVectorPatch patch : patches) {
             if (patch instanceof LeafSyntheticVectorPath leaf) {
                 String key = extractRelativePath(rootPath, leaf.fullPath());
-                map.put(key, leaf.value());
+                XContentMapValues.insertValue(key, map, leaf.value(), false);
             } else if (patch instanceof NestedSyntheticVectorPath nested) {
                 String nestedPath = extractRelativePath(rootPath, nested.fullPath());
                 List<Map<?, ?>> nestedMaps = XContentMapValues.extractNestedSources(nestedPath, map);