[CALCITE-7592] Add expression support for FETCH

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

@@ -35,6 +35,7 @@
 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.RexVisitorImpl;
 import org.apache.calcite.sql.SqlKind;
@@ -405,16 +406,26 @@ protected CassandraLimitRule(CassandraLimitRuleConfig config) {
       super(config);
     }
 
-    public RelNode convert(EnumerableLimit limit) {
+    public @Nullable RelNode convert(EnumerableLimit limit) {
+      final RexLiteral fetch =
+          limit.fetch == null
+              ? null
+              : EnumerableLimit.reduceFetchToIntLiteral(limit.getCluster(), limit.fetch);
+      if (limit.fetch != null && fetch == null) {
+        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,22 @@ 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 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

@@ -31,12 +31,20 @@
 import org.apache.calcite.rel.metadata.RelMdDistribution;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexExecutor;
 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.util.BuiltInMethod;
+import org.apache.calcite.util.NumberUtil;
+import org.apache.calcite.util.Util;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /** Relational expression that applies a limit and/or offset to its input. */
@@ -54,6 +62,7 @@
       @Nullable RexNode offset,
       @Nullable RexNode fetch) {
     super(cluster, traitSet, input);
+    validateLiteralFetch(fetch);
     this.offset = offset;
     this.fetch = fetch;
     assert getConvention() instanceof EnumerableConvention;
@@ -111,7 +120,7 @@
       v =
           builder.append("fetch",
               Expressions.call(v, BuiltInMethod.TAKE.method,
-                  getExpression(fetch)));
+                  getExpressionForFetch(fetch, implementor, builder)));
     }
 
     builder.add(Expressions.return_(null, v));
@@ -133,4 +142,160 @@
       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, "toIntFetch",
+          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(
+          toIntFetch(((RexLiteral) rexNode).getValueAs(Number.class)));
+    } else {
+      final Expression expression =
+          RexToLixTranslator.forAggregation(implementor.getTypeFactory(),
+              builder, null, implementor.getConformance())
+              .translate(rexNode);
+      return Expressions.call(EnumerableLimit.class, "toIntFetch",
+          Expressions.convert_(Expressions.box(expression), Number.class));
+    }
+  }
+
+  /** Converts a FETCH expression result to the range supported by Enumerable. */
+  public static int toIntFetch(@Nullable Number value) {
+    final BigDecimal decimal = validateFetchValue(value);
+    final int result;
+    try {
+      result = decimal.intValueExact();
+    } catch (ArithmeticException e) {
+      throw new IllegalArgumentException("FETCH value " + value
+          + " is out of range; expected a value between 0 and "
+          + Integer.MAX_VALUE, e);
+    }
+    if (result < 0) {
+      throw new IllegalArgumentException("FETCH value " + value
+          + " is out of range; expected a value between 0 and "
+          + Integer.MAX_VALUE);
+    }
+    return result;
+  }
+
+  /** Converts a FETCH expression result to the range supported by a long. */
+  public static long toLongFetch(@Nullable Number value) {
+    final BigDecimal decimal = validateFetchValue(value);
+    final long result;
+    try {
+      result = decimal.longValueExact();
+    } catch (ArithmeticException e) {
+      throw new IllegalArgumentException("FETCH value " + value
+          + " is out of range; expected a value between 0 and "
+          + Long.MAX_VALUE, e);
+    }
+    if (result < 0) {
+      throw new IllegalArgumentException("FETCH value " + value
+          + " is out of range; expected a value between 0 and "
+          + Long.MAX_VALUE);
+    }
+    return result;
+  }
+
+  private static BigDecimal validateFetchValue(@Nullable Number value) {
+    if (value == null) {
+      throw new IllegalArgumentException("FETCH expression evaluated to NULL");
+    }
+    final BigDecimal decimal = NumberUtil.toBigDecimal(value);
+    if (decimal == null) {
+      throw new IllegalArgumentException("FETCH value is not numeric: " + value);
+    }
+    try {
+      decimal.toBigIntegerExact();
+    } catch (ArithmeticException e) {
+      throw new IllegalArgumentException("FETCH value " + value
+          + " is not an integer", e);
+    }
+    return decimal;
+  }
+
+  static void validateLiteralFetch(@Nullable RexNode fetch) {
+    if (fetch instanceof RexLiteral) {
+      final Number value = ((RexLiteral) fetch).getValueAs(Number.class);
+      final BigDecimal decimal = NumberUtil.toBigDecimal(value);
+      if (decimal != null && decimal.signum() < 0) {
+        toIntFetch(value);
+      }
+    }
+  }
+
+  /** Reduces a constant FETCH expression to a validated literal. */
+  public static @Nullable RexLiteral reduceFetchToLiteral(
+      RelOptCluster cluster, RexNode fetch) {
+    final RexLiteral literal;
+    if (fetch instanceof RexLiteral) {
+      literal = (RexLiteral) fetch;
+    } else {
+      if (!RexUtil.isConstant(fetch)
+          || !RexUtil.isDeterministic(fetch)
+          || RexUtil.containsDynamicFunction(fetch)
+          || containsDynamicParam(fetch)) {
+        return null;
+      }
+      final RexExecutor executor =
+          Util.first(cluster.getPlanner().getExecutor(), RexUtil.EXECUTOR);
+      final List<RexNode> reducedValues = new ArrayList<>(1);
+      executor.reduce(cluster.getRexBuilder(),
+          Collections.singletonList(fetch), reducedValues);
+      final RexNode reduced = reducedValues.get(0);
+      if (!(reduced instanceof RexLiteral)) {
+        return null;
+      }
+      literal = (RexLiteral) reduced;
+    }
+    final Number value = literal.getValueAs(Number.class);
+    if (validateFetchValue(value).signum() < 0) {
+      throw new IllegalArgumentException(
+          "FETCH value " + value + " is out of range; expected a non-negative value");
+    }
+    return literal;
+  }
+
+  /** Reduces a constant FETCH expression to a validated integer literal. */
+  public static @Nullable RexLiteral reduceFetchToIntLiteral(
+      RelOptCluster cluster, RexNode fetch) {
+    final RexLiteral literal = reduceFetchToLiteral(cluster, fetch);
+    if (literal == null) {
+      return null;
+    }
+    final int value = toIntFetch(literal.getValueAs(Number.class));
+    return cluster.getRexBuilder().makeExactLiteral(BigDecimal.valueOf(value));
+  }
+
+  /** Reduces a constant FETCH expression to a validated long literal. */
+  public static @Nullable RexLiteral reduceFetchToLongLiteral(
+      RelOptCluster cluster, RexNode fetch) {
+    final RexLiteral literal = reduceFetchToLiteral(cluster, fetch);
+    if (literal == null) {
+      return null;
+    }
+    final long value = toLongFetch(literal.getValueAs(Number.class));
+    return cluster.getRexBuilder().makeExactLiteral(BigDecimal.valueOf(value));
+  }
+
+  private static boolean containsDynamicParam(RexNode node) {
+    try {
+      node.accept(
+          new RexVisitorImpl<Void>(true) {
+            @Override public Void visitDynamicParam(RexDynamicParam dynamicParam) {
+              throw Util.FoundOne.NULL;
+            }
+          });
+      return false;
+    } catch (Util.FoundOne e) {
+      return true;
+    }
+  }
 }

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

@@ -31,6 +31,7 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 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 +53,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();
   }
@@ -100,7 +102,7 @@ public static EnumerableLimitSort create(
     if (this.fetch == null) {
       fetchVal = Expressions.constant(Integer.MAX_VALUE);
     } else {
-      fetchVal = getExpression(this.fetch);
+      fetchVal = getExpressionForFetch(this.fetch, implementor, builder);
     }
 
     final Expression offsetVal;

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

@@ -16,6 +16,7 @@
  */
 package org.apache.calcite.adapter.jdbc;
 
+import org.apache.calcite.adapter.enumerable.EnumerableLimit;
 import org.apache.calcite.linq4j.Queryable;
 import org.apache.calcite.linq4j.tree.Expression;
 import org.apache.calcite.plan.Contexts;
@@ -54,6 +55,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 +769,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 = EnumerableLimit.reduceFetchToLiteral(sort.getCluster(), sort.fetch);
+        if (fetch == null) {
+          return null;
+        }
+      }
       final RelTraitSet traitSet = sort.getTraitSet().replace(out);
 
       final RelNode input;
@@ -779,7 +792,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);
     }
   }