Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 71 additions & 2 deletions core/src/main/java/io/substrait/expression/FieldReference.java
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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<ReferenceSegment> 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.
*
* <p>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.
*
* <p>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<Type> resolveType(Type rootType, List<ReferenceSegment> segments) {
if (segments.isEmpty()) {
return Optional.empty();
}
Type resolved = rootType;
for (int i = segments.size() - 1; i >= 0; i--) {
Optional<Type> 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.
*
* <p>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<Type> resolveSegmentType(ReferenceSegment segment, Type type) {
if (segment instanceof StructField && type instanceof Type.Struct) {
int offset = ((StructField) segment).offset();
List<Type> 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<Type, RuntimeException> {

Expand Down
19 changes: 19 additions & 0 deletions core/src/main/java/io/substrait/relation/CopyOnWriteUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,25 @@ public static <T> Optional<T> or(Optional<T> left, Supplier<? extends Optional<T
}
}

/**
* A {@link Supplier} that is allowed to throw the exception type of the visitation it runs
* within. Used to scope a rewrite so that state set up around it is always torn down.
*
* @param <T> the type of the supplied value
* @param <E> the exception type that may be thrown
*/
@FunctionalInterface
public interface ThrowingSupplier<T, E extends Exception> {

/**
* 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.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -424,39 +426,108 @@ protected Optional<Expression.MultiOrListRecord> visitMultiOrListRecord(
@Override
public Optional<Expression> visit(FieldReference fieldReference, EmptyVisitationContext context)
throws E {
Optional<Expression> 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.
*
* <p>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.
*
* <p>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.
*
* <p>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<FieldReference> visitFieldReference(
FieldReference fieldReference, EmptyVisitationContext context) throws E {
if (fieldReference.inputExpression().isPresent()) {
Optional<Expression> 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<FieldReference> 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> 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<Expression> 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<Expression> 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());
}

@Override
public Optional<Expression> visit(
Expression.InPredicate inPredicate, EmptyVisitationContext context) throws E {
Optional<Rel> haystack = inPredicate.haystack().accept(getRelCopyOnWriteVisitor(), context);
// The needles are evaluated in the current scope; only the haystack is a subquery boundary.
Optional<List<Expression>> needles = visitExprList(inPredicate.needles(), context);
Optional<Rel> haystack =
getRelCopyOnWriteVisitor()
.inSubqueryScope(
() -> inPredicate.haystack().accept(getRelCopyOnWriteVisitor(), context));

if (allEmpty(haystack, needles)) {
return Optional.empty();
Expand Down Expand Up @@ -523,15 +594,6 @@ protected Optional<List<Expression>> visitExprList(
return transformList(exprs, context, (e, c) -> e.accept(this, c));
}

private Optional<Expression> visitOptionalExpression(
Optional<Expression> 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.
*
Expand Down
Loading
Loading