From 3767c7ecc5df9c9bcdcc2f86d71980209be776d0 Mon Sep 17 00:00:00 2001 From: Beetle brank <120192315+beetle0915@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:31:09 +0800 Subject: [PATCH 1/5] [FLINK-40344][table] Support selecting computed grouping expressions Rewrite computed grouping expressions to reference aggregate output fields and add regression tests. Generated-by: OpenAI Codex (GPT-5) --- .../flink/table/api/internal/TableImpl.java | 36 +++++++--- .../operations/AggregateQueryOperation.java | 2 + .../utils/OperationExpressionsUtils.java | 71 +++++++++++++++---- .../utils/OperationTreeBuilder.java | 9 ++- .../QueryOperationSqlSerializationTest.java | 50 +++++++++++++ 5 files changed, 141 insertions(+), 27 deletions(-) diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java index 8bdce4a023f762..3eba7d24bb574a 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java +++ b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java @@ -587,11 +587,13 @@ 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 groupingExpressions = table.preprocessExpressions(groupKeys); + QueryOperation aggregateOperation = + table.operationTreeBuilder.aggregate( + groupingExpressions, extracted.getAggregations(), table.operationTree); + + return table.projectAfterAggregation( + extracted.getProjections(), groupingExpressions, aggregateOperation); } @Override @@ -623,11 +625,13 @@ 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 groupingExpressions = table.preprocessExpressions(groupKeys); + QueryOperation aggregateOperation = + table.operationTreeBuilder.aggregate( + groupingExpressions, aggregateFunction, table.operationTree); + + return table.projectAfterAggregation( + table.preprocessExpressions(fields), groupingExpressions, aggregateOperation); } } @@ -950,6 +954,18 @@ private TableImpl createTable(QueryOperation operation) { return new TableImpl(tableEnvironment, operation, operationTreeBuilder, lookupResolver); } + private Table projectAfterAggregation( + List projections, + List groupingExpressions, + QueryOperation aggregateOperation) { + List rewrittenProjections = + OperationExpressionsUtils.replaceGroupingExpressions( + projections, + groupingExpressions, + aggregateOperation.getResolvedSchema().getColumnNames()); + return createTable(operationTreeBuilder.project(rewrittenProjections, aggregateOperation)); + } + private List preprocessExpressions(List expressions) { return preprocessExpressions(expressions.toArray(new Expression[0])); } diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/AggregateQueryOperation.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/AggregateQueryOperation.java index 147850e686e89f..44e217d43d10a9 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/AggregateQueryOperation.java +++ b/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( diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationExpressionsUtils.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationExpressionsUtils.java index 81b4daf380c182..9cb4cfc9b8237d 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationExpressionsUtils.java +++ b/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 extractedExpressionToFieldName = + new LinkedHashMap<>(splitter.aggregates); + extractedExpressionToFieldName.putAll(splitter.properties); + List projections = expressions.stream() .map( expr -> expr.accept( - new AggregationAndPropertiesReplacer( - splitter.aggregates, splitter.properties))) + new ExpressionReplacer( + extractedExpressionToFieldName))) .collect(Collectors.toList()); List 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. + * + *

Aggregate outputs place grouping fields before aggregate fields. Therefore, grouping + * expressions and output field names are matched by position. + */ + public static List replaceGroupingExpressions( + List projections, + List groupingExpressions, + List aggregateOutputFieldNames) { + if (groupingExpressions.size() > aggregateOutputFieldNames.size()) { + throw new IllegalArgumentException( + "The aggregate output does not contain all grouping expressions."); + } + + Map 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 nameExpressions(Map 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 { + private static class ExpressionReplacer extends ApiExpressionDefaultVisitor { - private final Map aggregates; - private final Map properties; + private final Map expressionToFieldName; - private AggregationAndPropertiesReplacer( - Map aggregates, Map properties) { - this.aggregates = aggregates; - this.properties = properties; + private ExpressionReplacer(Map 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 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 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> { @Override diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java index 6eb66edfb28495..30935454bc68da 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java +++ b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java @@ -247,8 +247,13 @@ public QueryOperation aggregate( List aggregates, QueryOperation child) { - ExpressionResolver resolver = getAggResolver(child, groupingExpressions); - List resolvedGroupings = resolver.resolve(groupingExpressions); + // Computed grouping expressions need stable names for subsequent projections. + List namedGroupingExpressions = + addAliasToTheCallInAggregate( + child.getResolvedSchema().getColumnNames(), groupingExpressions); + + ExpressionResolver resolver = getAggResolver(child, namedGroupingExpressions); + List resolvedGroupings = resolver.resolve(namedGroupingExpressions); List resolvedAggregates = resolver.resolve(aggregates); return aggregateOperationFactory.createAggregate( diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java index d3175a773bb0f1..c1180a2ece5334 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java +++ b/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; @@ -143,6 +144,49 @@ 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 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 +208,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( From 9a8550e7ea48c6abeb3c17484e6db59ad6d2c025 Mon Sep 17 00:00:00 2001 From: Beetle brank <120192315+beetle0915@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:35:41 +0800 Subject: [PATCH 2/5] [FLINK-40344][table] Expand grouping column functions before rewriting --- .../org/apache/flink/table/api/internal/TableImpl.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java index 3eba7d24bb574a..a310049e38a9db 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java +++ b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/TableImpl.java @@ -587,7 +587,9 @@ public Table select(Expression... fields) { "Window properties can only be used on windowed tables."); } - List groupingExpressions = table.preprocessExpressions(groupKeys); + List groupingExpressions = + table.operationTreeBuilder.expandExpressions( + table.preprocessExpressions(groupKeys), table.operationTree); QueryOperation aggregateOperation = table.operationTreeBuilder.aggregate( groupingExpressions, extracted.getAggregations(), table.operationTree); @@ -625,7 +627,9 @@ private AggregatedTableImpl( @Override public Table select(Expression... fields) { - List groupingExpressions = table.preprocessExpressions(groupKeys); + List groupingExpressions = + table.operationTreeBuilder.expandExpressions( + table.preprocessExpressions(groupKeys), table.operationTree); QueryOperation aggregateOperation = table.operationTreeBuilder.aggregate( groupingExpressions, aggregateFunction, table.operationTree); From 2dc4c87336376027568791388502408f373a296b Mon Sep 17 00:00:00 2001 From: Beetle brank <120192315+beetle0915@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:37:24 +0800 Subject: [PATCH 3/5] [FLINK-40344][table] Expand grouping column functions before rewriting --- .../flink/table/operations/utils/OperationTreeBuilder.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java index 30935454bc68da..871ded7f383c18 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java +++ b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java @@ -241,6 +241,10 @@ public QueryOperation dropColumns(List fieldLists, QueryOperation ch return project(finalFields, child, false); } + + public List expandExpressions(List expressions, QueryOperation child) { + return getResolver(child).resolveExpanding(expressions); + } public QueryOperation aggregate( List groupingExpressions, From f3c39dedcf244d8fefdb68bd8f5b65f32d2672e2 Mon Sep 17 00:00:00 2001 From: Beetle brank <120192315+beetle0915@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:39:38 +0800 Subject: [PATCH 4/5] [FLINK-40344][table] Expand grouping column functions before rewriting Add test for column function grouping to verify expanded fields are retained in the result schema. --- .../api/QueryOperationSqlSerializationTest.java | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java index c1180a2ece5334..5e2a1776a4d1f3 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/api/QueryOperationSqlSerializationTest.java @@ -46,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; @@ -178,6 +179,19 @@ void testComputedGroupingExpressionCanBeSelectedAfterAggregate() { 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); From cd82e7c5e09616abbfa9026d714ffee300b94eda Mon Sep 17 00:00:00 2001 From: Beetle brank <120192315+beetle0915@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:46:26 +0800 Subject: [PATCH 5/5] [FLINK-40344][table] Fix formatting --- .../flink/table/operations/utils/OperationTreeBuilder.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java index 871ded7f383c18..02030ab42ce450 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java +++ b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java @@ -241,7 +241,7 @@ public QueryOperation dropColumns(List fieldLists, QueryOperation ch return project(finalFields, child, false); } - + public List expandExpressions(List expressions, QueryOperation child) { return getResolver(child).resolveExpanding(expressions); }