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/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..9efc4c022 100644 --- a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java @@ -9,6 +9,8 @@ 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.List; import java.util.Optional; @@ -424,30 +426,95 @@ 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); + } + + /** + * 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. + * + *

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 + * @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(); + } + // 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); + FieldReference.resolveType(inputExpression.get().getType(), fieldReference.segments()) + .ifPresent(rewritten::type); + return Optional.of(rewritten.build()); + } + return retypeRootReference(fieldReference); + } - if (allEmpty(inputExpression)) { + /** + * 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(); + } + Type.Struct rootType = + getRelCopyOnWriteVisitor() + .inputTypeStepsOut(fieldReference.outerReferenceStepsOut().orElse(0)); + if (rootType == null) { + return Optional.empty(); + } + // 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(FieldReference.builder().inputExpression(inputExpression).build()); + return Optional.of( + ImmutableFieldReference.builder().from(fieldReference).type(type.get()).build()); } @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 +522,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 +594,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..7b40941f0 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,21 @@ 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 = - transformList(hashJoin.getKeys(), context, this::visitComparisonJoinKey); + transformList( + hashJoin.getKeys(), + context, + (key, c) -> visitComparisonJoinKey(key, leftType, rightType, c)); 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 +643,21 @@ 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 = - transformList(mergeJoin.getKeys(), context, this::visitComparisonJoinKey); + transformList( + mergeJoin.getKeys(), + context, + (key, c) -> visitComparisonJoinKey(key, leftType, rightType, c)); 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 +679,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 +700,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 +761,93 @@ 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. + * + *

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 + */ + 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 +864,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,27 +874,33 @@ 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); } /** * 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 new file mode 100644 index 000000000..cc334cd08 --- /dev/null +++ b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java @@ -0,0 +1,660 @@ +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.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; +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) { + 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()) + .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 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 = + 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()); + } +}