diff --git a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java index 8d8d48bd078a..148f43028604 100644 --- a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java +++ b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java @@ -2502,9 +2502,10 @@ private RelBuilder pruneAggregateInputFieldsAndDeduplicateAggCalls( newProjects.add(project.getProjects().get(i)); builder.add(project.getRowType().getFieldList().get(i)); } + r = - project.copy(cluster.traitSet(), project.getInput(), newProjects, - builder.build()); + project.copy(project.getTraitSet().apply(targetMapping), project.getInput(), + newProjects, builder.build()); } else { groupSetAfterPruning = groupSet; groupSetsAfterPruning = groupSets; diff --git a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java index 7f8f42c51b24..24b06aedee62 100644 --- a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java +++ b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java @@ -547,6 +547,31 @@ private void checkSimplify(UnaryOperator transform, assertThat(root, matcher); } + /** Test case for + * + * [CALCITE-6340] RelBuilder always creates Project with Convention.NONE during aggregate_.. + */ + @Test void testPruneProjectInputOfAggregatePreservesTraitSet() { + final RelBuilder builder = createBuilder(config -> config.withPruneInputOfAggregate(true)); + + // This issue only occurs when projecting more columns than there are fields and putting + // an aggregate over that projection. + RelNode root = + builder.scan("DEPT") + .adoptConvention(EnumerableConvention.INSTANCE) + .project(builder.alias(builder.field(0), "a"), + builder.alias(builder.field(1), "b"), + builder.alias(builder.field(2), "c"), + builder.alias(builder.field(1), "d")) + .aggregate(builder.groupKey(0, 1, 2, 1), + builder.aggregateCall(SqlStdOperatorTable.SUM, + builder.field(0))) + .build(); + + // Verify that the project under the aggregate kept the EnumerableConvention.INSTANCE trait. + assertTrue(root.getInput(0).getTraitSet().contains(EnumerableConvention.INSTANCE)); + } + @Test void testScanFilterOr() { // Equivalent SQL: // SELECT *