From a1ad89a9377453ec4e184df794abf8c6b88b26da Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Mon, 13 Jul 2026 21:07:03 +0200 Subject: [PATCH] feat(core): add TopNRel physical operator Add a TopN POJO (a SingleInputRel) modeling the TopNRel physical relation introduced in the Substrait spec v0.88.0: expression-based offset/count, a tie-handling mode (ROWS_ONLY/WITH_TIES), and one or more sort fields. Map the mode to/from the proto FetchMode enum, wire the relation visitor (RelVisitor, AbstractRelVisitor, RelCopyOnWriteVisitor), add both proto converters (RelProtoConverter and ProtoRelConverter), a SubstraitBuilder DSL helper, and a proto round-trip test. Engine mappings for isthmus (Calcite) and Spark are intentionally out of scope: those integrations map only logical operators today, so how to map this physical operator is deferred. SqlKindFromRel and the example SubstraitStringify get the minimal visit(TopN) needed to keep the direct RelVisitor implementors compiling. Closes #806 --- .../io/substrait/dsl/SubstraitBuilder.java | 62 ++++++++ .../relation/AbstractRelVisitor.java | 6 + .../substrait/relation/ProtoRelConverter.java | 40 +++++ .../relation/RelCopyOnWriteVisitor.java | 22 +++ .../substrait/relation/RelProtoConverter.java | 19 +++ .../io/substrait/relation/RelVisitor.java | 11 ++ .../io/substrait/relation/physical/TopN.java | 118 +++++++++++++++ .../type/proto/TopNRelRoundtripTest.java | 140 ++++++++++++++++++ .../examples/util/SubstraitStringify.java | 18 +++ .../io/substrait/isthmus/SqlKindFromRel.java | 6 + 10 files changed, 442 insertions(+) create mode 100644 core/src/main/java/io/substrait/relation/physical/TopN.java create mode 100644 core/src/test/java/io/substrait/type/proto/TopNRelRoundtripTest.java diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 39fc4fc61..98f21878a 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -38,8 +38,10 @@ import io.substrait.relation.VirtualTableScan; import io.substrait.relation.physical.ComparisonJoinKey; import io.substrait.relation.physical.HashJoin; +import io.substrait.relation.physical.ImmutableTopN; import io.substrait.relation.physical.MergeJoin; import io.substrait.relation.physical.NestedLoopJoin; +import io.substrait.relation.physical.TopN; import io.substrait.type.NamedStruct; import io.substrait.type.Type; import io.substrait.type.TypeCreator; @@ -883,6 +885,66 @@ private Sort sort( return Sort.builder().input(input).sortFields(condition).remap(remap).build(); } + /** + * Creates a top-N relation that sorts and then returns a limited number of rows, using {@link + * TopN.Mode#ROWS_ONLY}. + * + * @param sortFieldFn function to derive the sort fields from the input relation + * @param offset the number of leading rows to skip + * @param count the maximum number of rows to return + * @param input the input relation + * @return a new {@link TopN} relation + */ + public TopN topN( + Function> sortFieldFn, + long offset, + long count, + Rel input) { + return topN(sortFieldFn, offset, count, TopN.Mode.ROWS_ONLY, Optional.empty(), input); + } + + /** + * Creates a top-N relation that sorts and then returns a limited number of rows with a specific + * tie-handling mode. + * + * @param sortFieldFn function to derive the sort fields from the input relation + * @param offset the number of leading rows to skip + * @param count the maximum number of rows to return + * @param mode the tie-handling mode + * @param input the input relation + * @return a new {@link TopN} relation + */ + public TopN topN( + Function> sortFieldFn, + long offset, + long count, + TopN.Mode mode, + Rel input) { + return topN(sortFieldFn, offset, count, mode, Optional.empty(), input); + } + + private TopN topN( + Function> sortFieldFn, + long offset, + long count, + TopN.Mode mode, + Optional remap, + Rel input) { + Iterable sortFields = sortFieldFn.apply(input); + ImmutableTopN.Builder builder = + TopN.builder() + .input(input) + .sortFields(sortFields) + .count(i64(count)) + .mode(mode) + .remap(remap); + // A null offset is treated as 0, so only emit an offset expression when it is non-zero. + if (offset != 0) { + builder.offset(i64(offset)); + } + return builder.build(); + } + // Expressions /** diff --git a/core/src/main/java/io/substrait/relation/AbstractRelVisitor.java b/core/src/main/java/io/substrait/relation/AbstractRelVisitor.java index d39a96bb2..79b19ffa5 100644 --- a/core/src/main/java/io/substrait/relation/AbstractRelVisitor.java +++ b/core/src/main/java/io/substrait/relation/AbstractRelVisitor.java @@ -8,6 +8,7 @@ import io.substrait.relation.physical.RoundRobinExchange; import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TopN; import io.substrait.util.VisitationContext; /** @@ -178,4 +179,9 @@ public O visit(BroadcastExchange exchange, C context) throws E { public O visit(RoundRobinExchange exchange, C context) throws E { return visitFallback(exchange, context); } + + @Override + public O visit(TopN topN, C context) throws E { + return visitFallback(topN, context); + } } diff --git a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java index a58e2e3b4..ac009896a 100644 --- a/core/src/main/java/io/substrait/relation/ProtoRelConverter.java +++ b/core/src/main/java/io/substrait/relation/ProtoRelConverter.java @@ -36,6 +36,7 @@ import io.substrait.proto.ReadRel; import io.substrait.proto.SetRel; import io.substrait.proto.SortRel; +import io.substrait.proto.TopNRel; import io.substrait.proto.UpdateRel; import io.substrait.proto.WriteRel; import io.substrait.relation.extensions.EmptyDetail; @@ -51,6 +52,7 @@ import io.substrait.relation.physical.ImmutableRoundRobinExchange; import io.substrait.relation.physical.ImmutableScatterExchange; import io.substrait.relation.physical.ImmutableSingleBucketExchange; +import io.substrait.relation.physical.ImmutableTopN; import io.substrait.relation.physical.MergeJoin; import io.substrait.relation.physical.MultiBucketExchange; import io.substrait.relation.physical.NestedLoopJoin; @@ -58,6 +60,7 @@ import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; import io.substrait.relation.physical.TargetType; +import io.substrait.relation.physical.TopN; import io.substrait.type.NamedStruct; import io.substrait.type.Type; import io.substrait.type.proto.ProtoTypeConverter; @@ -203,6 +206,8 @@ public Rel from(io.substrait.proto.Rel rel) { return newUpdate(rel.getUpdate()); case EXCHANGE: return newExchange(rel.getExchange()); + case TOP_N: + return newTopN(rel.getTopN()); default: throw new UnsupportedOperationException("Unsupported RelTypeCase of " + relType); } @@ -969,6 +974,41 @@ protected Sort newSort(SortRel rel) { return builder.build(); } + /** + * Converts the corresponding protobuf message to its POJO representation. + * + * @param rel the protobuf value to convert + * @return the converted result + */ + protected TopN newTopN(TopNRel rel) { + Rel input = from(rel.getInput()); + ProtoExpressionConverter converter = + new ProtoExpressionConverter(lookup, extensions, input.getRecordType(), this); + ImmutableTopN.Builder builder = + TopN.builder() + .input(input) + .sortFields( + rel.getSortsList().stream() + .map(converter::fromSortField) + .collect(Collectors.toList())) + .mode(TopN.Mode.fromProto(rel.getMode())); + if (rel.hasOffset()) { + builder.offset(converter.from(rel.getOffset())); + } + if (rel.hasCount()) { + 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(); + } + /** * Converts the corresponding protobuf message to its POJO representation. * diff --git a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java index d5c19ddb1..96bbc961e 100644 --- a/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java +++ b/core/src/main/java/io/substrait/relation/RelCopyOnWriteVisitor.java @@ -17,6 +17,7 @@ import io.substrait.relation.physical.RoundRobinExchange; import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TopN; import io.substrait.util.EmptyVisitationContext; import java.util.List; import java.util.Optional; @@ -416,6 +417,27 @@ public Optional visit(Sort sort, EmptyVisitationContext context) throws E { .build()); } + @Override + public Optional visit(TopN topN, EmptyVisitationContext context) throws E { + Optional input = topN.getInput().accept(this, context); + Optional> sortFields = + transformList(topN.getSortFields(), context, this::visitSortField); + Optional offset = visitOptionalExpression(topN.getOffset(), context); + Optional count = visitOptionalExpression(topN.getCount(), context); + + if (allEmpty(input, sortFields, offset, count)) { + return Optional.empty(); + } + return Optional.of( + TopN.builder() + .from(topN) + .input(input.orElse(topN.getInput())) + .sortFields(sortFields.orElse(topN.getSortFields())) + .offset(or(offset, topN::getOffset)) + .count(or(count, topN::getCount)) + .build()); + } + @Override public Optional visit(Cross cross, EmptyVisitationContext context) throws E { Optional left = cross.getLeft().accept(this, context); diff --git a/core/src/main/java/io/substrait/relation/RelProtoConverter.java b/core/src/main/java/io/substrait/relation/RelProtoConverter.java index 391c8e66b..a7d237b8d 100644 --- a/core/src/main/java/io/substrait/relation/RelProtoConverter.java +++ b/core/src/main/java/io/substrait/relation/RelProtoConverter.java @@ -42,6 +42,7 @@ import io.substrait.proto.SetRel; import io.substrait.proto.SortField; import io.substrait.proto.SortRel; +import io.substrait.proto.TopNRel; import io.substrait.proto.UpdateRel; import io.substrait.proto.WriteRel; import io.substrait.relation.files.FileOrFiles; @@ -56,6 +57,7 @@ import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; import io.substrait.relation.physical.TargetType; +import io.substrait.relation.physical.TopN; import io.substrait.type.proto.TypeProtoConverter; import io.substrait.util.EmptyVisitationContext; import java.util.ArrayList; @@ -856,6 +858,23 @@ public Rel visit(Sort sort, EmptyVisitationContext context) throws RuntimeExcept return Rel.newBuilder().setSort(builder).build(); } + @Override + public Rel visit(TopN topN, EmptyVisitationContext context) throws RuntimeException { + TopNRel.Builder builder = + TopNRel.newBuilder() + .setCommon(common(topN)) + .setInput(toProto(topN.getInput())) + .addAllSorts(toProtoS(topN.getSortFields())) + .setMode(topN.getMode().toProto()); + + topN.getOffset().ifPresent(offset -> builder.setOffset(toProto(offset))); + topN.getCount().ifPresent(count -> builder.setCount(toProto(count))); + + topN.getExtension() + .ifPresent(ae -> builder.setAdvancedExtension(extensionProtoConverter.toProto(ae))); + return Rel.newBuilder().setTopN(builder).build(); + } + @Override public Rel visit(Cross cross, EmptyVisitationContext context) throws RuntimeException { CrossRel.Builder builder = diff --git a/core/src/main/java/io/substrait/relation/RelVisitor.java b/core/src/main/java/io/substrait/relation/RelVisitor.java index 18a4eb328..58bca9c90 100644 --- a/core/src/main/java/io/substrait/relation/RelVisitor.java +++ b/core/src/main/java/io/substrait/relation/RelVisitor.java @@ -8,6 +8,7 @@ import io.substrait.relation.physical.RoundRobinExchange; import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TopN; import io.substrait.util.VisitationContext; /** @@ -318,4 +319,14 @@ public interface RelVisitor * @throws E on visit failure */ O visit(BroadcastExchange exchange, C context) throws E; + + /** + * Visit a physical top-N relation. + * + * @param topN the top-N node + * @param context visitation context + * @return visit result + * @throws E on visit failure + */ + O visit(TopN topN, C context) throws E; } diff --git a/core/src/main/java/io/substrait/relation/physical/TopN.java b/core/src/main/java/io/substrait/relation/physical/TopN.java new file mode 100644 index 000000000..addd38b43 --- /dev/null +++ b/core/src/main/java/io/substrait/relation/physical/TopN.java @@ -0,0 +1,118 @@ +package io.substrait.relation.physical; + +import io.substrait.expression.Expression; +import io.substrait.proto.FetchMode; +import io.substrait.relation.HasExtension; +import io.substrait.relation.RelVisitor; +import io.substrait.relation.SingleInputRel; +import io.substrait.type.Type; +import io.substrait.util.VisitationContext; +import java.util.List; +import java.util.Optional; +import org.immutables.value.Value; + +/** + * A physical relation that combines a sort and a fetch, retaining only the number of records + * required to produce a limited, ordered output. Optionally supports {@code WITH TIES} semantics + * via its {@link Mode}. + */ +@Value.Immutable +public abstract class TopN extends SingleInputRel implements HasExtension { + + /** + * Returns the sort fields defining the ordering, applied in order of precedence. At least one + * sort field is required. + * + * @return the sort fields + */ + public abstract List getSortFields(); + + /** + * Returns the expression evaluating to the number of leading rows to skip. When empty, no rows + * are skipped (equivalent to an offset of 0). + * + * @return the optional offset expression + */ + public abstract Optional getOffset(); + + /** + * Returns the expression evaluating to the maximum number of rows to return. When empty, all + * remaining rows are returned. + * + * @return the optional row count expression + */ + public abstract Optional getCount(); + + /** + * Returns the tie-handling mode for rows tied with the last returned row. Defaults to {@link + * Mode#ROWS_ONLY}. + * + * @return the fetch mode + */ + @Value.Default + public Mode getMode() { + return Mode.ROWS_ONLY; + } + + /** Determines how a {@link TopN} handles rows tied with the last returned row. */ + public enum Mode { + /** Unspecified. Producers should set one of the other modes. */ + UNSPECIFIED(FetchMode.FETCH_MODE_UNSPECIFIED), + /** Return only the requested number of rows. */ + ROWS_ONLY(FetchMode.FETCH_MODE_ROWS_ONLY), + /** Include additional rows tied with the last row, per the sort fields. */ + WITH_TIES(FetchMode.FETCH_MODE_WITH_TIES); + + private final FetchMode proto; + + Mode(FetchMode proto) { + this.proto = proto; + } + + /** + * Returns the protobuf representation of this mode. + * + * @return the proto fetch mode + */ + public FetchMode toProto() { + return proto; + } + + /** + * Returns the {@link Mode} matching the given protobuf fetch mode. + * + * @param proto the proto fetch mode + * @return the matching mode + * @throws IllegalArgumentException if the mode is not recognized + */ + public static Mode fromProto(FetchMode proto) { + for (Mode v : values()) { + if (v.proto == proto) { + return v; + } + } + + throw new IllegalArgumentException("Unknown mode: " + proto); + } + } + + @Override + protected Type.Struct deriveRecordType() { + return getInput().getRecordType(); + } + + @Override + public O accept( + RelVisitor visitor, C context) throws E { + return visitor.visit(this, context); + } + + /** + * Creates a builder for {@link TopN}. + * + * @return a new builder + */ + public static ImmutableTopN.Builder builder() { + return ImmutableTopN.builder(); + } +} diff --git a/core/src/test/java/io/substrait/type/proto/TopNRelRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/TopNRelRoundtripTest.java new file mode 100644 index 000000000..119c0791b --- /dev/null +++ b/core/src/test/java/io/substrait/type/proto/TopNRelRoundtripTest.java @@ -0,0 +1,140 @@ +package io.substrait.type.proto; + +import io.substrait.TestBase; +import io.substrait.expression.Expression; +import io.substrait.relation.Rel; +import io.substrait.relation.physical.TopN; +import java.util.Arrays; +import java.util.Collections; +import org.junit.jupiter.api.Test; + +class TopNRelRoundtripTest extends TestBase { + + final Rel baseTable = + sb.namedScan( + Collections.singletonList("topn_test_table"), + Arrays.asList("id", "amount", "name", "category"), + Arrays.asList(R.I64, R.FP64, R.STRING, R.STRING)); + + private Expression.SortField sortField(int field, Expression.SortDirection direction) { + return Expression.SortField.builder() + .expr(sb.fieldReference(baseTable, field)) + .direction(direction) + .build(); + } + + @Test + void rowsOnlyWithCount() { + Rel topN = + TopN.builder() + .input(baseTable) + .addSortFields(sortField(0, Expression.SortDirection.ASC_NULLS_FIRST)) + .count(sb.i64(10)) + .build(); + + verifyRoundTrip(topN); + } + + @Test + void offsetAndCount() { + Rel topN = + TopN.builder() + .input(baseTable) + .addSortFields(sortField(1, Expression.SortDirection.DESC_NULLS_LAST)) + .offset(sb.i64(5)) + .count(sb.i64(20)) + .build(); + + verifyRoundTrip(topN); + } + + @Test + void countAbsentReturnsAll() { + // No count set: signals ALL. Offset present. + Rel topN = + TopN.builder() + .input(baseTable) + .addSortFields(sortField(0, Expression.SortDirection.ASC_NULLS_FIRST)) + .offset(sb.i64(3)) + .build(); + + verifyRoundTrip(topN); + } + + @Test + void withTiesMode() { + Rel topN = + TopN.builder() + .input(baseTable) + .addSortFields(sortField(1, Expression.SortDirection.DESC_NULLS_FIRST)) + .count(sb.i64(15)) + .mode(TopN.Mode.WITH_TIES) + .build(); + + verifyRoundTrip(topN); + } + + @Test + void defaultModeIsRowsOnly() { + // No mode set: defaults to ROWS_ONLY and must round-trip. + Rel topN = + TopN.builder() + .input(baseTable) + .addSortFields(sortField(2, Expression.SortDirection.ASC_NULLS_LAST)) + .count(sb.i64(1)) + .build(); + + verifyRoundTrip(topN); + } + + @Test + void multipleSortFields() { + Rel topN = + TopN.builder() + .input(baseTable) + .addSortFields( + sortField(3, Expression.SortDirection.ASC_NULLS_FIRST), + sortField(1, Expression.SortDirection.DESC_NULLS_LAST)) + .offset(sb.i64(2)) + .count(sb.i64(8)) + .mode(TopN.Mode.WITH_TIES) + .build(); + + verifyRoundTrip(topN); + } + + @Test + void viaBuilderHelper() { + Rel topN = + sb.topN( + input -> Arrays.asList(sortField(0, Expression.SortDirection.ASC_NULLS_FIRST)), + 0, + 5, + baseTable); + + verifyRoundTrip(topN); + } + + @Test + void nestedUnderTopN() { + Rel inner = + TopN.builder() + .input(baseTable) + .addSortFields(sortField(0, Expression.SortDirection.ASC_NULLS_FIRST)) + .count(sb.i64(100)) + .build(); + + Rel outer = + TopN.builder() + .input(inner) + .addSortFields( + Expression.SortField.builder() + .expr(sb.fieldReference(inner, 1)) + .direction(Expression.SortDirection.DESC_NULLS_LAST) + .build()) + .count(sb.i64(10)) + .build(); + + verifyRoundTrip(outer); + } +} diff --git a/examples/substrait-spark/src/main/java/io/substrait/examples/util/SubstraitStringify.java b/examples/substrait-spark/src/main/java/io/substrait/examples/util/SubstraitStringify.java index b5240f819..4beb416e4 100644 --- a/examples/substrait-spark/src/main/java/io/substrait/examples/util/SubstraitStringify.java +++ b/examples/substrait-spark/src/main/java/io/substrait/examples/util/SubstraitStringify.java @@ -34,6 +34,7 @@ import io.substrait.relation.physical.RoundRobinExchange; import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TopN; import io.substrait.type.NamedStruct; import io.substrait.util.EmptyVisitationContext; import java.util.ArrayList; @@ -426,4 +427,21 @@ public String visit(BroadcastExchange exchange, EmptyVisitationContext context) StringBuilder sb = getIndent().append("broadcastExchange:: "); return getOutdent(sb); } + + @Override + public String visit(TopN topN, EmptyVisitationContext context) throws RuntimeException { + StringBuilder sb = getIndent().append("TopN:: ").append(getRemap(topN)); + topN.getSortFields() + .forEach( + sf -> { + ExpressionStringify expr = new ExpressionStringify(indent); + sb.append(sf.expr().accept(expr, context)).append(" ").append(sf.direction()); + }); + topN.getInputs() + .forEach( + i -> { + sb.append(i.accept(this, context)); + }); + return getOutdent(sb); + } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java b/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java index 87c959211..6459ebf75 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java @@ -31,6 +31,7 @@ import io.substrait.relation.physical.RoundRobinExchange; import io.substrait.relation.physical.ScatterExchange; import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TopN; import io.substrait.util.EmptyVisitationContext; import org.apache.calcite.sql.SqlKind; @@ -264,4 +265,9 @@ public SqlKind visit(BroadcastExchange exchange, EmptyVisitationContext context) throws RuntimeException { return SqlKind.OTHER_DDL; } + + @Override + public SqlKind visit(TopN topN, EmptyVisitationContext context) throws RuntimeException { + return QUERY_KIND; + } }