[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

@@ -126,6 +126,9 @@ public static BigDecimal numberToBigDecimal(Object value, String kind) {
    * the configured rounding policy. */
   public static BigDecimal numberToBigDecimal(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 safeToRepeat =
+          RexUtil.isDeterministic(sort.fetch);
+      if (sort.offset == null && safeToRepeat) {
         inputFetch = sort.fetch;
-      } else if (sort.fetch instanceof RexLiteral && sort.offset instanceof RexLiteral) {
+      } else if (safeToRepeat
+          && 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,8 +43,25 @@
  * {@link org.apache.calcite.rel.core.Sort}.
  */
 public class SortNode extends AbstractSingleNode<Sort> {
+  private final @Nullable Scalar fetchScalar;
+  private final @Nullable Context fetchContext;
+  private final FetchOffsetRoundingPolicy fetchOffsetRoundingPolicy;
+
   public SortNode(Compiler compiler, Sort rel) {
     super(compiler, rel);
+    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;
   }
 
   private static int getValueAsInt(RexNode node) {
@@ -49,10 +74,7 @@
         rel.offset == null
             ? 0
             : getValueAsInt(rel.offset);
-    final int fetch =
-        rel.fetch == null
-            ? -1
-            : getValueAsInt(rel.fetch);
+    final @Nullable BigDecimal fetch = getFetch();
     // In pure limit mode. No sort required.
     Row row;
   loop:
@@ -63,9 +85,12 @@
           break loop;
         }
       }
-      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,9 +104,11 @@
         list.add(row);
       }
       list.sort(comparator());
-      final int end = fetch < 0 || offset + fetch > list.size()
+      final int available = Math.max(list.size() - offset, 0);
+      final int end = fetch == null
+          || fetch.compareTo(BigDecimal.valueOf(available)) >= 0
           ? list.size()
-          : offset + fetch;
+          : offset + fetchRowCount(fetch);
       for (int i = offset; i < end; i++) {
         sink.send(list.get(i));
       }
@@ -116,4 +143,23 @@
       };
     }
   }
+
+  private @Nullable BigDecimal getFetch() {
+    if (rel.fetch == null) {
+      return null;
+    }
+    final @Nullable Object value;
+    if (rel.fetch instanceof RexLiteral) {
+      value = RexLiteral.bigDecimalValue(rel.fetch);
+    } else {
+      value =
+          requireNonNull(fetchScalar, "fetchScalar")
+              .execute(requireNonNull(fetchContext, "fetchContext"));
+    }
+    return EnumUtils.numberToBigDecimal(value, "FETCH", fetchOffsetRoundingPolicy);
+  }
+
+  private static int fetchRowCount(BigDecimal fetch) {
+    return fetch.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,12 @@ 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, rowCount);
     return limit < rowCount ? limit : rowCount;
   }
 
@@ -129,10 +131,12 @@ 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, rowCount);
     return limit < rowCount ? limit : rowCount;
   }
 
@@ -174,7 +178,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);
         }
       }
     }

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

@@ -165,10 +165,12 @@ public Double getRowCount(Calc rel, RelMetadataQuery mq) {
       return null;
     }
 
-    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;
   }
 
@@ -178,10 +180,12 @@ public Double getRowCount(Calc rel, RelMetadataQuery mq) {
       return null;
     }
 
-    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;
   }