fix(isthmus): fix nullability for grouping sets

Author: ZorinAnton

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

@@ -21,17 +21,24 @@ public abstract class Aggregate extends SingleInputRel implements HasExtension {
 
   @Override
   protected Type.Struct deriveRecordType() {
-    return TypeCreator.REQUIRED.struct(
-        Stream.concat(
-            // unique grouping expressions
-            getGroupings().stream()
-                .flatMap(g -> g.getExpressions().stream())
-                .collect(Collectors.toCollection(LinkedHashSet::new))
-                .stream()
-                .map(Expression::getType),
+    boolean isGroupingSet = getGroupings().size() > 1;
 
-            // measures
-            getMeasures().stream().map(t -> t.getFunction().getType())));
+    Stream<Type> groupingTypes =
+        getGroupings().stream()
+            .flatMap(g -> g.getExpressions().stream())
+            .collect(Collectors.toCollection(LinkedHashSet::new))
+            .stream()
+            .map(
+                expr -> {
+                  if (isGroupingSet) {
+                    return TypeCreator.asNullable(expr.getType());
+                  }
+                  return expr.getType();
+                });
+
+    Stream<Type> measureTypes = getMeasures().stream().map(t -> t.getFunction().getType());
+
+    return TypeCreator.REQUIRED.struct(Stream.concat(groupingTypes, measureTypes));
   }
 
   @Override

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

@@ -1,9 +1,11 @@
 package io.substrait.isthmus;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 import io.substrait.isthmus.utils.SetUtils;
+import io.substrait.plan.Plan;
 import io.substrait.relation.Set;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelRoot;
@@ -14,6 +16,14 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 public class Substrait2SqlTest extends PlanTestBase {
+  private void assertSqlRoundTripViaPojoAndProto(String inputSql) {
+    Plan plan =
+        assertDoesNotThrow(() -> toSubstraitPlan(inputSql, TPCH_CATALOG), "SQL to Substrait POJO");
+    assertDoesNotThrow(() -> toSql(plan), "Substrait POJO to SQL");
+    io.substrait.proto.Plan proto =
+        assertDoesNotThrow(() -> toProto(plan), "Substrait POJO to Substrait PROTO");
+    assertDoesNotThrow(() -> toSql(proto), "Substrait PROTO to SQL");
+  }
 
   @Test
   public void simpleTest() throws Exception {
@@ -84,6 +94,12 @@ public void simpleTestGroupingSets() throws Exception {
         "select sum(l_discount) from lineitem group by grouping sets ((l_orderkey, L_COMMITDATE), (l_orderkey, L_COMMITDATE, l_linestatus), l_shipdate, ())");
   }
 
+  @Test
+  void testRollup() {
+    assertSqlRoundTripViaPojoAndProto(
+        "select charcol from (select charcol, count(*) from (values('a')) as t(charcol) group by rollup(charcol))");
+  }
+
   @Test
   public void simpleTestAggFilter() throws Exception {
     assertSqlSubstraitRelRoundTrip(

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

@@ -15,7 +15,7 @@
 public class TpcdsQueryTest extends PlanTestBase {
   private static final Set<Integer> toSubstraitExclusions = Set.of(9, 27, 36, 70, 86);
   private static final Set<Integer> fromSubstraitPojoExclusions = Set.of(1, 30, 81);
-  private static final Set<Integer> fromSubstraitProtoExclusions = Set.of(1, 30, 67, 81);
+  private static final Set<Integer> fromSubstraitProtoExclusions = Set.of(1, 30, 81);
 
   static IntStream testCases() {
     return IntStream.rangeClosed(1, 99).filter(n -> !toSubstraitExclusions.contains(n));