fix(isthmus)!: keep the declared column order of an aggregate over grouping sets

isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java

@@ -7,7 +7,9 @@
 import io.substrait.relation.Aggregate;
 import io.substrait.relation.Project;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
@@ -119,11 +121,18 @@ public static class PreCalciteAggregateTransformer {
     // New expressions to include in the project before the aggregate
     private final List<Expression> newExpressions;
 
+    // The field reference each grouping expression was projected out to. A field grouped on by
+    // several grouping sets is one column of the aggregate's output, so it has to stay one column
+    // of the project underneath it: two copies of it would each be missing from a grouping set,
+    // and Calcite would make both of them nullable.
+    private final Map<Expression, Expression> projectedGroupingExpressions;
+
     // Tracks the offset of the next expression added
     private int expressionOffset;
 
     private PreCalciteAggregateTransformer(Aggregate aggregate) {
       this.newExpressions = new ArrayList<>();
+      this.projectedGroupingExpressions = new HashMap<>();
       this.expressionOffset = aggregate.getInput().getRecordType().fields().size();
     }
 
@@ -193,7 +202,9 @@ private Aggregate.Measure updateMeasure(Aggregate.Measure measure) {
 
     private Aggregate.Grouping updateGrouping(Aggregate.Grouping grouping) {
       List<Expression> newGroupingExpressions =
-          grouping.getExpressions().stream().map(this::projectOut).collect(Collectors.toList());
+          grouping.getExpressions().stream()
+              .map(expr -> projectedGroupingExpressions.computeIfAbsent(expr, this::projectOut))
+              .collect(Collectors.toList());
       return Aggregate.Grouping.builder().expressions(newGroupingExpressions).build();
     }
 

isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java

@@ -50,9 +50,11 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Locale;
@@ -363,7 +365,10 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
             .collect(java.util.stream.Collectors.toList());
 
     Optional<Remap> remap = aggregate.getRemap();
-    final int lastFieldIndex = groupExprs.size() + aggregateCalls.size();
+    // A field grouped on by several sets is one column of the relation, so the grouping-set index
+    // sits after the distinct grouping expressions, not after every mention of them.
+    final int groupColumnCount = new LinkedHashSet<>(groupExprs).size();
+    final int lastFieldIndex = groupColumnCount + aggregateCalls.size();
 
     // map grouping set index if it is not removed via remap
     final boolean emitDirect = remap.isEmpty();
@@ -383,7 +388,9 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
               RelCollations.EMPTY,
               typeConverter.toCalcite(typeFactory, TypeCreator.REQUIRED.I64),
               null));
-      final int groupingCallIndex = aggregateCalls.size() - 1;
+      // The call was appended, so it is the last column of the converted aggregate: the grouping
+      // columns come first, then the calls.
+      final int groupingCallIndex = groupColumnCount + aggregateCalls.size() - 1;
       if (groupingSetIndexGetsRemapped) {
         List<Integer> remapList = new LinkedList<>(remap.get().indices());
         for (int i = 0; i < remapList.size(); i++) {
@@ -411,7 +418,65 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
     RelNode node = aggregateBuilder.push(child).aggregate(groupKey, aggregateCalls).build();
     // Not applyRelCommon: the mapping applied here is the one rewritten above, not the one the
     // relation carries.
-    return applyOutputNames(applyRemap(node, remap), aggregate, child);
+    return applyOutputNames(
+        applyRemap(node, inConvertedGroupingOrder(remap, groupExprs, aggregateCalls.size())),
+        aggregate,
+        child);
+  }
+
+  /**
+   * Returns the emit mapping of a converted aggregate with its indices translated from the order
+   * the relation declares its output in to the order the converted aggregate emits it.
+   *
+   * <p>Substrait takes the grouping columns of an aggregate to be the distinct grouping expressions
+   * in the order they first appear across its grouping sets. Calcite takes them from a bit set, so
+   * it emits them ordered by field index. A relation whose grouping sets first mention field 1 and
+   * then field 0 declares them in that order, and its emit mapping indexes that order, while the
+   * aggregate underneath emits field 0 first.
+   *
+   * <p>An aggregate that emits directly and declares an order Calcite does not produce gets a
+   * mapping it did not carry, which is what puts the columns back in the declared order.
+   *
+   * @param remap the emit mapping the relation carries, indexing its declared output
+   * @param groupExprs the converted grouping expressions, in declared order, with duplicates
+   * @param callCount the number of aggregate calls, including any grouping-set index
+   * @return the mapping to apply to the converted aggregate
+   */
+  private static Optional<Remap> inConvertedGroupingOrder(
+      Optional<Remap> remap, List<RexNode> groupExprs, int callCount) {
+    List<RexNode> declared = new ArrayList<>(new LinkedHashSet<>(groupExprs));
+    // Calcite emits the grouping columns in the order they sit in the aggregate's input: a field
+    // reference where its field sits, and anything else -- an outer reference, which the transform
+    // above leaves alone -- in the projection Calcite adds after them, in the order it was
+    // declared. Sorting is stable, so giving the second kind one key keeps that order among them.
+    List<RexNode> converted =
+        declared.stream()
+            .sorted(
+                Comparator.comparingInt(
+                    expr ->
+                        expr instanceof RexInputRef
+                            ? ((RexInputRef) expr).getIndex()
+                            : Integer.MAX_VALUE))
+            .collect(Collectors.toList());
+    if (converted.equals(declared)) {
+      return remap;
+    }
+    List<Integer> declaredToConverted = new ArrayList<>();
+    for (RexNode expression : declared) {
+      declaredToConverted.add(converted.indexOf(expression));
+    }
+    for (int call = 0; call < callCount; call++) {
+      declaredToConverted.add(declared.size() + call);
+    }
+    return Optional.of(
+        Remap.of(
+            remap
+                .map(
+                    mapping ->
+                        mapping.indices().stream()
+                            .map(declaredToConverted::get)
+                            .collect(Collectors.toList()))
+                .orElse(declaredToConverted)));
   }
 
   /**

isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java

@@ -40,6 +40,7 @@
 import io.substrait.type.TypeCreator;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -414,32 +415,33 @@ public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) {
         Aggregate.builder().input(input).addAllGroupings(groupings).addAllMeasures(aggCalls);
 
     if (groupings.size() > 1) {
+      // Substrait declares the grouping columns of an aggregate as the distinct grouping
+      // expressions in the order they first appear across its grouping sets, while Calcite emits
+      // them ordered by field index. Where the two differ, the emit mapping carries the reordering,
+      // so that a parent converted from the same Calcite plan finds its columns where it left them.
+      List<Integer> groupingRemap = calciteGroupingOrder(groupings);
+
       // remove the grouping set index if there was no explicit GROUP_ID() function call
       if (groupIdCalls.isEmpty()) {
-        builder.remap(Remap.offset(0, groupingFieldCount + aggCalls.size()));
+        List<Integer> remap = new ArrayList<>(groupingRemap);
+        for (int call = 0; call < aggCalls.size(); call++) {
+          remap.add(groupingFieldCount + call);
+        }
+        builder.remap(Remap.of(remap));
       } else {
-        // remap grouping set index at the field positions where the GROUP_ID() function calls were.
-        // Use the non-distinct total here: when grouping sets share expressions the aggregate
-        // output
-        // contains one slot per (groupingSet × expression) entry, not one per distinct expression.
-        final int groupingFieldCountWithDuplicates =
-            Math.toIntExact(groupings.stream().flatMap(g -> g.getExpressions().stream()).count());
+        // remap grouping set index at the field positions where the GROUP_ID() function calls were
         final int filterAggCallCount = aggCalls.size();
-        final Integer groupingSetIndex = groupingFieldCountWithDuplicates + filterAggCallCount;
+        final Integer groupingSetIndex = groupingFieldCount + filterAggCallCount;
 
-        final List<Integer> remap =
-            IntStream.range(0, groupingFieldCountWithDuplicates)
-                .mapToObj(i -> i)
-                .collect(Collectors.toCollection(ArrayList::new));
+        final List<Integer> remap = new ArrayList<>(groupingRemap);
 
         for (int i = 0; i < aggregate.getAggCallList().size(); i++) {
           AggregateCall aggCall = aggregate.getAggCallList().get(i);
           if (filteredAggCalls.contains(aggCall)) {
             remap.add(
-                i + groupingFieldCountWithDuplicates,
-                filteredAggCalls.indexOf(aggCall) + groupingFieldCountWithDuplicates);
+                i + groupingFieldCount, filteredAggCalls.indexOf(aggCall) + groupingFieldCount);
           } else if (groupIdCalls.contains(aggCall)) {
-            remap.add(i + groupingFieldCountWithDuplicates, groupingSetIndex);
+            remap.add(i + groupingFieldCount, groupingSetIndex);
           } else {
             // this should never get triggered
             throw new IllegalStateException(
@@ -497,6 +499,42 @@ public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) {
         .build();
   }
 
+  /**
+   * Returns, for each grouping column of the converted Calcite aggregate, the position that column
+   * holds in the output the Substrait aggregate declares.
+   *
+   * <p>Substrait takes the grouping columns to be the distinct grouping expressions in the order
+   * they first appear across the grouping sets; Calcite takes them from a bit set and so emits them
+   * ordered by field index. Reading the result as an emit mapping presents the aggregate's output
+   * in Calcite's order.
+   *
+   * @param groupings the grouping sets of the converted aggregate
+   * @return the declared position of each grouping column, in the order Calcite emits them
+   */
+  private static List<Integer> calciteGroupingOrder(List<Grouping> groupings) {
+    List<Expression> declared =
+        groupings.stream()
+            .flatMap(grouping -> grouping.getExpressions().stream())
+            .distinct()
+            .collect(Collectors.toList());
+    return declared.stream()
+        .sorted(Comparator.comparingInt(SubstraitRelVisitor::groupingFieldOffset))
+        .map(declared::indexOf)
+        .collect(Collectors.toList());
+  }
+
+  /**
+   * Returns the field the given grouping expression references.
+   *
+   * @param expression a grouping expression, as built by {@link #fromGroupSet(ImmutableBitSet,
+   *     Rel)}
+   * @return the offset of the field it references
+   */
+  private static int groupingFieldOffset(Expression expression) {
+    FieldReference reference = (FieldReference) expression;
+    return ((FieldReference.StructField) reference.segments().get(0)).offset();
+  }
+
   Aggregate.Grouping fromGroupSet(ImmutableBitSet bitSet, Rel input) {
     List<Expression> references =
         bitSet.asList().stream()