fix(core): use new aggregate grouping behavior in RelProtoConverter

Signed-off-by: Niels Pardon <par@zurich.ibm.com>

Author: nielspardon

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

@@ -155,12 +155,21 @@ private io.substrait.proto.Expression.FieldReference toProto(FieldReference fiel
 
   @Override
   public Rel visit(Aggregate aggregate, EmptyVisitationContext context) throws RuntimeException {
-    AggregateRel.Builder builder =
+    final List<Expression> uniqueGroupingExpressions =
+        aggregate.getGroupings().stream()
+            .flatMap(g -> g.getExpressions().stream())
+            .distinct()
+            .collect(Collectors.toList());
+
+    final AggregateRel.Builder builder =
         AggregateRel.newBuilder()
             .setInput(toProto(aggregate.getInput()))
             .setCommon(common(aggregate))
+            .addAllGroupingExpressions(toProto(uniqueGroupingExpressions))
             .addAllGroupings(
-                aggregate.getGroupings().stream().map(this::toProto).collect(Collectors.toList()))
+                aggregate.getGroupings().stream()
+                    .map(g -> toProto(g, uniqueGroupingExpressions))
+                    .collect(Collectors.toList()))
             .addAllMeasures(
                 aggregate.getMeasures().stream().map(this::toProto).collect(Collectors.toList()));
 
@@ -203,9 +212,13 @@ private AggregateRel.Measure toProto(Aggregate.Measure measure) {
     return builder.build();
   }
 
-  private AggregateRel.Grouping toProto(Aggregate.Grouping grouping) {
+  private AggregateRel.Grouping toProto(
+      Aggregate.Grouping grouping, List<Expression> uniqueGroupingExpressions) {
     return AggregateRel.Grouping.newBuilder()
-        .addAllGroupingExpressions(toProto(grouping.getExpressions()))
+        .addAllExpressionReferences(
+            grouping.getExpressions().stream()
+                .map(e -> uniqueGroupingExpressions.indexOf(e))
+                .collect(Collectors.toList()))
         .build();
   }
 

core/src/test/java/io/substrait/relation/AggregateRelTest.java

@@ -11,6 +11,7 @@
 import io.substrait.proto.Plan;
 import io.substrait.proto.ReadRel;
 import io.substrait.proto.Rel;
+import io.substrait.util.EmptyVisitationContext;
 import org.junit.jupiter.api.Test;
 
 class AggregateRelTest extends TestBase {
@@ -90,6 +91,24 @@ void testDeprecatedGroupingExpressionConversion() {
     Aggregate agg = (Aggregate) resultRel;
     assertEquals(1, agg.getGroupings().size());
     assertEquals(2, agg.getGroupings().get(0).getExpressions().size());
+
+    Rel roundtripRel = relProtoConverter.visit(agg, EmptyVisitationContext.INSTANCE);
+    assertTrue(roundtripRel.hasAggregate());
+
+    AggregateRel roundtripAggr = roundtripRel.getAggregate();
+    assertEquals(1, roundtripAggr.getGroupingsCount(), "grouping count should be 1");
+    assertEquals(
+        2,
+        roundtripAggr.getGroupingExpressionsCount(),
+        "grouping expressions count of aggregate should be 2");
+    assertEquals(
+        0,
+        roundtripAggr.getGroupings(0).getGroupingExpressionsCount(),
+        "grouping expressions count of grouping should be 0");
+    assertEquals(
+        2,
+        roundtripAggr.getGroupings(0).getExpressionReferencesCount(),
+        "expression reference count of grouping should be 2");
   }
 
   @Test
@@ -127,6 +146,24 @@ void testAggregateWithSingleGrouping() {
     Aggregate agg = (Aggregate) resultRel;
     assertEquals(1, agg.getGroupings().size());
     assertEquals(2, agg.getGroupings().get(0).getExpressions().size());
+
+    Rel roundtripRel = relProtoConverter.visit(agg, EmptyVisitationContext.INSTANCE);
+    assertTrue(roundtripRel.hasAggregate());
+
+    AggregateRel roundtripAgg = roundtripRel.getAggregate();
+    assertEquals(1, roundtripAgg.getGroupingsCount(), "grouping count should be 1");
+    assertEquals(
+        2,
+        roundtripAgg.getGroupingExpressionsCount(),
+        "grouping expressions count of aggregate should be 2");
+    assertEquals(
+        0,
+        roundtripAgg.getGroupings(0).getGroupingExpressionsCount(),
+        "grouping expressions count of grouping should be 0");
+    assertEquals(
+        2,
+        roundtripAgg.getGroupings(0).getExpressionReferencesCount(),
+        "expression reference count of grouping should be 2");
   }
 
   @Test
@@ -169,5 +206,27 @@ void testAggregateWithMultipleGroupings() {
     assertEquals(2, agg.getGroupings().size());
     assertEquals(2, agg.getGroupings().get(0).getExpressions().size());
     assertEquals(1, agg.getGroupings().get(1).getExpressions().size());
+
+    Rel roundtripRel = relProtoConverter.visit(agg, EmptyVisitationContext.INSTANCE);
+    assertTrue(roundtripRel.hasAggregate());
+
+    AggregateRel roundtripAgg = roundtripRel.getAggregate();
+    assertEquals(2, roundtripAgg.getGroupingsCount(), "grouping count should be 2");
+    assertEquals(
+        2,
+        roundtripAgg.getGroupingExpressionsCount(),
+        "grouping expressions count of aggregate should be 2");
+    assertEquals(
+        0,
+        roundtripAgg.getGroupings(0).getGroupingExpressionsCount(),
+        "grouping expressions count of grouping should be 0");
+    assertEquals(
+        2,
+        roundtripAgg.getGroupings(0).getExpressionReferencesCount(),
+        "expression reference count of grouping should be 2");
+    assertEquals(
+        1,
+        roundtripAgg.getGroupings(1).getExpressionReferencesCount(),
+        "expression reference count of grouping should be 2");
   }
 }