Merge branch 'main' into gordon.hamilton/aggregateGroupingNewSubstraitForm

vbarua · web-flow · commit dfd50ebeffbb · 2025-10-01T19:51:24.000-07:00
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -52,7 +52,7 @@ allprojects {
         googleJavaFormat()
         removeUnusedImports()
         trimTrailingWhitespace()
-        removeWildcardImports()
+        forbidWildcardImports()
       }
     }
   }
diff --git a/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java b/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java
@@ -2,6 +2,7 @@
 
 import io.substrait.expression.Expression;
 import io.substrait.expression.proto.ProtoExpressionConverter;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.ExtensionCollector;
 import io.substrait.extension.ExtensionLookup;
 import io.substrait.extension.ImmutableExtensionLookup;
@@ -23,7 +24,7 @@ public class ProtoExtendedExpressionConverter {
           new ExtensionCollector(), SimpleExtension.ExtensionCollection.builder().build());
 
   public ProtoExtendedExpressionConverter() {
-    this(SimpleExtension.loadDefaults());
+    this(DefaultExtensionCatalog.DEFAULT_COLLECTION);
   }
 
   public ProtoExtendedExpressionConverter(SimpleExtension.ExtensionCollection extensionCollection) {
diff --git a/core/src/main/java/io/substrait/extension/DefaultExtensionCatalog.java b/core/src/main/java/io/substrait/extension/DefaultExtensionCatalog.java
@@ -1,5 +1,9 @@
 package io.substrait.extension;
 
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
 public class DefaultExtensionCatalog {
   public static final String FUNCTIONS_AGGREGATE_APPROX = "/functions_aggregate_approx.yaml";
   public static final String FUNCTIONS_AGGREGATE_GENERIC = "/functions_aggregate_generic.yaml";
@@ -14,4 +18,28 @@ public class DefaultExtensionCatalog {
   public static final String FUNCTIONS_ROUNDING_DECIMAL = "/functions_rounding_decimal.yaml";
   public static final String FUNCTIONS_SET = "/functions_set.yaml";
   public static final String FUNCTIONS_STRING = "/functions_string.yaml";
+
+  public static final SimpleExtension.ExtensionCollection DEFAULT_COLLECTION =
+      loadDefaultCollection();
+
+  private static SimpleExtension.ExtensionCollection loadDefaultCollection() {
+    List<String> defaultFiles =
+        Arrays.asList(
+                "boolean",
+                "aggregate_generic",
+                "aggregate_approx",
+                "arithmetic_decimal",
+                "arithmetic",
+                "comparison",
+                "datetime",
+                "logarithmic",
+                "rounding",
+                "rounding_decimal",
+                "string")
+            .stream()
+            .map(c -> String.format("/functions_%s.yaml", c))
+            .collect(Collectors.toList());
+
+    return SimpleExtension.load(defaultFiles);
+  }
 }
diff --git a/core/src/main/java/io/substrait/extension/SimpleExtension.java b/core/src/main/java/io/substrait/extension/SimpleExtension.java
@@ -22,7 +22,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.UncheckedIOException;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -701,27 +700,6 @@ public ExtensionCollection merge(ExtensionCollection extensionCollection) {
     }
   }
 
-  public static ExtensionCollection loadDefaults() {
-    List<String> defaultFiles =
-        Arrays.asList(
-                "boolean",
-                "aggregate_generic",
-                "aggregate_approx",
-                "arithmetic_decimal",
-                "arithmetic",
-                "comparison",
-                "datetime",
-                "logarithmic",
-                "rounding",
-                "rounding_decimal",
-                "string")
-            .stream()
-            .map(c -> String.format("/functions_%s.yaml", c))
-            .collect(Collectors.toList());
-
-    return load(defaultFiles);
-  }
-
   public static ExtensionCollection load(List<String> resourcePaths) {
     if (resourcePaths.isEmpty()) {
       throw new IllegalArgumentException("Require at least one resource path.");
diff --git a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java
@@ -1,5 +1,6 @@
 package io.substrait.plan;
 
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.ExtensionLookup;
 import io.substrait.extension.ImmutableExtensionLookup;
 import io.substrait.extension.SimpleExtension;
@@ -16,7 +17,7 @@ public class ProtoPlanConverter {
   protected final SimpleExtension.ExtensionCollection extensionCollection;
 
   public ProtoPlanConverter() {
-    this(SimpleExtension.loadDefaults());
+    this(DefaultExtensionCatalog.DEFAULT_COLLECTION);
   }
 
   public ProtoPlanConverter(SimpleExtension.ExtensionCollection extensionCollection) {
diff --git a/core/src/main/java/io/substrait/relation/Aggregate.java b/core/src/main/java/io/substrait/relation/Aggregate.java
@@ -21,17 +21,42 @@ 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),
-
-            // measures
-            getMeasures().stream().map(t -> t.getFunction().getType())));
+    // If there's only one grouping set (or none), the nullability rule doesn't apply.
+    if (getGroupings().size() <= 1) {
+      final Stream<Type> groupingTypes =
+          getGroupings().stream()
+              .flatMap(g -> g.getExpressions().stream())
+              .map(Expression::getType);
+      final Stream<Type> measureTypes = getMeasures().stream().map(t -> t.getFunction().getType());
+      return TypeCreator.REQUIRED.struct(Stream.concat(groupingTypes, measureTypes));
+    }
+
+    final LinkedHashSet<Expression> uniqueGroupingExpressions =
+        getGroupings().stream()
+            .flatMap(g -> g.getExpressions().stream())
+            .collect(Collectors.toCollection(LinkedHashSet::new));
+
+    // For each unique grouping expression, determine its final nullability based on the spec.
+    final Stream<Type> groupingTypes =
+        uniqueGroupingExpressions.stream()
+            .map(
+                expr -> {
+                  // the code below implements the following statement from the spec
+                  // (https://substrait.io/relations/logical_relations/#aggregate-operation):
+                  // "The values for the grouping expression columns that are not
+                  // part of the grouping set for a particular record will be set to null."
+                  final boolean appearsInAllSets =
+                      getGroupings().stream().allMatch(g -> g.getExpressions().contains(expr));
+                  if (appearsInAllSets) {
+                    return expr.getType();
+                  } else {
+                    return TypeCreator.asNullable(expr.getType());
+                  }
+                });
+
+    final Stream<Type> measureTypes = getMeasures().stream().map(t -> t.getFunction().getType());
+
+    return TypeCreator.REQUIRED.struct(Stream.concat(groupingTypes, measureTypes));
   }
 
   @Override
diff --git a/core/src/main/java/io/substrait/relation/ProtoAggregateFunctionConverter.java b/core/src/main/java/io/substrait/relation/ProtoAggregateFunctionConverter.java
@@ -5,6 +5,7 @@
 import io.substrait.expression.FunctionArg;
 import io.substrait.expression.FunctionOption;
 import io.substrait.expression.proto.ProtoExpressionConverter;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.ExtensionLookup;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.type.proto.ProtoTypeConverter;
@@ -24,7 +25,7 @@ public class ProtoAggregateFunctionConverter {
 
   public ProtoAggregateFunctionConverter(
       ExtensionLookup lookup, ProtoExpressionConverter protoExpressionConverter) {
-    this(lookup, SimpleExtension.loadDefaults(), protoExpressionConverter);
+    this(lookup, DefaultExtensionCatalog.DEFAULT_COLLECTION, protoExpressionConverter);
   }
 
   public ProtoAggregateFunctionConverter(
diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java
@@ -3,6 +3,7 @@
 import io.substrait.expression.Expression;
 import io.substrait.expression.proto.ProtoExpressionConverter;
 import io.substrait.extension.AdvancedExtension;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.ExtensionLookup;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.hint.Hint;
@@ -55,7 +56,7 @@ public class ProtoRelConverter {
   private final ProtoTypeConverter protoTypeConverter;
 
   public ProtoRelConverter(ExtensionLookup lookup) {
-    this(lookup, SimpleExtension.loadDefaults());
+    this(lookup, DefaultExtensionCatalog.DEFAULT_COLLECTION);
   }
 
   public ProtoRelConverter(ExtensionLookup lookup, SimpleExtension.ExtensionCollection extensions) {
diff --git a/core/src/test/java/io/substrait/TestBase.java b/core/src/test/java/io/substrait/TestBase.java
@@ -3,6 +3,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import io.substrait.dsl.SubstraitBuilder;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.ExtensionCollector;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.relation.ProtoRelConverter;
@@ -13,7 +14,7 @@
 public abstract class TestBase {
 
   protected static final SimpleExtension.ExtensionCollection defaultExtensionCollection =
-      SimpleExtension.loadDefaults();
+      DefaultExtensionCatalog.DEFAULT_COLLECTION;
 
   protected TypeCreator R = TypeCreator.REQUIRED;
 
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -6,8 +6,8 @@ graal = "25.0.0"
 graal-plugin = "0.11.0"
 gradle-extensions = "2.0.0"
 guava = "33.5.0-jre"
-httpclient5 = "5.5"
-immutables = "2.11.3"
+httpclient5 = "5.5.1"
+immutables = "2.11.4"
 jackson = "2.20.0"
 jreleaser = "1.20.0"
 json-smart = "2.6.0"
@@ -20,10 +20,10 @@ reflections = "0.9.12"
 scala-library = "2.12.20"
 scalatest = "3.2.19"
 scalatestplus-junit5 = "3.2.19.0"
-shadow = "9.1.0"
+shadow = "9.2.2"
 slf4j = "2.0.17"
 spark = "3.4.4"
-spotless = "7.2.1"
+spotless = "8.0.0"
 
 [libraries]
 antlr4 = { module = "org.antlr:antlr4", version.ref = "antlr" }
diff --git a/isthmus-cli/src/main/java/io/substrait/isthmus/cli/IsthmusEntryPoint.java b/isthmus-cli/src/main/java/io/substrait/isthmus/cli/IsthmusEntryPoint.java
@@ -4,7 +4,7 @@
 import com.google.protobuf.Message;
 import com.google.protobuf.TextFormat;
 import com.google.protobuf.util.JsonFormat;
-import io.substrait.extension.SimpleExtension;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.isthmus.FeatureBoard;
 import io.substrait.isthmus.ImmutableFeatureBoard;
 import io.substrait.isthmus.SqlExpressionToSubstrait;
@@ -87,7 +87,7 @@ public Integer call() throws Exception {
     // Isthmus image is parsing SQL Expression if that argument is defined
     if (sqlExpressions != null) {
       SqlExpressionToSubstrait converter =
-          new SqlExpressionToSubstrait(featureBoard, SimpleExtension.loadDefaults());
+          new SqlExpressionToSubstrait(featureBoard, DefaultExtensionCatalog.DEFAULT_COLLECTION);
       ExtendedExpression extendedExpression = converter.convert(sqlExpressions, createStatements);
       printMessage(extendedExpression);
     } else { // by default Isthmus image are parsing SQL Query
diff --git a/isthmus/src/main/java/io/substrait/isthmus/SqlConverterBase.java b/isthmus/src/main/java/io/substrait/isthmus/SqlConverterBase.java
@@ -1,5 +1,6 @@
 package io.substrait.isthmus;
 
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.SimpleExtension;
 import org.apache.calcite.config.CalciteConnectionConfig;
 import org.apache.calcite.config.CalciteConnectionProperty;
@@ -20,7 +21,7 @@
 
 class SqlConverterBase {
   protected static final SimpleExtension.ExtensionCollection EXTENSION_COLLECTION =
-      SimpleExtension.loadDefaults();
+      DefaultExtensionCatalog.DEFAULT_COLLECTION;
 
   final RelDataTypeFactory factory;
   final RelOptCluster relOptCluster;
diff --git a/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java b/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java
@@ -6,6 +6,7 @@
 import com.google.common.collect.ImmutableList;
 import io.substrait.expression.Expression;
 import io.substrait.expression.ExpressionCreator;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.isthmus.SubstraitRelNodeConverter.Context;
 import io.substrait.isthmus.expression.ExpressionRexConverter;
@@ -22,7 +23,7 @@
 public class CalciteCallTest extends CalciteObjs {
 
   private static final SimpleExtension.ExtensionCollection EXTENSION_COLLECTION =
-      SimpleExtension.loadDefaults();
+      DefaultExtensionCatalog.DEFAULT_COLLECTION;
   private final ScalarFunctionConverter functionConverter =
       new ScalarFunctionConverter(EXTENSION_COLLECTION.scalarFunctions(), type);
   private final RexExpressionConverter rexExpressionConverter =
diff --git a/isthmus/src/test/java/io/substrait/isthmus/RelCopyOnWriteVisitorTest.java b/isthmus/src/test/java/io/substrait/isthmus/RelCopyOnWriteVisitorTest.java
@@ -6,6 +6,7 @@
 
 import io.substrait.expression.AggregateFunctionInvocation;
 import io.substrait.expression.Expression;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.isthmus.sql.SubstraitSqlDialect;
 import io.substrait.plan.Plan;
@@ -249,7 +250,8 @@ private static class ReplaceCountDistinctWithApprox {
     private final ReplaceCountDistinctWithApproxVisitor visitor;
 
     public ReplaceCountDistinctWithApprox() {
-      visitor = new ReplaceCountDistinctWithApproxVisitor(SimpleExtension.loadDefaults());
+      visitor =
+          new ReplaceCountDistinctWithApproxVisitor(DefaultExtensionCatalog.DEFAULT_COLLECTION);
     }
 
     public Optional<Plan> modify(Plan plan) {
diff --git a/isthmus/src/test/java/io/substrait/isthmus/Substrait2SqlTest.java b/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(
diff --git a/isthmus/src/test/java/io/substrait/isthmus/TpcdsQueryTest.java b/isthmus/src/test/java/io/substrait/isthmus/TpcdsQueryTest.java
@@ -16,7 +16,7 @@ public class TpcdsQueryTest extends PlanTestBase {
   private static final Set<Integer> alternateForms = Set.of(27, 36, 70, 86);
   private static final Set<Integer> toSubstraitExclusions = Set.of(9);
   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));
diff --git a/spark/.gitignore b/spark/.gitignore
@@ -0,0 +1,4 @@
+metastore_db
+spark-warehouse
+/src/test/resources/write-a.csv
+derby.log
diff --git a/spark/build.gradle.kts b/spark/build.gradle.kts
@@ -118,6 +118,7 @@ dependencies {
   implementation(libs.scala.library)
   api(libs.spark.core)
   api(libs.spark.sql)
+  implementation(libs.spark.hive)
   implementation(libs.spark.catalyst)
   implementation(libs.slf4j.api)
 
@@ -150,6 +151,9 @@ tasks {
   test {
     dependsOn(":core:shadowJar")
     useJUnitPlatform { includeEngines("scalatest") }
-    jvmArgs("--add-exports=java.base/sun.nio.ch=ALL-UNNAMED")
+    jvmArgs(
+      "--add-exports=java.base/sun.nio.ch=ALL-UNNAMED",
+      "--add-opens=java.base/java.net=ALL-UNNAMED",
+    )
   }
 }
diff --git a/spark/src/main/scala/io/substrait/spark/SparkExtension.scala b/spark/src/main/scala/io/substrait/spark/SparkExtension.scala
diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala
diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala
diff --git a/spark/src/test/scala/io/substrait/spark/HiveTableSuite.scala b/spark/src/test/scala/io/substrait/spark/HiveTableSuite.scala

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@ allprojects {`
`52`	`52`	`googleJavaFormat()`
`53`	`53`	`removeUnusedImports()`
`54`	`54`	`trimTrailingWhitespace()`
`55`		`- removeWildcardImports()`
	`55`	`+ forbidWildcardImports()`
`56`	`56`	`}`
`57`	`57`	`}`
`58`	`58`	`}`