SQL: selecting a literal from grouped by query generates error

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

@@ -144,6 +144,26 @@ countDistinctAlias
 SELECT COUNT(DISTINCT hire_date) AS count FROM test_emp;
 countDistinctAndCountSimpleWithAlias
 SELECT COUNT(*) cnt, COUNT(DISTINCT first_name) as names, gender FROM test_emp GROUP BY gender ORDER BY gender;
+aliasedCountWithFunctionFilterAndGroupBy
+SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender ORDER BY gender;
+countWithFunctionFilterAndGroupBy
+SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender ORDER BY gender;
+aliasedCountWithMultiFunctionFilterAndGroupBy
+SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender ORDER BY gender;
+countWithMultiFunctionFilterAndGroupBy
+SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender ORDER BY gender;
+aliasedCountWithFunctionFilterAndMultiGroupBy
+SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender, salary ORDER BY gender;
+countWithFunctionFilterAndMultiGroupBy
+SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender, salary ORDER BY gender;
+aliasedCountWithMultiFunctionFilterAndMultiGroupBy
+SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender, salary ORDER BY gender;
+countWithMultiFunctionFilterAndMultiGroupBy
+SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender, salary ORDER BY gender;
+aliasedCountLiteralColumnWithFunctionFilterAndMultiGroupBy
+SELECT 1, gender as g, COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY g, salary ORDER BY gender;
+aliasedCountLiteralColumnWithFunctionFilterAndMultiGroupByWithFunction
+SELECT 1, gender as g, COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY g, YEAR(birth_date) ORDER BY gender, YEAR(birth_date);
 
 aggCountAliasAndWhereClauseMultiGroupBy
 SELECT gender g, languages l, COUNT(*) c FROM "test_emp" WHERE emp_no < 10020 GROUP BY gender, languages ORDER BY gender, languages;
@@ -563,6 +583,22 @@ SELECT MIN(salary) min, MAX(salary) max, gender g, languages l, COUNT(*) c FROM
 // group by with literal
 implicitGroupByWithLiteral
 SELECT 10, MAX("salary") FROM test_emp;
+literalWithGroupBy
+SELECT 1 FROM test_emp GROUP BY gender;
+literalsWithGroupBy
+SELECT 1, 2 FROM test_emp GROUP BY gender;
+aliasedLiteralWithGroupBy
+SELECT 1 AS s FROM test_emp GROUP BY gender;
+aliasedLiteralsWithGroupBy
+SELECT 1 AS s, 2 FROM test_emp GROUP BY gender;
+literalsWithMultipleGroupBy
+SELECT 1, 2 FROM test_emp GROUP BY gender, salary;
+divisionLiteralsAdditionWithMultipleGroupBy
+SELECT 144 / 12 AS division, 1, 2 AS x, 1 + 2 AS addition FROM test_emp GROUP BY gender, salary;
+aliasedLiteralsWithMultipleGroupBy
+SELECT 1 as s, 2 FROM test_emp GROUP BY gender, salary;
+aliasedLiteralsWithMultipleGroupByWithFunction
+SELECT 1 as s, 2 FROM test_emp GROUP BY gender, YEAR(birth_date);
 implicitGroupByWithLiterals
 SELECT 10, 'foo', MAX("salary"), 20, 'bar' FROM test_emp;
 groupByWithLiteral

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

@@ -341,6 +341,7 @@ else if (plan instanceof Aggregate) {
                     List<Expression> groupings = a.groupings();
                     List<Expression> newGroupings = new ArrayList<>();
                     AttributeMap<Expression> resolved = Expressions.aliases(a.aggregates());
+
                     boolean changed = false;
                     for (Expression grouping : groupings) {
                         if (grouping instanceof UnresolvedAttribute) {
@@ -618,7 +619,7 @@ protected LogicalPlan rule(LogicalPlan plan) {
                 for (Order or : o.order()) {
                     maybeResolved.add(or.resolved() ? or : tryResolveExpression(or, child));
                 }
-                
+
                 Stream<Order> referencesStream = maybeResolved.stream()
                         .filter(Expression::resolved);
 
@@ -629,7 +630,7 @@ protected LogicalPlan rule(LogicalPlan plan) {
                 // a + 1 in SELECT is actually Alias("a + 1", a + 1) and translates to ReferenceAttribute
                 // in the output. However it won't match the unnamed a + 1 despite being the same expression
                 // so explicitly compare the source
-                
+
                 // if there's a match, remove the item from the reference stream
                 if (Expressions.hasReferenceAttribute(child.outputSet())) {
                     final Map<Attribute, Expression> collectRefs = new LinkedHashMap<>();
@@ -720,6 +721,18 @@ protected LogicalPlan rule(LogicalPlan plan) {
                 }
             }
 
+            // Try to resolve aggregates and groupings based on the child plan
+            if (plan instanceof Aggregate) {
+                Aggregate a = (Aggregate) plan;
+                LogicalPlan child = a.child();
+                List<Expression> newGroupings = new ArrayList<>(a.groupings().size());
+                a.groupings().forEach(e -> newGroupings.add(tryResolveExpression(e, child)));
+                List<NamedExpression> newAggregates = new ArrayList<>(a.aggregates().size());
+                a.aggregates().forEach(e -> newAggregates.add(tryResolveExpression(e, child)));
+                if (newAggregates.equals(a.aggregates()) == false || newGroupings.equals(a.groupings()) == false) {
+                    return new Aggregate(a.source(), child, newGroupings, newAggregates);
+                }
+            }
             return plan;
         }
 

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

@@ -1130,7 +1130,8 @@ protected LogicalPlan rule(LogicalPlan plan) {
 
             plan.forEachDown(a -> {
                 List<Object> values = extractConstants(a.aggregates());
-                if (values.size() == a.aggregates().size() && isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) {
+                if (values.size() == a.aggregates().size() && a.groupings().isEmpty()
+                    && isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) {
                     optimizedPlan.set(new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray())));
                 }
             }, Aggregate.class);

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

@@ -393,20 +393,21 @@ protected PhysicalPlan rule(AggregateExec a) {
             }
             return a;
         }
-        
+
         static EsQueryExec fold(AggregateExec a, EsQueryExec exec) {
-            
+
             QueryContainer queryC = exec.queryContainer();
-            
+
             // track aliases defined in the SELECT and used inside GROUP BY
             // SELECT x AS a ... GROUP BY a
             Map<Attribute, Expression> aliasMap = new LinkedHashMap<>();
+            String id = null;
             for (NamedExpression ne : a.aggregates()) {
                 if (ne instanceof Alias) {
                     aliasMap.put(ne.toAttribute(), ((Alias) ne).child());
                 }
             }
-            
+
             if (aliasMap.isEmpty() == false) {
                 Map<Attribute, Expression> newAliases = new LinkedHashMap<>(queryC.aliases());
                 newAliases.putAll(aliasMap);
@@ -451,7 +452,7 @@ static EsQueryExec fold(AggregateExec a, EsQueryExec exec) {
                     target = ((Alias) ne).child();
                 }
 
-                String id = Expressions.id(target);
+                id = Expressions.id(target);
 
                 // literal
                 if (target.foldable()) {
@@ -587,7 +588,14 @@ else if (target.foldable()) {
                     }
                 }
             }
-
+            // If we're only selecting literals, we have to still execute the aggregation to create
+            // the correct grouping buckets, in order to return the appropriate number of rows
+            if (a.aggregates().stream().allMatch(e -> e.anyMatch(Expression::foldable))) {
+                for (Expression grouping : a.groupings()) {
+                    GroupByKey matchingGroup = groupingContext.groupFor(grouping);
+                    queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, false), id);
+                }
+            }
             return new EsQueryExec(exec.source(), exec.index(), a.output(), queryC);
         }
 

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

@@ -16,7 +16,9 @@
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.FieldAttribute;
 import org.elasticsearch.xpack.ql.expression.Literal;
+import org.elasticsearch.xpack.ql.expression.function.aggregate.Count;
 import org.elasticsearch.xpack.ql.expression.gen.script.ScriptTemplate;
+import org.elasticsearch.xpack.ql.expression.predicate.operator.comparison.GreaterThan;
 import org.elasticsearch.xpack.ql.index.EsIndex;
 import org.elasticsearch.xpack.ql.index.IndexResolution;
 import org.elasticsearch.xpack.ql.plan.logical.Aggregate;
@@ -130,6 +132,34 @@ public void testTermEqualityAnalyzer() {
         assertEquals("value", tq.value());
     }
 
+    public void testAliasAndGroupByResolution(){
+        LogicalPlan p = plan("SELECT COUNT(*) AS c FROM test WHERE ABS(int) > 0 GROUP BY int");
+        assertTrue(p instanceof Aggregate);
+        var pc = ((Aggregate) p).child();
+        assertTrue(pc instanceof Filter);
+        Expression condition = ((Filter) pc).condition();
+        assertEquals(((GreaterThan) condition).functionName(), "GREATERTHAN");
+        List<Expression> groupings = ((Aggregate) p).groupings();
+        assertTrue(groupings.get(0).resolved());
+        var agg = ((Aggregate) p).aggregates();
+        assertEquals((agg.get(0)).name(), "c");
+        assertEquals(((Count) ((Alias) agg.get(0)).child()).functionName(), "COUNT");
+    }
+    public void testLiteralWithGroupBy(){
+        LogicalPlan p = plan("SELECT 1 as t, 2 FROM test GROUP BY int");
+        assertTrue(p instanceof Aggregate);
+        List<Expression> groupings = ((Aggregate) p).groupings();
+        assertTrue(groupings.size() == 1);
+        assertTrue(groupings.get(0).resolved());
+        assertTrue(groupings.get(0) instanceof FieldAttribute);
+        var aggs = ((Aggregate) p).aggregates();
+        assertTrue(aggs.size() == 2);
+        assertEquals((aggs.get(0)).name(), "t");
+        assertTrue(((Alias) aggs.get(0)).child() instanceof Literal);
+        assertEquals(((Alias) aggs.get(0)).child().toString(), "1");
+        assertEquals(((Alias) aggs.get(1)).child().toString(), "2");
+    }
+
     public void testTermEqualityNotAnalyzed() {
         LogicalPlan p = plan("SELECT some.string FROM test WHERE int = 5");
         assertTrue(p instanceof Project);