[CALCITE-7592] Add expression support for FETCH

core/src/main/codegen/templates/Parser.jj

@@ -691,7 +691,7 @@ SqlNode ExprOrJoinOrOrderedQuery(ExprContext exprContext) :
  *
  * <blockquote><pre>
  *    [ OFFSET start { ROW | ROWS } ]
- *    [ FETCH { FIRST | NEXT } [ count ] { ROW | ROWS } ONLY ]</pre>
+ *    [ FETCH { FIRST | NEXT } [ count | (expression) ] { ROW | ROWS } ONLY ]</pre>
  * </blockquote>
  */
 SqlNode OrderedQueryOrExpr(ExprContext exprContext) :
@@ -778,10 +778,26 @@ void FetchClause(SqlNode[] offsetFetch) :
 {
     // SQL:2008-style syntax. "OFFSET ... FETCH ...".
     // If you specify both LIMIT and FETCH, FETCH wins.
-    <FETCH> ( <FIRST> | <NEXT> ) offsetFetch[1] = UnsignedNumericLiteralOrParam()
+    <FETCH> ( <FIRST> | <NEXT> ) offsetFetch[1] = FetchCount()
     ( <ROW> | <ROWS> ) <ONLY>
 }
 
+/**
+ * Parses the row count of a FETCH clause. Expressions must be parenthesized.
+ */
+SqlNode FetchCount() :
+{
+    final SqlNode e;
+}
+{
+    (
+        e = UnsignedNumericLiteralOrParam()
+    |
+        <LPAREN> e = Expression(ExprContext.ACCEPT_NON_QUERY) <RPAREN>
+    )
+    { return e; }
+}
+
 /**
  * Parses a LIMIT clause in an ORDER BY expression.
  */

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumUtils.java

@@ -116,16 +116,19 @@
   /** Converts a FETCH or OFFSET runtime value to {@link BigDecimal}.
    *
    * <p>The value must be numeric and non-negative. */
-  public static BigDecimal numberToBigDecimal(Object value, String kind) {
+  public static BigDecimal numberToBigDecimal(@Nullable Object value, String kind) {
     return numberToBigDecimal(value, kind, FetchOffsetRoundingPolicy.NONE);
   }
 
   /** Converts a FETCH or OFFSET runtime value to {@link BigDecimal}.
    *
    * <p>The value must be numeric and non-negative. The result is adjusted by
    * the configured rounding policy. */
-  public static BigDecimal numberToBigDecimal(Object value, String kind,
+  public static BigDecimal numberToBigDecimal(@Nullable Object value, String kind,
       FetchOffsetRoundingPolicy roundingPolicy) {
+    if (value == null) {
+      throw new IllegalArgumentException(kind + " expression evaluated to NULL");
+    }
     if (!(value instanceof Number)) {
       throw new IllegalArgumentException(kind + " must be a number");
     }

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableLimit.java

@@ -106,30 +106,42 @@ public static EnumerableLimit create(final RelNode input, @Nullable RexNode offs
       v =
           builder.append("offset",
               Expressions.call(BuiltInMethod.SKIP_BIG_DECIMAL.method, v,
-                  getExpression(offset, "OFFSET", roundingPolicyExp)));
+                  getExpression(offset, "OFFSET", implementor, builder,
+                      roundingPolicyExp, false)));
     }
     if (fetch != null) {
       v =
           builder.append("fetch",
               Expressions.call(BuiltInMethod.TAKE_BIG_DECIMAL.method, v,
-                  getExpression(fetch, "FETCH", roundingPolicyExp)));
+                  getExpression(fetch, "FETCH", implementor, builder,
+                      roundingPolicyExp, true)));
     }
 
     builder.add(Expressions.return_(null, v));
     return implementor.result(physType, builder.toBlock());
   }
 
   static Expression getExpression(RexNode rexNode, String kind,
-      Expression roundingPolicy) {
+      EnumerableRelImplementor implementor, BlockBuilder builder,
+      Expression roundingPolicy, boolean translateExpression) {
     final Expression value;
     if (rexNode instanceof RexDynamicParam) {
       final RexDynamicParam param = (RexDynamicParam) rexNode;
       value =
           Expressions.call(DataContext.ROOT,
               BuiltInMethod.DATA_CONTEXT_GET.method,
               Expressions.constant("?" + param.getIndex()));
-    } else {
+    } else if (rexNode instanceof RexLiteral) {
       value = Expressions.constant(RexLiteral.bigDecimalValue(rexNode));
+    } else {
+      if (!translateExpression) {
+        throw new IllegalArgumentException(kind + " must be a literal or dynamic parameter");
+      }
+
+      value =
+          RexToLixTranslator.forAggregation(implementor.getTypeFactory(),
+              builder, null, implementor.getConformance())
+              .translate(rexNode);
     }
     return Expressions.call(
         BuiltInMethod.NUMBER_TO_BIG_DECIMAL_LIMIT.method,

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableLimitSort.java

@@ -104,14 +104,16 @@ public static EnumerableLimitSort create(
     if (this.fetch == null) {
       fetchVal = Expressions.constant(BigDecimal.valueOf(Integer.MAX_VALUE));
     } else {
-      fetchVal = getExpression(this.fetch, "FETCH", roundingPolicyExp);
+      fetchVal =
+          getExpression(this.fetch, "FETCH", implementor, builder, roundingPolicyExp, true);
     }
 
     final Expression offsetVal;
     if (this.offset == null) {
       offsetVal = Expressions.constant(BigDecimal.ZERO);
     } else {
-      offsetVal = getExpression(this.offset, "OFFSET", roundingPolicyExp);
+      offsetVal =
+          getExpression(this.offset, "OFFSET", implementor, builder, roundingPolicyExp, false);
     }
 
     builder.add(

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeUnionRule.java

@@ -29,6 +29,7 @@
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.tools.RelBuilder;
 import org.apache.calcite.util.ImmutableBitSet;
 
@@ -88,9 +89,13 @@ public EnumerableMergeUnionRule(Config config) {
     // Push down sort limit, if possible.
     RexNode inputFetch = null;
     if (sort.fetch != null) {
-      if (sort.offset == null) {
+      final boolean safeToReevaluate =
+          RexUtil.isDeterministic(sort.fetch);
+      if (sort.offset == null && safeToReevaluate) {
         inputFetch = sort.fetch;
-      } else if (sort.fetch instanceof RexLiteral && sort.offset instanceof RexLiteral) {
+      } else if (safeToReevaluate
+          && sort.fetch instanceof RexLiteral
+          && sort.offset instanceof RexLiteral) {
         inputFetch =
             call.builder().literal(RexLiteral.bigDecimalValue(sort.fetch)
                 .add(RexLiteral.bigDecimalValue(sort.offset)));

core/src/main/java/org/apache/calcite/interpreter/SortNode.java

@@ -16,14 +16,22 @@
  */
 package org.apache.calcite.interpreter;
 
+import org.apache.calcite.adapter.enumerable.EnumUtils;
+import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor;
+import org.apache.calcite.adapter.enumerable.FetchOffsetRoundingPolicy;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.util.Util;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Ordering;
 
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
@@ -35,37 +43,57 @@
  * {@link org.apache.calcite.rel.core.Sort}.
  */
 public class SortNode extends AbstractSingleNode<Sort> {
+  private final @Nullable Scalar offsetScalar;
+  private final @Nullable Context offsetContext;
+  private final @Nullable Scalar fetchScalar;
+  private final @Nullable Context fetchContext;
+  private final FetchOffsetRoundingPolicy fetchOffsetRoundingPolicy;
+
   public SortNode(Compiler compiler, Sort rel) {
     super(compiler, rel);
-  }
-
-  private static int getValueAsInt(RexNode node) {
-    return requireNonNull(((RexLiteral) node).getValueAs(Integer.class),
-        () -> "getValueAs(Integer.class) for " + node);
+    if (rel.offset != null && !(rel.offset instanceof RexLiteral)) {
+      this.offsetScalar = compiler.compile(ImmutableList.of(rel.offset), null);
+      this.offsetContext = compiler.createContext();
+    } else {
+      this.offsetScalar = null;
+      this.offsetContext = null;
+    }
+    if (rel.fetch != null && !(rel.fetch instanceof RexLiteral)) {
+      this.fetchScalar = compiler.compile(ImmutableList.of(rel.fetch), null);
+      this.fetchContext = compiler.createContext();
+    } else {
+      this.fetchScalar = null;
+      this.fetchContext = null;
+    }
+    final Object roundingPolicy = compiler.getDataContext()
+        .get(EnumerableRelImplementor.FETCH_OFFSET_ROUNDING_POLICY);
+    this.fetchOffsetRoundingPolicy =
+        roundingPolicy instanceof FetchOffsetRoundingPolicy
+            ? (FetchOffsetRoundingPolicy) roundingPolicy
+            : FetchOffsetRoundingPolicy.NONE;
   }
 
   @Override public void run() throws InterruptedException {
-    final int offset =
-        rel.offset == null
-            ? 0
-            : getValueAsInt(rel.offset);
-    final int fetch =
-        rel.fetch == null
-            ? -1
-            : getValueAsInt(rel.fetch);
+    final BigDecimal offset = getOffset();
+    final @Nullable BigDecimal fetch = getFetch();
     // In pure limit mode. No sort required.
     Row row;
   loop:
     if (rel.getCollation().getFieldCollations().isEmpty()) {
-      for (int i = 0; i < offset; i++) {
+      BigDecimal skipped = BigDecimal.ZERO;
+      while (skipped.compareTo(offset) < 0) {
         row = source.receive();
         if (row == null) {
           break loop;
         }
+        skipped = skipped.add(BigDecimal.ONE);
       }
-      if (fetch >= 0) {
-        for (int i = 0; i < fetch && (row = source.receive()) != null; i++) {
+      if (fetch != null) {
+        BigDecimal fetched = BigDecimal.ZERO;
+        while (fetched.compareTo(fetch) < 0
+            && (row = source.receive()) != null) {
           sink.send(row);
+          fetched = fetched.add(BigDecimal.ONE);
         }
       } else {
         while ((row = source.receive()) != null) {
@@ -79,10 +107,15 @@ private static int getValueAsInt(RexNode node) {
         list.add(row);
       }
       list.sort(comparator());
-      final int end = fetch < 0 || offset + fetch > list.size()
+      final int start = offset.compareTo(BigDecimal.valueOf(list.size())) >= 0
+          ? list.size()
+          : rowCount(offset);
+      final int available = list.size() - start;
+      final int end = fetch == null
+          || fetch.compareTo(BigDecimal.valueOf(available)) >= 0
           ? list.size()
-          : offset + fetch;
-      for (int i = offset; i < end; i++) {
+          : start + rowCount(fetch);
+      for (int i = start; i < end; i++) {
         sink.send(list.get(i));
       }
     }
@@ -116,4 +149,35 @@ private static Comparator<Row> comparator(RelFieldCollation fieldCollation) {
       };
     }
   }
+
+  private @Nullable BigDecimal getFetch() {
+    if (rel.fetch == null) {
+      return null;
+    }
+    return getValue(rel.fetch, fetchScalar, fetchContext, "FETCH");
+  }
+
+  private BigDecimal getOffset() {
+    if (rel.offset == null) {
+      return BigDecimal.ZERO;
+    }
+    return getValue(rel.offset, offsetScalar, offsetContext, "OFFSET");
+  }
+
+  private BigDecimal getValue(RexNode node, @Nullable Scalar scalar,
+      @Nullable Context context, String kind) {
+    final @Nullable Object value;
+    if (node instanceof RexLiteral) {
+      value = RexLiteral.bigDecimalValue(node);
+    } else {
+      value =
+          requireNonNull(scalar, () -> kind + " scalar")
+              .execute(requireNonNull(context, () -> kind + " context"));
+    }
+    return EnumUtils.numberToBigDecimal(value, kind, fetchOffsetRoundingPolicy);
+  }
+
+  private static int rowCount(BigDecimal value) {
+    return value.setScale(0, RoundingMode.CEILING).intValueExact();
+  }
 }

core/src/main/java/org/apache/calcite/rel/metadata/RelMdMaxRowCount.java

@@ -117,10 +117,12 @@ public Double getMaxRowCount(Sort rel, RelMetadataQuery mq) {
       rowCount = Double.POSITIVE_INFINITY;
     }
 
-    final double offset = literalValueApproximatedByDouble(rel.offset, 0D);
+    final double offset =
+        literalValueApproximatedByDouble(rel.offset, 0D);
     rowCount = Math.max(rowCount - offset, 0D);
 
-    final double limit = literalValueApproximatedByDouble(rel.fetch, rowCount);
+    final double limit =
+        literalValueApproximatedByDouble(rel.fetch, rowCount);
     return limit < rowCount ? limit : rowCount;
   }
 
@@ -130,10 +132,12 @@ public Double getMaxRowCount(EnumerableLimit rel, RelMetadataQuery mq) {
       rowCount = Double.POSITIVE_INFINITY;
     }
 
-    final double offset = literalValueApproximatedByDouble(rel.offset, 0D);
+    final double offset =
+        literalValueApproximatedByDouble(rel.offset, 0D);
     rowCount = Math.max(rowCount - offset, 0D);
 
-    final double limit = literalValueApproximatedByDouble(rel.fetch, rowCount);
+    final double limit =
+        literalValueApproximatedByDouble(rel.fetch, rowCount);
     return limit < rowCount ? limit : rowCount;
   }
 
@@ -214,7 +218,8 @@ public Double getMaxRowCount(RelSubset rel, RelMetadataQuery mq) {
       if (node instanceof Sort) {
         Sort sort = (Sort) node;
         if (sort.fetch instanceof RexLiteral) {
-          return literalValueApproximatedByDouble(sort.fetch, Double.POSITIVE_INFINITY);
+          return literalValueApproximatedByDouble(sort.fetch,
+              Double.POSITIVE_INFINITY);
         }
       }
     }

core/src/main/java/org/apache/calcite/rel/metadata/RelMdMinRowCount.java

@@ -116,10 +116,13 @@ public Double getMinRowCount(Sort rel, RelMetadataQuery mq) {
       rowCount = 0D;
     }
 
-    final double offset = literalValueApproximatedByDouble(rel.offset, 0D);
+    final double offset =
+        literalValueApproximatedByDouble(rel.offset, 0D);
     rowCount = Math.max(rowCount - offset, 0D);
 
-    final double limit = literalValueApproximatedByDouble(rel.fetch, rowCount);
+    final double limit =
+        literalValueApproximatedByDouble(rel.fetch,
+            rel.fetch == null ? rowCount : 0D);
     return limit < rowCount ? limit : rowCount;
   }
 
@@ -129,10 +132,13 @@ public Double getMinRowCount(EnumerableLimit rel, RelMetadataQuery mq) {
       rowCount = 0D;
     }
 
-    final double offset = literalValueApproximatedByDouble(rel.offset, 0D);
+    final double offset =
+        literalValueApproximatedByDouble(rel.offset, 0D);
     rowCount = Math.max(rowCount - offset, 0D);
 
-    final double limit = literalValueApproximatedByDouble(rel.fetch, rowCount);
+    final double limit =
+        literalValueApproximatedByDouble(rel.fetch,
+            rel.fetch == null ? rowCount : 0D);
     return limit < rowCount ? limit : rowCount;
   }
 
@@ -174,7 +180,8 @@ public Double getMinRowCount(RelSubset rel, RelMetadataQuery mq) {
       if (node instanceof Sort) {
         Sort sort = (Sort) node;
         if (sort.fetch instanceof RexLiteral) {
-          return literalValueApproximatedByDouble(sort.fetch, Double.POSITIVE_INFINITY);
+          return literalValueApproximatedByDouble(sort.fetch,
+              Double.POSITIVE_INFINITY);
         }
       }
     }