fix: apply feedback

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

Author: nielspardon

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

@@ -12,6 +12,8 @@
 import io.substrait.proto.ReadRel;
 import io.substrait.proto.Rel;
 import io.substrait.util.EmptyVisitationContext;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.junit.jupiter.api.Test;
 
 class AggregateRelTest extends TestBase {
@@ -58,6 +60,43 @@ public static io.substrait.proto.Expression createFieldReference(int col) {
     return Expression.newBuilder().setSelection(fieldRef1).build();
   }
 
+  /**
+   * Helper method to extract expression references from an AggregateRel.
+   *
+   * @param aggregateRel the AggregateRel to extract expression references from
+   * @return a list of lists, where each inner list contains the expression reference indices for a
+   *     grouping
+   */
+  private static List<List<Integer>> getExpressionReferences(AggregateRel aggregateRel) {
+    return aggregateRel.getGroupingsList().stream()
+        .map(
+            grouping ->
+                grouping.getExpressionReferencesList().stream().collect(Collectors.toList()))
+        .collect(Collectors.toList());
+  }
+
+  /**
+   * Helper method to extract deprecated grouping expressions from an AggregateRel.
+   *
+   * @param aggregateRel the AggregateRel to extract grouping expressions from
+   * @return a list of lists, where each inner list contains the grouping expressions for a grouping
+   */
+  private static List<List<Expression>> getGroupingExpressions(AggregateRel aggregateRel) {
+    return aggregateRel.getGroupingsList().stream()
+        .map(grouping -> grouping.getGroupingExpressionsList())
+        .collect(Collectors.toList());
+  }
+
+  /**
+   * Helper method to extract aggregate-level grouping expressions from an AggregateRel.
+   *
+   * @param aggregateRel the AggregateRel to extract grouping expressions from
+   * @return a list of expressions at the aggregate level
+   */
+  private static List<Expression> getAggregateGroupingExpressions(AggregateRel aggregateRel) {
+    return aggregateRel.getGroupingExpressionsList();
+  }
+
   @Test
   void testDeprecatedGroupingExpressionConversion() {
     Expression col1Ref = createFieldReference(0);
@@ -96,19 +135,12 @@ void testDeprecatedGroupingExpressionConversion() {
     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(
-        2,
-        roundtripAgg.getGroupings(0).getGroupingExpressionsCount(),
-        "grouping expressions count of grouping should be 2");
-    assertEquals(
-        2,
-        roundtripAgg.getGroupings(0).getExpressionReferencesCount(),
-        "expression reference count of grouping should be 2");
+    // Verify new expression_references structure
+    assertEquals(List.of(List.of(0, 1)), getExpressionReferences(roundtripAgg));
+    // Verify backward compatibility: deprecated grouping_expressions field is also populated
+    assertEquals(List.of(List.of(col1Ref, col2Ref)), getGroupingExpressions(roundtripAgg));
+    // Verify aggregate-level grouping_expressions field is populated
+    assertEquals(List.of(col1Ref, col2Ref), getAggregateGroupingExpressions(roundtripAgg));
   }
 
   @Test
@@ -151,19 +183,12 @@ void testAggregateWithSingleGrouping() {
     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(
-        2,
-        roundtripAgg.getGroupings(0).getGroupingExpressionsCount(),
-        "grouping expressions count of grouping should be 2");
-    assertEquals(
-        2,
-        roundtripAgg.getGroupings(0).getExpressionReferencesCount(),
-        "expression reference count of grouping should be 2");
+    // Verify new expression_references structure
+    assertEquals(List.of(List.of(0, 1)), getExpressionReferences(roundtripAgg));
+    // Verify backward compatibility: deprecated grouping_expressions field is also populated
+    assertEquals(List.of(List.of(col1Ref, col2Ref)), getGroupingExpressions(roundtripAgg));
+    // Verify aggregate-level grouping_expressions field is populated
+    assertEquals(List.of(col1Ref, col2Ref), getAggregateGroupingExpressions(roundtripAgg));
   }
 
   @Test
@@ -211,22 +236,12 @@ void testAggregateWithMultipleGroupings() {
     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(
-        2,
-        roundtripAgg.getGroupings(0).getGroupingExpressionsCount(),
-        "grouping expressions count of grouping should be 2");
-    assertEquals(
-        2,
-        roundtripAgg.getGroupings(0).getExpressionReferencesCount(),
-        "expression reference count of grouping should be 2");
+    // Verify new expression_references structure
+    assertEquals(List.of(List.of(0, 1), List.of(1)), getExpressionReferences(roundtripAgg));
+    // Verify backward compatibility: deprecated grouping_expressions field is also populated
     assertEquals(
-        1,
-        roundtripAgg.getGroupings(1).getExpressionReferencesCount(),
-        "expression reference count of grouping should be 2");
+        List.of(List.of(col1Ref, col2Ref), List.of(col2Ref)), getGroupingExpressions(roundtripAgg));
+    // Verify aggregate-level grouping_expressions field is populated
+    assertEquals(List.of(col1Ref, col2Ref), getAggregateGroupingExpressions(roundtripAgg));
   }
 }