feat: google.protobuf.Any handling (#135)

* refactor: mark visitation method as overriden for consistency
* refactor: extract newFileOrFiles
* feat: support Advanced Extensions from RelCommon
* feat: support Leaf/Single/Multi Rel extensions
* feat: support ExtensionTable
* feat: support Advanced Extensions from Rels
* refactor: extension -> commonExtension
* feat: allow extension rels to derive their own record types
* refactor: make rel extension details non-optional
* feat: allow extension table to derive its schema
* test: verify roundtripping of google.protobuf.Any fields
* test: verify default ProtoRelConvert behaviour when Any data is present
* feat: allow users to provide custom ProtoRelConverter implementations
* refactor: make extension table details non-optional
* fix: handle Detail data in StringHolderHandlingProtoRelConverter
* docs: document google.protobuf.Any handling code
* feat: allow for easier extension of ProtoRelConverter
* docs: difference between Optimization and Enhancement
* docs: correct optimization docs
* refactor: getRelExtension -> getExtension
* refactor: extract ToProto interface
* refactor: split out extension from Rel
* refactor: generalExtension -> relExtension

BREAKING CHANGE: Use commonExtension to indicate that these extensions
are associated with the RelCommon message

BREAKING CHANGE: Use relExtension to indicate that these extensions
are associated directly with Rels

Author: vbarua

core/src/main/java/io/substrait/io/substrait/extension/AdvancedExtension.java

@@ -1,6 +1,24 @@
 package io.substrait.io.substrait.extension;
 
-public interface AdvancedExtension {
+import io.substrait.relation.Extension;
+import java.util.Optional;
+import org.immutables.value.Value;
 
-  io.substrait.proto.AdvancedExtension toProto();
+@Value.Immutable
+public abstract class AdvancedExtension {
+
+  public abstract Optional<Extension.Optimization> getOptimization();
+
+  public abstract Optional<Extension.Enhancement> getEnhancement();
+
+  public io.substrait.proto.AdvancedExtension toProto() {
+    var builder = io.substrait.proto.AdvancedExtension.newBuilder();
+    getEnhancement().ifPresent(e -> builder.setEnhancement(e.toProto()));
+    getOptimization().ifPresent(e -> builder.setOptimization(e.toProto()));
+    return builder.build();
+  }
+
+  public static ImmutableAdvancedExtension.Builder builder() {
+    return ImmutableAdvancedExtension.builder();
+  }
 }

core/src/main/java/io/substrait/plan/ProtoPlanConverter.java

@@ -15,7 +15,7 @@ public class ProtoPlanConverter {
   static final org.slf4j.Logger logger =
       org.slf4j.LoggerFactory.getLogger(io.substrait.plan.ProtoPlanConverter.class);
 
-  private final SimpleExtension.ExtensionCollection extensionCollection;
+  protected final SimpleExtension.ExtensionCollection extensionCollection;
 
   public ProtoPlanConverter() throws IOException {
     this(SimpleExtension.loadDefaults());
@@ -25,9 +25,14 @@ public ProtoPlanConverter(SimpleExtension.ExtensionCollection extensionCollectio
     this.extensionCollection = extensionCollection;
   }
 
+  /** Override hook for providing custom {@link ProtoRelConverter} implementations */
+  protected ProtoRelConverter getProtoRelConverter(FunctionLookup functionLookup) {
+    return new ProtoRelConverter(functionLookup, this.extensionCollection);
+  }
+
   public Plan from(io.substrait.proto.Plan plan) {
     FunctionLookup functionLookup = ImmutableFunctionLookup.builder().from(plan).build();
-    ProtoRelConverter relConverter = new ProtoRelConverter(functionLookup, extensionCollection);
+    ProtoRelConverter relConverter = getProtoRelConverter(functionLookup);
     List<Plan.Root> roots = new ArrayList<>();
     for (PlanRel planRel : plan.getRelationsList()) {
       io.substrait.proto.RelRoot root = planRel.getRoot();

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

@@ -1,12 +1,11 @@
 package io.substrait.relation;
 
 import io.substrait.expression.Expression;
-import io.substrait.io.substrait.extension.AdvancedExtension;
 import io.substrait.type.NamedStruct;
 import io.substrait.type.Type;
 import java.util.Optional;
 
-public abstract class AbstractReadRel extends ZeroInputRel {
+public abstract class AbstractReadRel extends ZeroInputRel implements HasExtension {
 
   public abstract NamedStruct getInitialSchema();
 
@@ -15,8 +14,6 @@ public abstract class AbstractReadRel extends ZeroInputRel {
   // TODO:
   // public abstract Optional<MaskExpression>
 
-  public abstract Optional<AdvancedExtension> getGeneralExtension();
-
   @Override
   protected final Type.Struct deriveRecordType() {
     return getInitialSchema().struct();

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

@@ -59,7 +59,28 @@ public OUTPUT visit(VirtualTableScan virtualTableScan) throws EXCEPTION {
     return visitFallback(virtualTableScan);
   }
 
+  @Override
   public OUTPUT visit(Cross cross) throws EXCEPTION {
     return visitFallback(cross);
   }
+
+  @Override
+  public OUTPUT visit(ExtensionLeaf extensionLeaf) throws EXCEPTION {
+    return visitFallback(extensionLeaf);
+  }
+
+  @Override
+  public OUTPUT visit(ExtensionSingle extensionSingle) throws EXCEPTION {
+    return visitFallback(extensionSingle);
+  }
+
+  @Override
+  public OUTPUT visit(ExtensionMulti extensionMulti) throws EXCEPTION {
+    return visitFallback(extensionMulti);
+  }
+
+  @Override
+  public OUTPUT visit(ExtensionTable extensionTable) throws EXCEPTION {
+    return visitFallback(extensionTable);
+  }
 }

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

@@ -12,7 +12,7 @@
 import org.immutables.value.Value;
 
 @Value.Immutable
-public abstract class Aggregate extends SingleInputRel {
+public abstract class Aggregate extends SingleInputRel implements HasExtension {
   static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(Aggregate.class);
 
   public abstract List<Grouping> getGroupings();

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

@@ -6,7 +6,7 @@
 import org.immutables.value.Value;
 
 @Value.Immutable
-public abstract class Cross extends BiRel {
+public abstract class Cross extends BiRel implements HasExtension {
 
   @Override
   protected Type.Struct deriveRecordType() {

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

@@ -0,0 +1,54 @@
+package io.substrait.relation;
+
+import io.substrait.type.NamedStruct;
+import io.substrait.type.Type;
+import java.util.List;
+
+/** Contains tag interfaces for handling {@link com.google.protobuf.Any} types within Substrait. */
+public class Extension {
+
+  /**
+   * Optimization associated with an {@link io.substrait.proto.AdvancedExtension}
+   *
+   * <p>An optimization is helpful information that don't influence semantics. May be ignored by a
+   * consumer.
+   */
+  public interface Optimization extends ToProto<com.google.protobuf.Any> {}
+
+  /**
+   * Enhancement associated with an {@link io.substrait.proto.AdvancedExtension}
+   *
+   * <p>An enhancement alter semantics. Cannot be ignored by a consumer.
+   */
+  public interface Enhancement extends ToProto<com.google.protobuf.Any> {}
+
+  public interface LeafRelDetail extends ToProto<com.google.protobuf.Any> {
+    /**
+     * @return the record layout for the associated {@link ExtensionLeaf} relation
+     */
+    Type.Struct deriveRecordType();
+  }
+
+  public interface SingleRelDetail extends ToProto<com.google.protobuf.Any> {
+    /**
+     * @param input to the associated {@link ExtensionSingle} relation
+     * @return the record layout for the associated {@link ExtensionSingle} relation
+     */
+    Type.Struct deriveRecordType(Rel input);
+  }
+
+  public interface MultiRelDetail extends ToProto<com.google.protobuf.Any> {
+    /**
+     * @param inputs to the associated {@link ExtensionMulti} relation
+     * @return the record layout for the associated {@link ExtensionMulti} relation
+     */
+    Type.Struct deriveRecordType(List<Rel> inputs);
+  }
+
+  public interface ExtensionTableDetail extends ToProto<com.google.protobuf.Any> {
+    /**
+     * @return the table schema for the associated {@link ExtensionTable} relation
+     */
+    NamedStruct deriveSchema();
+  }
+}

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

@@ -0,0 +1,20 @@
+package io.substrait.relation;
+
+import org.immutables.value.Value;
+
+@Value.Immutable
+public abstract class ExtensionLeaf extends ZeroInputRel {
+
+  public abstract Extension.LeafRelDetail getDetail();
+
+  @Override
+  public <O, E extends Exception> O accept(RelVisitor<O, E> visitor) throws E {
+    return visitor.visit(this);
+  }
+
+  public static ImmutableExtensionLeaf.Builder from(Extension.LeafRelDetail detail) {
+    return ImmutableExtensionLeaf.builder()
+        .detail(detail)
+        .deriveRecordType(detail.deriveRecordType());
+  }
+}

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

@@ -0,0 +1,30 @@
+package io.substrait.relation;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.immutables.value.Value;
+
+@Value.Immutable
+public abstract class ExtensionMulti extends AbstractRel {
+
+  public abstract Extension.MultiRelDetail getDetail();
+
+  @Override
+  public <O, E extends Exception> O accept(RelVisitor<O, E> visitor) throws E {
+    return visitor.visit(this);
+  }
+
+  public static ImmutableExtensionMulti.Builder from(
+      Extension.MultiRelDetail detail, Rel... inputs) {
+    return from(detail, Arrays.stream(inputs).collect(Collectors.toList()));
+  }
+
+  public static ImmutableExtensionMulti.Builder from(
+      Extension.MultiRelDetail detail, List<Rel> inputs) {
+    return ImmutableExtensionMulti.builder()
+        .addAllInputs(inputs)
+        .detail(detail)
+        .deriveRecordType(detail.deriveRecordType(inputs));
+  }
+}

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

@@ -0,0 +1,21 @@
+package io.substrait.relation;
+
+import org.immutables.value.Value;
+
+@Value.Immutable
+public abstract class ExtensionSingle extends SingleInputRel {
+
+  public abstract Extension.SingleRelDetail getDetail();
+
+  @Override
+  public <O, E extends Exception> O accept(RelVisitor<O, E> visitor) throws E {
+    return visitor.visit(this);
+  }
+
+  public static ImmutableExtensionSingle.Builder from(Extension.SingleRelDetail detail, Rel input) {
+    return ImmutableExtensionSingle.builder()
+        .input(input)
+        .detail(detail)
+        .deriveRecordType(detail.deriveRecordType(input));
+  }
+}