[CALCITE-6128] RelBuilder.limit should apply offset and fetch to prev…

…ious Sort operator, if possible

The RelBuilder.limit(offset, fetch) method should apply offset
and fetch to the previous Sort operator, possible. If I call
RelBuilder.sort to create a Sort with an offset but no fetch,
then I call RelBuilder.limit(0, fetch) to set a fetch, it
currently creates two LogicalSort nodes but should just set
the fetch of the LogicalSort that already exists.

Close apache#3538

core/src/main/java/org/apache/calcite/tools/RelBuilder.java

@@ -161,6 +161,7 @@
 import static org.apache.calcite.rel.rules.AggregateRemoveRule.canFlattenStatic;
 import static org.apache.calcite.sql.SqlKind.UNION;
 import static org.apache.calcite.util.Static.RESOURCE;
+import static org.apache.calcite.util.Util.first;
 
 import static java.util.Objects.requireNonNull;
 
@@ -204,7 +205,7 @@ protected RelBuilder(@Nullable Context context, RelOptCluster cluster,
     final RexExecutor executor =
         context.maybeUnwrap(RexExecutor.class)
             .orElse(
-                Util.first(cluster.getPlanner().getExecutor(),
+                first(cluster.getPlanner().getExecutor(),
                     RexUtil.EXECUTOR));
     final RelOptPredicateList predicates = RelOptPredicateList.EMPTY;
     this.simplifier =
@@ -3323,7 +3324,11 @@ private static <E> ImmutableList<ImmutableList<E>> copy(
     return builder.build();
   }
 
-  /** Creates a limit without a sort. */
+  /** Creates a limit and/or offset without a sort.
+   *
+   * @param offset Number of rows to skip; non-positive means don't skip any
+   * @param fetch Maximum number of rows to fetch; negative means no limit
+   */
   public RelBuilder limit(int offset, int fetch) {
     return sortLimit(offset, fetch, ImmutableList.of());
   }
@@ -3434,11 +3439,16 @@ public RelBuilder sortLimit(@Nullable RexNode offsetNode, @Nullable RexNode fetc
       RelNode top = peek();
       if (top instanceof Sort) {
         final Sort sort2 = (Sort) top;
-        if (sort2.offset == null && sort2.fetch == null) {
+        if ((offsetNode == null || sort2.offset == null)
+            && (fetchNode == null || sort2.fetch == null)) {
+          // We're not trying to replace something that's already set - an
+          // offset in a sort that already has an offset, or a fetch in a sort
+          // that already has a fetch - and so we can merge them.
           replaceTop(sort2.getInput());
           final RelNode sort =
               struct.sortFactory.createSort(peek(), sort2.collation,
-                  offsetNode, fetchNode);
+                  first(offsetNode, sort2.offset),
+                  first(fetchNode, sort2.fetch));
           replaceTop(sort);
           return this;
         }
@@ -3482,7 +3492,7 @@ private static RelFieldCollation collation(RexNode node,
     switch (node.getKind()) {
     case INPUT_REF:
       return new RelFieldCollation(((RexInputRef) node).getIndex(), direction,
-          Util.first(nullDirection, direction.defaultNullDirection()));
+          first(nullDirection, direction.defaultNullDirection()));
     case DESCENDING:
       return collation(((RexCall) node).getOperands().get(0),
           RelFieldCollation.Direction.DESCENDING,
@@ -3497,7 +3507,7 @@ private static RelFieldCollation collation(RexNode node,
       final int fieldIndex = extraNodes.size();
       extraNodes.add(node);
       return new RelFieldCollation(fieldIndex, direction,
-          Util.first(nullDirection, direction.defaultNullDirection()));
+          first(nullDirection, direction.defaultNullDirection()));
     }
   }
 

core/src/test/java/org/apache/calcite/test/RelBuilderTest.java

@@ -3671,6 +3671,51 @@ private static RelBuilder assertSize(RelBuilder b,
         hasTree(expectedNoSimplify));
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6128">[CALCITE-6128]
+   * RelBuilder.sortLimit should compose offset and fetch</a>. */
+  @Test void testSortOffsetLimit() {
+    // Equivalent SQL:
+    //   SELECT *
+    //   FROM emp
+    //   ORDER BY deptno OFFSET 2 LIMIT 3
+
+    // Case 1. Set sort+offset, then set fetch.
+    final Function<RelBuilder, RelNode> f = b ->
+        b.scan("EMP")
+            .sortLimit(2, -1, b.field("DEPTNO")) // ORDER BY deptno OFFSET 2
+            .limit(-1, 3) // LIMIT 3
+            .build();
+    final String expected = ""
+        + "LogicalSort(sort0=[$7], dir0=[ASC], offset=[2], fetch=[3])\n"
+        + "  LogicalTableScan(table=[[scott, EMP]])\n";
+    assertThat(f.apply(createBuilder()), hasTree(expected));
+    assertThat(f.apply(createBuilder(c -> c.withSimplifyLimit(true))),
+        hasTree(expected));
+
+    // Case 2. Set sort, then offset, then fetch. Same effect as case 1.
+    final Function<RelBuilder, RelNode> f2 = b ->
+        b.scan("EMP")
+            .sort(b.field("DEPTNO")) // ORDER BY deptno
+            .limit(2, -1) // OFFSET 2
+            .limit(-1, 3) // LIMIT 3
+            .build();
+    assertThat(f2.apply(createBuilder()), hasTree(expected));
+    assertThat(f2.apply(createBuilder(c -> c.withSimplifyLimit(true))),
+        hasTree(expected));
+
+    // Case 3. Set sort, then fetch, then offset. Same effect as case 1 & 2.
+    final Function<RelBuilder, RelNode> f3 = b ->
+        b.scan("EMP")
+            .sort(b.field("DEPTNO")) // ORDER BY deptno
+            .limit(-1, 3) // LIMIT 3
+            .limit(2, -1) // OFFSET 2
+            .build();
+    assertThat(f3.apply(createBuilder()), hasTree(expected));
+    assertThat(f3.apply(createBuilder(c -> c.withSimplifyLimit(true))),
+        hasTree(expected));
+  }
+
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-1610">[CALCITE-1610]
    * RelBuilder sort-combining optimization treats aliases incorrectly</a>. */