diff --git a/isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java b/isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java index f8695db64..61e7c176e 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java +++ b/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. + * + *

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

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 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 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 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 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: - * - *

+ * 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 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(); } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 0138300ad..c404d9c86 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/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 = aggregate.getRemap(); - final int lastFieldIndex = groupExprs.size() + aggregateCalls.size(); + final Optional 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 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. + * + *

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

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))); } /** @@ -1191,10 +1246,10 @@ protected RelNode applyRelCommon(RelNode relNode, Rel rel, RelNode... inputs) { * inputs are compared against. * *

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

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 diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java index 8a434e998..853be3c9e 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; @@ -392,9 +393,7 @@ public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) { } // Number of distinct grouping-expression output fields produced by the aggregate. - // Used by the no-GROUP_ID remap and the LITERAL_AGG project wrapper below. - // The GROUP_ID remap branch intentionally uses a non-distinct count instead (see comment - // there). + // Used by both remap branches and by the LITERAL_AGG project wrapper below. final int groupingFieldCount = Math.toIntExact( groupings.stream().flatMap(g -> g.getExpressions().stream()).distinct().count()); @@ -414,32 +413,35 @@ public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) { Aggregate.builder().input(input).addAllGroupings(groupings).addAllMeasures(aggCalls); if (groupings.size() > 1) { + // substrait-java declares the grouping columns of an aggregate as the distinct grouping + // expressions in the order they first appear across its grouping sets -- the reconstruction + // it puts in place of the shared grouping-expression list the spec orders them by, which the + // POJO cannot hold -- 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-java takes the grouping columns to be the distinct grouping expressions in the + * order they first appear across the grouping sets, reconstructing the shared list the spec + * orders them by; 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..6cdbd1f95 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -4,13 +4,26 @@ 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.Project; 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.SqlKind; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.tools.RelBuilder; +import org.apache.calcite.util.ImmutableBitSet; import org.junit.jupiter.api.Test; class ComplexAggregateTest extends PlanTestBase { @@ -214,6 +227,314 @@ 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 = + 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, 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. + 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, 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))), ((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)); + } + + /** + * Every other fixture here swaps two columns, and a transposition is its own inverse, so a + * mapping replaced by the one that undoes it would go unnoticed. These two sets mention fields 0 + * and 3 before 1 and 2, which makes the mapping a three-cycle: the relation declares (a, d, b, c) + * where the aggregate underneath emits (a, b, c, d), and the inverse would declare (a, c, d, b). + */ + @Test + void anAggregateOverGroupingSetsInANonSwapOrderKeepsTheDeclaredOrder() { + Rel aggregate = + sb.aggregate( + input -> List.of(sb.grouping(input, 0, 3), sb.grouping(input, 1, 2)), + input -> List.of(), + Optional.empty(), + sb.namedScan( + List.of("foo"), + List.of("a", "b", "c", "d"), + List.of(R.I64, R.STRING, R.FP64, R.BOOLEAN))); + + RelNode relNode = substraitToCalcite.convert(aggregate); + + assertEquals( + "LogicalProject(a=[$0], d=[$3], b=[$1], c=[$2], $f4=[0:BIGINT])\n" + + " LogicalAggregate(group=[{0, 1, 2, 3}], groups=[[{0, 3}, {1, 2}]])\n" + + " LogicalTableScan(table=[[foo]])\n", + RelOptUtil.toString(relNode)); + } + + /** The same shape in the other direction, asserted on the mapping the conversion produces. */ + @Test + void groupingSetsInANonSwapOrderGiveAMappingThatIsNotItsOwnInverse() { + RelBuilder relBuilder = new RelCreator(TPCH_CATALOG).createRelBuilder(); + RelNode scan = relBuilder.scan("LINEITEM").build(); + RelNode calciteAggregate = + LogicalAggregate.create( + scan, + List.of(), + ImmutableBitSet.of(0, 1, 2, 3), + List.of(ImmutableBitSet.of(0, 3), ImmutableBitSet.of(1, 2)), + List.of()); + + Rel rel = + SubstraitRelVisitor.convert(RelRoot.of(calciteAggregate, SqlKind.SELECT), converterProvider) + .getInput(); + + assertEquals(Optional.of(Rel.Remap.of(List.of(0, 2, 3, 1))), ((Aggregate) rel).getRemap()); + } + + /** + * A lone grouping set gives every mention of an expression a column of its own -- {@code + * Aggregate.deriveRecordType} dedups the grouping expressions only across several sets -- while + * Calcite's grouping bit set cannot hold a field twice. So a repeated field reaches Calcite as + * two columns of the projection the conversion puts underneath the aggregate. + */ + @Test + void aGroupingSetNamingAFieldTwiceKeepsAColumnPerMention() { + Rel scan = + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING)); + + RelNode relNode = + substraitToCalcite.convert( + sb.aggregate(input -> sb.grouping(input, 2, 0, 2), input -> List.of(), scan)); + + assertRowMatch(relNode.getRowType(), R.STRING, R.I64, R.STRING); + } + + /** The same repeat under an emit mapping, whose indices count the columns the aggregate holds. */ + @Test + void aGroupingSetNamingAFieldTwiceUnderAnEmitMappingKeepsAColumnPerMention() { + Rel scan = + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING)); + + RelNode relNode = + substraitToCalcite.convert( + sb.aggregate( + input -> List.of(sb.grouping(input, 2, 0, 2)), + input -> List.of(), + Optional.of(Rel.Remap.of(List.of(2, 0))), + scan)); + + assertRowMatch(relNode.getRowType(), R.STRING, R.STRING); + } + + /** A repeat the grouping fields are in ascending order for reaches Calcite the same way. */ + @Test + void anAscendingGroupingSetNamingAFieldTwiceKeepsAColumnPerMention() { + Rel scan = + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING)); + + RelNode relNode = + substraitToCalcite.convert( + sb.aggregate(input -> sb.grouping(input, 0, 2, 2), input -> List.of(), scan)); + + assertRowMatch(relNode.getRowType(), R.I64, R.STRING, R.STRING); + } + + /** And so does a repeat of an expression the transformer has to project out anyway. */ + @Test + void aGroupingSetNamingAnExpressionTwiceKeepsAColumnPerMention() { + Rel scan = + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING)); + + RelNode relNode = + substraitToCalcite.convert( + sb.aggregate( + input -> + sb.grouping( + sb.add(sb.fieldReference(input, 0), sb.i64(42)), + sb.add(sb.fieldReference(input, 0), sb.i64(42))), + input -> List.of(), + scan)); + + assertRowMatch(relNode.getRowType(), R.I64, R.I64); + } + + /** + * The order the fields are grouped in is carried by the emit mapping, so it is no longer a reason + * to rewrite the input: the aggregate reads the relation it was given, and the declared order is + * a projection above it rather than below. + */ + @Test + void outOfOrderGroupingKeysLeaveTheInputAlone() { + Rel scan = + sb.namedScan(List.of("foo"), List.of("a", "b", "c"), List.of(R.I64, R.I64, R.STRING)); + + RelNode relNode = + substraitToCalcite.convert( + sb.aggregate(input -> sb.grouping(input, 2, 0), input -> List.of(), scan)); + + assertEquals( + "LogicalProject(c=[$1], a=[$0])\n" + + " LogicalAggregate(group=[{0, 2}])\n" + + " LogicalTableScan(table=[[foo]])\n", + RelOptUtil.toString(relNode)); + } + @Test void outOfOrderGroupingKeysHaveCorrectCalciteType() { Rel rel = @@ -224,4 +545,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..a8f380022 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)), @@ -266,11 +267,35 @@ void leavesTheNamesOfAnotherRelationAloneWhenItsOwnOperatorIsElided() { assertEquals(List.of("inner"), substraitToCalcite.convert(filter).getRowType().getFieldNames()); } + @Test + void namesAnAggregateWhoseGroupingColumnsCalciteOrdersDifferently() { + // The grouping sets first mention field 1 and then field 0, so the relation declares its + // grouping columns as (b, a) where the aggregate underneath emits (a, b). The emit mapping the + // conversion adds puts them back in the declared order, which is what lets the names be bound + // by position at all. The mapping drops the grouping-set index, so every remaining column is + // one the relation declares. + Rel scan3 = + sb.namedScan(List.of("t3"), List.of("a", "b", "c"), List.of(R.I64, N.STRING, R.FP64)); + Rel aggregate = + sb.aggregate( + input -> List.of(sb.grouping(input, 1), sb.grouping(input, 0)), + input -> List.of(sb.count(input, 0)), + Optional.of(Rel.Remap.of(List.of(0, 1, 2))), + scan3); + + RelNode named = + substraitToCalcite.convert( + aggregate.withHint( + Optional.of(Hint.builder().addOutputNames("k_b", "k_a", "n").build()))); + + assertEquals(List.of("k_b", "k_a", "n"), named.getRowType().getFieldNames()); + } + @Test void dropsNamesWhereTheColumnsAreNotTheRelationsColumns() { - // An aggregate over several grouping sets orders its grouping columns by first appearance, - // where Calcite orders them by group key: the record type reads (b, a, count, index) and the - // converted node (a, b, count, ...), so the names would land on columns the plan does not name. + // Same aggregate with the grouping-set index emitted: the relation types it i32 where the + // GROUP_ID call the conversion appends is i64, so the fourth column is not the fourth column + // the relation declares and the names would land on a column the plan does not name. Rel scan3 = sb.namedScan(List.of("t3"), List.of("a", "b", "c"), List.of(R.I64, N.STRING, R.FP64)); Rel aggregate =