SQL: Fix issue with aliased subqueries and GROUP BY (#73233) (#73403)

Previously, when a subquery was used with an alias in combination with
a nested GROUP BY, the collapsing of the nested queries into a flattened
`Aggregate` query, lead to wrong attribute qualifier on the external
projection, which was still referencing the removed subquery. e.g.:

For the following query:
```
SELECT languages FROM (
    SELECT languages FROM test_emp GROUP BY languages
) AS subquery
```
The `languages` of the top level SELECT, was qualified with `subquery`
which was removed during the flattening optimisation leading to
Exception of not being able to resolve the refenced group:
`test_emp.languages`.

Fix this behaviour by introducing a new rule which precedes the
`PruneSubqueryAliases` rules and updates the `qualifier` for the
`FieldAttributes`.

Fixes: #69263
(cherry picked from commit 1a2a3df)

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

@@ -155,3 +155,5 @@ selectGroupByOrderByOrderByLimit
 SELECT * FROM (SELECT max(salary) AS max, languages FROM test_emp GROUP BY languages ORDER BY max ASC) ORDER BY max DESC NULLS FIRST LIMIT 4;
 selectGroupByOrderByOrderByLimitNulls
 SELECT * FROM (SELECT max(salary) AS max, languages FROM test_emp GROUP BY languages ORDER BY max ASC NULLS LAST) ORDER BY max DESC NULLS FIRST LIMIT 4;
+selectGroupByWithAliasedSubQuery
+SELECT max, languages FROM (SELECT max(salary) AS max, languages FROM test_emp GROUP BY languages ORDER BY max ASC NULLS LAST) AS subquery;

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

@@ -124,7 +124,8 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 //new ImplicitCasting()
                 );
         Batch finish = new Batch("Finish Analysis",
-                new PruneSubqueryAliases(),
+                new ReplaceSubQueryAliases(), // Should be run before pruning SubqueryAliases
+                new PruneSubQueryAliases(),
                 new AddMissingEqualsToBoolField(),
                 CleanAliases.INSTANCE
                 );
@@ -1216,7 +1217,33 @@ private Expression implicitCast(Expression e) {
         }
     }
 
-    public static class PruneSubqueryAliases extends AnalyzerRule<SubQueryAlias> {
+    public static class ReplaceSubQueryAliases extends AnalyzerRule<UnaryPlan> {
+
+        @Override
+        protected LogicalPlan rule(UnaryPlan plan) {
+            if (plan.child() instanceof SubQueryAlias) {
+                SubQueryAlias a = (SubQueryAlias) plan.child();
+                return plan.transformExpressionsDown(FieldAttribute.class, f -> {
+                   if (f.qualifier() != null && f.qualifier().equals(a.alias())) {
+                       // Find the underlying concrete relation (EsIndex) and its name as the new qualifier
+                       List<LogicalPlan> children = a.collectFirstChildren(p -> p instanceof EsRelation);
+                       if (children.isEmpty() == false) {
+                           return f.withQualifier(((EsRelation) children.get(0)).index().name());
+                       }
+                   }
+                   return f;
+                });
+            }
+            return plan;
+        }
+
+        @Override
+        protected boolean skipResolved() {
+            return false;
+        }
+    }
+
+    public static class PruneSubQueryAliases extends AnalyzerRule<SubQueryAlias> {
 
         @Override
         protected LogicalPlan rule(SubQueryAlias alias) {

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java

@@ -54,7 +54,8 @@
 import org.elasticsearch.xpack.ql.type.EsField;
 import org.elasticsearch.xpack.ql.util.CollectionUtils;
 import org.elasticsearch.xpack.ql.util.StringUtils;
-import org.elasticsearch.xpack.sql.analysis.analyzer.Analyzer.PruneSubqueryAliases;
+import org.elasticsearch.xpack.sql.analysis.analyzer.Analyzer.ReplaceSubQueryAliases;
+import org.elasticsearch.xpack.sql.analysis.analyzer.Analyzer.PruneSubQueryAliases;
 import org.elasticsearch.xpack.sql.expression.function.aggregate.Avg;
 import org.elasticsearch.xpack.sql.expression.function.aggregate.ExtendedStats;
 import org.elasticsearch.xpack.sql.expression.function.aggregate.First;
@@ -144,6 +145,7 @@
 import static org.elasticsearch.xpack.sql.type.SqlDataTypes.DATE;
 import static org.elasticsearch.xpack.sql.util.DateUtils.UTC;
 import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 
 public class OptimizerTests extends ESTestCase {
@@ -174,13 +176,25 @@ private static FieldAttribute getFieldAttribute(String name) {
         return new FieldAttribute(EMPTY, name, new EsField(name + "f", INTEGER, emptyMap(), true));
     }
 
-    public void testPruneSubqueryAliases() {
+    public void testPruneSubQueryAliases() {
         ShowTables s = new ShowTables(EMPTY, null, null, false);
         SubQueryAlias plan = new SubQueryAlias(EMPTY, s, "show");
-        LogicalPlan result = new PruneSubqueryAliases().apply(plan);
+        LogicalPlan result = new PruneSubQueryAliases().apply(plan);
         assertEquals(result, s);
     }
 
+    public void testReplaceSubQueryAliases() {
+        FieldAttribute firstField = new FieldAttribute(EMPTY, "field", new EsField("field", BYTE, emptyMap(), true));
+        EsRelation relation = new EsRelation(EMPTY, new EsIndex("table", emptyMap()), false);
+        Aggregate agg = new Aggregate(EMPTY, relation, Collections.singletonList(firstField), Collections.singletonList(firstField));
+        SubQueryAlias subquery = new SubQueryAlias(EMPTY, agg, "subquery");
+        Project project = new Project(EMPTY, subquery, Collections.singletonList(firstField.withQualifier("subquery")));
+        LogicalPlan result = new ReplaceSubQueryAliases().apply(project);
+        assertThat(result, instanceOf(Project.class));
+        assertThat(((Project) result).projections().get(0), instanceOf(FieldAttribute.class));
+        assertEquals("table", ((FieldAttribute) ((Project) result).projections().get(0)).qualifier());
+    }
+
     public void testCombineProjections() {
         // a
         Alias a = new Alias(EMPTY, "a", FIVE);

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

@@ -28,7 +28,6 @@
 import org.hamcrest.Matcher;
 import org.junit.BeforeClass;
 
-import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -54,12 +53,8 @@ private static class TestContext {
             planner = new Planner();
         }
 
-        LogicalPlan plan(String sql, ZoneId zoneId) {
-            return analyzer.analyze(parser.createStatement(sql, zoneId), true);
-        }
-
         LogicalPlan plan(String sql) {
-            return plan(sql, DateUtils.UTC);
+            return analyzer.analyze(parser.createStatement(sql, DateUtils.UTC), true);
         }
 
         PhysicalPlan optimizeAndPlan(String sql) {

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

@@ -66,13 +66,15 @@
 import org.elasticsearch.xpack.sql.optimizer.Optimizer;
 import org.elasticsearch.xpack.sql.parser.SqlParser;
 import org.elasticsearch.xpack.sql.plan.physical.EsQueryExec;
+import org.elasticsearch.xpack.sql.plan.physical.LocalExec;
 import org.elasticsearch.xpack.sql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.sql.planner.QueryFolder.FoldAggregate.GroupingContext;
 import org.elasticsearch.xpack.sql.planner.QueryTranslator.QueryTranslation;
 import org.elasticsearch.xpack.sql.proto.SqlTypedParamValue;
 import org.elasticsearch.xpack.sql.querydsl.agg.AggFilter;
 import org.elasticsearch.xpack.sql.querydsl.agg.GroupByDateHistogram;
 import org.elasticsearch.xpack.sql.querydsl.container.MetricAggRef;
+import org.elasticsearch.xpack.sql.session.SingletonExecutable;
 import org.elasticsearch.xpack.sql.stats.Metrics;
 import org.elasticsearch.xpack.sql.types.SqlTypesTests;
 import org.elasticsearch.xpack.sql.util.DateUtils;
@@ -1288,4 +1290,41 @@ public void testAddMissingEqualsToNestedBoolField() {
 
         assertEquals(expectedCondition, condition);
     }
+
+    // Subqueries
+    /////////////////////
+    public void testMultiLevelSubqueryWithoutRelation1() {
+        PhysicalPlan p = optimizeAndPlan(
+                "SELECT int FROM (" +
+                "  SELECT int FROM (" +
+                "    SELECT 1 AS int" +
+                "  ) AS subq1" +
+                ") AS subq2");
+        assertThat(p, instanceOf(LocalExec.class));
+        LocalExec le = (LocalExec) p;
+        assertThat(le.executable(), instanceOf(SingletonExecutable.class));
+        assertEquals(1, le.executable().output().size());
+        assertEquals("int", le.executable().output().get(0).name());
+    }
+
+    public void testMultiLevelSubqueryWithoutRelation2() {
+        PhysicalPlan p = optimizeAndPlan(
+                "SELECT i, string FROM (" +
+                "  SELECT * FROM (" +
+                "    SELECT int as i, str AS string FROM (" +
+                "      SELECT * FROM (" +
+                "        SELECT int, s AS str FROM (" +
+                "          SELECT 1 AS int, 'foo' AS s" +
+                "        ) AS subq1" +
+                "      )" +
+                "    ) AS subq2" +
+                "  ) AS subq3" +
+                ")");
+        assertThat(p, instanceOf(LocalExec.class));
+        LocalExec le = (LocalExec) p;
+        assertThat(le.executable(), instanceOf(SingletonExecutable.class));
+        assertEquals(2, le.executable().output().size());
+        assertEquals("i", le.executable().output().get(0).name());
+        assertEquals("string", le.executable().output().get(1).name());
+    }
 }

x-pack/plugin/sql/src/test/resources/org/elasticsearch/xpack/sql/planner/querytranslator_subqueries_tests.txt

@@ -178,6 +178,52 @@ SELECT j AS k FROM (
 ) GROUP BY k;
 ;
 
+MultiLevelAliasedSubqueryGroupBy1
+SELECT int FROM (
+    SELECT int FROM (
+        SELECT int FROM test GROUP BY int
+    ) AS subq1
+) AS subq2;
+;
+
+MultiLevelAliasedSubqueryGroupBy2
+SELECT int FROM (
+    SELECT int FROM (
+        SELECT int FROM test
+    GROUP BY int) AS subq1
+) AS subq2;
+;
+
+MultiLevelAliasedSubqueryGroupBy3
+SELECT * FROM (
+    SELECT int FROM (
+        SELECT * FROM (
+            SELECT int FROM test
+        GROUP BY int) AS subq1
+    ) AS subq2
+)AS subq3;
+;
+
+MultiLevelAliasedSubqueryGroupBy4
+SELECT subq3.int FROM (
+    SELECT int FROM (
+        SELECT subq1.int FROM (
+            SELECT int FROM test
+        GROUP BY int) AS subq1
+    ) AS subq2
+)AS subq3;
+;
+
+MultiLevelAliasedSubqueryGroupBy5
+SELECT subq3.int FROM (
+    SELECT subq2.i AS int FROM (
+        SELECT subq1.int AS i FROM (
+            SELECT int FROM test
+        GROUP BY int) AS subq1
+    ) AS subq2
+)AS subq3;
+;
+
 SubqueryGroupByFilterAndOrderByByRealiased
 SELECT g as h FROM (
     SELECT date AS f, int AS g FROM test