fix(core): resolve reference segments totally, and scope join keys per side

Two defects in the field-reference retyping, both found reviewing it.

Segment resolution was guarded one segment deep. The guard checked only that the
outermost segment selected a field the new record type has, so a reference into a
nested type that had been reshaped reached the derivation anyway and threw — a bare
IndexOutOfBoundsException out of a struct-field finder, or an IllegalArgumentException
or UnsupportedOperationException for a list or map segment. That contradicted the
guard's own promise to leave an unresolvable reference alone, and it made a narrowing
rewrite fail where it previously produced a stale type.

Resolution is now total: FieldReference.resolveType reports the type a chain of
segments selects, or nothing, at any depth and for every kind of segment. It lives
beside the finders whose rules it mirrors, so the two cannot drift apart unnoticed
without the parity test that pins them going red. It mirrors those rules exactly,
including the two asymmetries that matter: a list element offset is not bounds
checked, because the length of a list is not part of its type, and a map key type is
compared exactly, nullability included. The visitor's one-deep guard and its defensive
segment copy both go, and the same guarantee now covers a reference rooted at another
expression, which had no guard at all.

Join keys were retyped against the wrong scope. The offsets of a hash or merge join
key are relative to the side of the join the key selects from, not to the two inputs
combined — proto conversion types each side with its own converter, and only the
condition, post-join filter and residual expression use the combined type. Retyping
both sides against the combined type silently resolved a right-side offset to a left
column. Each side is now rewritten against its own input, which changes the signature
of visitComparisonJoinKey; the method could not previously return a usable reference at
all, so nothing can have depended on the old one.

Two things this deliberately does not do: the off-by-one bound check in
StructFieldFinder stays, because the exception it produces is part of what
ProtoExpressionConverter reports for a malformed plan; and MergeJoin.deriveRecordType
reads its right input for both sides, which is a separate bug in a relation this change
only passes through.

Author: nielspardon

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

@@ -480,7 +480,7 @@ public FieldReference constructOnRoot(Type.Struct struct) {
    * Creates a field reference rooted at an expression and navigating through the given segments.
    *
    * @param expression the expression to reference into
-   * @param segments the navigation segments, outermost first
+   * @param segments the navigation segments, innermost first
    * @return the field reference
    */
   public static FieldReference ofExpression(
@@ -509,13 +509,82 @@ private static FieldReference of(
    * Creates a field reference rooted at a struct and navigating through the given segments.
    *
    * @param struct the root struct type
-   * @param segments the navigation segments, outermost first
+   * @param segments the navigation segments, innermost first
    * @return the field reference
    */
   public static FieldReference ofRoot(Type.Struct struct, List<ReferenceSegment> segments) {
     return of(struct, null, segments);
   }
 
+  /**
+   * Resolves the type that the given reference segments select out of the given type, or reports
+   * that they select nothing out of it.
+   *
+   * <p>The segments are given innermost first, in the order {@link #segments()} holds them, and are
+   * applied outermost first: the last segment selects out of {@code rootType}, the one before it
+   * out of the type that segment selected, and so on inwards. The given list is not modified.
+   *
+   * <p>Resolution is total. A segment that does not fit the type it is applied to, at any depth — a
+   * struct field offset the struct does not have, a list element or a map key on a type that is not
+   * a list or a map, a map key whose type is not the map's key type — yields an empty result
+   * instead of throwing. This lets a caller that re-derives the type cached on a reference, against
+   * a type that has since changed, tell a reference that no longer resolves from a failure of its
+   * own work, and leave such a reference as it is. An empty segment list selects nothing and also
+   * yields an empty result, matching {@link #ofRoot} and {@link #ofExpression}, which build no
+   * reference for it.
+   *
+   * @param rootType the type the outermost segment selects out of
+   * @param segments the navigation segments, innermost first
+   * @return the type the segments select, or empty if they do not all resolve against {@code
+   *     rootType}
+   */
+  public static Optional<Type> resolveType(Type rootType, List<ReferenceSegment> segments) {
+    if (segments.isEmpty()) {
+      return Optional.empty();
+    }
+    Type resolved = rootType;
+    for (int i = segments.size() - 1; i >= 0; i--) {
+      Optional<Type> selected = resolveSegmentType(segments.get(i), resolved);
+      if (!selected.isPresent()) {
+        return Optional.empty();
+      }
+      resolved = selected.get();
+    }
+    return Optional.of(resolved);
+  }
+
+  /**
+   * Resolves the type a single segment selects out of the given type, or reports that it selects
+   * nothing out of it.
+   *
+   * <p>This mirrors the type each segment derives when it is applied: a struct field selects the
+   * field at its offset, a list element selects the element type whatever its offset, as the length
+   * of a list is not part of its type, and a map key selects the value type of a map whose key type
+   * it matches, nullability included. A segment applied to a type that is not the container it
+   * navigates into, and any other kind of segment, select nothing.
+   */
+  private static Optional<Type> resolveSegmentType(ReferenceSegment segment, Type type) {
+    if (segment instanceof StructField && type instanceof Type.Struct) {
+      int offset = ((StructField) segment).offset();
+      List<Type> fields = ((Type.Struct) type).fields();
+      return offset >= 0 && offset < fields.size()
+          ? Optional.of(fields.get(offset))
+          : Optional.empty();
+    }
+    if (segment instanceof ListElement && type instanceof Type.ListType) {
+      return Optional.of(((Type.ListType) type).elementType());
+    }
+    if (segment instanceof MapKey && type instanceof Type.Map) {
+      // The type of the key literal is only read once the type is known to be a map, which is also
+      // the only case in which applying the segment reads it.
+      Type.Map map = (Type.Map) type;
+      return map.key().equals(((MapKey) segment).key().getType())
+          ? Optional.of(map.value())
+          : Optional.empty();
+    }
+    return Optional.empty();
+  }
+
   private static class StructFieldFinder
       extends TypeVisitor.TypeThrowsVisitor<Type, RuntimeException> {
 

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

@@ -12,7 +12,6 @@
 import io.substrait.expression.ImmutableFieldReference;
 import io.substrait.type.Type;
 import io.substrait.util.EmptyVisitationContext;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
@@ -435,11 +434,21 @@ public Optional<Expression> visit(FieldReference fieldReference, EmptyVisitation
    * cached type from the root it resolves against. Re-deriving the type is what keeps references
    * correct when a relation's input is replaced by one emitting a different record type.
    *
+   * <p>A reference that no longer resolves against its root, at any segment depth, is left exactly
+   * as it is, including the type it has cached. Such a reference makes the relation tree invalid
+   * whatever this visitor does with it, so re-deriving its type is not turned into a failure of the
+   * rewrite.
+   *
    * <p>The type of a reference to a lambda parameter, or of an outer reference identified by the
    * {@link io.substrait.relation.Rel#getRelAnchor() rel anchor} of the relation it is rooted on, is
    * left as it is: neither resolves against a relation in the enclosing scopes tracked during the
    * traversal.
    *
+   * <p>Override this rather than {@link #visit(FieldReference, EmptyVisitationContext)} to change
+   * how references are rewritten: this is what the positions that hold a reference rather than an
+   * arbitrary expression — a {@link io.substrait.relation.physical.ScatterExchange}'s fields and a
+   * {@link io.substrait.relation.physical.ComparisonJoinKey}'s sides — are rewritten through.
+   *
    * @param fieldReference the field reference to visit
    * @param context the visitation context
    * @return Optional containing the modified field reference, or empty if no changes
@@ -453,12 +462,12 @@ public Optional<FieldReference> visitFieldReference(
       if (!inputExpression.isPresent()) {
         return Optional.empty();
       }
+      // The reference is returned rewritten even when its type could not be re-derived: the
+      // expression it is rooted at changed, and that change is not lost because the type is stale.
       ImmutableFieldReference.Builder rewritten =
           ImmutableFieldReference.builder().from(fieldReference).inputExpression(inputExpression);
-      if (!fieldReference.segments().isEmpty()) {
-        rewritten.type(
-            FieldReference.ofExpression(inputExpression.get(), segmentsOf(fieldReference)).type());
-      }
+      FieldReference.resolveType(inputExpression.get().getType(), fieldReference.segments())
+          .ifPresent(rewritten::type);
       return Optional.of(rewritten.build());
     }
     return retypeRootReference(fieldReference);
@@ -476,40 +485,21 @@ private Optional<FieldReference> retypeRootReference(FieldReference fieldReferen
     Type.Struct rootType =
         getRelCopyOnWriteVisitor()
             .inputTypeStepsOut(fieldReference.outerReferenceStepsOut().orElse(0));
-    // A rewrite that drops fields can leave a reference selecting a field its input no longer has.
-    // The resulting relation tree is invalid either way, so leave the reference as it is rather
-    // than turn re-deriving its type into a failure of the rewrite.
-    if (rootType == null || !selectsFieldOf(fieldReference, rootType)) {
+    if (rootType == null) {
       return Optional.empty();
     }
-    Type type = FieldReference.ofRoot(rootType, segmentsOf(fieldReference)).type();
-    if (type.equals(fieldReference.type())) {
+    // A rewrite that drops a field, or that reshapes a nested one, can leave a reference selecting
+    // something its input no longer has. The resulting relation tree is invalid either way, so
+    // leave
+    // the reference as it is rather than turn re-deriving its type into a failure of the rewrite.
+    // Resolution is total, at every segment depth and for every kind of segment, so this is a
+    // decision the rewrite makes rather than an exception it has to recover from.
+    Optional<Type> type = FieldReference.resolveType(rootType, fieldReference.segments());
+    if (!type.isPresent() || type.get().equals(fieldReference.type())) {
       return Optional.empty();
     }
-    return Optional.of(ImmutableFieldReference.builder().from(fieldReference).type(type).build());
-  }
-
-  /**
-   * Returns whether the outermost segment of the given reference selects a field that the given
-   * record type has. The segments are held innermost first, so the outermost one is the last.
-   */
-  private static boolean selectsFieldOf(FieldReference fieldReference, Type.Struct rootType) {
-    List<FieldReference.ReferenceSegment> segments = fieldReference.segments();
-    if (segments.isEmpty()) {
-      return false;
-    }
-    FieldReference.ReferenceSegment outermost = segments.get(segments.size() - 1);
-    return outermost instanceof FieldReference.StructField
-        && ((FieldReference.StructField) outermost).offset() < rootType.fields().size();
-  }
-
-  /**
-   * Returns the segments of the given reference as a mutable list, which is what {@link
-   * FieldReference#ofRoot} and {@link FieldReference#ofExpression} require: they reverse the list
-   * they are given.
-   */
-  private static List<FieldReference.ReferenceSegment> segmentsOf(FieldReference fieldReference) {
-    return new ArrayList<>(fieldReference.segments());
+    return Optional.of(
+        ImmutableFieldReference.builder().from(fieldReference).type(type.get()).build());
   }
 
   @Override

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

@@ -609,12 +609,15 @@ public Optional<Rel> visit(ExtensionTable extensionTable, EmptyVisitationContext
   public Optional<Rel> visit(HashJoin hashJoin, EmptyVisitationContext context) throws E {
     Optional<Rel> left = hashJoin.getLeft().accept(this, context);
     Optional<Rel> right = hashJoin.getRight().accept(this, context);
+    Type.Struct leftType = recordTypeOf(left.orElse(hashJoin.getLeft()));
+    Type.Struct rightType = recordTypeOf(right.orElse(hashJoin.getRight()));
     Type.Struct inputType =
         recordTypeOf(left.orElse(hashJoin.getLeft()), right.orElse(hashJoin.getRight()));
     Optional<List<ComparisonJoinKey>> keys =
-        inInputScope(
-            inputType,
-            () -> transformList(hashJoin.getKeys(), context, this::visitComparisonJoinKey));
+        transformList(
+            hashJoin.getKeys(),
+            context,
+            (key, c) -> visitComparisonJoinKey(key, leftType, rightType, c));
     Optional<Expression> postFilter =
         inInputScope(
             inputType, () -> visitOptionalExpression(hashJoin.getPostJoinFilter(), context));
@@ -640,12 +643,15 @@ public Optional<Rel> visit(HashJoin hashJoin, EmptyVisitationContext context) th
   public Optional<Rel> visit(MergeJoin mergeJoin, EmptyVisitationContext context) throws E {
     Optional<Rel> left = mergeJoin.getLeft().accept(this, context);
     Optional<Rel> right = mergeJoin.getRight().accept(this, context);
+    Type.Struct leftType = recordTypeOf(left.orElse(mergeJoin.getLeft()));
+    Type.Struct rightType = recordTypeOf(right.orElse(mergeJoin.getRight()));
     Type.Struct inputType =
         recordTypeOf(left.orElse(mergeJoin.getLeft()), right.orElse(mergeJoin.getRight()));
     Optional<List<ComparisonJoinKey>> keys =
-        inInputScope(
-            inputType,
-            () -> transformList(mergeJoin.getKeys(), context, this::visitComparisonJoinKey));
+        transformList(
+            mergeJoin.getKeys(),
+            context,
+            (key, c) -> visitComparisonJoinKey(key, leftType, rightType, c));
     Optional<Expression> postFilter =
         inInputScope(
             inputType, () -> visitOptionalExpression(mergeJoin.getPostJoinFilter(), context));
@@ -761,6 +767,8 @@ protected Optional<ConsistentPartitionWindow.WindowRelFunctionInvocation> visitW
    * Returns the record type that the root field references of a relation's own expressions resolve
    * against: the record types of the given inputs, concatenated in order.
    *
+   * <p>Only the fields of the result are ever read, so its own nullability is not meaningful.
+   *
    * @param inputs the relations the expressions are evaluated over, in field order
    * @return the combined record type
    */
@@ -872,15 +880,27 @@ public Optional<FieldReference> visitFieldReference(
   /**
    * Rewrites a comparison join key, returning a new one if either side changed.
    *
+   * <p>Each side is rewritten against its own input, because the field offsets of a join key are
+   * relative to the side of the join they select from — unlike those of a join condition or
+   * post-join filter, which are relative to the two inputs combined.
+   *
    * @param key the comparison join key to rewrite
+   * @param leftType the record type the key's left side selects from
+   * @param rightType the record type the key's right side selects from
    * @param context the visitation context
    * @return the rewritten comparison join key, or empty if unchanged
    * @throws E if the visit fails
    */
   public Optional<ComparisonJoinKey> visitComparisonJoinKey(
-      ComparisonJoinKey key, EmptyVisitationContext context) throws E {
-    Optional<FieldReference> left = visitFieldReference(key.getLeft(), context);
-    Optional<FieldReference> right = visitFieldReference(key.getRight(), context);
+      ComparisonJoinKey key,
+      Type.Struct leftType,
+      Type.Struct rightType,
+      EmptyVisitationContext context)
+      throws E {
+    Optional<FieldReference> left =
+        inInputScope(leftType, () -> visitFieldReference(key.getLeft(), context));
+    Optional<FieldReference> right =
+        inInputScope(rightType, () -> visitFieldReference(key.getRight(), context));
     if (allEmpty(left, right)) {
       return Optional.empty();
     }