SQL: Resolve attributes recursively for improved subquery support

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeMap.java

@@ -6,6 +6,8 @@
  */
 package org.elasticsearch.xpack.ql.expression;
 
+import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
+
 import java.util.AbstractSet;
 import java.util.Collection;
 import java.util.Iterator;
@@ -266,15 +268,25 @@ public E getOrDefault(Object key, E defaultValue) {
         return defaultValue;
     }
 
-    @SuppressWarnings({ "unchecked" })
+    public E resolve(Object key) {
+        return resolve(key, null);
+    }
+
     public E resolve(Object key, E defaultValue) {
-        AttributeMap<Object> map = (AttributeMap<Object>) this;
-        Object candidate = defaultValue;
-        Object value = null;
-        while ((value = map.getOrDefault(key, NOT_FOUND)) != NOT_FOUND && value != key) {
-            key = candidate = value;
+        E value = defaultValue;
+        E candidate = null;
+        int allowedLookups = 1000;
+        while ((candidate = get(key)) != null || containsKey(key)) {
+            // instead of circling around, return
+            if (candidate == key) {
+                return candidate;
+            }
+            if (--allowedLookups == 0) {
+                throw new QlIllegalArgumentException("Potential cycle detected");
+            }
+            key = value = candidate;
         }
-        return (E) candidate;
+        return value;
     }
 
     @Override

x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java

@@ -7,6 +7,7 @@
 package org.elasticsearch.xpack.ql.expression;
 
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 
@@ -77,14 +78,48 @@ public void testResolve() {
         builder.put(threeAliasAlias.toAttribute(), threeAliasAlias.child());
         AttributeMap<Object> map = builder.build();
 
-        assertEquals(of("one"), map.resolve(one, null));
-        assertEquals("two", map.resolve(two, null));
-        assertEquals(of("three"), map.resolve(three, null));
-        assertEquals(of("three"), map.resolve(threeAlias, null));
-        assertEquals(of("three"), map.resolve(threeAliasAlias, null));
+        assertEquals(of("one"), map.resolve(one));
+        assertEquals("two", map.resolve(two));
+        assertEquals(of("three"), map.resolve(three));
+        assertEquals(of("three"), map.resolve(threeAlias));
+        assertEquals(of("three"), map.resolve(threeAliasAlias));
+        assertEquals(of("three"), map.resolve(threeAliasAlias, threeAlias));
         Attribute four = a("four");
         assertEquals("not found", map.resolve(four, "not found"));
-        assertNull(map.resolve(four, null));
+        assertNull(map.resolve(four));
+        assertEquals(four, map.resolve(four, four));
+    }
+
+    public void testResolveOneHopCycle() {
+        AttributeMap.Builder<Object> builder = AttributeMap.builder();
+        Attribute a = fieldAttribute("a", DataTypes.INTEGER);
+        Attribute b = fieldAttribute("b", DataTypes.INTEGER);
+        builder.put(a, a);
+        builder.put(b, a);
+        AttributeMap<Object> map = builder.build();
+
+        assertEquals(a, map.resolve(a, "default"));
+        assertEquals(a, map.resolve(b, "default"));
+        assertEquals("default", map.resolve("non-existing-key", "default"));
+    }
+
+    public void testResolveMultiHopCycle() {
+        AttributeMap.Builder<Object> builder = AttributeMap.builder();
+        Attribute a = fieldAttribute("a", DataTypes.INTEGER);
+        Attribute b = fieldAttribute("b", DataTypes.INTEGER);
+        Attribute c = fieldAttribute("c", DataTypes.INTEGER);
+        Attribute d = fieldAttribute("d", DataTypes.INTEGER);
+        builder.put(a, b);
+        builder.put(b, c);
+        builder.put(c, d);
+        builder.put(d, a);
+        AttributeMap<Object> map = builder.build();
+
+        // note: multi hop cycles should not happen, unless we have a 
+        // bug in the code that populates the AttributeMaps
+        expectThrows(QlIllegalArgumentException.class, () -> {
+            assertEquals(a, map.resolve(a, c));
+        });
     }
 
     private Alias createIntParameterAlias(int index, int value) {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java

@@ -411,7 +411,7 @@ private static boolean checkGroupByHavingHasOnlyAggs(Expression e, Set<Expressio
 
         // resolve FunctionAttribute to backing functions
         if (e instanceof ReferenceAttribute) {
-            e = attributeRefs.resolve(e, null);
+            e = attributeRefs.resolve(e);
         }
 
         // scalar functions can be a binary tree
@@ -584,7 +584,7 @@ private static boolean checkGroupMatch(Expression e, Node<?> source, List<Expres
 
         // resolve FunctionAttribute to backing functions
         if (e instanceof ReferenceAttribute) {
-            e = attributeRefs.resolve(e, null);
+            e = attributeRefs.resolve(e);
         }
 
         // scalar functions can be a binary tree

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

@@ -300,7 +300,7 @@ static class PruneOrderByNestedFields extends OptimizerRule<Project> {
         private void findNested(Expression exp, AttributeMap<Function> functions, Consumer<FieldAttribute> onFind) {
             exp.forEachUp(e -> {
                 if (e instanceof ReferenceAttribute) {
-                    Function f = functions.resolve(e, null);
+                    Function f = functions.resolve(e);
                     if (f != null) {
                         findNested(f, functions, onFind);
                     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java

@@ -710,7 +710,7 @@ protected PhysicalPlan rule(OrderExec plan) {
 
                     // if it's a reference, get the target expression
                     if (orderExpression instanceof ReferenceAttribute) {
-                        orderExpression = qContainer.aliases().resolve(orderExpression, null);
+                        orderExpression = qContainer.aliases().resolve(orderExpression);
                     }
                     String lookup = Expressions.id(orderExpression);
                     GroupByKey group = qContainer.findGroupForAgg(lookup);

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java

@@ -376,7 +376,7 @@ static Query rewriteToContainNestedField(@Nullable Query query, Source source, S
     // replace function/operators's input with references
     private Tuple<QueryContainer, FieldExtraction> resolvedTreeComputingRef(ScalarFunction function, Attribute attr) {
         Pipe proc = null;
-        if ((proc = scalarFunctions.resolve(attr, null)) == null) {
+        if ((proc = scalarFunctions.resolve(attr)) == null) {
             proc = function.asPipe();
             scalarFunctions = AttributeMap.builder(scalarFunctions).put(attr, proc).build();
         }