From 617bfa04c90956c650384fde77ded5ca1c467828 Mon Sep 17 00:00:00 2001
From: James Duong <james.duong@improving.com>
Date: Thu, 11 Apr 2024 15:26:24 -0700
Subject: [PATCH] [CALCITE-6340] RelBuilder always creates Project with
 Convention.NONE during aggregate_

Calculate the post-pruning RelTraitSet on the projection using TraitSet#apply(Mapping)
---
 .../org/apache/calcite/tools/RelBuilder.java  |  5 ++--
 .../apache/calcite/test/RelBuilderTest.java   | 25 +++++++++++++++++++
 2 files changed, 28 insertions(+), 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
index 8d8d48bd078..148f4302860 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 7f8f42c51b2..24b06aedee6 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<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 *