Refactor IngestDocument's resolve method (elastic#125051)

Author: joegallo

server/src/main/java/org/elasticsearch/ingest/IngestCtxMap.java

@@ -17,15 +17,17 @@
 
 /**
  * Map containing ingest source and metadata.
- *
- * The Metadata values in {@link IngestDocument.Metadata} are validated when put in the map.
- * _index, _id and _routing must be a String or null
- * _version_type must be a lower case VersionType or null
- * _version must be representable as a long without loss of precision or null
- * _dynamic_templates must be a map
- * _if_seq_no must be a long or null
- * _if_primary_term must be a long or null
- *
+ * <p>
+ * The Metadata values in {@link IngestDocument.Metadata} are validated when put in the map:
+ * <ul>
+ *   <li>{@code _index}, {@code _id} and {@code _routing} must be a String or null</li>
+ *   <li>{@code _version_type} must be a lower case VersionType or null</li>
+ *   <li>{@code _version} must be representable as a long without loss of precision or null</li>
+ *   <li>{@code _dynamic_templates} must be a map</li>
+ *   <li>{@code _if_seq_no} must be a long or null</li>
+ *   <li>{@code _if_primary_term} must be a long or null</li>
+ * </ul>
+ * <p>
  * The map is expected to be used by processors, server code should the typed getter and setters where possible.
  */
 final class IngestCtxMap extends CtxMap<IngestDocMetadata> {

server/src/main/java/org/elasticsearch/ingest/IngestDocument.java

@@ -193,17 +193,14 @@ public <T> T getFieldValue(String path, Class<T> clazz) {
     public <T> T getFieldValue(String path, Class<T> clazz, boolean ignoreMissing) {
         final FieldPath fieldPath = FieldPath.of(path);
         Object context = fieldPath.initialContext(this);
-        for (String pathElement : fieldPath.pathElements) {
-            ResolveResult result = resolve(pathElement, path, context);
-            if (result.wasSuccessful) {
-                context = result.resolvedObject;
-            } else if (ignoreMissing && hasField(path) == false) {
-                return null;
-            } else {
-                throw new IllegalArgumentException(result.errorMessage);
-            }
+        ResolveResult result = resolve(fieldPath.pathElements, fieldPath.pathElements.length, path, context);
+        if (result.wasSuccessful) {
+            return cast(path, result.resolvedObject, clazz);
+        } else if (ignoreMissing) {
+            return null;
+        } else {
+            throw new IllegalArgumentException(result.errorMessage);
         }
-        return cast(path, context, clazz);
     }
 
     /**
@@ -266,6 +263,8 @@ public boolean hasField(String path, boolean failOutOfRange) {
             String pathElement = fieldPath.pathElements[i];
             if (context == null) {
                 return false;
+            } else if (context instanceof IngestCtxMap map) { // optimization: handle IngestCtxMap separately from Map
+                context = map.get(pathElement);
             } else if (context instanceof Map<?, ?> map) {
                 context = map.get(pathElement);
             } else if (context instanceof List<?> list) {
@@ -292,6 +291,8 @@ public boolean hasField(String path, boolean failOutOfRange) {
         String leafKey = fieldPath.pathElements[fieldPath.pathElements.length - 1];
         if (context == null) {
             return false;
+        } else if (context instanceof IngestCtxMap map) { // optimization: handle IngestCtxMap separately from Map
+            return map.containsKey(leafKey);
         } else if (context instanceof Map<?, ?> map) {
             return map.containsKey(leafKey);
         } else if (context instanceof List<?> list) {
@@ -322,18 +323,22 @@ public boolean hasField(String path, boolean failOutOfRange) {
     public void removeField(String path) {
         final FieldPath fieldPath = FieldPath.of(path);
         Object context = fieldPath.initialContext(this);
-        for (int i = 0; i < fieldPath.pathElements.length - 1; i++) {
-            ResolveResult result = resolve(fieldPath.pathElements[i], path, context);
-            if (result.wasSuccessful) {
-                context = result.resolvedObject;
-            } else {
-                throw new IllegalArgumentException(result.errorMessage);
-            }
+        ResolveResult result = resolve(fieldPath.pathElements, fieldPath.pathElements.length - 1, path, context);
+        if (result.wasSuccessful) {
+            context = result.resolvedObject;
+        } else {
+            throw new IllegalArgumentException(result.errorMessage);
         }
 
         String leafKey = fieldPath.pathElements[fieldPath.pathElements.length - 1];
         if (context == null) {
             throw new IllegalArgumentException(Errors.cannotRemove(path, leafKey, null));
+        } else if (context instanceof IngestCtxMap map) { // optimization: handle IngestCtxMap separately from Map
+            if (map.containsKey(leafKey)) {
+                map.remove(leafKey);
+            } else {
+                throw new IllegalArgumentException(Errors.notPresent(path, leafKey));
+            }
         } else if (context instanceof Map<?, ?> map) {
             if (map.containsKey(leafKey)) {
                 map.remove(leafKey);
@@ -357,33 +362,48 @@ public void removeField(String path) {
         }
     }
 
-    private static ResolveResult resolve(String pathElement, String fullPath, Object context) {
-        if (context == null) {
-            return ResolveResult.error(Errors.cannotResolve(fullPath, pathElement, null));
-        } else if (context instanceof Map<?, ?>) {
-            @SuppressWarnings("unchecked")
-            Map<String, Object> map = (Map<String, Object>) context;
-            Object object = map.getOrDefault(pathElement, NOT_FOUND); // getOrDefault is faster than containsKey + get
-            if (object == NOT_FOUND) {
-                return ResolveResult.error(Errors.notPresent(fullPath, pathElement));
-            } else {
-                return ResolveResult.success(object);
-            }
-        } else if (context instanceof List<?> list) {
-            int index;
-            try {
-                index = Integer.parseInt(pathElement);
-            } catch (NumberFormatException e) {
-                return ResolveResult.error(Errors.notInteger(fullPath, pathElement));
-            }
-            if (index < 0 || index >= list.size()) {
-                return ResolveResult.error(Errors.outOfBounds(fullPath, index, list.size()));
+    /**
+     * Resolves the path elements (up to the limit) within the context. The result of such resolution can either be successful,
+     * or can indicate a failure.
+     */
+    private static ResolveResult resolve(final String[] pathElements, final int limit, final String fullPath, Object context) {
+        for (int i = 0; i < limit; i++) {
+            String pathElement = pathElements[i];
+            if (context == null) {
+                return ResolveResult.error(Errors.cannotResolve(fullPath, pathElement, null));
+            } else if (context instanceof IngestCtxMap map) { // optimization: handle IngestCtxMap separately from Map
+                Object object = map.getOrDefault(pathElement, NOT_FOUND); // getOrDefault is faster than containsKey + get
+                if (object == NOT_FOUND) {
+                    return ResolveResult.error(Errors.notPresent(fullPath, pathElement));
+                } else {
+                    context = object;
+                }
+            } else if (context instanceof Map<?, ?>) {
+                @SuppressWarnings("unchecked")
+                Map<String, Object> map = (Map<String, Object>) context;
+                Object object = map.getOrDefault(pathElement, NOT_FOUND); // getOrDefault is faster than containsKey + get
+                if (object == NOT_FOUND) {
+                    return ResolveResult.error(Errors.notPresent(fullPath, pathElement));
+                } else {
+                    context = object;
+                }
+            } else if (context instanceof List<?> list) {
+                int index;
+                try {
+                    index = Integer.parseInt(pathElement);
+                } catch (NumberFormatException e) {
+                    return ResolveResult.error(Errors.notInteger(fullPath, pathElement));
+                }
+                if (index < 0 || index >= list.size()) {
+                    return ResolveResult.error(Errors.outOfBounds(fullPath, index, list.size()));
+                } else {
+                    context = list.get(index);
+                }
             } else {
-                return ResolveResult.success(list.get(index));
+                return ResolveResult.error(Errors.cannotResolve(fullPath, pathElement, context));
             }
-        } else {
-            return ResolveResult.error(Errors.cannotResolve(fullPath, pathElement, context));
         }
+        return ResolveResult.success(context);
     }
 
     /**
@@ -518,6 +538,15 @@ private void setFieldValue(String path, Object value, boolean append, boolean al
             String pathElement = fieldPath.pathElements[i];
             if (context == null) {
                 throw new IllegalArgumentException(Errors.cannotResolve(path, pathElement, null));
+            } else if (context instanceof IngestCtxMap map) { // optimization: handle IngestCtxMap separately from Map
+                Object object = map.getOrDefault(pathElement, NOT_FOUND); // getOrDefault is faster than containsKey + get
+                if (object == NOT_FOUND) {
+                    Map<Object, Object> newMap = new HashMap<>();
+                    map.put(pathElement, newMap);
+                    context = newMap;
+                } else {
+                    context = object;
+                }
             } else if (context instanceof Map<?, ?>) {
                 @SuppressWarnings("unchecked")
                 Map<String, Object> map = (Map<String, Object>) context;
@@ -549,6 +578,22 @@ private void setFieldValue(String path, Object value, boolean append, boolean al
         String leafKey = fieldPath.pathElements[fieldPath.pathElements.length - 1];
         if (context == null) {
             throw new IllegalArgumentException(Errors.cannotSet(path, leafKey, null));
+        } else if (context instanceof IngestCtxMap map) { // optimization: handle IngestCtxMap separately from Map
+            if (append) {
+                Object object = map.getOrDefault(leafKey, NOT_FOUND); // getOrDefault is faster than containsKey + get
+                if (object == NOT_FOUND) {
+                    List<Object> list = new ArrayList<>();
+                    appendValues(list, value);
+                    map.put(leafKey, list);
+                } else {
+                    Object list = appendValues(object, value, allowDuplicates);
+                    if (list != object) {
+                        map.put(leafKey, list);
+                    }
+                }
+                return;
+            }
+            map.put(leafKey, value);
         } else if (context instanceof Map<?, ?>) {
             @SuppressWarnings("unchecked")
             Map<String, Object> map = (Map<String, Object>) context;