From f2d5c65797bd66fe0ca4136bb09df6f3b616d76c Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Wed, 5 Aug 2026 14:10:05 +0200 Subject: [PATCH 1/2] fix(core): re-derive field reference types when copy-on-write replaces a relation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A FieldReference caches the type of the field it references. RelCopyOnWriteVisitor copied those references over verbatim, so replacing a subtree with one that emits a different record type left every reference above it carrying the type of the relation that is no longer there — and the record types derived from those references, such as a Project's, were wrong in turn. The visitor now tracks the record type that each relation's own expressions resolve against and re-derives the cached type of every field reference it rewrites from it. Inputs are rewritten before their relation's expressions so the scope is the type a replaced input emits, and enclosing scopes are tracked per subquery boundary so an outer reference re-derives against the relation it steps out to. References that resolve against something else are left as they are: a lambda parameter, an outer reference identified by rel anchor, the filter of a read relation, and a reference a rewrite has left selecting a field its input no longer has. The types cached on function invocations are still not re-derived — that needs the function declarations, which the visitor does not have — so a relation whose record type comes from a measure or window function can still be stale. Whether these types should be cached on the POJOs at all is the broader question the issue raises. Three neighbouring bugs surfaced while wiring this up and are fixed as well: - visitFieldReference built its replacement without copying the original, dropping the segments and the type. Since the type is mandatory, rewriting any reference rooted at an expression threw instead of returning the rewritten reference. - The input of a ConsistentPartitionWindow was never visited, so no subtree beneath a window relation could be replaced. - MultiBucketExchange reported no change when only its expression had been rewritten, discarding that rewrite. Closes #185 --- .../substrait/relation/CopyOnWriteUtils.java | 19 ++ .../ExpressionCopyOnWriteVisitor.java | 112 ++++++-- .../relation/RelCopyOnWriteVisitor.java | 268 +++++++++++++++--- .../relation/RelCopyOnWriteVisitorTest.java | 253 +++++++++++++++++ 4 files changed, 585 insertions(+), 67 deletions(-) create mode 100644 core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java diff --git a/core/src/main/java/io/substrait/relation/CopyOnWriteUtils.java b/core/src/main/java/io/substrait/relation/CopyOnWriteUtils.java index adf8a6183..7e3b96fba 100644 --- a/core/src/main/java/io/substrait/relation/CopyOnWriteUtils.java +++ b/core/src/main/java/io/substrait/relation/CopyOnWriteUtils.java @@ -37,6 +37,25 @@ public static Optional or(Optional left, Supplier the type of the supplied value + * @param the exception type that may be thrown + */ + @FunctionalInterface + public interface ThrowingSupplier { + + /** + * Supplies a value. + * + * @return the supplied value + * @throws E if producing the value fails + */ + T get() throws E; + } + /** * Functional interface for transforming values during copy-on-write operations. * diff --git a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java index 1a8923083..daa8c83ed 100644 --- a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java @@ -9,7 +9,10 @@ import io.substrait.expression.FieldReference; import io.substrait.expression.FunctionArg; import io.substrait.expression.ImmutableExpression; +import io.substrait.expression.ImmutableFieldReference; +import io.substrait.type.Type; import io.substrait.util.EmptyVisitationContext; +import java.util.ArrayList; import java.util.List; import java.util.Optional; @@ -424,30 +427,104 @@ protected Optional visitMultiOrListRecord( @Override public Optional visit(FieldReference fieldReference, EmptyVisitationContext context) throws E { - Optional inputExpression = - visitOptionalExpression(fieldReference.inputExpression(), context); + return visitFieldReference(fieldReference, context).map(Expression.class::cast); + } - if (allEmpty(inputExpression)) { + /** + * Visits a field reference, rewriting the expression it is rooted at, if any, and re-deriving its + * cached type from the root it resolves against. Re-deriving the type is what keeps references + * correct when a relation's input is replaced by one emitting a different record type. + * + *

The type of a reference to a lambda parameter, or of an outer reference identified by the + * {@link io.substrait.relation.Rel#getRelAnchor() rel anchor} of the relation it is rooted on, is + * left as it is: neither resolves against a relation in the enclosing scopes tracked during the + * traversal. + * + * @param fieldReference the field reference to visit + * @param context the visitation context + * @return Optional containing the modified field reference, or empty if no changes + * @throws E if an error occurs during visitation + */ + public Optional visitFieldReference( + FieldReference fieldReference, EmptyVisitationContext context) throws E { + if (fieldReference.inputExpression().isPresent()) { + Optional inputExpression = + fieldReference.inputExpression().get().accept(this, context); + if (!inputExpression.isPresent()) { + return Optional.empty(); + } + ImmutableFieldReference.Builder rewritten = + ImmutableFieldReference.builder().from(fieldReference).inputExpression(inputExpression); + if (!fieldReference.segments().isEmpty()) { + rewritten.type( + FieldReference.ofExpression(inputExpression.get(), segmentsOf(fieldReference)).type()); + } + return Optional.of(rewritten.build()); + } + return retypeRootReference(fieldReference); + } + + /** + * Re-derives the type of a reference into the record type of the relation it is rooted on, which + * may have been replaced by one emitting a different record type. + */ + private Optional retypeRootReference(FieldReference fieldReference) { + if (fieldReference.isLambdaParameterReference() + || fieldReference.outerReferenceRelReference().isPresent()) { return Optional.empty(); } - return Optional.of(FieldReference.builder().inputExpression(inputExpression).build()); + Type.Struct rootType = + getRelCopyOnWriteVisitor() + .inputTypeStepsOut(fieldReference.outerReferenceStepsOut().orElse(0)); + // A rewrite that drops fields can leave a reference selecting a field its input no longer has. + // The resulting relation tree is invalid either way, so leave the reference as it is rather + // than turn re-deriving its type into a failure of the rewrite. + if (rootType == null || !selectsFieldOf(fieldReference, rootType)) { + return Optional.empty(); + } + Type type = FieldReference.ofRoot(rootType, segmentsOf(fieldReference)).type(); + if (type.equals(fieldReference.type())) { + return Optional.empty(); + } + return Optional.of(ImmutableFieldReference.builder().from(fieldReference).type(type).build()); + } + + /** + * Returns whether the outermost segment of the given reference selects a field that the given + * record type has. The segments are held innermost first, so the outermost one is the last. + */ + private static boolean selectsFieldOf(FieldReference fieldReference, Type.Struct rootType) { + List segments = fieldReference.segments(); + if (segments.isEmpty()) { + return false; + } + FieldReference.ReferenceSegment outermost = segments.get(segments.size() - 1); + return outermost instanceof FieldReference.StructField + && ((FieldReference.StructField) outermost).offset() < rootType.fields().size(); + } + + /** + * Returns the segments of the given reference as a mutable list, which is what {@link + * FieldReference#ofRoot} and {@link FieldReference#ofExpression} require: they reverse the list + * they are given. + */ + private static List segmentsOf(FieldReference fieldReference) { + return new ArrayList<>(fieldReference.segments()); } @Override public Optional visit( Expression.SetPredicate setPredicate, EmptyVisitationContext context) throws E { - return setPredicate - .tuples() - .accept(getRelCopyOnWriteVisitor(), context) + return getRelCopyOnWriteVisitor() + .inSubqueryScope(() -> setPredicate.tuples().accept(getRelCopyOnWriteVisitor(), context)) .map(tuple -> Expression.SetPredicate.builder().from(setPredicate).tuples(tuple).build()); } @Override public Optional visit( Expression.ScalarSubquery scalarSubquery, EmptyVisitationContext context) throws E { - return scalarSubquery - .input() - .accept(getRelCopyOnWriteVisitor(), context) + return getRelCopyOnWriteVisitor() + .inSubqueryScope(() -> scalarSubquery.input().accept(getRelCopyOnWriteVisitor(), context)) .map( input -> Expression.ScalarSubquery.builder().from(scalarSubquery).input(input).build()); } @@ -455,8 +532,12 @@ public Optional visit( @Override public Optional visit( Expression.InPredicate inPredicate, EmptyVisitationContext context) throws E { - Optional haystack = inPredicate.haystack().accept(getRelCopyOnWriteVisitor(), context); + // The needles are evaluated in the current scope; only the haystack is a subquery boundary. Optional> needles = visitExprList(inPredicate.needles(), context); + Optional haystack = + getRelCopyOnWriteVisitor() + .inSubqueryScope( + () -> inPredicate.haystack().accept(getRelCopyOnWriteVisitor(), context)); if (allEmpty(haystack, needles)) { return Optional.empty(); @@ -523,15 +604,6 @@ protected Optional> visitExprList( return transformList(exprs, context, (e, c) -> e.accept(this, c)); } - private Optional visitOptionalExpression( - Optional optExpr, EmptyVisitationContext context) throws E { - // not using optExpr.map to allow us to propagate the EXCEPTION nicely - if (optExpr.isPresent()) { - return optExpr.get().accept(this, context); - } - return Optional.empty(); - } - /** * Visits a list of function arguments. * diff --git a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java index 301ba4dac..2198b907e 100644 --- a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java @@ -18,7 +18,11 @@ import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; import io.substrait.relation.physical.TopN; +import io.substrait.type.Type; +import io.substrait.type.TypeCreator; import io.substrait.util.EmptyVisitationContext; +import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Optional; import java.util.function.Function; @@ -28,12 +32,38 @@ * overriding a visitor method. The traversal will include relations inside of subquery expressions. * By default, no subtree substitution will be performed. However, if a visit method is overridden * to return a non-empty optional value, then that value will replace the relation in the tree. + * + *

Replacing a subtree may change the record type it emits. Because a {@link FieldReference} + * caches the type of the field it references, the references in the expressions of the relations + * above the replaced subtree would otherwise be left with a stale type. To avoid that, this visitor + * tracks the record type that each relation's own expressions resolve against and re-derives the + * cached type of every field reference it rewrites from the input the reference resolves against. + * The types cached on function invocations are not re-derived; deriving those requires the function + * declarations, which this visitor does not have. + * + *

Tracking that scope makes a visitor instance stateful for the duration of a traversal, so an + * instance must not be used to visit several relation trees concurrently. */ public class RelCopyOnWriteVisitor implements RelVisitor, EmptyVisitationContext, E> { private final ExpressionCopyOnWriteVisitor expressionCopyOnWriteVisitor; + /** + * The record types that the root {@link FieldReference}s of the expressions being rewritten + * resolve against, innermost last. An entry is {@code null} when the expressions it covers do not + * resolve against an input record type, as for the filter of a read relation, which resolves + * against the schema being read. + */ + private final List inputTypes = new ArrayList<>(); + + /** + * The record types of the enclosing scopes, one entry per subquery boundary crossed. An outer + * reference stepping out {@code stepsOut} levels resolves against the entry {@code stepsOut} from + * the top. + */ + private final List outerInputTypes = new ArrayList<>(); + /** Creates a visitor using a default expression visitor bound to this relation visitor. */ public RelCopyOnWriteVisitor() { this.expressionCopyOnWriteVisitor = new ExpressionCopyOnWriteVisitor<>(this); @@ -70,10 +100,13 @@ protected ExpressionCopyOnWriteVisitor getExpressionCopyOnWriteVisitor() { @Override public Optional visit(Aggregate aggregate, EmptyVisitationContext context) throws E { Optional input = aggregate.getInput().accept(this, context); + Type.Struct inputType = recordTypeOf(input.orElse(aggregate.getInput())); Optional> groupings = - transformList(aggregate.getGroupings(), context, this::visitGrouping); + inInputScope( + inputType, () -> transformList(aggregate.getGroupings(), context, this::visitGrouping)); Optional> measures = - transformList(aggregate.getMeasures(), context, this::visitMeasure); + inInputScope( + inputType, () -> transformList(aggregate.getMeasures(), context, this::visitMeasure)); if (allEmpty(input, groupings, measures)) { return Optional.empty(); @@ -155,8 +188,11 @@ protected Optional visitAggregateFunction( @Override public Optional visit(Fetch fetch, EmptyVisitationContext context) throws E { Optional input = fetch.getInput().accept(this, context); - Optional offset = visitOptionalExpression(fetch.getOffset(), context); - Optional count = visitOptionalExpression(fetch.getCount(), context); + Type.Struct inputType = recordTypeOf(input.orElse(fetch.getInput())); + Optional offset = + inInputScope(inputType, () -> visitOptionalExpression(fetch.getOffset(), context)); + Optional count = + inInputScope(inputType, () -> visitOptionalExpression(fetch.getCount(), context)); if (allEmpty(input, offset, count)) { return Optional.empty(); @@ -174,7 +210,9 @@ public Optional visit(Fetch fetch, EmptyVisitationContext context) throws E public Optional visit(Filter filter, EmptyVisitationContext context) throws E { Optional input = filter.getInput().accept(this, context); Optional condition = - filter.getCondition().accept(getExpressionCopyOnWriteVisitor(), context); + inInputScope( + recordTypeOf(input.orElse(filter.getInput())), + () -> filter.getCondition().accept(getExpressionCopyOnWriteVisitor(), context)); if (allEmpty(input, condition)) { return Optional.empty(); @@ -191,8 +229,12 @@ public Optional visit(Filter filter, EmptyVisitationContext context) throws public Optional visit(Join join, EmptyVisitationContext context) throws E { Optional left = join.getLeft().accept(this, context); Optional right = join.getRight().accept(this, context); - Optional condition = visitOptionalExpression(join.getCondition(), context); - Optional postFilter = visitOptionalExpression(join.getPostJoinFilter(), context); + Type.Struct inputType = + recordTypeOf(left.orElse(join.getLeft()), right.orElse(join.getRight())); + Optional condition = + inInputScope(inputType, () -> visitOptionalExpression(join.getCondition(), context)); + Optional postFilter = + inInputScope(inputType, () -> visitOptionalExpression(join.getPostJoinFilter(), context)); if (allEmpty(left, right, condition, postFilter)) { return Optional.empty(); @@ -211,9 +253,13 @@ public Optional visit(Join join, EmptyVisitationContext context) throws E { public Optional visit(LateralJoin lateralJoin, EmptyVisitationContext context) throws E { Optional left = lateralJoin.getLeft().accept(this, context); Optional right = lateralJoin.getRight().accept(this, context); - Optional condition = visitOptionalExpression(lateralJoin.getCondition(), context); + Type.Struct inputType = + recordTypeOf(left.orElse(lateralJoin.getLeft()), right.orElse(lateralJoin.getRight())); + Optional condition = + inInputScope(inputType, () -> visitOptionalExpression(lateralJoin.getCondition(), context)); Optional postFilter = - visitOptionalExpression(lateralJoin.getPostJoinFilter(), context); + inInputScope( + inputType, () -> visitOptionalExpression(lateralJoin.getPostJoinFilter(), context)); if (allEmpty(left, right, condition, postFilter)) { return Optional.empty(); @@ -236,7 +282,8 @@ public Optional visit(Set set, EmptyVisitationContext context) throws E { @Override public Optional visit(NamedScan namedScan, EmptyVisitationContext context) throws E { - Optional filter = visitOptionalExpression(namedScan.getFilter(), context); + Optional filter = + outsideInputScope(() -> visitOptionalExpression(namedScan.getFilter(), context)); if (allEmpty(filter)) { return Optional.empty(); @@ -247,7 +294,8 @@ public Optional visit(NamedScan namedScan, EmptyVisitationContext context) @Override public Optional visit(LocalFiles localFiles, EmptyVisitationContext context) throws E { - Optional filter = visitOptionalExpression(localFiles.getFilter(), context); + Optional filter = + outsideInputScope(() -> visitOptionalExpression(localFiles.getFilter(), context)); if (allEmpty(filter)) { return Optional.empty(); @@ -259,7 +307,10 @@ public Optional visit(LocalFiles localFiles, EmptyVisitationContext context @Override public Optional visit(Project project, EmptyVisitationContext context) throws E { Optional input = project.getInput().accept(this, context); - Optional> expressions = visitExprList(project.getExpressions(), context); + Optional> expressions = + inInputScope( + recordTypeOf(input.orElse(project.getInput())), + () -> visitExprList(project.getExpressions(), context)); if (allEmpty(input, expressions)) { return Optional.empty(); @@ -329,10 +380,14 @@ protected Optional visitTransformExpression( @Override public Optional visit(NamedUpdate update, EmptyVisitationContext context) throws E { Optional condition = - update.getCondition().accept(getExpressionCopyOnWriteVisitor(), context); + outsideInputScope( + () -> update.getCondition().accept(getExpressionCopyOnWriteVisitor(), context)); Optional> transformations = - transformList(update.getTransformations(), context, this::visitTransformExpression); + outsideInputScope( + () -> + transformList( + update.getTransformations(), context, this::visitTransformExpression)); if (allEmpty(condition, transformations)) { return Optional.empty(); @@ -350,7 +405,9 @@ public Optional visit(NamedUpdate update, EmptyVisitationContext context) t public Optional visit(ScatterExchange exchange, EmptyVisitationContext context) throws E { Optional input = exchange.getInput().accept(this, context); Optional> fields = - transformList(exchange.getFields(), context, this::visitFieldReference); + inInputScope( + recordTypeOf(input.orElse(exchange.getInput())), + () -> transformList(exchange.getFields(), context, this::visitFieldReference)); if (allEmpty(input, fields)) { return Optional.empty(); @@ -370,7 +427,9 @@ public Optional visit(SingleBucketExchange exchange, EmptyVisitationContext Optional input = exchange.getInput().accept(this, context); Optional expression = - exchange.getExpression().accept(getExpressionCopyOnWriteVisitor(), context); + inInputScope( + recordTypeOf(input.orElse(exchange.getInput())), + () -> exchange.getExpression().accept(getExpressionCopyOnWriteVisitor(), context)); if (allEmpty(input, expression)) { return Optional.empty(); @@ -389,9 +448,11 @@ public Optional visit(MultiBucketExchange exchange, EmptyVisitationContext throws E { Optional input = exchange.getInput().accept(this, context); Optional expression = - exchange.getExpression().accept(getExpressionCopyOnWriteVisitor(), context); + inInputScope( + recordTypeOf(input.orElse(exchange.getInput())), + () -> exchange.getExpression().accept(getExpressionCopyOnWriteVisitor(), context)); - if (allEmpty(input)) { + if (allEmpty(input, expression)) { return Optional.empty(); } @@ -435,7 +496,9 @@ public Optional visit(BroadcastExchange exchange, EmptyVisitationContext co public Optional visit(Sort sort, EmptyVisitationContext context) throws E { Optional input = sort.getInput().accept(this, context); Optional> sortFields = - transformList(sort.getSortFields(), context, this::visitSortField); + inInputScope( + recordTypeOf(input.orElse(sort.getInput())), + () -> transformList(sort.getSortFields(), context, this::visitSortField)); if (allEmpty(input, sortFields)) { return Optional.empty(); @@ -451,10 +514,14 @@ public Optional visit(Sort sort, EmptyVisitationContext context) throws E { @Override public Optional visit(TopN topN, EmptyVisitationContext context) throws E { Optional input = topN.getInput().accept(this, context); + Type.Struct inputType = recordTypeOf(input.orElse(topN.getInput())); Optional> sortFields = - transformList(topN.getSortFields(), context, this::visitSortField); - Optional offset = visitOptionalExpression(topN.getOffset(), context); - Optional count = visitOptionalExpression(topN.getCount(), context); + inInputScope( + inputType, () -> transformList(topN.getSortFields(), context, this::visitSortField)); + Optional offset = + inInputScope(inputType, () -> visitOptionalExpression(topN.getOffset(), context)); + Optional count = + inInputScope(inputType, () -> visitOptionalExpression(topN.getCount(), context)); if (allEmpty(input, sortFields, offset, count)) { return Optional.empty(); @@ -488,7 +555,8 @@ public Optional visit(Cross cross, EmptyVisitationContext context) throws E @Override public Optional visit(VirtualTableScan virtualTableScan, EmptyVisitationContext context) throws E { - Optional filter = visitOptionalExpression(virtualTableScan.getFilter(), context); + Optional filter = + outsideInputScope(() -> visitOptionalExpression(virtualTableScan.getFilter(), context)); if (allEmpty(filter)) { return Optional.empty(); @@ -524,7 +592,8 @@ public Optional visit(ExtensionMulti extensionMulti, EmptyVisitationContext @Override public Optional visit(ExtensionTable extensionTable, EmptyVisitationContext context) throws E { - Optional filter = visitOptionalExpression(extensionTable.getFilter(), context); + Optional filter = + outsideInputScope(() -> visitOptionalExpression(extensionTable.getFilter(), context)); if (allEmpty(filter)) { return Optional.empty(); @@ -540,12 +609,18 @@ public Optional visit(ExtensionTable extensionTable, EmptyVisitationContext public Optional visit(HashJoin hashJoin, EmptyVisitationContext context) throws E { Optional left = hashJoin.getLeft().accept(this, context); Optional right = hashJoin.getRight().accept(this, context); + Type.Struct inputType = + recordTypeOf(left.orElse(hashJoin.getLeft()), right.orElse(hashJoin.getRight())); Optional> keys = - transformList(hashJoin.getKeys(), context, this::visitComparisonJoinKey); + inInputScope( + inputType, + () -> transformList(hashJoin.getKeys(), context, this::visitComparisonJoinKey)); Optional postFilter = - visitOptionalExpression(hashJoin.getPostJoinFilter(), context); + inInputScope( + inputType, () -> visitOptionalExpression(hashJoin.getPostJoinFilter(), context)); Optional residual = - visitOptionalExpression(hashJoin.getResidualExpression(), context); + inInputScope( + inputType, () -> visitOptionalExpression(hashJoin.getResidualExpression(), context)); if (allEmpty(left, right, keys, postFilter, residual)) { return Optional.empty(); @@ -565,12 +640,18 @@ public Optional visit(HashJoin hashJoin, EmptyVisitationContext context) th public Optional visit(MergeJoin mergeJoin, EmptyVisitationContext context) throws E { Optional left = mergeJoin.getLeft().accept(this, context); Optional right = mergeJoin.getRight().accept(this, context); + Type.Struct inputType = + recordTypeOf(left.orElse(mergeJoin.getLeft()), right.orElse(mergeJoin.getRight())); Optional> keys = - transformList(mergeJoin.getKeys(), context, this::visitComparisonJoinKey); + inInputScope( + inputType, + () -> transformList(mergeJoin.getKeys(), context, this::visitComparisonJoinKey)); Optional postFilter = - visitOptionalExpression(mergeJoin.getPostJoinFilter(), context); + inInputScope( + inputType, () -> visitOptionalExpression(mergeJoin.getPostJoinFilter(), context)); Optional residual = - visitOptionalExpression(mergeJoin.getResidualExpression(), context); + inInputScope( + inputType, () -> visitOptionalExpression(mergeJoin.getResidualExpression(), context)); if (allEmpty(left, right, keys, postFilter, residual)) { return Optional.empty(); @@ -592,7 +673,10 @@ public Optional visit(NestedLoopJoin nestedLoopJoin, EmptyVisitationContext Optional left = nestedLoopJoin.getLeft().accept(this, context); Optional right = nestedLoopJoin.getRight().accept(this, context); Optional condition = - nestedLoopJoin.getCondition().accept(getExpressionCopyOnWriteVisitor(), context); + inInputScope( + recordTypeOf( + left.orElse(nestedLoopJoin.getLeft()), right.orElse(nestedLoopJoin.getRight())), + () -> nestedLoopJoin.getCondition().accept(getExpressionCopyOnWriteVisitor(), context)); if (allEmpty(left, right, condition)) { return Optional.empty(); @@ -610,24 +694,34 @@ public Optional visit(NestedLoopJoin nestedLoopJoin, EmptyVisitationContext public Optional visit( ConsistentPartitionWindow consistentPartitionWindow, EmptyVisitationContext context) throws E { + Optional input = consistentPartitionWindow.getInput().accept(this, context); + Type.Struct inputType = recordTypeOf(input.orElse(consistentPartitionWindow.getInput())); Optional> windowFunctions = - transformList( - consistentPartitionWindow.getWindowFunctions(), context, this::visitWindowRelFunction); + inInputScope( + inputType, + () -> + transformList( + consistentPartitionWindow.getWindowFunctions(), + context, + this::visitWindowRelFunction)); Optional> partitionExpressions = - transformList( - consistentPartitionWindow.getPartitionExpressions(), - context, - (t, c) -> t.accept(getExpressionCopyOnWriteVisitor(), c)); + inInputScope( + inputType, + () -> visitExprList(consistentPartitionWindow.getPartitionExpressions(), context)); Optional> sorts = - transformList(consistentPartitionWindow.getSorts(), context, this::visitSortField); + inInputScope( + inputType, + () -> + transformList(consistentPartitionWindow.getSorts(), context, this::visitSortField)); - if (allEmpty(windowFunctions, partitionExpressions, sorts)) { + if (allEmpty(input, windowFunctions, partitionExpressions, sorts)) { return Optional.empty(); } return Optional.of( ConsistentPartitionWindow.builder() .from(consistentPartitionWindow) + .input(input.orElse(consistentPartitionWindow.getInput())) .partitionExpressions( partitionExpressions.orElse(consistentPartitionWindow.getPartitionExpressions())) .sorts(sorts.orElse(consistentPartitionWindow.getSorts())) @@ -661,6 +755,91 @@ protected Optional visitW .build()); } + // input scope tracking + + /** + * Returns the record type that the root field references of a relation's own expressions resolve + * against: the record types of the given inputs, concatenated in order. + * + * @param inputs the relations the expressions are evaluated over, in field order + * @return the combined record type + */ + protected static Type.Struct recordTypeOf(Rel... inputs) { + return TypeCreator.REQUIRED.struct( + Arrays.stream(inputs).flatMap(input -> input.getRecordType().fields().stream())); + } + + /** + * Runs the given rewrite of a relation's own expressions with the record type that their root + * field references resolve against, so that the cached type of a rewritten reference can be + * re-derived from it. Inputs must be rewritten before calling this, and their rewritten + * record type passed in, so that references pick up the type a replaced input emits. + * + * @param the type of the rewrite's result + * @param inputType the record type the expressions resolve against, or {@code null} if they do + * not resolve against an input record type + * @param rewrite the expression rewrite to run + * @return the result of the rewrite + * @throws E if the rewrite fails + */ + protected T inInputScope( + Type.Struct inputType, CopyOnWriteUtils.ThrowingSupplier rewrite) throws E { + inputTypes.add(inputType); + try { + return rewrite.get(); + } finally { + inputTypes.remove(inputTypes.size() - 1); + } + } + + /** + * Runs the given rewrite of expressions that do not resolve against an input record type, such as + * the filter of a read relation, whose references resolve against the schema being read. + * + * @param the type of the rewrite's result + * @param rewrite the expression rewrite to run + * @return the result of the rewrite + * @throws E if the rewrite fails + */ + protected T outsideInputScope(CopyOnWriteUtils.ThrowingSupplier rewrite) throws E { + return inInputScope(null, rewrite); + } + + /** + * Records the scope currently being rewritten as an enclosing one for the duration of the given + * rewrite, so that an outer reference within it resolves against the right relation. + */ + T inSubqueryScope(CopyOnWriteUtils.ThrowingSupplier rewrite) throws E { + outerInputTypes.add(currentInputType()); + // The relations within the subquery set their own scope as they are visited. Entering the + // subquery with no scope keeps one of them that hosts no expressions from leaking the enclosing + // scope into the expressions it contains. + inputTypes.add(null); + try { + return rewrite.get(); + } finally { + inputTypes.remove(inputTypes.size() - 1); + outerInputTypes.remove(outerInputTypes.size() - 1); + } + } + + /** + * Returns the record type that a field reference stepping out of {@code stepsOut} subquery levels + * resolves against, or {@code null} if it is not known. + */ + Type.Struct inputTypeStepsOut(int stepsOut) { + if (stepsOut <= 0) { + return currentInputType(); + } + int index = outerInputTypes.size() - stepsOut; + return index < 0 ? null : outerInputTypes.get(index); + } + + /** Returns the record type the expressions being rewritten resolve against, if it is known. */ + private Type.Struct currentInputType() { + return inputTypes.isEmpty() ? null : inputTypes.get(inputTypes.size() - 1); + } + // utilities /** @@ -677,7 +856,8 @@ protected Optional> visitExprList( } /** - * Rewrites a field reference, returning a new one if its input expression changed. + * Rewrites a field reference, returning a new one if the expression it is rooted at changed or + * its cached type no longer matches the input it resolves against. * * @param fieldReference the field reference to rewrite * @param context the visitation context @@ -686,13 +866,7 @@ protected Optional> visitExprList( */ public Optional visitFieldReference( FieldReference fieldReference, EmptyVisitationContext context) throws E { - Optional inputExpression = - visitOptionalExpression(fieldReference.inputExpression(), context); - if (allEmpty(inputExpression)) { - return Optional.empty(); - } - - return Optional.of(FieldReference.builder().inputExpression(inputExpression).build()); + return getExpressionCopyOnWriteVisitor().visitFieldReference(fieldReference, context); } /** diff --git a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java new file mode 100644 index 000000000..dfd765d49 --- /dev/null +++ b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java @@ -0,0 +1,253 @@ +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.assertInstanceOf; + +import io.substrait.TestBase; +import io.substrait.expression.Expression; +import io.substrait.expression.FieldReference; +import io.substrait.expression.FunctionArg; +import io.substrait.expression.WindowBound; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.SimpleExtension; +import io.substrait.relation.physical.ScatterExchange; +import io.substrait.type.NamedStruct; +import io.substrait.type.Type; +import io.substrait.util.EmptyVisitationContext; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Test; + +/** + * Covers the re-derivation of the type cached on a {@link FieldReference} when a copy-on-write + * visitation replaces a relation with one that emits a different record type. + */ +class RelCopyOnWriteVisitorTest extends TestBase { + + private Rel scan(String table, Type... columnTypes) { + return sb.namedScan( + Arrays.asList(table), + Arrays.asList("a", "b", "c").subList(0, columnTypes.length), + Arrays.asList(columnTypes)); + } + + /** + * Rewrites the given relation tree, replacing every {@link NamedScan} in it with one that reads + * the given column types instead of the ones it was built with. + */ + private static Rel replaceScanTypes(Rel rel, Type... columnTypes) { + RelCopyOnWriteVisitor visitor = + new RelCopyOnWriteVisitor() { + @Override + public Optional visit(NamedScan namedScan, EmptyVisitationContext context) { + return Optional.of( + NamedScan.builder() + .from(namedScan) + .initialSchema( + NamedStruct.of(namedScan.getInitialSchema().names(), R.struct(columnTypes))) + .build()); + } + }; + return rel.accept(visitor, EmptyVisitationContext.INSTANCE) + .orElseThrow(() -> new AssertionError("expected the visitation to replace the scan")); + } + + private static List argumentTypes(List arguments) { + return arguments.stream() + .map(argument -> ((Expression) argument).getType()) + .collect(Collectors.toList()); + } + + private static List argumentTypes(Expression expression) { + return argumentTypes( + assertInstanceOf(Expression.ScalarFunctionInvocation.class, expression).arguments()); + } + + @Test + void projectExpression() { + Rel plan = sb.project(input -> sb.fieldReferences(input, 0), scan("t", R.I64)); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getExpressions().get(0).getType()); + // A project derives its record type from its expressions, so it follows the reference. + assertEquals(R.struct(N.I64, N.I64), rewritten.getRecordType()); + } + + @Test + void filterCondition() { + Rel plan = sb.filter(input -> sb.fieldReference(input, 0), scan("t", R.I64)); + + Filter rewritten = assertInstanceOf(Filter.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getCondition().getType()); + } + + @Test + void sortField() { + Rel plan = sb.sort(input -> sb.sortFields(input, 0), scan("t", R.I64)); + + Sort rewritten = assertInstanceOf(Sort.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getSortFields().get(0).expr().getType()); + } + + @Test + void aggregateGroupingAndMeasure() { + Rel input = scan("t", R.I64); + Rel plan = + Aggregate.builder() + .input(input) + .addGroupings(sb.grouping(input, 0)) + .addMeasures(sb.max(sb.fieldReference(input, 0))) + .build(); + + Aggregate rewritten = assertInstanceOf(Aggregate.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getGroupings().get(0).getExpressions().get(0).getType()); + assertEquals( + Arrays.asList(N.I64), + argumentTypes(rewritten.getMeasures().get(0).getFunction().arguments())); + } + + @Test + void joinConditionSpansBothInputs() { + Rel plan = + sb.innerJoin( + inputs -> sb.equal(sb.fieldReference(inputs, 0), sb.fieldReference(inputs, 1)), + scan("l", R.I64), + scan("r", R.I64)); + + Join rewritten = assertInstanceOf(Join.class, replaceScanTypes(plan, N.I64)); + assertEquals( + Arrays.asList(N.I64, N.I64), argumentTypes(rewritten.getCondition().orElseThrow())); + } + + @Test + void exchangeFields() { + Rel input = scan("t", R.I64); + Rel plan = + ScatterExchange.builder() + .input(input) + .partitionCount(2) + .addFields(sb.fieldReference(input, 0)) + .build(); + + ScatterExchange rewritten = + assertInstanceOf(ScatterExchange.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getFields().get(0).getType()); + } + + @Test + void windowRelationInputIsTraversedAndReferencesRetyped() { + SimpleExtension.WindowFunctionVariant lead = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + Rel input = scan("t", R.I64); + Rel plan = + ConsistentPartitionWindow.builder() + .input(input) + .addPartitionExpressions(sb.fieldReference(input, 0)) + .sorts(sb.sortFields(input, 0)) + .addWindowFunctions( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(lead) + .addArguments(sb.fieldReference(input, 0)) + .outputType(R.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .boundsType(Expression.WindowBoundsType.RANGE) + .lowerBound(WindowBound.Unbounded.UNBOUNDED) + .upperBound(WindowBound.Following.CURRENT_ROW) + .build()) + .build(); + + ConsistentPartitionWindow rewritten = + assertInstanceOf(ConsistentPartitionWindow.class, replaceScanTypes(plan, N.I64)); + // The input of a window relation used not to be traversed at all. + assertEquals(R.struct(N.I64), rewritten.getInput().getRecordType()); + assertEquals(N.I64, rewritten.getPartitionExpressions().get(0).getType()); + assertEquals(N.I64, rewritten.getSorts().get(0).expr().getType()); + assertEquals( + Arrays.asList(N.I64), argumentTypes(rewritten.getWindowFunctions().get(0).arguments())); + } + + @Test + void referenceRootedAtAnotherExpression() { + // A reference into a struct-typed column is rooted at the reference to that column rather than + // at the input relation, so its type comes from the rewritten root expression. + Rel input = scan("t", R.struct(R.I64)); + Rel plan = + sb.project( + in -> Arrays.asList(FieldReference.newStructReference(0, sb.fieldReference(in, 0))), + input); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, R.struct(N.I64))); + FieldReference reference = + assertInstanceOf(FieldReference.class, rewritten.getExpressions().get(0)); + assertEquals(N.I64, reference.getType()); + assertEquals(R.struct(N.I64), reference.inputExpression().orElseThrow().getType()); + // Rewriting the root must not drop the segments the reference navigates through. + assertEquals(1, reference.segments().size()); + } + + @Test + void outerReferenceInCorrelatedSubquery() { + Rel correlated = + sb.filter( + in -> + sb.equal( + sb.fieldReference(in, 0), + FieldReference.newRootStructOuterReference(0, R.I64, 1)), + scan("inner", R.I64)); + Rel plan = sb.filter(in -> sb.exists(correlated), scan("outer", R.I64)); + + Filter rewritten = assertInstanceOf(Filter.class, replaceScanTypes(plan, N.I64)); + Expression.SetPredicate exists = + assertInstanceOf(Expression.SetPredicate.class, rewritten.getCondition()); + Filter innerFilter = assertInstanceOf(Filter.class, exists.tuples()); + // The first argument resolves against the subquery's own input, the second steps out one level + // to the relation the subquery is correlated with. + assertEquals(Arrays.asList(N.I64, N.I64), argumentTypes(innerFilter.getCondition())); + } + + @Test + void readRelationFilterIsNotRetypedAgainstTheEnclosingScope() { + // The filter of a read relation resolves against the schema being read, so the record type of + // the enclosing relation's input must not be applied to it. + Rel scan = + NamedScan.builder() + .from((NamedScan) scan("t", R.I64, R.STRING)) + .filter(FieldReference.newRootStructReference(1, R.STRING)) + .build(); + Rel plan = sb.project(input -> sb.fieldReferences(input, 0), scan); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, N.I64, R.STRING)); + NamedScan rewrittenScan = assertInstanceOf(NamedScan.class, rewritten.getInput()); + assertEquals(R.STRING, rewrittenScan.getFilter().orElseThrow().getType()); + } + + @Test + void referenceBeyondTheNewInputIsLeftAlone() { + // A rewrite that drops a column leaves the reference to it selecting a field the input no + // longer has. Its type cannot be derived, and that must not fail the rewrite. + Rel plan = sb.project(input -> sb.fieldReferences(input, 1), scan("t", R.I64, R.STRING)); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, N.I64)); + assertEquals(R.STRING, rewritten.getExpressions().get(0).getType()); + } + + @Test + void unchangedPlanIsNotCopied() { + Rel plan = + sb.project( + input -> sb.fieldReferences(input, 0), + sb.filter(input -> sb.fieldReference(input, 0), scan("t", R.I64))); + + // Re-deriving the reference types must not by itself report the plan as changed: for a + // visitation that replaces nothing, the derived types are the ones already cached. + assertFalse( + plan.accept(new RelCopyOnWriteVisitor(), EmptyVisitationContext.INSTANCE) + .isPresent()); + } +} From d86b8b10007a9c22d774607f7601ff774ff50612 Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Wed, 5 Aug 2026 15:14:53 +0200 Subject: [PATCH 2/2] fix(core): resolve reference segments totally, and scope join keys per side MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two defects in the field-reference retyping, both found reviewing it. Segment resolution was guarded one segment deep. The guard checked only that the outermost segment selected a field the new record type has, so a reference into a nested type that had been reshaped reached the derivation anyway and threw — a bare IndexOutOfBoundsException out of a struct-field finder, or an IllegalArgumentException or UnsupportedOperationException for a list or map segment. That contradicted the guard's own promise to leave an unresolvable reference alone, and it made a narrowing rewrite fail where it previously produced a stale type. Resolution is now total: FieldReference.resolveType reports the type a chain of segments selects, or nothing, at any depth and for every kind of segment. It lives beside the finders whose rules it mirrors, so the two cannot drift apart unnoticed without the parity test that pins them going red. It mirrors those rules exactly, including the two asymmetries that matter: a list element offset is not bounds checked, because the length of a list is not part of its type, and a map key type is compared exactly, nullability included. The visitor's one-deep guard and its defensive segment copy both go, and the same guarantee now covers a reference rooted at another expression, which had no guard at all. Join keys were retyped against the wrong scope. The offsets of a hash or merge join key are relative to the side of the join the key selects from, not to the two inputs combined — proto conversion types each side with its own converter, and only the condition, post-join filter and residual expression use the combined type. Retyping both sides against the combined type silently resolved a right-side offset to a left column. Each side is now rewritten against its own input, which changes the signature of visitComparisonJoinKey; the method could not previously return a usable reference at all, so nothing can have depended on the old one. Two things this deliberately does not do: the off-by-one bound check in StructFieldFinder stays, because the exception it produces is part of what ProtoExpressionConverter reports for a malformed plan; and MergeJoin.deriveRecordType reads its right input for both sides, which is a separate bug in a relation this change only passes through. --- .../substrait/expression/FieldReference.java | 73 ++- .../ExpressionCopyOnWriteVisitor.java | 60 +-- .../relation/RelCopyOnWriteVisitor.java | 38 +- .../FieldReferenceResolveTypeTest.java | 201 ++++++++ .../relation/RelCopyOnWriteVisitorTest.java | 433 +++++++++++++++++- 5 files changed, 746 insertions(+), 59 deletions(-) create mode 100644 core/src/test/java/io/substrait/expression/FieldReferenceResolveTypeTest.java diff --git a/core/src/main/java/io/substrait/expression/FieldReference.java b/core/src/main/java/io/substrait/expression/FieldReference.java index 92e726fe3..a8ceeaff9 100644 --- a/core/src/main/java/io/substrait/expression/FieldReference.java +++ b/core/src/main/java/io/substrait/expression/FieldReference.java @@ -480,7 +480,7 @@ public FieldReference constructOnRoot(Type.Struct struct) { * Creates a field reference rooted at an expression and navigating through the given segments. * * @param expression the expression to reference into - * @param segments the navigation segments, outermost first + * @param segments the navigation segments, innermost first * @return the field reference */ public static FieldReference ofExpression( @@ -509,13 +509,82 @@ private static FieldReference of( * Creates a field reference rooted at a struct and navigating through the given segments. * * @param struct the root struct type - * @param segments the navigation segments, outermost first + * @param segments the navigation segments, innermost first * @return the field reference */ public static FieldReference ofRoot(Type.Struct struct, List segments) { return of(struct, null, segments); } + /** + * Resolves the type that the given reference segments select out of the given type, or reports + * that they select nothing out of it. + * + *

The segments are given innermost first, in the order {@link #segments()} holds them, and are + * applied outermost first: the last segment selects out of {@code rootType}, the one before it + * out of the type that segment selected, and so on inwards. The given list is not modified. + * + *

Resolution is total. A segment that does not fit the type it is applied to, at any depth — a + * struct field offset the struct does not have, a list element or a map key on a type that is not + * a list or a map, a map key whose type is not the map's key type — yields an empty result + * instead of throwing. This lets a caller that re-derives the type cached on a reference, against + * a type that has since changed, tell a reference that no longer resolves from a failure of its + * own work, and leave such a reference as it is. An empty segment list selects nothing and also + * yields an empty result, matching {@link #ofRoot} and {@link #ofExpression}, which build no + * reference for it. + * + * @param rootType the type the outermost segment selects out of + * @param segments the navigation segments, innermost first + * @return the type the segments select, or empty if they do not all resolve against {@code + * rootType} + */ + public static Optional resolveType(Type rootType, List segments) { + if (segments.isEmpty()) { + return Optional.empty(); + } + Type resolved = rootType; + for (int i = segments.size() - 1; i >= 0; i--) { + Optional selected = resolveSegmentType(segments.get(i), resolved); + if (!selected.isPresent()) { + return Optional.empty(); + } + resolved = selected.get(); + } + return Optional.of(resolved); + } + + /** + * Resolves the type a single segment selects out of the given type, or reports that it selects + * nothing out of it. + * + *

This mirrors the type each segment derives when it is applied: a struct field selects the + * field at its offset, a list element selects the element type whatever its offset, as the length + * of a list is not part of its type, and a map key selects the value type of a map whose key type + * it matches, nullability included. A segment applied to a type that is not the container it + * navigates into, and any other kind of segment, select nothing. + */ + private static Optional resolveSegmentType(ReferenceSegment segment, Type type) { + if (segment instanceof StructField && type instanceof Type.Struct) { + int offset = ((StructField) segment).offset(); + List fields = ((Type.Struct) type).fields(); + return offset >= 0 && offset < fields.size() + ? Optional.of(fields.get(offset)) + : Optional.empty(); + } + if (segment instanceof ListElement && type instanceof Type.ListType) { + return Optional.of(((Type.ListType) type).elementType()); + } + if (segment instanceof MapKey && type instanceof Type.Map) { + // The type of the key literal is only read once the type is known to be a map, which is also + // the only case in which applying the segment reads it. + Type.Map map = (Type.Map) type; + return map.key().equals(((MapKey) segment).key().getType()) + ? Optional.of(map.value()) + : Optional.empty(); + } + return Optional.empty(); + } + private static class StructFieldFinder extends TypeVisitor.TypeThrowsVisitor { diff --git a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java index daa8c83ed..9efc4c022 100644 --- a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java @@ -12,7 +12,6 @@ import io.substrait.expression.ImmutableFieldReference; import io.substrait.type.Type; import io.substrait.util.EmptyVisitationContext; -import java.util.ArrayList; import java.util.List; import java.util.Optional; @@ -435,11 +434,21 @@ public Optional visit(FieldReference fieldReference, EmptyVisitation * cached type from the root it resolves against. Re-deriving the type is what keeps references * correct when a relation's input is replaced by one emitting a different record type. * + *

A reference that no longer resolves against its root, at any segment depth, is left exactly + * as it is, including the type it has cached. Such a reference makes the relation tree invalid + * whatever this visitor does with it, so re-deriving its type is not turned into a failure of the + * rewrite. + * *

The type of a reference to a lambda parameter, or of an outer reference identified by the * {@link io.substrait.relation.Rel#getRelAnchor() rel anchor} of the relation it is rooted on, is * left as it is: neither resolves against a relation in the enclosing scopes tracked during the * traversal. * + *

Override this rather than {@link #visit(FieldReference, EmptyVisitationContext)} to change + * how references are rewritten: this is what the positions that hold a reference rather than an + * arbitrary expression — a {@link io.substrait.relation.physical.ScatterExchange}'s fields and a + * {@link io.substrait.relation.physical.ComparisonJoinKey}'s sides — are rewritten through. + * * @param fieldReference the field reference to visit * @param context the visitation context * @return Optional containing the modified field reference, or empty if no changes @@ -453,12 +462,12 @@ public Optional visitFieldReference( if (!inputExpression.isPresent()) { return Optional.empty(); } + // The reference is returned rewritten even when its type could not be re-derived: the + // expression it is rooted at changed, and that change is not lost because the type is stale. ImmutableFieldReference.Builder rewritten = ImmutableFieldReference.builder().from(fieldReference).inputExpression(inputExpression); - if (!fieldReference.segments().isEmpty()) { - rewritten.type( - FieldReference.ofExpression(inputExpression.get(), segmentsOf(fieldReference)).type()); - } + FieldReference.resolveType(inputExpression.get().getType(), fieldReference.segments()) + .ifPresent(rewritten::type); return Optional.of(rewritten.build()); } return retypeRootReference(fieldReference); @@ -476,40 +485,21 @@ private Optional retypeRootReference(FieldReference fieldReferen Type.Struct rootType = getRelCopyOnWriteVisitor() .inputTypeStepsOut(fieldReference.outerReferenceStepsOut().orElse(0)); - // A rewrite that drops fields can leave a reference selecting a field its input no longer has. - // The resulting relation tree is invalid either way, so leave the reference as it is rather - // than turn re-deriving its type into a failure of the rewrite. - if (rootType == null || !selectsFieldOf(fieldReference, rootType)) { + if (rootType == null) { return Optional.empty(); } - Type type = FieldReference.ofRoot(rootType, segmentsOf(fieldReference)).type(); - if (type.equals(fieldReference.type())) { + // A rewrite that drops a field, or that reshapes a nested one, can leave a reference selecting + // something its input no longer has. The resulting relation tree is invalid either way, so + // leave + // the reference as it is rather than turn re-deriving its type into a failure of the rewrite. + // Resolution is total, at every segment depth and for every kind of segment, so this is a + // decision the rewrite makes rather than an exception it has to recover from. + Optional type = FieldReference.resolveType(rootType, fieldReference.segments()); + if (!type.isPresent() || type.get().equals(fieldReference.type())) { return Optional.empty(); } - return Optional.of(ImmutableFieldReference.builder().from(fieldReference).type(type).build()); - } - - /** - * Returns whether the outermost segment of the given reference selects a field that the given - * record type has. The segments are held innermost first, so the outermost one is the last. - */ - private static boolean selectsFieldOf(FieldReference fieldReference, Type.Struct rootType) { - List segments = fieldReference.segments(); - if (segments.isEmpty()) { - return false; - } - FieldReference.ReferenceSegment outermost = segments.get(segments.size() - 1); - return outermost instanceof FieldReference.StructField - && ((FieldReference.StructField) outermost).offset() < rootType.fields().size(); - } - - /** - * Returns the segments of the given reference as a mutable list, which is what {@link - * FieldReference#ofRoot} and {@link FieldReference#ofExpression} require: they reverse the list - * they are given. - */ - private static List segmentsOf(FieldReference fieldReference) { - return new ArrayList<>(fieldReference.segments()); + return Optional.of( + ImmutableFieldReference.builder().from(fieldReference).type(type.get()).build()); } @Override diff --git a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java index 2198b907e..7b40941f0 100644 --- a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java @@ -609,12 +609,15 @@ public Optional visit(ExtensionTable extensionTable, EmptyVisitationContext public Optional visit(HashJoin hashJoin, EmptyVisitationContext context) throws E { Optional left = hashJoin.getLeft().accept(this, context); Optional right = hashJoin.getRight().accept(this, context); + Type.Struct leftType = recordTypeOf(left.orElse(hashJoin.getLeft())); + Type.Struct rightType = recordTypeOf(right.orElse(hashJoin.getRight())); Type.Struct inputType = recordTypeOf(left.orElse(hashJoin.getLeft()), right.orElse(hashJoin.getRight())); Optional> keys = - inInputScope( - inputType, - () -> transformList(hashJoin.getKeys(), context, this::visitComparisonJoinKey)); + transformList( + hashJoin.getKeys(), + context, + (key, c) -> visitComparisonJoinKey(key, leftType, rightType, c)); Optional postFilter = inInputScope( inputType, () -> visitOptionalExpression(hashJoin.getPostJoinFilter(), context)); @@ -640,12 +643,15 @@ public Optional visit(HashJoin hashJoin, EmptyVisitationContext context) th public Optional visit(MergeJoin mergeJoin, EmptyVisitationContext context) throws E { Optional left = mergeJoin.getLeft().accept(this, context); Optional right = mergeJoin.getRight().accept(this, context); + Type.Struct leftType = recordTypeOf(left.orElse(mergeJoin.getLeft())); + Type.Struct rightType = recordTypeOf(right.orElse(mergeJoin.getRight())); Type.Struct inputType = recordTypeOf(left.orElse(mergeJoin.getLeft()), right.orElse(mergeJoin.getRight())); Optional> keys = - inInputScope( - inputType, - () -> transformList(mergeJoin.getKeys(), context, this::visitComparisonJoinKey)); + transformList( + mergeJoin.getKeys(), + context, + (key, c) -> visitComparisonJoinKey(key, leftType, rightType, c)); Optional postFilter = inInputScope( inputType, () -> visitOptionalExpression(mergeJoin.getPostJoinFilter(), context)); @@ -761,6 +767,8 @@ protected Optional visitW * Returns the record type that the root field references of a relation's own expressions resolve * against: the record types of the given inputs, concatenated in order. * + *

Only the fields of the result are ever read, so its own nullability is not meaningful. + * * @param inputs the relations the expressions are evaluated over, in field order * @return the combined record type */ @@ -872,15 +880,27 @@ public Optional visitFieldReference( /** * Rewrites a comparison join key, returning a new one if either side changed. * + *

Each side is rewritten against its own input, because the field offsets of a join key are + * relative to the side of the join they select from — unlike those of a join condition or + * post-join filter, which are relative to the two inputs combined. + * * @param key the comparison join key to rewrite + * @param leftType the record type the key's left side selects from + * @param rightType the record type the key's right side selects from * @param context the visitation context * @return the rewritten comparison join key, or empty if unchanged * @throws E if the visit fails */ public Optional visitComparisonJoinKey( - ComparisonJoinKey key, EmptyVisitationContext context) throws E { - Optional left = visitFieldReference(key.getLeft(), context); - Optional right = visitFieldReference(key.getRight(), context); + ComparisonJoinKey key, + Type.Struct leftType, + Type.Struct rightType, + EmptyVisitationContext context) + throws E { + Optional left = + inInputScope(leftType, () -> visitFieldReference(key.getLeft(), context)); + Optional right = + inInputScope(rightType, () -> visitFieldReference(key.getRight(), context)); if (allEmpty(left, right)) { return Optional.empty(); } diff --git a/core/src/test/java/io/substrait/expression/FieldReferenceResolveTypeTest.java b/core/src/test/java/io/substrait/expression/FieldReferenceResolveTypeTest.java new file mode 100644 index 000000000..d45c73b3c --- /dev/null +++ b/core/src/test/java/io/substrait/expression/FieldReferenceResolveTypeTest.java @@ -0,0 +1,201 @@ +package io.substrait.expression; + +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.TestBase; +import io.substrait.expression.FieldReference.ListElement; +import io.substrait.expression.FieldReference.MapKey; +import io.substrait.expression.FieldReference.ReferenceSegment; +import io.substrait.expression.FieldReference.StructField; +import io.substrait.type.Type; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Pins {@link FieldReference#resolveType} against the reference-building factories it mirrors. + * + *

{@code resolveType} reports the type a chain of segments selects without throwing when it + * selects nothing, while {@link FieldReference#ofRoot} and {@link FieldReference#ofExpression} + * throw in that case. The two must agree on which chains select something, so that a + * caller re-deriving a cached type can tell a reference that no longer resolves from a failure of + * its own work. Asserting the equivalence rather than hard-coding expectations is what keeps {@code + * resolveType} from drifting away from the segment derivation rules it duplicates. + */ +class FieldReferenceResolveTypeTest extends TestBase { + + /** Segments are held innermost first, the order {@link FieldReference#segments()} uses. */ + private static List segments(ReferenceSegment... innermostFirst) { + return Arrays.asList(innermostFirst); + } + + private static MapKey key(String value) { + return MapKey.of(Expression.StrLiteral.builder().value(value).build()); + } + + static Stream cases() { + return Stream.of( + // struct field, single segment + Arguments.of("struct field in range", R.struct(R.I64), segments(StructField.of(0))), + Arguments.of( + "struct field, second column", R.struct(R.I64, R.STRING), segments(StructField.of(1))), + Arguments.of("struct field past the end", R.struct(R.I64), segments(StructField.of(2))), + Arguments.of( + "struct field at the field count", R.struct(R.I64), segments(StructField.of(1))), + Arguments.of("negative struct field", R.struct(R.I64), segments(StructField.of(-1))), + // struct field, nested + Arguments.of( + "nested struct field in range", + R.struct(R.struct(R.I64, R.STRING)), + segments(StructField.of(1), StructField.of(0))), + Arguments.of( + "nested struct field past the end", + R.struct(R.struct(R.I64)), + segments(StructField.of(1), StructField.of(0))), + Arguments.of( + "three struct fields deep, in range", + R.struct(R.struct(R.struct(R.I64, R.STRING))), + segments(StructField.of(1), StructField.of(0), StructField.of(0))), + Arguments.of( + "three struct fields deep, innermost gone", + R.struct(R.struct(R.struct(R.I64))), + segments(StructField.of(1), StructField.of(0), StructField.of(0))), + // container kind mismatches under a struct field + Arguments.of( + "struct field on a list", + R.struct(R.list(R.I64)), + segments(StructField.of(0), StructField.of(0))), + Arguments.of( + "struct field on a map", + R.struct(R.map(R.STRING, R.I64)), + segments(StructField.of(0), StructField.of(0))), + Arguments.of( + "struct field on a scalar", + R.struct(R.I64), + segments(StructField.of(0), StructField.of(0))), + // list element + Arguments.of( + "list element on a list", + R.struct(R.list(R.I64)), + segments(ListElement.of(0), StructField.of(0))), + Arguments.of( + "list element offset is not bounds checked", + R.struct(R.list(R.I64)), + segments(ListElement.of(7), StructField.of(0))), + Arguments.of( + "list element on a struct", + R.struct(R.struct(R.I64)), + segments(ListElement.of(0), StructField.of(0))), + Arguments.of( + "list element as the outermost segment", R.struct(R.I64), segments(ListElement.of(0))), + // map key + Arguments.of( + "map key matching the key type", + R.struct(R.map(R.STRING, R.I64)), + segments(key("k"), StructField.of(0))), + Arguments.of( + "map key differing in nullability", + R.struct(R.map(N.STRING, R.I64)), + segments(key("k"), StructField.of(0))), + Arguments.of( + "map key of the wrong type", + R.struct(R.map(R.I64, R.I64)), + segments(key("k"), StructField.of(0))), + Arguments.of( + "map key on a list", R.struct(R.list(R.I64)), segments(key("k"), StructField.of(0))), + Arguments.of("map key as the outermost segment", R.struct(R.I64), segments(key("k"))), + // degenerate + Arguments.of("no segments", R.struct(R.I64), Collections.emptyList())); + } + + /** + * Resolves via {@link FieldReference#ofRoot}, mapping its failure modes — every exception it can + * throw, and the null it returns for an empty segment chain — onto "selects nothing". + */ + private static Optional viaOfRoot(Type.Struct rootType, List segments) { + try { + FieldReference reference = FieldReference.ofRoot(rootType, new ArrayList<>(segments)); + return reference == null ? Optional.empty() : Optional.of(reference.type()); + } catch (RuntimeException e) { + return Optional.empty(); + } + } + + /** + * The same, via {@link FieldReference#ofExpression} rooted at an expression of {@code rootType}. + */ + private static Optional viaOfExpression(Type rootType, List segments) { + Expression root = FieldReference.newRootStructReference(0, rootType); + try { + FieldReference reference = FieldReference.ofExpression(root, new ArrayList<>(segments)); + return reference == null ? Optional.empty() : Optional.of(reference.type()); + } catch (RuntimeException e) { + return Optional.empty(); + } + } + + @ParameterizedTest(name = "{0}") + @MethodSource("cases") + void agreesWithOfRoot(String name, Type.Struct rootType, List segments) { + assertEquals(viaOfRoot(rootType, segments), FieldReference.resolveType(rootType, segments)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("cases") + void agreesWithOfExpression(String name, Type.Struct rootType, List segments) { + assertEquals( + viaOfExpression(rootType, segments), FieldReference.resolveType(rootType, segments)); + } + + @Test + void resolvesAgainstRootsThatAreNotStructs() { + // ofRoot only accepts a struct, but a reference rooted at an expression can navigate into a + // list or a map directly, so resolveType has to accept any type as the root. + assertEquals( + Optional.of(R.I64), FieldReference.resolveType(R.list(R.I64), segments(ListElement.of(0)))); + assertEquals( + Optional.of(R.I64), FieldReference.resolveType(R.map(R.STRING, R.I64), segments(key("k")))); + assertFalse(FieldReference.resolveType(R.I64, segments(StructField.of(0))).isPresent()); + } + + @Test + void doesNotModifyTheGivenSegments() { + // ofRoot and ofExpression reverse the list they are given in place, which is why they cannot be + // handed FieldReference.segments() directly. resolveType must not have that requirement. + List segments = + Collections.unmodifiableList( + Arrays.asList(StructField.of(1), StructField.of(0))); + + assertEquals( + Optional.of(R.STRING), + FieldReference.resolveType(R.struct(R.struct(R.I64, R.STRING)), segments)); + assertEquals(StructField.of(1), segments.get(0)); + assertEquals(StructField.of(0), segments.get(1)); + } + + @Test + void resolvesTheSegmentsOfAReferenceItselfUnchanged() { + // The end the whole method exists for: taking segments() straight off a reference and resolving + // them against a record type, with no defensive copy at the call site. + FieldReference reference = + FieldReference.newRootStructReference(0, R.struct(R.I64, R.STRING)).dereferenceStruct(1); + + assertEquals(2, reference.segments().size()); + assertEquals( + Optional.of(R.STRING), + FieldReference.resolveType(R.struct(R.struct(R.I64, R.STRING)), reference.segments())); + assertTrue( + FieldReference.resolveType(R.struct(R.struct(R.I64, R.STRING)), reference.segments()) + .isPresent()); + assertEquals(2, reference.segments().size()); + } +} diff --git a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java index dfd765d49..cc334cd08 100644 --- a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java +++ b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java @@ -7,11 +7,19 @@ import io.substrait.TestBase; import io.substrait.expression.Expression; import io.substrait.expression.FieldReference; +import io.substrait.expression.FieldReference.ListElement; +import io.substrait.expression.FieldReference.MapKey; +import io.substrait.expression.FieldReference.StructField; import io.substrait.expression.FunctionArg; import io.substrait.expression.WindowBound; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.SimpleExtension; +import io.substrait.relation.physical.ComparisonJoinKey; +import io.substrait.relation.physical.HashJoin; +import io.substrait.relation.physical.MultiBucketExchange; import io.substrait.relation.physical.ScatterExchange; +import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TopN; import io.substrait.type.NamedStruct; import io.substrait.type.Type; import io.substrait.util.EmptyVisitationContext; @@ -39,22 +47,54 @@ private Rel scan(String table, Type... columnTypes) { * the given column types instead of the ones it was built with. */ private static Rel replaceScanTypes(Rel rel, Type... columnTypes) { - RelCopyOnWriteVisitor visitor = - new RelCopyOnWriteVisitor() { - @Override - public Optional visit(NamedScan namedScan, EmptyVisitationContext context) { - return Optional.of( - NamedScan.builder() - .from(namedScan) - .initialSchema( - NamedStruct.of(namedScan.getInitialSchema().names(), R.struct(columnTypes))) - .build()); - } - }; - return rel.accept(visitor, EmptyVisitationContext.INSTANCE) + return rel.accept(scanTypeReplacer(null, columnTypes), EmptyVisitationContext.INSTANCE) .orElseThrow(() -> new AssertionError("expected the visitation to replace the scan")); } + /** The same, replacing only the {@link NamedScan} that reads the given table. */ + private static Rel replaceScanTypesOf(Rel rel, String table, Type... columnTypes) { + return rel.accept(scanTypeReplacer(table, columnTypes), EmptyVisitationContext.INSTANCE) + .orElseThrow(() -> new AssertionError("expected the visitation to replace the scan")); + } + + private static RelCopyOnWriteVisitor scanTypeReplacer( + String table, Type... columnTypes) { + return new RelCopyOnWriteVisitor() { + @Override + public Optional visit(NamedScan namedScan, EmptyVisitationContext context) { + if (table != null && !namedScan.getNames().equals(Arrays.asList(table))) { + return Optional.empty(); + } + return Optional.of( + NamedScan.builder() + .from(namedScan) + .initialSchema( + NamedStruct.of(namedScan.getInitialSchema().names(), R.struct(columnTypes))) + .build()); + } + }; + } + + /** + * Projects a reference that navigates the given segments into the single column of a scan of + * {@code columnType}, then rewrites the plan with that column replaced by {@code + * replacementType}, and returns the resulting reference. The segments are given outermost first, + * the order they are navigated in. + */ + private FieldReference rewriteNestedReference( + Type columnType, Type replacementType, FieldReference.ReferenceSegment... segments) { + Rel input = scan("t", columnType); + FieldReference navigated = sb.fieldReference(input, 0); + for (FieldReference.ReferenceSegment segment : segments) { + navigated = segment.apply(navigated); + } + FieldReference reference = navigated; + Rel plan = sb.project(in -> Arrays.asList(reference), input); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, replacementType)); + return assertInstanceOf(FieldReference.class, rewritten.getExpressions().get(0)); + } + private static List argumentTypes(List arguments) { return arguments.stream() .map(argument -> ((Expression) argument).getType()) @@ -237,6 +277,373 @@ void referenceBeyondTheNewInputIsLeftAlone() { assertEquals(R.STRING, rewritten.getExpressions().get(0).getType()); } + @Test + void nestedStructFieldIsRetyped() { + FieldReference reference = + rewriteNestedReference(R.struct(R.I64), R.struct(N.I64), StructField.of(0)); + assertEquals(N.I64, reference.getType()); + assertEquals(2, reference.segments().size()); + } + + @Test + void nestedStructFieldTheInputNoLongerHasIsLeftAlone() { + // Resolution has to hold at every depth, not just for the segment that selects out of the + // record type: this reference selects a field of a struct column that has since shrunk. + FieldReference reference = + rewriteNestedReference(R.struct(R.I64, R.STRING), R.struct(R.I64), StructField.of(1)); + assertEquals(R.STRING, reference.getType()); + assertEquals(2, reference.segments().size()); + } + + @Test + void threeSegmentsDeepIsRetyped() { + FieldReference reference = + rewriteNestedReference( + R.struct(R.struct(R.I64, R.STRING)), + R.struct(R.struct(R.I64, N.STRING)), + StructField.of(0), + StructField.of(1)); + assertEquals(N.STRING, reference.getType()); + assertEquals(3, reference.segments().size()); + } + + @Test + void threeSegmentsDeepWithTheInnermostGoneIsLeftAlone() { + FieldReference reference = + rewriteNestedReference( + R.struct(R.struct(R.I64, R.STRING)), + R.struct(R.struct(R.I64)), + StructField.of(0), + StructField.of(1)); + assertEquals(R.STRING, reference.getType()); + assertEquals(3, reference.segments().size()); + } + + @Test + void structFieldSegmentOnAContainerIsLeftAlone() { + // The column keeps its field count but stops being a struct, so the inner segment no longer + // applies. Retyping it against the container's element or value type would be wrong. + assertEquals( + R.STRING, + rewriteNestedReference(R.struct(R.STRING), R.list(R.I64), StructField.of(0)).getType()); + assertEquals( + R.STRING, + rewriteNestedReference(R.struct(R.STRING), R.map(R.STRING, R.I64), StructField.of(0)) + .getType()); + assertEquals( + R.STRING, rewriteNestedReference(R.struct(R.STRING), R.I64, StructField.of(0)).getType()); + } + + @Test + void listElementSegmentIsRetyped() { + assertEquals( + N.I64, rewriteNestedReference(R.list(R.I64), R.list(N.I64), ListElement.of(0)).getType()); + // The length of a list is not part of its type, so the offset never affects the result and is + // deliberately not bounds checked. + assertEquals( + N.I64, rewriteNestedReference(R.list(R.I64), R.list(N.I64), ListElement.of(7)).getType()); + } + + @Test + void listElementSegmentOnANonListIsLeftAlone() { + assertEquals( + R.STRING, + rewriteNestedReference(R.list(R.STRING), R.struct(R.I64), ListElement.of(0)).getType()); + } + + @Test + void mapKeySegmentIsRetyped() { + assertEquals( + N.I64, + rewriteNestedReference( + R.map(R.STRING, R.I64), R.map(R.STRING, N.I64), MapKey.of(sb.str("k"))) + .getType()); + } + + @Test + void mapKeySegmentWhoseKeyTypeNoLongerMatchesIsLeftAlone() { + // The derivation compares the key type exactly, nullability included, so a map whose key became + // nullable no longer accepts this segment. + assertEquals( + R.STRING, + rewriteNestedReference( + R.map(R.STRING, R.STRING), R.map(N.STRING, R.I64), MapKey.of(sb.str("k"))) + .getType()); + } + + @Test + void mapKeySegmentOnANonMapIsLeftAlone() { + assertEquals( + R.STRING, + rewriteNestedReference(R.map(R.STRING, R.STRING), R.list(R.I64), MapKey.of(sb.str("k"))) + .getType()); + } + + @Test + void referenceThatDoesNotStartAtAStructFieldIsLeftAlone() { + // A root reference selects out of the input's record type, which is a struct, so a reference + // whose outermost segment is a list element or a map key cannot resolve against it. + for (FieldReference.ReferenceSegment outermost : + Arrays.asList(ListElement.of(0), MapKey.of(sb.str("k")))) { + Rel plan = + sb.project( + in -> + Arrays.asList( + FieldReference.builder().addSegments(outermost).type(R.STRING).build()), + scan("t", R.I64)); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, N.I64)); + assertEquals(R.STRING, rewritten.getExpressions().get(0).getType()); + } + } + + @Test + void referenceRootedAtAnotherExpressionThatNoLongerResolves() { + Rel input = scan("t", R.struct(R.I64, R.STRING)); + Rel plan = + sb.project( + in -> Arrays.asList(FieldReference.newStructReference(1, sb.fieldReference(in, 0))), + input); + + Project rewritten = assertInstanceOf(Project.class, replaceScanTypes(plan, R.struct(R.I64))); + FieldReference reference = + assertInstanceOf(FieldReference.class, rewritten.getExpressions().get(0)); + // The expression it is rooted at was rewritten, so the reference is rewritten too, but its own + // type is kept because the field it selects is gone. + assertEquals(R.STRING, reference.getType()); + assertEquals(R.struct(R.I64), reference.inputExpression().orElseThrow().getType()); + assertEquals(1, reference.segments().size()); + } + + @Test + void hashJoinKeysAreRetypedAgainstTheirOwnSide() { + Rel left = scan("l", R.I64, R.I64); + Rel right = scan("r", R.STRING); + // A join key's offsets are relative to the side it selects from, so the right key's offset is 0 + // even though that column is the third of the joined output. + Rel plan = + HashJoin.builder() + .left(left) + .right(right) + .joinType(HashJoin.JoinType.INNER) + .addKeys( + ComparisonJoinKey.of( + sb.fieldReference(left, 0), + sb.fieldReference(right, 0), + ComparisonJoinKey.SimpleComparisonType.EQ)) + .build(); + + HashJoin rewritten = assertInstanceOf(HashJoin.class, replaceScanTypesOf(plan, "r", N.STRING)); + ComparisonJoinKey key = rewritten.getKeys().get(0); + assertEquals(R.I64, key.getLeft().getType()); + // Resolving the right key against the two inputs combined would have found a left column here. + assertEquals(N.STRING, key.getRight().getType()); + } + + @Test + void hashJoinFiltersAreRetypedAgainstBothInputs() { + Rel left = scan("l", R.I64); + Rel right = scan("r", R.I64); + List inputs = Arrays.asList(left, right); + Rel plan = + HashJoin.builder() + .left(left) + .right(right) + .joinType(HashJoin.JoinType.INNER) + .postJoinFilter(sb.equal(sb.fieldReference(inputs, 0), sb.fieldReference(inputs, 1))) + .residualExpression( + sb.equal(sb.fieldReference(inputs, 1), sb.fieldReference(inputs, 0))) + .build(); + + HashJoin rewritten = assertInstanceOf(HashJoin.class, replaceScanTypesOf(plan, "r", N.I64)); + // Unlike the keys, these resolve against the concatenation of the two inputs. + assertEquals( + Arrays.asList(R.I64, N.I64), argumentTypes(rewritten.getPostJoinFilter().orElseThrow())); + assertEquals( + Arrays.asList(N.I64, R.I64), + argumentTypes(rewritten.getResidualExpression().orElseThrow())); + } + + @Test + void lateralJoinConditionSpansBothInputs() { + Rel left = scan("l", R.I64); + Rel right = scan("r", R.I64); + List inputs = Arrays.asList(left, right); + Rel plan = + LateralJoin.builder() + .left(left) + .right(right) + .joinType(Join.JoinType.INNER) + .relAnchor(1) + .condition(sb.equal(sb.fieldReference(inputs, 0), sb.fieldReference(inputs, 1))) + .build(); + + LateralJoin rewritten = + assertInstanceOf(LateralJoin.class, replaceScanTypesOf(plan, "l", N.I64)); + assertEquals( + Arrays.asList(N.I64, R.I64), argumentTypes(rewritten.getCondition().orElseThrow())); + } + + @Test + void anchorBasedOuterReferenceIsLeftAlone() { + // A lateral join's right input references the current left row by the join's rel anchor rather + // than by stepping out of subquery scopes. Resolving an anchor needs the whole plan, which this + // visitor does not track, so such a reference keeps its type and can be left stale. + Rel left = scan("l", R.I64); + Rel plan = + LateralJoin.builder() + .left(left) + .right( + sb.filter( + in -> + sb.equal( + sb.fieldReference(in, 0), + FieldReference.newRootStructOuterReferenceByRelReference(0, R.I64, 1)), + scan("r", R.I64))) + .joinType(Join.JoinType.INNER) + .relAnchor(1) + .build(); + + LateralJoin rewritten = + assertInstanceOf(LateralJoin.class, replaceScanTypesOf(plan, "l", N.I64)); + Filter right = assertInstanceOf(Filter.class, rewritten.getRight()); + assertEquals(Arrays.asList(R.I64, R.I64), argumentTypes(right.getCondition())); + } + + @Test + void topNSortFieldAndCount() { + Rel input = scan("t", R.I64); + Rel plan = + TopN.builder() + .input(input) + .sortFields(sb.sortFields(input, 0)) + .count(sb.fieldReference(input, 0)) + .build(); + + TopN rewritten = assertInstanceOf(TopN.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getSortFields().get(0).expr().getType()); + assertEquals(N.I64, rewritten.getCount().orElseThrow().getType()); + } + + @Test + void fetchOffsetAndCount() { + Rel input = scan("t", R.I64); + Rel plan = + Fetch.builder() + .input(input) + .offset(sb.fieldReference(input, 0)) + .count(sb.fieldReference(input, 0)) + .build(); + + Fetch rewritten = assertInstanceOf(Fetch.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getOffset().orElseThrow().getType()); + assertEquals(N.I64, rewritten.getCount().orElseThrow().getType()); + } + + @Test + void singleBucketExchangeExpression() { + Rel input = scan("t", R.I64); + Rel plan = + SingleBucketExchange.builder() + .input(input) + .partitionCount(2) + .expression(sb.fieldReference(input, 0)) + .build(); + + SingleBucketExchange rewritten = + assertInstanceOf(SingleBucketExchange.class, replaceScanTypes(plan, N.I64)); + assertEquals(N.I64, rewritten.getExpression().getType()); + } + + @Test + void multiBucketExchangeExpressionIsRetypedWithAnUnchangedInput() { + // Nothing here replaces the input; the reference simply carries a type its input does not emit. + // A guard that only asked whether the input had changed would discard the retyped expression. + Rel plan = + MultiBucketExchange.builder() + .input(scan("t", R.I64)) + .partitionCount(2) + .constrainedToCount(true) + .expression(FieldReference.newRootStructReference(0, N.I64)) + .build(); + + MultiBucketExchange rewritten = + assertInstanceOf( + MultiBucketExchange.class, + plan.accept( + new RelCopyOnWriteVisitor(), EmptyVisitationContext.INSTANCE) + .orElseThrow(() -> new AssertionError("expected the expression to be retyped"))); + assertEquals(R.I64, rewritten.getExpression().getType()); + } + + @Test + void outerReferenceInScalarSubquery() { + Rel correlated = + sb.filter( + in -> + sb.equal( + sb.fieldReference(in, 0), + FieldReference.newRootStructOuterReference(0, R.I64, 1)), + scan("inner", R.I64)); + Rel plan = + sb.filter( + in -> sb.equal(sb.fieldReference(in, 0), sb.scalarSubquery(correlated, R.I64)), + scan("outer", R.I64)); + + Filter rewritten = assertInstanceOf(Filter.class, replaceScanTypes(plan, N.I64)); + Expression.ScalarSubquery subquery = + assertInstanceOf( + Expression.ScalarSubquery.class, + assertInstanceOf(Expression.ScalarFunctionInvocation.class, rewritten.getCondition()) + .arguments() + .get(1)); + Filter innerFilter = assertInstanceOf(Filter.class, subquery.input()); + assertEquals(Arrays.asList(N.I64, N.I64), argumentTypes(innerFilter.getCondition())); + } + + @Test + void outerReferenceInInPredicateHaystack() { + Rel correlated = + sb.filter( + in -> + sb.equal( + sb.fieldReference(in, 0), + FieldReference.newRootStructOuterReference(0, R.I64, 1)), + scan("inner", R.I64)); + Rel plan = + sb.filter(in -> sb.inPredicate(correlated, sb.fieldReference(in, 0)), scan("outer", R.I64)); + + Filter rewritten = assertInstanceOf(Filter.class, replaceScanTypes(plan, N.I64)); + Expression.InPredicate inPredicate = + assertInstanceOf(Expression.InPredicate.class, rewritten.getCondition()); + // The needles are evaluated in the enclosing scope; only the haystack is a subquery boundary. + assertEquals(N.I64, inPredicate.needles().get(0).getType()); + Filter innerFilter = assertInstanceOf(Filter.class, inPredicate.haystack()); + assertEquals(Arrays.asList(N.I64, N.I64), argumentTypes(innerFilter.getCondition())); + } + + @Test + void aVisitorCanBeReusedForASecondTraversal() { + // The scope bookkeeping is pushed and popped around every rewrite, so a traversal leaves no + // residue behind that would mistype the next one. + RelCopyOnWriteVisitor visitor = scanTypeReplacer(null, N.I64); + Rel plan = + sb.project( + input -> sb.fieldReferences(input, 0), + sb.filter(input -> sb.fieldReference(input, 0), scan("t", R.I64))); + + for (int traversal = 0; traversal < 2; traversal++) { + Project rewritten = + assertInstanceOf( + Project.class, + plan.accept(visitor, EmptyVisitationContext.INSTANCE) + .orElseThrow(() -> new AssertionError("expected the scan to be replaced"))); + assertEquals(N.I64, rewritten.getExpressions().get(0).getType()); + assertEquals( + N.I64, assertInstanceOf(Filter.class, rewritten.getInput()).getCondition().getType()); + } + } + @Test void unchangedPlanIsNotCopied() { Rel plan =