feat(planner): Support GROUP BY and ORDER BY ordinals in MV query rewriting (prestodb#27422) (prestodb#27422)

Summary:
Queries using GROUP BY/ORDER BY ordinals (e.g. `GROUP BY 1`) silently
fell back to the base table because the MV optimizer runs before the
analyzer resolves ordinals to column references. Fix by resolving
ordinals to SELECT expressions during MV validation and passing them
through unchanged during rewriting.


```
== RELEASE NOTES ==

General Changes
* Add support for ``GROUP BY`` and ``ORDER BY`` ordinal references in
  materialized view query rewriting. Previously, queries like
  ``SELECT a, SUM(b) FROM t GROUP BY 1`` would silently skip
  materialized view optimization.
```


Differential Revision: D97920227

Pulled By: ceekay47

Author: ceekay47

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/MaterializedViewQueryOptimizer.java

@@ -56,6 +56,7 @@
 import com.facebook.presto.sql.tree.Join;
 import com.facebook.presto.sql.tree.Lateral;
 import com.facebook.presto.sql.tree.LogicalBinaryExpression;
+import com.facebook.presto.sql.tree.LongLiteral;
 import com.facebook.presto.sql.tree.Node;
 import com.facebook.presto.sql.tree.OrderBy;
 import com.facebook.presto.sql.tree.QualifiedName;
@@ -109,6 +110,7 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.sql.relational.Expressions.call;
 import static com.facebook.presto.util.AnalyzerUtil.createParsingOptions;
+import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
@@ -495,18 +497,38 @@ protected Node visitQuerySpecification(QuerySpecification node, Void context)
                 removablePrefix = Optional.of(new Identifier(baseTable.getName().toString()));
             }
             if (node.getGroupBy().isPresent()) {
+                List<SelectItem> selectItems = node.getSelect().getSelectItems();
                 ImmutableSet.Builder<Expression> expressionsInGroupByBuilder = ImmutableSet.builder();
                 for (GroupingElement element : node.getGroupBy().get().getGroupingElements()) {
                     element = removeGroupingElementPrefix(element, removablePrefix);
                     Optional<Set<Expression>> groupByOfMaterializedView = materializedViewInfo.getGroupBy();
                     if (groupByOfMaterializedView.isPresent()) {
                         for (Expression expression : element.getExpressions()) {
-                            if (!groupByOfMaterializedView.get().contains(expression) || !materializedViewInfo.getBaseToViewColumnMap().containsKey(expression)) {
+                            // Resolve ordinal references (e.g. GROUP BY 1) to the corresponding SELECT expression
+                            Expression resolved = expression;
+                            if (expression instanceof LongLiteral) {
+                                int ordinal = toIntExact(((LongLiteral) expression).getValue());
+                                if (ordinal < 1 || ordinal > selectItems.size()) {
+                                    throw new SemanticException(NOT_SUPPORTED, expression, "GROUP BY ordinal %d is out of range (1 to %d)", ordinal, selectItems.size());
+                                }
+                                SelectItem selectItem = selectItems.get(ordinal - 1);
+                                if (selectItem instanceof SingleColumn) {
+                                    resolved = removeExpressionPrefix(((SingleColumn) selectItem).getExpression(), removablePrefix);
+                                }
+                                else {
+                                    throw new IllegalStateException("GROUP BY ordinal references non-single-column select item");
+                                }
+                            }
+                            if (!groupByOfMaterializedView.get().contains(resolved) || !materializedViewInfo.getBaseToViewColumnMap().containsKey(resolved)) {
                                 throw new IllegalStateException(format("Grouping element %s is not present in materialized view groupBy field", element));
                             }
+                            // Store the resolved expression so visitSingleColumn can match against it
+                            expressionsInGroupByBuilder.add(resolved);
                         }
                     }
-                    expressionsInGroupByBuilder.addAll(element.getExpressions());
+                    else {
+                        expressionsInGroupByBuilder.addAll(element.getExpressions());
+                    }
                 }
                 expressionsInGroupBy = Optional.of(expressionsInGroupByBuilder.build());
             }
@@ -715,8 +737,11 @@ protected Node visitOrderBy(OrderBy node, Void context)
             ImmutableList.Builder<SortItem> rewrittenOrderBy = ImmutableList.builder();
             for (SortItem sortItem : node.getSortItems()) {
                 sortItem = removeSortItemPrefix(sortItem, removablePrefix);
-                if (!materializedViewInfo.getBaseToViewColumnMap().containsKey(sortItem.getSortKey())) {
-                    throw new IllegalStateException(format("Sort key %s is not present in materialized view select fields", sortItem.getSortKey()));
+                // Ordinal references (e.g. ORDER BY 3) refer to SELECT items which are already validated
+                Expression sortKey = sortItem.getSortKey();
+                if (!(sortKey instanceof LongLiteral)
+                        && !materializedViewInfo.getBaseToViewColumnMap().containsKey(sortKey)) {
+                    throw new IllegalStateException(format("Sort key %s is not present in materialized view select fields", sortKey));
                 }
                 rewrittenOrderBy.add((SortItem) process(sortItem, context));
             }
@@ -726,15 +751,26 @@ protected Node visitOrderBy(OrderBy node, Void context)
         @Override
         protected Node visitSortItem(SortItem node, Void context)
         {
-            return new SortItem((Expression) process(node.getSortKey(), context), node.getOrdering(), node.getNullOrdering());
+            Expression sortKey = node.getSortKey();
+            // Ordinal references (e.g. ORDER BY 1) refer to SELECT positions which are already rewritten; pass through unchanged
+            if (sortKey instanceof LongLiteral) {
+                return node;
+            }
+            return new SortItem((Expression) process(sortKey, context), node.getOrdering(), node.getNullOrdering());
         }
 
         @Override
         protected Node visitSimpleGroupBy(SimpleGroupBy node, Void context)
         {
             ImmutableList.Builder<Expression> rewrittenSimpleGroupBy = ImmutableList.builder();
             for (Expression column : node.getExpressions()) {
-                rewrittenSimpleGroupBy.add((Expression) process(removeExpressionPrefix(column, removablePrefix), context));
+                // Ordinal references (e.g. GROUP BY 1) refer to SELECT positions which are already rewritten; pass through unchanged
+                if (column instanceof LongLiteral) {
+                    rewrittenSimpleGroupBy.add(column);
+                }
+                else {
+                    rewrittenSimpleGroupBy.add((Expression) process(removeExpressionPrefix(column, removablePrefix), context));
+                }
             }
             return new SimpleGroupBy(rewrittenSimpleGroupBy.build());
         }

presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestMaterializedViewQueryOptimizer.java

@@ -255,6 +255,36 @@ public void testWithOrderBy()
         assertOptimizedQuery(baseQuerySql, expectedRewrittenSql, originalViewSql, BASE_TABLE_1, VIEW_1);
     }
 
+    @Test
+    public void testWithGroupByOrdinals()
+    {
+        String originalViewSql = format("SELECT a as mv_a, b, c as mv_c FROM %s", BASE_TABLE_1);
+        String baseQuerySql = format("SELECT SUM(a * b), MAX(a + b), c FROM %s GROUP BY 3", BASE_TABLE_1);
+        String expectedRewrittenSql = format("SELECT SUM(mv_a * b), MAX(mv_a + b), mv_c as c FROM %s GROUP BY 3", VIEW_1);
+
+        assertOptimizedQuery(baseQuerySql, expectedRewrittenSql, originalViewSql, BASE_TABLE_1, VIEW_1);
+    }
+
+    @Test
+    public void testWithOrderByOrdinals()
+    {
+        String originalViewSql = format("SELECT a as mv_a, b, c as mv_c FROM %s", BASE_TABLE_1);
+        String baseQuerySql = format("SELECT a, b, c FROM %s ORDER BY 3 ASC, 2 DESC, 1", BASE_TABLE_1);
+        String expectedRewrittenSql = format("SELECT mv_a as a, b, mv_c as c FROM %s ORDER BY 3 ASC, 2 DESC, 1", VIEW_1);
+
+        assertOptimizedQuery(baseQuerySql, expectedRewrittenSql, originalViewSql, BASE_TABLE_1, VIEW_1);
+    }
+
+    @Test
+    public void testWithGroupByAndOrderByOrdinals()
+    {
+        String originalViewSql = format("SELECT MAX(a) as mv_max_a, b FROM %s GROUP BY b", BASE_TABLE_1);
+        String baseQuerySql = format("SELECT MAX(a), b FROM %s GROUP BY 2 ORDER BY 1 DESC, 2 ASC", BASE_TABLE_1);
+        String expectedRewrittenSql = format("SELECT MAX(mv_max_a), b FROM %s GROUP BY 2 ORDER BY 1 DESC, 2 ASC", VIEW_1);
+
+        assertOptimizedQuery(baseQuerySql, expectedRewrittenSql, originalViewSql, BASE_TABLE_1, VIEW_1);
+    }
+
     @Test
     public void testWithNoMatchingBaseTable()
     {