feat(core): add TopNRel physical operator

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<Rel, Iterable<? extends Expression.SortField>> 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<Rel, Iterable<? extends Expression.SortField>> sortFieldFn,
+      long offset,
+      long count,
+      TopN.Mode mode,
+      Rel input) {
+    return topN(sortFieldFn, offset, count, mode, Optional.empty(), input);
+  }
+
+  private TopN topN(
+      Function<Rel, Iterable<? extends Expression.SortField>> sortFieldFn,
+      long offset,
+      long count,
+      TopN.Mode mode,
+      Optional<Rel.Remap> remap,
+      Rel input) {
+    Iterable<? extends Expression.SortField> 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
 
   /**

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

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,13 +52,15 @@
 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;
 import io.substrait.relation.physical.RoundRobinExchange;
 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.
    *

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<Rel> visit(Sort sort, EmptyVisitationContext context) throws E {
             .build());
   }
 
+  @Override
+  public Optional<Rel> visit(TopN topN, EmptyVisitationContext context) throws E {
+    Optional<Rel> input = topN.getInput().accept(this, context);
+    Optional<List<Expression.SortField>> sortFields =
+        transformList(topN.getSortFields(), context, this::visitSortField);
+    Optional<Expression> offset = visitOptionalExpression(topN.getOffset(), context);
+    Optional<Expression> 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<Rel> visit(Cross cross, EmptyVisitationContext context) throws E {
     Optional<Rel> left = cross.getLeft().accept(this, context);

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 =

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<O, C extends VisitationContext, E extends Exception>
    * @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;
 }

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<Expression.SortField> 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<Expression> 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<Expression> 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, C extends VisitationContext, E extends Exception> O accept(
+      RelVisitor<O, C, E> 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();
+  }
+}