fix: optimization singular -> list

substrait-io · Jul 16, 2024 · 031dfd4 · 031dfd4
1 parent d6d24eb
commit 031dfd4
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git 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<Extension.Optimization> getOptimization();
+  public abstract List<Extension.Optimization> getOptimizations();
 
   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()));
+    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
@@ -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/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
@@ -52,15 +52,15 @@ 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");
 
   final AdvancedExtension relExtension =
       AdvancedExtension.builder()
           .enhancement(new StringHolder("REL ENHANCEMENT"))
-          .optimization(new StringHolder("REL OPTIMIZATION"))
+          .addOptimizations(new StringHolder("REL OPTIMIZATION"))
           .build();
 
   @Override