SQL: Fix issues with GROUP BY queries (#41964)

Translate to an agg query even if only literals are selected,
so that the correct number of rows is returned (number of buckets).

Fix issue with key only in GROUP BY (not in select) and WHERE clause:
Resolve aggregates and groupings based on the child plan which holds
the info info for all the fields of the underlying table.

Fixes: #41951
Fixes: #41413
(cherry picked from commit 45b85809678b34a448639a420b97e25436ae851f)

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 @@ public class Analyzer extends RuleExecutor<LogicalPlan> {
                     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) {
@@ -720,6 +721,18 @@ public class Analyzer extends RuleExecutor<LogicalPlan> {
                 }
             }
 
+            // 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 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
 
             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

@@ -401,6 +401,7 @@ class QueryFolder extends RuleExecutor<PhysicalPlan> {
             // 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());
@@ -451,7 +452,7 @@ class QueryFolder extends RuleExecutor<PhysicalPlan> {
                     target = ((Alias) ne).child();
                 }
 
-                String id = Expressions.id(target);
+                id = Expressions.id(target);
 
                 // literal
                 if (target.foldable()) {
@@ -587,7 +588,14 @@ class QueryFolder extends RuleExecutor<PhysicalPlan> {
                     }
                 }
             }
-
+            // 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.Alias;
 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 class QueryTranslatorTests extends ESTestCase {
         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);
+        Aggregate a = (Aggregate) p;
+        LogicalPlan pc = ((Aggregate) p).child();
+        assertTrue(pc instanceof Filter);
+        Expression condition = ((Filter) pc).condition();
+        assertEquals("GREATERTHAN", ((GreaterThan) condition).functionName());
+        List<Expression> groupings = a.groupings();
+        assertTrue(groupings.get(0).resolved());
+        assertEquals("c", a.aggregates().get(0).name());
+        assertEquals("COUNT", ((Count) ((Alias) a.aggregates().get(0)).child()).functionName());
+    }
+    public void testLiteralWithGroupBy(){
+        LogicalPlan p = plan("SELECT 1 as t, 2 FROM test GROUP BY int");
+        assertTrue(p instanceof Aggregate);
+        Aggregate a = (Aggregate) p;
+        List<Expression> groupings = a.groupings();
+        assertEquals(1, groupings.size());
+        assertTrue(groupings.get(0).resolved());
+        assertTrue(groupings.get(0) instanceof FieldAttribute);
+        assertEquals(2, a.aggregates().size());
+        assertEquals("t", a.aggregates().get(0).name());
+        assertTrue(((Alias) a.aggregates().get(0)).child() instanceof Literal);
+        assertEquals("1", ((Alias) a.aggregates().get(0)).child().toString());
+        assertEquals("2", ((Alias) a.aggregates().get(1)).child().toString());
+    }
+
     public void testTermEqualityNotAnalyzed() {
         LogicalPlan p = plan("SELECT some.string FROM test WHERE int = 5");
         assertTrue(p instanceof Project);