From c21feb6950c8469ef3c7a08081751286cb5e5b86 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Wed, 26 Aug 2026 01:18:55 +0300 Subject: [PATCH 1/4] fix(isthmus)!: keep the declared column order of an aggregate over grouping sets Substrait takes the grouping columns of an aggregate to be the distinct grouping expressions in the order they first appear across its grouping sets, while Calcite takes them from a bit set and emits them ordered by field index. Neither direction accounted for that, so a plan whose sets first mention field 1 and then field 0 changed meaning on the way through: a reference to the aggregate's first column reached the column Calcite had put there instead, with its type quietly changing along with it. Both directions now carry the difference in the emit mapping. On the way in, the mapping a relation carries is translated into the order the converted aggregate emits, and a relation that emits directly gets the mapping that puts its columns back in the declared order. On the way out, the mapping presents the aggregate's output in Calcite's order, so a parent converted from the same Calcite plan finds its columns where it left them. Neither adds a relation the plan did not have, so a plan that already agrees with Calcite round-trips unchanged. The pre-aggregate projection now reuses one column for a field grouped on by several sets. Two copies of it would each be missing from a grouping set, and Calcite would make both nullable. Closes #1159 BREAKING CHANGE: an aggregate over several grouping sets is now emitted with an emit mapping that presents its output in the order the plan it came from had, and a plan carrying such a mapping is read that way. Consumers that assumed the grouping columns were ordered by field index will see them in the order the grouping sets declare. --- .../isthmus/PreCalciteAggregateValidator.java | 13 +- .../isthmus/SubstraitRelNodeConverter.java | 58 +++++++- .../isthmus/SubstraitRelVisitor.java | 68 +++++++-- .../isthmus/ComplexAggregateTest.java | 132 ++++++++++++++++++ 4 files changed, 254 insertions(+), 17 deletions(-) diff --git a/isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java b/isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java index f8695db64..159a855b9 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java +++ b/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 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 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 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(); } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index ad41b6803..c525b5d56 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/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; @@ -411,7 +413,61 @@ 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. + * + *

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. + * + *

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 inConvertedGroupingOrder( + Optional remap, List groupExprs, int callCount) { + List declared = new ArrayList<>(new LinkedHashSet<>(groupExprs)); + if (!declared.stream().allMatch(RexInputRef.class::isInstance)) { + // The conversion projects expressions that are not field references below the aggregate, in + // the order they were declared, and groups over that projection, so the orders agree. + return remap; + } + List converted = + declared.stream() + .sorted(Comparator.comparingInt(expr -> ((RexInputRef) expr).getIndex())) + .collect(Collectors.toList()); + if (converted.equals(declared)) { + return remap; + } + List 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))); } /** diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java index 8a434e998..c3c400187 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java +++ b/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 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 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 remap = - IntStream.range(0, groupingFieldCountWithDuplicates) - .mapToObj(i -> i) - .collect(Collectors.toCollection(ArrayList::new)); + final List 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. + * + *

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 calciteGroupingOrder(List groupings) { + List 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 references = bitSet.asList().stream() diff --git a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java index c198b1f27..6a534e733 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -10,7 +10,15 @@ import io.substrait.relation.Rel; import io.substrait.type.Type; import java.util.List; +import java.util.Optional; +import org.apache.calcite.rel.RelCollations; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelRoot; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; import org.junit.jupiter.api.Test; class ComplexAggregateTest extends PlanTestBase { @@ -214,6 +222,130 @@ void handleOutOfOrderGroupingArguments() { validateAggregateTransformation(rel, expectedFinal); } + @Test + void outOfOrderGroupingSetsHaveCorrectCalciteType() { + // Each grouping set holds one field and is trivially in order, but the aggregate declares + // field 2 before field 0, while Calcite emits its grouping columns in ascending field order. + Rel rel = + sb.aggregate( + input -> List.of(sb.grouping(input, 2), sb.grouping(input, 0)), + input -> List.of(), + Optional.of(Rel.Remap.of(List.of(0, 1))), + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING))); + + RelNode relNode = substraitToCalcite.convert(rel); + + assertRowMatch(relNode.getRowType(), N.STRING, N.I64); + } + + @Test + void groupingFieldSharedBySetsStaysOneColumn() { + // Field 2 is grouped on twice. It is one column of the aggregate's output, so it has to stay + // one column of the project the conversion puts underneath it. + Rel rel = + sb.aggregate( + input -> List.of(sb.grouping(input, 2, 0), sb.grouping(input, 2)), + input -> List.of(), + Optional.of(Rel.Remap.of(List.of(0, 1))), + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING))); + + RelNode relNode = substraitToCalcite.convert(rel); + + assertRowMatch(relNode.getRowType(), R.STRING, N.I64); + } + + @Test + void aReferenceOverOutOfOrderGroupingSetsReachesTheColumnItNames() { + Rel aggregate = + sb.aggregate( + input -> List.of(sb.grouping(input, 2), sb.grouping(input, 0)), + input -> List.of(), + Optional.empty(), + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING))); + // Field 0 of the aggregate is the field it groups on first, the string. + Rel project = + io.substrait.relation.Project.builder() + .input(aggregate) + .remap(Rel.Remap.offset(3, 1)) + .addExpressions(sb.fieldReference(aggregate, 0)) + .build(); + + RelNode relNode = substraitToCalcite.convert(project); + + assertRowMatch(relNode.getRowType(), N.STRING); + } + + @Test + void anAggregateOverOutOfOrderGroupingSetsRoundTrips() { + // The grouping columns survive the trip in the order the aggregate declares them, rather than + // in the order Calcite happens to emit them. Only those columns are compared: the grouping-set + // index comes back as an i64, because the conversion builds Calcite's GROUP_ID call as a + // BIGINT and Calcite folds it to a literal of that type, which is a separate difference. + Rel aggregate = + sb.aggregate( + input -> List.of(sb.grouping(input, 2), sb.grouping(input, 0)), + input -> List.of(), + Optional.empty(), + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING))); + + RelNode relNode = substraitToCalcite.convert(aggregate); + Rel converted = + SubstraitRelVisitor.convert( + RelRoot.of(relNode, org.apache.calcite.sql.SqlKind.SELECT), converterProvider) + .getInput(); + + List declared = aggregate.getRecordType().fields(); + List roundTripped = converted.getRecordType().fields(); + assertEquals(declared.size(), roundTripped.size()); + assertEquals(declared.subList(0, 2), roundTripped.subList(0, 2)); + } + + @Test + void anExplicitGroupIdCallKeepsTheDeclaredColumnOrder() { + // Calcite folds GROUP_ID() into a literal wherever it can work out the answer, so a plan that + // still carries the call has to be built rather than parsed. Its grouping sets mention field 3 + // before field 2, which is the order the converted relation has to declare its columns in -- + // the shape a query whose grouping sets are followed by another key produces. + org.apache.calcite.tools.RelBuilder relBuilder = + new RelCreator(TPCH_CATALOG).createRelBuilder(); + RelNode scan = relBuilder.scan("LINEITEM").build(); + AggregateCall groupId = + AggregateCall.create( + SqlStdOperatorTable.GROUP_ID, + false, + false, + false, + List.of(), + List.of(), + -1, + null, + RelCollations.EMPTY, + typeFactory.createSqlType(SqlTypeName.BIGINT), + null); + RelNode calciteAggregate = + LogicalAggregate.create( + scan, + List.of(), + ImmutableBitSet.of(0, 1, 2, 3), + List.of(ImmutableBitSet.of(0, 1, 3), ImmutableBitSet.of(2, 3)), + List.of(groupId)); + + Rel rel = + SubstraitRelVisitor.convert( + RelRoot.of(calciteAggregate, org.apache.calcite.sql.SqlKind.SELECT), + converterProvider) + .getInput(); + + // What the relation says it emits is what the Calcite aggregate it came from emits. The + // grouping-set index is left out of the comparison: Calcite types its GROUP_ID column BIGINT + // while Substrait gives the aggregate an i32 one, which is a difference of its own. + List emitted = rel.getRecordType().fields(); + assertEquals(5, emitted.size()); + assertRowMatch( + typeFactory.createStructType(calciteAggregate.getRowType().getFieldList().subList(0, 4)), + emitted.subList(0, 4)); + } + @Test void outOfOrderGroupingKeysHaveCorrectCalciteType() { Rel rel = From 462c43ca71cd816f5b6aff4508b220255c7897fb Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Wed, 26 Aug 2026 16:13:54 +0300 Subject: [PATCH 2/4] fix(isthmus): place a grouping column that is not a field reference The translation from declared to emitted grouping order gave up when a grouping expression was not a field reference into the aggregate's input, on the grounds that anything else is projected below the aggregate in declared order. That holds for what transformToValidCalciteAggregate rewrites, but not for an outer reference: it passes the validator, is left alone, and Calcite projects it itself, after the input's own fields. A plan that groups on one before a field of its input then read the wrong column. Order the columns by where they sit in the aggregate's input instead -- a field reference where its field is, anything else after them all, in declared order, which a stable sort keeps. --- .../isthmus/SubstraitRelNodeConverter.java | 16 +++++---- .../isthmus/ComplexAggregateTest.java | 36 +++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index c525b5d56..6f0491150 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -440,14 +440,18 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti private static Optional inConvertedGroupingOrder( Optional remap, List groupExprs, int callCount) { List declared = new ArrayList<>(new LinkedHashSet<>(groupExprs)); - if (!declared.stream().allMatch(RexInputRef.class::isInstance)) { - // The conversion projects expressions that are not field references below the aggregate, in - // the order they were declared, and groups over that projection, so the orders agree. - return remap; - } + // 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 converted = declared.stream() - .sorted(Comparator.comparingInt(expr -> ((RexInputRef) expr).getIndex())) + .sorted( + Comparator.comparingInt( + expr -> + expr instanceof RexInputRef + ? ((RexInputRef) expr).getIndex() + : Integer.MAX_VALUE)) .collect(Collectors.toList()); if (converted.equals(declared)) { return remap; diff --git a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java index 6a534e733..78f075023 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -4,6 +4,7 @@ import io.substrait.expression.AggregateFunctionInvocation; import io.substrait.expression.Expression; +import io.substrait.expression.FieldReference; import io.substrait.expression.ImmutableAggregateFunctionInvocation; import io.substrait.relation.Aggregate; import io.substrait.relation.NamedScan; @@ -356,4 +357,39 @@ void outOfOrderGroupingKeysHaveCorrectCalciteType() { RelNode relNode = substraitToCalcite.convert(rel); assertRowMatch(relNode.getRowType(), R.STRING, R.I64); } + + /** + * A grouping expression that is not a field reference into the aggregate's input -- an outer + * reference, which the pre-Calcite transform leaves alone rather than projecting out -- is put by + * Calcite in a projection after the input's own fields, so it is emitted last however early the + * aggregate declares it. The emit mapping has to follow it there. + */ + @Test + void outOfOrderGroupingSetsOverAnOuterReference() { + Rel outer = sb.namedScan(List.of("bar"), List.of("x"), List.of(R.I64)).withRelAnchor(1); + Rel inner = + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING)); + + Aggregate aggregate = + Aggregate.builder() + .input(inner) + .addGroupings( + Aggregate.Grouping.builder() + .addExpressions( + FieldReference.newRootStructOuterReferenceByRelReference(0, R.I64, 1)) + .build()) + .addGroupings( + Aggregate.Grouping.builder().addExpressions(sb.fieldReference(inner, 2)).build()) + // The first grouping column the aggregate declares, which is the outer reference. + .remap(Rel.Remap.of(List.of(0))) + .build(); + + Rel root = + sb.project( + input -> List.of(sb.scalarSubquery(aggregate, N.I64)), Rel.Remap.of(List.of(1)), outer); + + RelNode relNode = substraitToCalcite.convert(root); + + assertRowMatch(relNode.getRowType(), N.I64); + } } From 6bb070ee5f499ccf3615cf064d5173c9be6e18b0 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Thu, 27 Aug 2026 18:33:43 +0300 Subject: [PATCH 3/4] fix(isthmus): put the grouping-set index of an aggregate on its own column The mapping that keeps the index replaced it with aggregateCalls.size() - 1, an index into the aggregate calls rather than into the aggregate's output, so the column that came back was a copy of a grouping column. The index the relation declares for that column counted every mention of a grouping expression, so a field grouped on by several sets shifted it past the end: an aggregate with such a field and a mapping that keeps the index threw ArrayIndexOutOfBoundsException. Both counts are now over the distinct grouping columns, which is what the record type holds. --- .../isthmus/SubstraitRelNodeConverter.java | 9 +++- .../isthmus/ComplexAggregateTest.java | 48 +++++++++++++++++++ .../io/substrait/isthmus/OutputNamesTest.java | 9 ++-- 3 files changed, 60 insertions(+), 6 deletions(-) diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 6f0491150..67d159492 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -365,7 +365,10 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti .collect(java.util.stream.Collectors.toList()); Optional 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(); @@ -385,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 remapList = new LinkedList<>(remap.get().indices()); for (int i = 0; i < remapList.size(); i++) { diff --git a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java index 78f075023..a57f01208 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -12,6 +12,7 @@ import io.substrait.type.Type; import java.util.List; import java.util.Optional; +import org.apache.calcite.plan.RelOptUtil; import org.apache.calcite.rel.RelCollations; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelRoot; @@ -255,6 +256,53 @@ void groupingFieldSharedBySetsStaysOneColumn() { assertRowMatch(relNode.getRowType(), R.STRING, N.I64); } + /** + * A relation that keeps its grouping-set index maps it to the column the conversion adds for it, + * which sits after the grouping columns and the measures. Calcite folds the {@code GROUP_ID} call + * into a literal, so that is what the column holds -- which value it holds is a separate question + * from which column it is. + */ + @Test + void theGroupingSetIndexIsTheColumnTheConversionAddedForIt() { + Rel aggregate = + sb.aggregate( + input -> List.of(sb.grouping(input, 2), sb.grouping(input, 0)), + input -> List.of(sb.count(input, 0)), + Optional.of(Rel.Remap.of(List.of(0, 1, 2, 3))), + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING))); + + RelNode relNode = substraitToCalcite.convert(aggregate); + + assertEquals( + "LogicalProject(c=[$1], a=[$0], $f2=[$2], $f3=[0:BIGINT])\n" + + " LogicalAggregate(group=[{0, 2}], groups=[[{0}, {2}]], agg#0=[COUNT($0)])\n" + + " LogicalTableScan(table=[[foo]])\n", + RelOptUtil.toString(relNode)); + } + + /** + * Field 0 is grouped on by both sets and is one column of the output, so the grouping-set index + * is the fourth column and not the fifth. Counting every mention of a grouping expression put it + * past the end, and the mapping then kept an index the converted aggregate did not have. + */ + @Test + void aGroupingFieldSharedBySetsLeavesTheGroupingSetIndexWhereItIs() { + Rel aggregate = + sb.aggregate( + input -> List.of(sb.grouping(input, 0, 2), sb.grouping(input, 0)), + input -> List.of(sb.count(input, 0)), + Optional.of(Rel.Remap.of(List.of(0, 1, 2, 3))), + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING))); + + RelNode relNode = substraitToCalcite.convert(aggregate); + + assertEquals( + "LogicalProject(a=[$0], c=[$1], $f2=[$2], $f3=[0:BIGINT])\n" + + " LogicalAggregate(group=[{0, 2}], groups=[[{0, 2}, {0}]], agg#0=[COUNT($0)])\n" + + " LogicalTableScan(table=[[foo]])\n", + RelOptUtil.toString(relNode)); + } + @Test void aReferenceOverOutOfOrderGroupingSetsReachesTheColumnItNames() { Rel aggregate = diff --git a/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java b/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java index 1c4fb0028..d8b78460d 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java @@ -169,10 +169,11 @@ private Rel twoColumnProject() { @Test void leavesAnAggregateThatEmitsDirectlyAlone() { - // The conversion of an aggregate over several grouping sets ends in a projection, but that - // projection carries the grouping-set index rather than this relation's emit mapping, and the - // columns underneath it are ordered by Calcite's group key rather than by the relation's own - // record type. Names are dropped rather than pinned onto columns chosen by something else. + // The conversion of an aggregate over several grouping sets ends in a projection that carries + // the grouping-set index. Its other columns are the relation's own, in the declared order, but + // that one comes back as Calcite's folded GROUP_ID literal -- a BIGINT where the relation + // declares an i32 -- so the names are dropped rather than pinned onto a column whose type the + // plan does not describe. Rel aggregate = sb.aggregate( input -> List.of(sb.grouping(input, 0), sb.grouping(input, 1)), From 227ae0dfc8e3ec9af16c4177c8d0b67de5b41680 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Fri, 28 Aug 2026 11:28:17 +0300 Subject: [PATCH 4/4] test(isthmus): assert the emit mapping rather than the column types Three of the four grouping columns in this fixture are BIGINT, so comparing types cannot show a permutation among them. The mapping is what carries the declared order, so it is asserted directly. --- .../java/io/substrait/isthmus/ComplexAggregateTest.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java index a57f01208..988618ace 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -385,6 +385,14 @@ void anExplicitGroupIdCallKeepsTheDeclaredColumnOrder() { converterProvider) .getInput(); + // The mapping is what carries the difference, and it is asserted directly: the sets mention + // fields 0, 1 and 3 before 2, so the relation declares them in that order, while the aggregate + // underneath emits them by field index. Types alone would not show it -- three of these four + // columns are BIGINT. + assertEquals( + Optional.of(Rel.Remap.of(List.of(0, 1, 3, 2, 4))), + ((io.substrait.relation.Aggregate) rel).getRemap()); + // What the relation says it emits is what the Calcite aggregate it came from emits. The // grouping-set index is left out of the comparison: Calcite types its GROUP_ID column BIGINT // while Substrait gives the aggregate an i32 one, which is a difference of its own.