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..67d159492 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; @@ -363,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(); @@ -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 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. + * + *

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)); + // 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 -> + expr instanceof RexInputRef + ? ((RexInputRef) expr).getIndex() + : Integer.MAX_VALUE)) + .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..988618ace 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -4,13 +4,23 @@ 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; import io.substrait.relation.Rel; 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; +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 +224,185 @@ 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); + } + + /** + * 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 = + 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(); + + // 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. + 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 = @@ -224,4 +413,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); + } } 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)),