From 291e3716361ce7feeebfb2bb4aa5ce309de6ca19 Mon Sep 17 00:00:00 2001
From: Julian Hyde <jhyde@apache.org>
Date: Mon, 2 Mar 2020 16:29:48 -0800
Subject: [PATCH] [CALCITE-3839] After calling RelBuilder.aggregate, cannot
 lookup field by name

---
 .../org/apache/calcite/tools/RelBuilder.java  | 12 +++++++---
 .../apache/calcite/test/RelBuilderTest.java   | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+), 3 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 2f4e3a5a082..b4283d47178 100644
--- a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
+++ b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
@@ -1715,6 +1715,7 @@ public RelBuilder aggregate(GroupKey groupKey, Iterable<AggCall> aggCalls) {
       assert groupSet.contains(set);
     }
 
+    List<Field> inFields = frame.fields;
     if (config.pruneInputOfAggregate()
         && r instanceof Project) {
       final Set<Integer> fieldsUsed =
@@ -1744,6 +1745,7 @@ public RelBuilder aggregate(GroupKey groupKey, Iterable<AggCall> aggCalls) {
         for (AggregateCall aggregateCall : oldAggregateCalls) {
           aggregateCalls.add(aggregateCall.transform(targetMapping));
         }
+        inFields = Mappings.permute(inFields, targetMapping.inverse());
 
         final Project project = (Project) r;
         final List<RexNode> newProjects = new ArrayList<>();
@@ -1760,7 +1762,7 @@ public RelBuilder aggregate(GroupKey groupKey, Iterable<AggCall> aggCalls) {
 
     if (!config.dedupAggregateCalls() || Util.isDistinct(aggregateCalls)) {
       return aggregate_(groupSet, groupSets, r, aggregateCalls,
-          registrar.extraNodes, frame.fields);
+          registrar.extraNodes, inFields);
     }
 
     // There are duplicate aggregate calls. Rebuild the list to eliminate
@@ -1782,7 +1784,7 @@ public RelBuilder aggregate(GroupKey groupKey, Iterable<AggCall> aggCalls) {
       projects.add(Pair.of(groupSet.cardinality() + i, aggregateCall.name));
     }
     aggregate_(groupSet, groupSets, r, distinctAggregateCalls,
-        registrar.extraNodes, frame.fields);
+        registrar.extraNodes, inFields);
     final List<RexNode> fields = projects.stream()
         .map(p -> p.right == null ? field(p.left)
             : alias(field(p.left), p.right))
@@ -1795,7 +1797,7 @@ public RelBuilder aggregate(GroupKey groupKey, Iterable<AggCall> aggCalls) {
   private RelBuilder aggregate_(ImmutableBitSet groupSet,
       ImmutableList<ImmutableBitSet> groupSets, RelNode input,
       List<AggregateCall> aggregateCalls, List<RexNode> extraNodes,
-      ImmutableList<Field> inFields) {
+      List<Field> inFields) {
     final RelNode aggregate =
         struct.aggregateFactory.createAggregate(input,
             ImmutableList.of(), groupSet, groupSets, aggregateCalls);
@@ -3004,6 +3006,10 @@ private Frame(RelNode rel) {
       this.fields = builder.build();
     }
 
+    @Override public String toString() {
+      return rel + ": " + fields;
+    }
+
     private static String deriveAlias(RelNode rel) {
       if (rel instanceof TableScan) {
         final List<String> names = rel.getTable().getQualifiedName();
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 a6038f1ed8e..09d8aba9ce6 100644
--- a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java
@@ -1081,6 +1081,28 @@ private RexNode caseCall(RelBuilder b, RexNode ref, RexNode... nodes) {
     assertThat(root, hasTree(expected));
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-3839">[CALCITE-3839]
+   * After calling RelBuilder.aggregate, cannot lookup field by name</a>. */
+  @Test public void testAggregateAndThenProjectNamedField() {
+    final Function<RelBuilder, RelNode> f = builder ->
+        builder.scan("EMP")
+            .project(builder.field("EMPNO"), builder.field("ENAME"),
+                builder.field("SAL"))
+            .aggregate(builder.groupKey(builder.field("ENAME")),
+                builder.sum(builder.field("SAL")))
+            // Before [CALCITE-3839] was fixed, the following line gave
+            // 'field [ENAME] not found'
+            .project(builder.field("ENAME"))
+            .build();
+    final String expected = ""
+        + "LogicalProject(ENAME=[$0])\n"
+        + "  LogicalAggregate(group=[{0}], agg#0=[SUM($1)])\n"
+        + "    LogicalProject(ENAME=[$1], SAL=[$5])\n"
+        + "      LogicalTableScan(table=[[scott, EMP]])\n";
+    assertThat(f.apply(createBuilder(c -> c)), hasTree(expected));
+  }
+
   /** Tests that {@link RelBuilder#aggregate} eliminates duplicate aggregate
    * calls and creates a {@code Project} to compensate. */
   @Test public void testAggregateEliminatesDuplicateCalls() {