diff --git a/core/src/main/java/io/substrait/extension/AdvancedExtension.java b/core/src/main/java/io/substrait/extension/AdvancedExtension.java index bb4490efd..0d7278ef4 100644 --- a/core/src/main/java/io/substrait/extension/AdvancedExtension.java +++ b/core/src/main/java/io/substrait/extension/AdvancedExtension.java @@ -1,20 +1,21 @@ package io.substrait.extension; import io.substrait.relation.Extension; +import java.util.List; import java.util.Optional; import org.immutables.value.Value; @Value.Immutable public abstract class AdvancedExtension { - public abstract Optional getOptimization(); + public abstract List getOptimizations(); public abstract Optional 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())); + getOptimizations().forEach(e -> builder.addOptimization(e.toProto())); return builder.build(); } diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java index 69cb0cdf6..2d66c5c4c 100644 --- a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java +++ b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java @@ -658,15 +658,18 @@ private AdvancedExtension advancedExtension( if (advancedExtension.hasEnhancement()) { builder.enhancement(enhancementFromAdvancedExtension(advancedExtension.getEnhancement())); } - if (advancedExtension.hasOptimization()) { - builder.optimization(optimizationFromAdvancedExtension(advancedExtension.getOptimization())); - } + advancedExtension + .getOptimizationList() + .forEach( + optimization -> + builder.addOptimizations(optimizationFromAdvancedExtension(optimization))); + return builder.build(); } /** * Override to provide a custom converter for {@link - * io.substrait.proto.AdvancedExtension#getOptimization()} data + * io.substrait.proto.AdvancedExtension#getOptimizationList()} ()} data */ protected Extension.Optimization optimizationFromAdvancedExtension(com.google.protobuf.Any any) { return new EmptyOptimization(); diff --git a/core/src/main/java/io/substrait/relation/extensions/EmptyOptimization.java b/core/src/main/java/io/substrait/relation/extensions/EmptyOptimization.java index 74cc40ad4..46a11888d 100644 --- a/core/src/main/java/io/substrait/relation/extensions/EmptyOptimization.java +++ b/core/src/main/java/io/substrait/relation/extensions/EmptyOptimization.java @@ -1,11 +1,12 @@ package io.substrait.relation.extensions; import com.google.protobuf.Any; +import io.substrait.proto.AdvancedExtension; import io.substrait.relation.Extension; /** - * Default type to which {@link io.substrait.proto.AdvancedExtension#getOptimization()} data is - * converted to by the {@link io.substrait.relation.ProtoRelConverter} + * Default type to which {@link AdvancedExtension#getOptimizationList()} data is converted to by the + * {@link io.substrait.relation.ProtoRelConverter} */ public class EmptyOptimization implements Extension.Optimization { @Override diff --git a/core/src/test/java/io/substrait/relation/ProtoRelConverterTest.java b/core/src/test/java/io/substrait/relation/ProtoRelConverterTest.java index 1f16e29fa..51e0dfae3 100644 --- a/core/src/test/java/io/substrait/relation/ProtoRelConverterTest.java +++ b/core/src/test/java/io/substrait/relation/ProtoRelConverterTest.java @@ -36,12 +36,12 @@ Rel relWithExtension(AdvancedExtension advancedExtension) { Rel emptyAdvancedExtension = relWithExtension(AdvancedExtension.builder().build()); Rel advancedExtensionWithOptimization = - relWithExtension(AdvancedExtension.builder().optimization(OPTIMIZED).build()); + relWithExtension(AdvancedExtension.builder().addOptimizations(OPTIMIZED).build()); Rel advancedExtensionWithEnhancement = relWithExtension(AdvancedExtension.builder().enhancement(ENHANCED).build()); Rel advancedExtensionWithEnhancementAndOptimization = relWithExtension( - AdvancedExtension.builder().enhancement(ENHANCED).optimization(OPTIMIZED).build()); + AdvancedExtension.builder().enhancement(ENHANCED).addOptimizations(OPTIMIZED).build()); @Test void emptyAdvancedExtension() { diff --git a/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java index a99f0ede2..06ed71dec 100644 --- a/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java +++ b/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java @@ -52,7 +52,7 @@ public class ExtensionRoundtripTest extends TestBase { final AdvancedExtension commonExtension = AdvancedExtension.builder() .enhancement(new StringHolder("COMMON ENHANCEMENT")) - .optimization(new StringHolder("COMMON OPTIMIZATION")) + .addOptimizations(new StringHolder("COMMON OPTIMIZATION")) .build(); final StringHolder detail = new StringHolder("DETAIL"); @@ -60,7 +60,7 @@ public class ExtensionRoundtripTest extends TestBase { final AdvancedExtension relExtension = AdvancedExtension.builder() .enhancement(new StringHolder("REL ENHANCEMENT")) - .optimization(new StringHolder("REL OPTIMIZATION")) + .addOptimizations(new StringHolder("REL OPTIMIZATION")) .build(); @Override diff --git a/substrait b/substrait index fcf0f762a..a68c1ac62 160000 --- a/substrait +++ b/substrait @@ -1 +1 @@ -Subproject commit fcf0f762a3d39eabadf2f2b507ff803aadcb123e +Subproject commit a68c1ac62f92d703da624cb8ac0cef854dd2b35f