feat(core,isthmus): support grouping set index in Aggregate (#565)

Signed-off-by: Niels Pardon <[email protected]>
Co-authored-by: Mark S. Lewis <[email protected]>

Author: nielspardon

core/src/main/java/io/substrait/relation/Aggregate.java

@@ -56,7 +56,13 @@ protected Type.Struct deriveRecordType() {
 
     final Stream<Type> measureTypes = getMeasures().stream().map(t -> t.getFunction().getType());
 
-    return TypeCreator.REQUIRED.struct(Stream.concat(groupingTypes, measureTypes));
+    // an aggregate relation with more than one grouping set receives an extra i32 column on the
+    // right-hand side per spec:
+    // https://substrait.io/relations/logical_relations/#aggregate-operation
+    final Stream<Type> groupingSetIndex = Stream.of(TypeCreator.REQUIRED.I32);
+
+    return TypeCreator.REQUIRED.struct(
+        Stream.concat(Stream.concat(groupingTypes, measureTypes), groupingSetIndex));
   }
 
   @Override

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

@@ -34,15 +34,18 @@
 import io.substrait.relation.NamedWrite;
 import io.substrait.relation.Project;
 import io.substrait.relation.Rel;
+import io.substrait.relation.Rel.Remap;
 import io.substrait.relation.Set;
 import io.substrait.relation.Sort;
 import io.substrait.relation.VirtualTableScan;
 import io.substrait.type.NamedStruct;
+import io.substrait.type.TypeCreator;
 import io.substrait.util.VisitationContext;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -74,6 +77,7 @@
 import org.apache.calcite.rex.RexSlot;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.parser.SqlParser;
 import org.apache.calcite.tools.Frameworks;
 import org.apache.calcite.tools.RelBuilder;
@@ -348,8 +352,43 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
         aggregate.getMeasures().stream()
             .map(measure -> fromMeasure(measure, context))
             .collect(java.util.stream.Collectors.toList());
+
+    Optional<Remap> remap = aggregate.getRemap();
+    final int lastFieldIndex = groupExprs.size() + aggregateCalls.size();
+
+    // map grouping set index if it is not removed via remap
+    final boolean emitDirect = remap.isEmpty();
+    final boolean groupingSetIndexGetsRemapped =
+        remap.map(r -> r.indices().contains(lastFieldIndex)).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 int groupingCallIndex = aggregateCalls.size() - 1;
+      if (groupingSetIndexGetsRemapped) {
+        List<Integer> remapList = new LinkedList<>(remap.get().indices());
+        for (int i = 0; i < remapList.size(); i++) {
+          if (remapList.get(i).equals(lastFieldIndex)) {
+            // replace last field index with field index of the GROUP_ID() function call
+            remapList.set(i, groupingCallIndex);
+          }
+        }
+        remap = Optional.of(Remap.of(remapList));
+      }
+    }
+
     RelNode node = relBuilder.push(child).aggregate(groupKey, aggregateCalls).build();
-    return applyRemap(node, aggregate.getRemap());
+    return applyRemap(node, remap);
   }
 
   private AggregateCall fromMeasure(Aggregate.Measure measure, Context context) {

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

@@ -24,6 +24,7 @@
 import io.substrait.relation.EmptyScan;
 import io.substrait.relation.Fetch;
 import io.substrait.relation.Filter;
+import io.substrait.relation.ImmutableAggregate;
 import io.substrait.relation.ImmutableFetch;
 import io.substrait.relation.ImmutableMeasure.Builder;
 import io.substrait.relation.Join;
@@ -34,6 +35,7 @@
 import io.substrait.relation.NamedWrite;
 import io.substrait.relation.Project;
 import io.substrait.relation.Rel;
+import io.substrait.relation.Rel.Remap;
 import io.substrait.relation.Set;
 import io.substrait.relation.Sort;
 import io.substrait.relation.VirtualTableScan;
@@ -45,6 +47,7 @@
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelFieldCollation.Direction;
@@ -58,6 +61,7 @@
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexFieldAccess;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.util.ImmutableBitSet;
 import org.immutables.value.Value;
 
@@ -262,16 +266,65 @@ public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) {
     List<Grouping> groupings =
         sets.filter(s -> s != null).map(s -> fromGroupSet(s, input)).collect(Collectors.toList());
 
-    List<Measure> aggCalls =
+    // get GROUP_ID() function calls
+    List<AggregateCall> groupIdCalls =
+        aggregate.getAggCallList().stream()
+            .filter(c -> c.getAggregation().equals(SqlStdOperatorTable.GROUP_ID))
+            .collect(Collectors.toList());
+
+    List<AggregateCall> filteredAggCalls =
         aggregate.getAggCallList().stream()
+            // remove GROUP_ID() function calls
+            .filter(c -> !groupIdCalls.contains(c))
+            .collect(Collectors.toList());
+
+    List<Measure> aggCalls =
+        filteredAggCalls.stream()
             .map(c -> fromAggCall(aggregate.getInput(), input.getRecordType(), c))
             .collect(Collectors.toList());
 
-    return Aggregate.builder()
-        .input(input)
-        .addAllGroupings(groupings)
-        .addAllMeasures(aggCalls)
-        .build();
+    ImmutableAggregate.Builder builder =
+        Aggregate.builder().input(input).addAllGroupings(groupings).addAllMeasures(aggCalls);
+
+    if (groupings.size() > 1) {
+      // remove the grouping set index if there was no explicit GROUP_ID() function call
+      if (groupIdCalls.isEmpty()) {
+        int groupingExprSize =
+            Math.toIntExact(
+                groupings.stream().flatMap(g -> g.getExpressions().stream()).distinct().count());
+        builder.remap(Remap.offset(0, groupingExprSize + aggCalls.size()));
+      } else {
+        // remap grouping set index at the field positions where the GROUP_ID() function calls were
+        final int groupingFieldCount =
+            Math.toIntExact(groupings.stream().flatMap(g -> g.getExpressions().stream()).count());
+        final int filterAggCallCount = aggCalls.size();
+        final Integer groupingSetIndex = groupingFieldCount + filterAggCallCount;
+
+        final List<Integer> remap =
+            IntStream.range(0, groupingFieldCount)
+                .mapToObj(i -> i)
+                .collect(Collectors.toCollection(ArrayList::new));
+
+        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));
+      }
+    }
+
+    return builder.build();
   }
 
   Aggregate.Grouping fromGroupSet(ImmutableBitSet bitSet, Rel input) {

isthmus/src/test/java/io/substrait/isthmus/Substrait2SqlTest.java

@@ -92,6 +92,14 @@ public void simpleTestGroupingSets() throws Exception {
         "select sum(l_discount) from lineitem group by grouping sets ((l_orderkey, L_COMMITDATE), l_shipdate, ()), l_linestatus");
     assertSqlSubstraitRelRoundTrip(
         "select sum(l_discount) from lineitem group by grouping sets ((l_orderkey, L_COMMITDATE), (l_orderkey, L_COMMITDATE, l_linestatus), l_shipdate, ())");
+
+    // GROUP_ID()
+    assertSqlSubstraitRelRoundTrip(
+        "select sum(l_discount), group_id() from lineitem group by grouping sets ((l_orderkey, L_COMMITDATE), l_shipdate)");
+    assertSqlSubstraitRelRoundTrip(
+        "select group_id(), sum(l_discount) from lineitem group by grouping sets ((l_orderkey, L_COMMITDATE), l_shipdate)");
+    assertSqlSubstraitRelRoundTrip(
+        "select group_id(), sum(l_discount), group_id() from lineitem group by grouping sets ((l_orderkey, L_COMMITDATE), l_shipdate)");
   }
 
   @Test