feat: check that VirtualTableScan field names correspond to schema

substrait-io · Jul 18, 2024 · 8bc5997 · 8bc5997
1 parent ada8d0b
commit 8bc5997
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/core/src/main/java/io/substrait/relation/VirtualTableScan.java b/core/src/main/java/io/substrait/relation/VirtualTableScan.java
@@ -24,13 +24,16 @@ public abstract class VirtualTableScan extends AbstractReadRel {
   @Value.Check
   protected void check() {
     var names = getInitialSchema().names();
+
+    assert names.size()
+        == NamedFieldCountingTypeVisitor.countNames(this.getInitialSchema().struct());
     var rows = getRows();
 
     assert rows.size() > 0
         && names.stream().noneMatch(s -> s == null)
         && rows.stream().noneMatch(r -> r == null)
         && rows.stream()
-            .allMatch(r -> r.getType().accept(new NamedFieldCountingTypeVisitor()) == names.size());
+            .allMatch(r -> NamedFieldCountingTypeVisitor.countNames(r.getType()) == names.size());
   }
 
   @Override
@@ -44,6 +47,14 @@ public static ImmutableVirtualTableScan.Builder builder() {
 
   private static class NamedFieldCountingTypeVisitor
       implements TypeVisitor<Integer, RuntimeException> {
+
+    private static final NamedFieldCountingTypeVisitor VISITOR =
+        new NamedFieldCountingTypeVisitor();
+
+    private static Integer countNames(Type type) {
+      return type.accept(VISITOR);
+    }
+
     @Override
     public Integer visit(Type.Bool type) throws RuntimeException {
       return 0;

diff --git a/core/src/test/java/io/substrait/relation/VirtualTableScanTest.java b/core/src/test/java/io/substrait/relation/VirtualTableScanTest.java
@@ -39,8 +39,8 @@ void check() {
                     R.struct(
                         R.STRING,
                         R.struct(R.STRING, R.STRING),
-                        R.list(R.STRING),
-                        R.map(R.STRING, R.STRING))))
+                        R.list(R.struct(R.STRING)),
+                        R.map(R.struct(R.STRING), R.struct(R.STRING)))))
             .addRows(
                 struct(
                     false,