SQL: Resolve attributes recursively for improved subquery support (#69765) (#70325)

Previously we did not resolve the attributes recursively which meant that if a field or expression was re-aliased multiple times (through multiple levels of subqueries), the aliases were only resolved one level down. This led to failed query translation because `ReferenceAttribute`s were pointing to non-existing attributes during query translation.

For example the query

```sql
SELECT i AS j FROM ( SELECT int AS i FROM test) ORDER BY j
```

failed during translation because the `OrderBy` resolved the `j` ReferenceAttribute to another `i` ReferenceAttribute that was later removed by an Optimization:

```
OrderBy[[Order[j{r}#4,ASC,LAST]]]                                             ! OrderBy[[Order[i{r}#2,ASC,LAST]]]
\_Project[[j]]                                                                = \_Project[[j]]
  \_Project[[i]]                                                              !   \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..]
    \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..] ! 
```

By resolving the `Attributes` recursively both `j{r}` and `i{r}` will resolve to `test.int{f}` above:

```
OrderBy[[Order[test.int{f}#22,ASC,LAST]]]                                     = OrderBy[[Order[test.int{f}#22,ASC,LAST]]]
\_Project[[j]]                                                                = \_Project[[j]]
  \_Project[[i]]                                                              !   \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..]
    \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..] ! 
 ```

The scope of recursive resolution depends on how the `AttributeMap` is constructed and populated.

Fixes #67237

Author: Andras Palinkas

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;
@@ -148,6 +150,8 @@ public static final <E> AttributeMap<E> emptyAttributeMap() {
         return EMPTY;
     }
 
+    private static final Object NOT_FOUND = new Object();
+
     private final Map<AttributeWrapper, E> delegate;
     private Set<Attribute> keySet = null;
     private Collection<E> values = null;
@@ -238,14 +242,14 @@ public boolean isEmpty() {
     @Override
     public boolean containsKey(Object key) {
         if (key instanceof NamedExpression) {
-            return delegate.keySet().contains(new AttributeWrapper(((NamedExpression) key).toAttribute()));
+            return delegate.containsKey(new AttributeWrapper(((NamedExpression) key).toAttribute()));
         }
         return false;
     }
 
     @Override
     public boolean containsValue(Object value) {
-        return delegate.values().contains(value);
+        return delegate.containsValue(value);
     }
 
     @Override
@@ -258,10 +262,32 @@ public E get(Object key) {
 
     @Override
     public E getOrDefault(Object key, E defaultValue) {
-        E e;
-        return (((e = get(key)) != null) || containsKey(key))
-            ? e
-            : defaultValue;
+        if (key instanceof NamedExpression) {
+            return delegate.getOrDefault(new AttributeWrapper(((NamedExpression) key).toAttribute()), defaultValue);
+        }
+        return defaultValue;
+    }
+
+    public E resolve(Object key) {
+        return resolve(key, null);
+    }
+
+    public E resolve(Object key, E defaultValue) {
+        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 = candidate;
+            value = 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;
 
@@ -18,6 +19,8 @@
 import java.util.Set;
 
 import static java.util.stream.Collectors.toList;
+import static org.elasticsearch.xpack.ql.TestUtils.fieldAttribute;
+import static org.elasticsearch.xpack.ql.TestUtils.of;
 import static org.hamcrest.Matchers.arrayContaining;
 import static org.hamcrest.Matchers.arrayWithSize;
 import static org.hamcrest.Matchers.contains;
@@ -62,6 +65,64 @@ public void testAttributeMapWithSameAliasesCanResolveAttributes() {
         assertTrue(newAttributeMap.get(param2.toAttribute()) == param2.child());
     }
 
+    public void testResolve() {
+        AttributeMap.Builder<Object> builder = AttributeMap.builder();
+        Attribute one = a("one");
+        Attribute two = fieldAttribute("two", DataTypes.INTEGER);
+        Attribute three = fieldAttribute("three", DataTypes.INTEGER);
+        Alias threeAlias = new Alias(Source.EMPTY, "three_alias", three);
+        Alias threeAliasAlias = new Alias(Source.EMPTY, "three_alias_alias", threeAlias);
+        builder.put(one, of("one"));
+        builder.put(two, "two");
+        builder.put(three, of("three"));
+        builder.put(threeAlias.toAttribute(), threeAlias.child());
+        builder.put(threeAliasAlias.toAttribute(), threeAliasAlias.child());
+        AttributeMap<Object> map = builder.build();
+        
+        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));
+        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) {
         Source source = new Source(1, index * 5, "?");
         Literal literal = new Literal(source, value, DataTypes.INTEGER);

x-pack/plugin/sql/qa/server/src/main/resources/subselect.sql-spec

@@ -16,10 +16,34 @@ basicGroupBy
 SELECT gender FROM (SELECT first_name AS f, last_name, gender FROM test_emp) GROUP BY gender ORDER BY gender ASC;
 basicGroupByAlias
 SELECT g FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) GROUP BY g ORDER BY g ASC;
-// @AwaitsFix(bugUrl = "follow-up to https://github.com/elastic/elasticsearch/pull/67216")
-basicGroupByWithFilterByAlias-Ignore
+basicGroupByWithFilterByAlias
 SELECT g FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) WHERE g IS NOT NULL GROUP BY g ORDER BY g ASC;
+basicGroupByRealiased
+SELECT g AS h FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) GROUP BY g ORDER BY g DESC NULLS last;
+basicGroupByRealiasedTwice
+SELECT g AS h FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) GROUP BY h ORDER BY h DESC NULLS last;
+basicOrderByRealiasedField
+SELECT g AS h FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) ORDER BY g DESC NULLS first;
+
+groupAndOrderByRealiasedExpression
+SELECT emp_group AS e, min_high_salary AS s
+FROM (
+    SELECT emp_no % 2 AS emp_group, MIN(salary) AS min_high_salary 
+    FROM test_emp
+    WHERE salary > 50000 
+    GROUP BY emp_group
+)
+ORDER BY e DESC;
+
+multiLevelSelectStar
+SELECT * FROM (SELECT * FROM ( SELECT * FROM test_emp ));
+
+multiLevelSelectStarWithAlias
+SELECT * FROM (SELECT * FROM ( SELECT * FROM test_emp ) b) c;
 
+// AwaitsFix: https://github.com/elastic/elasticsearch/issues/69758
+filterAfterGroupBy-Ignore
+SELECT s2 AS s3 FROM (SELECT s AS s2 FROM ( SELECT salary AS s FROM test_emp) GROUP BY s2) WHERE s2 < 5 ORDER BY s3 DESC NULLS last;
 
 countAndComplexCondition
 SELECT COUNT(*) as c FROM (SELECT * FROM test_emp WHERE gender IS NOT NULL) WHERE ABS(salary) > 0 GROUP BY gender ORDER BY gender;

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

@@ -638,12 +638,12 @@ protected LogicalPlan doRule(LogicalPlan plan) {
                     AttributeMap.Builder<Expression> builder = AttributeMap.builder();
                     // collect aliases
                     child.forEachUp(p -> p.forEachExpressionUp(Alias.class, a -> builder.put(a.toAttribute(), a.child())));
-                    final Map<Attribute, Expression> collectRefs = builder.build();
+                    final AttributeMap<Expression> collectRefs = builder.build();
 
                     referencesStream = referencesStream.filter(r -> {
                         for (Attribute attr : child.outputSet()) {
                             if (attr instanceof ReferenceAttribute) {
-                                Expression source = collectRefs.getOrDefault(attr, attr);
+                                Expression source = collectRefs.resolve(attr, attr);
                                 // found a match, no need to resolve it further
                                 // so filter it out
                                 if (source.equals(r.child())) {

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

@@ -316,10 +316,11 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
                 Map<Expression, Node<?>> missing = new LinkedHashMap<>();
 
                 o.order().forEach(oe -> {
-                    Expression e = oe.child();
+                    final Expression e = oe.child();
+                    final Expression resolvedE = attributeRefs.resolve(e, e);
 
                     // aggregates are allowed
-                    if (Functions.isAggregate(attributeRefs.getOrDefault(e, e))) {
+                    if (Functions.isAggregate(resolvedE)) {
                         return;
                     }
 
@@ -340,8 +341,12 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
                     // e.g.: if "GROUP BY f2(f1(field))" you can "ORDER BY f4(f3(f2(f1(field))))"
                     //
                     // Also, make sure to compare attributes directly
-                    if (e.anyMatch(expression -> Expressions.anyMatch(groupingAndMatchingAggregatesAliases,
-                        g -> expression.semanticEquals(expression instanceof Attribute ? Expressions.attribute(g) : g)))) {
+                    if (resolvedE.anyMatch(expression -> Expressions.anyMatch(groupingAndMatchingAggregatesAliases,
+                        g -> {
+                            Expression resolvedG = attributeRefs.resolve(g, g);
+                            resolvedG = expression instanceof Attribute ? Expressions.attribute(resolvedG) : resolvedG;
+                            return expression.semanticEquals(resolvedG);
+                        }))) {
                         return;
                     }
 
@@ -406,7 +411,7 @@ private static boolean checkGroupByHavingHasOnlyAggs(Expression e, Set<Expressio
 
         // resolve FunctionAttribute to backing functions
         if (e instanceof ReferenceAttribute) {
-            e = attributeRefs.get(e);
+            e = attributeRefs.resolve(e);
         }
 
         // scalar functions can be a binary tree
@@ -484,7 +489,7 @@ private static boolean onlyRawFields(Iterable<? extends Expression> expressions,
 
         expressions.forEach(e -> e.forEachDown(c -> {
             if (c instanceof ReferenceAttribute) {
-                c = attributeRefs.getOrDefault(c, c);
+                c = attributeRefs.resolve(c, c);
             }
             if (c instanceof Function) {
                 localFailures.add(fail(c, "No functions allowed (yet); encountered [{}]", c.sourceText()));
@@ -579,7 +584,7 @@ private static boolean checkGroupMatch(Expression e, Node<?> source, List<Expres
 
         // resolve FunctionAttribute to backing functions
         if (e instanceof ReferenceAttribute) {
-            e = attributeRefs.get(e);
+            e = attributeRefs.resolve(e);
         }
 
         // scalar functions can be a binary tree
@@ -668,7 +673,7 @@ private static void checkFilterOnAggs(LogicalPlan p, Set<Failure> localFailures,
             LogicalPlan filterChild = filter.child();
             if (filterChild instanceof Aggregate == false) {
                 filter.condition().forEachDown(Expression.class, e -> {
-                    if (Functions.isAggregate(attributeRefs.getOrDefault(e, e))) {
+                    if (Functions.isAggregate(attributeRefs.resolve(e, e))) {
                         if (filterChild instanceof Project) {
                             filter.condition().forEachDown(FieldAttribute.class,
                                 f -> localFailures.add(fail(e, "[{}] field must appear in the GROUP BY clause or in an aggregate function",
@@ -690,7 +695,7 @@ private static void checkFilterOnGrouping(LogicalPlan p, Set<Failure> localFailu
         if (p instanceof Filter) {
             Filter filter = (Filter) p;
             filter.condition().forEachDown(Expression.class, e -> {
-                if (Functions.isGrouping(attributeRefs.getOrDefault(e, e))) {
+                if (Functions.isGrouping(attributeRefs.resolve(e, e))) {
                     localFailures
                         .add(fail(e, "Cannot filter on grouping function [{}], use its argument instead", Expressions.name(e)));
                 }
@@ -717,7 +722,7 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
             }
         };
         Consumer<Expression> checkForNested = e ->
-            attributeRefs.getOrDefault(e, e).forEachUp(FieldAttribute.class, matchNested);
+            attributeRefs.resolve(e, e).forEachUp(FieldAttribute.class, matchNested);
         Consumer<ScalarFunction> checkForNestedInFunction = f -> f.arguments().forEach(
             arg -> arg.forEachUp(FieldAttribute.class, matchNested));
 
@@ -739,7 +744,7 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
 
         // check in where (scalars not allowed)
         p.forEachDown(Filter.class, f -> f.condition().forEachUp(e ->
-            attributeRefs.getOrDefault(e, e).forEachUp(ScalarFunction.class, sf -> {
+            attributeRefs.resolve(e, e).forEachUp(ScalarFunction.class, sf -> {
                 if (sf instanceof BinaryComparison == false &&
                     sf instanceof IsNull == false &&
                     sf instanceof IsNotNull == false &&
@@ -758,7 +763,7 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
 
         // check in order by (scalars not allowed)
         p.forEachDown(OrderBy.class, ob -> ob.order().forEach(o -> o.forEachUp(e ->
-            attributeRefs.getOrDefault(e, e).forEachUp(ScalarFunction.class, checkForNestedInFunction)
+            attributeRefs.resolve(e, e).forEachUp(ScalarFunction.class, checkForNestedInFunction)
         )));
         if (nested.isEmpty() == false) {
             localFailures.add(

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

@@ -222,8 +222,8 @@ public LogicalPlan apply(LogicalPlan plan) {
             AttributeMap.Builder<Expression> builder = AttributeMap.builder();
             // collect aliases
             plan.forEachExpressionUp(Alias.class, a -> builder.put(a.toAttribute(), a.child()));
-            final Map<Attribute, Expression> collectRefs = builder.build();
-            java.util.function.Function<ReferenceAttribute, Expression> replaceReference = r -> collectRefs.getOrDefault(r, r);
+            final AttributeMap<Expression> collectRefs = builder.build();
+            java.util.function.Function<ReferenceAttribute, Expression> replaceReference = r -> collectRefs.resolve(r, r);
 
             plan = plan.transformUp(p -> {
                 // non attribute defining plans get their references removed
@@ -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.get(e);
+                    Function f = functions.resolve(e);
                     if (f != null) {
                         findNested(f, functions, onFind);
                     }
@@ -578,7 +578,7 @@ private List<NamedExpression> combineProjections(List<? extends NamedExpression>
             // replace any matching attribute with a lower alias (if there's a match)
             // but clean-up non-top aliases at the end
             for (NamedExpression ne : upper) {
-                NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.getOrDefault(a, a));
+                NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.resolve(a, a));
                 replaced.add((NamedExpression) CleanAliases.trimNonTopLevelAliases(replacedExp));
             }
             return replaced;

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

@@ -600,10 +600,13 @@ else if (target.foldable()) {
                 else {
                     GroupByKey matchingGroup = null;
                     if (groupingContext != null) {
+                        target = queryC.aliases().resolve(target, target);
+                        id = Expressions.id(target);
                         matchingGroup = groupingContext.groupFor(target);
                         Check.notNull(matchingGroup, "Cannot find group [{}]", Expressions.name(ne));
 
-                        queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, isDateBased(ne.dataType())), id);
+                        queryC = queryC.addColumn(
+                            new GroupByRef(matchingGroup.id(), null, isDateBased(ne.dataType())), id);
                     }
                     // fallback
                     else {
@@ -707,7 +710,7 @@ protected PhysicalPlan rule(OrderExec plan) {
 
                     // if it's a reference, get the target expression
                     if (orderExpression instanceof ReferenceAttribute) {
-                        orderExpression = qContainer.aliases().get(orderExpression);
+                        orderExpression = qContainer.aliases().resolve(orderExpression);
                     }
                     String lookup = Expressions.id(orderExpression);
                     GroupByKey group = qContainer.findGroupForAgg(lookup);