From 0c13fdd22c8cf7cf5800fb757bf8a9324e5d0e9f Mon Sep 17 00:00:00 2001
From: Costin Leau <costin.leau@gmail.com>
Date: Tue, 23 Oct 2018 21:52:55 +0300
Subject: [PATCH] SQL: Verifier allows aliases aggregates for sorting

Improve Verifier rule that prevented grouping, aliases inside aggregates
to not be accepted for ordering.

Close #34607
---
 .../xpack/sql/analysis/analyzer/Verifier.java    | 16 +++++++++++++++-
 .../analyzer/VerifierErrorMessagesTests.java     | 13 +++++++++++++
 x-pack/qa/sql/src/main/resources/agg.sql-spec    |  2 ++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java
index e5ab3ce082b71..32d57175114d8 100644
--- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java
+++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java
@@ -6,6 +6,7 @@
 package org.elasticsearch.xpack.sql.analysis.analyzer;
 
 import org.elasticsearch.xpack.sql.capabilities.Unresolvable;
+import org.elasticsearch.xpack.sql.expression.Alias;
 import org.elasticsearch.xpack.sql.expression.Attribute;
 import org.elasticsearch.xpack.sql.expression.AttributeSet;
 import org.elasticsearch.xpack.sql.expression.Exists;
@@ -249,8 +250,21 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
                         return;
                     }
 
+                    // take aliases declared inside the aggregates which point to the grouping (but are not included in there)
+                    // to correlate them to the order
+                    List<Expression> groupingAndMatchingAggregatesAliases = new ArrayList<>(a.groupings());
+                    
+                    a.aggregates().forEach(as -> {
+                        if (as instanceof Alias) {
+                            Alias al = (Alias) as;
+                            if (Expressions.anyMatch(a.groupings(), g -> Expressions.equalsAsAttribute(al.child(), g))) {
+                                groupingAndMatchingAggregatesAliases.add(al);
+                            }
+                        }
+                    });
+                    
                     // make sure to compare attributes directly
-                    if (Expressions.anyMatch(a.groupings(),
+                    if (Expressions.anyMatch(groupingAndMatchingAggregatesAliases,
                             g -> e.semanticEquals(e instanceof Attribute ? Expressions.attribute(g) : g))) {
                         return;
                     }
diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerifierErrorMessagesTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerifierErrorMessagesTests.java
index c193dcfd5461f..b10972110750e 100644
--- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerifierErrorMessagesTests.java
+++ b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerifierErrorMessagesTests.java
@@ -11,6 +11,7 @@
 import org.elasticsearch.xpack.sql.analysis.index.IndexResolution;
 import org.elasticsearch.xpack.sql.expression.function.FunctionRegistry;
 import org.elasticsearch.xpack.sql.parser.SqlParser;
+import org.elasticsearch.xpack.sql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.sql.type.EsField;
 import org.elasticsearch.xpack.sql.type.TypesTests;
 
@@ -34,6 +35,13 @@ private String verify(IndexResolution getIndexResult, String sql) {
         return e.getMessage().substring(header.length());
     }
 
+    private LogicalPlan accepted(String sql) {
+        Map<String, EsField> mapping = TypesTests.loadMapping("mapping-multi-field-with-nested.json");
+        EsIndex test = new EsIndex("test", mapping);
+        Analyzer analyzer = new Analyzer(new FunctionRegistry(), IndexResolution.valid(test), TimeZone.getTimeZone("UTC"));
+        return analyzer.analyze(parser.createStatement(sql), true);
+    }
+
     public void testMissingIndex() {
         assertEquals("1:17: Unknown index [missing]", verify(IndexResolution.notFound("missing"), "SELECT foo FROM missing"));
     }
@@ -110,6 +118,11 @@ public void testGroupByOrderByNonGrouped() {
                 verify("SELECT MAX(int) FROM test GROUP BY text ORDER BY bool"));
     }
 
+    public void testGroupByOrderByAliasedInSelectAllowed() {
+        LogicalPlan lp = accepted("SELECT text t FROM test GROUP BY text ORDER BY t");
+        assertNotNull(lp);
+    }
+
     public void testGroupByOrderByScalarOverNonGrouped() {
         assertEquals("1:50: Cannot order by non-grouped column [YEAR(date [UTC])], expected [text]",
                 verify("SELECT MAX(int) FROM test GROUP BY text ORDER BY YEAR(date)"));
diff --git a/x-pack/qa/sql/src/main/resources/agg.sql-spec b/x-pack/qa/sql/src/main/resources/agg.sql-spec
index 2c6248059f5fb..d3033e16f1c7b 100644
--- a/x-pack/qa/sql/src/main/resources/agg.sql-spec
+++ b/x-pack/qa/sql/src/main/resources/agg.sql-spec
@@ -35,6 +35,8 @@ groupByOnNumberWithWhereAndLimit
 SELECT emp_no e FROM "test_emp" WHERE emp_no < 10020 GROUP BY emp_no ORDER BY emp_no DESC LIMIT 1;
 groupByOnNumberOnAlias
 SELECT emp_no e FROM "test_emp" WHERE emp_no < 10020 GROUP BY e ORDER BY emp_no DESC;
+groupByOnNumberWithAliasInSelect
+SELECT emp_no e FROM "test_emp" WHERE emp_no < 10020 GROUP BY e ORDER BY e DESC;
 
 // group by scalar
 groupByAddScalar