fix(core): retain reference scope during dereference

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

@@ -161,10 +161,10 @@ public FieldReference dereferenceStruct(int index) {
 
   private FieldReference dereference(Type newType, ReferenceSegment nextSegment) {
     return ImmutableFieldReference.builder()
+        .from(this)
         .type(newType)
-        .addSegments(nextSegment)
+        .segments(Collections.singletonList(nextSegment))
         .addAllSegments(segments())
-        .inputExpression(inputExpression())
         .build();
   }
 

core/src/test/java/io/substrait/expression/FieldReferenceDereferenceTest.java

@@ -0,0 +1,105 @@
+package io.substrait.expression;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.substrait.TestBase;
+import io.substrait.type.Type;
+import java.util.List;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+class FieldReferenceDereferenceTest extends TestBase {
+
+  enum ReferenceScope {
+    ROOT,
+    EXPRESSION,
+    OUTER_STEPS,
+    OUTER_ANCHOR,
+    LAMBDA_CURRENT,
+    LAMBDA_OUTER
+  }
+
+  @ParameterizedTest
+  @EnumSource(ReferenceScope.class)
+  void structDereferencePreservesScope(ReferenceScope scope) {
+    FieldReference reference = reference(scope, R.struct(R.BOOLEAN, N.I64));
+
+    assertDereference(
+        reference, reference.dereferenceStruct(1), N.I64, FieldReference.StructField.of(1));
+  }
+
+  @ParameterizedTest
+  @EnumSource(ReferenceScope.class)
+  void listDereferencePreservesScope(ReferenceScope scope) {
+    FieldReference reference = reference(scope, R.list(N.I64));
+
+    assertDereference(
+        reference, reference.dereferenceList(2), N.I64, FieldReference.ListElement.of(2));
+  }
+
+  @ParameterizedTest
+  @EnumSource(ReferenceScope.class)
+  void mapDereferencePreservesScope(ReferenceScope scope) {
+    FieldReference reference = reference(scope, R.map(R.STRING, N.I64));
+    Expression.Literal key = ExpressionCreator.string(false, "key");
+
+    assertDereference(
+        reference, reference.dereferenceMap(key), N.I64, FieldReference.MapKey.of(key));
+  }
+
+  private FieldReference reference(ReferenceScope scope, Type type) {
+    ImmutableFieldReference.Builder builder =
+        FieldReference.builder().type(type).addSegments(FieldReference.StructField.of(1));
+    switch (scope) {
+      case EXPRESSION:
+        builder.inputExpression(
+            Expression.DynamicParameter.builder()
+                .type(R.struct(R.BOOLEAN, type))
+                .parameterReference(0)
+                .build());
+        break;
+      case OUTER_STEPS:
+        builder.outerReferenceStepsOut(2);
+        break;
+      case OUTER_ANCHOR:
+        builder.outerReferenceRelReference(7);
+        break;
+      case LAMBDA_CURRENT:
+        builder.lambdaParameterReferenceStepsOut(0);
+        break;
+      case LAMBDA_OUTER:
+        builder.lambdaParameterReferenceStepsOut(2);
+        break;
+      case ROOT:
+        break;
+      default:
+        throw new IllegalArgumentException("Unexpected reference scope: " + scope);
+    }
+    return builder.build();
+  }
+
+  private void assertDereference(
+      FieldReference original,
+      FieldReference dereferenced,
+      Type expectedType,
+      FieldReference.ReferenceSegment nextSegment) {
+    assertEquals(expectedType, dereferenced.getType());
+    assertEquals(List.of(nextSegment, original.segments().get(0)), dereferenced.segments());
+    assertEquals(original.inputExpression(), dereferenced.inputExpression());
+    assertEquals(original.outerReferenceStepsOut(), dereferenced.outerReferenceStepsOut());
+    assertEquals(original.outerReferenceRelReference(), dereferenced.outerReferenceRelReference());
+    assertEquals(
+        original.lambdaParameterReferenceStepsOut(),
+        dereferenced.lambdaParameterReferenceStepsOut());
+
+    io.substrait.proto.Expression.FieldReference originalProto =
+        expressionProtoConverter.toProto(original).getSelection();
+    io.substrait.proto.Expression.FieldReference dereferencedProto =
+        expressionProtoConverter.toProto(dereferenced).getSelection();
+    assertEquals(originalProto.getRootTypeCase(), dereferencedProto.getRootTypeCase());
+    assertEquals(originalProto.getOuterReference(), dereferencedProto.getOuterReference());
+    assertEquals(
+        originalProto.getLambdaParameterReference(),
+        dereferencedProto.getLambdaParameterReference());
+  }
+}

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

@@ -0,0 +1,58 @@
+package io.substrait.isthmus;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.substrait.isthmus.sql.SubstraitCreateStatementParser;
+import io.substrait.plan.PlanProtoConverter;
+import io.substrait.proto.Expression;
+import io.substrait.proto.FilterRel;
+import io.substrait.proto.Plan;
+import org.apache.calcite.prepare.Prepare;
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.junit.jupiter.api.Test;
+
+class CorrelatedNestedFieldTest {
+
+  @Test
+  void nestedOuterFieldKeepsItsCorrelationAnchor() throws SqlParseException {
+    Prepare.CatalogReader catalog =
+        SubstraitCreateStatementParser.processCreateStatementsToCatalog(
+            "CREATE TABLE outer_table (id INTEGER NOT NULL, s ROW(v INTEGER NOT NULL) NOT NULL);"
+                + "CREATE TABLE inner_table (id INTEGER NOT NULL, s ROW(v INTEGER NOT NULL) NOT NULL)");
+    Plan plan =
+        new PlanProtoConverter()
+            .toProto(
+                new SqlToSubstrait()
+                    .convert(
+                        "SELECT o.id FROM outer_table o WHERE EXISTS"
+                            + " (SELECT 1 FROM inner_table i WHERE i.id = o.s.v)",
+                        catalog));
+
+    FilterRel outerFilter =
+        plan.getRelations(0).getRoot().getInput().getProject().getInput().getFilter();
+    FilterRel innerFilter =
+        outerFilter.getCondition().getSubquery().getSetPredicate().getTuples().getFilter();
+    Expression.FieldReference outerField =
+        innerFilter.getCondition().getScalarFunction().getArguments(1).getValue().getSelection();
+
+    assertTrue(outerFilter.getInput().getRead().getCommon().hasRelAnchor());
+    assertTrue(outerField.hasOuterReference());
+    assertTrue(outerField.getOuterReference().hasRelReference());
+    assertEquals(
+        outerFilter.getInput().getRead().getCommon().getRelAnchor(),
+        outerField.getOuterReference().getRelReference());
+    assertEquals(1, outerField.getDirectReference().getStructField().getField());
+    assertTrue(outerField.getDirectReference().getStructField().hasChild());
+    assertEquals(
+        0, outerField.getDirectReference().getStructField().getChild().getStructField().getField());
+    assertTrue(
+        innerFilter
+            .getCondition()
+            .getScalarFunction()
+            .getArguments(0)
+            .getValue()
+            .getSelection()
+            .hasRootReference());
+  }
+}