feat: support Nested Structs (#579)

Author: gord02

core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java

@@ -151,6 +151,11 @@ public O visit(Expression.StructLiteral expr, C context) throws E {
     return visitFallback(expr, context);
   }
 
+  @Override
+  public O visit(Expression.NestedStruct expr, C context) throws E {
+    return visitFallback(expr, context);
+  }
+
   @Override
   public O visit(Expression.Switch expr, C context) throws E {
     return visitFallback(expr, context);

core/src/main/java/io/substrait/expression/Expression.java

@@ -32,6 +32,13 @@ default boolean nullable() {
     }
   }
 
+  interface Nested extends Expression {
+    @Value.Default
+    default boolean nullable() {
+      return false;
+    }
+  }
+
   <R, C extends VisitationContext, E extends Throwable> R accept(
       ExpressionVisitor<R, C, E> visitor, C context) throws E;
 
@@ -662,6 +669,30 @@ public <R, C extends VisitationContext, E extends Throwable> R accept(
     }
   }
 
+  @Value.Immutable
+  abstract class NestedStruct implements Nested {
+    public abstract List<Expression> fields();
+
+    @Override
+    public Type getType() {
+      return Type.withNullability(nullable())
+          .struct(
+              fields().stream()
+                  .map(Expression::getType)
+                  .collect(java.util.stream.Collectors.toList()));
+    }
+
+    public static ImmutableExpression.NestedStruct.Builder builder() {
+      return ImmutableExpression.NestedStruct.builder();
+    }
+
+    @Override
+    public <R, C extends VisitationContext, E extends Throwable> R accept(
+        ExpressionVisitor<R, C, E> visitor, C context) throws E {
+      return visitor.visit(this, context);
+    }
+  }
+
   @Value.Immutable
   abstract class UserDefinedLiteral implements Literal {
     public abstract ByteString value();

core/src/main/java/io/substrait/expression/ExpressionCreator.java

@@ -286,6 +286,15 @@ public static Expression.StructLiteral struct(
     return Expression.StructLiteral.builder().nullable(nullable).addAllFields(values).build();
   }
 
+  public static Expression.NestedStruct nestedStruct(
+      boolean nullable, Iterable<Expression> fields) {
+    return Expression.NestedStruct.builder().nullable(nullable).addAllFields(fields).build();
+  }
+
+  public static Expression.NestedStruct nestedStruct(boolean nullable, Expression... fields) {
+    return Expression.NestedStruct.builder().nullable(nullable).addFields(fields).build();
+  }
+
   public static Expression.UserDefinedLiteral userDefinedLiteral(
       boolean nullable, String urn, String name, Any value) {
     return Expression.UserDefinedLiteral.builder()

core/src/main/java/io/substrait/expression/ExpressionVisitor.java

@@ -62,6 +62,8 @@ public interface ExpressionVisitor<R, C extends VisitationContext, E extends Thr
 
   R visit(Expression.StructLiteral expr, C context) throws E;
 
+  R visit(Expression.NestedStruct expr, C context) throws E;
+
   R visit(Expression.UserDefinedLiteral expr, C context) throws E;
 
   R visit(Expression.Switch expr, C context) throws E;

core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java

@@ -75,6 +75,12 @@ private Expression lit(Consumer<Expression.Literal.Builder> consumer) {
     return Expression.newBuilder().setLiteral(builder).build();
   }
 
+  private Expression nested(Consumer<Expression.Nested.Builder> consumer) {
+    Expression.Nested.Builder builder = Expression.Nested.newBuilder();
+    consumer.accept(builder);
+    return Expression.newBuilder().setNested(builder).build();
+  }
+
   @Override
   public Expression visit(
       io.substrait.expression.Expression.BoolLiteral expr, EmptyVisitationContext context) {
@@ -357,6 +363,18 @@ public Expression visit(
         });
   }
 
+  @Override
+  public Expression visit(
+      io.substrait.expression.Expression.NestedStruct expr, EmptyVisitationContext context) {
+    return nested(
+        bldr -> {
+          List<Expression> values =
+              expr.fields().stream().map(this::toProto).collect(Collectors.toList());
+          bldr.setStruct(Expression.Nested.Struct.newBuilder().addAllFields(values))
+              .setNullable(expr.nullable());
+        });
+  }
+
   @Override
   public Expression visit(
       io.substrait.expression.Expression.UserDefinedLiteral expr, EmptyVisitationContext context) {

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java

@@ -502,6 +502,18 @@ public Expression.Literal from(io.substrait.proto.Expression.Literal literal) {
     }
   }
 
+  public Expression.StructLiteral from(io.substrait.proto.Expression.Literal.Struct struct) {
+    return Expression.StructLiteral.builder()
+        .fields(struct.getFieldsList().stream().map(this::from).collect(Collectors.toList()))
+        .build();
+  }
+
+  public Expression.NestedStruct from(io.substrait.proto.Expression.Nested.Struct struct) {
+    return Expression.NestedStruct.builder()
+        .fields(struct.getFieldsList().stream().map(this::from).collect(Collectors.toList()))
+        .build();
+  }
+
   private static List<FunctionArg> fromFunctionArgumentList(
       int argumentsCount,
       FunctionArg.ProtoFrom argVisitor,

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

@@ -203,6 +203,15 @@ public Optional<Expression> visit(Expression.StructLiteral expr, EmptyVisitation
     return visitLiteral(expr);
   }
 
+  @Override
+  public Optional<Expression> visit(Expression.NestedStruct expr, EmptyVisitationContext context)
+      throws E {
+    Optional<List<Expression>> expressions = visitExprList(expr.fields(), context);
+    return expressions.map(
+        expressionList ->
+            Expression.NestedStruct.builder().from(expr).fields(expressionList).build());
+  }
+
   @Override
   public Optional<Expression> visit(
       Expression.UserDefinedLiteral expr, EmptyVisitationContext context) throws E {

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

@@ -188,7 +188,7 @@ public Rel from(io.substrait.proto.Rel rel) {
   protected Rel newRead(ReadRel rel) {
     if (rel.hasVirtualTable()) {
       ReadRel.VirtualTable virtualTable = rel.getVirtualTable();
-      if (virtualTable.getValuesCount() == 0) {
+      if (virtualTable.getValuesCount() == 0 && virtualTable.getExpressionsCount() == 0) {
         return newEmptyScan(rel);
       } else {
         return newVirtualTable(rel);
@@ -596,20 +596,50 @@ protected FileOrFiles newFileOrFiles(ReadRel.LocalFiles.FileOrFiles file) {
     return builder.build();
   }
 
+  /**
+   * Converts StructLiteral instances to NestedStruct for VirtualTableScan. This is a convenience
+   * method for migrating from the legacy StructLiteral-based VirtualTable API to the new
+   * NestedStruct-based API.
+   *
+   * @param nullable whether the resulting NestedStruct instances should be nullable
+   * @param structs the StructLiteral instances to convert
+   * @return a list of NestedStruct instances with the same field structure
+   */
+  private static List<Expression.NestedStruct> nestedStruct(
+      boolean nullable, Expression.StructLiteral... structs) {
+    List<Expression.NestedStruct> nestedStructs = new ArrayList<>();
+    for (Expression.StructLiteral struct : structs) {
+      nestedStructs.add(
+          Expression.NestedStruct.builder()
+              .nullable(nullable)
+              .addAllFields(struct.fields())
+              .build());
+    }
+    return nestedStructs;
+  }
+
   protected VirtualTableScan newVirtualTable(ReadRel rel) {
     ReadRel.VirtualTable virtualTable = rel.getVirtualTable();
+    // If both values and expressions are set, raise an error
+    if (virtualTable.getValuesCount() > 0 && virtualTable.getExpressionsCount() > 0) {
+      throw new IllegalArgumentException(
+          "VirtualTable cannot have both values and expressions set");
+    }
     NamedStruct virtualTableSchema = newNamedStruct(rel);
     ProtoExpressionConverter converter =
         new ProtoExpressionConverter(lookup, extensions, virtualTableSchema.struct(), this);
-    List<Expression.StructLiteral> structLiterals = new ArrayList<>(virtualTable.getValuesCount());
+
+    List<Expression.NestedStruct> expressions =
+        new ArrayList<>(virtualTable.getValuesCount() + virtualTable.getExpressionsCount());
+
+    // We cannot have a null row in VirtualTable, therefore we set the nullability to false
+    // nullability is also not supported at the Expression.Nested.Struct level
     for (io.substrait.proto.Expression.Literal.Struct struct : virtualTable.getValuesList()) {
-      structLiterals.add(
-          Expression.StructLiteral.builder()
-              .fields(
-                  struct.getFieldsList().stream()
-                      .map(converter::from)
-                      .collect(java.util.stream.Collectors.toList()))
-              .build());
+      expressions.addAll(nestedStruct(false, converter.from(struct)));
+    }
+
+    for (io.substrait.proto.Expression.Nested.Struct expr : virtualTable.getExpressionsList()) {
+      expressions.add(converter.from(expr));
     }
 
     ImmutableVirtualTableScan.Builder builder =
@@ -619,7 +649,7 @@ protected VirtualTableScan newVirtualTable(ReadRel rel) {
                     rel.hasBestEffortFilter() ? converter.from(rel.getBestEffortFilter()) : null))
             .filter(Optional.ofNullable(rel.hasFilter() ? converter.from(rel.getFilter()) : null))
             .initialSchema(NamedStruct.fromProto(rel.getBaseSchema(), protoTypeConverter))
-            .rows(structLiterals);
+            .rows(expressions);
 
     builder
         .commonExtension(optionalAdvancedExtension(rel.getCommon()))

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

@@ -56,6 +56,7 @@
 import io.substrait.relation.physical.TargetType;
 import io.substrait.type.proto.TypeProtoConverter;
 import io.substrait.util.EmptyVisitationContext;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -749,17 +750,21 @@ public Rel visit(Cross cross, EmptyVisitationContext context) throws RuntimeExce
   @Override
   public Rel visit(VirtualTableScan virtualTableScan, EmptyVisitationContext context)
       throws RuntimeException {
+    List<io.substrait.proto.Expression.Nested.Struct> structs = new ArrayList<>();
+    for (Expression.NestedStruct row : virtualTableScan.getRows()) {
+      structs.add(
+          io.substrait.proto.Expression.Nested.Struct.newBuilder()
+              .addAllFields(
+                  row.fields().stream()
+                      .map(this::toProto)
+                      .collect(java.util.stream.Collectors.toList()))
+              .build());
+    }
+
     ReadRel.Builder builder =
         ReadRel.newBuilder()
             .setCommon(common(virtualTableScan))
-            .setVirtualTable(
-                ReadRel.VirtualTable.newBuilder()
-                    .addAllValues(
-                        virtualTableScan.getRows().stream()
-                            .map(this::toProto)
-                            .map(t -> t.getLiteral().getStruct())
-                            .collect(Collectors.toList()))
-                    .build())
+            .setVirtualTable(ReadRel.VirtualTable.newBuilder().addAllExpressions(structs).build())
             .setBaseSchema(virtualTableScan.getInitialSchema().toProto(typeProtoConverter));
 
     virtualTableScan.getFilter().ifPresent(f -> builder.setFilter(toProto(f)));

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

@@ -5,15 +5,16 @@
 import io.substrait.type.TypeVisitor;
 import io.substrait.util.VisitationContext;
 import java.util.List;
+import java.util.Objects;
 import org.immutables.value.Value;
 
 @Value.Immutable
 public abstract class VirtualTableScan extends AbstractReadRel {
 
-  public abstract List<Expression.StructLiteral> getRows();
+  public abstract List<Expression.NestedStruct> getRows();
 
   /**
-   *
+   * Checks the following invariants when construction a VirtualTableScan
    *
    * <ul>
    *   <li>non-empty rowset
@@ -29,15 +30,28 @@ protected void check() {
 
     assert names.size()
         == NamedFieldCountingTypeVisitor.countNames(this.getInitialSchema().struct());
-    List<Expression.StructLiteral> rows = getRows();
-
-    assert rows.size() > 0
-        && names.stream().noneMatch(s -> s == null)
-        && rows.stream().noneMatch(r -> r == null)
+    List<Expression.NestedStruct> rows = getRows();
+
+    // At the PROTOBUF layer, the Nested.Struct message does not carry nullability information.
+    // Nullability is attached to the Nested message, which can contain a Nested.Struct.
+    // The NestedStruct POJO flattens the Nested and Nested.Struct messages together, allowing the
+    // nullability of a NestedStruct to be set directly.
+    //
+    // HOWEVER, the VirtualTable message contains a list of Nested.Struct messages, and as such
+    // the nullability cannot be set at the protobuf layer. To avoid users attaching meaningless
+    // nullability information in the POJOs, we restrict the nullability of NestedStructs to false
+    // when used in VirtualTableScans.
+    for (Expression.NestedStruct row : rows) {
+      assert !row.nullable();
+    }
+
+    assert !rows.isEmpty()
+        && names.stream().noneMatch(Objects::isNull)
+        && rows.stream().noneMatch(Objects::isNull)
         && rows.stream()
             .allMatch(r -> NamedFieldCountingTypeVisitor.countNames(r.getType()) == names.size());
 
-    for (Expression.StructLiteral row : rows) {
+    for (Expression.NestedStruct row : rows) {
       validateRowConformsToSchema(row);
     }
   }
@@ -48,10 +62,10 @@ protected void check() {
    * @param row the row to validate
    * @throws AssertionError if the row does not conform to the schema
    */
-  private void validateRowConformsToSchema(Expression.StructLiteral row) {
+  private void validateRowConformsToSchema(Expression.NestedStruct row) {
     Type.Struct schemaStruct = getInitialSchema().struct();
     List<Type> schemaFieldTypes = schemaStruct.fields();
-    List<Expression.Literal> rowFields = row.fields();
+    List<Expression> rowFields = row.fields();
 
     assert rowFields.size() == schemaFieldTypes.size()
         : String.format(