diff --git a/core/src/main/java/io/substrait/expression/FieldReference.java b/core/src/main/java/io/substrait/expression/FieldReference.java index 92e726fe3..e6a7a21f2 100644 --- a/core/src/main/java/io/substrait/expression/FieldReference.java +++ b/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(); } diff --git a/core/src/test/java/io/substrait/expression/FieldReferenceDereferenceTest.java b/core/src/test/java/io/substrait/expression/FieldReferenceDereferenceTest.java new file mode 100644 index 000000000..31913788e --- /dev/null +++ b/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()); + } +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/CorrelatedNestedFieldTest.java b/isthmus/src/test/java/io/substrait/isthmus/CorrelatedNestedFieldTest.java new file mode 100644 index 000000000..2287bde01 --- /dev/null +++ b/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()); + } +}