diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 0f58da6af..788bc3fb9 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -16,7 +16,9 @@ import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.SimpleExtension; import io.substrait.function.ToTypeString; +import io.substrait.plan.ImmutableExecutionBehavior; import io.substrait.plan.Plan; +import io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode; import io.substrait.relation.AbstractWriteRel; import io.substrait.relation.Aggregate; import io.substrait.relation.Aggregate.Measure; @@ -1523,13 +1525,54 @@ public Plan.Root root(Rel rel) { } /** - * Creates a plan from a plan root. + * Creates a plan from a plan root with default execution behavior. + * + *

The plan is created with {@link VariableEvaluationMode#VARIABLE_EVALUATION_MODE_PER_PLAN} as + * the default variable evaluation mode. To specify a custom execution behavior, use {@link + * #plan(Plan.ExecutionBehavior, Plan.Root)} instead. * * @param root the plan root * @return a new {@link Plan} */ public Plan plan(Plan.Root root) { - return Plan.builder().addRoots(root).build(); + return plan( + ImmutableExecutionBehavior.builder() + .variableEvaluationMode(VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(), + root); + } + + /** + * Creates a plan from a plan root with custom execution behavior. + * + * @param executionBehavior the execution behavior for the plan + * @param root the plan root + * @return a new {@link Plan} + */ + public Plan plan(Plan.ExecutionBehavior executionBehavior, Plan.Root root) { + return Plan.builder().executionBehavior(executionBehavior).addRoots(root).build(); + } + + /** + * Creates a plan from multiple plan roots with custom execution behavior. + * + * @param executionBehavior the execution behavior for the plan + * @param roots the plan roots + * @return a new {@link Plan} + */ + public Plan plan(Plan.ExecutionBehavior executionBehavior, Plan.Root... roots) { + return Plan.builder().executionBehavior(executionBehavior).roots(Arrays.asList(roots)).build(); + } + + /** + * Creates a plan from multiple plan roots with custom execution behavior. + * + * @param executionBehavior the execution behavior for the plan + * @param roots the plan roots as an iterable + * @return a new {@link Plan} + */ + public Plan plan(Plan.ExecutionBehavior executionBehavior, Iterable roots) { + return Plan.builder().executionBehavior(executionBehavior).roots(roots).build(); } /** diff --git a/core/src/main/java/io/substrait/plan/Plan.java b/core/src/main/java/io/substrait/plan/Plan.java index 6436993f4..a3f897632 100644 --- a/core/src/main/java/io/substrait/plan/Plan.java +++ b/core/src/main/java/io/substrait/plan/Plan.java @@ -21,6 +21,38 @@ public Version getVersion() { public abstract Optional getAdvancedExtension(); + public abstract Optional getExecutionBehavior(); + + /** + * Validates that the execution behavior is properly configured. + * + *

This validation method ensures that: + * + *

+ * + * @throws IllegalArgumentException if the execution behavior is not present, or if the variable + * evaluation mode is set to {@link + * ExecutionBehavior.VariableEvaluationMode#VARIABLE_EVALUATION_MODE_UNSPECIFIED} + */ + @Value.Check + protected void check() { + if (!getExecutionBehavior().isPresent()) { + throw new IllegalArgumentException("ExecutionBehavior is required but was not set"); + } + ExecutionBehavior behavior = getExecutionBehavior().get(); + if (behavior.getVariableEvaluationMode() + == ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_UNSPECIFIED) { + throw new IllegalArgumentException( + "ExecutionBehavior requires a specified VariableEvaluationMode, but got: " + + behavior.getVariableEvaluationMode()); + } + } + public static ImmutablePlan.Builder builder() { return ImmutablePlan.builder(); } @@ -69,4 +101,19 @@ public static ImmutableRoot.Builder builder() { return ImmutableRoot.builder(); } } + + @Value.Immutable + public abstract static class ExecutionBehavior { + public abstract VariableEvaluationMode getVariableEvaluationMode(); + + public static ImmutableExecutionBehavior.Builder builder() { + return ImmutableExecutionBehavior.builder(); + } + + public enum VariableEvaluationMode { + VARIABLE_EVALUATION_MODE_UNSPECIFIED, + VARIABLE_EVALUATION_MODE_PER_PLAN, + VARIABLE_EVALUATION_MODE_PER_RECORD + } + } } diff --git a/core/src/main/java/io/substrait/plan/PlanProtoConverter.java b/core/src/main/java/io/substrait/plan/PlanProtoConverter.java index 03b2a4227..d463999a8 100644 --- a/core/src/main/java/io/substrait/plan/PlanProtoConverter.java +++ b/core/src/main/java/io/substrait/plan/PlanProtoConverter.java @@ -4,6 +4,8 @@ import io.substrait.extension.ExtensionCollector; import io.substrait.extension.ExtensionProtoConverter; import io.substrait.extension.SimpleExtension.ExtensionCollection; +import io.substrait.proto.ExecutionBehavior; +import io.substrait.proto.ExecutionBehavior.VariableEvaluationMode; import io.substrait.proto.Plan; import io.substrait.proto.PlanRel; import io.substrait.proto.Rel; @@ -61,6 +63,26 @@ public PlanProtoConverter( this.extensionProtoConverter = extensionProtoConverter; } + /** + * Converts a {@link io.substrait.plan.Plan} object to its protobuf representation. + * + *

This method performs a complete conversion of the Plan object, including: + * + *

+ * + *

The execution behavior is optional. If present, it will be converted using {@link + * #toProtoExecutionBehavior(io.substrait.plan.Plan.ExecutionBehavior)}. + * + * @param plan the Plan object to convert, must not be null + * @return the protobuf Plan representation + * @throws IllegalArgumentException if the plan contains invalid data + */ public Plan toProto(final io.substrait.plan.Plan plan) { final List planRels = new ArrayList<>(); final ExtensionCollector functionCollector = new ExtensionCollector(extensionCollection); @@ -97,6 +119,65 @@ public Plan toProto(final io.substrait.plan.Plan plan) { builder.setVersion(versionBuilder); + // Set execution behavior if present + if (plan.getExecutionBehavior().isPresent()) { + builder.setExecutionBehavior(toProtoExecutionBehavior(plan.getExecutionBehavior().get())); + } + return builder.build(); } + + /** + * Converts an {@link io.substrait.plan.Plan.ExecutionBehavior} to its protobuf representation. + * + *

This method converts the execution behavior configuration, including the variable evaluation + * mode, from the POJO representation to the protobuf format. + * + * @param executionBehavior the ExecutionBehavior to convert, must not be null + * @return the protobuf ExecutionBehavior representation + * @throws IllegalArgumentException if the variable evaluation mode is unknown + */ + private ExecutionBehavior toProtoExecutionBehavior( + final io.substrait.plan.Plan.ExecutionBehavior executionBehavior) { + return ExecutionBehavior.newBuilder() + .setVariableEvalMode( + toProtoVariableEvaluationMode(executionBehavior.getVariableEvaluationMode())) + .build(); + } + + /** + * Converts a {@link io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode} to its + * protobuf representation. + * + *

Supported modes: + * + *

+ * + * @param mode the VariableEvaluationMode to convert, must not be null + * @return the protobuf VariableEvaluationMode representation + * @throws IllegalArgumentException if the mode is unknown or not supported + */ + private VariableEvaluationMode toProtoVariableEvaluationMode( + final io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode mode) { + switch (mode) { + case VARIABLE_EVALUATION_MODE_UNSPECIFIED: + return VariableEvaluationMode.VARIABLE_EVALUATION_MODE_UNSPECIFIED; + case VARIABLE_EVALUATION_MODE_PER_PLAN: + return VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN; + case VARIABLE_EVALUATION_MODE_PER_RECORD: + return VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_RECORD; + default: + throw new IllegalArgumentException("Unknown VariableEvaluationMode: " + mode); + } + } } diff --git a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java index 2e5334b3e..cb719b6c9 100644 --- a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java +++ b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java @@ -66,6 +66,29 @@ protected ProtoRelConverter getProtoRelConverter(final ExtensionLookup functionL return new ProtoRelConverter(functionLookup, this.extensionCollection, protoExtensionConverter); } + /** + * Converts a protobuf {@link io.substrait.proto.Plan} to a {@link Plan} POJO. + * + *

This method performs a complete conversion from the protobuf representation, including: + * + *

+ * + *

Note: While execution behavior is optional in the protobuf message, the Plan + * validation will still fail if ExecutionBehavior is not set or if it contains + * VARIABLE_EVALUATION_MODE_UNSPECIFIED. Ensure the protobuf Plan has a valid execution behavior + * before conversion. + * + * @param plan the protobuf Plan to convert, must not be null + * @return the converted Plan POJO + * @throws IllegalArgumentException if the plan contains invalid data or if the execution behavior + * validation fails + */ public Plan from(io.substrait.proto.Plan plan) { ExtensionLookup functionLookup = ImmutableExtensionLookup.builder().from(plan).build(); ProtoRelConverter relConverter = getProtoRelConverter(functionLookup); @@ -92,15 +115,81 @@ public Plan from(io.substrait.proto.Plan plan) { versionBuilder.producer(Optional.of(plan.getVersion().getProducer())); } - return Plan.builder() - .roots(roots) - .expectedTypeUrls(plan.getExpectedTypeUrlsList()) - .advancedExtension( - Optional.ofNullable( - plan.hasAdvancedExtensions() - ? protoExtensionConverter.fromProto(plan.getAdvancedExtensions()) - : null)) - .version(versionBuilder.build()) + ImmutablePlan.Builder planBuilder = + Plan.builder() + .roots(roots) + .expectedTypeUrls(plan.getExpectedTypeUrlsList()) + .advancedExtension( + Optional.ofNullable( + plan.hasAdvancedExtensions() + ? protoExtensionConverter.fromProto(plan.getAdvancedExtensions()) + : null)) + .version(versionBuilder.build()); + + // Set execution behavior if present + if (plan.hasExecutionBehavior()) { + planBuilder.executionBehavior(fromProtoExecutionBehavior(plan.getExecutionBehavior())); + } + + return planBuilder.build(); + } + + /** + * Converts a protobuf {@link io.substrait.proto.ExecutionBehavior} to its POJO representation. + * + *

This method converts the execution behavior configuration from the protobuf format to the + * POJO representation, including the variable evaluation mode. + * + * @param executionBehavior the protobuf ExecutionBehavior to convert, must not be null + * @return the POJO ExecutionBehavior representation + * @throws IllegalArgumentException if the variable evaluation mode is unknown or UNRECOGNIZED + */ + private io.substrait.plan.Plan.ExecutionBehavior fromProtoExecutionBehavior( + final io.substrait.proto.ExecutionBehavior executionBehavior) { + return io.substrait.plan.Plan.ExecutionBehavior.builder() + .variableEvaluationMode( + fromProtoVariableEvaluationMode(executionBehavior.getVariableEvalMode())) .build(); } + + /** + * Converts a protobuf {@link io.substrait.proto.ExecutionBehavior.VariableEvaluationMode} to its + * POJO representation. + * + *

Supported modes: + * + *

+ * + * @param mode the protobuf VariableEvaluationMode to convert, must not be null + * @return the POJO VariableEvaluationMode representation + * @throws IllegalArgumentException if the mode is UNRECOGNIZED or not supported + */ + private io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode + fromProtoVariableEvaluationMode( + final io.substrait.proto.ExecutionBehavior.VariableEvaluationMode mode) { + switch (mode) { + case VARIABLE_EVALUATION_MODE_UNSPECIFIED: + return io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_UNSPECIFIED; + case VARIABLE_EVALUATION_MODE_PER_PLAN: + return io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_PER_PLAN; + case VARIABLE_EVALUATION_MODE_PER_RECORD: + return io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_PER_RECORD; + case UNRECOGNIZED: + default: + throw new IllegalArgumentException("Unknown VariableEvaluationMode: " + mode); + } + } } diff --git a/core/src/test/java/io/substrait/plan/PlanConverterTest.java b/core/src/test/java/io/substrait/plan/PlanConverterTest.java index ab0a321b4..0c2a7f9ba 100644 --- a/core/src/test/java/io/substrait/plan/PlanConverterTest.java +++ b/core/src/test/java/io/substrait/plan/PlanConverterTest.java @@ -1,7 +1,11 @@ package io.substrait.plan; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import io.substrait.expression.Expression; import io.substrait.expression.ExpressionCreator; @@ -21,13 +25,25 @@ import org.junit.jupiter.api.Test; class PlanConverterTest { + private final PlanProtoConverter toProtoConverter = new PlanProtoConverter(); + private final ProtoPlanConverter fromProtoConverter = new ProtoPlanConverter(); + + private static Plan.ExecutionBehavior defaultExecutionBehavior() { + return ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(); + } + @Test void emptyAdvancedExtensionTest() { - final Plan plan = Plan.builder().advancedExtension(AdvancedExtension.builder().build()).build(); - final PlanProtoConverter toProtoConverter = new PlanProtoConverter(); + final Plan plan = + Plan.builder() + .executionBehavior(defaultExecutionBehavior()) + .advancedExtension(AdvancedExtension.builder().build()) + .build(); final io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(plan); - final ProtoPlanConverter fromProtoConverter = new ProtoPlanConverter(); final Plan plan2 = fromProtoConverter.from(protoPlan); assertEquals(plan, plan2); @@ -39,9 +55,9 @@ void enhancementOnlyAdvancedExtensionWithoutExtensionProtoConverter() { final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension(AdvancedExtension.builder().enhancement(enhanced).build()) .build(); - final PlanProtoConverter toProtoConverter = new PlanProtoConverter(); assertThrows( UnsupportedOperationException.class, @@ -55,13 +71,13 @@ void enhancementOnlyAdvancedExtensionWithExtensionProtoConverter() { final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension(AdvancedExtension.builder().enhancement(enhanced).build()) .build(); final PlanProtoConverter toProtoConverter = new PlanProtoConverter(new StringHolderHandlingExtensionProtoConverter()); final io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(plan); - final ProtoPlanConverter fromProtoConverter = new ProtoPlanConverter(); assertThrows( UnsupportedOperationException.class, () -> fromProtoConverter.from(protoPlan), @@ -74,6 +90,7 @@ void enhancementOnlyAdvancedExtensionWithExtensionProtoConverterAndProtoExtensio final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension(AdvancedExtension.builder().enhancement(enhanced).build()) .build(); final PlanProtoConverter toProtoConverter = @@ -93,9 +110,9 @@ void optimizationOnlyAdvancedExtensionWithoutExtensionProtoConverter() { final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension(AdvancedExtension.builder().addOptimizations(optimized).build()) .build(); - final PlanProtoConverter toProtoConverter = new PlanProtoConverter(); assertThrows( UnsupportedOperationException.class, @@ -109,13 +126,13 @@ void optimizationOnlyAdvancedExtensionWithExtensionProtoConverter() { final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension(AdvancedExtension.builder().addOptimizations(optimized).build()) .build(); final PlanProtoConverter toProtoConverter = new PlanProtoConverter(new StringHolderHandlingExtensionProtoConverter()); final io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(plan); - final ProtoPlanConverter fromProtoConverter = new ProtoPlanConverter(); assertThrows( UnsupportedOperationException.class, () -> fromProtoConverter.from(protoPlan), @@ -128,6 +145,7 @@ void optimizationOnlyAdvancedExtensionWithExtensionProtoConverterAndProtoExtensi final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension(AdvancedExtension.builder().addOptimizations(optimized).build()) .build(); final PlanProtoConverter toProtoConverter = @@ -148,6 +166,7 @@ void advancedExtensionWithEnhancementAndOptimization() { final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .advancedExtension( AdvancedExtension.builder() .enhancement(enhanced) @@ -172,6 +191,7 @@ void planIncludingRelationWithAdvancedExtension() { final Plan plan = Plan.builder() + .executionBehavior(defaultExecutionBehavior()) .addRoots( Root.builder() .input( @@ -304,9 +324,12 @@ void nestedUserDefinedTypesShareExtensionCollector() { false, nullablePointLiteral, pointLiteral, vectorOfPointLiteral)) .build(); - Plan plan = Plan.builder().addRoots(Root.builder().input(virtualTable).build()).build(); + Plan plan = + Plan.builder() + .executionBehavior(defaultExecutionBehavior()) + .addRoots(Root.builder().input(virtualTable).build()) + .build(); - PlanProtoConverter toProtoConverter = new PlanProtoConverter(); io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(plan); assertEquals(1, protoPlan.getExtensionUrnsCount(), "Should have exactly 1 extension URN"); @@ -319,4 +342,316 @@ void nestedUserDefinedTypesShareExtensionCollector() { Plan roundTrippedPlan = fromProtoConverter.from(protoPlan); assertEquals(plan, roundTrippedPlan, "Plan should roundtrip correctly"); } + + /** + * Conversion of Plan with VARIABLE_EVALUATION_MODE_PER_PLAN to protobuf. + * + *

Verifies that a Plan with ExecutionBehavior set to PER_PLAN mode is correctly converted to + * protobuf format, including the execution behavior field. + */ + @Test + void testToProtoWithExecutionBehaviorPerPlan() { + // Create a Plan with ExecutionBehavior set to PER_PLAN + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(); + + Plan plan = + ImmutablePlan.builder() + .executionBehavior(executionBehavior) + .roots(Collections.emptyList()) + .expectedTypeUrls(Collections.emptyList()) + .build(); + + // Convert to protobuf + io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(plan); + + // Verify the protobuf has execution behavior + assertTrue(protoPlan.hasExecutionBehavior(), "Protobuf Plan should have ExecutionBehavior"); + assertEquals( + io.substrait.proto.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_PER_PLAN, + protoPlan.getExecutionBehavior().getVariableEvalMode(), + "Variable evaluation mode should be PER_PLAN"); + } + + /** + * Conversion of Plan with VARIABLE_EVALUATION_MODE_PER_RECORD to protobuf. + * + *

Verifies that a Plan with ExecutionBehavior set to PER_RECORD mode is correctly converted to + * protobuf format. + */ + @Test + void testToProtoWithExecutionBehaviorPerRecord() { + // Create a Plan with ExecutionBehavior set to PER_RECORD + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_RECORD) + .build(); + + Plan plan = + ImmutablePlan.builder() + .executionBehavior(executionBehavior) + .roots(Collections.emptyList()) + .expectedTypeUrls(Collections.emptyList()) + .build(); + + // Convert to protobuf + io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(plan); + + // Verify the protobuf has execution behavior + assertTrue(protoPlan.hasExecutionBehavior(), "Protobuf Plan should have ExecutionBehavior"); + assertEquals( + io.substrait.proto.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_PER_RECORD, + protoPlan.getExecutionBehavior().getVariableEvalMode(), + "Variable evaluation mode should be PER_RECORD"); + } + + /** + * Conversion from protobuf Plan with ExecutionBehavior to POJO. + * + *

Verifies that a protobuf Plan with ExecutionBehavior is correctly converted to the POJO + * representation. + */ + @Test + void testFromProtoWithExecutionBehaviorPerPlan() { + // Create a protobuf Plan with ExecutionBehavior + io.substrait.proto.Plan protoPlan = + io.substrait.proto.Plan.newBuilder() + .setExecutionBehavior( + io.substrait.proto.ExecutionBehavior.newBuilder() + .setVariableEvalMode( + io.substrait.proto.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_PER_PLAN) + .build()) + .build(); + + // Convert to POJO + Plan plan = fromProtoConverter.from(protoPlan); + + // Verify the POJO has execution behavior + assertTrue( + plan.getExecutionBehavior().isPresent(), "Plan should have ExecutionBehavior present"); + assertEquals( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN, + plan.getExecutionBehavior().get().getVariableEvaluationMode(), + "Variable evaluation mode should be PER_PLAN"); + } + + /** + * Conversion from protobuf Plan with PER_RECORD mode to POJO. + * + *

Verifies that a protobuf Plan with PER_RECORD execution behavior is correctly converted. + */ + @Test + void testFromProtoWithExecutionBehaviorPerRecord() { + // Create a protobuf Plan with ExecutionBehavior set to PER_RECORD + io.substrait.proto.Plan protoPlan = + io.substrait.proto.Plan.newBuilder() + .setExecutionBehavior( + io.substrait.proto.ExecutionBehavior.newBuilder() + .setVariableEvalMode( + io.substrait.proto.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_PER_RECORD) + .build()) + .build(); + + // Convert to POJO + Plan plan = fromProtoConverter.from(protoPlan); + + // Verify the POJO has execution behavior + assertTrue( + plan.getExecutionBehavior().isPresent(), "Plan should have ExecutionBehavior present"); + assertEquals( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_RECORD, + plan.getExecutionBehavior().get().getVariableEvaluationMode(), + "Variable evaluation mode should be PER_RECORD"); + } + + /** + * Conversion from protobuf Plan without ExecutionBehavior. + * + *

Verifies that a protobuf Plan without ExecutionBehavior results in a POJO Plan that fails + * validation (since ExecutionBehavior is required). + */ + @Test + void testFromProtoWithoutExecutionBehaviorFailsValidation() { + // Create a protobuf Plan without ExecutionBehavior + io.substrait.proto.Plan protoPlan = io.substrait.proto.Plan.newBuilder().build(); + + // Attempt to convert to POJO - should fail validation + assertThrows( + IllegalArgumentException.class, + () -> fromProtoConverter.from(protoPlan), + "Conversion should fail when ExecutionBehavior is not set"); + } + + /** + * Test case 6: Conversion from protobuf Plan with UNSPECIFIED mode fails validation. + * + *

Verifies that a protobuf Plan with VARIABLE_EVALUATION_MODE_UNSPECIFIED results in a + * validation failure when converted to POJO. + */ + @Test + void testFromProtoWithUnspecifiedModeFailsValidation() { + // Create a protobuf Plan with UNSPECIFIED mode + io.substrait.proto.Plan protoPlan = + io.substrait.proto.Plan.newBuilder() + .setExecutionBehavior( + io.substrait.proto.ExecutionBehavior.newBuilder() + .setVariableEvalMode( + io.substrait.proto.ExecutionBehavior.VariableEvaluationMode + .VARIABLE_EVALUATION_MODE_UNSPECIFIED) + .build()) + .build(); + + // Attempt to convert to POJO - should fail validation + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> fromProtoConverter.from(protoPlan), + "Conversion should fail when VariableEvaluationMode is UNSPECIFIED"); + + assertTrue( + exception.getMessage().contains("VariableEvaluationMode"), + "Error message should mention VariableEvaluationMode"); + } + + /** + * Round-trip conversion with PER_PLAN mode. + * + *

Verifies that converting a Plan to protobuf and back preserves all data, including the + * execution behavior with PER_PLAN mode. + */ + @Test + void testRoundTripWithExecutionBehaviorPerPlan() { + // Create original Plan + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(); + + Plan originalPlan = + ImmutablePlan.builder() + .executionBehavior(executionBehavior) + .roots(Collections.emptyList()) + .expectedTypeUrls(Collections.emptyList()) + .build(); + + // Convert to protobuf and back + io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(originalPlan); + Plan roundTrippedPlan = fromProtoConverter.from(protoPlan); + + // Verify data integrity + assertTrue( + roundTrippedPlan.getExecutionBehavior().isPresent(), + "Round-tripped Plan should have ExecutionBehavior"); + assertEquals( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN, + roundTrippedPlan.getExecutionBehavior().get().getVariableEvaluationMode(), + "Variable evaluation mode should be preserved"); + assertEquals( + originalPlan.getRoots().size(), + roundTrippedPlan.getRoots().size(), + "Number of roots should be preserved"); + assertEquals( + originalPlan.getExpectedTypeUrls().size(), + roundTrippedPlan.getExpectedTypeUrls().size(), + "Number of expected type URLs should be preserved"); + } + + /** + * Round-trip conversion with PER_RECORD mode. + * + *

Verifies that converting a Plan with PER_RECORD mode to protobuf and back preserves the + * execution behavior correctly. + */ + @Test + void testRoundTripWithExecutionBehaviorPerRecord() { + // Create original Plan + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_RECORD) + .build(); + + Plan originalPlan = + ImmutablePlan.builder() + .executionBehavior(executionBehavior) + .roots(Collections.emptyList()) + .expectedTypeUrls(Collections.emptyList()) + .build(); + + // Convert to protobuf and back + io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(originalPlan); + Plan roundTrippedPlan = fromProtoConverter.from(protoPlan); + + // Verify data integrity + assertTrue( + roundTrippedPlan.getExecutionBehavior().isPresent(), + "Round-tripped Plan should have ExecutionBehavior"); + assertEquals( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_RECORD, + roundTrippedPlan.getExecutionBehavior().get().getVariableEvaluationMode(), + "Variable evaluation mode should be preserved"); + } + + /** + * Empty plan with only execution behavior. + * + *

Verifies that a minimal Plan with only execution behavior (no roots, no type URLs) can be + * successfully converted. + */ + @Test + void testRoundTripEmptyPlanWithExecutionBehavior() { + // Create minimal Plan + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(); + + Plan originalPlan = + ImmutablePlan.builder() + .executionBehavior(executionBehavior) + .roots(Collections.emptyList()) + .expectedTypeUrls(Collections.emptyList()) + .build(); + + // Verify conversion succeeds + assertDoesNotThrow( + () -> { + io.substrait.proto.Plan protoPlan = toProtoConverter.toProto(originalPlan); + Plan roundTrippedPlan = fromProtoConverter.from(protoPlan); + assertNotNull(roundTrippedPlan, "Round-tripped Plan should not be null"); + }, + "Empty plan with execution behavior should convert successfully"); + } + + /** + * Verify that protobuf without execution behavior field is handled. + * + *

Tests the edge case where a protobuf Plan doesn't have the execution behavior field set at + * all. + */ + @Test + void testFromProtoMissingExecutionBehaviorField() { + // Create protobuf Plan without setting execution behavior + io.substrait.proto.Plan protoPlan = io.substrait.proto.Plan.newBuilder().build(); + + // Verify hasExecutionBehavior returns false + assertFalse( + protoPlan.hasExecutionBehavior(), "Protobuf Plan should not have execution behavior field"); + + // Conversion should fail validation + assertThrows( + IllegalArgumentException.class, + () -> fromProtoConverter.from(protoPlan), + "Conversion should fail when execution behavior is missing"); + } } diff --git a/core/src/test/java/io/substrait/plan/PlanValidationTest.java b/core/src/test/java/io/substrait/plan/PlanValidationTest.java new file mode 100644 index 000000000..906074ed6 --- /dev/null +++ b/core/src/test/java/io/substrait/plan/PlanValidationTest.java @@ -0,0 +1,120 @@ +package io.substrait.plan; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.Test; + +/** + * Test cases for the Plan validation method that ensures ExecutionBehavior is properly configured. + */ +class PlanValidationTest { + + /** + * Test case 1: Valid execution behavior with proper variable evaluation mode. + * + *

This test verifies that a Plan with a properly configured ExecutionBehavior (with + * VariableEvaluationMode set to a valid value other than UNSPECIFIED) is successfully created + * without throwing any exceptions. + */ + @Test + void testValidExecutionBehavior_PerPlan() { + // Create a valid ExecutionBehavior with VARIABLE_EVALUATION_MODE_PER_PLAN + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(); + + // Create a Plan with the valid ExecutionBehavior + assertDoesNotThrow( + () -> { + ImmutablePlan.builder().executionBehavior(executionBehavior).build(); + }, + "Plan creation should succeed with valid ExecutionBehavior"); + } + + /** + * Test case 1b: Valid execution behavior with VARIABLE_EVALUATION_MODE_PER_RECORD. + * + *

This test verifies that a Plan with ExecutionBehavior set to + * VARIABLE_EVALUATION_MODE_PER_RECORD is also valid. + */ + @Test + void testValidExecutionBehavior_PerRecord() { + // Create a valid ExecutionBehavior with VARIABLE_EVALUATION_MODE_PER_RECORD + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_RECORD) + .build(); + + // Create a Plan with the valid ExecutionBehavior + assertDoesNotThrow( + () -> { + ImmutablePlan.builder().executionBehavior(executionBehavior).build(); + }, + "Plan creation should succeed with valid ExecutionBehavior"); + } + + /** + * Test case 2: Missing execution behavior (empty Optional). + * + *

This test verifies that attempting to create a Plan without setting the ExecutionBehavior + * field throws an IllegalArgumentException with an appropriate error message indicating that + * ExecutionBehavior is required. + */ + @Test + void testMissingExecutionBehavior() { + // Attempt to create a Plan without ExecutionBehavior + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> { + ImmutablePlan.builder().build(); + }, + "Plan creation should fail when ExecutionBehavior is not set"); + + // Verify the error message + assertEquals( + "ExecutionBehavior is required but was not set", + exception.getMessage(), + "Error message should indicate ExecutionBehavior is required"); + } + + /** + * Test case 3: Execution behavior with unspecified variable evaluation mode. + * + *

This test verifies that attempting to create a Plan with an ExecutionBehavior that has its + * VariableEvaluationMode set to VARIABLE_EVALUATION_MODE_UNSPECIFIED throws an + * IllegalArgumentException with an appropriate error message. + */ + @Test + void testUnspecifiedVariableEvaluationMode() { + // Create an ExecutionBehavior with VARIABLE_EVALUATION_MODE_UNSPECIFIED + Plan.ExecutionBehavior executionBehavior = + ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_UNSPECIFIED) + .build(); + + // Attempt to create a Plan with the invalid ExecutionBehavior + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> { + ImmutablePlan.builder().executionBehavior(executionBehavior).build(); + }, + "Plan creation should fail when VariableEvaluationMode is UNSPECIFIED"); + + // Verify the error message contains the expected information + String expectedMessage = + "ExecutionBehavior requires a specified VariableEvaluationMode, but got: " + + "VARIABLE_EVALUATION_MODE_UNSPECIFIED"; + assertEquals( + expectedMessage, + exception.getMessage(), + "Error message should indicate VariableEvaluationMode cannot be UNSPECIFIED"); + } +} diff --git a/core/src/test/resources/plan-roundtrip/complex-expected-plan.json b/core/src/test/resources/plan-roundtrip/complex-expected-plan.json index 9f4c47a18..99125c4f9 100644 --- a/core/src/test/resources/plan-roundtrip/complex-expected-plan.json +++ b/core/src/test/resources/plan-roundtrip/complex-expected-plan.json @@ -148,6 +148,9 @@ } ], "version": { - "minorNumber": 75 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" } } diff --git a/core/src/test/resources/plan-roundtrip/complex-input-plan.json b/core/src/test/resources/plan-roundtrip/complex-input-plan.json index 17a1984bd..682c296cc 100644 --- a/core/src/test/resources/plan-roundtrip/complex-input-plan.json +++ b/core/src/test/resources/plan-roundtrip/complex-input-plan.json @@ -133,6 +133,9 @@ } ], "version": { - "minorNumber": 75 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" } } diff --git a/core/src/test/resources/plan-roundtrip/simple-expected-plan.json b/core/src/test/resources/plan-roundtrip/simple-expected-plan.json index 5ae8daa3f..2edc69c1d 100644 --- a/core/src/test/resources/plan-roundtrip/simple-expected-plan.json +++ b/core/src/test/resources/plan-roundtrip/simple-expected-plan.json @@ -120,6 +120,9 @@ } ], "version": { - "minorNumber": 75 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" } } diff --git a/core/src/test/resources/plan-roundtrip/simple-input-plan.json b/core/src/test/resources/plan-roundtrip/simple-input-plan.json index b5f5af06e..845b68dab 100644 --- a/core/src/test/resources/plan-roundtrip/simple-input-plan.json +++ b/core/src/test/resources/plan-roundtrip/simple-input-plan.json @@ -105,6 +105,9 @@ } ], "version": { - "minorNumber": 75 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" } } diff --git a/core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json b/core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json index ff6b8ab0e..b60e6dd7b 100644 --- a/core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json +++ b/core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json @@ -92,6 +92,9 @@ } ], "version": { - "minorNumber": 75 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" } } diff --git a/core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json b/core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json index 202814110..d8300a067 100644 --- a/core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json +++ b/core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json @@ -79,6 +79,9 @@ } ], "version": { - "minorNumber": 75 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java index b05d864bf..1f7c60222 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java +++ b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java @@ -12,6 +12,8 @@ import io.substrait.isthmus.expression.SqlArrayValueConstructorCallConverter; import io.substrait.isthmus.expression.SqlMapValueConstructorCallConverter; import io.substrait.isthmus.expression.WindowFunctionConverter; +import io.substrait.plan.ImmutableExecutionBehavior; +import io.substrait.plan.Plan; import io.substrait.relation.Rel; import java.util.ArrayList; import java.util.List; @@ -64,6 +66,8 @@ public class ConverterProvider { /** Converter for Substrait types to Calcite types and vice versa. */ protected TypeConverter typeConverter; + protected final Plan.ExecutionBehavior executionBehavior; + /** * Creates a ConverterProvider with default extension collection and type factory. Uses {@link * DefaultExtensionCatalog#DEFAULT_COLLECTION} and {@link SubstraitTypeSystem#TYPE_FACTORY}. @@ -115,12 +119,36 @@ public ConverterProvider( AggregateFunctionConverter afc, WindowFunctionConverter wfc, TypeConverter tc) { + this(typeFactory, extensions, sfc, afc, wfc, tc, createDefaultExecutionBehavior()); + } + + public ConverterProvider( + RelDataTypeFactory typeFactory, + SimpleExtension.ExtensionCollection extensions, + ScalarFunctionConverter sfc, + AggregateFunctionConverter afc, + WindowFunctionConverter wfc, + TypeConverter tc, + Plan.ExecutionBehavior executionBehavior) { this.typeFactory = typeFactory; this.extensions = extensions; this.scalarFunctionConverter = sfc; this.aggregateFunctionConverter = afc; this.windowFunctionConverter = wfc; this.typeConverter = tc; + this.executionBehavior = executionBehavior; + } + + /** + * Creates the default execution behavior with PER_PLAN variable evaluation mode. + * + * @return the default execution behavior + */ + private static Plan.ExecutionBehavior createDefaultExecutionBehavior() { + return ImmutableExecutionBehavior.builder() + .variableEvaluationMode( + Plan.ExecutionBehavior.VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build(); } // SQL to Calcite Processing @@ -351,4 +379,18 @@ public WindowFunctionConverter getWindowFunctionConverter() { public TypeConverter getTypeConverter() { return typeConverter; } + + /** + * Returns the execution behavior for plans created by this converter. + * + *

The default execution behavior uses {@link + * Plan.ExecutionBehavior.VariableEvaluationMode#VARIABLE_EVALUATION_MODE_PER_PLAN}, which + * evaluates variables once per plan execution. This can be customized by providing a different + * execution behavior through the constructor. + * + * @return the execution behavior to use when creating plans + */ + public Plan.ExecutionBehavior getExecutionBehavior() { + return executionBehavior; + } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SqlToSubstrait.java b/isthmus/src/main/java/io/substrait/isthmus/SqlToSubstrait.java index 0fb7cf1c3..765dd8dfd 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SqlToSubstrait.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SqlToSubstrait.java @@ -47,6 +47,7 @@ public Plan convert(final String sqlStatements, final Prepare.CatalogReader cata throws SqlParseException { Builder builder = io.substrait.plan.Plan.builder(); builder.version(Version.builder().from(Version.DEFAULT_VERSION).producer("isthmus").build()); + builder.executionBehavior(converterProvider.getExecutionBehavior()); // TODO: consider case in which one sql passes conversion while others don't SubstraitSqlToCalcite.convertQueries(sqlStatements, catalogReader, operatorTable).stream() @@ -73,6 +74,7 @@ public Plan convert( throws SqlParseException { Builder builder = io.substrait.plan.Plan.builder(); builder.version(Version.builder().from(Version.DEFAULT_VERSION).producer("isthmus").build()); + builder.executionBehavior(converterProvider.getExecutionBehavior()); final SqlParser.Config sqlParserConfig = sqlDialect.configureParser( diff --git a/isthmus/src/test/java/io/substrait/isthmus/AutomaticDynamicFunctionMappingRoundtripTest.java b/isthmus/src/test/java/io/substrait/isthmus/AutomaticDynamicFunctionMappingRoundtripTest.java index 9538421e5..03f0c8f0f 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/AutomaticDynamicFunctionMappingRoundtripTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/AutomaticDynamicFunctionMappingRoundtripTest.java @@ -63,13 +63,9 @@ void testMultipleVariantsOfUnmappedFunctionRoundtrip() { sb.remap(2, 3), // The inputs are indices 0, 1. The two scalarFn outputs are 2, 3. table); - // Build plan with output field names + // Build plan with output field names and default execution behavior Plan plan = - Plan.builder() - .roots( - List.of( - Plan.Root.builder().input(project).names(List.of("ts_str", "dt_str")).build())) - .build(); + sb.plan(Plan.Root.builder().input(project).names(List.of("ts_str", "dt_str")).build()); // Use PlanTestBase helper method for comprehensive roundtrip testing assertFullRoundTrip(plan.getRoots().get(0)); diff --git a/isthmus/src/test/resources/lambdas/basic-lambda.json b/isthmus/src/test/resources/lambdas/basic-lambda.json index 0a542c74e..85b70cc00 100644 --- a/isthmus/src/test/resources/lambdas/basic-lambda.json +++ b/isthmus/src/test/resources/lambdas/basic-lambda.json @@ -1,7 +1,9 @@ { "version": { - "majorNumber": 0, - "minorNumber": 79 + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" }, "extensionUrns": [ { diff --git a/isthmus/src/test/resources/lambdas/lambda-field-ref.json b/isthmus/src/test/resources/lambdas/lambda-field-ref.json index a16657663..db4fd21d7 100644 --- a/isthmus/src/test/resources/lambdas/lambda-field-ref.json +++ b/isthmus/src/test/resources/lambdas/lambda-field-ref.json @@ -1,4 +1,10 @@ { + "version": { + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" + }, "extensionUrns": [ { "extensionUrnAnchor": 1, diff --git a/isthmus/src/test/resources/lambdas/lambda-with-function.json b/isthmus/src/test/resources/lambdas/lambda-with-function.json index f11451160..6252a2f63 100644 --- a/isthmus/src/test/resources/lambdas/lambda-with-function.json +++ b/isthmus/src/test/resources/lambdas/lambda-with-function.json @@ -1,4 +1,10 @@ { + "version": { + "minorNumber": 87 + }, + "executionBehavior": { + "variableEvalMode": "VARIABLE_EVALUATION_MODE_PER_PLAN" + }, "extensionUrns": [ { "extensionUrnAnchor": 1, diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala index d55e4c395..2d4d9fb88 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala @@ -47,7 +47,8 @@ import io.substrait.expression.{Expression => SExpression, ExpressionCreator} import io.substrait.expression.Expression.StructLiteral import io.substrait.extension.ExtensionCollector import io.substrait.hint.Hint -import io.substrait.plan.Plan +import io.substrait.plan.{ImmutableExecutionBehavior, Plan} +import io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode import io.substrait.relation.AbstractDdlRel.{DdlObject, DdlOp} import io.substrait.relation.AbstractWriteRel.{CreateMode, OutputMode, WriteOp} import io.substrait.relation.RelProtoConverter @@ -707,6 +708,10 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { .from(Plan.Version.DEFAULT_VERSION) .producer("substrait-spark") .build()) + .executionBehavior( + ImmutableExecutionBehavior.builder() + .variableEvaluationMode(VariableEvaluationMode.VARIABLE_EVALUATION_MODE_PER_PLAN) + .build()) .roots( Collections.singletonList( Plan.Root diff --git a/substrait b/substrait index 2aaae7c0d..7cceb8369 160000 --- a/substrait +++ b/substrait @@ -1 +1 @@ -Subproject commit 2aaae7c0d369a23b5b469badce0644097f1769a1 +Subproject commit 7cceb8369cd533134431979835124de5634a4a47