From 1bd5e351d429351c5db3e9f09dd8be30e90ff21d Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 7 Sep 2026 16:30:45 +0300 Subject: [PATCH] fix(isthmus)!: preserve grouping-set indices through Calcite conversion Substrait defines the grouping-set index as the zero-based position in the declared set list, with type i32 ([spec v0.102.0](https://github.com/substrait-io/substrait/blob/v0.102.0/site/docs/relations/logical_relations.md#aggregate-operation)). Calcite's GROUP_ID returns zero for distinct sets, and its normalized set order can change which set an index identifies after a round trip. Build the index from grouping-key membership and the occurrence of duplicate sets. For declared sets (b), (a), (b), the corresponding rows retain indices 0, 1, and 2. Calcite may keep its normalized set order; a projection preserves the index visible to parent relations. Translate SQL GROUPING/GROUPING_ID through Substrait's implicit index so these values also survive conversion back to Substrait. Calcite 1.42.0 still has an execution limitation with mixed empty and non-empty grouping sets on empty input: a row for an empty set can be lost. The same failure occurs with a Calcite-only plan. Closes #1182 Closes #1209 BREAKING CHANGE: The grouping-set index produced by Substrait-to-Calcite conversion now has INTEGER (i32) type instead of BIGINT. Callers that assume a BIGINT output column must use the i32 schema declared by Substrait. --- .../isthmus/SubstraitRelNodeConverter.java | 137 ++++++++-- .../isthmus/SubstraitRelVisitor.java | 174 +++++-------- .../isthmus/ComplexAggregateTest.java | 41 ++- .../isthmus/GroupingSetIndexTest.java | 246 ++++++++++++++++++ .../io/substrait/isthmus/OutputNamesTest.java | 20 +- .../SubstraitRelNodeConverterTest.java | 27 +- 6 files changed, 485 insertions(+), 160 deletions(-) create mode 100644 isthmus/src/test/java/io/substrait/isthmus/GroupingSetIndexTest.java diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 72d319364..78fe3ac9c 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -373,21 +373,8 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti final boolean emitDirect = remap.isEmpty(); final boolean groupingSetIndexGetsRemapped = remap.map(r -> r.indices().contains(groupingSetIndex)).orElse(false); - if (aggregate.getGroupings().size() > 1 && (emitDirect || groupingSetIndexGetsRemapped)) { - aggregateCalls.add( - AggregateCall.create( - SqlStdOperatorTable.GROUP_ID, - false, - false, - false, - Collections.emptyList(), - Collections.emptyList(), - -1, - null, - RelCollations.EMPTY, - typeConverter.toCalcite(typeFactory, TypeCreator.REQUIRED.I64), - null)); - } + final boolean needsGroupingSetIndex = + aggregate.getGroupings().size() > 1 && (emitDirect || groupingSetIndexGetsRemapped); exitUncorrelatedScope(context, Aggregate.class); @@ -401,15 +388,131 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti ? relBuilder.transform(config -> config.withDedupAggregateCalls(false)) : relBuilder; - RelNode node = aggregateBuilder.push(child).aggregate(groupKey, aggregateCalls).build(); + aggregateBuilder.push(child); + RelNode node = + needsGroupingSetIndex + ? aggregateWithGroupingSetIndex( + aggregateBuilder, groupExprLists, aggregateCalls, groupingSetIndex) + : aggregateBuilder.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, inConvertedGroupingOrder(remap, groupExprs, aggregateCalls.size())), + applyRemap( + node, + inConvertedGroupingOrder( + remap, groupExprs, aggregateCalls.size() + (needsGroupingSetIndex ? 1 : 0))), aggregate, child); } + /** + * Materializes the declared grouping-set ordinal. GROUPING distinguishes membership even when a + * grouped value is null; GROUP_ID distinguishes repetitions of the same set. Calcite may sort the + * sets and expand duplicates into UNION ALL branches, but these two values retain their meaning + * through those rewrites. + */ + private RelNode aggregateWithGroupingSetIndex( + RelBuilder builder, + List> groupExprLists, + List aggregateCalls, + int groupingSetIndex) { + List> sets = + groupExprLists.stream().map(HashSet::new).collect(Collectors.toList()); + // Use one indicator per varying key, rather than a single bit mask with a fixed width. + List indicators = + groupExprLists.stream() + .flatMap(Collection::stream) + .distinct() + .filter(key -> sets.stream().anyMatch(set -> !set.contains(key))) + .collect(Collectors.toList()); + boolean duplicates = new HashSet<>(sets).size() != sets.size(); + // Materialize non-field keys before constructing typed calls. RelBuilder's expression-based + // AggCall API re-infers stored measure types and uses placeholder types during duplicate-set + // expansion; the AggregateCall API preserves the types from fromMeasure. + List extraKeys = + groupExprLists.stream() + .flatMap(Collection::stream) + .distinct() + .filter(key -> !(key instanceof RexInputRef)) + .collect(Collectors.toList()); + int inputFieldCount = builder.peek().getRowType().getFieldCount(); + if (!extraKeys.isEmpty()) { + builder.projectPlus(extraKeys); + } + Map keyFields = new HashMap<>(); + for (List keys : groupExprLists) { + for (RexNode key : keys) { + keyFields.put( + key, + key instanceof RexInputRef + ? key + : builder.field(inputFieldCount + extraKeys.indexOf(key))); + } + } + List> fieldSets = + groupExprLists.stream() + .map(keys -> keys.stream().map(keyFields::get).collect(Collectors.toList())) + .collect(Collectors.toList()); + RelBuilder.GroupKey groupKey = builder.groupKey(keyFields.values(), fieldSets); + List calls = new ArrayList<>(aggregateCalls); + for (RexNode key : indicators) { + calls.add( + groupingCall( + SqlStdOperatorTable.GROUPING, + List.of(((RexInputRef) keyFields.get(key)).getIndex()))); + } + if (duplicates) { + calls.add(groupingCall(SqlStdOperatorTable.GROUP_ID, List.of())); + } + builder.aggregate(groupKey, calls); + + List cases = new ArrayList<>(); + Map, Integer> occurrences = new HashMap<>(); + RelDataType indexType = typeConverter.toCalcite(typeFactory, TypeCreator.REQUIRED.I32); + for (int index = 0; index < sets.size() - 1; index++) { + java.util.Set set = sets.get(index); + List conditions = new ArrayList<>(); + for (int indicator = 0; indicator < indicators.size(); indicator++) { + RexNode field = builder.field(groupingSetIndex + indicator); + conditions.add( + builder.equals( + field, + rexBuilder.makeLiteral( + set.contains(indicators.get(indicator)) ? 0L : 1L, field.getType()))); + } + if (duplicates) { + int occurrence = occurrences.merge(set, 1, Integer::sum) - 1; + RexNode field = builder.field(groupingSetIndex + indicators.size()); + conditions.add( + builder.equals(field, rexBuilder.makeLiteral((long) occurrence, field.getType()))); + } + cases.add(builder.and(conditions)); + cases.add(rexBuilder.makeLiteral(index, indexType)); + } + cases.add(rexBuilder.makeLiteral(sets.size() - 1, indexType)); + List output = new ArrayList<>(); + for (int field = 0; field < groupingSetIndex; field++) { + output.add(builder.field(field)); + } + output.add(rexBuilder.makeCall(SqlStdOperatorTable.CASE, cases)); + return builder.project(output).build(); + } + + private AggregateCall groupingCall(SqlAggFunction function, List arguments) { + return AggregateCall.create( + function, + false, + false, + false, + List.of(), + arguments, + -1, + null, + RelCollations.EMPTY, + typeConverter.toCalcite(typeFactory, TypeCreator.REQUIRED.I64), + null); + } + /** * 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. diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java index 479b6fb8f..7e8b200bf 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java @@ -66,7 +66,6 @@ import org.apache.calcite.rex.RexNode; import org.apache.calcite.rex.RexUtil; import org.apache.calcite.sql.SqlKind; -import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.util.ImmutableBitSet; import org.immutables.value.Value; @@ -376,131 +375,98 @@ public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) { List groupings = sets.filter(s -> s != null).map(s -> fromGroupSet(s, input)).collect(Collectors.toList()); - // get GROUP_ID() function calls - java.util.Set groupIdCalls = - aggregate.getAggCallList().stream() - .filter(c -> c.getAggregation().equals(SqlStdOperatorTable.GROUP_ID)) - .collect(Collectors.toSet()); - - // get LITERAL_AGG() function calls — injected by SubQueryRemoveRule (CALCITE-6945) as a - // null-presence indicator; they carry a RexLiteral in rexList and have no Substrait binding. - java.util.Set literalAggCalls = - aggregate.getAggCallList().stream() - .filter(c -> c.getAggregation().getKind() == SqlKind.LITERAL_AGG) - .collect(Collectors.toSet()); - - if (!literalAggCalls.isEmpty() && groupings.size() > 1) { + List calls = aggregate.getAggCallList(); + boolean hasGrouping = + calls.stream().anyMatch(c -> c.getAggregation().getKind() == SqlKind.GROUPING); + boolean hasSpecialCalls = calls.stream().anyMatch(SubstraitRelVisitor::isGroupingOrLiteralCall); + boolean hasLiteral = + calls.stream().anyMatch(c -> c.getAggregation().getKind() == SqlKind.LITERAL_AGG); + if (hasLiteral && groupings.size() > 1) { throw new UnsupportedOperationException( "LITERAL_AGG combined with GROUPING SETS / CUBE / ROLLUP is not supported"); } - // Number of distinct grouping-expression output fields produced by the aggregate. - // 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()); - - List filteredAggCalls = - aggregate.getAggCallList().stream() - // remove GROUP_ID() and LITERAL_AGG() function calls - .filter(c -> !groupIdCalls.contains(c) && !literalAggCalls.contains(c)) - .collect(Collectors.toList()); - - List aggCalls = - filteredAggCalls.stream() + List measures = + calls.stream() + .filter(c -> !isGroupingOrLiteralCall(c)) .map(c -> fromAggCall(aggregate.getInput(), input.getRecordType(), c)) .collect(Collectors.toList()); - ImmutableAggregate.Builder builder = - Aggregate.builder().input(input).addAllGroupings(groupings).addAllMeasures(aggCalls); - + Aggregate.builder().input(input).addAllGroupings(groupings).addAllMeasures(measures); + List mapping = new ArrayList<>(calciteGroupingOrder(groupings)); + int groupingFieldCount = mapping.size(); + int groupingSetIndex = groupingFieldCount + measures.size(); + for (int call = 0; call < measures.size(); call++) { + mapping.add(groupingFieldCount + call); + } 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()) { - 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 - final int filterAggCallCount = aggCalls.size(); - final Integer groupingSetIndex = groupingFieldCount + filterAggCallCount; - - 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 + groupingFieldCount, filteredAggCalls.indexOf(aggCall) + groupingFieldCount); - } else if (groupIdCalls.contains(aggCall)) { - remap.add(i + groupingFieldCount, groupingSetIndex); - } else { - // this should never get triggered - throw new IllegalStateException( - "encountered AggregateCall that is neither in filteredAggCalls nor in groupIdCalls" - + aggCall); - } - } - - builder.remap(Remap.of(remap)); + // Keep the implicit ordinal only while deriving GROUPING values. The project below restores + // the original Calcite schema; ordinary aggregates emit just their keys and measures. + if (hasGrouping) { + mapping.add(groupingSetIndex); } + builder.remap(Remap.of(mapping)); } - Rel aggRel = builder.build(); - - if (literalAggCalls.isEmpty()) { + if (!hasSpecialCalls) { return aggRel; } - // Wrap the aggregate in a Project that replaces LITERAL_AGG output positions with their - // literal values and passes through all other fields via FieldReference. - // - // The aggregate output schema is: [grouping fields..., real agg measures...] - // The full output schema requested is: [grouping fields..., all agg calls (in original order)] - // For each position in the original agg call list: - // - real measure → FieldReference into the aggregate output - // - LITERAL_AGG → the literal value from aggCall.rexList - final int realAggCount = aggCalls.size(); - final int totalAggOutputFields = groupingFieldCount + realAggCount; - - // Build the project expression list: grouping fields first, then one expression per original - // agg call in declaration order. - List projectExprs = new ArrayList<>(); - for (int i = 0; i < groupingFieldCount; i++) { - projectExprs.add(FieldReference.newInputRelReference(i, aggRel)); + List output = new ArrayList<>(); + for (int field = 0; field < groupingFieldCount; field++) { + output.add(FieldReference.newInputRelReference(field, aggRel)); } - int realAggIndex = groupingFieldCount; // tracks next real-measure field index in aggRel output - for (AggregateCall aggCall : aggregate.getAggCallList()) { - if (literalAggCalls.contains(aggCall)) { - // Convert the RexLiteral stored in rexList to a Substrait literal expression - RexNode rexLiteral = Iterables.getOnlyElement(aggCall.rexList); - projectExprs.add(toExpression(rexLiteral)); - } else if (!groupIdCalls.contains(aggCall)) { - // real measure: pass through by reference - projectExprs.add(FieldReference.newInputRelReference(realAggIndex, aggRel)); - realAggIndex++; + int measureIndex = groupingFieldCount; + for (AggregateCall call : calls) { + switch (call.getAggregation().getKind()) { + case GROUPING: + output.add(groupingValue(call, aggregate.getGroupSets(), aggRel, groupingSetIndex)); + break; + case GROUP_ID: + // A Calcite Aggregate holds distinct, sorted grouping sets. RelBuilder expands any + // repetitions into UNION ALL branches before they reach this visitor. + output.add(ExpressionCreator.i64(false, 0)); + break; + case LITERAL_AGG: + output.add(toExpression(Iterables.getOnlyElement(call.rexList))); + break; + default: + output.add(FieldReference.newInputRelReference(measureIndex++, aggRel)); } - // GROUP_ID calls are not present in the outer schema here (groupings.size() <= 1 branch); - // if groupings.size() > 1 they are handled by the remap above and should not appear here } - return Project.builder() - .remap(Remap.offset(totalAggOutputFields, projectExprs.size())) - .expressions(projectExprs) .input(aggRel) + .expressions(output) + .remap(Remap.offset(aggRel.getRecordType().fields().size(), output.size())) .build(); } + private static boolean isGroupingOrLiteralCall(AggregateCall call) { + SqlKind kind = call.getAggregation().getKind(); + return kind == SqlKind.GROUPING || kind == SqlKind.GROUP_ID || kind == SqlKind.LITERAL_AGG; + } + + /** Computes SQL's membership bit mask from Substrait's declared grouping-set ordinal. */ + private static Expression groupingValue( + AggregateCall call, List sets, Rel aggregate, int groupingSetIndex) { + List clauses = new ArrayList<>(); + Expression value = ExpressionCreator.i64(false, 0); + for (int index = 0; index < sets.size(); index++) { + long mask = 0; + for (int argument : call.getArgList()) { + mask = (mask << 1) | (sets.get(index).get(argument) ? 0 : 1); + } + value = ExpressionCreator.i64(false, mask); + if (index < sets.size() - 1) { + clauses.add(ExpressionCreator.switchClause(ExpressionCreator.i32(false, index), value)); + } + } + return sets.size() <= 1 + ? value + : ExpressionCreator.switchStatement( + FieldReference.newInputRelReference(groupingSetIndex, aggregate), value, clauses); + } + /** * Returns, for each grouping column of the converted Calcite aggregate, the position that column * holds in the output the Substrait aggregate declares. diff --git a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java index 6cdbd1f95..eeefecb8c 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ComplexAggregateTest.java @@ -1,6 +1,7 @@ package io.substrait.isthmus; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import io.substrait.expression.AggregateFunctionInvocation; import io.substrait.expression.Expression; @@ -261,9 +262,8 @@ void groupingFieldSharedBySetsStaysOneColumn() { /** * 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. + * which sits after the grouping columns and the measures. Its value identifies the declared set, + * independently of Calcite's ordering of grouping sets. */ @Test void theGroupingSetIndexIsTheColumnTheConversionAddedForIt() { @@ -277,8 +277,8 @@ void theGroupingSetIndexIsTheColumnTheConversionAddedForIt() { 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" + "LogicalProject(c=[$1], a=[$0], $f2=[$2], $f3=[CASE(AND(=($3, 0), =($4, 1)), 0, 1)])\n" + + " LogicalAggregate(group=[{0, 2}], groups=[[{0}, {2}]], agg#0=[COUNT($0)], agg#1=[GROUPING($2)], agg#2=[GROUPING($0)])\n" + " LogicalTableScan(table=[[foo]])\n", RelOptUtil.toString(relNode)); } @@ -300,8 +300,8 @@ void aGroupingFieldSharedBySetsLeavesTheGroupingSetIndexWhereItIs() { 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" + "LogicalProject(a=[$0], c=[$1], $f2=[$2], $f3=[CASE(=($3, 0), 0, 1)])\n" + + " LogicalAggregate(group=[{0, 2}], groups=[[{0, 2}, {0}]], agg#0=[COUNT($0)], agg#1=[GROUPING($2)])\n" + " LogicalTableScan(table=[[foo]])\n", RelOptUtil.toString(relNode)); } @@ -329,10 +329,7 @@ void aReferenceOverOutOfOrderGroupingSetsReachesTheColumnItNames() { @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. + // Both the declared column order and the i32 grouping-set index survive the conversion. Rel aggregate = sb.aggregate( input -> List.of(sb.grouping(input, 2), sb.grouping(input, 0)), @@ -348,7 +345,7 @@ void anAggregateOverOutOfOrderGroupingSetsRoundTrips() { List declared = aggregate.getRecordType().fields(); List roundTripped = converted.getRecordType().fields(); assertEquals(declared.size(), roundTripped.size()); - assertEquals(declared.subList(0, 2), roundTripped.subList(0, 2)); + assertEquals(declared, roundTripped); } @Test @@ -388,16 +385,12 @@ void anExplicitGroupIdCallKeepsTheDeclaredColumnOrder() { // 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)); + Project project = assertInstanceOf(Project.class, rel); + Aggregate converted = assertInstanceOf(Aggregate.class, project.getInput()); + assertEquals(Optional.of(Rel.Remap.of(List.of(0, 1, 3, 2))), converted.getRemap()); + assertEquals( + 0, assertInstanceOf(Expression.I64Literal.class, project.getExpressions().get(4)).value()); + assertRowMatch(calciteAggregate.getRowType(), rel.getRecordType().fields()); } /** @@ -421,8 +414,8 @@ void anAggregateOverGroupingSetsInANonSwapOrderKeepsTheDeclaredOrder() { 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" + "LogicalProject(a=[$0], d=[$3], b=[$1], c=[$2], $f4=[CASE(AND(=($4, 0), =($5, 0), =($6, 1), =($7, 1)), 0, 1)])\n" + + " LogicalAggregate(group=[{0, 1, 2, 3}], groups=[[{0, 3}, {1, 2}]], agg#0=[GROUPING($0)], agg#1=[GROUPING($3)], agg#2=[GROUPING($1)], agg#3=[GROUPING($2)])\n" + " LogicalTableScan(table=[[foo]])\n", RelOptUtil.toString(relNode)); } diff --git a/isthmus/src/test/java/io/substrait/isthmus/GroupingSetIndexTest.java b/isthmus/src/test/java/io/substrait/isthmus/GroupingSetIndexTest.java new file mode 100644 index 000000000..8f8a2a5b3 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/GroupingSetIndexTest.java @@ -0,0 +1,246 @@ +package io.substrait.isthmus; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import io.substrait.expression.ExpressionCreator; +import io.substrait.relation.Rel; +import io.substrait.type.NamedStruct; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelRoot; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.tools.RelRunners; +import org.apache.calcite.util.ImmutableBitSet; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +class GroupingSetIndexTest extends PlanTestBase { + + @Test + void indexFollowsDeclaredSetOrder() throws SQLException { + Rel aggregate = aggregate(List.of(List.of(1), List.of(0), List.of()), input()); + assertRowsAndRoundTrip( + aggregate, + List.of( + row(10, null, 2L, 0), + row(null, 1, 1L, 1), + row(null, 2, 1L, 1), + row(null, null, 2L, 2))); + } + + @Test + void nestedGroupingSetsWithoutMeasuresKeepTheirIndices() throws SQLException { + Rel aggregate = + sb.aggregate( + rel -> List.of(sb.grouping(rel, 0), sb.grouping(rel, 0, 1)), + rel -> List.of(), + Optional.empty(), + input()); + assertRowsAndRoundTrip( + aggregate, List.of(row(1, null, 0), row(2, null, 0), row(1, 10, 1), row(2, 10, 1))); + } + + @Test + void computedGroupingKeysKeepTheirIndices() throws SQLException { + Rel aggregate = + sb.aggregate( + rel -> + List.of( + sb.grouping(sb.add(sb.fieldReference(rel, 0), sb.i32(1))), + sb.grouping(rel, 1), + sb.grouping(sb.add(sb.fieldReference(rel, 0), sb.i32(1)))), + rel -> List.of(sb.countStar()), + Optional.empty(), + input()); + assertRowsAndRoundTrip( + aggregate, + List.of( + row(2, null, 1L, 0), + row(3, null, 1L, 0), + row(null, 10, 2L, 1), + row(2, null, 1L, 2), + row(3, null, 1L, 2))); + } + + @Test + void duplicateSetsKeepTheirOwnIndices() throws SQLException { + Rel aggregate = + aggregate(List.of(List.of(1), List.of(0), List.of(1), List.of(), List.of()), input()); + assertRowsAndRoundTrip( + aggregate, + List.of( + row(10, null, 2L, 0), + row(null, 1, 1L, 1), + row(null, 2, 1L, 1), + row(10, null, 2L, 2), + row(null, null, 2L, 3), + row(null, null, 2L, 4))); + } + + @Test + void nullKeysDoNotIdentifyGroupingSets() throws SQLException { + Rel input = + virtualTable( + NamedStruct.of(List.of("a", "b"), R.struct(N.I32, N.I32)), + List.of(ExpressionCreator.i32(true, 1), ExpressionCreator.typedNull(N.I32)), + List.of(ExpressionCreator.typedNull(N.I32), ExpressionCreator.typedNull(N.I32))); + assertRowsAndRoundTrip( + aggregate(List.of(List.of(1), List.of(0), List.of()), input), + List.of( + row(null, null, 2L, 0), + row(null, 1, 1L, 1), + row(null, null, 1L, 1), + row(null, null, 2L, 2))); + } + + @Test + void emptyInputStillProducesEachEmptySet() throws SQLException { + Rel input = virtualTable(NamedStruct.of(List.of("a", "b"), R.struct(R.I32, R.I32))); + assertRowsAndRoundTrip( + aggregate(List.of(List.of(), List.of()), input), List.of(row(0L, 0), row(0L, 1))); + } + + @Test + @Disabled("Calcite 1.42.0 drops an empty group on empty input with mixed grouping sets") + void mixedGroupingSetsOnEmptyInputKeepEveryEmptySet() throws SQLException { + Rel input = virtualTable(NamedStruct.of(List.of("a", "b"), R.struct(R.I32, R.I32))); + assertRowsAndRoundTrip( + aggregate(List.of(List.of(), List.of(0), List.of()), input), + List.of(row(null, 0L, 0), row(null, 0L, 2))); + } + + @Test + void emitCanReorderAndRepeatTheIndex() throws SQLException { + Rel aggregate = aggregate(List.of(List.of(1), List.of(0)), input()); + assertRowsAndRoundTrip( + aggregate.withRemap(Optional.of(Rel.Remap.of(List.of(3, 2, 3)))), + List.of(row(0, 2L, 0), row(1, 1L, 1), row(1, 1L, 1))); + } + + @Test + void aFilterSelectsTheDeclaredOccurrenceOfASet() throws SQLException { + Rel aggregate = aggregate(List.of(List.of(1), List.of(0), List.of(1)), input()); + assertRowsAndRoundTrip( + sb.filter(rel -> sb.equal(sb.fieldReference(rel, 3), sb.i32(2)), aggregate), + List.of(row(10, null, 2L, 2))); + } + + @Test + void omittingTheIndexKeepsDuplicateRows() throws SQLException { + Rel aggregate = aggregate(List.of(List.of(1), List.of(0), List.of(1)), input()); + assertRowsAndRoundTrip( + aggregate.withRemap(Optional.of(Rel.Remap.of(List.of(0, 1, 2)))), + List.of(row(10, null, 2L), row(10, null, 2L), row(null, 1, 1L), row(null, 2, 1L))); + } + + @Test + void indexIsNotLimitedByTheWidthOfAGroupingMask() throws SQLException { + int width = 65; + Rel input = + virtualTable( + NamedStruct.of( + IntStream.range(0, width).mapToObj(i -> "c" + i).collect(Collectors.toList()), + R.struct( + IntStream.range(0, width).mapToObj(i -> R.I32).collect(Collectors.toList()))), + IntStream.range(0, width).mapToObj(sb::i32).collect(Collectors.toList())); + Rel aggregate = + aggregate( + List.of(IntStream.range(0, width).boxed().collect(Collectors.toList()), List.of()), + input); + assertRowsAndRoundTrip( + aggregate.withRemap(Optional.of(Rel.Remap.of(List.of(width + 1)))), + List.of(row(0), row(1))); + } + + @Test + void groupingMasksPreserveArgumentOrderAndMeasurePositions() throws SQLException { + builder.push(substraitToCalcite.convert(input())); + RelNode calcite = + builder + .aggregate( + builder.groupKey( + ImmutableBitSet.of(0, 1), + List.of(ImmutableBitSet.of(0), ImmutableBitSet.of(1), ImmutableBitSet.of())), + builder.aggregateCall( + SqlStdOperatorTable.GROUPING, builder.field(1), builder.field(0)), + builder.countStar("n"), + builder.aggregateCall( + SqlStdOperatorTable.GROUPING_ID, builder.field(0), builder.field(1))) + .build(); + Rel exported = + SubstraitRelVisitor.convert(RelRoot.of(calcite, SqlKind.SELECT), converterProvider) + .getInput(); + assertRowsAndRoundTrip( + exported, + List.of( + row(1, null, 2L, 1L, 1L), + row(2, null, 2L, 1L, 1L), + row(null, 10, 1L, 2L, 2L), + row(null, null, 3L, 2L, 3L))); + } + + private Rel input() { + return virtualTable( + NamedStruct.of(List.of("a", "b"), R.struct(R.I32, R.I32)), + List.of(sb.i32(1), sb.i32(10)), + List.of(sb.i32(2), sb.i32(10))); + } + + private Rel aggregate(List> sets, Rel input) { + return sb.aggregate( + rel -> + sets.stream() + .map(set -> sb.grouping(rel, set.stream().mapToInt(Integer::intValue).toArray())) + .collect(Collectors.toList()), + rel -> List.of(sb.countStar()), + Optional.empty(), + input); + } + + private void assertRowsAndRoundTrip(Rel rel, List> expected) throws SQLException { + RelNode calcite = substraitToCalcite.convert(rel); + assertEquals( + multiset(expected), multiset(execute(calcite)), () -> RelOptUtil.toString(calcite)); + assertRowMatch(calcite.getRowType(), rel.getRecordType().fields()); + Rel exported = + SubstraitRelVisitor.convert(RelRoot.of(calcite, SqlKind.SELECT), converterProvider) + .getInput(); + assertEquals(rel.getRecordType(), exported.getRecordType()); + assertEquals(multiset(expected), multiset(execute(substraitToCalcite.convert(exported)))); + } + + private static List row(Object... values) { + return Arrays.asList(values); + } + + private static Map, Long> multiset(List> rows) { + return rows.stream().collect(Collectors.groupingBy(Function.identity(), Collectors.counting())); + } + + private static List> execute(RelNode rel) throws SQLException { + try (PreparedStatement statement = RelRunners.run(rel); + ResultSet result = statement.executeQuery()) { + List> rows = new ArrayList<>(); + while (result.next()) { + List row = new ArrayList<>(); + for (int column = 1; column <= result.getMetaData().getColumnCount(); column++) { + row.add(result.getObject(column)); + } + rows.add(row); + } + return rows; + } + } +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java b/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java index a8f380022..bdaeb62fa 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java @@ -168,12 +168,9 @@ private Rel twoColumnProject() { } @Test - void leavesAnAggregateThatEmitsDirectlyAlone() { - // 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. + void namesAnAggregateThatEmitsItsGroupingSetIndexDirectly() { + // The derived grouping-set index has the declared i32 type, so all output names apply. + Rel aggregate = sb.aggregate( input -> List.of(sb.grouping(input, 0), sb.grouping(input, 1)), @@ -181,13 +178,12 @@ void leavesAnAggregateThatEmitsDirectlyAlone() { Optional.empty(), scan); - RelNode plain = substraitToCalcite.convert(aggregate); RelNode node = substraitToCalcite.convert( aggregate.withHint( Optional.of(Hint.builder().addOutputNames("k1", "k2", "n", "g").build()))); - assertEquals(plain.getRowType().getFieldNames(), node.getRowType().getFieldNames()); + assertEquals(List.of("k1", "k2", "n", "g"), node.getRowType().getFieldNames()); } @Test @@ -292,10 +288,7 @@ void namesAnAggregateWhoseGroupingColumnsCalciteOrdersDifferently() { } @Test - void dropsNamesWhereTheColumnsAreNotTheRelationsColumns() { - // 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. + void namesAnAggregateThatEmitsItsGroupingSetIndexThroughAMapping() { Rel scan3 = sb.namedScan(List.of("t3"), List.of("a", "b", "c"), List.of(R.I64, N.STRING, R.FP64)); Rel aggregate = @@ -305,13 +298,12 @@ void dropsNamesWhereTheColumnsAreNotTheRelationsColumns() { Optional.of(Rel.Remap.of(List.of(0, 1, 2, 3))), scan3); - RelNode plain = substraitToCalcite.convert(aggregate); RelNode named = substraitToCalcite.convert( aggregate.withHint( Optional.of(Hint.builder().addOutputNames("k_b", "k_a", "n", "gs").build()))); - assertEquals(plain.getRowType().getFieldNames(), named.getRowType().getFieldNames()); + assertEquals(List.of("k_b", "k_a", "n", "gs"), named.getRowType().getFieldNames()); } private Rel hintedInnerProject() { diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java index 364ca20ea..98f2928eb 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java @@ -2,6 +2,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; @@ -596,6 +597,27 @@ void differentDeclaredTypesOnDuplicateMeasuresArePreserved() { assertEquals(2, ((org.apache.calcite.rel.core.Aggregate) relNode).getAggCallList().size()); } + @Test + void groupingSetIndicatorsPreserveDeclaredMeasureTypes() { + Rel input = sb.namedScan(List.of("example"), List.of("a", "g"), List.of(R.I32, R.STRING)); + Rel aggregate = + sb.aggregate( + i -> List.of(sb.grouping(i, 0), sb.grouping(i, 1), sb.grouping(i, 0)), + i -> + List.of( + withOutputType(sb.sum(i, 0), R.I64), withOutputType(sb.sum(i, 0), R.FP64)), + Optional.empty(), + input); + + assertRowMatch( + substraitToCalcite.convert(aggregate).getRowType(), + N.I32, + N.STRING, + R.I64, + R.FP64, + R.I32); + } + @Test void declaredTypeSurvivesAggregateRollupRule() { Rel input = sb.namedScan(List.of("example"), List.of("a", "g"), List.of(R.I32, R.STRING)); @@ -660,8 +682,11 @@ void explicitEmptyGroupingAmongGroupingSets() { // top, so the plans are not structurally identical — but the aggregate itself keeps both // groupings, the empty one included, and the measure's declared type. Rel back = SubstraitRelVisitor.convert(relNode, converterProvider); + while (back instanceof io.substrait.relation.Project) { + back = ((io.substrait.relation.Project) back).getInput(); + } io.substrait.relation.Aggregate aggregateBack = - (io.substrait.relation.Aggregate) ((io.substrait.relation.Project) back).getInput(); + assertInstanceOf(io.substrait.relation.Aggregate.class, back); assertEquals(aggregate.getGroupings(), aggregateBack.getGroupings()); assertEquals(aggregate.getMeasures(), aggregateBack.getMeasures()); }