diff --git a/AGENTS.md b/AGENTS.md index 2392dbda6..004211e2a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,6 +114,19 @@ Proto conversion is split into two directions, and the class name tells you whic to one proto value, order the deprecated ones **last** so the reverse lookup returns the canonical one (see `Join.JoinType`). Adding an enum value is an edit to the enum, not the converters. +- **`RelCommon` data** (emit mapping, hint, rel anchor, common advanced extension) is converted by + one helper per direction: `ProtoRelConverter.applyRelCommon(rel, relCommon)` on the way in and + `RelProtoConverter.common(Rel)` on the way out. Both are still *called* per relation, so both + halves are easy to forget: a new `newXxx` must route its result through `applyRelCommon`, and a + new `visit` must call `.setCommon(common(rel))`. `RelCommonRoundtripTest` covers every + `RelVisitor` type and fails when one has no sample — but it is keyed on POJO types, not on proto + `oneof` cases, so a new case mapping to an existing POJO type slips through. `applyRelCommon` + runs *after* `build()`, so a relation with a `@Value.Check` on one of these fields must also set + it on its builder (see `newLateralJoin`). `UpdateRel` is the only *modeled* relation message with + no `common` field (`ReferenceRel` has none either, but no POJO models it); the POJO accessors live + on `Rel`, so `NamedUpdate` can still *hold* this data — + and an emit mapping would change its `getRecordType()` — so `RelProtoConverter.checkNoRelCommon` + rejects it on the way out rather than silently serializing a plan with a different schema. - POJO types are created with `TypeCreator.REQUIRED` / `TypeCreator.NULLABLE`. ## Building and testing diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 141ecf473..e70016348 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -822,6 +822,14 @@ public NamedUpdate namedUpdate( /** * Creates a named update relation that updates rows in a table with output field remapping. * + *

Note that {@code UpdateRel} has no {@code common} field as of spec v0.99.0, so it has + * nowhere to carry the emit mapping. The resulting relation therefore cannot be serialized: + * {@link io.substrait.relation.RelProtoConverter} rejects it rather than emitting a plan whose + * schema silently differs from the relation's own {@link + * io.substrait.relation.Rel#getRecordType()}. Until the spec gives {@code UpdateRel} a {@code + * common} field, use {@link #namedUpdate(Iterable, Iterable, List, Expression, boolean)} for + * relations that have to survive a round trip. + * * @param tableName the qualified name of the table to update * @param columnNames the names of the columns in the table * @param transformations the list of transformation expressions to apply diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java index ed7f08e2c..e4045be93 100644 --- a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java +++ b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java @@ -362,16 +362,10 @@ protected NamedWrite newNamedWrite(final WriteRel rel) { .outputMode(NamedWrite.OutputMode.fromProto(rel.getOutput())) .operation(NamedWrite.WriteOp.fromProto(rel.getOp())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); - if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -393,16 +387,10 @@ protected Rel newExtensionWrite(final WriteRel rel) { .outputMode(NamedWrite.OutputMode.fromProto(rel.getOutput())) .operation(NamedWrite.WriteOp.fromProto(rel.getOp())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); - if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -438,17 +426,13 @@ protected NamedDdl newNamedDdl(final DdlRel rel) { .tableDefaults(tableDefaults(rel.getTableDefaults(), tableSchema)) .operation(NamedDdl.DdlOp.fromProto(rel.getOp())) .object(NamedDdl.DdlObject.fromProto(rel.getObject())) - .viewDefinition(optionalViewDefinition(rel)) - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); + .viewDefinition(optionalViewDefinition(rel)); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -468,17 +452,13 @@ protected ExtensionDdl newExtensionDdl(final DdlRel rel) { .tableDefaults(tableDefaults(rel.getTableDefaults(), tableSchema)) .operation(ExtensionDdl.DdlOp.fromProto(rel.getOp())) .object(ExtensionDdl.DdlObject.fromProto(rel.getObject())) - .viewDefinition(optionalViewDefinition(rel)) - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); + .viewDefinition(optionalViewDefinition(rel)); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -554,6 +534,10 @@ protected Rel newNamedUpdate(UpdateRel rel) { if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } + // No applyRelCommon here: alone among the modeled relation messages, UpdateRel carries no + // `common` field (spec v0.99.0), so an update relation cannot express an emit mapping, a hint, + // a rel anchor or a common advanced extension. (ReferenceRel has none either, but no POJO + // models it.) return builder.build(); } @@ -571,15 +555,10 @@ protected Filter newFilter(FilterRel rel) { .condition( new ProtoExpressionConverter(lookup, extensions, input.getRecordType(), this) .from(rel.getCondition())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -621,13 +600,7 @@ protected NamedStruct newNamedStruct(io.substrait.proto.NamedStruct namedStruct) */ protected ExtensionLeaf newExtensionLeaf(ExtensionLeafRel rel) { Extension.LeafRelDetail detail = detailFromExtensionLeafRel(rel.getDetail()); - ImmutableExtensionLeaf.Builder builder = - ExtensionLeaf.from(detail) - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); - return builder.build(); + return applyRelCommon(ExtensionLeaf.from(detail).build(), rel.getCommon()); } /** @@ -639,13 +612,7 @@ protected ExtensionLeaf newExtensionLeaf(ExtensionLeafRel rel) { protected ExtensionSingle newExtensionSingle(ExtensionSingleRel rel) { Extension.SingleRelDetail detail = detailFromExtensionSingleRel(rel.getDetail()); Rel input = from(rel.getInput()); - ImmutableExtensionSingle.Builder builder = - ExtensionSingle.from(detail, input) - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); - return builder.build(); + return applyRelCommon(ExtensionSingle.from(detail, input).build(), rel.getCommon()); } /** @@ -657,16 +624,7 @@ protected ExtensionSingle newExtensionSingle(ExtensionSingleRel rel) { protected ExtensionMulti newExtensionMulti(ExtensionMultiRel rel) { Extension.MultiRelDetail detail = detailFromExtensionMultiRel(rel.getDetail()); List inputs = rel.getInputsList().stream().map(this::from).collect(Collectors.toList()); - ImmutableExtensionMulti.Builder builder = - ExtensionMulti.from(detail, inputs) - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); - if (rel.hasDetail()) { - builder.detail(detailFromExtensionMultiRel(rel.getDetail())); - } - return builder.build(); + return applyRelCommon(ExtensionMulti.from(detail, inputs).build(), rel.getCommon()); } /** @@ -697,15 +655,10 @@ protected NamedScan newNamedScan(ReadRel rel) { : null)) .projection(optionalMaskExpression(rel)); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -721,16 +674,11 @@ protected ExtensionTable newExtensionTable(final ReadRel rel) { final ImmutableExtensionTable.Builder builder = ExtensionTable.from(detail).initialSchema(namedStruct); - builder - .projection(optionalMaskExpression(rel)) - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); + builder.projection(optionalMaskExpression(rel)); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -765,15 +713,10 @@ protected LocalFiles newLocalFiles(ReadRel rel) { : null)) .projection(optionalMaskExpression(rel)); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -851,15 +794,10 @@ protected VirtualTableScan newVirtualTable(ReadRel rel) { .rows(expressions) .projection(optionalMaskExpression(rel)); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -881,15 +819,10 @@ protected Fetch newFetch(FetchRel rel) { builder.count(converter.from(rel.getCountExpr())); } - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -910,15 +843,10 @@ protected Project newProject(ProjectRel rel) { .map(converter::from) .collect(java.util.stream.Collectors.toList())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -957,12 +885,7 @@ protected Expand newExpand(ExpandRel rel) { }) .collect(java.util.stream.Collectors.toList())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1009,15 +932,10 @@ protected Aggregate newAggregate(AggregateRel rel) { ImmutableAggregate.Builder builder = Aggregate.builder().input(input).groupings(groupings).measures(measures); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1043,15 +961,10 @@ protected Sort newSort(SortRel rel) { .build()) .collect(java.util.stream.Collectors.toList())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1079,14 +992,10 @@ protected TopN newTopN(TopNRel rel) { builder.count(converter.from(rel.getCount())); } - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1113,15 +1022,10 @@ protected Join newJoin(JoinRel rel) { Optional.ofNullable( rel.hasPostJoinFilter() ? converter.from(rel.getPostJoinFilter()) : null)); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1132,10 +1036,10 @@ protected Join newJoin(JoinRel rel) { */ protected Rel newLateralJoin(LateralJoinRel rel) { Rel left = from(rel.getLeft()); + Optional relAnchor = optionalRelAnchor(rel.getCommon()); // The right input's outer references resolve to the current left row via this join's anchor, so // register that scope before converting the right input (which is where those references live). - optionalRelAnchor(rel.getCommon()) - .ifPresent(anchor -> anchorScopes.put(anchor, left.getRecordType())); + relAnchor.ifPresent(anchor -> anchorScopes.put(anchor, left.getRecordType())); Rel right = from(rel.getRight()); Type.Struct leftStruct = left.getRecordType(); Type.Struct rightStruct = right.getRecordType(); @@ -1152,17 +1056,16 @@ protected Rel newLateralJoin(LateralJoinRel rel) { .joinType(Join.JoinType.fromProto(rel.getType())) .postJoinFilter( Optional.ofNullable( - rel.hasPostJoinFilter() ? converter.from(rel.getPostJoinFilter()) : null)); + rel.hasPostJoinFilter() ? converter.from(rel.getPostJoinFilter()) : null)) + // A lateral join validates that it carries an anchor at construction time, so the + // anchor has to be set here rather than being left to applyRelCommon, which only runs + // after build() (it then sees the same value and skips it). + .relAnchor(relAnchor); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1176,14 +1079,10 @@ protected Rel newCross(CrossRel rel) { Rel right = from(rel.getRight()); ImmutableCross.Builder builder = Cross.builder().left(left).right(right); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1200,15 +1099,10 @@ protected Set newSet(SetRel rel) { ImmutableSet.Builder builder = Set.builder().inputs(inputs).setOp(Set.SetOp.fromProto(rel.getOp())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1244,15 +1138,10 @@ protected Rel newHashJoin(HashJoinRel rel) { rel.hasResidualExpression() ? unionConverter.from(rel.getResidualExpression()) : null)); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1289,15 +1178,10 @@ protected Rel newMergeJoin(MergeJoinRel rel) { ? unionConverter.from(rel.getResidualExpression()) : null)); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1367,15 +1251,10 @@ protected NestedLoopJoin newNestedLoopJoin(NestedLoopJoinRel rel) { : Expression.BoolLiteral.builder().value(true).build()) .joinType(NestedLoopJoin.JoinType.fromProto(rel.getType())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1411,15 +1290,10 @@ protected ConsistentPartitionWindow newConsistentPartitionWindow( .sorts(sortFields) .windowFunctions(windowRelFunctions); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1471,15 +1345,10 @@ protected ScatterExchange newScatterExchange(ExchangeRel rel) { .partitionCount(rel.getPartitionCount()) .targets(targets); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1502,15 +1371,10 @@ protected SingleBucketExchange newSingleBucketExchange(ExchangeRel rel) { .targets(targets) .expression(protoExprConverter.from(rel.getSingleTarget().getExpression())); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1534,15 +1398,10 @@ protected MultiBucketExchange newMultiBucketExchange(ExchangeRel rel) { .expression(protoExprConverter.from(rel.getMultiTarget().getExpression())) .constrainedToCount(rel.getMultiTarget().getConstrainedToCount()); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1563,15 +1422,10 @@ protected RoundRobinExchange newRoundRobinExchange(ExchangeRel rel) { .targets(targets) .exact(rel.getRoundRobin().getExact()); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1591,15 +1445,10 @@ protected BroadcastExchange newBroadcastExchange(ExchangeRel rel) { .partitionCount(rel.getPartitionCount()) .targets(targets); - builder - .commonExtension(optionalAdvancedExtension(rel.getCommon())) - .remap(optionalRelmap(rel.getCommon())) - .hint(optionalHint(rel.getCommon())) - .relAnchor(optionalRelAnchor(rel.getCommon())); if (rel.hasAdvancedExtension()) { builder.extension(protoExtensionConverter.fromProto(rel.getAdvancedExtension())); } - return builder.build(); + return applyRelCommon(builder.build(), rel.getCommon()); } /** @@ -1626,6 +1475,55 @@ protected AbstractExchangeRel.ExchangeTarget newExchangeTarget( return builder.build(); } + /** + * Copies every {@link io.substrait.proto.RelCommon} field that the POJO model represents — the + * {@link Rel#getRemap() emit mapping}, the {@link Rel#getCommonExtension() common extension}, the + * {@link Rel#getHint() hint} and the {@link Rel#getRelAnchor() relation anchor} — onto a freshly + * converted relation. + * + *

This is the proto → POJO counterpart of {@link RelProtoConverter}'s {@code + * common(io.substrait.relation.Rel)} and exists so that adding a {@code RelCommon} field means + * touching one method rather than every {@code newXxx} converter. Every {@code newXxx} method + * whose protobuf message carries a {@code common} field must route its result through here. + * + *

Each field is only copied when it actually differs from what {@code rel} already carries. + * That keeps this method a true no-op for a {@code common { direct {} }} message — including for + * custom, non-Immutables {@link Rel} implementations, which inherit {@code Rel}'s throwing {@code + * withXxx} defaults and would otherwise fail on a relation that carries no common data at all. + * + *

Note that the {@code newXxx} method has already called {@code build()} by the time this + * runs, so a {@code @Value.Check} that requires one of these fields (see {@link LateralJoin}) + * cannot be satisfied here — such a relation must also set that field on its builder. + * + * @param the relation type, preserved so callers keep their concrete return type + * @param rel the relation to copy the common fields onto + * @param relCommon the protobuf value to convert + * @return a copy of {@code rel} carrying the converted common fields, or {@code rel} itself when + * it already carries them + */ + @SuppressWarnings("unchecked") + protected R applyRelCommon(R rel, io.substrait.proto.RelCommon relCommon) { + // Every Immutables-generated withXxx returns the concrete relation type, so R is preserved. + Rel result = rel; + Optional relAnchor = optionalRelAnchor(relCommon); + if (!relAnchor.equals(result.getRelAnchor())) { + result = result.withRelAnchor(relAnchor); + } + Optional remap = optionalRelmap(relCommon); + if (!remap.equals(result.getRemap())) { + result = result.withRemap(remap); + } + Optional commonExtension = optionalAdvancedExtension(relCommon); + if (!commonExtension.equals(result.getCommonExtension())) { + result = result.withCommonExtension(commonExtension); + } + Optional hint = optionalHint(relCommon); + if (!hint.equals(result.getHint())) { + result = result.withHint(hint); + } + return (R) result; + } + /** * Converts the corresponding protobuf message to its POJO representation. * diff --git a/core/src/main/java/io/substrait/relation/Rel.java b/core/src/main/java/io/substrait/relation/Rel.java index ca4a9b5e6..cd6c3ea96 100644 --- a/core/src/main/java/io/substrait/relation/Rel.java +++ b/core/src/main/java/io/substrait/relation/Rel.java @@ -73,6 +73,61 @@ default Rel withRelAnchor(Optional relAnchor) { getClass() + " does not support setting a relation anchor"); } + /** + * Returns a copy of this relation with its {@link #getRemap() output mapping} set to the given + * optional value ({@link Optional#empty()} clears it, making the relation emit directly). + * + *

Like {@link #withRelAnchor(Optional)}, this is overridden by the generated Immutables {@code + * withRemap(Optional)} and provides a type-agnostic way to set the remap on an arbitrary + * relation. Custom {@link Rel} implementations that are not Immutables-backed inherit this + * throwing default. + * + *

The {@code ? extends} wildcard is load-bearing: Immutables emits it for attribute types that + * are not {@code final} (unlike {@code withRelAnchor(Optional)}), and no + * {@code @Override} links the two, so narrowing this parameter type would silently leave the + * throwing default in place on every generated relation. + * + * @param remap the output mapping to set, or empty to clear it + * @return a copy of this relation carrying the given output mapping + */ + default Rel withRemap(Optional remap) { + throw new UnsupportedOperationException( + getClass() + " does not support setting an output mapping"); + } + + /** + * Returns a copy of this relation with its {@link #getCommonExtension() common extension} set to + * the given optional value ({@link Optional#empty()} clears it). + * + *

Like {@link #withRelAnchor(Optional)}, this is overridden by the generated Immutables {@code + * withCommonExtension(Optional)} and provides a type-agnostic way to set the extension on an + * arbitrary relation. Custom {@link Rel} implementations that are not Immutables-backed inherit + * this throwing default. + * + * @param commonExtension the extension to set, or empty to clear it + * @return a copy of this relation carrying the given common extension + */ + default Rel withCommonExtension(Optional commonExtension) { + throw new UnsupportedOperationException( + getClass() + " does not support setting a common extension"); + } + + /** + * Returns a copy of this relation with its {@link #getHint() hint} set to the given optional + * value ({@link Optional#empty()} clears it). + * + *

Like {@link #withRelAnchor(Optional)}, this is overridden by the generated Immutables {@code + * withHint(Optional)} and provides a type-agnostic way to set the hint on an arbitrary relation. + * Custom {@link Rel} implementations that are not Immutables-backed inherit this throwing + * default. + * + * @param hint the hint to set, or empty to clear it + * @return a copy of this relation carrying the given hint + */ + default Rel withHint(Optional hint) { + throw new UnsupportedOperationException(getClass() + " does not support setting a hint"); + } + /** * Returns the record type (schema) produced by this relation. * diff --git a/core/src/main/java/io/substrait/relation/RelProtoConverter.java b/core/src/main/java/io/substrait/relation/RelProtoConverter.java index 0096a872f..090955075 100644 --- a/core/src/main/java/io/substrait/relation/RelProtoConverter.java +++ b/core/src/main/java/io/substrait/relation/RelProtoConverter.java @@ -613,6 +613,7 @@ public Rel visit(ExtensionDdl ddl, EmptyVisitationContext context) throws Runtim @Override public Rel visit(NamedUpdate update, EmptyVisitationContext context) throws RuntimeException { + checkNoRelCommon(update); UpdateRel.Builder builder = UpdateRel.newBuilder() .setNamedTable(NamedTable.newBuilder().addAllNames(update.getNames())) @@ -927,6 +928,31 @@ public Rel visit(ExtensionMulti extensionMulti, EmptyVisitationContext context) return Rel.newBuilder().setExtensionMulti(builder).build(); } + /** + * Rejects {@link RelCommon} data on a relation whose protobuf message has no {@code common} + * field. {@code UpdateRel} is the only modeled message without one (spec v0.99.0), so an update + * relation cannot express an emit mapping, a hint, a rel anchor or a common advanced extension. + * ({@code ReferenceRel} has no {@code common} field either, but it points at another subtree + * instead of producing output of its own, and no POJO models it.) The POJO model exposes those + * accessors on every {@link io.substrait.relation.Rel}, and an emit mapping additionally changes + * {@link io.substrait.relation.Rel#getRecordType()} — serializing such a relation would silently + * produce a plan with a different schema, so fail instead. + * + * @param rel the relation about to be serialized + */ + private static void checkNoRelCommon(io.substrait.relation.Rel rel) { + if (rel.getRemap().isPresent() + || rel.getHint().isPresent() + || rel.getRelAnchor().isPresent() + || rel.getCommonExtension().isPresent()) { + throw new UnsupportedOperationException( + "Relation " + + rel.getClass() + + " carries RelCommon data (emit mapping, hint, rel anchor or common advanced " + + "extension) but its protobuf message has no common field to carry it"); + } + } + private RelCommon common(io.substrait.relation.Rel rel) { RelCommon.Builder builder = RelCommon.newBuilder(); rel.getCommonExtension() @@ -949,6 +975,9 @@ private RelCommon common(io.substrait.relation.Rel rel) { hint.getAlias().ifPresent(hintBuilder::setAlias); hintBuilder.addAllOutputNames(hint.getOutputNames()); + hint.getExtension() + .ifPresent(ae -> hintBuilder.setAdvancedExtension(extensionProtoConverter.toProto(ae))); + if (hint.getStats().isPresent()) { io.substrait.hint.Hint.Stats stats = hint.getStats().get(); Stats.Builder statsBuilder = Stats.newBuilder(); diff --git a/core/src/test/java/io/substrait/type/proto/RelCommonRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/RelCommonRoundtripTest.java new file mode 100644 index 000000000..b6acceb32 --- /dev/null +++ b/core/src/test/java/io/substrait/type/proto/RelCommonRoundtripTest.java @@ -0,0 +1,449 @@ +package io.substrait.type.proto; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.TestBase; +import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; +import io.substrait.expression.WindowBound; +import io.substrait.extension.AdvancedExtension; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.ExtensionLookup; +import io.substrait.extension.SimpleExtension; +import io.substrait.hint.Hint; +import io.substrait.proto.RelCommon; +import io.substrait.relation.AbstractWriteRel; +import io.substrait.relation.ConsistentPartitionWindow; +import io.substrait.relation.Expand; +import io.substrait.relation.ExtensionDdl; +import io.substrait.relation.ExtensionLeaf; +import io.substrait.relation.ExtensionMulti; +import io.substrait.relation.ExtensionSingle; +import io.substrait.relation.ExtensionTable; +import io.substrait.relation.ExtensionWrite; +import io.substrait.relation.Join; +import io.substrait.relation.LateralJoin; +import io.substrait.relation.LocalFiles; +import io.substrait.relation.NamedDdl; +import io.substrait.relation.NamedUpdate; +import io.substrait.relation.ProtoRelConverter; +import io.substrait.relation.Rel; +import io.substrait.relation.RelVisitor; +import io.substrait.relation.Set; +import io.substrait.relation.SingleInputRel; +import io.substrait.relation.VirtualTableScan; +import io.substrait.relation.extensions.EmptyDetail; +import io.substrait.relation.physical.BroadcastExchange; +import io.substrait.relation.physical.HashJoin; +import io.substrait.relation.physical.MergeJoin; +import io.substrait.relation.physical.MultiBucketExchange; +import io.substrait.relation.physical.NestedLoopJoin; +import io.substrait.relation.physical.RoundRobinExchange; +import io.substrait.relation.physical.ScatterExchange; +import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.type.NamedStruct; +import io.substrait.type.Type; +import io.substrait.util.VisitationContext; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Optional; +import java.util.stream.Collectors; +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 that every relation type carries the {@link io.substrait.proto.RelCommon} data the POJO + * model represents — the {@link Rel#getHint() hint}, {@link Rel#getRemap() emit mapping}, {@link + * Rel#getCommonExtension() common extension} and {@link Rel#getRelAnchor() rel anchor} — through a + * POJO → proto → POJO round trip. + * + *

{@link #everyRelationTypeIsCovered()} keeps the sample set exhaustive: it fails when a + * relation is added to {@link RelVisitor} without a sample here, so a new relation cannot silently + * miss its {@code RelCommon} wiring. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class RelCommonRoundtripTest extends TestBase { + + /** + * Relations whose protobuf message has no {@code common} field and which therefore cannot carry + * any {@code RelCommon} data at all. {@code UpdateRel} is the only modeled message without one + * (spec v0.99.0); {@code ReferenceRel} has no {@code common} field either, but no POJO models it, + * so it cannot reach this set. + */ + static final java.util.Set> WITHOUT_REL_COMMON = + Collections.singleton(NamedUpdate.class); + + static final Hint HINT = + Hint.builder() + .extension(AdvancedExtension.builder().build()) + .alias("an_alias") + .addAllOutputNames(Arrays.asList("name1", "name2")) + .stats(Hint.Stats.builder().rowCount(42).recordSize(13).build()) + .runtimeConstraint(Hint.RuntimeConstraint.builder().build()) + .addLoadedComputations( + Hint.LoadedComputation.builder() + .computationId(1) + .computationType(Hint.ComputationType.COMPUTATION_TYPE_HASHTABLE) + .build()) + .addSavedComputations( + Hint.SavedComputation.builder() + .computationId(2) + .computationType(Hint.ComputationType.COMPUTATION_TYPE_BLOOM_FILTER) + .build()) + .build(); + + static final AdvancedExtension COMMON_EXTENSION = AdvancedExtension.builder().build(); + + static final int REL_ANCHOR = 11; + + final Rel left = + sb.namedScan( + Arrays.asList("left_table"), Arrays.asList("a", "b"), Arrays.asList(R.I64, R.STRING)); + + final Rel right = + sb.namedScan( + Arrays.asList("right_table"), Arrays.asList("c", "d"), Arrays.asList(R.I64, R.STRING)); + + final NamedStruct schema = NamedStruct.of(Arrays.asList("a", "b"), R.struct(R.I64, R.STRING)); + + @ParameterizedTest(name = "{0}") + @MethodSource("samples") + void relCommonRoundtrips(String relationType, Rel rel) { + if (WITHOUT_REL_COMMON.contains(relationType(rel))) { + // The message has nowhere to carry the data, so serialization must fail loudly instead of + // emitting a plan that silently lost it (an emit mapping even changes the record type). + // Each field is rejected on its own, so none of them can start being dropped silently. + for (Rel sample : + Arrays.asList( + rel.withRemap(Optional.of(reversedRemap(rel))), + rel.withHint(Optional.of(HINT)), + rel.withRelAnchor(Optional.of(REL_ANCHOR)), + rel.withCommonExtension(Optional.of(COMMON_EXTENSION)), + withRelCommon(rel))) { + assertThrows( + UnsupportedOperationException.class, + () -> relProtoConverter.toProto(sample), + relationType); + } + // Without any RelCommon data the relation still round-trips. + verifyRoundTrip(rel); + return; + } + verifyRoundTrip(withRelCommon(rel)); + } + + @Test + void everyRelationTypeIsCovered() { + java.util.Set> visitable = + Arrays.stream(RelVisitor.class.getDeclaredMethods()) + .filter(method -> "visit".equals(method.getName())) + .map(method -> method.getParameterTypes()[0]) + .collect(Collectors.toCollection(HashSet::new)); + + java.util.Set> sampled = + allRelationTypes().stream() + .map(RelCommonRoundtripTest::relationType) + .collect(Collectors.toCollection(HashSet::new)); + + assertEquals(visitable, sampled, "every relation type needs a RelCommon round-trip sample"); + } + + @Test + void samplesCarryTheDataUnderTest() { + // Guards the round-trip assertions above against silently asserting on empty optionals. + for (Rel rel : allRelationTypes()) { + Rel sample = withRelCommon(rel); + assertEquals(Optional.of(HINT), sample.getHint()); + assertEquals(Optional.of(COMMON_EXTENSION), sample.getCommonExtension()); + assertTrue(sample.getRemap().isPresent()); + assertTrue(sample.getRelAnchor().isPresent()); + } + } + + @Test + void applyRelCommonLeavesACustomRelAloneWhenThereIsNothingToApply() { + // A custom, non-Immutables Rel inherits Rel's throwing withXxx defaults. applyRelCommon is a + // documented extension point that every newXxx must route through, so it must not fail on a + // relation whose common message carries no data. + ApplyRelCommonConverter converter = new ApplyRelCommonConverter(functionCollector, extensions); + Rel custom = new PassThroughRel(left); + + assertEquals(custom, converter.applyRelCommon(custom, RelCommon.getDefaultInstance())); + assertEquals( + custom, + converter.applyRelCommon( + custom, + RelCommon.newBuilder().setDirect(RelCommon.Direct.getDefaultInstance()).build())); + } + + /** Exposes {@link ProtoRelConverter#applyRelCommon} so the test can call it directly. */ + static final class ApplyRelCommonConverter extends ProtoRelConverter { + ApplyRelCommonConverter( + ExtensionLookup lookup, SimpleExtension.ExtensionCollection extensions) { + super(lookup, extensions); + } + + @Override + public R applyRelCommon(R rel, RelCommon relCommon) { + return super.applyRelCommon(rel, relCommon); + } + } + + /** A minimal hand-written {@link Rel} that inherits {@code Rel}'s throwing {@code withXxx}. */ + static final class PassThroughRel extends SingleInputRel { + private final Rel input; + + PassThroughRel(Rel input) { + this.input = input; + } + + @Override + public Rel getInput() { + return input; + } + + @Override + protected Type.Struct deriveRecordType() { + return input.getRecordType(); + } + + @Override + public Optional getRemap() { + return Optional.empty(); + } + + @Override + public Optional getCommonExtension() { + return Optional.empty(); + } + + @Override + public Optional getHint() { + return Optional.empty(); + } + + @Override + public Optional getRelAnchor() { + return Optional.empty(); + } + + @Override + public O accept( + RelVisitor visitor, C context) { + throw new UnsupportedOperationException("not visitable"); + } + } + + Stream samples() { + return allRelationTypes().stream() + .map(rel -> Arguments.of(relationType(rel).getSimpleName(), rel)); + } + + /** + * Stamps the full set of {@code RelCommon} data onto {@code rel} using the type-agnostic {@code + * Rel.withXxx} copy methods. An existing rel anchor is kept because a {@link LateralJoin}'s right + * input resolves its outer references against it. + */ + static Rel withRelCommon(Rel rel) { + Optional relAnchor = + rel.getRelAnchor().isPresent() ? rel.getRelAnchor() : Optional.of(REL_ANCHOR); + return rel.withHint(Optional.of(HINT)) + .withCommonExtension(Optional.of(COMMON_EXTENSION)) + .withRemap(Optional.of(reversedRemap(rel))) + .withRelAnchor(relAnchor); + } + + /** + * An emit mapping that reverses {@code rel}'s fields. Reversing rather than using the identity + * makes the assertions sensitive to the order and count of {@code RelCommon.Emit.output_mapping}, + * not just to its presence. + */ + static Rel.Remap reversedRemap(Rel rel) { + List indices = new ArrayList<>(); + for (int i = rel.getRecordType().fields().size() - 1; i >= 0; i--) { + indices.add(i); + } + return Rel.Remap.of(indices); + } + + /** The declared relation type of {@code rel}, i.e. the type its generated Immutable extends. */ + static Class relationType(Rel rel) { + return rel.getClass().getSuperclass(); + } + + /** One sample of every relation type reachable through {@link RelVisitor}. */ + List allRelationTypes() { + List rels = new ArrayList<>(); + + // Read relations + rels.add(left); + rels.add( + VirtualTableScan.builder() + .initialSchema(schema) + .addRows( + Expression.NestedStruct.builder() + .addFields(ExpressionCreator.i64(false, 1)) + .addFields(ExpressionCreator.string(false, "one")) + .build()) + .build()); + rels.add(LocalFiles.builder().initialSchema(schema).build()); + // A real schema rather than EmptyDetail's empty one, so the reversed emit mapping is non-empty + // and the assertions cover the order of RelCommon.Emit.output_mapping for this relation too. + rels.add(ExtensionTable.from(new EmptyDetail()).initialSchema(schema).build()); + + // Logical relations + rels.add(sb.filter(input -> sb.equal(sb.fieldReference(input, 0), sb.i64(1)), left)); + rels.add(sb.limit(10, left)); + rels.add(sb.project(input -> Arrays.asList(sb.fieldReference(input, 0)), left)); + rels.add( + sb.aggregate( + input -> sb.grouping(input, 0), input -> Arrays.asList(sb.count(input, 0)), left)); + rels.add(sb.sort(input -> sb.sortFields(input, 0), left)); + rels.add( + sb.innerJoin( + inputs -> sb.equal(sb.fieldReference(inputs, 0), sb.fieldReference(inputs, 2)), + left, + right)); + rels.add(sb.cross(left, right)); + rels.add(sb.set(Set.SetOp.UNION_ALL, left, right)); + rels.add( + sb.expand( + input -> + Arrays.asList( + Expand.ConsistentField.builder() + .expression(sb.fieldReference(input, 0)) + .build()), + left)); + rels.add( + LateralJoin.builder() + .left(left) + .right(right) + .joinType(Join.JoinType.INNER) + .relAnchor(7) + .build()); + rels.add(consistentPartitionWindow()); + + // Physical relations + rels.add(sb.topN(input -> sb.sortFields(input, 0), 0, 10, left)); + rels.add(sb.hashJoin(Arrays.asList(0), Arrays.asList(0), HashJoin.JoinType.INNER, left, right)); + rels.add( + sb.mergeJoin(Arrays.asList(0), Arrays.asList(0), MergeJoin.JoinType.INNER, left, right)); + rels.add( + sb.nestedLoopJoin( + inputs -> sb.equal(sb.fieldReference(inputs, 0), sb.fieldReference(inputs, 2)), + NestedLoopJoin.JoinType.INNER, + left, + right)); + rels.add(BroadcastExchange.builder().input(left).partitionCount(1).build()); + rels.add(RoundRobinExchange.builder().input(left).exact(true).partitionCount(1).build()); + rels.add( + ScatterExchange.builder() + .input(left) + .addFields(sb.fieldReference(left, 0)) + .partitionCount(1) + .build()); + rels.add( + SingleBucketExchange.builder() + .input(left) + .expression(sb.fieldReference(left, 0)) + .partitionCount(1) + .build()); + rels.add( + MultiBucketExchange.builder() + .input(left) + .expression(sb.fieldReference(left, 0)) + .constrainedToCount(true) + .partitionCount(1) + .build()); + + // Write, DDL and update relations + rels.add( + sb.namedWrite( + Arrays.asList("target_table"), + Arrays.asList("a", "b"), + AbstractWriteRel.WriteOp.INSERT, + AbstractWriteRel.CreateMode.REPLACE_IF_EXISTS, + AbstractWriteRel.OutputMode.NO_OUTPUT, + left)); + rels.add( + ExtensionWrite.builder() + .input(left) + .detail(new EmptyDetail()) + .tableSchema(schema) + .operation(ExtensionWrite.WriteOp.INSERT) + .createMode(ExtensionWrite.CreateMode.APPEND_IF_EXISTS) + .outputMode(ExtensionWrite.OutputMode.NO_OUTPUT) + .build()); + rels.add( + NamedDdl.builder() + .names(Arrays.asList("target_table")) + .tableSchema(schema) + .tableDefaults(tableDefaults()) + .operation(NamedDdl.DdlOp.CREATE) + .object(NamedDdl.DdlObject.TABLE) + .build()); + rels.add( + ExtensionDdl.builder() + .detail(new EmptyDetail()) + .tableSchema(schema) + .tableDefaults(tableDefaults()) + .operation(ExtensionDdl.DdlOp.ALTER) + .object(ExtensionDdl.DdlObject.TABLE) + .build()); + rels.add( + sb.namedUpdate( + Arrays.asList("target_table"), + Arrays.asList("a"), + Arrays.asList( + NamedUpdate.TransformExpression.builder() + .columnTarget(0) + .transformation(sb.i64(1)) + .build()), + sb.bool(true), + false)); + + // Extension relations + rels.add(ExtensionLeaf.from(new EmptyDetail()).build()); + rels.add(ExtensionSingle.from(new EmptyDetail(), left).build()); + rels.add(ExtensionMulti.from(new EmptyDetail(), Arrays.asList(left, right)).build()); + + return rels; + } + + private Expression.StructLiteral tableDefaults() { + return ExpressionCreator.struct( + false, ExpressionCreator.i64(false, 1), ExpressionCreator.string(false, "one")); + } + + private Rel consistentPartitionWindow() { + SimpleExtension.WindowFunctionVariant lead = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + return ConsistentPartitionWindow.builder() + .input(left) + .addWindowFunctions( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(lead) + .arguments(Arrays.asList(sb.fieldReference(left, 0))) + .outputType(R.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .lowerBound(WindowBound.Unbounded.UNBOUNDED) + .upperBound(WindowBound.Following.CURRENT_ROW) + .boundsType(Expression.WindowBoundsType.RANGE) + .build()) + .addPartitionExpressions(sb.fieldReference(left, 1)) + .sorts(sb.sortFields(left, 0)) + .build(); + } +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java index b6f18d045..f61c1be26 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubtraitRelVisitorExtensionTest.java @@ -131,6 +131,21 @@ public Rel withRelAnchor(final Optional relAnchor) { return new SubstraitRepeatRel(input.withRelAnchor(relAnchor), repeatCount); } + @Override + public Rel withRemap(final Optional remap) { + return new SubstraitRepeatRel(input.withRemap(remap), repeatCount); + } + + @Override + public Rel withCommonExtension(final Optional commonExtension) { + return new SubstraitRepeatRel(input.withCommonExtension(commonExtension), repeatCount); + } + + @Override + public Rel withHint(final Optional hint) { + return new SubstraitRepeatRel(input.withHint(hint), repeatCount); + } + @Override public O accept( final RelVisitor visitor, final C context) throws E { @@ -291,4 +306,26 @@ void customRelSupportsRelAnchor() { // Clearing the anchor works too. assertFalse(anchored.withRelAnchor(Optional.empty()).getRelAnchor().isPresent()); } + + @Test + void customRelSupportsTheRestOfTheRelCommonContract() { + // The remaining RelCommon copy methods are delegated the same way, so a custom Rel can be + // handed to code that stamps an emit mapping, a hint or a common extension onto an arbitrary + // relation without hitting Rel's throwing defaults. + final SubstraitBuilder sb = new SubstraitBuilder(); + final Rel scan = sb.namedScan(List.of("t"), List.of("a"), List.of(TypeCreator.REQUIRED.I64)); + final SubstraitRepeatRel repeat = new SubstraitRepeatRel(scan, 3); + + final Rel remapped = repeat.withRemap(Optional.of(Rel.Remap.of(List.of(0)))); + assertTrue(remapped instanceof SubstraitRepeatRel); + assertEquals(List.of(0), remapped.getRemap().orElseThrow(AssertionError::new).indices()); + + final Hint hint = Hint.builder().alias("an_alias").build(); + assertEquals(Optional.of(hint), repeat.withHint(Optional.of(hint)).getHint()); + + final AdvancedExtension extension = AdvancedExtension.builder().build(); + assertEquals( + Optional.of(extension), + repeat.withCommonExtension(Optional.of(extension)).getCommonExtension()); + } }