[CALCITE-6340] RelBuilder always creates Project with Convention.NONE…

… during aggregate_ Calculate the post-pruning RelTraitSet on the projection using TraitSet#apply(Mapping)
asolimando · Apr 12, 2024 · 617bfa0 · 617bfa0
1 parent cc1d46a
commit 617bfa0
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git 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
@@ -547,6 +547,31 @@ private void checkSimplify(UnaryOperator<RelBuilder.Config> transform,
     assertThat(root, matcher);
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/projects/CALCITE/issues/CALCITE-6340">
+   * [CALCITE-6340] RelBuilder always creates Project with Convention.NONE during aggregate_.</a>.
+   */
+  @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 *