[FLINK-40344][table] Support selecting computed grouping expressions

flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java

@@ -587,11 +587,15 @@ public Table select(Expression... fields) {
                         "Window properties can only be used on windowed tables.");
             }
 
-            return table.createTable(
-                    table.operationTreeBuilder.project(
-                            extracted.getProjections(),
-                            table.operationTreeBuilder.aggregate(
-                                    groupKeys, extracted.getAggregations(), table.operationTree)));
+            List<Expression> groupingExpressions =
+                    table.operationTreeBuilder.expandExpressions(
+                            table.preprocessExpressions(groupKeys), table.operationTree);
+            QueryOperation aggregateOperation =
+                    table.operationTreeBuilder.aggregate(
+                            groupingExpressions, extracted.getAggregations(), table.operationTree);
+
+            return table.projectAfterAggregation(
+                    extracted.getProjections(), groupingExpressions, aggregateOperation);
         }
 
         @Override
@@ -623,11 +627,15 @@ private AggregatedTableImpl(
 
         @Override
         public Table select(Expression... fields) {
-            return table.createTable(
-                    table.operationTreeBuilder.project(
-                            Arrays.asList(fields),
-                            table.operationTreeBuilder.aggregate(
-                                    groupKeys, aggregateFunction, table.operationTree)));
+            List<Expression> groupingExpressions =
+                    table.operationTreeBuilder.expandExpressions(
+                            table.preprocessExpressions(groupKeys), table.operationTree);
+            QueryOperation aggregateOperation =
+                    table.operationTreeBuilder.aggregate(
+                            groupingExpressions, aggregateFunction, table.operationTree);
+
+            return table.projectAfterAggregation(
+                    table.preprocessExpressions(fields), groupingExpressions, aggregateOperation);
         }
     }
 
@@ -950,6 +958,18 @@ private TableImpl createTable(QueryOperation operation) {
         return new TableImpl(tableEnvironment, operation, operationTreeBuilder, lookupResolver);
     }
 
+    private Table projectAfterAggregation(
+            List<Expression> projections,
+            List<Expression> groupingExpressions,
+            QueryOperation aggregateOperation) {
+        List<Expression> rewrittenProjections =
+                OperationExpressionsUtils.replaceGroupingExpressions(
+                        projections,
+                        groupingExpressions,
+                        aggregateOperation.getResolvedSchema().getColumnNames());
+        return createTable(operationTreeBuilder.project(rewrittenProjections, aggregateOperation));
+    }
+
     private List<Expression> preprocessExpressions(List<Expression> expressions) {
         return preprocessExpressions(expressions.toArray(new Expression[0]));
     }

flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/AggregateQueryOperation.java

@@ -101,6 +101,8 @@ private String getGroupingExprs(SqlFactory sqlFactory) {
             return "1";
         } else {
             return groupingExpressions.stream()
+                    // Alias declarations define output names in SELECT but are invalid in GROUP BY.
+                    .map(OperationExpressionsUtils::unwrapAlias)
                     .map(
                             expr ->
                                     OperationExpressionsUtils.scopeReferencesWithAlias(

flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationExpressionsUtils.java

@@ -126,13 +126,17 @@ public static CategorizedExpressions extractAggregationsAndProperties(
         AggregationAndPropertiesSplitter splitter = new AggregationAndPropertiesSplitter();
         expressions.forEach(expr -> expr.accept(splitter));
 
+        Map<Expression, String> extractedExpressionToFieldName =
+                new LinkedHashMap<>(splitter.aggregates);
+        extractedExpressionToFieldName.putAll(splitter.properties);
+
         List<Expression> projections =
                 expressions.stream()
                         .map(
                                 expr ->
                                         expr.accept(
-                                                new AggregationAndPropertiesReplacer(
-                                                        splitter.aggregates, splitter.properties)))
+                                                new ExpressionReplacer(
+                                                        extractedExpressionToFieldName)))
                         .collect(Collectors.toList());
 
         List<Expression> aggregates = nameExpressions(splitter.aggregates);
@@ -141,6 +145,36 @@ public static CategorizedExpressions extractAggregationsAndProperties(
         return new CategorizedExpressions(projections, aggregates, properties);
     }
 
+    /**
+     * Replaces grouping expressions in projections with references to the corresponding aggregate
+     * output fields.
+     *
+     * <p>Aggregate outputs place grouping fields before aggregate fields. Therefore, grouping
+     * expressions and output field names are matched by position.
+     */
+    public static List<Expression> replaceGroupingExpressions(
+            List<Expression> projections,
+            List<Expression> groupingExpressions,
+            List<String> aggregateOutputFieldNames) {
+        if (groupingExpressions.size() > aggregateOutputFieldNames.size()) {
+            throw new IllegalArgumentException(
+                    "The aggregate output does not contain all grouping expressions.");
+        }
+
+        Map<Expression, String> groupingExpressionToFieldName = new LinkedHashMap<>();
+        for (int i = 0; i < groupingExpressions.size(); i++) {
+            groupingExpressionToFieldName.put(
+                    groupingExpressions.get(i), aggregateOutputFieldNames.get(i));
+        }
+
+        return projections.stream()
+                .map(
+                        projection ->
+                                projection.accept(
+                                        new ExpressionReplacer(groupingExpressionToFieldName)))
+                .collect(Collectors.toList());
+    }
+
     private static List<Expression> nameExpressions(Map<Expression, String> expressions) {
         return expressions.entrySet().stream()
                 .map(entry -> unresolvedCall(AS, entry.getKey(), valueLiteral(entry.getValue())))
@@ -179,16 +213,12 @@ protected Void defaultMethod(Expression expression) {
         }
     }
 
-    private static class AggregationAndPropertiesReplacer
-            extends ApiExpressionDefaultVisitor<Expression> {
+    private static class ExpressionReplacer extends ApiExpressionDefaultVisitor<Expression> {
 
-        private final Map<Expression, String> aggregates;
-        private final Map<Expression, String> properties;
+        private final Map<Expression, String> expressionToFieldName;
 
-        private AggregationAndPropertiesReplacer(
-                Map<Expression, String> aggregates, Map<Expression, String> properties) {
-            this.aggregates = aggregates;
-            this.properties = properties;
+        private ExpressionReplacer(Map<Expression, String> expressionToFieldName) {
+            this.expressionToFieldName = expressionToFieldName;
         }
 
         @Override
@@ -204,10 +234,9 @@ public Expression visit(CallExpression call) {
 
         @Override
         public Expression visit(UnresolvedCallExpression unresolvedCall) {
-            if (aggregates.get(unresolvedCall) != null) {
-                return unresolvedRef(aggregates.get(unresolvedCall));
-            } else if (properties.get(unresolvedCall) != null) {
-                return unresolvedRef(properties.get(unresolvedCall));
+            String fieldName = expressionToFieldName.get(unresolvedCall);
+            if (fieldName != null) {
+                return unresolvedRef(fieldName);
             }
 
             final List<Expression> args =
@@ -219,7 +248,8 @@ public Expression visit(UnresolvedCallExpression unresolvedCall) {
 
         @Override
         protected Expression defaultMethod(Expression expression) {
-            return expression;
+            String fieldName = expressionToFieldName.get(expression);
+            return fieldName == null ? expression : unresolvedRef(fieldName);
         }
     }
 
@@ -265,6 +295,17 @@ public static Optional<String> extractName(Expression expression) {
         return expression.accept(extractNameVisitor);
     }
 
+    /** Returns the underlying expression if the given expression declares an alias. */
+    public static ResolvedExpression unwrapAlias(ResolvedExpression expression) {
+        if (expression instanceof CallExpression) {
+            CallExpression call = (CallExpression) expression;
+            if (call.getFunctionDefinition() == AS) {
+                return call.getResolvedChildren().get(0);
+            }
+        }
+        return expression;
+    }
+
     private static class ExtractNameVisitor extends ApiExpressionDefaultVisitor<Optional<String>> {
 
         @Override

flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java

@@ -242,13 +242,22 @@ public QueryOperation dropColumns(List<Expression> fieldLists, QueryOperation ch
         return project(finalFields, child, false);
     }
 
+    public List<Expression> expandExpressions(List<Expression> expressions, QueryOperation child) {
+        return getResolver(child).resolveExpanding(expressions);
+    }
+
     public QueryOperation aggregate(
             List<Expression> groupingExpressions,
             List<Expression> aggregates,
             QueryOperation child) {
 
-        ExpressionResolver resolver = getAggResolver(child, groupingExpressions);
-        List<ResolvedExpression> resolvedGroupings = resolver.resolve(groupingExpressions);
+        // Computed grouping expressions need stable names for subsequent projections.
+        List<Expression> namedGroupingExpressions =
+                addAliasToTheCallInAggregate(
+                        child.getResolvedSchema().getColumnNames(), groupingExpressions);
+
+        ExpressionResolver resolver = getAggResolver(child, namedGroupingExpressions);
+        List<ResolvedExpression> resolvedGroupings = resolver.resolve(namedGroupingExpressions);
         List<ResolvedExpression> resolvedAggregates = resolver.resolve(aggregates);
 
         return aggregateOperationFactory.createAggregate(

flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java

@@ -21,6 +21,7 @@
 import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.streaming.api.graph.StreamGraph;
 import org.apache.flink.table.api.internal.TableEnvironmentImpl;
+import org.apache.flink.table.expressions.Expression;
 import org.apache.flink.table.expressions.SqlFactory;
 import org.apache.flink.table.functions.FunctionDefinition;
 import org.apache.flink.table.operations.CollectModifyOperation;
@@ -45,6 +46,7 @@
 
 import static org.apache.flink.table.api.Expressions.$;
 import static org.apache.flink.table.api.Expressions.lit;
+import static org.apache.flink.table.api.Expressions.withColumns;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -143,6 +145,62 @@ void testProctimeWindowGeneratedSqlPlans() {
         assertThatCode(sqlTable::explain).doesNotThrowAnyException();
     }
 
+    @Test
+    void testComputedGroupingExpressionCanBeSelected() {
+        final TableEnvironment env = setupEnv(QueryOperationTestPrograms.AGGREGATE_QUERY_OPERATION);
+        final Expression key = $("a").isGreater(10);
+
+        final Table result = env.from("s").groupBy(key).select(key, $("a").count().as("n"));
+
+        assertGeneratedSqlCanBeParsed(env, result);
+    }
+
+    @Test
+    void testStructurallyEquivalentComputedGroupingExpressionCanBeSelected() {
+        final TableEnvironment env = setupEnv(QueryOperationTestPrograms.AGGREGATE_QUERY_OPERATION);
+
+        final Table result =
+                env.from("s")
+                        .groupBy($("a").isGreater(10))
+                        .select($("a").isGreater(10).as("flag"), $("a").count().as("n"));
+
+        assertThat(result.getResolvedSchema().getColumnNames()).containsExactly("flag", "n");
+        assertGeneratedSqlCanBeParsed(env, result);
+    }
+
+    @Test
+    void testComputedGroupingExpressionCanBeSelectedAfterAggregate() {
+        final TableEnvironment env = setupEnv(QueryOperationTestPrograms.AGGREGATE_QUERY_OPERATION);
+        final Expression key = $("a").isGreater(10);
+
+        final Table result =
+                env.from("s").groupBy(key).aggregate($("a").count().as("n")).select(key, $("n"));
+
+        assertGeneratedSqlCanBeParsed(env, result);
+    }
+
+    @Test
+    void testColumnFunctionGroupingKeepsExpandedFields() {
+        final TableEnvironment env = setupEnv(QueryOperationTestPrograms.AGGREGATE_QUERY_OPERATION);
+
+        final Table result =
+                env.from("s")
+                        .groupBy(withColumns(1, 2))
+                        .select(withColumns(1, 2), $("a").count().as("n"));
+
+        assertThat(result.getResolvedSchema().getColumnNames()).containsExactly("a", "b", "n");
+        assertGeneratedSqlCanBeParsed(env, result);
+    }
+
+    @Test
+    void testNonGroupingExpressionStillCannotBeSelected() {
+        final TableEnvironment env = setupEnv(QueryOperationTestPrograms.AGGREGATE_QUERY_OPERATION);
+
+        assertThatThrownBy(() -> env.from("s").groupBy($("a").isGreater(10)).select($("a")))
+                .isInstanceOf(ValidationException.class)
+                .hasMessageContaining("Cannot resolve field [a]");
+    }
+
     @Test
     void testProctimePropertyOfEventTimeWindowCannotBeExpressedInWindowingTvfSyntax() {
         final TableEnvironment env =
@@ -164,6 +222,12 @@ void testProctimePropertyOfEventTimeWindowCannotBeExpressedInWindowingTvfSyntax(
                                 + "expressed in windowing-TVF syntax.");
     }
 
+    private static void assertGeneratedSqlCanBeParsed(TableEnvironment env, Table table) {
+        final String generatedSql =
+                table.getQueryOperation().asSerializableString(new InlineFunctionSqlFactory());
+        assertThatCode(() -> env.sqlQuery(generatedSql)).doesNotThrowAnyException();
+    }
+
     private static TableEnvironment setupEnv(TableTestProgram program) {
         final TableEnvironment env =
                 TableEnvironment.create(