[CALCITE-7592] Add expression support for FETCH

cassandra/src/main/java/org/apache/calcite/adapter/cassandra/CassandraRules.java

@@ -35,7 +35,9 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.validate.SqlValidatorUtil;
@@ -405,16 +407,33 @@
       super(config);
     }
 
-    public RelNode convert(EnumerableLimit limit) {
+    public @Nullable RelNode convert(EnumerableLimit limit) {
+      final RexLiteral fetch =
+          limit.fetch == null
+              ? null
+              : RexUtil.reduceFetchToLiteral(limit.getCluster(), limit.fetch);
+      if (limit.fetch != null && fetch == null) {
+        return null;
+      }
+      if (fetch != null) {
+        try {
+          RexLiteral.bigDecimalValue(fetch).intValueExact();
+        } catch (ArithmeticException e) {
+          return null;
+        }
+      }
       final RelTraitSet traitSet =
           limit.getTraitSet().replace(CassandraRel.CONVENTION);
       return new CassandraLimit(limit.getCluster(), traitSet,
-        convert(limit.getInput(), CassandraRel.CONVENTION), limit.offset, limit.fetch);
+          convert(limit.getInput(), CassandraRel.CONVENTION), limit.offset, fetch);
     }
 
     @Override public void onMatch(RelOptRuleCall call) {
       EnumerableLimit limit = call.rel(0);
-      call.transformTo(convert(limit));
+      final RelNode converted = convert(limit);
+      if (converted != null) {
+        call.transformTo(converted);
+      }
     }
 
     /** Deprecated in favor of CassandraLimitRuleConfig. */

cassandra/src/test/java/org/apache/calcite/test/CassandraAdapterTest.java

@@ -127,6 +127,34 @@ static void load(CqlSession session) {
         .explainContains("CassandraLimit(fetch=[8])\n");
   }
 
+  @Test void testFetchExpression() {
+    CalciteAssert.that()
+        .with(TWISSANDRA)
+        .query("select \"tweet_id\" from \"userline\" "
+            + "where \"username\" = '!PUBLIC!' "
+            + "fetch next (1 + abs(-2)) rows only")
+        .returnsCount(3)
+        .explainContains("CassandraLimit(fetch=[3])\n");
+    CalciteAssert.that()
+        .with(TWISSANDRA)
+        .query("select \"tweet_id\" from \"userline\" "
+            + "where \"username\" = '!PUBLIC!' "
+            + "fetch next (0 - 1) rows only")
+        .throws_("FETCH value -1 is out of range");
+  }
+
+  @Test void testFetchExpressionBeyondIntegerRange() {
+    CalciteAssert.that()
+        .with(TWISSANDRA)
+        .query("select \"tweet_id\" from \"userline\" "
+            + "where \"username\" = '!PUBLIC!' "
+            + "fetch next "
+            + "(cast(3000000000 as bigint) + 1) rows only")
+        .returnsCount(146)
+        .explainContains("EnumerableLimit(fetch=[+(3000000000:BIGINT, 1)])\n"
+            + "  CassandraToEnumerableConverter\n");
+  }
+
   @Test void testSortLimit() {
     CalciteAssert.that()
         .with(TWISSANDRA)

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

@@ -690,7 +690,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) :
@@ -777,10 +777,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/EnumerableLimit.java

@@ -17,6 +17,11 @@
 package org.apache.calcite.adapter.enumerable;
 
 import org.apache.calcite.DataContext;
+import org.apache.calcite.linq4j.AbstractEnumerable;
+import org.apache.calcite.linq4j.Enumerable;
+import org.apache.calcite.linq4j.EnumerableDefaults;
+import org.apache.calcite.linq4j.Enumerator;
+import org.apache.calcite.linq4j.function.Function1;
 import org.apache.calcite.linq4j.tree.BlockBuilder;
 import org.apache.calcite.linq4j.tree.Expression;
 import org.apache.calcite.linq4j.tree.Expressions;
@@ -33,10 +38,14 @@
 import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.util.BuiltInMethod;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.util.Comparator;
 import java.util.List;
 
 /** Relational expression that applies a limit and/or offset to its input. */
@@ -54,6 +63,7 @@
       @Nullable RexNode offset,
       @Nullable RexNode fetch) {
     super(cluster, traitSet, input);
+    validateLiteralFetch(fetch);
     this.offset = offset;
     this.fetch = fetch;
     assert getConvention() instanceof EnumerableConvention;
@@ -110,8 +120,8 @@
     if (fetch != null) {
       v =
           builder.append("fetch",
-              Expressions.call(v, BuiltInMethod.TAKE.method,
-                  getExpression(fetch)));
+              Expressions.call(EnumerableLimit.class, "take", v,
+                  getExpressionForFetch(fetch, implementor, builder)));
     }
 
     builder.add(Expressions.return_(null, v));
@@ -127,10 +137,105 @@
               Expressions.constant("?" + param.getIndex())),
           Integer.class);
     } else {
-      // TODO: Enumerable runtime only supports INT types for FETCH and OFFSET, not BIGINT types.
-      //  Currently, using BIGINT types for execution will result in an error message.
+      // TODO: Enumerable runtime only supports INT types for OFFSET, not BIGINT types.
+      //  Currently, using BIGINT types for OFFSET will result in an error message.
       //  This issue needs to be fixed. For more information, see CALCITE-7156.
       return Expressions.constant(RexLiteral.intValue(rexNode));
     }
   }
+
+  static Expression getExpressionForFetch(RexNode rexNode,
+      EnumerableRelImplementor implementor, BlockBuilder builder) {
+    if (rexNode instanceof RexDynamicParam) {
+      final RexDynamicParam param = (RexDynamicParam) rexNode;
+      return Expressions.call(EnumerableLimit.class, "toFetchValue",
+          Expressions.convert_(
+              Expressions.call(DataContext.ROOT,
+                  BuiltInMethod.DATA_CONTEXT_GET.method,
+                  Expressions.constant("?" + param.getIndex())),
+              Number.class));
+    } else if (rexNode instanceof RexLiteral) {
+      return Expressions.constant(
+          toFetchValue(((RexLiteral) rexNode).getValueAs(Number.class)));
+    } else {
+      final Expression expression =
+          RexToLixTranslator.forAggregation(implementor.getTypeFactory(),
+              builder, null, implementor.getConformance())
+              .translate(rexNode);
+      return Expressions.call(EnumerableLimit.class, "toFetchValue",
+          Expressions.convert_(Expressions.box(expression), Number.class));
+    }
+  }
+
+  /** Converts a FETCH expression result to Calcite's canonical representation. */
+  public static BigDecimal toFetchValue(@Nullable Number value) {
+    return RexUtil.validateFetchValue(value);
+  }
+
+  /** Applies a FETCH value without narrowing it to {@code int} or {@code long}. */
+  public static <T> Enumerable<T> take(Enumerable<T> source, BigDecimal fetch) {
+    final BigInteger count = fetch.toBigIntegerExact();
+    return new AbstractEnumerable<T>() {
+      @Override public Enumerator<T> enumerator() {
+        final Enumerator<T> input = source.enumerator();
+        return new Enumerator<T>() {
+          private BigInteger remaining = count;
+          private boolean done;
+
+          @Override public T current() {
+            return input.current();
+          }
+
+          @Override public boolean moveNext() {
+            if (done) {
+              return false;
+            }
+            if (remaining.signum() == 0 || !input.moveNext()) {
+              done = true;
+              return false;
+            }
+            // Preserve take(int)'s eager evaluation of the current row.
+            input.current();
+            remaining = remaining.subtract(BigInteger.ONE);
+            return true;
+          }
+
+          @Override public void reset() {
+            input.reset();
+            remaining = count;
+            done = false;
+          }
+
+          @Override public void close() {
+            input.close();
+          }
+        };
+      }
+    };
+  }
+
+  /** Sorts and applies FETCH while preserving an arbitrary-precision value. */
+  public static <T, TKey> Enumerable<T> orderBy(Enumerable<T> source,
+      Function1<T, TKey> keySelector, @Nullable Comparator<TKey> comparator,
+      int offset, @Nullable BigDecimal fetch) {
+    if (fetch != null
+        && fetch.compareTo(BigDecimal.valueOf(Integer.MAX_VALUE)) <= 0
+        && comparator != null) {
+      return EnumerableDefaults.orderBy(source, keySelector, comparator,
+          offset, fetch.intValueExact());
+    }
+    Enumerable<T> result = EnumerableDefaults.orderBy(source, keySelector, comparator);
+    if (offset > 0) {
+      result = result.skip(offset);
+    }
+    return fetch == null ? result : take(result, fetch);
+  }
+
+  static void validateLiteralFetch(@Nullable RexNode fetch) {
+    if (fetch instanceof RexLiteral) {
+      final Number value = ((RexLiteral) fetch).getValueAs(Number.class);
+      toFetchValue(value);
+    }
+  }
+
 }

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

@@ -25,12 +25,14 @@
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rex.RexNode;
-import org.apache.calcite.util.BuiltInMethod;
 import org.apache.calcite.util.Pair;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import java.math.BigDecimal;
+
 import static org.apache.calcite.adapter.enumerable.EnumerableLimit.getExpression;
+import static org.apache.calcite.adapter.enumerable.EnumerableLimit.getExpressionForFetch;
 
 /**
  * Implementation of {@link org.apache.calcite.rel.core.Sort} in
@@ -52,6 +54,7 @@ public EnumerableLimitSort(
       @Nullable RexNode offset,
       @Nullable RexNode fetch) {
     super(cluster, traitSet, input, collation, offset, fetch);
+    EnumerableLimit.validateLiteralFetch(fetch);
     assert this.getConvention() instanceof EnumerableConvention;
     assert this.getConvention() == input.getConvention();
   }
@@ -98,9 +101,9 @@ public static EnumerableLimitSort create(
 
     final Expression fetchVal;
     if (this.fetch == null) {
-      fetchVal = Expressions.constant(Integer.MAX_VALUE);
+      fetchVal = Expressions.constant(null, BigDecimal.class);
     } else {
-      fetchVal = getExpression(this.fetch);
+      fetchVal = getExpressionForFetch(this.fetch, implementor, builder);
     }
 
     final Expression offsetVal;
@@ -112,17 +115,13 @@ public static EnumerableLimitSort create(
 
     builder.add(
         Expressions.return_(null,
-            Expressions.call(BuiltInMethod.ORDER_BY_WITH_FETCH_AND_OFFSET.method,
+            Expressions.call(EnumerableLimit.class, "orderBy",
                 Expressions.list(childExp,
                         builder.append("keySelector", pair.left))
                     .appendIfNotNull(
                         builder.appendIfNotNull("comparator", pair.right))
-                    .appendIfNotNull(
-                        builder.appendIfNotNull("offset",
-                            Expressions.constant(offsetVal)))
-                    .appendIfNotNull(
-                        builder.appendIfNotNull("fetch",
-                            Expressions.constant(fetchVal))))));
+                    .append(builder.append("offset", offsetVal))
+                    .append(builder.append("fetch", fetchVal)))));
     return implementor.result(physType, builder.toBlock());
   }
 }

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/adapter/jdbc/JdbcRules.java

@@ -54,6 +54,7 @@
 import org.apache.calcite.rel.rel2sql.SqlImplementor;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexProgram;
@@ -767,7 +768,18 @@ protected JdbcSortRule(Config config) {
      *                            JDBC convention
      * @return A new JdbcSort
      */
-    public RelNode convert(Sort sort, boolean convertInputTraits) {
+    public @Nullable RelNode convert(Sort sort, boolean convertInputTraits) {
+      final RexNode fetch;
+      if (sort.fetch == null
+          || sort.fetch instanceof RexLiteral
+          || sort.fetch instanceof RexDynamicParam) {
+        fetch = sort.fetch;
+      } else {
+        fetch = RexUtil.reduceFetchToLiteral(sort.getCluster(), sort.fetch);
+        if (fetch == null) {
+          return null;
+        }
+      }
       final RelTraitSet traitSet = sort.getTraitSet().replace(out);
 
       final RelNode input;
@@ -779,7 +791,7 @@ public RelNode convert(Sort sort, boolean convertInputTraits) {
       }
 
       return new JdbcSort(sort.getCluster(), traitSet,
-          input, sort.getCollation(), sort.offset, sort.fetch);
+          input, sort.getCollation(), sort.offset, fetch);
     }
   }