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,10 @@
 import io.substrait.relation.Aggregate;
 import io.substrait.relation.Project;
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
@@ -34,7 +37,28 @@ public static boolean isValidCalciteAggregate(Aggregate aggregate) {
     return aggregate.getMeasures().stream()
             .allMatch(PreCalciteAggregateValidator::isValidCalciteMeasure)
         && aggregate.getGroupings().stream()
-            .allMatch(PreCalciteAggregateValidator::isValidCalciteGrouping);
+            .allMatch(PreCalciteAggregateValidator::isValidCalciteGrouping)
+        && aLoneGroupingSetNamesEachExpressionOnce(aggregate);
+  }
+
+  /**
+   * Checks that an aggregate holding one grouping set does not name an expression in it twice.
+   *
+   * <p>Calcite holds a grouping set in an {@link org.apache.calcite.util.ImmutableBitSet}, which
+   * cannot hold a field twice, while a lone grouping set gives each mention a column of its own --
+   * {@code Aggregate.deriveRecordType} dedups the grouping expressions only across several sets. So
+   * a repeat has to reach Calcite as two columns of the relation underneath, which is what the
+   * transformer makes of it.
+   *
+   * @param aggregate the aggregate relation
+   * @return {@code true} if valid, {@code false} otherwise
+   */
+  private static boolean aLoneGroupingSetNamesEachExpressionOnce(Aggregate aggregate) {
+    if (aggregate.getGroupings().size() != 1) {
+      return true;
+    }
+    List<Expression> expressions = aggregate.getGroupings().get(0).getExpressions();
+    return new HashSet<>(expressions).size() == expressions.size();
   }
 
   /**
@@ -60,32 +84,19 @@ private static boolean isValidCalciteMeasure(Aggregate.Measure measure) {
   }
 
   /**
-   * Checks if an {@link Aggregate.Grouping} uses only {@link FieldReference}s and ensures grouping
-   * fields are in ascending order.
+   * Checks if an {@link Aggregate.Grouping} uses only {@link FieldReference}s.
+   *
+   * <p>The order the fields are grouped in is not a reason to rewrite the aggregate. Calcite holds
+   * a grouping set in an {@link org.apache.calcite.util.ImmutableBitSet} and emits its grouping
+   * columns in ascending field order whatever order they were declared in, so a plan grouping on
+   * (0, 2, 1) reaches Calcite as (0, 1, 2); the conversion carries the declared order in the emit
+   * mapping instead.
    *
    * @param grouping the aggregate grouping to validate
    * @return {@code true} if valid, {@code false} otherwise
    */
   private static boolean isValidCalciteGrouping(Aggregate.Grouping grouping) {
-    if (!grouping.getExpressions().stream().allMatch(e -> isSimpleFieldReference(e))) {
-      return false;
-    }
-
-    // Calcite stores grouping fields in an ImmutableBitSet and does not track the order of the
-    // grouping fields. The output record shape that Calcite generates ALWAYS has the groupings in
-    // ascending field order. This causes issues with Substrait in cases where the grouping fields
-    // in Substrait are not defined in ascending order.
-
-    // For example, if a grouping is defined as (0, 2, 1) in Substrait, Calcite will output it as
-    // (0, 1, 2), which means that the Calcite output will no longer line up with the expectations
-    // of the Substrait plan.
-
-    List<Integer> groupingFields =
-        grouping.getExpressions().stream()
-            .map(expr -> getFieldRefOffset((FieldReference) expr))
-            .collect(Collectors.toList());
-
-    return isOrdered(groupingFields);
+    return grouping.getExpressions().stream().allMatch(e -> isSimpleFieldReference(e));
   }
 
   private static boolean isSimpleFieldReference(FunctionArg e) {
@@ -97,43 +108,39 @@ private static boolean isSimpleFieldReference(FunctionArg e) {
     return segments.size() == 1 && segments.get(0) instanceof FieldReference.StructField;
   }
 
-  private static int getFieldRefOffset(FieldReference fr) {
-    return ((FieldReference.StructField) fr.segments().get(0)).offset();
-  }
-
-  private static boolean isOrdered(List<Integer> list) {
-    for (int i = 1; i < list.size(); i++) {
-      if (list.get(i - 1) > list.get(i)) {
-        return false;
-      }
-    }
-    return true;
-  }
-
   /**
-   * Transforms invalid aggregates into Calcite-compatible form by projecting non-field expressions
-   * and reordering groupings.
+   * Transforms invalid aggregates into Calcite-compatible form by projecting out the grouping
+   * expressions Calcite cannot hold as they are.
    */
   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, kept only where the
+    // aggregate's own record type shares a column between mentions of one expression: with several
+    // grouping sets a field grouped on by two of them is one column of the output, so it has to
+    // stay one column of the project underneath -- two copies of it would each be missing from a
+    // set, and Calcite would make both nullable. A lone grouping set gives every mention a column
+    // of its own, so there the map is not consulted.
+    private final Map<Expression, Expression> projectedGroupingExpressions;
+
+    private final boolean groupingColumnsAreShared;
+
     // Tracks the offset of the next expression added
     private int expressionOffset;
 
     private PreCalciteAggregateTransformer(Aggregate aggregate) {
       this.newExpressions = new ArrayList<>();
+      this.projectedGroupingExpressions = new HashMap<>();
+      this.groupingColumnsAreShared = aggregate.getGroupings().size() > 1;
       this.expressionOffset = aggregate.getInput().getRecordType().fields().size();
     }
 
     /**
-     * Rewrites an {@link Aggregate} so that it can be converted to Calcite by:
-     *
-     * <ul>
-     *   <li>Projecting non-field references before aggregation
-     *   <li>Ensuring groupings are in ascending order
-     * </ul>
+     * Rewrites an {@link Aggregate} so that it can be converted to Calcite by projecting the
+     * grouping expressions and the measures' non-field arguments out before the aggregation, so
+     * that each is a field reference of its own.
      *
      * @param aggregate the original Substrait aggregate
      * @return a transformed Calcite-compatible aggregate
@@ -193,7 +200,13 @@ 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 ->
+                      groupingColumnsAreShared
+                          ? projectedGroupingExpressions.computeIfAbsent(expr, this::projectOut)
+                          : projectOut(expr))
+              .collect(Collectors.toList());
       return Aggregate.Grouping.builder().expressions(newGroupingExpressions).build();
     }
 

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

@@ -50,10 +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.LinkedList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -361,13 +362,17 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
             .map(measure -> fromMeasure(measure, context, child, hasEmptyGroup))
             .collect(java.util.stream.Collectors.toList());
 
-    Optional<Remap> remap = aggregate.getRemap();
-    final int lastFieldIndex = groupExprs.size() + aggregateCalls.size();
+    final Optional<Remap> remap = aggregate.getRemap();
+    // 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 groupingSetIndex = groupColumnCount + aggregateCalls.size();
 
-    // map grouping set index if it is not removed via remap
+    // The index is a column of the converted aggregate only where the relation emits it: an
+    // aggregate that maps its output away does not need the call at all.
     final boolean emitDirect = remap.isEmpty();
     final boolean groupingSetIndexGetsRemapped =
-        remap.map(r -> r.indices().contains(lastFieldIndex)).orElse(false);
+        remap.map(r -> r.indices().contains(groupingSetIndex)).orElse(false);
     if (aggregate.getGroupings().size() > 1 && (emitDirect || groupingSetIndexGetsRemapped)) {
       aggregateCalls.add(
           AggregateCall.create(
@@ -382,17 +387,6 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
               RelCollations.EMPTY,
               typeConverter.toCalcite(typeFactory, TypeCreator.REQUIRED.I64),
               null));
-      final int groupingCallIndex = aggregateCalls.size() - 1;
-      if (groupingSetIndexGetsRemapped) {
-        List<Integer> remapList = new LinkedList<>(remap.get().indices());
-        for (int i = 0; i < remapList.size(); i++) {
-          if (remapList.get(i).equals(lastFieldIndex)) {
-            // replace last field index with field index of the GROUP_ID() function call
-            remapList.set(i, groupingCallIndex);
-          }
-        }
-        remap = Optional.of(Remap.of(remapList));
-      }
     }
 
     exitUncorrelatedScope(context, Aggregate.class);
@@ -410,7 +404,68 @@ 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-java takes the grouping columns of an aggregate to be the distinct grouping
+   * expressions in the order they first appear across its grouping sets. The spec orders them by
+   * the relation's shared grouping-expression list, which each set's expression references index
+   * into; the POJO models a per-set expression list and cannot hold that list's order, so
+   * first-appearance is the reconstruction this library reads and writes. 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)));
   }
 
   /**
@@ -1191,10 +1246,10 @@ protected RelNode applyRelCommon(RelNode relNode, Rel rel, RelNode... inputs) {
    * inputs are compared against.
    *
    * <p>They are dropped as well where the columns of that projection are not the columns of the
-   * relation's record type, type by type. An aggregate over several grouping sets orders its
-   * grouping columns by first appearance where Calcite orders them by group key, so the two
-   * disagree on what the third column is, and binding the names by position would name columns the
-   * plan does not name.
+   * relation's record type, type by type. An aggregate over several grouping sets types its
+   * grouping-set index i32, where the GROUP_ID call standing for it is i64, so the two disagree on
+   * what the last column is and binding the names by position would name columns the plan does not
+   * name.
    *
    * <p>Only the names of the top-level fields are applied. The names of the fields nested inside
    * them belong to the type of the expression that produces the field, which a projection cannot