SQL: Remove restriction for single column grouping (#31818)

For historical reasons SQL restricts GROUP BY to only one field.
This commit removes the restriction and improves the test suite with
multi group by tests.

Close #31793

Author: costin

docs/reference/sql/language/syntax/select.asciidoc

@@ -177,6 +177,13 @@ And grouping by column expression (typically used along-side an alias):
 include-tagged::{sql-specs}/docs.csv-spec[groupByExpression]
 ----
 
+Or a mixture of the above:
+["source","sql",subs="attributes,callouts,macros"]
+----
+include-tagged::{sql-specs}/docs.csv-spec[groupByMulti]
+----
+
+
 When a `GROUP BY` clause is used in a `SELECT`, _all_ output expressions must be either aggregate functions or expressions used for grouping or derivatives of (otherwise there would be more than one possible value to return for each ungrouped column).
 
 To wit:

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Expressions.java

@@ -104,6 +104,14 @@ public static Attribute attribute(Expression e) {
         return null;
     }
 
+    public static boolean equalsAsAttribute(Expression left, Expression right) {
+        if (!left.semanticEquals(right)) {
+            Attribute l = attribute(left);
+            return (l != null && l.semanticEquals(attribute(right)));
+        }
+        return true;
+    }
+
     public static TypeResolution typeMustBe(Expression e, Predicate<Expression> predicate, String message) {
         return predicate.test(e) ? TypeResolution.TYPE_RESOLVED : new TypeResolution(message);
     }

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

@@ -112,7 +112,7 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 new ReplaceAggsWithStats(),
                 new PromoteStatsToExtendedStats(),
                 new ReplaceAggsWithPercentiles(),
-                new ReplceAggsWithPercentileRanks()
+                new ReplaceAggsWithPercentileRanks()
                 );
 
         Batch operators = new Batch("Operator Optimization",
@@ -132,7 +132,9 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 new PruneFilters(),
                 new PruneOrderBy(),
                 new PruneOrderByNestedFields(),
-                new PruneCast()
+                new PruneCast(),
+                // order by alignment of the aggs
+                new SortAggregateOnOrderBy()
                 // requires changes in the folding
                 // since the exact same function, with the same ID can appear in multiple places
                 // see https://github.com/elastic/x-pack-elasticsearch/issues/3527
@@ -612,7 +614,7 @@ protected LogicalPlan rule(LogicalPlan e) {
         }
     }
 
-    static class ReplceAggsWithPercentileRanks extends Rule<LogicalPlan, LogicalPlan> {
+    static class ReplaceAggsWithPercentileRanks extends Rule<LogicalPlan, LogicalPlan> {
 
         @Override
         public LogicalPlan apply(LogicalPlan p) {
@@ -828,6 +830,46 @@ protected LogicalPlan rule(OrderBy ob) {
         }
     }
 
+    /**
+     * Align the order in aggregate based on the order by.
+     */
+    static class SortAggregateOnOrderBy extends OptimizerRule<OrderBy> {
+
+        @Override
+        protected LogicalPlan rule(OrderBy ob) {
+            List<Order> order = ob.order();
+
+            // remove constants
+            List<Order> nonConstant = order.stream().filter(o -> !o.child().foldable()).collect(toList());
+
+            // if the sort points to an agg, change the agg order based on the order
+            if (ob.child() instanceof Aggregate) {
+                Aggregate a = (Aggregate) ob.child();
+                List<Expression> groupings = new ArrayList<>(a.groupings());
+                boolean orderChanged = false;
+
+                for (int orderIndex = 0; orderIndex < nonConstant.size(); orderIndex++) {
+                    Order o = nonConstant.get(orderIndex);
+                    Expression fieldToOrder = o.child();
+                    for (Expression group : a.groupings()) {
+                        if (Expressions.equalsAsAttribute(fieldToOrder, group)) {
+                            // move grouping in front
+                            groupings.remove(group);
+                            groupings.add(orderIndex, group);
+                            orderChanged = true;
+                        }
+                    }
+                }
+
+                if (orderChanged) {
+                    Aggregate newAgg = new Aggregate(a.location(), a.child(), groupings, a.aggregates());
+                    return new OrderBy(ob.location(), newAgg, ob.order());
+                }
+            }
+            return ob;
+        }
+    }
+
     static class CombineLimits extends OptimizerRule<Limit> {
 
         @Override

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

@@ -5,8 +5,6 @@
  */
 package org.elasticsearch.xpack.sql.planner;
 
-import org.elasticsearch.xpack.sql.expression.Expressions;
-import org.elasticsearch.xpack.sql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.sql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.sql.plan.physical.Unexecutable;
 import org.elasticsearch.xpack.sql.plan.physical.UnplannedExec;
@@ -71,23 +69,11 @@ static List<Failure> verifyMappingPlan(PhysicalPlan plan) {
                     failures.add(fail(e, "Unresolved expression"));
                 }
             });
-
-            if (p instanceof AggregateExec) {
-                forbidMultiFieldGroupBy((AggregateExec) p, failures);
-            }
         });
 
         return failures;
     }
 
-    private static void forbidMultiFieldGroupBy(AggregateExec a, List<Failure> failures) {
-        if (a.groupings().size() > 1) {
-            failures.add(fail(a.groupings().get(0), "Currently, only a single expression can be used with GROUP BY; please select one of "
-                    + Expressions.names(a.groupings())));
-        }
-    }
-
-
     static List<Failure> verifyExecutingPlan(PhysicalPlan plan) {
         List<Failure> failures = new ArrayList<>();
 

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

@@ -11,7 +11,7 @@
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.xpack.qa.sql.jdbc.CsvTestUtils.CsvTestCase;
 import org.elasticsearch.xpack.qa.sql.jdbc.DataLoader;
-import org.elasticsearch.xpack.qa.sql.jdbc.JdbcAssert;
+import org.elasticsearch.xpack.qa.sql.jdbc.JdbcTestUtils;
 import org.elasticsearch.xpack.qa.sql.jdbc.SpecBaseIntegrationTestCase;
 import org.elasticsearch.xpack.qa.sql.jdbc.SqlSpecTestCase;
 
@@ -68,8 +68,8 @@ protected void assertResults(ResultSet expected, ResultSet elastic) throws SQLEx
         //
         // uncomment this to printout the result set and create new CSV tests
         //
-        //JdbcTestUtils.logLikeCLI(elastic, log);
-        JdbcAssert.assertResultSets(expected, elastic, log, true);
+        JdbcTestUtils.logLikeCLI(elastic, log);
+        //JdbcAssert.assertResultSets(expected, elastic, log, true);
     }
 
     @Override