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

… during aggregate_ Propagate the trait set after pruning a projection under an aggregate in RelBuilder.
apache · Apr 12, 2024 · 6f4a096 · 6f4a096
1 parent e7c1f0c
commit 6f4a096
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
@@ -16,6 +16,7 @@
  */
 package org.apache.calcite.tools;
 
+import org.apache.calcite.linq4j.Nullness;
 import org.apache.calcite.linq4j.Ord;
 import org.apache.calcite.linq4j.function.Experimental;
 import org.apache.calcite.plan.Context;
@@ -27,6 +28,7 @@
 import org.apache.calcite.plan.RelOptSchema;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.plan.ViewExpanders;
 import org.apache.calcite.prepare.RelOptTableImpl;
 import org.apache.calcite.rel.RelCollation;
@@ -2502,8 +2504,12 @@ private RelBuilder pruneAggregateInputFieldsAndDeduplicateAggCalls(
           newProjects.add(project.getProjects().get(i));
           builder.add(project.getRowType().getFieldList().get(i));
         }
+        final RelTraitSet traitSet =
+            project.getConvention() != null ? cluster.traitSet().replace(
+                Nullness.castNonNull(project.getConvention())) : cluster.traitSet();
+
         r =
-            project.copy(cluster.traitSet(), project.getInput(), newProjects,
+            project.copy(traitSet, project.getInput(), newProjects,
                 builder.build());
       } else {
         groupSetAfterPruning = groupSet;

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 *