fix: make SubstraitRelVisitor to use initialSchema, as well as two tests

substrait-io · Jun 20, 2024 · 98ca9ad · 98ca9ad
1 parent d997326
commit 98ca9ad
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/core/src/test/java/io/substrait/type/proto/AggregateRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/AggregateRoundtripTest.java
@@ -13,6 +13,7 @@
 import io.substrait.relation.ProtoRelConverter;
 import io.substrait.relation.RelProtoConverter;
 import io.substrait.relation.VirtualTableScan;
+import io.substrait.type.NamedStruct;
 import io.substrait.type.TypeCreator;
 import java.io.IOException;
 import java.math.BigDecimal;
@@ -25,8 +26,11 @@ public class AggregateRoundtripTest extends TestBase {
   private void assertAggregateRoundtrip(Expression.AggregationInvocation invocation) {
     var expression = ExpressionCreator.decimal(false, BigDecimal.TEN, 10, 2);
     Expression.StructLiteral literal =
-        ImmutableExpression.StructLiteral.builder().from(expression).build();
-    var input = VirtualTableScan.builder().addRows(literal).build();
+        ImmutableExpression.StructLiteral.builder().addFields(expression).build();
+    var input = VirtualTableScan.builder()
+            .initialSchema(NamedStruct.of(Arrays.asList("decimal"), R.struct(R.decimal(10, 2))))
+            .addRows(literal)
+            .build();
     ExtensionCollector functionCollector = new ExtensionCollector();
     var to = new RelProtoConverter(functionCollector);
     var extensions = defaultExtensionCollection;

diff --git a/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java
@@ -74,6 +74,7 @@ protected void verifyRoundTrip(Rel rel) {
   void virtualTable() {
     Rel rel =
         VirtualTableScan.builder()
+            .initialSchema(NamedStruct.of(Collections.emptyList(), R.struct()))
             .addRows(Expression.StructLiteral.builder().fields(Collections.emptyList()).build())
             .commonExtension(commonExtension)
             .extension(relExtension)

diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java
@@ -136,7 +136,7 @@ public Rel visit(org.apache.calcite.rel.core.Values values) {
                   return ExpressionCreator.struct(false, fields);
                 })
             .collect(Collectors.toUnmodifiableList());
-    return VirtualTableScan.builder().addAllDfsNames(type.names()).addAllRows(structs).build();
+    return VirtualTableScan.builder().initialSchema(type).addAllRows(structs).build();
   }
 
   @Override