fix: harden id-based outer-reference conversion

Address review feedback on the id-based outer-reference work:

- core: OuterReferenceConverter clears currentInput inside a subquery
  scope so a correlated reference reached through a multi-input host
  (join/set) is rejected as unsupported instead of silently binding to
  the wrong relation; share one anchor counter across all plan roots so
  assigned rel_anchors are unique plan-wide.
- core: simplify FieldReference#check() to an int sum and drop
  fully-qualified names in OuterReferenceConverter.
- isthmus: reject duplicate rel_anchors on consume, surface an
  unresolvable rel_reference as UnsupportedOperationException
  (consistent with core), and report an actionable error when a custom
  Rel binding target cannot carry an anchor.

Add regression tests for multi-root anchor uniqueness, multi-input scope
rejection, duplicate-anchor rejection, and custom-Rel anchor support.

Author: nielspardon

core/src/main/java/io/substrait/expression/FieldReference.java

@@ -104,13 +104,10 @@ public <R, C extends VisitationContext, E extends Throwable> R accept(
    */
   @Value.Check
   protected void check() {
-    long formsSet =
-        java.util.stream.Stream.of(
-                outerReferenceStepsOut(),
-                outerReferenceRelReference(),
-                lambdaParameterReferenceStepsOut())
-            .filter(Optional::isPresent)
-            .count();
+    int formsSet =
+        (outerReferenceStepsOut().isPresent() ? 1 : 0)
+            + (outerReferenceRelReference().isPresent() ? 1 : 0)
+            + (lambdaParameterReferenceStepsOut().isPresent() ? 1 : 0);
     if (formsSet > 1) {
       throw new IllegalArgumentException(
           "FieldReference can set at most one of outerReferenceStepsOut, "

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

@@ -6,10 +6,13 @@
 import io.substrait.plan.Plan;
 import io.substrait.util.EmptyVisitationContext;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
+import java.util.function.Supplier;
 
 /**
  * Converts a Substrait plan between the two outer-reference resolution encodings:
@@ -88,15 +91,22 @@ public static Plan toStepsOut(Plan plan) {
   }
 
   private static Plan convert(Plan plan, Direction direction) {
+    // A single State (and therefore a single anchor counter) is shared across all roots so that
+    // assigned rel_anchors are unique plan-wide, as required by Rel#getRelAnchor(). The per-tree
+    // traversal state (scope stack, currentInput) is push/pop balanced and empty between roots.
+    State state = new State(direction);
     List<Plan.Root> roots = new ArrayList<>(plan.getRoots().size());
     for (Plan.Root root : plan.getRoots()) {
-      roots.add(Plan.Root.builder().from(root).input(convert(root.getInput(), direction)).build());
+      roots.add(Plan.Root.builder().from(root).input(convert(root.getInput(), state)).build());
     }
     return Plan.builder().from(plan).roots(roots).build();
   }
 
   private static Rel convert(Rel root, Direction direction) {
-    State state = new State(direction);
+    return convert(root, new State(direction));
+  }
+
+  private static Rel convert(Rel root, State state) {
     RelRewriter relRewriter = new RelRewriter(state);
     return root.accept(relRewriter, EmptyVisitationContext.INSTANCE).orElse(root);
   }
@@ -107,8 +117,9 @@ private static final class State {
 
     /**
      * Stack of enclosing scope relations, one entry per subquery boundary. A {@code null} entry
-     * marks a multi-input (unsupported) scope. {@code steps_out=N} resolves to the entry {@code N}
-     * from the top.
+     * marks a scope whose expressions are not hosted by a single-input {@link Filter}/{@link
+     * Project} (e.g. a multi-input join/set condition), which is unsupported. {@code steps_out=N}
+     * resolves to the entry {@code N} from the top.
      */
     final List<Rel> outerScopes = new ArrayList<>();
 
@@ -118,9 +129,15 @@ private static final class State {
     final Map<Rel, Integer> anchorByRel = new IdentityHashMap<>();
 
     /** Anchors that were resolved to a {@code steps_out} value and should be stripped from rels. */
-    final java.util.Set<Integer> resolvedAnchors = new java.util.HashSet<>();
+    final Set<Integer> resolvedAnchors = new HashSet<>();
 
-    /** The input relation whose expressions are currently being rewritten (RootReference scope). */
+    /**
+     * The input relation whose expressions are currently being rewritten, or {@code null} when the
+     * expressions being rewritten are not those of a single-input host (a {@link Filter} or {@link
+     * Project}). Pushed onto {@link #outerScopes} at each subquery boundary; a {@code null} scope
+     * is therefore what makes an outer reference into a multi-input (or otherwise unsupported)
+     * scope fail rather than bind to the wrong relation.
+     */
     Rel currentInput;
 
     int nextAnchor = 1;
@@ -209,7 +226,7 @@ public Optional<Rel> visit(NamedDdl ddl, EmptyVisitationContext context) {
      * Runs the given expression rewrite with {@code currentInput} temporarily set to {@code input},
      * so that subqueries entered while rewriting push the correct scope.
      */
-    private <T> T rewriteInScope(Rel input, java.util.function.Supplier<T> rewrite) {
+    private <T> T rewriteInScope(Rel input, Supplier<T> rewrite) {
       Rel saved = state.currentInput;
       state.currentInput = input;
       try {
@@ -297,8 +314,9 @@ private int anchorForStepsOut(int stepsOut) {
       Rel binding = state.outerScopes.get(index);
       if (binding == null) {
         throw new UnsupportedOperationException(
-            "Cannot assign a rel_anchor for an outer reference that resolves to a multi-input "
-                + "(e.g. join) scope; only single-input scopes are supported");
+            "Cannot assign a rel_anchor for an outer reference that resolves to a scope that is not "
+                + "the input of a single-input host (Filter/Project); multi-input scopes (e.g. a "
+                + "join/set condition) are not supported");
       }
       return state.allocateAnchor(binding);
     }
@@ -337,7 +355,7 @@ public Optional<Expression> visit(
         Expression.InPredicate inPredicate, EmptyVisitationContext context) {
       // needles are evaluated in the current (outer) scope; the haystack is the subquery boundary.
       Optional<List<Expression>> needles = visitExprList(inPredicate.needles(), context);
-      Optional<io.substrait.relation.Rel> haystack =
+      Optional<Rel> haystack =
           withSubqueryScope(
               () -> inPredicate.haystack().accept(getRelCopyOnWriteVisitor(), context));
 
@@ -353,9 +371,15 @@ public Optional<Expression> visit(
     }
 
     /** Pushes {@code currentInput} as a new subquery scope, runs the action, then pops it. */
-    private <T> T withSubqueryScope(java.util.function.Supplier<T> action) {
+    private <T> T withSubqueryScope(Supplier<T> action) {
       Rel savedInput = state.currentInput;
       state.outerScopes.add(state.currentInput);
+      // Clear currentInput for the duration of the subquery so that only a single-input host
+      // (Filter/Project, which re-sets it via rewriteInScope) contributes a non-null scope. Without
+      // this, a subquery reached through a multi-input host (join/set) or any other relation would
+      // inherit this scope's currentInput and silently bind an outer reference to the wrong
+      // relation instead of failing as an unsupported multi-input scope.
+      state.currentInput = null;
       try {
         return action.get();
       } finally {

core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java

@@ -8,6 +8,7 @@
 import io.substrait.TestBase;
 import io.substrait.expression.Expression;
 import io.substrait.expression.FieldReference;
+import io.substrait.plan.Plan;
 import io.substrait.relation.Rel.Remap;
 import io.substrait.type.TypeCreator;
 import java.util.List;
@@ -212,6 +213,96 @@ void unresolvableRelReferenceIsRejected() {
         UnsupportedOperationException.class, () -> OuterReferenceConverter.toStepsOut(dangling));
   }
 
+  @Test
+  void multiRootPlanAssignsPlanWideUniqueAnchors() {
+    // Two roots, each with a correlated subquery binding to a distinct relation. Anchors must be
+    // unique plan-wide (Rel#getRelAnchor()), so the second root's binding relation cannot reuse the
+    // first root's anchor.
+    Plan stepsOut =
+        Plan.builder()
+            // from(...) supplies a valid executionBehavior (and the first root); add the second.
+            .from(sb.plan(sb.root(oneStepPlanBoundTo(orderTableScan))))
+            .addRoots(sb.root(oneStepPlanBoundTo(nationTableScan)))
+            .build();
+
+    Plan idBased = OuterReferenceConverter.toIdBased(stepsOut);
+
+    int anchor0 =
+        ((Project) idBased.getRoots().get(0).getInput())
+            .getInput()
+            .getRelAnchor()
+            .orElseThrow(AssertionError::new);
+    int anchor1 =
+        ((Project) idBased.getRoots().get(1).getInput())
+            .getInput()
+            .getRelAnchor()
+            .orElseThrow(AssertionError::new);
+    assertNotEquals(anchor0, anchor1);
+
+    assertEquals(stepsOut, OuterReferenceConverter.toStepsOut(idBased));
+  }
+
+  @Test
+  void correlatedReferenceIntoJoinScopeIsRejected() {
+    // A correlated reference whose binding scope is a multi-input (join) relation, reached through
+    // a
+    // subquery nested in the join condition, is unsupported and must fail rather than silently bind
+    // the reference to the wrong (single-input) enclosing relation.
+    Rel plan =
+        sb.project(
+            input ->
+                List.of(
+                    sb.scalarSubquery(
+                        sb.innerJoin(
+                            joinInputs ->
+                                sb.equal(
+                                    sb.fieldReference(joinInputs, 0),
+                                    sb.scalarSubquery(
+                                        sb.project(
+                                            input3 -> List.of(sb.fieldReference(input3, 1)),
+                                            Remap.of(List.of(1)),
+                                            sb.filter(
+                                                input3 ->
+                                                    sb.equal(
+                                                        sb.fieldReference(input3, 0),
+                                                        // steps_out=1 resolves to the enclosing
+                                                        // join scope, which is multi-input.
+                                                        FieldReference.newRootStructOuterReference(
+                                                            0, TypeCreator.REQUIRED.I64, 1)),
+                                                customerTableScan)),
+                                        TypeCreator.NULLABLE.I64)),
+                            nationTableScan,
+                            orderTableScan),
+                        TypeCreator.NULLABLE.I64)),
+            Remap.of(List.of(2)),
+            orderTableScan);
+
+    assertThrows(
+        UnsupportedOperationException.class, () -> OuterReferenceConverter.toIdBased(plan));
+  }
+
+  /** A one-step correlated-subquery plan whose outer reference binds to {@code bindingScan}. */
+  private Rel oneStepPlanBoundTo(Rel bindingScan) {
+    return 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)),
+                        sb.filter(
+                            input2 ->
+                                sb.equal(
+                                    sb.fieldReference(input2, 0),
+                                    FieldReference.newRootStructOuterReference(
+                                        1, TypeCreator.REQUIRED.I64, 1)),
+                            customerTableScan)),
+                    TypeCreator.NULLABLE.I64)),
+        Remap.of(List.of(2, 3)),
+        bindingScan);
+  }
+
   /** Extracts the outer field reference from the one-step plan's scalar-subquery filter. */
   private FieldReference outerReferenceOf(Project outerProject) {
     Expression.ScalarSubquery subquery =

isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java

@@ -841,6 +841,15 @@ public static class Context implements VisitationContext {
     /** Maps a {@code rel_anchor} to the single {@link CorrelationId} minted for it. */
     private final Map<Integer, CorrelationId> correlationIdByAnchor = new HashMap<>();
 
+    /**
+     * Every {@code rel_anchor} that has entered a scope. Resolution keys {@link #scopeByAnchor} and
+     * {@link #correlationIdByAnchor} purely by anchor value, which is only sound if anchors are
+     * unique plan-wide (as required by {@link io.substrait.relation.Rel#getRelAnchor()}). This set
+     * lets {@link #enterScope} reject a plan that reuses an anchor for two distinct relations
+     * rather than silently mis-resolving references.
+     */
+    private final java.util.Set<Integer> seenAnchors = new HashSet<>();
+
     /** One correlation scope per enclosing relational operator. */
     private static final class Scope {
       final Map<Integer, RelDataType> rowTypeByAnchor = new HashMap<>();
@@ -867,6 +876,13 @@ public void enterScope(final AnchoredInput... inputs) {
       for (final AnchoredInput input : inputs) {
         if (input.anchor.isPresent()) {
           final int anchor = input.anchor.get();
+          if (!seenAnchors.add(anchor)) {
+            throw new UnsupportedOperationException(
+                "Duplicate rel_anchor="
+                    + anchor
+                    + "; rel_anchors must be unique plan-wide for id-based outer references to "
+                    + "resolve unambiguously");
+          }
           scope.rowTypeByAnchor.put(anchor, input.rowType);
           scopeByAnchor.put(anchor, scope);
         }
@@ -919,7 +935,10 @@ public CorrelationId correlationIdForAnchor(
     private Scope requireScope(final int anchor) {
       final Scope scope = scopeByAnchor.get(anchor);
       if (scope == null) {
-        throw new IllegalStateException(
+        // The anchor is not on the active scope stack: the referenced relation is not an enclosing
+        // single-input host. This includes forward references and shared subtrees reached via a
+        // ReferenceRel. Signalled as unsupported, consistent with OuterReferenceConverter.
+        throw new UnsupportedOperationException(
             "Outer reference rel_reference="
                 + anchor
                 + " has no enclosing relation with that anchor");

isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java

@@ -805,7 +805,20 @@ public Rel apply(RelNode r) {
     if (outerReferenceResolver != null) {
       Integer anchor = outerReferenceResolver.anchorForTarget(r);
       if (anchor != null) {
-        return rel.withRelAnchor(anchor);
+        try {
+          return rel.withRelAnchor(anchor);
+        } catch (UnsupportedOperationException e) {
+          // rel is the binding point of a correlated outer reference but cannot carry a rel_anchor
+          // (a custom, non-Immutables Rel that does not override withRelAnchor). Fail with context
+          // rather than propagating the bare "does not support setting a relation anchor" default.
+          throw new UnsupportedOperationException(
+              "Relation "
+                  + rel.getClass()
+                  + " is the binding point of an id-based outer reference but does not support "
+                  + "setting a rel_anchor; override Rel#withRelAnchor to convert correlated "
+                  + "subqueries into this relation type",
+              e);
+        }
       }
     }
     return rel;

isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java

@@ -1,7 +1,10 @@
 package io.substrait.isthmus;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import io.substrait.dsl.SubstraitBuilder;
 import io.substrait.extension.AdvancedExtension;
 import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.SimpleExtension;
@@ -10,6 +13,7 @@
 import io.substrait.relation.RelVisitor;
 import io.substrait.relation.SingleInputRel;
 import io.substrait.type.Type;
+import io.substrait.type.TypeCreator;
 import io.substrait.util.VisitationContext;
 import java.util.List;
 import java.util.Optional;
@@ -115,6 +119,18 @@ public Optional<Integer> getRelAnchor() {
       return input.getRelAnchor();
     }
 
+    @Override
+    public Rel withRelAnchor(final int relAnchor) {
+      // Delegate to the input, mirroring getRelAnchor(), so this custom Rel can serve as the
+      // binding point of an id-based outer reference instead of inheriting Rel's throwing default.
+      return new SubstraitRepeatRel(input.withRelAnchor(relAnchor), repeatCount);
+    }
+
+    @Override
+    public Rel withRelAnchor(final Optional<Integer> relAnchor) {
+      return new SubstraitRepeatRel(input.withRelAnchor(relAnchor), repeatCount);
+    }
+
     @Override
     public <O, C extends VisitationContext, E extends Exception> O accept(
         final RelVisitor<O, C, E> visitor, final C context) throws E {
@@ -257,4 +273,22 @@ void test() {
         findNode(rel, SubstraitRepeatRel.class),
         "substrait plan must contain SubstraitRepeatRel relation");
   }
+
+  @Test
+  void customRelSupportsRelAnchor() {
+    // A custom (non-Immutables) Rel completes the Rel contract by delegating withRelAnchor, so
+    // SubstraitRelVisitor#apply can stamp an id-based outer-reference anchor on it rather than
+    // hitting Rel's throwing withRelAnchor default.
+    final SubstraitBuilder sb = new SubstraitBuilder();
+    final Rel scan = sb.namedScan(List.of("t"), List.of("a"), List.of(TypeCreator.REQUIRED.I64));
+    final SubstraitRepeatRel repeat = new SubstraitRepeatRel(scan, 3);
+
+    final Rel anchored = repeat.withRelAnchor(7);
+    assertTrue(anchored instanceof SubstraitRepeatRel);
+    assertEquals(3, ((SubstraitRepeatRel) anchored).getRepeatCount());
+    assertEquals(7, anchored.getRelAnchor().orElseThrow(AssertionError::new));
+
+    // Clearing the anchor works too.
+    assertFalse(anchored.withRelAnchor(Optional.empty()).getRelAnchor().isPresent());
+  }
 }