IGNITE-28739 Add expression support for FETCH in the Calcite SQL engine

modules/calcite/pom.xml

@@ -341,6 +341,11 @@
                                     outputRoot="${project.build.directory}/generated-sources/fmpp"
                                     data="tdd(${project.build.directory}/codegen/config.fmpp), default: tdd(${project.build.directory}/codegen/default_config.fmpp)"
                                 />
+                                <patch
+                                    dir="${project.build.directory}/generated-sources/fmpp"
+                                    patchfile="${project.build.directory}/codegen/patches/select-fetch-expression.patch"
+                                    strip="1"
+                                />
                             </target>
                         </configuration>
                     </execution>

modules/calcite/src/main/codegen/includes/parserImpls.ftl

@@ -816,3 +816,17 @@ SqlDrop SqlDropView(Span s, boolean replace) :
         return SqlDdlNodes.dropView(s.end(this), ifExists, id);
     }
 }
+
+JAVACODE
+SqlNode FetchCount() {
+    SqlNode e;
+    if (getToken(1).kind == LPAREN) {
+        jj_consume_token(LPAREN);
+        e = Expression(ExprContext.ACCEPT_NON_QUERY);
+        jj_consume_token(RPAREN);
+    }
+    else
+        e = UnsignedNumericLiteralOrParam();
+
+    return e;
+}

modules/calcite/src/main/codegen/patches/select-fetch-expression.patch

@@ -0,0 +1,26 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+--- a/javacc/Parser.jj
++++ b/javacc/Parser.jj
+@@ -754,7 +754,7 @@
+ {
+     // 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>
+ }

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/LogicalRelImplementor.java

@@ -627,7 +627,7 @@ private boolean hasExchange(RelNode rel) {
     /** {@inheritDoc} */
     @Override public Node<Row> visit(IgniteLimit rel) {
         Supplier<Integer> offset = (rel.offset() == null) ? null : expressionFactory.execute(rel.offset());
-        Supplier<Integer> fetch = (rel.fetch() == null) ? null : expressionFactory.execute(rel.fetch());
+        Supplier<Number> fetch = (rel.fetch() == null) ? null : expressionFactory.execute(rel.fetch());
 
         LimitNode<Row> node = new LimitNode<>(ctx, rel.getRowType(), offset, fetch);
 
@@ -643,7 +643,7 @@ private boolean hasExchange(RelNode rel) {
         RelCollation collation = rel.getCollation();
 
         Supplier<Integer> offset = (rel.offset == null) ? null : expressionFactory.execute(rel.offset);
-        Supplier<Integer> fetch = (rel.fetch == null) ? null : expressionFactory.execute(rel.fetch);
+        Supplier<Number> fetch = (rel.fetch == null) ? null : expressionFactory.execute(rel.fetch);
 
         SortNode<Row> node = new SortNode<>(ctx, rel.getRowType(), expressionFactory.comparator(collation), offset,
             fetch);

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/LimitNode.java

@@ -20,6 +20,7 @@
 import java.util.function.Supplier;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
+import org.apache.ignite.internal.processors.query.calcite.util.RelNodeUtils;
 import org.apache.ignite.internal.util.typedef.F;
 import org.jetbrains.annotations.Nullable;
 
@@ -35,27 +36,22 @@ public class LimitNode<Row> extends AbstractNode<Row> implements SingleNode<Row>
     private int rowsProcessed;
 
     /** Fetch can be unset, in this case we need all rows. */
-    private @Nullable Supplier<Integer> fetchNode;
+    private final @Nullable Supplier<Number> fetchNode;
 
     /** Waiting results counter. */
     private int waiting;
 
-    /**
-     * Constructor.
-     *
-     * @param ctx Execution context.
-     * @param rowType Row type.
-     */
+    /** */
     public LimitNode(
         ExecutionContext<Row> ctx,
         RelDataType rowType,
-        Supplier<Integer> offsetNode,
-        Supplier<Integer> fetchNode
+        @Nullable Supplier<Integer> offsetNode,
+        @Nullable Supplier<Number> fetchNode
     ) {
         super(ctx, rowType);
 
         offset = offsetNode == null ? 0 : offsetNode.get();
-        fetch = fetchNode == null ? 0 : fetchNode.get();
+        fetch = RelNodeUtils.resolveFetch(fetchNode, "FETCH");
         this.fetchNode = fetchNode;
     }
 

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/SortNode.java

@@ -23,6 +23,7 @@
 import java.util.function.Supplier;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
+import org.apache.ignite.internal.processors.query.calcite.util.RelNodeUtils;
 import org.apache.ignite.internal.util.GridBoundedPriorityQueue;
 import org.apache.ignite.internal.util.typedef.F;
 import org.jetbrains.annotations.Nullable;
@@ -49,24 +50,19 @@ public class SortNode<Row> extends MemoryTrackingNode<Row> implements SingleNode
     /** Reverse-ordered rows in case of limited sort. */
     private List<Row> reversed;
 
-    /**
-     * @param ctx Execution context.
-     * @param comp Rows comparator.
-     * @param offset Offset.
-     * @param fetch Limit.
-     */
+    /** */
     public SortNode(
         ExecutionContext<Row> ctx, RelDataType rowType,
         Comparator<Row> comp,
         @Nullable Supplier<Integer> offset,
-        @Nullable Supplier<Integer> fetch
+        @Nullable Supplier<Number> fetch
     ) {
         super(ctx, rowType);
 
-        assert fetch == null || fetch.get() >= 0;
         assert offset == null || offset.get() >= 0;
 
-        limit = fetch == null ? -1 : fetch.get() + (offset == null ? 0 : offset.get());
+        limit = fetch == null ?
+            -1 : RelNodeUtils.resolveFetch(fetch, "FETCH") + (offset == null ? 0 : offset.get());
 
         if (limit < 0)
             rows = new PriorityQueue<>(comp);

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/prepare/IgnitePlanner.java

@@ -343,7 +343,11 @@ private static boolean isAsCall(SqlNode node) {
 
     /** {@inheritDoc} */
     @Override public RelRoot rel(SqlNode sql) {
-        SqlToRelConverter sqlToRelConverter = sqlToRelConverter(validator(), catalogReader, sqlToRelConverterCfg);
+        IgniteSqlValidator validator = (IgniteSqlValidator)validator();
+
+        validator.deriveLimitDynamicParameterTypes(sql);
+
+        SqlToRelConverter sqlToRelConverter = sqlToRelConverter(validator, catalogReader, sqlToRelConverterCfg);
 
         return sqlToRelConverter.convertQuery(sql, false, true);
     }

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/prepare/IgniteRelFieldTrimmer.java

@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.processors.query.calcite.prepare;
+
+import java.util.Set;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.type.RelDataTypeField;
+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.sql.validate.SqlValidator;
+import org.apache.calcite.sql2rel.RelFieldTrimmer;
+import org.apache.calcite.tools.RelBuilder;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.mapping.Mapping;
+import org.apache.calcite.util.mapping.Mappings;
+import org.jetbrains.annotations.Nullable;
+
+import static java.util.Collections.emptySet;
+
+/** Field trimmer that preserves expression-based FETCH nodes. */
+public class IgniteRelFieldTrimmer extends RelFieldTrimmer {
+    /**  */
+    IgniteRelFieldTrimmer(@Nullable SqlValidator validator, RelBuilder relBuilder) {
+        super(validator, relBuilder);
+    }
+
+    /** {@inheritDoc} */
+    @Override public TrimResult trimFields(
+        Sort sort,
+        ImmutableBitSet fieldsUsed,
+        Set<RelDataTypeField> extraFields
+    ) {
+        if (supportedByRelBuilder(sort.fetch))
+            return super.trimFields(sort, fieldsUsed, extraFields);
+
+        RelCollation collation = sort.getCollation();
+        RelNode input = sort.getInput();
+        int fieldCnt = sort.getRowType().getFieldCount();
+
+        ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
+
+        for (RelFieldCollation field : collation.getFieldCollations())
+            inputFieldsUsed.set(field.getFieldIndex());
+
+        TrimResult trimRes = trimChild(sort, input, inputFieldsUsed.build(), emptySet());
+        RelNode newInput = trimRes.left;
+        Mapping inputMapping = trimRes.right;
+
+        if (newInput == input && inputMapping.isIdentity() && fieldsUsed.cardinality() == fieldCnt)
+            return result(sort, Mappings.createIdentity(fieldCnt));
+
+        RelNode newSort = sort.copy(
+            sort.getTraitSet(),
+            newInput,
+            RexUtil.apply(inputMapping, collation),
+            sort.offset,
+            sort.fetch
+        );
+
+        return result(newSort, inputMapping, sort);
+    }
+
+    /**
+     * @param node Rex node.
+     * @return {@code true} if Calcite RelBuilder accepts the node for FETCH.
+     */
+    private static boolean supportedByRelBuilder(@Nullable RexNode node) {
+        return node == null || node instanceof RexLiteral || node instanceof RexDynamicParam;
+    }
+}

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/prepare/IgniteSqlToRelConvertor.java

@@ -46,6 +46,7 @@
 import org.apache.calcite.sql.validate.SqlValidator;
 import org.apache.calcite.sql.validate.SqlValidatorScope;
 import org.apache.calcite.sql.validate.SqlValidatorUtil;
+import org.apache.calcite.sql2rel.RelFieldTrimmer;
 import org.apache.calcite.sql2rel.SqlRexConvertletTable;
 import org.apache.calcite.sql2rel.SqlToRelConverter;
 import org.apache.calcite.tools.RelBuilder;
@@ -87,6 +88,11 @@ public IgniteSqlToRelConvertor(
             return super.convertQueryRecursive(qry, top, targetRowType);
     }
 
+    /** {@inheritDoc} */
+    @Override protected RelFieldTrimmer newFieldTrimmer() {
+        return new IgniteRelFieldTrimmer(validator, relBuilder);
+    }
+
     /** {@inheritDoc} */
     @Override protected RelNode convertInsert(SqlInsert call) {
         datasetStack.push(call);