[CALCITE-676] AssertionError in GROUPING SETS query

core/src/main/java/org/apache/calcite/rel/rules/AggregateProjectMergeRule.java

@@ -31,7 +31,9 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Planner rule that recognizes a {@link org.apache.calcite.rel.core.Aggregate}
@@ -67,31 +69,25 @@ public void onMatch(RelOptRuleCall call) {
   public static RelNode apply(Aggregate aggregate,
       Project project) {
     final List<Integer> newKeys = Lists.newArrayList();
+    final Map<Integer, Integer> map = new HashMap<>();
     for (int key : aggregate.getGroupSet()) {
       final RexNode rex = project.getProjects().get(key);
       if (rex instanceof RexInputRef) {
-        newKeys.add(((RexInputRef) rex).getIndex());
+        final int newKey = ((RexInputRef) rex).getIndex();
+        newKeys.add(newKey);
+        map.put(key, newKey);
       } else {
         // Cannot handle "GROUP BY expression"
         return null;
       }
     }
 
-    final ImmutableBitSet newGroupSet = ImmutableBitSet.of(newKeys);
-
+    final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
     ImmutableList<ImmutableBitSet> newGroupingSets = null;
     if (aggregate.indicator) {
-      ImmutableList.Builder<ImmutableBitSet> newGroupingSetsBuilder =
-              ImmutableList.builder();
-      for (ImmutableBitSet groupingSet: aggregate.getGroupSets()) {
-        final ImmutableBitSet.Builder newGroupingSet =
-                ImmutableBitSet.builder();
-        for (int c : groupingSet) {
-          newGroupingSet.set(newKeys.get(c));
-        }
-        newGroupingSetsBuilder.add(newGroupingSet.build());
-      }
-      newGroupingSets = newGroupingSetsBuilder.build();
+      newGroupingSets =
+          ImmutableBitSet.ORDERING.immutableSortedCopy(
+              ImmutableBitSet.permute(aggregate.getGroupSets(), map));
     }
 
     final ImmutableList.Builder<AggregateCall> aggCalls =
@@ -118,6 +114,7 @@ public static RelNode apply(Aggregate aggregate,
     // Add a project if the group set is not in the same order or
     // contains duplicates.
     RelNode rel = newAggregate;
+    //noinspection EqualsBetweenInconvertibleTypes
     if (!newGroupSet.toList().equals(newKeys)) {
       final List<Integer> posList = Lists.newArrayList();
       for (int newKey : newKeys) {

core/src/main/java/org/apache/calcite/util/ImmutableBitSet.java

@@ -34,7 +34,9 @@
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.SortedMap;
+import javax.annotation.Nonnull;
 
 /**
  * An immutable list of bits.
@@ -569,7 +571,7 @@ public List<Integer> asList() {
         return cardinality();
       }
 
-      @Override public Iterator<Integer> iterator() {
+      @Nonnull @Override public Iterator<Integer> iterator() {
         return ImmutableBitSet.this.iterator();
       }
     };
@@ -751,6 +753,27 @@ public BitSet toBitSet() {
     return BitSets.of(this);
   }
 
+  /** Permutes a bit set according to a given mapping. */
+  public ImmutableBitSet permute(Map<Integer, Integer> map) {
+    final Builder builder = builder();
+    for (int i = nextSetBit(0); i >= 0; i = nextSetBit(i + 1)) {
+      builder.set(map.get(i));
+    }
+    return builder.build();
+  }
+
+  /** Permutes a collection of bit sets according to a given mapping. */
+  public static Iterable<ImmutableBitSet> permute(
+      Iterable<ImmutableBitSet> bitSets,
+      final Map<Integer, Integer> map) {
+    return Iterables.transform(bitSets,
+        new Function<ImmutableBitSet, ImmutableBitSet>() {
+          public ImmutableBitSet apply(ImmutableBitSet bitSet) {
+            return bitSet.permute(map);
+          }
+        });
+  }
+
   /**
    * Setup equivalence Sets for each position. If i & j are equivalent then
    * they will have the same equivalence Set. The algorithm computes the

core/src/test/resources/sql/join.oq

@@ -134,4 +134,47 @@ EnumerableJoin(condition=[=($0, $1)], joinType=[inner])
       EnumerableTableScan(table=[[scott, EMP]])
 !plan
 
+select distinct dept.deptno
+from "scott".emp join "scott".dept using (deptno);
++--------+
+| DEPTNO |
++--------+
+|     10 |
+|     20 |
+|     30 |
++--------+
+(3 rows)
+
+!ok
+EnumerableAggregate(group=[{0}])
+  EnumerableCalc(expr#0..1=[{inputs}], DEPTNO=[$t0])
+    EnumerableJoin(condition=[=($0, $1)], joinType=[inner])
+      EnumerableCalc(expr#0..2=[{inputs}], DEPTNO=[$t0])
+        EnumerableTableScan(table=[[scott, DEPT]])
+      EnumerableCalc(expr#0..7=[{inputs}], DEPTNO=[$t7])
+        EnumerableTableScan(table=[[scott, EMP]])
+!plan
+
+# [CALCITE-676] AssertionError in GROUPING SETS query
+select emp.deptno as e, dept.deptno as d
+from "scott".emp join "scott".dept using (deptno)
+group by cube(emp.deptno, dept.deptno);
++----+----+
+| E  | D  |
++----+----+
+| 10 | 10 |
+| 10 |    |
+| 20 | 20 |
+| 20 |    |
+| 30 | 30 |
+| 30 |    |
+|    | 10 |
+|    | 20 |
+|    | 30 |
+|    |    |
++----+----+
+(10 rows)
+
+!ok
+
 # End join.oq