Add painless method getByPath, get value from nested collections with…

… dotted path

Given a nested structure composed of Lists and Maps, getByPath will return the value
keyed by path.  getByPath is a method on Lists and Maps.

The path is string Map keys and integer List indices separated by dot. An optional third
argument returns a default value if the path lookup fails due to a missing value.

Eg.
['key0': ['a', 'b'], 'key1': ['c', 'd']].getByPath('key1') = ['c', 'd']
['key0': ['a', 'b'], 'key1': ['c', 'd']].getByPath('key1.0') = 'c'
['key0': ['a', 'b'], 'key1': ['c', 'd']].getByPath('key2', 'x') = 'x'
[['key0': 'value0'], ['key1': 'value1']].getByPath('1.key1') = 'value1'

Throws IllegalArgumentException if an item cannot be found and a default is not given.
Throws NumberFormatException if a path element operating on a List is not an integer.

Fixes elastic#42769

modules/lang-painless/src/main/java/org/elasticsearch/painless/api/Augmentation.java

@@ -25,6 +25,7 @@
 import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.function.BiConsumer;
@@ -552,4 +553,78 @@ public static String[] splitOnToken(String receiver, String token, int limit) {
         // O(N) or faster depending on implementation
         return result.toArray(new String[0]);
     }
+
+    /**
+     * Access values in nested containers with a dot separated path.  Path elements are treated
+     * as strings for Maps and integers for Lists.  Throws 'IllegalArgumentException' if path does
+     * not exist. Throws 'NumberFormatException' if an path element for a List is not an integer.
+     */
+    private static Object getByPath(Object receiver, String path) {
+        Object value = getByPath(receiver, path, new missingValue());
+        if (value instanceof missingValue) {
+            throw new IllegalArgumentException("Could not find value at path [" + path + "]");
+        }
+        return value;
+    }
+
+    /**
+     * Same as 'getByPath(object, String)' but returns defaultValue if path does not exist.
+     */
+    private static Object getByPath(Object receiver, String path, Object defaultValue) {
+        Object next = receiver;
+        String[] elements = path.split("\\.");
+        for (int i = 0; i < elements.length; i ++) {
+            String element = elements[i];
+            if (next instanceof Map && (((Map) next).containsKey(element))) {
+                next = ((Map) next).get(element);
+            } else if (next instanceof List<?>) {
+                List<?> list = (List<?>)next;
+                try {
+                    int elemInt = Integer.parseInt(element);
+                    if (list.size() < elemInt) {
+                        return defaultValue;
+                    }
+                    next = list.get(elemInt);
+                } catch (NumberFormatException e) {
+                    String format = "Could not parse [%s] as a int index into list at path [%s] and index [%d]";
+                    throw new NumberFormatException(String.format(Locale.ROOT, format, element, path, i));
+                }
+            } else {
+                return defaultValue;
+            }
+        }
+        return next;
+    }
+
+
+    /**
+     * getByPath for Lists without default
+     */
+    public static <T> Object getByPath(List<T> receiver, String path) {
+        return getByPath((Object)receiver, path);
+    }
+
+    /**
+     * getByPath for Maps without default
+     */
+    public static <K,V> Object getByPath(Map<K,V> receiver, String path) {
+        return getByPath((Object)receiver, path);
+    }
+
+    /**
+     * getByPath for Lists with default
+     */
+    public static <T> Object getByPath(List<T> receiver, String path, Object defaultValue) {
+        return getByPath((Object)receiver, path, defaultValue);
+    }
+
+    /**
+     * getByPath for Maps with default
+     */
+    public static <K,V> Object getByPath(Map<K,V> receiver, String path, Object defaultValue) {
+        return getByPath((Object)receiver, path, defaultValue);
+    }
+
+    // missingValue indicates a value is missing in getByPath, used by 'getByPath(object, String)'
+    private static class missingValue {}
 }

modules/lang-painless/src/main/resources/org/elasticsearch/painless/spi/java.util.txt

@@ -126,6 +126,8 @@ class java.util.List {
   int org.elasticsearch.painless.api.Augmentation getLength()
   void sort(Comparator)
   List subList(int,int)
+  Object org.elasticsearch.painless.api.Augmentation getByPath(String)
+  Object org.elasticsearch.painless.api.Augmentation getByPath(String, Object)
 }
 
 class java.util.ListIterator {
@@ -161,6 +163,8 @@ class java.util.Map {
   void replaceAll(BiFunction)
   int size()
   Collection values()
+  Object org.elasticsearch.painless.api.Augmentation getByPath(String)
+  Object org.elasticsearch.painless.api.Augmentation getByPath(String, Object)
 
   # some adaptations of groovy methods
   List org.elasticsearch.painless.api.Augmentation collect(BiFunction)

modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java

@@ -20,8 +20,11 @@
 package org.elasticsearch.painless;
 
 import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.regex.Pattern;
 
@@ -232,11 +235,134 @@ public void testString_SplitOnToken() {
             new SplitCase("1\n1.1.\r\n1\r\n111", "\r\n"),
         };
         for (SplitCase split : cases) {
-            //System.out.println(String.format("Splitting '%s' by '%s' %d times", split.input, split.token, split.count));
             assertArrayEquals(
                 split.input.split(Pattern.quote(split.token), split.count),
                 (String[])exec("return \""+split.input+"\".splitOnToken(\""+split.token+"\", "+split.count+");")
             );
         }
     }
+
+    private static class GetByPathCase {
+        final String collection;
+        final String key;
+        final Object value;
+        final String defaultValue;
+        final boolean isDefaultValue;
+        final IllegalArgumentException illegalError;
+        final NumberFormatException numberError;
+
+        GetByPathCase(String collection, String key, Object value) {
+            this(collection, key, value, null, false, null);
+        }
+
+        GetByPathCase(String collection, String key, Exception error) {
+            this(collection, key, null, null, false, error);
+        }
+
+        GetByPathCase(String collection, String key, Object value, String defaultValue) {
+            this(collection, key, value, defaultValue, true, null);
+        }
+
+        GetByPathCase(String collection, String key, String defaultValue, Exception error) {
+            this(collection, key, null, defaultValue, true, error);
+        }
+
+        private GetByPathCase(String collection, String key, Object value, String defaultValue, boolean isDefaultValue, Exception error) {
+            this.collection = collection;
+            this.key = key;
+            this.value = value;
+            this.defaultValue = defaultValue;
+            this.isDefaultValue = isDefaultValue;
+            if (error == null) {
+                this.numberError = null;
+                this.illegalError = null;
+            } else {
+                if (error instanceof NumberFormatException) {
+                    this.numberError = (NumberFormatException)error;
+                    this.illegalError = null;
+                } else if (error instanceof IllegalArgumentException) {
+                    this.illegalError = (IllegalArgumentException)error;
+                    this.numberError = null;
+                } else {
+                    this.numberError = null;
+                    this.illegalError = null;
+                }
+            }
+        }
+
+        String Script() {
+            String format = "return %s.getByPath('%s'%s)";
+            String defaultValue = (this.isDefaultValue ? ", " + this.defaultValue : "");
+            return String.format(Locale.ROOT, format, this.collection, this.key, defaultValue);
+        }
+    }
+
+    private NumberFormatException numberFormat(String unparsable, String path, int index) {
+        String format = "Could not parse [%s] as a int index into list at path [%s] and index [%d]";
+        return new NumberFormatException(String.format(Locale.ROOT, format, unparsable, path, index));
+    }
+
+    private IllegalArgumentException illegalArg(String path) {
+        String format = "Could not find value at path [%s]";
+        return new IllegalArgumentException(String.format(Locale.ROOT, format, path));
+    }
+
+    public void testGetByPath() {
+        Map<String,String> k001 = new HashMap<>();
+        String k001Key = "k011";
+        k001.put(k001Key, "b");
+
+        List<String> ll2 = new ArrayList<>();
+        ll2.add("ll0");
+        String ll2Index1 = "ll1";
+        ll2.add(ll2Index1);
+
+        List<String> k1List = new ArrayList<>();
+        k1List.add("c");
+        k1List.add("d");
+
+        String mapMapList = "['k0': ['k01': [['k010': 'a'], ['k011': 'b']]], 'k1': ['q']]";
+        String listMapListList = "[['m0':'v0'],['m1':'v1'],['m2':['l0','l1', ['ll0', 'll1']]]]";
+
+        GetByPathCase[] cases = new GetByPathCase[] {
+          // basic
+          new GetByPathCase("['k0':'v0']", "k0", "v0"),
+          new GetByPathCase("['a','b','c','d']", "2", "c"),
+          new GetByPathCase("[['a','b'],['c','d']]", "1.k0", "'c'", numberFormat("k0", "1.k0", 1)),
+          new GetByPathCase("[['a','b'],['c','d']]", "1.k0", numberFormat("k0", "1.k0", 1)),
+          new GetByPathCase(mapMapList, "k0.k01.1.k012", illegalArg("k0.k01.1.k012")),
+
+          // nesting, map first
+          new GetByPathCase("['key0': ['a', 'b'], 'key1': ['c', 'd']]", "key1.0", "c"),
+          new GetByPathCase("['key0': ['a', 'b'], 'key1': ['c', 'd']]", "key1", k1List),
+          new GetByPathCase("['key0': ['a', 'b'], 'key1': ['c', 'd']]", "key2", "x", "'x'"),
+          new GetByPathCase("['k0': ['a','b','c','d'], 'k1': ['q']]", "k0.3", "d"),
+          new GetByPathCase(mapMapList, "k0.k01.1", k001),
+          new GetByPathCase(mapMapList, "k0.k01.1.k011", k001.get(k001Key)),
+          new GetByPathCase(mapMapList, "k0.k01.1.k012", "foo", "'foo'"),
+          new GetByPathCase(mapMapList, "k0.k01.1.k012", illegalArg("k0.k01.1.k012")),
+          new GetByPathCase(mapMapList, "k0.k01.k012", numberFormat("k012", "k0.k01.k012", 2)),
+          new GetByPathCase(mapMapList, "k0.k01.k012", "'q'", numberFormat("k012", "k0.k01.k012", 2)),
+
+          // nesting, list first
+          new GetByPathCase("[['key0': 'value0'], ['key1': 'value1']]", "1.key1", "value1"),
+          new GetByPathCase("[['a','b'],['c','d'],[['e','f'],['g','h']]]", "2.1.1", "h"),
+          new GetByPathCase(listMapListList, "2.m2.2.1", ll2Index1),
+          new GetByPathCase(listMapListList, "2.m2.2", ll2),
+          new GetByPathCase(listMapListList, "2.m2.12", illegalArg("2.m2.12")),
+          new GetByPathCase(listMapListList, "2.m2.a8", numberFormat("a8", "2.m2.a8", 2)),
+          new GetByPathCase(listMapListList, "2.m2.a8", "'r'", numberFormat("a8", "2.m2.a8", 2)),
+        };
+        for (GetByPathCase getCase: cases) {
+            if (getCase.illegalError != null) {
+                IllegalArgumentException illegal = expectScriptThrows(IllegalArgumentException.class, () -> exec(getCase.Script()));
+                assertEquals(getCase.illegalError.getMessage(), illegal.getMessage());
+            } else if (getCase.numberError != null) {
+                NumberFormatException number = expectScriptThrows(NumberFormatException.class, () -> exec(getCase.Script()));
+                assertEquals(getCase.numberError.getMessage(), number.getMessage());
+            } else {
+                assertEquals(getCase.value, exec(getCase.Script()));
+            }
+        }
+    }
 }