Resolve references recursively

Code simplification in the AttributeMap

Author: Andras Palinkas

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

@@ -238,30 +238,52 @@ 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);
     }
 
+    /**
+     * @param key {@link NamedExpression} to look up
+     * @return Looks up the key in the AttributeMap recursively, meaning if the value for this key
+     * is another key in the map, it will follow it. It will return the final value or null if the key
+     * does not exist in the map.
+     */
     @Override
     public E get(Object key) {
-        if (key instanceof NamedExpression) {
-            return delegate.get(new AttributeWrapper(((NamedExpression) key).toAttribute()));
-        }
-        return null;
+        return getOrDefault(key, null);
     }
 
+    /**
+     * @param key {@link NamedExpression} to look up
+     * @return Looks up the key in the AttributeMap recursively, meaning if the value for this key
+     * is another key in the map, it will follow it. It will return the final value or the
+     * specified default value if the key does not exist in the map.
+     */
     @Override
     public E getOrDefault(Object key, E defaultValue) {
-        E e;
-        return (((e = get(key)) != null) || containsKey(key))
-            ? e
-            : defaultValue;
+        E candidate = defaultValue;
+        E value = null;
+        while (((value = lookup(key)) != null || containsKey(key)) && value != key) {
+            key = candidate = value;
+        }
+        return candidate;
+    }
+
+    /**
+     * @param key {@link NamedExpression} to look up
+     * @return Result of the non-recursive lookup of the key in the map.
+     */
+    private E lookup(Object key) {
+        if (key instanceof NamedExpression) {
+            return delegate.get(new AttributeWrapper(((NamedExpression) key).toAttribute()));
+        }
+        return null;
     }
 
     @Override

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

@@ -17,6 +17,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;
@@ -61,6 +63,35 @@ public void testAttributeMapWithSameAliasesCanResolveAttributes() {
         assertTrue(newAttributeMap.get(param2.toAttribute()) == param2.child());
     }
 
+    public void testResolveRecursively() {
+        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.get(one));
+        assertEquals(map.get(one), map.getOrDefault(one, null));
+        assertEquals("two", map.get(two));
+        assertEquals(map.get(two), map.getOrDefault(two, null));
+        assertEquals(of("three"), map.get(three));
+        assertEquals(map.get(three), map.getOrDefault(three, null));
+        assertEquals(map.get(three), map.getOrDefault(threeAlias, null));
+        assertEquals(map.get(three), map.get(threeAlias));
+        assertEquals(map.get(three), map.getOrDefault(threeAliasAlias, null));
+        assertEquals(map.get(three), map.get(threeAliasAlias));
+        Attribute four = a("four");
+        assertEquals("not found", map.getOrDefault(four, "not found"));
+        assertNull(map.get(four));
+    }
+
     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,28 @@ 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;
 
+// 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/Verifier.java

@@ -317,9 +317,10 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
 
                 o.order().forEach(oe -> {
                     Expression e = oe.child();
+                    e = attributeRefs.getOrDefault(e, e);
 
                     // aggregates are allowed
-                    if (Functions.isAggregate(attributeRefs.getOrDefault(e, e))) {
+                    if (Functions.isAggregate(e)) {
                         return;
                     }
 
@@ -341,7 +342,7 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
                     //
                     // Also, make sure to compare attributes directly
                     if (e.anyMatch(expression -> Expressions.anyMatch(groupingAndMatchingAggregatesAliases,
-                        g -> expression.semanticEquals(expression instanceof Attribute ? Expressions.attribute(g) : g)))) {
+                        g -> expression.semanticEquals(attributeRefs.getOrDefault(Expressions.attribute(g), g))))) {
                         return;
                     }
 

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

@@ -131,8 +131,8 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 );
 
         Batch refs = new Batch("Replace References", Limiter.ONCE,
-                new ReplaceReferenceAttributeWithSource()
-                );
+            new ReplaceReferenceAttributeWithSource()
+        );
 
         Batch operators = new Batch("Operator Optimization",
                 // combining
@@ -222,7 +222,7 @@ 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();
+            final AttributeMap<Expression> collectRefs = builder.build();
             java.util.function.Function<ReferenceAttribute, Expression> replaceReference = r -> collectRefs.getOrDefault(r, r);
 
             plan = plan.transformUp(p -> {

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

@@ -600,10 +600,12 @@ else if (target.foldable()) {
                 else {
                     GroupByKey matchingGroup = null;
                     if (groupingContext != null) {
-                        matchingGroup = groupingContext.groupFor(target);
+                        final Expression resolvedTarget = queryC.aliases().getOrDefault(target, target);
+                        matchingGroup = groupingContext.groupFor(resolvedTarget);
                         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())), Expressions.id(resolvedTarget));
                     }
                     // fallback
                     else {

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerifierErrorMessagesTests.java

@@ -517,7 +517,14 @@ public void testGroupByOrderByScalarOverNonGrouped_WithHaving() {
 
     public void testGroupByHavingNonGrouped() {
         assertEquals("1:48: Cannot use HAVING filter on non-aggregate [int]; use WHERE instead",
-                error("SELECT AVG(int) FROM test GROUP BY text HAVING int > 10"));
+            error("SELECT AVG(int) FROM test GROUP BY bool HAVING int > 10"));
+        accept("SELECT AVG(int) FROM test GROUP BY bool HAVING AVG(int) > 2");
+    }
+    
+    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/69758")
+    public void testGroupByWhereSubselect() {
+        accept("SELECT b, a FROM (SELECT bool as b, AVG(int) as a FROM test GROUP BY bool) WHERE b = false");
+        accept("SELECT b, a FROM (SELECT bool as b, AVG(int) as a FROM test GROUP BY bool HAVING AVG(int) > 2) WHERE b = false");
     }
 
     public void testGroupByAggregate() {

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryTranslatorTests.java

@@ -2599,8 +2599,7 @@ public void testSubqueryFilterOrderByAlias() throws Exception {
             "WHERE i IS NOT NULL " +
             "ORDER BY i");
     }
-
-    @AwaitsFix(bugUrl = "follow-up to https://github.com/elastic/elasticsearch/pull/67216")
+    
     public void testSubqueryGroupByFilterAndOrderByByAlias() throws Exception {
         PhysicalPlan p = optimizeAndPlan("SELECT i FROM " +
             "( SELECT int AS i FROM test ) " +
@@ -2658,4 +2657,29 @@ public void testSubqueryWithAliasOrderByAlias() throws Exception {
             "( SELECT int AS i FROM test ) AS s " +
             "ORDER BY s.i > 10");
     }
+    
+    public void testReferenceResolutionInSubqueries() {
+        optimizeAndPlan("SELECT i AS j FROM ( SELECT int AS i FROM test) ORDER BY j");
+        optimizeAndPlan("SELECT j AS k FROM (SELECT i AS j FROM ( SELECT int AS i FROM test)) ORDER BY k");
+        optimizeAndPlan("SELECT int_group AS g, min_date AS d " +
+            "FROM (" + 
+            "    SELECT int % 2 AS int_group, MIN(date) AS min_date " +
+            "    FROM test WHERE date > '1970-01-01'::datetime GROUP BY int_group" + 
+            ") " +
+            "ORDER BY d DESC");
+        optimizeAndPlan("SELECT int_group AS g, min_date AS d " +
+            "FROM (" + 
+            "    SELECT int % 2 AS int_group, MIN(date) AS min_date " +
+            "    FROM test WHERE date > '1970-01-01'::datetime GROUP BY int_group " + 
+            ")" +
+            "ORDER BY g DESC");
+        optimizeAndPlan("SELECT i AS j FROM ( SELECT int AS i FROM test) GROUP BY j");
+        optimizeAndPlan("SELECT j AS k FROM (SELECT i AS j FROM ( SELECT int AS i FROM test)) GROUP BY k");
+        optimizeAndPlan("SELECT g FROM (SELECT date AS f, int AS g FROM test) WHERE g IS NOT NULL GROUP BY g ORDER BY g ASC");
+    }
+    
+    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/69758")
+    public void testFilterAfterGroupBy() {
+        optimizeAndPlan("SELECT j AS k FROM (SELECT i AS j FROM ( SELECT int AS i FROM test) GROUP BY j) WHERE j < 5");
+    }
 }