diff --git a/core/src/main/java/io/substrait/expression/proto/DetachedSubtrees.java b/core/src/main/java/io/substrait/expression/proto/DetachedSubtrees.java new file mode 100644 index 000000000..00e6c3585 --- /dev/null +++ b/core/src/main/java/io/substrait/expression/proto/DetachedSubtrees.java @@ -0,0 +1,133 @@ +package io.substrait.expression.proto; + +import io.substrait.proto.Expression.RexTypeCase; +import io.substrait.proto.ExtendedExpression; +import io.substrait.proto.PlanRel; +import java.util.BitSet; +import java.util.Collections; +import java.util.List; +import java.util.function.Function; + +/** + * The detached subtrees of one {@code PlanRel} or {@code ExtendedExpression}, resolving the {@code + * Expression.detached_expression_ordinal} references in it against its {@code detached_expressions} + * (spec v0.105.0). + * + *

Detachment is an encoding of the expression tree, so an ordinal is resolved by converting its + * entry in place of the reference: the entry binds to the input schema and correlation scope of the + * referencing site. Resolution enforces the rules the spec places on the portion of the table + * reachable from the container: an ordinal is in range, refers to an entry that is not itself an + * ordinal, is referenced at most once, and is not part of a cycle. Entries that nothing references + * are neither converted nor validated. + * + *

An instance records which entries it has resolved, so it serves one conversion of one + * container. The converted model inlines every resolved entry and keeps no table. The spec forbids + * a consumer that retains an opaque {@code google.protobuf.Any} payload from dropping the entries + * that payload's ordinals refer to, so a converter must not retain such a payload uninterpreted + * when it holds ordinals. + */ +public final class DetachedSubtrees { + private static final DetachedSubtrees NONE = new DetachedSubtrees(Collections.emptyList(), false); + + private final List entries; + private final boolean inContainer; + private final BitSet referenced = new BitSet(); + private final BitSet resolving = new BitSet(); + + private DetachedSubtrees(List entries, boolean inContainer) { + this.entries = entries; + this.inContainer = inContainer; + } + + /** + * Returns the resolver for expressions converted outside any {@code PlanRel} or {@code + * ExtendedExpression}, where every detached expression ordinal is invalid. + * + * @return a resolver that rejects every ordinal + */ + public static DetachedSubtrees none() { + return NONE; + } + + /** + * Creates a resolver over the detached subtrees of a {@code PlanRel}. + * + * @param planRel the container whose ordinals are resolved + * @return a resolver for that container + */ + public static DetachedSubtrees of(PlanRel planRel) { + return new DetachedSubtrees(planRel.getDetachedExpressionsList(), true); + } + + /** + * Creates a resolver over the detached subtrees of an {@code ExtendedExpression}. + * + * @param extendedExpression the container whose ordinals are resolved + * @return a resolver for that container + */ + public static DetachedSubtrees of(ExtendedExpression extendedExpression) { + return new DetachedSubtrees(extendedExpression.getDetachedExpressionsList(), true); + } + + /** + * Returns the {@code detached_expressions} entries a retained opaque payload may still refer to: + * every entry up to the last one this resolver has not resolved. A resolved entry is referenced + * only from the tree that inlined it, so dropping the trailing run of resolved entries removes no + * payload's target and renumbers nothing. + * + * @return the leading entries to retain, indexed by ordinal + */ + public List unresolvedPrefix() { + int end = entries.size(); + while (end > 0 && referenced.get(end - 1)) { + end--; + } + return entries.subList(0, end); + } + + /** + * Converts the {@code detached_expressions} entry an ordinal refers to. + * + * @param ordinal the {@code detached_expression_ordinal}, an unsigned 32-bit value + * @param converter converts the entry, resolving any ordinals within it through this resolver + * @param the converted type + * @return the converted entry + * @throws IllegalArgumentException if the reference breaks one of the spec's rules + */ + public T resolveExpression( + int ordinal, Function converter) { + long index = Integer.toUnsignedLong(ordinal); + if (!inContainer) { + throw new IllegalArgumentException( + String.format( + "Detached expression ordinal %d is outside a PlanRel or ExtendedExpression", index)); + } + if (index >= entries.size()) { + throw new IllegalArgumentException( + String.format( + "Detached expression ordinal %d is out of range for %d detached expressions", + index, entries.size())); + } + int i = (int) index; + if (resolving.get(i)) { + throw new IllegalArgumentException( + String.format("Detached expression %d references itself through a cycle", i)); + } + if (referenced.get(i)) { + throw new IllegalArgumentException( + String.format("Detached expression %d is referenced more than once", i)); + } + io.substrait.proto.Expression entry = entries.get(i); + if (entry.getRexTypeCase() == RexTypeCase.DETACHED_EXPRESSION_ORDINAL) { + throw new IllegalArgumentException( + String.format("Detached expression %d is itself a detached expression ordinal", i)); + } + referenced.set(i); + resolving.set(i); + try { + return converter.apply(entry); + } finally { + resolving.clear(i); + } + } +} diff --git a/core/src/main/java/io/substrait/expression/proto/NestingDetacher.java b/core/src/main/java/io/substrait/expression/proto/NestingDetacher.java new file mode 100644 index 000000000..39a2c1958 --- /dev/null +++ b/core/src/main/java/io/substrait/expression/proto/NestingDetacher.java @@ -0,0 +1,306 @@ +package io.substrait.expression.proto; + +import com.google.protobuf.CodedInputStream; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Descriptors.FieldDescriptor; +import com.google.protobuf.Message; +import com.google.protobuf.WireFormat; +import io.substrait.proto.Expression.RexTypeCase; +import io.substrait.proto.ExtendedExpression; +import io.substrait.proto.PlanRel; +import java.io.IOException; +import java.util.ArrayList; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; + +/** + * Bounds the protobuf nesting depth of a {@code PlanRel} or {@code ExtendedExpression} by moving + * expression subtrees into its {@code detached_expressions} and referring to them by {@code + * Expression.detached_expression_ordinal} (spec v0.105.0). + * + *

Protobuf parsers limit how deeply messages may nest; protobuf-java, for one, rejects a message + * nested more than {@value #DEFAULT_MAX_DEPTH} levels below the root it parses. A container already + * within the limit is returned unchanged. Otherwise, an expression that does not fit at its + * position is detached once it lies deeper than half the limit, and its subtree is processed again + * from its new position directly below the container. If that leaves the container too deep, for + * instance because the deep part is a chain of relations below a shallow subquery, detachment is + * retried with every expression below the container eligible. Ordinals are assigned in the order + * the subtrees are detached, so a detached entry refers only to entries after it. + * + *

Detachment shortens expression nesting only, not nesting within a chain of relations, a + * literal or type, or a {@code google.protobuf.Any} payload. A container that still exceeds the + * limit after both attempts is returned unchanged, so that it does not need a consumer that reads + * detached expressions without gaining one that parses it. + */ +public final class NestingDetacher { + /** The default parse recursion limit of protobuf-java, and of the C++ and upb runtimes. */ + public static final int DEFAULT_MAX_DEPTH = 100; + + private static final NestingDetacher DEFAULT = new NestingDetacher(DEFAULT_MAX_DEPTH, false); + private static final NestingDetacher DISABLED = new NestingDetacher(Integer.MAX_VALUE, false); + private static final NestingDetacher EVERY_EXPRESSION = new NestingDetacher(0, true); + + /** Nesting depth of a {@code PlanRel} below the {@code Plan} it is parsed as part of. */ + private static final int PLAN_REL_DEPTH = 1; + + /** The message-typed fields of each message type, the only ones that add nesting. */ + private static final Map> MESSAGE_FIELDS = + new ConcurrentHashMap<>(); + + private final int maxDepth; + private final boolean everyExpression; + + private NestingDetacher(int maxDepth, boolean everyExpression) { + this.maxDepth = maxDepth; + this.everyExpression = everyExpression; + } + + /** + * Returns a detacher that keeps a plan within the default parse recursion limit of {@value + * #DEFAULT_MAX_DEPTH}. + * + * @return the default detacher + */ + public static NestingDetacher defaultDetacher() { + return DEFAULT; + } + + /** + * Returns a detacher that never detaches an expression. + * + * @return a detacher that leaves every container unchanged + */ + public static NestingDetacher disabled() { + return DISABLED; + } + + /** + * Returns a detacher that detaches every expression below the container, however shallow, so that + * a consumer's handling of detached expressions can be exercised on any plan. + * + * @return a detacher that detaches every expression it can + */ + public static NestingDetacher everyExpression() { + return EVERY_EXPRESSION; + } + + /** + * Returns a detacher for a consumer that parses up to {@code maxDepth} levels of nesting below + * the root message, the root itself being at depth 0. + * + * @param maxDepth the deepest nesting a consumer accepts + * @return a detacher for that limit + */ + public static NestingDetacher withMaxDepth(int maxDepth) { + if (maxDepth < 0) { + throw new IllegalArgumentException("maxDepth must not be negative: " + maxDepth); + } + return new NestingDetacher(maxDepth, false); + } + + /** + * Detaches expressions of a {@code PlanRel} that is parsed as part of a {@code Plan}. + * + * @param planRel the relation to bound + * @return {@code planRel} itself if it fits, otherwise a copy with detached expressions + */ + public PlanRel detach(PlanRel planRel) { + return bound(planRel, PLAN_REL_DEPTH, planRel.getDetachedExpressionsList()); + } + + /** + * Detaches expressions of an {@code ExtendedExpression} that is parsed as the root message. + * + * @param extendedExpression the extended expression to bound + * @return {@code extendedExpression} itself if it fits, otherwise a copy with detached + * expressions + */ + public ExtendedExpression detach(ExtendedExpression extendedExpression) { + return bound(extendedExpression, 0, extendedExpression.getDetachedExpressionsList()); + } + + private M bound( + M container, int containerDepth, List existing) { + int entryDepth = containerDepth + 1; + if (everyExpression) { + return new Pass(containerDepth, entryDepth, existing).run(container); + } + if (maxDepth == Integer.MAX_VALUE || fitsWhole(container, containerDepth)) { + return container; + } + for (int threshold : new int[] {Math.max(maxDepth / 2, entryDepth), entryDepth}) { + M detached = new Pass(containerDepth, threshold, existing).run(container); + if (fitsWhole(detached, containerDepth)) { + return detached; + } + } + return container; + } + + private boolean fitsWhole(Message container, int containerDepth) { + return maxDepth >= containerDepth && within(container, maxDepth - containerDepth); + } + + /** One detachment of one container. */ + private final class Pass { + private final int containerDepth; + private final int threshold; + private final int existingEntries; + private final List entries; + private final Map heights = new IdentityHashMap<>(); + + /** + * @param threshold the depth an expression must lie below to be detached + */ + Pass(int containerDepth, int threshold, List existing) { + this.containerDepth = containerDepth; + this.threshold = threshold; + this.existingEntries = existing.size(); + this.entries = new ArrayList<>(existing); + } + + @SuppressWarnings("unchecked") + M run(M container) { + FieldDescriptor detachedField = + container.getDescriptorForType().findFieldByName("detached_expressions"); + Message.Builder builder = container.toBuilder().clearField(detachedField); + // An existing entry is not moved, so the ordinals referring to it keep their target. + for (int i = 0; i < existingEntries; i++) { + entries.set(i, (io.substrait.proto.Expression) children(entries.get(i), entryDepth())); + } + for (FieldDescriptor field : messageFields(container)) { + if (!field.equals(detachedField)) { + copyRewritten(container, builder, field, containerDepth + 1); + } + } + for (io.substrait.proto.Expression entry : entries) { + builder.addRepeatedField(detachedField, entry); + } + return (M) builder.build(); + } + + private int entryDepth() { + return containerDepth + 1; + } + + private boolean fits(Message message, int depth) { + return height(message) <= maxDepth - depth; + } + + /** Returns the message at {@code depth}, with whatever it contains detached as needed. */ + private Message rewrite(Message message, int depth) { + if (fits(message, depth)) { + return message; + } + if (message instanceof io.substrait.proto.Expression + && depth > threshold + && ((io.substrait.proto.Expression) message).getRexTypeCase() + != RexTypeCase.DETACHED_EXPRESSION_ORDINAL) { + int ordinal = entries.size(); + entries.add(null); + entries.set(ordinal, (io.substrait.proto.Expression) children(message, entryDepth())); + return io.substrait.proto.Expression.newBuilder() + .setDetachedExpressionOrdinal(ordinal) + .build(); + } + return children(message, depth); + } + + /** Returns the message at {@code depth} with each child message rewritten. */ + private Message children(Message message, int depth) { + if (fits(message, depth)) { + return message; + } + Message.Builder builder = message.toBuilder(); + for (FieldDescriptor field : messageFields(message)) { + copyRewritten(message, builder, field, depth + 1); + } + return builder.build(); + } + + private void copyRewritten( + Message message, Message.Builder builder, FieldDescriptor field, int childDepth) { + if (field.isRepeated()) { + for (int i = 0; i < message.getRepeatedFieldCount(field); i++) { + builder.setRepeatedField( + field, i, rewrite((Message) message.getRepeatedField(field, i), childDepth)); + } + } else if (message.hasField(field)) { + builder.setField(field, rewrite((Message) message.getField(field), childDepth)); + } + } + + /** Returns how many levels of messages are nested below {@code message}. */ + private int height(Message message) { + Integer known = heights.get(message); + if (known != null) { + return known; + } + int height = 0; + for (FieldDescriptor field : messageFields(message)) { + if (field.isRepeated()) { + for (int i = 0; i < message.getRepeatedFieldCount(field); i++) { + height = Math.max(height, 1 + height((Message) message.getRepeatedField(field, i))); + } + } else if (message.hasField(field)) { + height = Math.max(height, 1 + height((Message) message.getField(field))); + } + } + heights.put(message, height); + return height; + } + } + + /** + * Returns whether at most {@code levels} levels of messages are nested below {@code message}. + * + *

Every container goes through this check, so it scans the serialized form, which the + * generated code writes without reflection, rather than walking the message tree. + */ + private static boolean within(Message message, int levels) { + try { + return within( + message.getDescriptorForType(), + CodedInputStream.newInstance(message.toByteArray()), + levels); + } catch (IOException e) { + throw new IllegalStateException( + "Failed to scan a serialized " + message.getDescriptorForType().getFullName(), e); + } + } + + private static boolean within(Descriptor type, CodedInputStream input, int levels) + throws IOException { + while (!input.isAtEnd()) { + int tag = input.readTag(); + if (WireFormat.getTagWireType(tag) == WireFormat.WIRETYPE_LENGTH_DELIMITED) { + FieldDescriptor field = type.findFieldByNumber(WireFormat.getTagFieldNumber(tag)); + if (field != null && field.getJavaType() == FieldDescriptor.JavaType.MESSAGE) { + if (levels == 0) { + return false; + } + int limit = input.pushLimit(input.readRawVarint32()); + if (!within(field.getMessageType(), input, levels - 1)) { + return false; + } + input.popLimit(limit); + continue; + } + } + input.skipField(tag); + } + return true; + } + + private static List messageFields(Message message) { + return MESSAGE_FIELDS.computeIfAbsent( + message.getDescriptorForType(), + descriptor -> + descriptor.getFields().stream() + .filter(field -> field.getJavaType() == FieldDescriptor.JavaType.MESSAGE) + .collect(Collectors.toList())); + } +} diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java index 847ab34b8..ce4d1f589 100644 --- a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java +++ b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java @@ -41,9 +41,16 @@ public class ProtoExpressionConverter { private final ProtoRelConverter protoRelConverter; private final LambdaBuilder lambdaBuilder = new LambdaBuilder(); + /** The container's detached subtrees, or {@code null} to ask {@link #protoRelConverter}. */ + private final DetachedSubtrees detachedSubtrees; + /** * Creates a converter resolving functions and types against the given extensions. * + *

Detached ordinals resolve against the {@link ProtoRelConverter#detachedSubtrees()} that + * {@code relConverter} holds when an ordinal is reached, and are rejected when it is {@code + * null}. + * * @param lookup used to resolve function and type references * @param extensions the extension collection providing definitions * @param rootType the root struct type that field references are resolved against @@ -54,11 +61,42 @@ public ProtoExpressionConverter( SimpleExtension.ExtensionCollection extensions, Type.Struct rootType, ProtoRelConverter relConverter) { + this(lookup, extensions, rootType, relConverter, null); + } + + /** + * Creates a converter resolving functions and types against the given extensions, and detached + * expression ordinals against the given container. + * + * @param lookup used to resolve function and type references + * @param extensions the extension collection providing definitions + * @param rootType the root struct type that field references are resolved against + * @param relConverter converter for nested relations (e.g. subqueries) + * @param detachedSubtrees the detached subtrees of the containing {@code PlanRel} or {@code + * ExtendedExpression}, or {@code null} to resolve against those {@code relConverter} holds + * when an ordinal is reached + */ + public ProtoExpressionConverter( + ExtensionLookup lookup, + SimpleExtension.ExtensionCollection extensions, + Type.Struct rootType, + ProtoRelConverter relConverter, + DetachedSubtrees detachedSubtrees) { this.lookup = lookup; this.extensions = extensions; this.rootType = Objects.requireNonNull(rootType, "rootType"); this.protoTypeConverter = new ProtoTypeConverter(lookup, extensions); this.protoRelConverter = relConverter; + this.detachedSubtrees = detachedSubtrees; + } + + private DetachedSubtrees detachedSubtrees() { + if (detachedSubtrees != null) { + return detachedSubtrees; + } + return protoRelConverter == null + ? DetachedSubtrees.none() + : protoRelConverter.detachedSubtrees(); } /** @@ -364,6 +402,9 @@ public Type visit(Type.Struct type) throws RuntimeException { + ecv.getExecutionContextVariableTypeCase()); } } + case DETACHED_EXPRESSION_ORDINAL: + return detachedSubtrees() + .resolveExpression(expr.getDetachedExpressionOrdinal(), this::from); default: throw new IllegalArgumentException("Unknown type: " + expr.getRexTypeCase()); } diff --git a/core/src/main/java/io/substrait/extendedexpression/ExtendedExpression.java b/core/src/main/java/io/substrait/extendedexpression/ExtendedExpression.java index 3ea90c6df..daa777514 100644 --- a/core/src/main/java/io/substrait/extendedexpression/ExtendedExpression.java +++ b/core/src/main/java/io/substrait/extendedexpression/ExtendedExpression.java @@ -44,6 +44,19 @@ public abstract class ExtendedExpression { */ public abstract Optional getAdvancedExtension(); + /** + * Returns the {@code detached_expressions} entries that ordinals inside the {@link + * #getAdvancedExtension() advanced extension} may still refer to, at their original ordinals. + * + *

The advanced extension is kept as the raw proto, so an ordinal inside it is not resolved; + * the spec requires its target to keep its ordinal. Reading an extended expression that has an + * advanced extension retains every entry up to the last one no referred expression resolved, and + * writing one places them back at the same ordinals. Empty otherwise. + * + * @return the retained detached expressions, indexed by ordinal + */ + public abstract List getRetainedDetachedExpressions(); + /** * Creates a builder for {@link ExtendedExpression}. * diff --git a/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java b/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java index 45d6f8e68..5a9020e1a 100644 --- a/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java +++ b/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java @@ -1,6 +1,7 @@ package io.substrait.extendedexpression; import io.substrait.expression.proto.ExpressionProtoConverter; +import io.substrait.expression.proto.NestingDetacher; import io.substrait.extension.ExtensionCollector; import io.substrait.proto.ExpressionReference; import io.substrait.proto.ExtendedExpression; @@ -13,6 +14,29 @@ * io.substrait.proto.ExtendedExpression} */ public class ExtendedExpressionProtoConverter { + private final NestingDetacher nestingDetacher; + + /** + * Creates a converter that uses {@link NestingDetacher#defaultDetacher()}, which leaves an + * extended expression unchanged unless its expressions nest too deeply for protobuf parsers with + * default limits. + */ + public ExtendedExpressionProtoConverter() { + this(NestingDetacher.defaultDetacher()); + } + + /** + * Creates a converter that bounds the nesting depth of its output with the given detacher. + * + * @param nestingDetacher bounds the nesting depth of each {@link ExtendedExpression}, must not be + * null + */ + public ExtendedExpressionProtoConverter(NestingDetacher nestingDetacher) { + if (nestingDetacher == null) { + throw new IllegalArgumentException("NestingDetacher is required"); + } + this.nestingDetacher = nestingDetacher; + } /** * Converts an {@link io.substrait.extendedexpression.ExtendedExpression} into its protobuf @@ -71,6 +95,8 @@ public ExtendedExpression toProto( if (extendedExpression.getAdvancedExtension().isPresent()) { builder.setAdvancedExtensions(extendedExpression.getAdvancedExtension().get()); } - return builder.build(); + // Ahead of any newly detached entry, so the advanced extension's ordinals keep their targets. + builder.addAllDetachedExpressions(extendedExpression.getRetainedDetachedExpressions()); + return nestingDetacher.detach(builder.build()); } } diff --git a/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java b/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java index 55ec5e6b0..d98308208 100644 --- a/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java +++ b/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java @@ -1,6 +1,7 @@ package io.substrait.extendedexpression; import io.substrait.expression.Expression; +import io.substrait.expression.proto.DetachedSubtrees; import io.substrait.expression.proto.ProtoExpressionConverter; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.ExtensionCollector; @@ -59,9 +60,10 @@ public ExtendedExpression from(io.substrait.proto.ExtendedExpression extendedExp io.substrait.type.NamedStruct namedStruct = io.substrait.type.NamedStruct.fromProto(baseSchemaProto, protoTypeConverter); + DetachedSubtrees detachedSubtrees = DetachedSubtrees.of(extendedExpression); ProtoExpressionConverter protoExpressionConverter = new ProtoExpressionConverter( - functionLookup, this.extensionCollection, namedStruct.struct(), null); + functionLookup, this.extensionCollection, namedStruct.struct(), null, detachedSubtrees); List expressionReferences = new ArrayList<>(); @@ -111,6 +113,10 @@ public ExtendedExpression from(io.substrait.proto.ExtendedExpression extendedExp ? extendedExpression.getAdvancedExtensions() : null)) .baseSchema(namedStruct); + if (extendedExpression.hasAdvancedExtensions()) { + // The advanced extension stays raw, so the entries its ordinals may name keep their ordinals. + builder.retainedDetachedExpressions(detachedSubtrees.unresolvedPrefix()); + } return builder.build(); } diff --git a/core/src/main/java/io/substrait/plan/PlanProtoConverter.java b/core/src/main/java/io/substrait/plan/PlanProtoConverter.java index 61e9ff9b3..b94e8eeac 100644 --- a/core/src/main/java/io/substrait/plan/PlanProtoConverter.java +++ b/core/src/main/java/io/substrait/plan/PlanProtoConverter.java @@ -1,5 +1,6 @@ package io.substrait.plan; +import io.substrait.expression.proto.NestingDetacher; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.ExtensionCollector; import io.substrait.extension.ExtensionProtoConverter; @@ -21,6 +22,11 @@ public class PlanProtoConverter { @NonNull private final ExtensionCollection extensionCollection; + @NonNull private final NestingDetacher nestingDetacher; + + /** The spec release that introduced detached expressions, as major, minor and patch. */ + private static final int[] DETACHED_EXPRESSIONS_SINCE = {0, 105, 0}; + /** Default constructor. */ public PlanProtoConverter() { this(DefaultExtensionCatalog.DEFAULT_COLLECTION); @@ -53,6 +59,27 @@ public PlanProtoConverter(@NonNull final ExtensionProtoConverter extension public PlanProtoConverter( @NonNull final ExtensionCollection extensionCollection, @NonNull final ExtensionProtoConverter extensionProtoConverter) { + this(extensionCollection, extensionProtoConverter, NestingDetacher.defaultDetacher()); + } + + /** + * Constructor with custom {@link ExtensionCollection}, {@link ExtensionProtoConverter} and {@link + * NestingDetacher}. + * + *

The other constructors use {@link NestingDetacher#defaultDetacher()}, which leaves a plan + * unchanged unless its expressions nest too deeply for protobuf parsers with default limits. + * + * @param extensionCollection custom {@link ExtensionCollection} to use, must not be null + * @param extensionProtoConverter custom {@link ExtensionProtoConverter} to use, must not be null + * @param nestingDetacher bounds the nesting depth of each {@link PlanRel}, must not be null + */ + public PlanProtoConverter( + @NonNull final ExtensionCollection extensionCollection, + @NonNull final ExtensionProtoConverter extensionProtoConverter, + @NonNull final NestingDetacher nestingDetacher) { + if (nestingDetacher == null) { + throw new IllegalArgumentException("NestingDetacher is required"); + } if (extensionCollection == null) { throw new IllegalArgumentException("ExtensionCollection is required"); } @@ -61,6 +88,7 @@ public PlanProtoConverter( } this.extensionCollection = extensionCollection; this.extensionProtoConverter = extensionProtoConverter; + this.nestingDetacher = nestingDetacher; } /** @@ -72,18 +100,22 @@ public PlanProtoConverter( */ public Plan toProto(final io.substrait.plan.Plan plan) { final List planRels = new ArrayList<>(); + boolean detached = false; final ExtensionCollector functionCollector = new ExtensionCollector(extensionCollection); for (final io.substrait.plan.Plan.Root root : plan.getRoots()) { final Rel input = new RelProtoConverter(functionCollector, extensionProtoConverter) .toProto(root.getInput()); - planRels.add( + final PlanRel planRel = PlanRel.newBuilder() .setRoot( io.substrait.proto.RelRoot.newBuilder() .setInput(input) .addAllNames(root.getNames())) - .build()); + .build(); + final PlanRel bounded = nestingDetacher.detach(planRel); + detached |= bounded != planRel; + planRels.add(bounded); } final Plan.Builder builder = Plan.newBuilder() @@ -103,6 +135,14 @@ public Plan toProto(final io.substrait.plan.Plan plan) { plan.getVersion().getGitHash().ifPresent(gh -> versionBuilder.setGitHash(gh)); plan.getVersion().getProducer().ifPresent(p -> versionBuilder.setProducer(p)); + if (detached && precedesDetachedExpressions(versionBuilder)) { + // The git hash identifies the earlier release, so it no longer describes the plan. + versionBuilder + .setMajorNumber(DETACHED_EXPRESSIONS_SINCE[0]) + .setMinorNumber(DETACHED_EXPRESSIONS_SINCE[1]) + .setPatchNumber(DETACHED_EXPRESSIONS_SINCE[2]) + .clearGitHash(); + } builder.setVersion(versionBuilder); @@ -112,6 +152,22 @@ public Plan toProto(final io.substrait.plan.Plan plan) { return builder.build(); } + /** + * Whether a plan declaring {@code version} predates detached expressions, which a {@link PlanRel} + * rewritten by the {@link NestingDetacher} uses, so the declared version has to be raised. + */ + private static boolean precedesDetachedExpressions(final Version.Builder version) { + final int[] declared = { + version.getMajorNumber(), version.getMinorNumber(), version.getPatchNumber() + }; + for (int i = 0; i < declared.length; i++) { + if (declared[i] != DETACHED_EXPRESSIONS_SINCE[i]) { + return declared[i] < DETACHED_EXPRESSIONS_SINCE[i]; + } + } + return false; + } + /** * Converts an {@link io.substrait.plan.Plan.ExecutionBehavior} to its protobuf representation. * diff --git a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java index 629665f2e..a39ba9d65 100644 --- a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java +++ b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java @@ -9,7 +9,6 @@ import io.substrait.plan.Plan.ExecutionBehavior.VariableEvaluationMode; import io.substrait.proto.PlanRel; import io.substrait.relation.ProtoRelConverter; -import io.substrait.relation.Rel; import java.util.ArrayList; import java.util.List; import java.util.Optional; @@ -95,9 +94,7 @@ public Plan from(io.substrait.proto.Plan plan) { ProtoRelConverter relConverter = getProtoRelConverter(functionLookup); List roots = new ArrayList<>(); for (PlanRel planRel : plan.getRelationsList()) { - io.substrait.proto.RelRoot root = planRel.getRoot(); - Rel rel = relConverter.from(root.getInput()); - roots.add(Plan.Root.builder().input(rel).names(root.getNamesList()).build()); + roots.add(relConverter.from(planRel)); } ImmutableVersion.Builder versionBuilder = diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java index 83b3d9e38..04a454919 100644 --- a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java +++ b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java @@ -3,6 +3,7 @@ import io.substrait.expression.Expression; import io.substrait.expression.FieldReference; import io.substrait.expression.MaskExpression; +import io.substrait.expression.proto.DetachedSubtrees; import io.substrait.expression.proto.ProtoExpressionConverter; import io.substrait.expression.proto.ProtoMaskExpressionConverter; import io.substrait.extension.AdvancedExtension; @@ -102,6 +103,9 @@ public class ProtoRelConverter { */ private final Map anchorScopes = new HashMap<>(); + /** The detached subtrees of the {@code PlanRel} being converted, if any. */ + private DetachedSubtrees detachedSubtrees = DetachedSubtrees.none(); + /** * Constructor with custom {@link ExtensionLookup}. * @@ -188,6 +192,73 @@ public Rel from(io.substrait.proto.Rel rel) { return converted; } + /** + * Converts the root of a {@code PlanRel}, resolving the detached ordinals in it against that + * {@code PlanRel}'s detached subtrees. + * + * @param planRel the protobuf {@code PlanRel} to convert, which must hold a root + * @return the converted root + * @throws IllegalArgumentException if {@code planRel} holds a bare relation rather than a root + */ + public Plan.Root from(io.substrait.proto.PlanRel planRel) { + if (!planRel.hasRoot()) { + throw new IllegalArgumentException( + "Only PlanRel roots are supported, got " + planRel.getRelTypeCase()); + } + DetachedSubtrees enclosing = enter(planRel); + try { + return from(planRel.getRoot()); + } finally { + this.detachedSubtrees = enclosing; + } + } + + /** + * Converts the bare relation of a {@code PlanRel}, resolving the detached ordinals in it against + * that {@code PlanRel}'s detached subtrees. + * + * @param planRel the protobuf {@code PlanRel} to convert, which must hold a bare relation + * @return the converted relation + * @throws IllegalArgumentException if {@code planRel} holds a root rather than a bare relation + */ + public Rel fromRel(io.substrait.proto.PlanRel planRel) { + if (!planRel.hasRel()) { + throw new IllegalArgumentException( + "Expected a PlanRel holding a bare relation, got " + planRel.getRelTypeCase()); + } + DetachedSubtrees enclosing = enter(planRel); + try { + return from(planRel.getRel()); + } finally { + this.detachedSubtrees = enclosing; + } + } + + /** + * Resolves ordinals against {@code planRel} from now on, returning what they resolved against. + */ + private DetachedSubtrees enter(io.substrait.proto.PlanRel planRel) { + DetachedSubtrees enclosing = this.detachedSubtrees; + this.detachedSubtrees = DetachedSubtrees.of(planRel); + return enclosing; + } + + /** + * Returns the detached subtrees that expressions converted for the current relation resolve their + * ordinals against: those of the {@code PlanRel} being converted by {@link + * #from(io.substrait.proto.PlanRel)} or {@link #fromRel(io.substrait.proto.PlanRel)}, or a + * resolver that rejects every ordinal outside such a conversion. + * + *

The resolver records each entry it resolves so that it can reject a second reference to it, + * so it is meant for the expression converters of the current conversion only: resolving an + * ordinal through it from elsewhere makes the conversion's own reference to that entry fail. + * + * @return the current detached subtrees + */ + public DetachedSubtrees detachedSubtrees() { + return detachedSubtrees; + } + private Rel dispatch(io.substrait.proto.Rel rel) { io.substrait.proto.Rel.RelTypeCase relType = rel.getRelTypeCase(); switch (relType) { diff --git a/core/src/test/java/io/substrait/type/proto/DetachedExpressionRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/DetachedExpressionRoundtripTest.java new file mode 100644 index 000000000..b05059d41 --- /dev/null +++ b/core/src/test/java/io/substrait/type/proto/DetachedExpressionRoundtripTest.java @@ -0,0 +1,550 @@ +package io.substrait.type.proto; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.protobuf.Descriptors.FieldDescriptor; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.Message; +import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; +import io.substrait.expression.FieldReference; +import io.substrait.expression.LambdaBuilder; +import io.substrait.expression.proto.NestingDetacher; +import io.substrait.expression.proto.ProtoExpressionConverter; +import io.substrait.extendedexpression.ExtendedExpression; +import io.substrait.extendedexpression.ExtendedExpressionProtoConverter; +import io.substrait.extendedexpression.ImmutableExpressionReference; +import io.substrait.extendedexpression.ImmutableExtendedExpression; +import io.substrait.extendedexpression.ProtoExtendedExpressionConverter; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.ExtensionProtoConverter; +import io.substrait.plan.Plan; +import io.substrait.plan.PlanProtoConverter; +import io.substrait.plan.ProtoPlanConverter; +import io.substrait.proto.FilterRel; +import io.substrait.proto.PlanRel; +import io.substrait.relation.Filter; +import io.substrait.relation.NamedScan; +import io.substrait.relation.ProtoRelConverter; +import io.substrait.relation.Rel; +import io.substrait.relation.Rel.Remap; +import io.substrait.type.NamedStruct; +import io.substrait.utils.RelSamples; +import io.substrait.utils.StringHolderHandlingProtoRelConverter; +import io.substrait.utils.StringHolderHandlingRelProtoConverter; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Verifies detached expression subtrees (spec v0.105.0): the writer moves expressions into {@code + * detached_expressions} only when they nest too deeply, the reader resolves the ordinals back into + * the same expressions, and the reader enforces the rules the spec places on the references. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class DetachedExpressionRoundtripTest extends StringHolderRoundtripTestBase { + + static final NestingDetacher DETACH_ALL = NestingDetacher.everyExpression(); + + /** Deep enough that the nested additions exceed protobuf-java's parse recursion limit. */ + static final int DEEP = 200; + + final Map, Rel> relSamples = new RelSamples(sb, extensions).samples(); + + final NamedScan scan = sb.namedScan(List.of("t"), List.of("a"), List.of(R.I32)); + + final Rel orders = + sb.namedScan(List.of("orders"), List.of("o_orderkey", "o_custkey"), List.of(R.I64, R.I64)); + + final Rel customer = + sb.namedScan(List.of("customer"), List.of("c_custkey", "c_nationkey"), List.of(R.I64, R.I64)); + + Stream samples() { + return relSamples.entrySet().stream() + .map(sample -> Arguments.of(sample.getKey().getSimpleName(), sample.getValue())); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("samples") + void everyRelationRoundtripsWithEveryExpressionDetached(String relationType, Rel rel) { + io.substrait.proto.Rel protoRel = + new StringHolderHandlingRelProtoConverter(functionCollector).toProto(rel); + PlanRel detached = + DETACH_ALL.detach( + PlanRel.newBuilder() + .setRoot(io.substrait.proto.RelRoot.newBuilder().setInput(protoRel)) + .build()); + + if (containsExpression(protoRel)) { + assertTrue(detached.getDetachedExpressionsCount() > 0, "no expression was detached"); + assertNotEquals(protoRel, detached.getRoot().getInput()); + } + Rel returned = + new StringHolderHandlingProtoRelConverter(functionCollector, extensions) + .from(detached) + .getInput(); + assertEquals(rel, returned); + } + + @Test + void aShallowPlanIsLeftUnchanged() { + Plan plan = sb.plan(sb.root(sb.project(input -> List.of(nestedAdditions(input, 3)), scan))); + io.substrait.proto.Plan proto = new PlanProtoConverter().toProto(plan); + assertEquals(0, proto.getRelations(0).getDetachedExpressionsCount()); + assertEquals(proto, planConverter(NestingDetacher.disabled()).toProto(plan)); + } + + @Test + void aDeepPlanIsDetachedToParseWithDefaultLimits() throws InvalidProtocolBufferException { + Plan plan = sb.plan(sb.root(sb.project(input -> List.of(nestedAdditions(input, DEEP)), scan))); + + byte[] undetached = planConverter(NestingDetacher.disabled()).toProto(plan).toByteArray(); + assertThrows( + InvalidProtocolBufferException.class, () -> io.substrait.proto.Plan.parseFrom(undetached)); + + io.substrait.proto.Plan proto = + io.substrait.proto.Plan.parseFrom(new PlanProtoConverter().toProto(plan).toByteArray()); + assertTrue(proto.getRelations(0).getDetachedExpressionsCount() > 0); + assertTrue(depth(proto) <= NestingDetacher.DEFAULT_MAX_DEPTH, () -> "depth " + depth(proto)); + assertEquals(plan, new ProtoPlanConverter().from(proto)); + } + + @Test + void aDeepExtendedExpressionIsDetachedToParseWithDefaultLimits() + throws InvalidProtocolBufferException { + ExtendedExpression extendedExpression = + ImmutableExtendedExpression.builder() + .baseSchema(NamedStruct.of(List.of("a"), R.struct(R.I32))) + .addReferredExpressions( + ImmutableExpressionReference.builder() + .expression(nestedAdditions(scan, DEEP)) + .addOutputNames("deep") + .build()) + .build(); + + io.substrait.proto.ExtendedExpression proto = + io.substrait.proto.ExtendedExpression.parseFrom( + new ExtendedExpressionProtoConverter().toProto(extendedExpression).toByteArray()); + assertTrue(proto.getDetachedExpressionsCount() > 0); + assertTrue(depth(proto) <= NestingDetacher.DEFAULT_MAX_DEPTH, () -> "depth " + depth(proto)); + assertEquals(extendedExpression, new ProtoExtendedExpressionConverter().from(proto)); + + io.substrait.proto.ExtendedExpression all = + new ExtendedExpressionProtoConverter(DETACH_ALL).toProto(extendedExpression); + assertEquals(extendedExpression, new ProtoExtendedExpressionConverter().from(all)); + } + + @Test + void aContainerWithinTheLimitIsReturnedAsIs() { + PlanRel planRel = + PlanRel.newBuilder() + .setRoot( + io.substrait.proto.RelRoot.newBuilder() + .setInput(relProtoConverter.toProto(filterOn(sb.bool(true))))) + .build(); + assertSame(planRel, NestingDetacher.defaultDetacher().detach(planRel)); + } + + @Test + void anUnreferencedEntryIsNeitherConvertedNorValidated() { + PlanRel planRel = + DETACH_ALL.detach(planRel(filterOn(sb.equal(sb.fieldReference(scan, 0), sb.i32(1))))); + PlanRel withUnreachable = + planRel.toBuilder() + .addDetachedExpressions(ordinal(0)) + .addDetachedExpressions(io.substrait.proto.Expression.getDefaultInstance()) + .build(); + assertEquals(protoRelConverter.from(planRel), protoRelConverter.from(withUnreachable)); + } + + @Test + void anOrdinalOutOfRangeIsRejected() { + assertRejected("out of range", condition(ordinal(1)), List.of(literal())); + assertRejected("out of range", condition(ordinal(0)), List.of()); + } + + @Test + void anEntryReferencedTwiceIsRejected() { + assertRejected( + "referenced more than once", condition(equal(ordinal(0), ordinal(0))), List.of(literal())); + } + + @Test + void aCycleIsRejected() { + assertRejected( + "cycle", + condition(ordinal(0)), + List.of(equal(ordinal(1), literal()), equal(ordinal(0), literal()))); + } + + @Test + void anEntryThatIsItselfAnOrdinalIsRejected() { + assertRejected( + "itself a detached expression ordinal", + condition(ordinal(0)), + List.of(ordinal(1), literal())); + } + + @Test + void anOrdinalOutsideAContainerIsRejected() { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, () -> protoRelConverter.from(condition(ordinal(0)))); + assertTrue(error.getMessage().contains("outside a PlanRel"), error.getMessage()); + } + + @Test + void aShallowSubqueryOverALongChainIsDetachedByTheFallback() + throws InvalidProtocolBufferException { + Rel subquery = filters(scan, 28); + Rel project = + sb.project( + input -> List.of(sb.scalarSubquery(subquery, R.I32)), Remap.of(List.of(1)), scan); + Plan plan = sb.plan(sb.root(filters(project, 18))); + + io.substrait.proto.Plan undetached = planConverter(NestingDetacher.disabled()).toProto(plan); + assertTrue(depth(undetached) > NestingDetacher.DEFAULT_MAX_DEPTH); + io.substrait.proto.Plan proto = + io.substrait.proto.Plan.parseFrom(new PlanProtoConverter().toProto(plan).toByteArray()); + assertTrue(depth(proto) <= NestingDetacher.DEFAULT_MAX_DEPTH, () -> "depth " + depth(proto)); + assertEquals(plan, new ProtoPlanConverter().from(proto)); + } + + @Test + void aContainerThatCannotBeMadeToFitIsReturnedAsIs() { + PlanRel planRel = planRel(filters(scan, 60)); + assertTrue(depth(planRel) > NestingDetacher.DEFAULT_MAX_DEPTH); + assertSame(planRel, NestingDetacher.defaultDetacher().detach(planRel)); + } + + @Test + void detachingRaisesAnEarlierDeclaredVersion() { + Plan.Version earlier = + Plan.Version.builder() + .major(0) + .minor(90) + .patch(0) + .gitHash(Optional.of("0123abc")) + .producer(Optional.of("a producer")) + .build(); + Plan deep = + Plan.builder() + .from(sb.plan(sb.root(sb.project(in -> List.of(nestedAdditions(in, DEEP)), scan)))) + .version(earlier) + .build(); + io.substrait.proto.Version written = new PlanProtoConverter().toProto(deep).getVersion(); + assertEquals( + List.of(0, 105, 0), + List.of(written.getMajorNumber(), written.getMinorNumber(), written.getPatchNumber())); + assertEquals("", written.getGitHash()); + assertEquals("a producer", written.getProducer()); + + Plan shallow = + Plan.builder() + .from(sb.plan(sb.root(sb.project(in -> List.of(nestedAdditions(in, 3)), scan)))) + .version(earlier) + .build(); + io.substrait.proto.Version unchanged = new PlanProtoConverter().toProto(shallow).getVersion(); + assertEquals(90, unchanged.getMinorNumber()); + assertEquals("0123abc", unchanged.getGitHash()); + } + + @Test + void anOpaqueAdvancedExtensionKeepsItsDetachedTargets() { + io.substrait.proto.ExtendedExpression shallow = + new ExtendedExpressionProtoConverter() + .toProto( + ImmutableExtendedExpression.builder() + .baseSchema(NamedStruct.of(List.of("a"), R.struct(R.I32))) + .addReferredExpressions( + ImmutableExpressionReference.builder() + .expression(sb.equal(sb.fieldReference(scan, 0), sb.i32(1))) + .addOutputNames("eq") + .build()) + .build()); + io.substrait.proto.Expression payloadTarget = literal(); + // The payload names entry 0; the referred expression names entry 1, which is resolved. + io.substrait.proto.ExtendedExpression withPayload = + shallow.toBuilder() + .setReferredExpr(0, shallow.getReferredExpr(0).toBuilder().setExpression(ordinal(1))) + .addDetachedExpressions(payloadTarget) + .addDetachedExpressions(shallow.getReferredExpr(0).getExpression()) + .setAdvancedExtensions( + io.substrait.proto.AdvancedExtension.newBuilder() + .setEnhancement(com.google.protobuf.Any.pack(ordinal(0)))) + .build(); + + ExtendedExpression read = new ProtoExtendedExpressionConverter().from(withPayload); + assertEquals(List.of(payloadTarget), read.getRetainedDetachedExpressions()); + + io.substrait.proto.ExtendedExpression written = + new ExtendedExpressionProtoConverter().toProto(read); + assertEquals(List.of(payloadTarget), written.getDetachedExpressionsList()); + assertEquals(read, new ProtoExtendedExpressionConverter().from(written)); + + io.substrait.proto.ExtendedExpression allDetached = + new ExtendedExpressionProtoConverter(DETACH_ALL).toProto(read); + assertEquals(payloadTarget, allDetached.getDetachedExpressions(0)); + assertEquals(read, new ProtoExtendedExpressionConverter().from(allDetached)); + } + + @Test + void aBarePlanRelIsReadWithItsDetachedSubtrees() { + Rel rel = filterOn(sb.equal(sb.fieldReference(scan, 0), sb.i32(1))); + PlanRel detached = + DETACH_ALL.detach(PlanRel.newBuilder().setRel(relProtoConverter.toProto(rel)).build()); + assertTrue(detached.getDetachedExpressionsCount() > 0); + assertEquals(rel, protoRelConverter.fromRel(detached)); + assertThrows(IllegalArgumentException.class, () -> protoRelConverter.from(detached)); + } + + @Test + void anExpressionConverterBuiltBeforeTheConversionResolvesAgainstIt() { + ProtoRelConverter reader = + new ProtoRelConverter(functionCollector, extensions) { + // Built before any PlanRel is converted, as a subclass field would be. + final ProtoExpressionConverter early = + new ProtoExpressionConverter( + functionCollector, extensions, scan.getRecordType(), this); + + @Override + protected Filter newFilter(FilterRel rel) { + return Filter.builder() + .input(from(rel.getInput())) + .condition(early.from(rel.getCondition())) + .build(); + } + }; + Rel rel = filterOn(sb.equal(sb.fieldReference(scan, 0), sb.i32(1))); + assertEquals(rel, reader.from(DETACH_ALL.detach(planRel(rel))).getInput()); + } + + @Test + void aPlanRelWithoutARootIsRejected() { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> + protoRelConverter.from( + PlanRel.newBuilder().setRel(relProtoConverter.toProto(scan)).build())); + assertTrue(error.getMessage().contains("Only PlanRel roots"), error.getMessage()); + } + + @Test + void aLambdaBodyResolvesInTheLambdaScope() { + LambdaBuilder lb = new LambdaBuilder(); + Rel input = sb.namedScan(List.of("t"), List.of("xs", "k"), List.of(R.list(R.I64), R.I64)); + assertDetachedRoundTrip( + sb.project( + in -> + List.of( + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_LIST, + "transform:list_func", + R.list(R.I64), + sb.fieldReference(in, 0), + lb.lambda( + List.of(R.I64), + p -> sb.add(sb.add(p.ref(0), sb.fieldReference(in, 1)), sb.i64(1))))), + Remap.of(List.of(2)), + input)); + } + + @Test + void aNestedLambdaBodyResolvesInItsOwnScope() { + LambdaBuilder lb = new LambdaBuilder(); + Rel input = sb.namedScan(List.of("t"), List.of("xs"), List.of(R.list(R.list(R.I64)))); + assertDetachedRoundTrip( + sb.project( + in -> + List.of( + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_LIST, + "transform:list_func", + R.list(R.list(R.I64)), + sb.fieldReference(in, 0), + lb.lambda( + List.of(R.list(R.I64)), + outer -> + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_LIST, + "transform:list_func", + R.list(R.I64), + outer.ref(0), + lb.lambda( + List.of(R.I64), + inner -> sb.add(inner.ref(0), sb.i64(1))))))), + Remap.of(List.of(1)), + input)); + } + + @Test + void aStepsOutReferenceResolvesAgainstTheEnclosingQuery() { + Rel subquery = + correlated(FieldReference.newRootStructOuterReference(1, orders.getRecordType(), 1)); + assertDetachedRoundTrip( + sb.project( + input -> List.of(sb.fieldReference(input, 0), sb.scalarSubquery(subquery, R.I64)), + Remap.of(List.of(2, 3)), + orders)); + } + + @Test + void aRelReferenceResolvesAgainstTheAnchoredRelation() { + Rel anchored = orders.withRelAnchor(1); + Rel subquery = + correlated( + FieldReference.newRootStructOuterReferenceByRelReference( + 1, anchored.getRecordType(), 1)); + assertDetachedRoundTrip( + sb.project( + input -> List.of(sb.fieldReference(input, 0), sb.scalarSubquery(subquery, R.I64)), + Remap.of(List.of(2, 3)), + anchored)); + } + + @Test + void inPredicateNeedlesAndACorrelatedHaystackResolveInTheirScopes() { + Rel haystack = + sb.project( + input -> List.of(sb.fieldReference(input, 0)), + Remap.of(List.of(2)), + sb.filter( + input -> + sb.equal( + sb.fieldReference(input, 1), + FieldReference.newRootStructOuterReference(1, orders.getRecordType(), 1)), + customer)); + assertDetachedRoundTrip( + sb.filter( + input -> + sb.inPredicate( + haystack, + sb.add(sb.fieldReference(input, 0), sb.i64(1)), + sb.fieldReference(input, 1)), + orders)); + } + + /** Round-trips {@code rel} as a plan with every expression detached. */ + private void assertDetachedRoundTrip(Rel rel) { + Plan plan = sb.plan(sb.root(rel)); + io.substrait.proto.Plan proto = planConverter(DETACH_ALL).toProto(plan); + assertTrue(proto.getRelations(0).getDetachedExpressionsCount() > 0); + assertEquals(plan, new ProtoPlanConverter().from(proto)); + } + + /** A subquery over customer whose filter compares c_custkey with the given outer reference. */ + private Rel correlated(FieldReference outerReference) { + return sb.project( + input -> List.of(sb.fieldReference(input, 1)), + Remap.of(List.of(2)), + sb.filter(input -> sb.equal(sb.fieldReference(input, 0), outerReference), customer)); + } + + private void assertRejected( + String message, + io.substrait.proto.Rel rel, + List detachedExpressions) { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> + protoRelConverter.from( + PlanRel.newBuilder() + .setRoot(io.substrait.proto.RelRoot.newBuilder().setInput(rel)) + .addAllDetachedExpressions(detachedExpressions) + .build())); + assertTrue(error.getMessage().contains(message), error.getMessage()); + } + + /** {@code count} filters stacked on {@code input}. */ + private Rel filters(Rel input, int count) { + Rel rel = input; + for (int i = 0; i < count; i++) { + int value = i; + rel = sb.filter(in -> sb.equal(sb.fieldReference(in, 0), sb.i32(value)), rel); + } + return rel; + } + + private Rel filterOn(Expression condition) { + return sb.filter(input -> condition, scan); + } + + private PlanRel planRel(Rel rel) { + return PlanRel.newBuilder() + .setRoot(io.substrait.proto.RelRoot.newBuilder().setInput(relProtoConverter.toProto(rel))) + .build(); + } + + /** A filter over {@link #scan} with the given proto condition. */ + private io.substrait.proto.Rel condition(io.substrait.proto.Expression condition) { + io.substrait.proto.Rel rel = relProtoConverter.toProto(filterOn(sb.bool(true))); + return rel.toBuilder().setFilter(rel.getFilter().toBuilder().setCondition(condition)).build(); + } + + /** An {@code equal:any_any} of the two proto operands, shaped as the writer emits it. */ + private io.substrait.proto.Expression equal( + io.substrait.proto.Expression left, io.substrait.proto.Expression right) { + io.substrait.proto.Expression template = + expressionProtoConverter.toProto(sb.equal(sb.i32(1), sb.i32(1))); + io.substrait.proto.Expression.ScalarFunction.Builder function = + template.getScalarFunction().toBuilder(); + function.getArgumentsBuilder(0).setValue(left); + function.getArgumentsBuilder(1).setValue(right); + return template.toBuilder().setScalarFunction(function).build(); + } + + private io.substrait.proto.Expression literal() { + return expressionProtoConverter.toProto(sb.i32(1)); + } + + private static io.substrait.proto.Expression ordinal(int ordinal) { + return io.substrait.proto.Expression.newBuilder().setDetachedExpressionOrdinal(ordinal).build(); + } + + private Expression nestedAdditions(Rel input, int depth) { + Expression expression = sb.fieldReference(input, 0); + for (int i = 0; i < depth; i++) { + expression = sb.add(expression, ExpressionCreator.i32(false, i)); + } + return expression; + } + + private static PlanProtoConverter planConverter(NestingDetacher detacher) { + return new PlanProtoConverter( + DefaultExtensionCatalog.DEFAULT_COLLECTION, new ExtensionProtoConverter<>(), detacher); + } + + private static boolean containsExpression(Message message) { + if (message instanceof io.substrait.proto.Expression) { + return true; + } + return children(message).anyMatch(DetachedExpressionRoundtripTest::containsExpression); + } + + /** Returns how many levels of messages nest below {@code message}, as a parser counts them. */ + private static int depth(Message message) { + return children(message).mapToInt(child -> 1 + depth(child)).max().orElse(0); + } + + private static Stream children(Message message) { + return message.getAllFields().entrySet().stream() + .filter(field -> field.getKey().getJavaType() == FieldDescriptor.JavaType.MESSAGE) + .flatMap( + field -> + field.getKey().isRepeated() + ? ((List) field.getValue()).stream().map(Message.class::cast) + : Stream.of((Message) field.getValue())); + } +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 3c0fc4d57..a5d671457 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -31,7 +31,7 @@ spark-3-5 = "3.5.4" spark-4-0 = "4.0.2" spotless = "8.10.3" # substrait-packaging artifacts, versioned by the Substrait spec release they are generated from. -substrait-packaging = "0.104.0" +substrait-packaging = "0.105.0" testcontainers = "2.0.5" validator = "3.0.8" diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java index 98649e283..932a8006f 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java @@ -119,8 +119,8 @@ void sumAndAvgDeriveWhatTheExtensionsDeclare() { /** * The same in a query: isthmus's own SUM and AVG take their types from the type system, where - * they used to repeat the argument's. The average stays nullable, the right answer for an empty - * group, although avg:dec declares a required result. + * they used to repeat the argument's. The average is nullable, the right answer for an empty + * group, as avg:dec declares. */ @Test void sumAndAvgInAQueryTakeTheTypeSystemsTypes() throws Exception { diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala index dc0d09c4f..d50f1ca72 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala @@ -45,6 +45,7 @@ import io.substrait.{proto, relation} import io.substrait.debug.TreePrinter import io.substrait.expression.{Expression => SExpression, ExpressionCreator} import io.substrait.expression.Expression.StructLiteral +import io.substrait.expression.proto.NestingDetacher import io.substrait.extension.ExtensionCollector import io.substrait.hint.Hint import io.substrait.plan.{ImmutableExecutionBehavior, Plan} @@ -899,14 +900,14 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging with Predic val extensionCollector = new ExtensionCollector val relProtoConverter = new RelProtoConverter(extensionCollector) + val planRel = proto.PlanRel + .newBuilder() + .setRel(substraitRel + .accept(relProtoConverter, EmptyVisitationContext.INSTANCE)) + .build() val builder = proto.Plan .newBuilder() - .addRelations( - proto.PlanRel - .newBuilder() - .setRel(substraitRel - .accept(relProtoConverter, EmptyVisitationContext.INSTANCE)) - ) + .addRelations(NestingDetacher.defaultDetacher().detach(planRel)) extensionCollector.addExtensionsToPlan(builder) builder.build().toByteArray }