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; + } }