diff --git a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java index 13961775d..a991ea6b8 100644 --- a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java @@ -392,7 +392,7 @@ public Optional visit(MultiBucketExchange exchange, EmptyVisitationContext Optional expression = exchange.getExpression().accept(getExpressionCopyOnWriteVisitor(), context); - if (allEmpty(input)) { + if (allEmpty(input, expression)) { return Optional.empty(); } @@ -611,6 +611,7 @@ public Optional visit(NestedLoopJoin nestedLoopJoin, EmptyVisitationContext public Optional visit( ConsistentPartitionWindow consistentPartitionWindow, EmptyVisitationContext context) throws E { + Optional input = consistentPartitionWindow.getInput().accept(this, context); Optional> windowFunctions = transformList( consistentPartitionWindow.getWindowFunctions(), context, this::visitWindowRelFunction); @@ -622,13 +623,14 @@ public Optional visit( Optional> sorts = transformList(consistentPartitionWindow.getSorts(), context, this::visitSortField); - if (allEmpty(windowFunctions, partitionExpressions, sorts)) { + if (allEmpty(input, windowFunctions, partitionExpressions, sorts)) { return Optional.empty(); } return Optional.of( ConsistentPartitionWindow.builder() .from(consistentPartitionWindow) + .input(input.orElse(consistentPartitionWindow.getInput())) .partitionExpressions( partitionExpressions.orElse(consistentPartitionWindow.getPartitionExpressions())) .sorts(sorts.orElse(consistentPartitionWindow.getSorts())) diff --git a/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java b/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java index 3c70bee82..fd3c779ea 100644 --- a/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java +++ b/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java @@ -332,6 +332,74 @@ void outerReferenceInsideWindowBoundOffsetIsConverted() { assertEquals(stepsOut, OuterReferenceConverter.toStepsOut(idBased)); } + @Test + void outerReferenceUnderAWindowRelationIsConverted() { + // The correlated filter sits in the window relation's input rather than in its bounds, so the + // reference is reachable only once the rewrite descends into that input. Asserting on the + // id-based plan rather than on a round trip, which was an identity in both directions before. + SimpleExtension.WindowFunctionVariant declaration = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + + Rel correlatedFilter = + sb.filter( + input -> + sb.equal( + sb.fieldReference(input, 0), + FieldReference.newRootStructOuterReference(1, TypeCreator.REQUIRED.I64, 1)), + customerTableScan); + Rel window = + ConsistentPartitionWindow.builder() + .input(correlatedFilter) + .windowFunctions( + List.of( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(declaration) + .arguments(List.of(sb.fieldReference(correlatedFilter, 0))) + .outputType(TypeCreator.NULLABLE.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .lowerBound(WindowBound.UNBOUNDED) + .upperBound(WindowBound.CURRENT_ROW) + .boundsType(Expression.WindowBoundsType.RANGE) + .build())) + .build(); + + Rel stepsOut = + sb.project( + input -> + List.of( + sb.fieldReference(input, 0), + sb.scalarSubquery( + sb.project( + input2 -> List.of(sb.fieldReference(input2, 1)), + Remap.of(List.of(1)), + window), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2, 3)), + orderTableScan); + + Rel idBased = OuterReferenceConverter.toIdBased(stepsOut); + + assertNotEquals(stepsOut, idBased); + + Project outerProject = (Project) idBased; + assertEquals(1, outerProject.getInput().getRelAnchor().orElseThrow(AssertionError::new)); + + Expression.ScalarSubquery subquery = + (Expression.ScalarSubquery) outerProject.getExpressions().get(1); + Filter filter = + (Filter) ((ConsistentPartitionWindow) ((Project) subquery.input()).getInput()).getInput(); + Expression.ScalarFunctionInvocation equal = + (Expression.ScalarFunctionInvocation) filter.getCondition(); + FieldReference outerRef = (FieldReference) equal.arguments().get(1); + assertEquals(1, outerRef.outerReferenceRelReference().orElseThrow(AssertionError::new)); + assertFalse(outerRef.outerReferenceStepsOut().isPresent()); + + assertEquals(stepsOut, OuterReferenceConverter.toStepsOut(idBased)); + } + /** A one-step correlated-subquery plan whose outer reference binds to {@code bindingScan}. */ private Rel oneStepPlanBoundTo(Rel bindingScan) { return sb.project( diff --git a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java index 7856a81d1..f4a56b062 100644 --- a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java +++ b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java @@ -1,16 +1,23 @@ package io.substrait.relation; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; import io.substrait.TestBase; import io.substrait.expression.Expression; +import io.substrait.expression.FieldReference; import io.substrait.expression.WindowBound; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.SimpleExtension; +import io.substrait.relation.physical.MultiBucketExchange; import io.substrait.util.EmptyVisitationContext; +import io.substrait.utils.RelSamples; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.Optional; +import java.util.stream.Collectors; import org.junit.jupiter.api.Test; class RelCopyOnWriteVisitorTest extends TestBase { @@ -35,29 +42,33 @@ public Optional visitLiteral(Expression.Literal literal) { }); } - @Test - void consistentPartitionWindowBoundOffsetsAreRewritten() { + /** A window over {@code input} whose one function carries the given bounds. */ + private ConsistentPartitionWindow windowOver(Rel input, WindowBound lower, WindowBound upper) { SimpleExtension.WindowFunctionVariant declaration = extensions.getWindowFunction( SimpleExtension.FunctionAnchor.of( DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + return ConsistentPartitionWindow.builder() + .input(input) + .windowFunctions( + Arrays.asList( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(declaration) + .outputType(R.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .lowerBound(lower) + .upperBound(upper) + .boundsType(Expression.WindowBoundsType.RANGE) + .build())) + .build(); + } + + @Test + void consistentPartitionWindowBoundOffsetsAreRewritten() { Rel input = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I64)); ConsistentPartitionWindow window = - ConsistentPartitionWindow.builder() - .input(input) - .windowFunctions( - Arrays.asList( - ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() - .declaration(declaration) - .arguments(Collections.emptyList()) - .outputType(R.I64) - .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) - .invocation(Expression.AggregationInvocation.ALL) - .lowerBound(WindowBound.Preceding.of(sb.i32(5))) - .upperBound(WindowBound.Following.of(sb.i32(7))) - .boundsType(Expression.WindowBoundsType.RANGE) - .build())) - .build(); + windowOver(input, WindowBound.Preceding.of(sb.i32(5)), WindowBound.Following.of(sb.i32(7))); Optional rewritten = window.accept(negateI32LiteralsVisitor(), EmptyVisitationContext.INSTANCE); @@ -78,30 +89,137 @@ void consistentPartitionWindowBoundOffsetsAreRewritten() { @Test void consistentPartitionWindowWithUnchangedBoundsIsNotCopied() { - SimpleExtension.WindowFunctionVariant declaration = - extensions.getWindowFunction( - SimpleExtension.FunctionAnchor.of( - DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); Rel input = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I64)); ConsistentPartitionWindow window = - ConsistentPartitionWindow.builder() - .input(input) - .windowFunctions( - Arrays.asList( - ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() - .declaration(declaration) - .arguments(Collections.emptyList()) - .outputType(R.I64) - .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) - .invocation(Expression.AggregationInvocation.ALL) - .lowerBound(WindowBound.UNBOUNDED) - .upperBound(WindowBound.CURRENT_ROW) - .boundsType(Expression.WindowBoundsType.RANGE) - .build())) - .build(); + windowOver(input, WindowBound.UNBOUNDED, WindowBound.CURRENT_ROW); assertEquals( Optional.empty(), window.accept(negateI32LiteralsVisitor(), EmptyVisitationContext.INSTANCE)); } + + /** A rewrite that applies only below a window relation comes back with the input replaced. */ + @Test + void consistentPartitionWindowRewritesItsInput() { + Rel scan = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I32)); + Rel input = sb.project(in -> Arrays.asList(sb.i32(5)), sb.remap(1), scan); + ConsistentPartitionWindow window = + windowOver(input, WindowBound.UNBOUNDED, WindowBound.CURRENT_ROW); + + Optional rewritten = + window.accept(negateI32LiteralsVisitor(), EmptyVisitationContext.INSTANCE); + + ConsistentPartitionWindow expected = + ConsistentPartitionWindow.builder() + .from(window) + .input(sb.project(in -> Arrays.asList(sb.i32(-5)), sb.remap(1), scan)) + .build(); + assertEquals(Optional.of(expected), rewritten); + } + + /** + * The same guard, one relation over: a rewrite that touches only the exchange's own expression + * comes back, where the input alone used to decide whether anything changed. + */ + @Test + void multiBucketExchangeRewritesItsExpression() { + Rel scan = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I32)); + MultiBucketExchange exchange = + MultiBucketExchange.builder() + .input(scan) + .expression(sb.i32(5)) + .constrainedToCount(true) + .partitionCount(1) + .build(); + + Optional rewritten = + exchange.accept(negateI32LiteralsVisitor(), EmptyVisitationContext.INSTANCE); + + assertEquals( + Optional.of(MultiBucketExchange.builder().from(exchange).expression(sb.i32(-5)).build()), + rewritten); + } + + /** + * Every relation that has inputs hands back a rewrite made below it, so a relation added without + * visiting its input fails here rather than dropping rewrites silently. Driven by the shared + * samples, whose own test keeps them exhaustive over the model. + */ + @Test + void everyRelationWithInputsPropagatesARewriteBelowIt() { + // The two the visitor refuses outright rather than descending into. + List> notVisited = Arrays.asList(Expand.class, ExtensionWrite.class); + RelCopyOnWriteVisitor renameScans = + new RelCopyOnWriteVisitor() { + @Override + public Optional visit(NamedScan namedScan, EmptyVisitationContext context) { + return Optional.of( + NamedScan.builder() + .from(namedScan) + .names(Collections.singletonList("renamed")) + .build()); + } + }; + + new RelSamples(sb, extensions) + .samples() + .forEach( + (type, rel) -> { + if (rel.getInputs().isEmpty() || notVisited.contains(type)) { + return; + } + assertTrue( + rel.accept(renameScans, EmptyVisitationContext.INSTANCE).isPresent(), + type.getSimpleName()); + }); + } + + /** + * The sweep above exercises only the input branch of each guard: a rewrite that reaches an + * expression and leaves the input alone has to come back too, which is the second half of this + * fix. The relations whose samples carry a field reference are pinned by name, so a relation + * whose visit computes an expression and then decides from its input drops out of the set rather + * than passing quietly -- {@code MultiBucketExchange} does exactly that on the base. + */ + @Test + void everyRelationCarryingAnExpressionPropagatesARewriteInIt() { + List> notVisited = Arrays.asList(Expand.class, ExtensionWrite.class); + RelCopyOnWriteVisitor markEveryReference = + new RelCopyOnWriteVisitor<>( + relVisitor -> + new ExpressionCopyOnWriteVisitor(relVisitor) { + @Override + public Optional visit( + FieldReference reference, EmptyVisitationContext context) { + return Optional.of(reference); + } + }); + + List rewritten = new ArrayList<>(); + new RelSamples(sb, extensions) + .samples() + .forEach( + (type, rel) -> { + if (rel.getInputs().isEmpty() || notVisited.contains(type)) { + return; + } + if (rel.accept(markEveryReference, EmptyVisitationContext.INSTANCE).isPresent()) { + rewritten.add(type.getSimpleName()); + } + }); + + assertEquals( + Arrays.asList( + "Aggregate", + "ConsistentPartitionWindow", + "Filter", + "Join", + "MultiBucketExchange", + "NestedLoopJoin", + "Project", + "SingleBucketExchange", + "Sort", + "TopN"), + rewritten.stream().sorted().collect(Collectors.toList())); + } }