diff --git a/core/src/main/java/io/substrait/expression/FieldReference.java b/core/src/main/java/io/substrait/expression/FieldReference.java index 4d525bcbc..92e726fe3 100644 --- a/core/src/main/java/io/substrait/expression/FieldReference.java +++ b/core/src/main/java/io/substrait/expression/FieldReference.java @@ -40,10 +40,31 @@ public abstract class FieldReference implements Expression { /** * Returns the number of subquery levels stepped out of for an outer reference, if applicable. * + *

This offset-based mechanism is used for tree-shaped plans. For plans where a relation is + * shared via a {@code ReferenceRel} (making the reference target ambiguous), use the id-based + * {@link #outerReferenceRelReference()} instead. The two are mutually exclusive alternatives — + * they map to a single protobuf {@code oneof}, so at most one may be set. + * * @return the optional number of steps out */ public abstract Optional outerReferenceStepsOut(); + /** + * Returns the plan-wide unique {@code relAnchor} of the relation this outer reference is rooted + * on, if applicable. + * + *

This id-based mechanism resolves outer references unambiguously in DAG-shaped plans where a + * relation is shared via a {@code ReferenceRel} and the offset-based {@link + * #outerReferenceStepsOut()} would be ambiguous. The value must match a {@link + * io.substrait.relation.Rel#getRelAnchor()} defined elsewhere in the plan. + * + *

This and {@link #outerReferenceStepsOut()} are mutually exclusive alternatives — they map to + * a single protobuf {@code oneof}, so at most one may be set. + * + * @return the optional referenced {@code relAnchor} + */ + public abstract Optional outerReferenceRelReference(); + /** * Returns the number of lambda nesting levels stepped out of for a lambda parameter reference, if * applicable. @@ -73,16 +94,24 @@ public R accept( } /** - * Validates that a field reference is not simultaneously an outer reference and a lambda - * parameter reference. + * Validates that at most one reference form is set. The offset-based outer reference ({@link + * #outerReferenceStepsOut()}), the id-based outer reference ({@link + * #outerReferenceRelReference()}) and the lambda parameter reference ({@link + * #lambdaParameterReferenceStepsOut()}) are mutually exclusive; the two outer-reference forms in + * particular map to a single protobuf {@code oneof} and cannot be combined. * - * @throws IllegalArgumentException if both step-out values are set + * @throws IllegalArgumentException if more than one reference form is set */ @Value.Check protected void check() { - if (outerReferenceStepsOut().isPresent() && lambdaParameterReferenceStepsOut().isPresent()) { + int formsSet = + (outerReferenceStepsOut().isPresent() ? 1 : 0) + + (outerReferenceRelReference().isPresent() ? 1 : 0) + + (lambdaParameterReferenceStepsOut().isPresent() ? 1 : 0); + if (formsSet > 1) { throw new IllegalArgumentException( - "FieldReference cannot have both outerReferenceStepsOut and lambdaParameterReferenceStepsOut set"); + "FieldReference can set at most one of outerReferenceStepsOut, " + + "outerReferenceRelReference and lambdaParameterReferenceStepsOut"); } } @@ -95,16 +124,19 @@ public boolean isSimpleRootReference() { return segments().size() == 1 && !inputExpression().isPresent() && !outerReferenceStepsOut().isPresent() + && !outerReferenceRelReference().isPresent() && !lambdaParameterReferenceStepsOut().isPresent(); } /** - * Returns whether this reference steps out into an enclosing (outer) query. + * Returns whether this reference steps out into an enclosing (outer) query, via either the + * offset-based ({@link #outerReferenceStepsOut()}) or id-based ({@link + * #outerReferenceRelReference()}) mechanism. * * @return {@code true} if this is an outer reference */ public boolean isOuterReference() { - return outerReferenceStepsOut().orElse(0) > 0; + return outerReferenceStepsOut().orElse(0) > 0 || outerReferenceRelReference().isPresent(); } /** @@ -234,6 +266,24 @@ public static FieldReference newRootStructOuterReference( .build(); } + /** + * Creates an id-based reference to a field of an enclosing (outer) query's root struct, resolved + * via the referenced relation's {@link Rel#getRelAnchor()} rather than a subquery-level offset. + * + * @param index the struct field index + * @param knownType the known type of the referenced field + * @param relReference the {@code relAnchor} of the relation this field reference is rooted on + * @return the field reference + */ + public static FieldReference newRootStructOuterReferenceByRelReference( + int index, Type knownType, int relReference) { + return ImmutableFieldReference.builder() + .addSegments(StructField.of(index)) + .type(knownType) + .outerReferenceRelReference(relReference) + .build(); + } + /** * Creates a reference to a field of a single input relation by overall field index. * diff --git a/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java b/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java index 0ebf340dc..0feba3f47 100644 --- a/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java +++ b/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java @@ -657,6 +657,12 @@ public Expression visit(FieldReference expr, EmptyVisitationContext context) { if (expr.inputExpression().isPresent()) { out.setExpression(toProto(expr.inputExpression().get())); + } else if (expr.outerReferenceRelReference().isPresent()) { + // steps_out and rel_reference are mutually exclusive (a single OuterReference oneof); a + // FieldReference carries at most one, enforced by FieldReference#check(). + out.setOuterReference( + io.substrait.proto.Expression.FieldReference.OuterReference.newBuilder() + .setRelReference(expr.outerReferenceRelReference().get())); } else if (expr.outerReferenceStepsOut().isPresent()) { out.setOuterReference( io.substrait.proto.Expression.FieldReference.OuterReference.newBuilder() diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java index 76a65a803..b6384a648 100644 --- a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java +++ b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java @@ -88,10 +88,23 @@ public FieldReference from(io.substrait.proto.Expression.FieldReference referenc return FieldReference.ofRoot( rootType, getDirectReferenceSegments(reference.getDirectReference())); case OUTER_REFERENCE: - return FieldReference.newRootStructOuterReference( - reference.getDirectReference().getStructField().getField(), - rootType, - reference.getOuterReference().getStepsOut()); + { + io.substrait.proto.Expression.FieldReference.OuterReference outerReference = + reference.getOuterReference(); + int field = reference.getDirectReference().getStructField().getField(); + switch (outerReference.getOuterReferenceTypeCase()) { + case STEPS_OUT: + return FieldReference.newRootStructOuterReference( + field, rootType, outerReference.getStepsOut()); + case REL_REFERENCE: + return FieldReference.newRootStructOuterReferenceByRelReference( + field, rootType, outerReference.getRelReference()); + case OUTERREFERENCETYPE_NOT_SET: + default: + throw new IllegalArgumentException( + "Unhandled outer reference type: " + outerReference.getOuterReferenceTypeCase()); + } + } case LAMBDA_PARAMETER_REFERENCE: { io.substrait.proto.Expression.FieldReference.LambdaParameterReference lambdaParamRef = diff --git a/core/src/main/java/io/substrait/relation/OuterReferenceConverter.java b/core/src/main/java/io/substrait/relation/OuterReferenceConverter.java new file mode 100644 index 000000000..999352a8c --- /dev/null +++ b/core/src/main/java/io/substrait/relation/OuterReferenceConverter.java @@ -0,0 +1,391 @@ +package io.substrait.relation; + +import io.substrait.expression.Expression; +import io.substrait.expression.FieldReference; +import io.substrait.expression.ImmutableFieldReference; +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: + * + *

+ * + *

The two forms are semantically equivalent for tree-shaped plans and are mutually exclusive on + * the wire (they share a single protobuf {@code oneof}). This class rewrites every outer reference + * from one form to the other so an integration can work in a single encoding regardless of which + * form a plan arrives in. Anchor assignment requires plan-wide context, which is why this lives in + * core rather than being duplicated per integration. + * + *

Supported scope: outer references whose binding relation is the input of a single-input + * expression host ({@link Filter}, {@link Project}). This covers correlated scalar/IN/EXISTS + * subqueries as produced by the SQL integrations. Multi-input scopes (a correlated reference into a + * join/set condition) and shared subtrees ({@code ReferenceRel}) are not representable as a single + * binding relation and cause an {@link UnsupportedOperationException}. + * + *

Anchor allocation starts at {@code 1}; callers must not pass {@code toIdBased} a plan + * that already carries {@code rel_anchor}s. + */ +public final class OuterReferenceConverter { + + private OuterReferenceConverter() {} + + private enum Direction { + TO_ID, + TO_STEPS_OUT + } + + /** + * Rewrites offset-based outer references ({@code steps_out}) as id-based references ({@code + * rel_reference}), assigning a {@link Rel#getRelAnchor() rel anchor} to each referenced relation. + * + * @param root the relation tree to convert + * @return an equivalent tree using id-based outer references + */ + public static Rel toIdBased(Rel root) { + return convert(root, Direction.TO_ID); + } + + /** + * Rewrites id-based outer references ({@code rel_reference}) as offset-based references ({@code + * steps_out}), removing the {@link Rel#getRelAnchor() rel anchors} they resolved to. + * + * @param root the relation tree to convert + * @return an equivalent tree using offset-based outer references + */ + public static Rel toStepsOut(Rel root) { + return convert(root, Direction.TO_STEPS_OUT); + } + + /** + * Applies {@link #toIdBased(Rel)} to every root of the given plan. + * + * @param plan the plan to convert + * @return an equivalent plan using id-based outer references + */ + public static Plan toIdBased(Plan plan) { + return convert(plan, Direction.TO_ID); + } + + /** + * Applies {@link #toStepsOut(Rel)} to every root of the given plan. + * + * @param plan the plan to convert + * @return an equivalent plan using offset-based outer references + */ + public static Plan toStepsOut(Plan plan) { + return convert(plan, Direction.TO_STEPS_OUT); + } + + 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 roots = new ArrayList<>(plan.getRoots().size()); + for (Plan.Root root : plan.getRoots()) { + 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) { + 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); + } + + /** Mutable state shared between the paired relation and expression rewriters. */ + private static final class State { + final Direction direction; + + /** + * Stack of enclosing scope relations, one entry per subquery boundary. A {@code null} entry + * 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 outerScopes = new ArrayList<>(); + + /** + * Anchors assigned to binding relations while converting to the id-based form (by identity). + */ + final Map anchorByRel = new IdentityHashMap<>(); + + /** Anchors that were resolved to a {@code steps_out} value and should be stripped from rels. */ + final Set resolvedAnchors = new HashSet<>(); + + /** + * 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; + + State(Direction direction) { + this.direction = direction; + } + + int allocateAnchor(Rel binding) { + Integer existing = anchorByRel.get(binding); + if (existing != null) { + return existing; + } + int anchor = nextAnchor++; + anchorByRel.put(binding, anchor); + return anchor; + } + } + + /** + * Relation rewriter that maintains the scope stack and stamps/strips {@code rel_anchor}s. Only + * single-input expression hosts are handled explicitly; every other relation is traversed by the + * copy-on-write base class. + */ + private static final class RelRewriter extends RelCopyOnWriteVisitor { + private final State state; + + RelRewriter(State state) { + super(relVisitor -> new ExprRewriter(relVisitor, state)); + this.state = state; + } + + @Override + public Optional visit(Filter filter, EmptyVisitationContext context) { + // Expressions first, so a reference binding to this filter's input is discovered before the + // input is (re)built and (for the id-based direction) stamped with its anchor. + Optional condition = + rewriteInScope( + filter.getInput(), + () -> filter.getCondition().accept(getExpressionCopyOnWriteVisitor(), context)); + Rel newInput = rewriteInput(filter.getInput()); + + if (!condition.isPresent() && newInput == filter.getInput()) { + return Optional.empty(); + } + return Optional.of( + Filter.builder() + .from(filter) + .input(newInput) + .condition(condition.orElse(filter.getCondition())) + .build()); + } + + @Override + public Optional visit(Project project, EmptyVisitationContext context) { + Optional> expressions = + rewriteInScope( + project.getInput(), () -> visitExprList(project.getExpressions(), context)); + Rel newInput = rewriteInput(project.getInput()); + + if (!expressions.isPresent() && newInput == project.getInput()) { + return Optional.empty(); + } + return Optional.of( + Project.builder() + .from(project) + .input(newInput) + .expressions(expressions.orElse(project.getExpressions())) + .build()); + } + + @Override + public Optional visit(NamedDdl ddl, EmptyVisitationContext context) { + // The copy-on-write base throws for DDL; a view definition is an independent (top-level) + // query that may itself contain correlated subqueries, so traverse it rather than throw. + if (!ddl.getViewDefinition().isPresent()) { + return Optional.empty(); + } + return ddl.getViewDefinition() + .get() + .accept(this, context) + .map(definition -> NamedDdl.builder().from(ddl).viewDefinition(definition).build()); + } + + /** + * 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 rewriteInScope(Rel input, Supplier rewrite) { + Rel saved = state.currentInput; + state.currentInput = input; + try { + return rewrite.get(); + } finally { + state.currentInput = saved; + } + } + + /** + * Rebuilds the input and applies the anchor stamp/strip decided while rewriting expressions. + */ + private Rel rewriteInput(Rel originalInput) { + Rel rebuilt = + originalInput.accept(this, EmptyVisitationContext.INSTANCE).orElse(originalInput); + if (state.direction == Direction.TO_ID) { + Integer anchor = state.anchorByRel.get(originalInput); + if (anchor != null) { + return rebuilt.withRelAnchor(anchor); + } + } else { + Optional anchor = originalInput.getRelAnchor(); + if (anchor.isPresent() && state.resolvedAnchors.contains(anchor.get())) { + return rebuilt.withRelAnchor(Optional.empty()); + } + } + return rebuilt; + } + } + + /** Expression rewriter that pushes/pops subquery scopes and rewrites outer field references. */ + private static final class ExprRewriter extends ExpressionCopyOnWriteVisitor { + private final State state; + + ExprRewriter(RelCopyOnWriteVisitor relVisitor, State state) { + super(relVisitor); + this.state = state; + } + + @Override + public Optional visit(FieldReference fieldReference, EmptyVisitationContext context) + throws RuntimeException { + if (state.direction == Direction.TO_ID + && fieldReference.outerReferenceStepsOut().isPresent()) { + int anchor = anchorForStepsOut(fieldReference.outerReferenceStepsOut().get()); + return Optional.of( + ImmutableFieldReference.builder() + .from(fieldReference) + .outerReferenceStepsOut(Optional.empty()) + .outerReferenceRelReference(anchor) + .build()); + } + if (state.direction == Direction.TO_STEPS_OUT + && fieldReference.outerReferenceRelReference().isPresent()) { + int stepsOut = stepsOutForAnchor(fieldReference.outerReferenceRelReference().get()); + return Optional.of( + ImmutableFieldReference.builder() + .from(fieldReference) + .outerReferenceRelReference(Optional.empty()) + .outerReferenceStepsOut(stepsOut) + .build()); + } + + // Non-outer references: preserve all attributes, rewriting only a nested input expression. + if (fieldReference.inputExpression().isPresent()) { + Optional newInput = + fieldReference.inputExpression().get().accept(this, context); + if (newInput.isPresent()) { + return Optional.of( + ImmutableFieldReference.builder() + .from(fieldReference) + .inputExpression(newInput) + .build()); + } + } + return Optional.empty(); + } + + private int anchorForStepsOut(int stepsOut) { + int index = state.outerScopes.size() - stepsOut; + if (index < 0 || index >= state.outerScopes.size()) { + throw new IllegalArgumentException( + "Outer reference steps_out=" + stepsOut + " exceeds the enclosing subquery depth"); + } + 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 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); + } + + private int stepsOutForAnchor(int relReference) { + for (int i = state.outerScopes.size() - 1; i >= 0; i--) { + Rel scope = state.outerScopes.get(i); + if (scope != null + && scope.getRelAnchor().isPresent() + && scope.getRelAnchor().get() == relReference) { + state.resolvedAnchors.add(relReference); + return state.outerScopes.size() - i; + } + } + throw new UnsupportedOperationException( + "Cannot resolve outer reference rel_reference=" + + relReference + + " to an enclosing scope; the referenced relation is not an ancestor (shared " + + "subtrees / ReferenceRel are not supported)"); + } + + @Override + public Optional visit( + Expression.ScalarSubquery scalarSubquery, EmptyVisitationContext context) { + return withSubqueryScope(() -> super.visit(scalarSubquery, context)); + } + + @Override + public Optional visit( + Expression.SetPredicate setPredicate, EmptyVisitationContext context) { + return withSubqueryScope(() -> super.visit(setPredicate, context)); + } + + @Override + public Optional visit( + Expression.InPredicate inPredicate, EmptyVisitationContext context) { + // needles are evaluated in the current (outer) scope; the haystack is the subquery boundary. + Optional> needles = visitExprList(inPredicate.needles(), context); + Optional haystack = + withSubqueryScope( + () -> inPredicate.haystack().accept(getRelCopyOnWriteVisitor(), context)); + + if (!needles.isPresent() && !haystack.isPresent()) { + return Optional.empty(); + } + return Optional.of( + Expression.InPredicate.builder() + .from(inPredicate) + .haystack(haystack.orElse(inPredicate.haystack())) + .needles(needles.orElse(inPredicate.needles())) + .build()); + } + + /** Pushes {@code currentInput} as a new subquery scope, runs the action, then pops it. */ + private T withSubqueryScope(Supplier 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 { + state.outerScopes.remove(state.outerScopes.size() - 1); + state.currentInput = savedInput; + } + } + } +} diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java index a58e2e3b4..68403a4eb 100644 --- a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java +++ b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java @@ -266,7 +266,8 @@ protected NamedWrite newNamedWrite(final WriteRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); @@ -296,7 +297,8 @@ protected Rel newExtensionWrite(final WriteRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); @@ -340,7 +342,8 @@ protected NamedDdl newNamedDdl(final DdlRel rel) { .viewDefinition(optionalViewDefinition(rel)) .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); @@ -369,7 +372,8 @@ protected ExtensionDdl newExtensionDdl(final DdlRel rel) { .viewDefinition(optionalViewDefinition(rel)) .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); @@ -471,7 +475,8 @@ protected Filter newFilter(FilterRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -521,7 +526,8 @@ protected ExtensionLeaf newExtensionLeaf(ExtensionLeafRel rel) { ExtensionLeaf.from(detail) .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); return builder.build(); } @@ -538,7 +544,8 @@ protected ExtensionSingle newExtensionSingle(ExtensionSingleRel rel) { ExtensionSingle.from(detail, input) .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); return builder.build(); } @@ -555,7 +562,8 @@ protected ExtensionMulti newExtensionMulti(ExtensionMultiRel rel) { ExtensionMulti.from(detail, inputs) .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasDetail()) { builder.detail(detailFromExtensionMultiRel(rel.getDetail())); } @@ -593,7 +601,8 @@ protected NamedScan newNamedScan(ReadRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -617,7 +626,8 @@ protected ExtensionTable newExtensionTable(final ReadRel rel) { .projection(optionalMaskExpression(rel)) .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -659,7 +669,8 @@ protected LocalFiles newLocalFiles(ReadRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -779,7 +790,8 @@ protected VirtualTableScan newVirtualTable(ReadRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -804,7 +816,8 @@ protected Fetch newFetch(FetchRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -832,7 +845,8 @@ protected Project newProject(ProjectRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -878,7 +892,8 @@ protected Expand newExpand(ExpandRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); return builder.build(); } @@ -929,7 +944,8 @@ protected Aggregate newAggregate(AggregateRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -962,7 +978,8 @@ protected Sort newSort(SortRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -996,7 +1013,8 @@ protected Join newJoin(JoinRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1016,7 +1034,8 @@ protected Rel newCross(CrossRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())); + .remap(optionalRelmap(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1040,7 +1059,8 @@ protected Set newSet(SetRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1084,7 +1104,8 @@ protected Rel newHashJoin(HashJoinRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1129,7 +1150,8 @@ protected Rel newMergeJoin(MergeJoinRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1221,7 +1243,8 @@ protected NestedLoopJoin newNestedLoopJoin(NestedLoopJoinRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1264,7 +1287,8 @@ protected ConsistentPartitionWindow newConsistentPartitionWindow( builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1323,7 +1347,8 @@ protected ScatterExchange newScatterExchange(ExchangeRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1353,7 +1378,8 @@ protected SingleBucketExchange newSingleBucketExchange(ExchangeRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1384,7 +1410,8 @@ protected MultiBucketExchange newMultiBucketExchange(ExchangeRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1412,7 +1439,8 @@ protected RoundRobinExchange newRoundRobinExchange(ExchangeRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1439,7 +1467,8 @@ protected BroadcastExchange newBroadcastExchange(ExchangeRel rel) { builder .commonExtension(optionalAdvancedExtension(rel.getCommon())) .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); + .hint(optionalHint(rel.getCommon())) + .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } @@ -1481,6 +1510,17 @@ protected static Optional optionalRelmap(io.substrait.proto.RelCommon relCommon.hasEmit() ? Rel.Remap.of(relCommon.getEmit().getOutputMappingList()) : null); } + /** + * Converts the {@link io.substrait.proto.RelCommon#getRelAnchor()} field to its POJO + * representation. + * + * @param relCommon the protobuf value to convert + * @return the converted result + */ + protected static Optional optionalRelAnchor(io.substrait.proto.RelCommon relCommon) { + return relCommon.hasRelAnchor() ? Optional.of(relCommon.getRelAnchor()) : Optional.empty(); + } + /** * Converts the corresponding protobuf message to its POJO representation. * diff --git a/core/src/main/java/io/substrait/relation/Rel.java b/core/src/main/java/io/substrait/relation/Rel.java index 74915e83b..ca4a9b5e6 100644 --- a/core/src/main/java/io/substrait/relation/Rel.java +++ b/core/src/main/java/io/substrait/relation/Rel.java @@ -25,6 +25,54 @@ public interface Rel { */ Optional getCommonExtension(); + /** + * Returns the plan-wide unique anchor identifying this relation, if set. + * + *

This corresponds to {@link io.substrait.proto.RelCommon#getRelAnchor()} and is required when + * this relation is the binding point for an id-based outer reference (see {@link + * io.substrait.expression.FieldReference#outerReferenceRelReference()}). When set it must be + * unique across all relations within a plan and {@code >= 1}. + * + * @return the optional relation anchor + */ + Optional getRelAnchor(); + + /** + * Returns a copy of this relation with its {@link #getRelAnchor() relation anchor} set to the + * given value. + * + *

Overridden by the generated Immutables {@code withRelAnchor(int)} on every concrete + * relation, this provides a type-agnostic way to stamp an anchor onto an arbitrary relation (used + * by {@link OuterReferenceConverter} when assigning anchors during {@code steps_out → + * rel_reference} conversion). Custom {@link Rel} implementations that are not Immutables-backed + * inherit this throwing default. + * + * @param relAnchor the plan-wide unique anchor to set (must be {@code >= 1}) + * @return a copy of this relation carrying the given anchor + */ + default Rel withRelAnchor(int relAnchor) { + throw new UnsupportedOperationException( + getClass() + " does not support setting a relation anchor"); + } + + /** + * Returns a copy of this relation with its {@link #getRelAnchor() relation anchor} set to the + * given optional value ({@link Optional#empty()} clears it). + * + *

Like {@link #withRelAnchor(int)}, this is overridden by the generated Immutables {@code + * withRelAnchor(Optional)} and provides a type-agnostic way to set or clear the anchor (used by + * {@link OuterReferenceConverter} when stripping anchors during {@code rel_reference → steps_out} + * conversion). Custom {@link Rel} implementations that are not Immutables-backed inherit this + * throwing default. + * + * @param relAnchor the anchor to set, or empty to clear it + * @return a copy of this relation carrying the given anchor + */ + default Rel withRelAnchor(Optional relAnchor) { + throw new UnsupportedOperationException( + getClass() + " does not support setting a relation anchor"); + } + /** * Returns the record type (schema) produced by this relation. * diff --git a/core/src/main/java/io/substrait/relation/RelProtoConverter.java b/core/src/main/java/io/substrait/relation/RelProtoConverter.java index 391c8e66b..aa1af0bee 100644 --- a/core/src/main/java/io/substrait/relation/RelProtoConverter.java +++ b/core/src/main/java/io/substrait/relation/RelProtoConverter.java @@ -942,6 +942,8 @@ private RelCommon common(io.substrait.relation.Rel rel) { .ifPresent( extension -> builder.setAdvancedExtension(extensionProtoConverter.toProto(extension))); + rel.getRelAnchor().ifPresent(builder::setRelAnchor); + io.substrait.relation.Rel.Remap remap = rel.getRemap().orElse(null); if (remap != null) { builder.setEmit(RelCommon.Emit.newBuilder().addAllOutputMapping(remap.indices())); diff --git a/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java b/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java new file mode 100644 index 000000000..b42dbaa66 --- /dev/null +++ b/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java @@ -0,0 +1,315 @@ +package io.substrait.relation; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +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; +import org.junit.jupiter.api.Test; + +/** + * Tests {@link OuterReferenceConverter}, which translates outer references between the offset-based + * ({@code steps_out}) and id-based ({@code rel_anchor}/{@code rel_reference}) encodings. + */ +class OuterReferenceConverterTest extends TestBase { + + private final Rel customerTableScan = + sb.namedScan( + List.of("customer"), + List.of("c_custkey", "c_nationkey"), + List.of(TypeCreator.REQUIRED.I64, TypeCreator.REQUIRED.I64)); + + private final Rel orderTableScan = + sb.namedScan( + List.of("orders"), + List.of("o_orderkey", "o_custkey"), + List.of(TypeCreator.REQUIRED.I64, TypeCreator.REQUIRED.I64)); + + private final Rel nationTableScan = + sb.namedScan( + List.of("nation"), + List.of("n_nationkey", "n_name"), + List.of(TypeCreator.REQUIRED.I64, TypeCreator.REQUIRED.STRING)); + + /** SELECT o_orderkey, (SELECT c_nationkey FROM customer WHERE c_custkey = orders.o_custkey). */ + private Rel oneStepPlan() { + 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)), + orderTableScan); + } + + /** + * Nested subquery whose innermost outer reference steps out two boundaries to the orders scan. + */ + private Rel twoStepPlan() { + 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(2)), + sb.filter( + input2 -> + sb.equal( + sb.fieldReference(input2, 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), + FieldReference.newRootStructOuterReference( + 1, TypeCreator.REQUIRED.I64, 2)), + customerTableScan)), + TypeCreator.NULLABLE.I64)), + nationTableScan)), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2, 3)), + orderTableScan); + } + + @Test + void oneStepToIdBasedStampsAnchorAndReference() { + Rel stepsOut = oneStepPlan(); + Rel idBased = OuterReferenceConverter.toIdBased(stepsOut); + + assertNotEquals(stepsOut, idBased); + + // The outer project's input (the orders scan) is the binding relation and gets anchor 1. + Project outerProject = (Project) idBased; + assertEquals(1, outerProject.getInput().getRelAnchor().orElseThrow(AssertionError::new)); + + // The correlated reference now uses rel_reference instead of steps_out. + FieldReference outerRef = outerReferenceOf(outerProject); + assertEquals(1, outerRef.outerReferenceRelReference().orElseThrow(AssertionError::new)); + assertFalse(outerRef.outerReferenceStepsOut().isPresent()); + } + + @Test + void oneStepRoundTripsBackToStepsOut() { + Rel stepsOut = oneStepPlan(); + assertEquals( + stepsOut, OuterReferenceConverter.toStepsOut(OuterReferenceConverter.toIdBased(stepsOut))); + } + + @Test + void twoStepRoundTripsThroughIdBased() { + Rel stepsOut = twoStepPlan(); + Rel idBased = OuterReferenceConverter.toIdBased(stepsOut); + + assertNotEquals(stepsOut, idBased); + // The innermost reference steps out two levels to the orders scan, which gets the anchor. + assertEquals(1, ((Project) idBased).getInput().getRelAnchor().orElseThrow(AssertionError::new)); + assertEquals(stepsOut, OuterReferenceConverter.toStepsOut(idBased)); + } + + @Test + void sameScopeReferencesShareAnchor() { + // Two correlated references at the same subquery boundary must resolve to the same anchor. + 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)), + sb.filter( + input2 -> + sb.and( + sb.equal( + sb.fieldReference(input2, 0), + FieldReference.newRootStructOuterReference( + 0, TypeCreator.REQUIRED.I64, 1)), + sb.equal( + sb.fieldReference(input2, 1), + FieldReference.newRootStructOuterReference( + 1, TypeCreator.REQUIRED.I64, 1))), + customerTableScan)), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2, 3)), + orderTableScan); + + Rel idBased = OuterReferenceConverter.toIdBased(stepsOut); + assertEquals(1, ((Project) idBased).getInput().getRelAnchor().orElseThrow(AssertionError::new)); + assertEquals(stepsOut, OuterReferenceConverter.toStepsOut(idBased)); + } + + @Test + void stepsOutBeyondDepthIsRejected() { + Rel invalid = + sb.project( + input -> + List.of( + 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), + // steps_out=5 exceeds the single enclosing boundary + FieldReference.newRootStructOuterReference( + 1, TypeCreator.REQUIRED.I64, 5)), + customerTableScan)), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2)), + orderTableScan); + + assertThrows(IllegalArgumentException.class, () -> OuterReferenceConverter.toIdBased(invalid)); + } + + @Test + void unresolvableRelReferenceIsRejected() { + // An id-based reference to an anchor that is not an enclosing scope cannot become steps_out. + Rel dangling = + sb.project( + input -> + List.of( + 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.newRootStructOuterReferenceByRelReference( + 1, TypeCreator.REQUIRED.I64, 999)), + customerTableScan)), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2)), + orderTableScan); + + assertThrows( + 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 = + (Expression.ScalarSubquery) outerProject.getExpressions().get(1); + Filter filter = (Filter) ((Project) subquery.input()).getInput(); + Expression.ScalarFunctionInvocation equal = + (Expression.ScalarFunctionInvocation) filter.getCondition(); + return (FieldReference) equal.arguments().get(1); + } +} diff --git a/core/src/test/java/io/substrait/type/proto/OuterReferenceRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/OuterReferenceRoundtripTest.java new file mode 100644 index 000000000..672467fc7 --- /dev/null +++ b/core/src/test/java/io/substrait/type/proto/OuterReferenceRoundtripTest.java @@ -0,0 +1,71 @@ +package io.substrait.type.proto; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +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.ImmutableFieldReference; +import io.substrait.expression.proto.ProtoExpressionConverter; +import io.substrait.type.Type; +import org.junit.jupiter.api.Test; + +/** + * Round-trip tests for the two outer-reference resolution mechanisms introduced with Substrait + * v0.89.0: the offset-based {@link FieldReference#outerReferenceStepsOut()} and the id-based {@link + * FieldReference#outerReferenceRelReference()}. + * + *

Outer references derive their type from the converter's root type on the way back from proto, + * so a converter configured with a matching root type is used instead of {@link + * TestBase#protoExpressionConverter}. + */ +class OuterReferenceRoundtripTest extends TestBase { + + final Type.Struct outerSchema = + Type.Struct.builder().nullable(false).addFields(R.I64, R.STRING).build(); + + final ProtoExpressionConverter protoExpressionConverterWithRoot = + new ProtoExpressionConverter(functionCollector, extensions, outerSchema, protoRelConverter); + + private void verifyOuterReferenceRoundTrip(FieldReference reference) { + io.substrait.proto.Expression proto = expressionProtoConverter.toProto(reference); + Expression returned = protoExpressionConverterWithRoot.from(proto); + assertEquals(reference, returned); + } + + @Test + void offsetBasedOuterReference() { + FieldReference reference = FieldReference.newRootStructOuterReference(0, outerSchema, 1); + + assertTrue(reference.isOuterReference()); + assertFalse(reference.isSimpleRootReference()); + verifyOuterReferenceRoundTrip(reference); + } + + @Test + void idBasedOuterReference() { + FieldReference reference = + FieldReference.newRootStructOuterReferenceByRelReference(1, outerSchema, 42); + + assertTrue(reference.isOuterReference()); + assertFalse(reference.isSimpleRootReference()); + verifyOuterReferenceRoundTrip(reference); + } + + /** + * The two outer-reference forms map to a single protobuf {@code oneof} and are therefore mutually + * exclusive: a {@link FieldReference} may not carry both at once. + */ + @Test + void outerReferenceFormsAreMutuallyExclusive() { + ImmutableFieldReference.Builder both = + ImmutableFieldReference.builder() + .from(FieldReference.newRootStructOuterReference(0, outerSchema, 3)) + .outerReferenceRelReference(42); + + assertThrows(IllegalArgumentException.class, both::build); + } +} diff --git a/core/src/test/java/io/substrait/type/proto/RelAnchorRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/RelAnchorRoundtripTest.java new file mode 100644 index 000000000..8377c2ea8 --- /dev/null +++ b/core/src/test/java/io/substrait/type/proto/RelAnchorRoundtripTest.java @@ -0,0 +1,59 @@ +package io.substrait.type.proto; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import io.substrait.TestBase; +import io.substrait.relation.Filter; +import io.substrait.relation.Project; +import io.substrait.relation.Rel; +import java.util.Arrays; +import java.util.Collections; +import org.junit.jupiter.api.Test; + +/** + * Round-trip tests for the {@link Rel#getRelAnchor()} field (Substrait v0.89.0 {@code + * RelCommon.rel_anchor}), the plan-wide unique identifier used as the binding point for id-based + * outer references. + */ +class RelAnchorRoundtripTest extends TestBase { + + final Rel baseTable = + sb.namedScan( + Collections.singletonList("test_table"), + Arrays.asList("id", "name"), + Arrays.asList(R.I64, R.STRING)); + + @Test + void relAnchorOnProject() { + Rel projection = + Project.builder() + .input(baseTable) + .relAnchor(7) + .addExpressions(sb.fieldReference(baseTable, 0)) + .build(); + + assertEquals(7, projection.getRelAnchor().orElseThrow(AssertionError::new)); + verifyRoundTrip(projection); + } + + @Test + void relAnchorOnFilter() { + Rel filter = + Filter.builder() + .input(baseTable) + .relAnchor(1) + .condition(sb.equal(sb.fieldReference(baseTable, 0), sb.fieldReference(baseTable, 0))) + .build(); + + verifyRoundTrip(filter); + } + + @Test + void relAnchorAbsentByDefault() { + Rel projection = + Project.builder().input(baseTable).addExpressions(sb.fieldReference(baseTable, 0)).build(); + + assertEquals(false, projection.getRelAnchor().isPresent()); + verifyRoundTrip(projection); + } +} diff --git a/isthmus/src/main/java/io/substrait/isthmus/OuterReferenceResolver.java b/isthmus/src/main/java/io/substrait/isthmus/OuterReferenceResolver.java index ca24409da..6b6ce6da8 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/OuterReferenceResolver.java +++ b/isthmus/src/main/java/io/substrait/isthmus/OuterReferenceResolver.java @@ -8,199 +8,131 @@ import org.apache.calcite.rel.core.CorrelationId; import org.apache.calcite.rel.core.Filter; import org.apache.calcite.rel.core.Project; -import org.apache.calcite.rex.RexCorrelVariable; -import org.apache.calcite.rex.RexFieldAccess; import org.apache.calcite.rex.RexNode; import org.apache.calcite.rex.RexShuttle; import org.apache.calcite.rex.RexSubQuery; import org.apache.calcite.rex.RexUtil.SubQueryCollector; /** - * Resolve correlated variables and compute a depth map for {@link RexFieldAccess}. + * Assigns id-based outer-reference anchors to a {@link RelNode} tree. * - *

Traverses a {@link RelNode} tree and: + *

Calcite models correlation with {@link CorrelationId}s, each bound by a single relation (a + * {@link Correlate}'s left input, or the input of a {@link Filter}/{@link Project} that declares + * the id). This maps directly onto Substrait's id-based outer references: the bound relation + * carries a {@code rel_anchor} ({@link io.substrait.relation.Rel#getRelAnchor()}) and every + * correlated field reference carries the matching {@code rel_reference} ({@link + * io.substrait.expression.FieldReference#outerReferenceRelReference()}). + * + *

This resolver walks the tree (including subquery relations) and produces: * *

* - * See OuterReferenceResolver.md for details on how the depth map is computed. + * Anchors are plan-wide unique and assigned per bound relation, so multiple correlation ids sharing + * a relation share its anchor. */ public class OuterReferenceResolver extends RelNodeVisitor { - private final Map nestedDepth; - private final Map fieldAccessDepthMap; + private final Map anchorByCorrelationId = new HashMap<>(); + private final Map anchorByTarget = new IdentityHashMap<>(); private final RexVisitor rexVisitor = new RexVisitor(this); - /** Creates a new resolver with empty depth tracking maps. */ - public OuterReferenceResolver() { - nestedDepth = new HashMap<>(); - fieldAccessDepthMap = new IdentityHashMap<>(); - } + private int nextAnchor = 1; + + /** Creates a new resolver with empty anchor maps. */ + public OuterReferenceResolver() {} /** - * Applies the resolver to a {@link RelNode} tree, computing the depth map. + * Walks a {@link RelNode} tree, assigning outer-reference anchors. * - * @param r the root relational node + * @param root the root relational node * @return the same node after traversal - * @throws RuntimeException if the visitor encounters an unrecoverable condition */ - public Map apply(RelNode r) { - reverseAccept(r); - return fieldAccessDepthMap; + public RelNode resolve(RelNode root) { + reverseAccept(root); + return root; } /** - * Visits a {@link Filter}, registering any correlation variables and visiting its condition. + * Returns the anchor a correlated field reference bound to the given correlation id must emit as + * its {@code rel_reference}. * - * @param filter the filter node - * @return the result of {@link RelNodeVisitor#visit(Filter)} - * @throws RuntimeException if traversal fails + * @param id the correlation id + * @return the anchor, or {@code null} if the id was not resolved */ - @Override - public RelNode visit(Filter filter) throws RuntimeException { - for (CorrelationId id : filter.getVariablesSet()) { - nestedDepth.putIfAbsent(id, 0); - } - // Also register any CorrelationIds referenced directly in the condition expression. - // This covers partially-decorrelated plans where a Calcite optimizer (e.g. HepPlanner) - // rewrites correlated IN-subqueries into joins but leaves scalar-aggregate subquery - // correlations as RexFieldAccess($corN.col) inside a Filter condition without a - // surrounding Correlate node. In that case getVariablesSet() is empty but the condition - // still contains RexCorrelVariable references whose IDs must be in nestedDepth before - // rexVisitor.visitFieldAccess() is called. - filter - .getCondition() - .accept( - new RexShuttle() { - @Override - public RexNode visitFieldAccess(RexFieldAccess fieldAccess) { - if (fieldAccess.getReferenceExpr() instanceof RexCorrelVariable) { - CorrelationId id = ((RexCorrelVariable) fieldAccess.getReferenceExpr()).id; - nestedDepth.putIfAbsent(id, 0); - } - return fieldAccess; - } - }); - filter.getCondition().accept(rexVisitor); - return super.visit(filter); + public Integer anchorForCorrelationId(CorrelationId id) { + return anchorByCorrelationId.get(id); } /** - * Visits a {@link Correlate}, handling correlation depth for both sides. - * - *

Special case: the right side is a correlated subquery in the rel tree (not a REX), so we - * manually adjust depth before/after visiting it. + * Returns the anchor to stamp on the given relation, if it is the binding point of any outer + * reference. * - * @param correlate the correlate (correlated join) node - * @return the correlate node - * @throws RuntimeException if traversal fails + * @param rel the relation + * @return the anchor, or {@code null} if the relation binds no outer reference */ - @Override - public RelNode visit(Correlate correlate) throws RuntimeException { - for (CorrelationId id : correlate.getVariablesSet()) { - nestedDepth.putIfAbsent(id, 0); - } - - apply(correlate.getLeft()); - - // Correlated join is a special case. The right-rel is a correlated sub-query but not a REX. So, - // the RexVisitor cannot be applied to it to correctly compute the depth map. Hence, we need to - // manually compute the depth map for the right-rel. - nestedDepth.replaceAll((k, v) -> v + 1); - - apply(correlate.getRight()); // look inside sub-queries - - nestedDepth.replaceAll((k, v) -> v - 1); + public Integer anchorForTarget(RelNode rel) { + return anchorByTarget.get(rel); + } - return correlate; + /** Binds {@code id} to {@code target}, allocating {@code target}'s anchor on first use. */ + private void bind(CorrelationId id, RelNode target) { + int anchor = anchorByTarget.computeIfAbsent(target, t -> nextAnchor++); + anchorByCorrelationId.put(id, anchor); } - /** - * Visits a generic {@link RelNode}, applying the resolver to all the node inputs. - * - * @param other the node to visit - * @return the node - * @throws RuntimeException if traversal fails - */ @Override - public RelNode visitOther(RelNode other) throws RuntimeException { - for (RelNode child : other.getInputs()) { - apply(child); + public RelNode visit(Filter filter) throws RuntimeException { + for (CorrelationId id : filter.getVariablesSet()) { + bind(id, filter.getInput()); } - return other; + filter.getCondition().accept(rexVisitor); + return visitOther(filter); } - /** - * Visits a {@link Project}, registering correlation variables and visiting any subqueries within - * its expressions. - * - * @param project the project node - * @return the result of {@link RelNodeVisitor#visit(Project)} - * @throws RuntimeException if traversal fails - */ @Override public RelNode visit(Project project) throws RuntimeException { for (CorrelationId id : project.getVariablesSet()) { - nestedDepth.putIfAbsent(id, 0); + bind(id, project.getInput()); } - for (RexSubQuery subQuery : SubQueryCollector.collect(project)) { subQuery.accept(rexVisitor); } + return visitOther(project); + } - return super.visit(project); + @Override + public RelNode visit(Correlate correlate) throws RuntimeException { + bind(correlate.getCorrelationId(), correlate.getLeft()); + reverseAccept(correlate.getLeft()); + reverseAccept(correlate.getRight()); + return correlate; } - /** Rex visitor used to track correlation depth within expressions and subqueries. */ + @Override + public RelNode visitOther(RelNode other) throws RuntimeException { + for (RelNode child : other.getInputs()) { + reverseAccept(child); + } + return other; + } + + /** Rex visitor that descends into subquery relations to reach nested correlations. */ private static class RexVisitor extends RexShuttle { final OuterReferenceResolver referenceResolver; - /** - * Creates a new Rex visitor bound to the given reference resolver. - * - * @param referenceResolver the parent resolver maintaining depth maps - */ RexVisitor(OuterReferenceResolver referenceResolver) { this.referenceResolver = referenceResolver; } - /** - * Increments correlation depth when entering a subquery and decrements when exiting. - * - * @param subQuery the subquery expression - * @return the same subquery - */ @Override public RexNode visitSubQuery(RexSubQuery subQuery) { - referenceResolver.nestedDepth.replaceAll((k, v) -> v + 1); - - referenceResolver.apply(subQuery.rel); // look inside sub-queries - - referenceResolver.nestedDepth.replaceAll((k, v) -> v - 1); + referenceResolver.reverseAccept(subQuery.rel); return subQuery; } - - /** - * Records depth for {@link RexFieldAccess} referencing a {@link RexCorrelVariable}. - * - * @param fieldAccess the field access expression - * @return the same field access - */ - @Override - public RexNode visitFieldAccess(RexFieldAccess fieldAccess) { - if (fieldAccess.getReferenceExpr() instanceof RexCorrelVariable) { - CorrelationId id = ((RexCorrelVariable) fieldAccess.getReferenceExpr()).id; - if (referenceResolver.nestedDepth.containsKey(id)) { - referenceResolver.fieldAccessDepthMap.put( - fieldAccess, referenceResolver.nestedDepth.get(id)); - } - } - return fieldAccess; - } } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 0aef61857..d6a176ff6 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -1,9 +1,6 @@ package io.substrait.isthmus; import com.google.common.collect.ImmutableList; -import com.google.common.collect.Range; -import com.google.common.collect.RangeMap; -import com.google.common.collect.TreeRangeMap; import io.substrait.expression.Expression; import io.substrait.expression.Expression.SortDirection; import io.substrait.expression.FunctionArg; @@ -28,6 +25,7 @@ import io.substrait.relation.NamedScan; import io.substrait.relation.NamedUpdate; import io.substrait.relation.NamedWrite; +import io.substrait.relation.OuterReferenceConverter; import io.substrait.relation.Project; import io.substrait.relation.Rel; import io.substrait.relation.Rel.Remap; @@ -42,15 +40,18 @@ import io.substrait.type.NamedStruct; import io.substrait.type.TypeCreator; import io.substrait.util.VisitationContext; +import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.Deque; +import java.util.HashMap; import java.util.HashSet; import java.util.LinkedList; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalLong; -import java.util.Stack; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -164,18 +165,19 @@ public static RelNode convert( .typeSystem(converterProvider.getTypeSystem()) .programs() .build()); - return relRoot.accept( - converterProvider.getSubstraitRelNodeConverter(relBuilder), Context.newContext()); + // Normalize any offset-based outer references (steps_out) to the id-based form (rel_anchor / + // rel_reference) so the conversion below resolves correlations purely by anchor. Plans that are + // already id-based are left unchanged. + return OuterReferenceConverter.toIdBased(relRoot) + .accept(converterProvider.getSubstraitRelNodeConverter(relBuilder), Context.newContext()); } @Override public RelNode visit(Filter filter, Context context) throws RuntimeException { RelNode input = filter.getInput().accept(this, context); - context.pushOuterRowType(input.getRowType()); + context.enterScope(AnchoredInput.of(filter.getInput().getRelAnchor(), input.getRowType())); RexNode filterCondition = filter.getCondition().accept(expressionRexConverter, context); - RelNode node = - relBuilder.push(input).filter(context.popCorrelationIds(), filterCondition).build(); - context.popOuterRowType(); + RelNode node = relBuilder.push(input).filter(context.exitScope(), filterCondition).build(); return applyRemap(node, filter.getRemap()); } @@ -193,7 +195,7 @@ public RelNode visit(LocalFiles localFiles, Context context) throws RuntimeExcep @Override public RelNode visit(Project project, Context context) throws RuntimeException { RelNode child = project.getInput().accept(this, context); - context.pushOuterRowType(child.getRowType()); + context.enterScope(AnchoredInput.of(project.getInput().getRelAnchor(), child.getRowType())); Stream directOutputs = IntStream.range(0, child.getRowType().getFieldCount()) @@ -206,11 +208,7 @@ public RelNode visit(Project project, Context context) throws RuntimeException { Stream.concat(directOutputs, exprs).collect(java.util.stream.Collectors.toList()); RelNode node = - relBuilder - .push(child) - .project(rexExprs, List.of(), false, context.popCorrelationIds()) - .build(); - context.popOuterRowType(); + relBuilder.push(child).project(rexExprs, List.of(), false, context.exitScope()).build(); return applyRemap(node, project.getRemap()); } @@ -228,19 +226,16 @@ public RelNode visit(Cross cross, Context context) throws RuntimeException { public RelNode visit(Join join, Context context) throws RuntimeException { RelNode left = join.getLeft().accept(this, context); RelNode right = join.getRight().accept(this, context); - context.pushOuterRowType(left.getRowType(), right.getRowType()); + context.enterScope( + AnchoredInput.of(join.getLeft().getRelAnchor(), left.getRowType()), + AnchoredInput.of(join.getRight().getRelAnchor(), right.getRowType())); RexNode condition = join.getCondition() .map(c -> c.accept(expressionRexConverter, context)) .orElse(relBuilder.literal(true)); JoinRelType joinType = asJoinRelType(join); RelNode node = - relBuilder - .push(left) - .push(right) - .join(joinType, condition, context.popCorrelationIds()) - .build(); - context.popOuterRowType(); + relBuilder.push(left).push(right).join(joinType, condition, context.exitScope()).build(); return applyRemap(node, join.getRemap()); } @@ -802,15 +797,64 @@ private RelNode applyRemap(RelNode relNode, Rel.Remap remap) { return relBuilder.push(relNode).project(rexList).build(); } - /** A shared context for the Substrait to RelNode conversion. */ + /** + * A relational input paired with the {@code rel_anchor} it may carry, used by {@link Context}. + */ + public static final class AnchoredInput { + final Optional anchor; + final RelDataType rowType; + + private AnchoredInput(Optional anchor, RelDataType rowType) { + this.anchor = anchor; + this.rowType = rowType; + } + + /** + * Creates an anchored input. + * + * @param anchor the input's {@code rel_anchor}, if any + * @param rowType the input's Calcite row type + * @return the anchored input + */ + public static AnchoredInput of(Optional anchor, RelDataType rowType) { + return new AnchoredInput(anchor, rowType); + } + } + + /** + * A shared context for the Substrait to RelNode conversion. + * + *

Correlated (outer) references are resolved by id: a relation that binds an outer reference + * carries a {@code rel_anchor}, and each reference carries the matching {@code rel_reference}. + * When a relational operator is converted, it enters a scope recording the anchors of its inputs; + * a reference then resolves against the enclosing scope owning its anchor, and the {@link + * CorrelationId} minted for that anchor is attached to the operator that owns it. + */ public static class Context implements VisitationContext { - /** Stack of outer row type range maps used to resolve correlated references. */ - protected final Stack> outerRowTypes = new Stack<>(); - /** Stack of correlation ids collected while visiting subqueries. */ - protected final Stack> correlationIds = new Stack<>(); + /** Stack of correlation scopes, innermost on top. */ + private final Deque scopes = new ArrayDeque<>(); + + /** Maps an in-scope {@code rel_anchor} to the scope that owns it. */ + private final Map scopeByAnchor = new HashMap<>(); + + /** Maps a {@code rel_anchor} to the single {@link CorrelationId} minted for it. */ + private final Map correlationIdByAnchor = new HashMap<>(); - private int subqueryDepth; + /** + * 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 seenAnchors = new HashSet<>(); + + /** One correlation scope per enclosing relational operator. */ + private static final class Scope { + final Map rowTypeByAnchor = new HashMap<>(); + final java.util.Set correlationIds = new HashSet<>(); + } /** * Creates a new {@link Context} instance. @@ -822,75 +866,84 @@ public static Context newContext() { } /** - * Adds the outer row types to the top of the stack of outer row types. - * - *

Row types are stored as a {@link RangeMap} with field indices as keys and the {@link - * RelDataType} row type containing the field at the field index by continuously numbering the - * field indices from 0 across all provided row types in the order the row types are passed as - * arguments. + * Enters a correlation scope for a relational operator, recording the {@code rel_anchor} (if + * any) carried by each of its inputs. * - * @param inputs the row types to add + * @param inputs the operator's inputs paired with their anchors */ - public void pushOuterRowType(final RelDataType... inputs) { - final RangeMap fieldRangeMap = TreeRangeMap.create(); - int begin = 0; - for (final RelDataType parent : inputs) { - final int end = begin + parent.getFieldCount(); - final Range range = Range.closedOpen(begin, end); - fieldRangeMap.put(range, parent); - begin = end; + public void enterScope(final AnchoredInput... inputs) { + final Scope scope = new Scope(); + 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); + } } - - outerRowTypes.push(fieldRangeMap); - this.correlationIds.push(new HashSet<>()); - } - - /** Pops the most recent outer row type from the stack. */ - public void popOuterRowType() { - outerRowTypes.pop(); + scopes.push(scope); } /** - * Returns the outer row type {@link RangeMap} walking up the given steps from the current - * subquery depth. + * Exits the innermost correlation scope, returning the correlation ids to attach to the + * operator that owns it. * - * @param stepsOut number of steps to walk up from the current subquery depth - * @return {@link RangeMap} with field indices as keys and the {@link RelDataType} row type - * containing the field at the field index + * @return the correlation ids resolved against this operator's inputs */ - public RangeMap getOuterRowTypeRangeMap(final Integer stepsOut) { - return this.outerRowTypes.get(subqueryDepth - stepsOut); + public java.util.Set exitScope() { + final Scope scope = scopes.pop(); + for (final Integer anchor : scope.rowTypeByAnchor.keySet()) { + scopeByAnchor.remove(anchor); + } + return scope.correlationIds; } /** - * Removes the correlation ids at the top of the stack. + * Returns the Calcite row type of the in-scope relation bearing the given {@code rel_anchor}. * - * @return the correlation ids removed from the top of the stack + * @param anchor the referenced {@code rel_anchor} + * @return the row type of the binding relation */ - public java.util.Set popCorrelationIds() { - return correlationIds.pop(); + public RelDataType getAnchorRowType(final int anchor) { + final Scope scope = requireScope(anchor); + return scope.rowTypeByAnchor.get(anchor); } /** - * Adds a {@link CorrelationId} to the subquery depth walking up the given steps from the - * current subquery depth. + * Returns the {@link CorrelationId} for the given {@code rel_anchor}, creating it on first use + * and attaching it to the scope that owns the anchor so the binding operator declares it. * - * @param stepsOut number of steps to walk up from the current subquery depth - * @param correlationId the {@link CorrelationId} to add + * @param anchor the referenced {@code rel_anchor} + * @param factory supplies a fresh correlation id when one has not yet been minted + * @return the correlation id for this anchor */ - public void addCorrelationId(final int stepsOut, final CorrelationId correlationId) { - final int index = subqueryDepth - stepsOut; - this.correlationIds.get(index).add(correlationId); - } - - /** Increments the current subquery depth. */ - public void incrementSubqueryDepth() { - this.subqueryDepth++; - } - - /** Decrements the current subquery depth. */ - public void decrementSubqueryDepth() { - this.subqueryDepth--; + public CorrelationId correlationIdForAnchor( + final int anchor, final java.util.function.Supplier factory) { + final Scope scope = requireScope(anchor); + final CorrelationId correlationId = + correlationIdByAnchor.computeIfAbsent(anchor, k -> factory.get()); + scope.correlationIds.add(correlationId); + return correlationId; + } + + private Scope requireScope(final int anchor) { + final Scope scope = scopeByAnchor.get(anchor); + if (scope == null) { + // 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"); + } + return scope; } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java index bf5eedb75..7e1a1e7f3 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java @@ -40,7 +40,6 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.Map; import java.util.Optional; import java.util.OptionalLong; import java.util.stream.Collectors; @@ -51,11 +50,11 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.CorrelationId; import org.apache.calcite.rel.core.JoinRelType; import org.apache.calcite.rel.core.TableModify; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; -import org.apache.calcite.rex.RexFieldAccess; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlKind; @@ -84,7 +83,7 @@ public class SubstraitRelVisitor extends RelNodeVisitor { /** Converter for Calcite {@link RelDataType} to Substrait {@link Type}. */ protected final TypeConverter typeConverter; - private Map fieldAccessDepthMap; + private OuterReferenceResolver outerReferenceResolver; /** * Creates a new SubstraitRelVisitor with the specified type factory and extensions. @@ -772,33 +771,57 @@ public Rel visitOther(RelNode other) { } /** - * Precomputes depth for outer field accesses used by correlated expressions. + * Assigns id-based outer-reference anchors for correlated expressions in the given tree. * * @param root Root Calcite node to analyze */ - protected void popFieldAccessDepthMap(RelNode root) { - final OuterReferenceResolver resolver = new OuterReferenceResolver(); - fieldAccessDepthMap = resolver.apply(root); + protected void resolveOuterReferences(RelNode root) { + outerReferenceResolver = new OuterReferenceResolver(); + outerReferenceResolver.resolve(root); } /** - * Returns the depth of a field access for correlated expressions. + * Returns the outer-reference anchor a correlated field reference bound to the given correlation + * id must emit. * - * @param fieldAccess Rex field access - * @return Depth value, or {@code null} if unknown + * @param correlationId the Calcite correlation id + * @return the anchor, or {@code null} if unknown */ - public Integer getFieldAccessDepth(RexFieldAccess fieldAccess) { - return fieldAccessDepthMap.get(fieldAccess); + public Integer getOuterReferenceAnchor(CorrelationId correlationId) { + return outerReferenceResolver == null + ? null + : outerReferenceResolver.anchorForCorrelationId(correlationId); } /** - * Applies the visitor to a Calcite {@link RelNode}. + * Applies the visitor to a Calcite {@link RelNode}, stamping the produced relation with its + * outer-reference {@code rel_anchor} when it is the binding point of a correlated reference. * * @param r Calcite node * @return Converted Substrait relation */ public Rel apply(RelNode r) { - return reverseAccept(r); + Rel rel = reverseAccept(r); + if (outerReferenceResolver != null) { + Integer anchor = outerReferenceResolver.anchorForTarget(r); + if (anchor != null) { + 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; } /** @@ -840,7 +863,9 @@ public static Plan.Root convert(RelRoot relRoot, SimpleExtension.ExtensionCollec */ public static Plan.Root convert(RelRoot relRoot, ConverterProvider converterProvider) { SubstraitRelVisitor visitor = converterProvider.getSubstraitRelVisitor(); - visitor.popFieldAccessDepthMap(relRoot.rel); + // Assign id-based outer-reference anchors; apply() then emits rel_reference for correlated + // references and stamps rel_anchor on their binding relations. + visitor.resolveOuterReferences(relRoot.rel); Rel rel = visitor.apply(relRoot.project()); // Avoid using the names from relRoot.validatedRowType because if there are @@ -874,7 +899,7 @@ public static Rel convert(RelNode relNode, SimpleExtension.ExtensionCollection e */ public static Rel convert(RelNode relNode, ConverterProvider converterProvider) { SubstraitRelVisitor visitor = converterProvider.getSubstraitRelVisitor(); - visitor.popFieldAccessDepthMap(relNode); + visitor.resolveOuterReferences(relNode); return visitor.apply(relNode); } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java index bec19102c..bea3c318b 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java @@ -2,6 +2,7 @@ import io.substrait.isthmus.SubstraitRelNodeConverter.Context; import io.substrait.plan.Plan; +import io.substrait.relation.OuterReferenceConverter; import io.substrait.relation.Rel; import io.substrait.util.EmptyVisitationContext; import java.util.ArrayList; @@ -63,6 +64,11 @@ public SubstraitToCalcite( * @return {@link RelNode} */ public RelNode convert(Rel rel) { + // Normalize any offset-based outer references (steps_out) to the id-based form (rel_anchor / + // rel_reference) so the conversion resolves correlations purely by anchor. Plans that are + // already id-based are unchanged. + rel = OuterReferenceConverter.toIdBased(rel); + RelBuilder relBuilder; if (catalogReader != null) { relBuilder = converterProvider.getRelBuilder(catalogReader.getRootSchema()); diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java index a7ace5d46..3ec528afe 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -1,8 +1,6 @@ package io.substrait.isthmus.expression; import com.google.common.collect.ImmutableList; -import com.google.common.collect.Range; -import com.google.common.collect.RangeMap; import io.substrait.expression.AbstractExpressionVisitor; import io.substrait.expression.EnumArg; import io.substrait.expression.Expression; @@ -644,9 +642,7 @@ private boolean isDistinct(Expression.WindowFunctionInvocation expr) { public RexNode visit(Expression.InPredicate expr, Context context) throws RuntimeException { List needles = expr.needles().stream().map(e -> e.accept(this, context)).collect(Collectors.toList()); - context.incrementSubqueryDepth(); RelNode rel = expr.haystack().accept(relNodeConverter, context); - context.decrementSubqueryDepth(); return RexSubQuery.in(rel, ImmutableList.copyOf(needles)); } @@ -737,27 +733,21 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti } else if (expr.isOuterReference()) { final ReferenceSegment segment = expr.segments().get(0); - final RexNode rexInputRef; if (segment instanceof FieldReference.StructField) { final FieldReference.StructField field = (FieldReference.StructField) segment; - final RangeMap fieldRangeMap = - context.getOuterRowTypeRangeMap(expr.outerReferenceStepsOut().get()); - final Range range = fieldRangeMap.getEntry(field.offset()).getKey(); - final int fieldOffset = field.offset() - range.lowerEndpoint(); - + // Resolve the id-based outer reference by anchor. Incoming offset-based (steps_out) plans + // are normalized to the id-based form at the conversion boundary, so rel_reference is set. + final int anchor = expr.outerReferenceRelReference().get(); + final RelDataType rowType = context.getAnchorRowType(anchor); final CorrelationId correlationId = - relNodeConverter.getRelBuilder().getCluster().createCorrel(); - context.addCorrelationId(expr.outerReferenceStepsOut().get(), correlationId); - rexInputRef = - rexBuilder.makeFieldAccess( - rexBuilder.makeCorrel(fieldRangeMap.get(field.offset()), correlationId), - fieldOffset); + context.correlationIdForAnchor( + anchor, () -> relNodeConverter.getRelBuilder().getCluster().createCorrel()); + return rexBuilder.makeFieldAccess( + rexBuilder.makeCorrel(rowType, correlationId), field.offset()); } else { throw new IllegalArgumentException("Unhandled type: " + segment); } - - return rexInputRef; } else if (expr.isLambdaParameterReference()) { // as of now calcite doesn't support nested lambda functions // https://github.com/substrait-io/substrait-java/issues/711 @@ -820,17 +810,13 @@ public RexNode visitEnumArg( @Override public RexNode visit(ScalarSubquery expr, Context context) throws RuntimeException { - context.incrementSubqueryDepth(); RelNode inputRelnode = expr.input().accept(relNodeConverter, context); - context.decrementSubqueryDepth(); return RexSubQuery.scalar(inputRelnode); } @Override public RexNode visit(SetPredicate expr, Context context) throws RuntimeException { - context.incrementSubqueryDepth(); RelNode inputRelnode = expr.tuples().accept(relNodeConverter, context); - context.decrementSubqueryDepth(); switch (expr.predicateOp()) { case PREDICATE_OP_EXISTS: return RexSubQuery.exists(inputRelnode); diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/RexExpressionConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/RexExpressionConverter.java index f82275a67..10d3766b4 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/RexExpressionConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/RexExpressionConverter.java @@ -14,6 +14,7 @@ import java.util.List; import java.util.Optional; import java.util.stream.Collectors; +import org.apache.calcite.rel.core.CorrelationId; import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexCorrelVariable; import org.apache.calcite.rex.RexDynamicParam; @@ -224,28 +225,22 @@ public Expression visitFieldAccess(RexFieldAccess fieldAccess) { switch (kind) { case CORREL_VARIABLE: { - Integer stepsOut = relVisitor.getFieldAccessDepth(fieldAccess); - // Substrait requires steps_out >= 1 for an outer reference - // (proto/substrait/algebra.proto, OuterReference.steps_out: "Must be >= 1."). - // A null or 0 depth means the correlation does not step out of any subquery, - // typically a plan that was only partially decorrelated (the owning Correlate - // was rewritten into a join but the $cor reference remains in this Filter). - // We can't represent that with steps_out, so fail clearly rather than emit - // an invalid steps_out=0 reference. Id-based outer-reference resolution - // would handle this and is the proper long-term fix. - if (stepsOut == null || stepsOut == 0) { + CorrelationId correlationId = ((RexCorrelVariable) fieldAccess.getReferenceExpr()).id; + Integer anchor = relVisitor.getOuterReferenceAnchor(correlationId); + if (anchor == null) { throw new UnsupportedOperationException( String.format( - "Cannot convert outer reference %s: it does not step out of any " - + "subquery (steps_out=%s). This usually indicates a plan that was " - + "only partially decorrelated; Substrait requires steps_out >= 1.", - fieldAccess, stepsOut)); + "Cannot convert outer reference %s: correlation %s has no binding relation. " + + "This usually indicates a plan that was only partially decorrelated.", + fieldAccess, correlationId)); } - return FieldReference.newRootStructOuterReference( + // Emit an id-based outer reference. The binding relation is stamped with this anchor by + // SubstraitRelVisitor#apply. + return FieldReference.newRootStructOuterReferenceByRelReference( fieldAccess.getField().getIndex(), typeConverter.toSubstrait(fieldAccess.getType()), - stepsOut); + anchor); } case ITEM: case INPUT_REF: diff --git a/isthmus/src/test/java/io/substrait/isthmus/OuterReferenceResolverTest.java b/isthmus/src/test/java/io/substrait/isthmus/OuterReferenceResolverTest.java index bde78a3c6..a01d266cf 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/OuterReferenceResolverTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/OuterReferenceResolverTest.java @@ -1,12 +1,8 @@ package io.substrait.isthmus; -import java.util.Map; -import java.util.Map.Entry; -import java.util.Optional; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.JoinRelType; import org.apache.calcite.rex.RexCorrelVariable; -import org.apache.calcite.rex.RexFieldAccess; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.parser.SqlParseException; import org.apache.calcite.tools.RelBuilder; @@ -17,23 +13,10 @@ class OuterReferenceResolverTest extends PlanTestBase { private final RelBuilder tpcDsRelBuilder = new RelCreator(TPCDS_CATALOG).createRelBuilder(); - private static void validateOuterRef( - final Map fieldAccessDepthMap, - final String refName, - final String colName, - final int depth) { - final Optional> entry = - fieldAccessDepthMap.entrySet().stream() - .filter(f -> f.getKey().getReferenceExpr().toString().equals(refName)) - .filter(f -> f.getKey().getField().getName().equals(colName)) - .filter(f -> f.getValue() == depth) - .findFirst(); - Assertions.assertTrue(entry.isPresent()); - } - - private static Map buildOuterFieldRefMap(final RelNode root) { + private static OuterReferenceResolver resolve(final RelNode root) { final OuterReferenceResolver resolver = new OuterReferenceResolver(); - return resolver.apply(root); + resolver.resolve(root); + return resolver; } @Test @@ -60,9 +43,10 @@ void lateralJoinQuery() throws SqlParseException { tpcDsRelBuilder.field("SS_CUSTOMER_SK")) .build(); - final Map fieldAccessDepthMap = buildOuterFieldRefMap(calciteRel); - Assertions.assertEquals(1, fieldAccessDepthMap.size()); - validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 1); + final OuterReferenceResolver resolver = resolve(calciteRel); + // The correlation binds to the Correlate's left input, which is stamped with the anchor emitted + // by references to $cor0. + Assertions.assertNotNull(resolver.anchorForCorrelationId(cor0.get().id)); } @Test @@ -89,9 +73,8 @@ void outerApplyQuery() throws SqlParseException { tpcDsRelBuilder.field("SS_CUSTOMER_SK")) .build(); - final Map fieldAccessDepthMap = buildOuterFieldRefMap(calciteRel); - Assertions.assertEquals(1, fieldAccessDepthMap.size()); - validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 1); + final OuterReferenceResolver resolver = resolve(calciteRel); + Assertions.assertNotNull(resolver.anchorForCorrelationId(cor0.get().id)); } @Test @@ -101,8 +84,8 @@ void nestedApplyJoinQuery() throws SqlParseException { // FROM store_sales CROSS APPLY // ( SELECT i_item_sk // FROM item CROSS APPLY - // ( SELECT p_promo_sk\ - // FROM promotion\ + // ( SELECT p_promo_sk + // FROM promotion // WHERE p_item_sk = i_item_sk AND p_item_sk = ss_item_sk ) // WHERE item.i_item_sk = store_sales.ss_item_sk ) @@ -137,11 +120,13 @@ void nestedApplyJoinQuery() throws SqlParseException { tpcDsRelBuilder.field("SS_CUSTOMER_SK")) .build(); - final Map fieldAccessDepthMap = buildOuterFieldRefMap(calciteRel); - Assertions.assertEquals(3, fieldAccessDepthMap.size()); - validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 1); - validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 2); - validateOuterRef(fieldAccessDepthMap, "$cor1", "I_ITEM_SK", 1); + final OuterReferenceResolver resolver = resolve(calciteRel); + final Integer anchor0 = resolver.anchorForCorrelationId(cor0.get().id); + final Integer anchor1 = resolver.anchorForCorrelationId(cor1.get().id); + // Each correlation binds to a distinct relation, so they receive distinct anchors. + Assertions.assertNotNull(anchor0); + Assertions.assertNotNull(anchor1); + Assertions.assertNotEquals(anchor0, anchor1); } /** @@ -149,39 +134,22 @@ void nestedApplyJoinQuery() throws SqlParseException { * *

When a Calcite optimizer (e.g. HepPlanner) rewrites a correlated IN-subquery into a join, it * can produce a {@link org.apache.calcite.rel.core.Filter} whose condition contains a {@link - * RexFieldAccess} referencing a {@link RexCorrelVariable} ({@code $cor0.field}), but whose {@link - * RelNode#getVariablesSet()} is empty — because the enclosing {@link - * org.apache.calcite.rel.core.Correlate} node has already been replaced by a regular join. + * org.apache.calcite.rex.RexFieldAccess} referencing a {@link RexCorrelVariable} ({@code + * $cor0.field}), but whose {@link RelNode#getVariablesSet()} is empty — because the enclosing + * {@link org.apache.calcite.rel.core.Correlate} node has already been replaced by a regular join. * - *

Before the fix, {@code OuterReferenceResolver.visit(Filter)} only iterated {@code - * filter.getVariablesSet()}, so the {@link org.apache.calcite.rel.core.CorrelationId} was never - * registered in {@code nestedDepth}. The subsequent {@code rexVisitor.visitFieldAccess()} call - * found {@code !nestedDepth.containsKey(id)} and silently skipped the access, producing an empty - * {@code fieldAccessDepthMap} instead of the expected entry. - * - *

The fix pre-scans the Filter condition for {@link RexCorrelVariable} references and - * registers their IDs before delegating to the rex visitor. + *

Such a correlation is never declared, so it has no binding relation. The resolver leaves it + * unbound and conversion fails with a clear error rather than emitting an invalid or misdirected + * outer reference. */ @Test void filterWithCorrelVariableButEmptyVariablesSet() throws SqlParseException { - // Build the partially-decorrelated pattern: - // JOIN (inner side has a Filter whose condition references $cor0 but variablesSet is empty) - // - // SQL equivalent of what a post-optimisation plan looks like: - // SELECT ss_sold_date_sk FROM store_sales - // JOIN item ON item.i_item_sk = store_sales.ss_item_sk ← decorrelated join - // WHERE item.i_item_sk = $cor0.SS_ITEM_SK ← Filter still has the correl ref - // - // We deliberately call .filter(condition) without passing the CorrelationId so that - // getVariablesSet() returns an empty set, reproducing the post-decorrelation shape. - final Holder cor0 = Holder.empty(); // Push STORE_SALES and capture the correlation variable for it. tpcDsRelBuilder.scan("tpcds", "STORE_SALES").variable(cor0::set); // Push ITEM on top, then build the filter condition while ITEM is the top-of-stack. - // The condition equates item.I_ITEM_SK to the correlated $cor0.SS_ITEM_SK. tpcDsRelBuilder.scan("tpcds", "ITEM"); RexNode correlCondition = tpcDsRelBuilder.equals( @@ -195,17 +163,15 @@ void filterWithCorrelVariableButEmptyVariablesSet() throws SqlParseException { .project(tpcDsRelBuilder.field("SS_SOLD_DATE_SK")) .build(); - // Resolver registers the dangling correlation (no longer silently dropped → no NPE). - final Map fieldAccessDepthMap = buildOuterFieldRefMap(calciteRel); - Assertions.assertEquals(1, fieldAccessDepthMap.size()); - validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 0); + // The dangling correlation is never declared, so it is not bound to any relation. + final OuterReferenceResolver resolver = resolve(calciteRel); + Assertions.assertNull(resolver.anchorForCorrelationId(cor0.get().id)); - // steps_out=0 is not a valid Substrait outer reference (proto requires >= 1), - // so conversion must fail clearly rather than emit an invalid plan. + // Conversion must fail clearly rather than emit a misdirected outer reference. UnsupportedOperationException ex = Assertions.assertThrows( UnsupportedOperationException.class, () -> SubstraitRelVisitor.convert(calciteRel, converterProvider)); - Assertions.assertTrue(ex.getMessage().contains("steps_out")); + Assertions.assertTrue(ex.getMessage().contains("no binding relation")); } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubqueryPlanTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubqueryPlanTest.java index 9e9aad657..2cbcb7773 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubqueryPlanTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubqueryPlanTest.java @@ -1,6 +1,7 @@ package io.substrait.isthmus; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -54,7 +55,9 @@ void existsCorrelatedSubquery() throws SqlParseException { .getSelection(); // l_orderkey assertEquals(0, correlatedCol.getDirectReference().getStructField().getField()); - assertEquals(1, correlatedCol.getOuterReference().getStepsOut()); + // Correlated references are now emitted in the id-based form (rel_reference) rather than + // steps_out; see OuterReferenceConverter. + assertTrue(correlatedCol.getOuterReference().hasRelReference()); } @Test @@ -97,7 +100,9 @@ void uniqueCorrelatedSubquery() throws IOException, SqlParseException { .getSelection(); // l_orderkey assertEquals(0, correlatedCol.getDirectReference().getStructField().getField()); - assertEquals(1, correlatedCol.getOuterReference().getStepsOut()); + // Correlated references are now emitted in the id-based form (rel_reference) rather than + // steps_out; see OuterReferenceConverter. + assertTrue(correlatedCol.getOuterReference().hasRelReference()); } @Test @@ -135,7 +140,9 @@ void inPredicateCorrelatedSubQuery() throws IOException, SqlParseException { .getSelection(); // l_partkey assertEquals(1, correlatedCol.getDirectReference().getStructField().getField()); - assertEquals(1, correlatedCol.getOuterReference().getStepsOut()); + // Correlated references are now emitted in the id-based form (rel_reference) rather than + // steps_out; see OuterReferenceConverter. + assertTrue(correlatedCol.getOuterReference().hasRelReference()); } @Test @@ -175,7 +182,9 @@ void notInPredicateCorrelatedSubquery() throws IOException, SqlParseException { .getSelection(); // l_partkey assertEquals(1, correlatedCol.getDirectReference().getStructField().getField()); - assertEquals(1, correlatedCol.getOuterReference().getStepsOut()); + // Correlated references are now emitted in the id-based form (rel_reference) rather than + // steps_out; see OuterReferenceConverter. + assertTrue(correlatedCol.getOuterReference().hasRelReference()); } @Test @@ -247,7 +256,7 @@ void existsNestedCorrelatedSubquery() throws IOException, SqlParseException { .getValue() .getSelection(); // p.p_partkey assertEquals(0, correlatedCol1.getDirectReference().getStructField().getField()); - assertEquals(2, correlatedCol1.getOuterReference().getStepsOut()); + assertTrue(correlatedCol1.getOuterReference().hasRelReference()); Expression.FieldReference correlatedCol2 = inner_subquery_cond2 @@ -256,7 +265,11 @@ void existsNestedCorrelatedSubquery() throws IOException, SqlParseException { .getValue() .getSelection(); // l.l_suppkey assertEquals(2, correlatedCol2.getDirectReference().getStructField().getField()); - assertEquals(1, correlatedCol2.getOuterReference().getStepsOut()); + // The two correlated columns reference different outer relations, so distinct rel anchors. + assertTrue(correlatedCol2.getOuterReference().hasRelReference()); + assertNotEquals( + correlatedCol1.getOuterReference().getRelReference(), + correlatedCol2.getOuterReference().getRelReference()); } @Test @@ -322,7 +335,7 @@ void nestedScalarCorrelatedSubquery() throws IOException, SqlParseException { .getValue() .getSelection(); // p.p_partkey assertEquals(0, correlatedCol1.getDirectReference().getStructField().getField()); - assertEquals(2, correlatedCol1.getOuterReference().getStepsOut()); + assertTrue(correlatedCol1.getOuterReference().hasRelReference()); Expression.FieldReference correlatedCol2 = inner_subquery_cond2 @@ -331,7 +344,11 @@ void nestedScalarCorrelatedSubquery() throws IOException, SqlParseException { .getValue() .getSelection(); // l.l_suppkey assertEquals(2, correlatedCol2.getDirectReference().getStructField().getField()); - assertEquals(1, correlatedCol2.getOuterReference().getStepsOut()); + // The two correlated columns reference different outer relations, so distinct rel anchors. + assertTrue(correlatedCol2.getOuterReference().hasRelReference()); + assertNotEquals( + correlatedCol1.getOuterReference().getRelReference(), + correlatedCol2.getOuterReference().getRelReference()); } @Test diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java index a3f234c47..b6f18d045 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java +++ b/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; @@ -110,6 +114,23 @@ public Optional getHint() { return input.getHint(); } + @Override + public Optional 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 relAnchor) { + return new SubstraitRepeatRel(input.withRelAnchor(relAnchor), repeatCount); + } + @Override public O accept( final RelVisitor visitor, final C context) throws E { @@ -146,7 +167,7 @@ public static Rel convert( final RelNode relNode, final SimpleExtension.ExtensionCollection extensions) { final SubstraitRelVisitorCustom visitor = new SubstraitRelVisitorCustom(relNode.getCluster().getTypeFactory(), extensions); - visitor.popFieldAccessDepthMap(relNode); + visitor.resolveOuterReferences(relNode); return visitor.apply(relNode); } } @@ -252,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()); + } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java b/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java index 9d089825c..88475a4df 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java @@ -1,6 +1,7 @@ package io.substrait.isthmus.expression; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import io.substrait.expression.FieldReference; import io.substrait.isthmus.PlanTestBase; @@ -87,6 +88,79 @@ void testOuterFieldReferenceOneStep() { SubstraitSqlDialect.toSql(calciteRel).getSql()); } + @Test + void testOuterFieldReferenceOneStepIdBased() { + /* + * Same correlated scalar subquery as testOuterFieldReferenceOneStep, but the outer reference is + * expressed in the id-based form (rel_reference) binding to the orders scan's rel_anchor. The + * consumer must resolve it to the same Calcite $cor0 correlation as the steps_out form. + */ + final Rel root = + 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.newRootStructOuterReferenceByRelReference( + 1, TypeCreator.REQUIRED.I64, 1)), + customerTableScan)), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2, 3)), + // The orders scan is the binding relation for the id-based outer reference. + orderTableScan.withRelAnchor(1)); + + final RelNode calciteRel = substraitToCalcite.convert(root); + + assertEquals( + "LogicalProject(variablesSet=[[$cor0]], o_orderkey0=[$0], $f3=[$SCALAR_QUERY({\n" + + "LogicalProject(c_nationkey=[$1])\n" + + " LogicalFilter(condition=[=($0, $cor0.o_custkey)])\n" + + " LogicalTableScan(table=[[customer]])\n" + + "})])\n" + + " LogicalTableScan(table=[[orders]])\n", + calciteRel.explain()); + } + + @Test + void duplicateRelAnchorIsRejected() { + /* + * A malformed id-based plan that reuses the same rel_anchor for two distinct relations (the + * outer orders scan and the inner customer scan). Resolving references purely by anchor value + * is only sound when anchors are unique plan-wide, so conversion must reject this rather than + * silently mis-resolve the outer reference. + */ + final Rel root = + 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.newRootStructOuterReferenceByRelReference( + 1, TypeCreator.REQUIRED.I64, 1)), + // duplicate anchor: the customer scan reuses the orders scan's + // rel_anchor (1) + customerTableScan.withRelAnchor(1))), + TypeCreator.NULLABLE.I64)), + Remap.of(List.of(2, 3)), + orderTableScan.withRelAnchor(1)); + + assertThrows(UnsupportedOperationException.class, () -> substraitToCalcite.convert(root)); + } + @Test void testOuterFieldReferenceTwoSteps() { /* diff --git a/substrait b/substrait index f09071699..d61a9403d 160000 --- a/substrait +++ b/substrait @@ -1 +1 @@ -Subproject commit f09071699edcb41d90e1080625e9ee25fce2b81c +Subproject commit d61a9403dc6fe0006b8b712b97b9d6a3bada1acf