[CALCITE-7662] Add expression support for OFFSET

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

@@ -671,13 +671,13 @@ SqlNode ExprOrJoinOrOrderedQuery(ExprContext exprContext) :
  *
  * <blockquote><pre>
  *    [ LIMIT { count | ALL } ]
- *    [ OFFSET start ]</pre>
+ *    [ OFFSET { start | expression } ]</pre>
  * </blockquote>
  *
  * <p>Trino syntax for limit:
  *
  * <blockquote><pre>
- *    [ OFFSET start ]
+ *    [ OFFSET { start | expression } ]
  *    [ LIMIT { count | ALL } ]</pre>
  * </blockquote>
  *
@@ -690,7 +690,7 @@ SqlNode ExprOrJoinOrOrderedQuery(ExprContext exprContext) :
  * <p>SQL:2008 syntax for limit:
  *
  * <blockquote><pre>
- *    [ OFFSET start { ROW | ROWS } ]
+ *    [ OFFSET { start | expression } { ROW | ROWS } ]
  *    [ FETCH { FIRST | NEXT } [ count | (expression) ] { ROW | ROWS } ONLY ]</pre>
  * </blockquote>
  */
@@ -765,7 +765,7 @@ void OffsetClause(Span s, SqlNode[] offsetFetch) :
     // ROW or ROWS is required in SQL:2008 but we make it optional
     // because it is not present in Postgres-style syntax.
     <OFFSET> { s.add(this); }
-    offsetFetch[0] = UnsignedNumericLiteralOrParam()
+    offsetFetch[0] = OffsetCount()
     [ <ROW> | <ROWS> ]
 }
 
@@ -782,6 +782,23 @@ void FetchClause(SqlNode[] offsetFetch) :
     ( <ROW> | <ROWS> ) <ONLY>
 }
 
+/**
+ * Parses the start value or expression of an OFFSET clause.
+ */
+SqlNode OffsetCount() :
+{
+    final SqlNode e;
+}
+{
+    // Unlike FETCH expressions, OFFSET expressions do not require parentheses
+    // and may start with a numeric literal or dynamic parameter. Therefore, a
+    // separate UnsignedNumericLiteralOrParam alternative would consume only
+    // the prefix of expressions such as "OFFSET 1 + 2" or "OFFSET ? + 1".
+    // Expression also covers standalone start values.
+    e = Expression(ExprContext.ACCEPT_NON_QUERY)
+    { return e; }
+}
+
 /**
  * Parses the row count of a FETCH clause. Expressions must be parenthesized.
  */

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

@@ -107,14 +107,14 @@ public static EnumerableLimit create(final RelNode input, @Nullable RexNode offs
           builder.append("offset",
               Expressions.call(BuiltInMethod.SKIP_BIG_DECIMAL.method, v,
                   getExpression(offset, "OFFSET", implementor, builder,
-                      roundingPolicyExp, false)));
+                      roundingPolicyExp)));
     }
     if (fetch != null) {
       v =
           builder.append("fetch",
               Expressions.call(BuiltInMethod.TAKE_BIG_DECIMAL.method, v,
                   getExpression(fetch, "FETCH", implementor, builder,
-                      roundingPolicyExp, true)));
+                      roundingPolicyExp)));
     }
 
     builder.add(Expressions.return_(null, v));
@@ -123,7 +123,7 @@ public static EnumerableLimit create(final RelNode input, @Nullable RexNode offs
 
   static Expression getExpression(RexNode rexNode, String kind,
       EnumerableRelImplementor implementor, BlockBuilder builder,
-      Expression roundingPolicy, boolean translateExpression) {
+      Expression roundingPolicy) {
     final Expression value;
     if (rexNode instanceof RexDynamicParam) {
       final RexDynamicParam param = (RexDynamicParam) rexNode;
@@ -134,10 +134,6 @@ static Expression getExpression(RexNode rexNode, String kind,
     } 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())

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

@@ -105,15 +105,15 @@ public static EnumerableLimitSort create(
       fetchVal = Expressions.constant(BigDecimal.valueOf(Integer.MAX_VALUE));
     } else {
       fetchVal =
-          getExpression(this.fetch, "FETCH", implementor, builder, roundingPolicyExp, true);
+          getExpression(this.fetch, "FETCH", implementor, builder, roundingPolicyExp);
     }
 
     final Expression offsetVal;
     if (this.offset == null) {
       offsetVal = Expressions.constant(BigDecimal.ZERO);
     } else {
       offsetVal =
-          getExpression(this.offset, "OFFSET", implementor, builder, roundingPolicyExp, false);
+          getExpression(this.offset, "OFFSET", implementor, builder, roundingPolicyExp);
     }
 
     builder.add(
@@ -125,10 +125,10 @@ public static EnumerableLimitSort create(
                         builder.appendIfNotNull("comparator", pair.right))
                     .appendIfNotNull(
                         builder.appendIfNotNull("offset",
-                            Expressions.constant(offsetVal)))
+                            offsetVal))
                     .appendIfNotNull(
                         builder.appendIfNotNull("fetch",
-                            Expressions.constant(fetchVal))))));
+                            fetchVal)))));
     return implementor.result(physType, builder.toBlock());
   }
 }

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

@@ -27,7 +27,6 @@
 import org.apache.calcite.rel.logical.LogicalUnion;
 import org.apache.calcite.rel.type.RelDataType;
 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;
@@ -90,15 +89,16 @@ public EnumerableMergeUnionRule(Config config) {
     RexNode inputFetch = null;
     if (sort.fetch != null) {
       final boolean safeToReevaluate =
-          RexUtil.isDeterministic(sort.fetch);
-      if (sort.offset == null && safeToReevaluate) {
-        inputFetch = sort.fetch;
-      } else if (safeToReevaluate
-          && sort.fetch instanceof RexLiteral
-          && sort.offset instanceof RexLiteral) {
-        inputFetch =
-            call.builder().literal(RexLiteral.bigDecimalValue(sort.fetch)
-                .add(RexLiteral.bigDecimalValue(sort.offset)));
+          RexUtil.isDeterministic(sort.fetch)
+              && (sort.offset == null || RexUtil.isDeterministic(sort.offset));
+      if (safeToReevaluate) {
+        if (sort.offset == null) {
+          inputFetch = sort.fetch;
+        } else {
+          inputFetch =
+              RexUtil.makeOffsetFetchSum(
+                  sort.getCluster().getRexBuilder(), sort.offset, sort.fetch);
+        }
       }
     }
 

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

@@ -117,7 +117,8 @@ public Double getMinRowCount(Sort rel, RelMetadataQuery mq) {
     }
 
     final double offset =
-        literalValueApproximatedByDouble(rel.offset, 0D);
+        literalValueApproximatedByDouble(rel.offset,
+            rel.offset == null ? 0D : rowCount);
     rowCount = Math.max(rowCount - offset, 0D);
 
     final double limit =
@@ -133,7 +134,8 @@ public Double getMinRowCount(EnumerableLimit rel, RelMetadataQuery mq) {
     }
 
     final double offset =
-        literalValueApproximatedByDouble(rel.offset, 0D);
+        literalValueApproximatedByDouble(rel.offset,
+            rel.offset == null ? 0D : rowCount);
     rowCount = Math.max(rowCount - offset, 0D);
 
     final double limit =

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

@@ -29,7 +29,6 @@
 import org.apache.calcite.rel.core.Union;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
-import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexLocalRef;
@@ -1057,7 +1056,9 @@ private static boolean alreadySmaller(RelMetadataQuery mq, RelNode input,
       }
     }
     final Double rowCount = mq.getMaxRowCount(input);
-    if (rowCount == null || offset instanceof RexDynamicParam || !(fetch instanceof RexLiteral)) {
+    if (rowCount == null
+        || (offset != null && !(offset instanceof RexLiteral))
+        || !(fetch instanceof RexLiteral)) {
       // Cannot be determined
       return false;
     }

core/src/main/java/org/apache/calcite/rel/rel2sql/RelToSqlConverter.java

@@ -1188,7 +1188,7 @@ public Result visit(Sort e) {
         SqlNodeList sortExps = exprList(builder.context, e.getSortExps());
         sqlSelect.setOrderBy(sortExps);
         if (e.offset != null) {
-          SqlNode offset = builder.context.toSql(null, e.offset);
+          SqlNode offset = toSqlOffset(e, builder.context);
           sqlSelect.setOffset(offset);
         }
         if (e.fetch != null) {
@@ -1253,10 +1253,17 @@ void offsetFetch(Sort e, Builder builder) {
       builder.setFetch(toSqlFetch(e, builder.context));
     }
     if (e.offset != null) {
-      builder.setOffset(builder.context.toSql(null, e.offset));
+      builder.setOffset(toSqlOffset(e, builder.context));
     }
   }
 
+  private static SqlNode toSqlOffset(Sort sort, Context context) {
+    final RexNode offset = requireNonNull(sort.offset, "offset");
+    final @Nullable RexLiteral reduced =
+        RexUtil.reduceOffsetToLiteral(sort.getCluster(), offset);
+    return context.toSql(null, reduced == null ? offset : reduced);
+  }
+
   private static SqlNode toSqlFetch(Sort sort, Context context) {
     final RexNode fetch = requireNonNull(sort.fetch, "fetch");
     final @Nullable RexLiteral reduced =

core/src/main/java/org/apache/calcite/rel/rules/SortJoinTransposeRule.java

@@ -30,7 +30,6 @@
 import org.apache.calcite.rel.metadata.RelMdUtil;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rex.RexBuilder;
-import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.tools.RelBuilderFactory;
@@ -106,8 +105,8 @@ public SortJoinTransposeRule(Class<? extends Sort> sortClass,
     final Join join = call.rel(1);
 
     // The pushed fetch is calculated from literal offset and fetch values.
-    if (sort.offset instanceof RexDynamicParam
-        || sort.fetch != null && !(sort.fetch instanceof RexLiteral)) {
+    if ((sort.offset != null && !(sort.offset instanceof RexLiteral))
+        || (sort.fetch != null && !(sort.fetch instanceof RexLiteral))) {
       return false;
     }
 

core/src/main/java/org/apache/calcite/rel/rules/SortUnionTransposeRule.java

@@ -23,6 +23,7 @@
 import org.apache.calcite.rel.core.Union;
 import org.apache.calcite.rel.metadata.RelMdUtil;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.tools.RelBuilderFactory;
 
@@ -67,12 +68,14 @@ public SortUnionTransposeRule(
   @Override public boolean matches(RelOptRuleCall call) {
     final Sort sort = call.rel(0);
     final Union union = call.rel(1);
-    // Re-evaluating a non-deterministic FETCH in every branch can produce a
-    // different limit from the top Sort.
+    // Re-evaluating a non-deterministic OFFSET or FETCH in every branch can
+    // produce a different limit from the top Sort.
     // There is a flag indicating if this rule should be applied when
     // Sort.fetch is null.
     return union.all
-        && sort.offset == null
+        && (sort.offset == null
+            || sort.fetch != null
+                && RexUtil.isDeterministic(sort.offset))
         && (sort.fetch == null
             || RexUtil.isDeterministic(sort.fetch))
         && (config.matchNullFetch() || sort.fetch != null);
@@ -81,18 +84,29 @@ public SortUnionTransposeRule(
   @Override public void onMatch(RelOptRuleCall call) {
     final Sort sort = call.rel(0);
     final Union union = call.rel(1);
+    // OFFSET cannot be pushed into each input independently. However, only
+    // the first OFFSET + FETCH rows of an input can contribute to the final
+    // result, so use that value as the input FETCH and retain the original
+    // OFFSET and FETCH in the top Sort.
+    final RexNode inputFetch;
+    if (sort.fetch == null || sort.offset == null) {
+      inputFetch = sort.fetch;
+    } else {
+      inputFetch =
+          RexUtil.makeOffsetFetchSum(sort.getCluster().getRexBuilder(), sort.offset, sort.fetch);
+    }
     List<RelNode> inputs = new ArrayList<>();
     // Thus we use 'ret' as a flag to identify if we have finished pushing the
     // sort past a union.
     boolean ret = true;
     final RelMetadataQuery mq = call.getMetadataQuery();
     for (RelNode input : union.getInputs()) {
       if (!RelMdUtil.checkInputForCollationAndLimit(mq, input,
-          sort.getCollation(), sort.offset, sort.fetch)) {
+          sort.getCollation(), null, inputFetch)) {
         ret = false;
         Sort branchSort =
             sort.copy(sort.getTraitSet(), input,
-                sort.getCollation(), sort.offset, sort.fetch);
+                sort.getCollation(), null, inputFetch);
         inputs.add(branchSort);
       } else {
         inputs.add(input);