SQL: Extend the multi dot field notation extraction to lists of values (#39823)

(cherry picked from commit 300ae48)

Author: astefan

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/extractor/FieldHitExtractor.java

@@ -138,13 +138,12 @@ private Object unwrapMultiValue(Object values) {
         throw new SqlIllegalArgumentException("Type {} (returned by [{}]) is not supported", values.getClass().getSimpleName(), fieldName);
     }
 
-    @SuppressWarnings("unchecked")
+    @SuppressWarnings({ "unchecked", "rawtypes" })
     Object extractFromSource(Map<String, Object> map) {
         Object value = null;
 
         // Used to avoid recursive method calls
-        // Holds the sub-maps in the document hierarchy that are pending to be inspected.
-        // along with the current index of the `path`.
+        // Holds the sub-maps in the document hierarchy that are pending to be inspected along with the current index of the `path`.
         Deque<Tuple<Integer, Map<String, Object>>> queue = new ArrayDeque<>();
         queue.add(new Tuple<>(-1, map));
 
@@ -160,6 +159,19 @@ Object extractFromSource(Map<String, Object> map) {
             for (int i = idx + 1; i < path.length; i++) {
                 sj.add(path[i]);
                 Object node = subMap.get(sj.toString());
+                
+                if (node instanceof List) {
+                    List listOfValues = (List) node;
+                    if (listOfValues.size() == 1) {
+                        // this is a List with a size of 1 e.g.: {"a" : [{"b" : "value"}]} meaning the JSON is a list with one element
+                        // or a list of values with one element e.g.: {"a": {"b" : ["value"]}}
+                        node = listOfValues.get(0);
+                    } else {
+                        // a List of elements with more than one value. Break early and let unwrapMultiValue deal with the list
+                        return unwrapMultiValue(node);
+                    }
+                }
+                
                 if (node instanceof Map) {
                     if (i < path.length - 1) {
                         // Add the sub-map to the queue along with the current path index

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/execution/search/extractor/FieldHitExtractorTests.java

@@ -28,6 +28,7 @@
 import java.util.function.Supplier;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
 import static org.hamcrest.Matchers.is;
 
@@ -267,6 +268,7 @@ public void testNestedFieldWithDotsWithNestedFieldWithDots() {
         assertEquals(value, fe.extractFromSource(map));
     }
 
+    @SuppressWarnings({ "rawtypes", "unchecked" })
     public void testNestedFieldsWithDotsAndRandomHiearachy() {
         String[] path = new String[100];
         StringJoiner sj = new StringJoiner(".");
@@ -288,12 +290,42 @@ public void testNestedFieldsWithDotsAndRandomHiearachy() {
             start = end;
         }
 
+        /*
+         * Randomize how many values the field to look for will have (1 - 3). It's not really relevant how many values there are in the list
+         * but that the list has one element or more than one.
+         * If it has one value, then randomize the way it's indexed: as a single-value array or not e.g.: "a":"value" or "a":["value"].
+         * If it has more than one value, it will always be an array e.g.: "a":["v1","v2","v3"].
+         */
+        int valuesCount = randomIntBetween(1, 3);
         Object value = randomValue();
+        if (valuesCount == 1) {
+            value = randomBoolean() ? singletonList(value) : value;
+        } else {
+            value = new ArrayList(valuesCount);
+            for(int i = 0; i < valuesCount; i++) {
+                ((List) value).add(randomValue());
+            }
+        }
+
+        // the path to the randomly generated fields path
+        StringBuilder expected = new StringBuilder(paths.get(paths.size() - 1));
+        // the actual value we will be looking for in the test at the end
         Map<String, Object> map = singletonMap(paths.get(paths.size() - 1), value);
+        // build the rest of the path and the expected path to check against in the error message
         for (int i = paths.size() - 2; i >= 0; i--) {
-            map = singletonMap(paths.get(i), map);
+            map = singletonMap(paths.get(i), randomBoolean() ? singletonList(map) : map);
+            expected.insert(0, paths.get(i) + ".");
+        }
+
+        if (valuesCount == 1) {
+            // if the number of generated values is 1, just check we return the correct value
+            assertEquals(value instanceof List ? ((List) value).get(0) : value, fe.extractFromSource(map));
+        } else {
+            // if we have an array with more than one value in it, check that we throw the correct exception and exception message
+            final Map<String, Object> map2 = Collections.unmodifiableMap(map);
+            SqlException ex = expectThrows(SqlException.class, () -> fe.extractFromSource(map2));
+            assertThat(ex.getMessage(), is("Arrays (returned by [" + expected + "]) are not supported"));
         }
-        assertEquals(value, fe.extractFromSource(map));
     }
 
     public void testExtractSourceIncorrectPathWithFieldWithDots() {
@@ -335,6 +367,51 @@ public void testFieldWithDotsAndSamePathButDifferentHierarchy() {
         SqlException ex = expectThrows(SqlException.class, () -> fe.extractFromSource(map));
         assertThat(ex.getMessage(), is("Multiple values (returned by [a.b.c.d.e.f.g]) are not supported"));
     }
+    
+    public void testFieldsWithSingleValueArrayAsSubfield() {
+        FieldHitExtractor fe = new FieldHitExtractor("a.b", null, false);
+        Object value = randomNonNullValue();
+        Map<String, Object> map = new HashMap<>();
+        // "a" : [{"b" : "value"}]
+        map.put("a", singletonList(singletonMap("b", value)));
+        assertEquals(value, fe.extractFromSource(map));
+    }
+    
+    public void testFieldsWithMultiValueArrayAsSubfield() {
+        FieldHitExtractor fe = new FieldHitExtractor("a.b", null, false);
+        Map<String, Object> map = new HashMap<>();
+        // "a" : [{"b" : "value1"}, {"b" : "value2"}]
+        map.put("a", asList(singletonMap("b", randomNonNullValue()), singletonMap("b", randomNonNullValue())));
+        SqlException ex = expectThrows(SqlException.class, () -> fe.extractFromSource(map));
+        assertThat(ex.getMessage(), is("Arrays (returned by [a.b]) are not supported"));
+    }
+    
+    public void testFieldsWithSingleValueArrayAsSubfield_TwoNestedLists() {
+        FieldHitExtractor fe = new FieldHitExtractor("a.b.c", null, false);
+        Object value = randomNonNullValue();
+        Map<String, Object> map = new HashMap<>();
+        // "a" : [{"b" : [{"c" : "value"}]}]
+        map.put("a", singletonList(singletonMap("b", singletonList(singletonMap("c", value)))));
+        assertEquals(value, fe.extractFromSource(map));
+    }
+    
+    public void testFieldsWithMultiValueArrayAsSubfield_ThreeNestedLists() {
+        FieldHitExtractor fe = new FieldHitExtractor("a.b.c", null, false);
+        Map<String, Object> map = new HashMap<>();
+        // "a" : [{"b" : [{"c" : ["value1", "value2"]}]}]
+        map.put("a", singletonList(singletonMap("b", singletonList(singletonMap("c", asList("value1", "value2"))))));
+        SqlException ex = expectThrows(SqlException.class, () -> fe.extractFromSource(map));
+        assertThat(ex.getMessage(), is("Arrays (returned by [a.b.c]) are not supported"));
+    }
+    
+    public void testFieldsWithSingleValueArrayAsSubfield_TwoNestedLists2() {
+        FieldHitExtractor fe = new FieldHitExtractor("a.b.c", null, false);
+        Object value = randomNonNullValue();
+        Map<String, Object> map = new HashMap<>();
+        // "a" : [{"b" : {"c" : ["value"]}]}]
+        map.put("a", singletonList(singletonMap("b", singletonMap("c", singletonList(value)))));
+        assertEquals(value, fe.extractFromSource(map));
+    }
 
     public void testObjectsForSourceValue() throws IOException {
         String fieldName = randomAlphaOfLength(5);