fix(core)!: visit a window relation's input when rewriting it

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

@@ -392,7 +392,7 @@ public Optional<Rel> visit(MultiBucketExchange exchange, EmptyVisitationContext
     Optional<Expression> expression =
         exchange.getExpression().accept(getExpressionCopyOnWriteVisitor(), context);
 
-    if (allEmpty(input)) {
+    if (allEmpty(input, expression)) {
       return Optional.empty();
     }
 
@@ -611,6 +611,7 @@ public Optional<Rel> visit(NestedLoopJoin nestedLoopJoin, EmptyVisitationContext
   public Optional<Rel> visit(
       ConsistentPartitionWindow consistentPartitionWindow, EmptyVisitationContext context)
       throws E {
+    Optional<Rel> input = consistentPartitionWindow.getInput().accept(this, context);
     Optional<List<ConsistentPartitionWindow.WindowRelFunctionInvocation>> windowFunctions =
         transformList(
             consistentPartitionWindow.getWindowFunctions(), context, this::visitWindowRelFunction);
@@ -622,13 +623,14 @@ public Optional<Rel> visit(
     Optional<List<Expression.SortField>> 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()))

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(

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<Expression> 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<Rel> 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<Rel> 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<Rel> 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<Class<?>> notVisited = Arrays.asList(Expand.class, ExtensionWrite.class);
+    RelCopyOnWriteVisitor<RuntimeException> renameScans =
+        new RelCopyOnWriteVisitor<RuntimeException>() {
+          @Override
+          public Optional<Rel> 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<Class<?>> notVisited = Arrays.asList(Expand.class, ExtensionWrite.class);
+    RelCopyOnWriteVisitor<RuntimeException> markEveryReference =
+        new RelCopyOnWriteVisitor<>(
+            relVisitor ->
+                new ExpressionCopyOnWriteVisitor<RuntimeException>(relVisitor) {
+                  @Override
+                  public Optional<Expression> visit(
+                      FieldReference reference, EmptyVisitationContext context) {
+                    return Optional.of(reference);
+                  }
+                });
+
+    List<String> 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()));
+  }
 }