IGNITE-28739 SQL Calcite: Add expression support for FETCH

modules/calcite/pom.xml

@@ -351,6 +351,14 @@
                                     outputRoot="${project.build.directory}/generated-sources/fmpp"
                                     data="tdd(${project.build.directory}/codegen/config.fmpp), default: tdd(${project.build.directory}/codegen/default_config.fmpp)"
                                 />
+                                <!-- TODO: https://issues.apache.org/jira/browse/CALCITE-7592
+                                     Remove this workaround after upgrading to Calcite 1.43. -->
+                                <replace
+                                    file="${project.build.directory}/generated-sources/fmpp/javacc/Parser.jj"
+                                    token="&lt;FETCH&gt; ( &lt;FIRST&gt; | &lt;NEXT&gt; ) offsetFetch[1] = UnsignedNumericLiteralOrParam()"
+                                    value="&lt;FETCH&gt; ( &lt;FIRST&gt; | &lt;NEXT&gt; ) offsetFetch[1] = FetchCount()"
+                                    failOnNoReplacements="true"
+                                />
                             </target>
                         </configuration>
                     </execution>

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

@@ -931,3 +931,19 @@ SqlNode SqlSelectForUpdate() :
     ]
     { return query; }
 }
+
+// TODO: https://issues.apache.org/jira/browse/CALCITE-7592
+//  Remove this method and the corresponding replacement in pom.xml after upgrading to Calcite 1.43.
+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/java/org/apache/ignite/internal/processors/query/calcite/exec/LogicalRelImplementor.java

@@ -41,6 +41,7 @@
 import org.apache.calcite.rel.core.Window;
 import org.apache.calcite.rel.type.RelDataType;
 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;
@@ -655,6 +656,10 @@ private boolean hasExchange(RelNode rel) {
         long offset = validateAndGetOffset(rel.offset, SortNode.OFFSET_DEFAULT);
         long fetch = validateAndGetFetch(rel.fetch, SortNode.FETCH_DEFAULT);
 
+        // Zero FETCH is enforced by the outer IgniteLimit, while SortNode accepts only positive FETCH values.
+        if (fetch == 0)
+            fetch = SortNode.FETCH_DEFAULT;
+
         SortNode<Row> node = new SortNode<>(ctx, rel.getRowType(), expressionFactory.comparator(collation), offset,
             fetch);
 
@@ -672,7 +677,7 @@ private long validateAndGetOffset(RexNode node, long defaultVal) {
 
     /** */
     private long validateAndGetFetch(RexNode node, long defaultVal) {
-        return node == null ? defaultVal : validateAndGetFetchOffsetParams(node, "fetch");
+        return node == null ? defaultVal : validateAndGetFetchOffsetParams(node, "fetch / limit");
     }
 
     /** {@inheritDoc} */
@@ -1072,6 +1077,11 @@ private long validateAndGetFetchOffsetParams(RexNode node, String op) {
         Supplier<Object> scalar = expressionFactory.execute(node);
         Object param = scalar.get();
 
+        if (param == null && !(node instanceof RexDynamicParam)) {
+            throw new IgniteSQLException(IgniteResource.INSTANCE.illegalFetchLimit(op).str(),
+                IgniteQueryErrorCode.UNEXPECTED_ELEMENT_TYPE);
+        }
+
         if (!(param instanceof Number)) {
             String actual = param == null ? "null" : param.getClass().getSimpleName();
             throw new IgniteSQLException(IgniteResource.INSTANCE.incorrectDynamicParameterType("BIGINT", actual).str(),

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/metadata/IgniteMdMinRowCount.java

@@ -0,0 +1,56 @@
+/*
+ * 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.metadata;
+
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.metadata.ReflectiveRelMetadataProvider;
+import org.apache.calcite.rel.metadata.RelMdMinRowCount;
+import org.apache.calcite.rel.metadata.RelMetadataProvider;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.util.BuiltInMethod;
+
+import static org.apache.calcite.rel.metadata.RelMdUtil.literalValueApproximatedByDouble;
+
+/** Minimum row count metadata compatible with expression-based FETCH and OFFSET. */
+// TODO: https://issues.apache.org/jira/browse/CALCITE-7592
+//  Remove this class and its registration in IgniteMetadata after upgrading to Calcite 1.43.
+@SuppressWarnings("unused") // Actually all methods are used by runtime generated classes.
+public class IgniteMdMinRowCount extends RelMdMinRowCount {
+    /** Metadata provider. */
+    public static final RelMetadataProvider SOURCE =
+        ReflectiveRelMetadataProvider.reflectiveSource(
+            BuiltInMethod.MIN_ROW_COUNT.method, new IgniteMdMinRowCount());
+
+    /** {@inheritDoc} */
+    @Override public Double getMinRowCount(Sort rel, RelMetadataQuery mq) {
+        Double rowCnt = mq.getMinRowCount(rel.getInput());
+
+        if (rowCnt == null)
+            rowCnt = 0D;
+
+        double offset = literalValueApproximatedByDouble(rel.offset,
+            rel.offset == null ? 0D : rowCnt);
+
+        rowCnt = Math.max(rowCnt - offset, 0D);
+
+        double limit = literalValueApproximatedByDouble(rel.fetch,
+            rel.fetch == null ? rowCnt : 0D);
+
+        return limit < rowCnt ? limit : rowCnt;
+    }
+}

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/metadata/IgniteMetadata.java

@@ -46,6 +46,7 @@ public class IgniteMetadata {
                 IgniteMdCumulativeCost.SOURCE,
                 IgniteMdNonCumulativeCost.SOURCE,
                 IgniteMdRowCount.SOURCE,
+                IgniteMdMinRowCount.SOURCE,
                 IgniteMdPredicates.SOURCE,
                 IgniteMdCollation.SOURCE,
                 IgniteMdSelectivity.SOURCE,

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

@@ -0,0 +1,94 @@
+/*
+ * 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. */
+// TODO: https://issues.apache.org/jira/browse/CALCITE-7592
+//  Remove this class and IgniteSqlToRelConvertor.newFieldTrimmer() after upgrading to Calcite 1.43.
+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));
+
+        RelCollation newCollation = RexUtil.apply(inputMapping, collation);
+
+        RelNode newSort = sort.copy(
+            sort.getTraitSet().replace(newCollation),
+            newInput,
+            newCollation,
+            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);

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

@@ -48,11 +48,13 @@
 import org.apache.calcite.sql.SqlNodeList;
 import org.apache.calcite.sql.SqlNumericLiteral;
 import org.apache.calcite.sql.SqlOperatorTable;
+import org.apache.calcite.sql.SqlOrderBy;
 import org.apache.calcite.sql.SqlSelect;
 import org.apache.calcite.sql.SqlUpdate;
 import org.apache.calcite.sql.SqlUtil;
 import org.apache.calcite.sql.SqlWindow;
 import org.apache.calcite.sql.dialect.CalciteSqlDialect;
+import org.apache.calcite.sql.fun.SqlCase;
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
 import org.apache.calcite.sql.type.SqlOperandTypeChecker;
@@ -267,10 +269,58 @@ private void validateTableModify(SqlNode table) {
     @Override protected void validateSelect(SqlSelect select, RelDataType targetRowType) {
         super.validateSelect(select, targetRowType);
 
-        validateFetchOffset(select.getFetch(), "fetch / limit");
+        validateFetch(select, "fetch / limit");
         validateFetchOffset(select.getOffset(), "offset");
     }
 
+    /** Validate fetch expression. */
+    // TODO: https://issues.apache.org/jira/browse/CALCITE-7592
+    //  Remove this method after upgrading to Calcite 1.43.
+    private void validateFetch(SqlSelect select, String clauseName) {
+        SqlNode fetch = select.getFetch();
+
+        if (fetch == null)
+            return;
+
+        if (SqlUtil.isNullLiteral(fetch, true))
+            throw newValidationError(fetch, IgniteResource.INSTANCE.illegalFetchLimit(clauseName));
+
+        validateFetchExpression(fetch, clauseName);
+
+        SqlValidatorScope scope = getEmptyScope();
+
+        inferUnknownTypes(typeFactory().createSqlType(SqlTypeName.DECIMAL), scope, fetch);
+        RelDataType type = deriveType(scope, fetch);
+
+        if (type.getSqlTypeName().getFamily() != SqlTypeFamily.NUMERIC && !(fetch instanceof SqlDynamicParam))
+            throw newValidationError(fetch, IgniteResource.INSTANCE.illegalFetchLimit(clauseName));
+
+        validateFetchOffset(fetch, clauseName);
+    }
+
+    /** Reject column references in a fetch expression. */
+    // TODO: https://issues.apache.org/jira/browse/CALCITE-7592
+    //  Remove this method after upgrading to Calcite 1.43.
+    private void validateFetchExpression(SqlNode node, String clauseName) {
+        if (node instanceof SqlIdentifier) {
+            if (makeNullaryCall((SqlIdentifier)node) == null)
+                throw newValidationError(node, IgniteResource.INSTANCE.illegalFetchLimit(clauseName));
+
+            return;
+        }
+
+        if (node instanceof SqlNodeList) {
+            for (SqlNode child : (SqlNodeList)node)
+                validateFetchExpression(child, clauseName);
+        }
+        else if (node instanceof SqlCall call) {
+            for (SqlNode child : call.getOperandList()) {
+                if (child != null)
+                    validateFetchExpression(child, clauseName);
+            }
+        }
+    }
+
     /**
      * Validate fetch/offset params restrictions.
      *
@@ -282,7 +332,12 @@ private void validateFetchOffset(@Nullable SqlNode n, String clauseName) {
             return;
 
         if (n instanceof SqlLiteral) {
-            BigDecimal offsetFetchLimit = ((SqlLiteral)n).bigDecimalValue();
+            SqlLiteral literal = (SqlLiteral)n;
+
+            if (literal.getTypeName().getFamily() != SqlTypeFamily.NUMERIC)
+                throw newValidationError(n, IgniteResource.INSTANCE.illegalFetchLimit(clauseName));
+
+            BigDecimal offsetFetchLimit = literal.bigDecimalValue();
 
             checkLimitOffset(offsetFetchLimit, n, clauseName);
         }
@@ -408,6 +463,21 @@ else if (call.getKind() == SqlKind.OVER && call.operand(1) instanceof SqlWindow)
 
     /** {@inheritDoc} */
     @Override protected SqlNode performUnconditionalRewrites(SqlNode node, boolean underFrom) {
+        if (node instanceof SqlOrderBy) {
+            SqlOrderBy orderBy = (SqlOrderBy)node;
+
+            if (orderBy.fetch != null) {
+                // SqlOrderBy does not implement setOperand(). Rewrite FETCH before Calcite visits its operands,
+                // otherwise rewrites such as COALESCE to CASE fail when Calcite tries to replace the FETCH operand.
+                SqlNode fetch = performUnconditionalRewrites(orderBy.fetch, false);
+
+                if (fetch != orderBy.fetch) {
+                    node = new SqlOrderBy(orderBy.getParserPosition(), orderBy.query, orderBy.orderList,
+                        orderBy.offset, fetch);
+                }
+            }
+        }
+
         // Workaround for https://issues.apache.org/jira/browse/CALCITE-4923
         if (node instanceof SqlSelect) {
             SqlSelect select = (SqlSelect)node;
@@ -688,7 +758,11 @@ private void validateVersionColumnDmlTarget(SqlIdentifier id) {
             && deriveDynamicParameterType((SqlDynamicParam)node, unknownType.equals(inferredType) ? nullType : inferredType) != null)
             return;
 
-        if (node instanceof SqlCall) {
+        if (node instanceof SqlCase) {
+            // SqlValidatorImpl assigns context-specific types to WHEN and result operands.
+            super.inferUnknownTypes(inferredType, scope, node);
+        }
+        else if (node instanceof SqlCall) {
             final SqlValidatorScope newScope = scopes.get(node);
 
             if (newScope != null)

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

@@ -43,6 +43,7 @@
 import org.apache.calcite.rel.rules.PruneEmptyRules;
 import org.apache.calcite.rel.rules.SetOpToFilterRule;
 import org.apache.calcite.rel.rules.SortRemoveRule;
+import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.tools.Program;
 import org.apache.calcite.tools.RuleSet;
 import org.apache.calcite.tools.RuleSets;
@@ -281,7 +282,9 @@ public enum PlannerPhase {
 
                     ((RelRule<?>)PruneEmptyRules.SORT_FETCH_ZERO_INSTANCE).config
                         .withOperandSupplier(b ->
-                            b.operand(LogicalSort.class).anyInputs())
+                            b.operand(LogicalSort.class)
+                                .predicate(sort -> sort.fetch instanceof RexLiteral)
+                                .anyInputs())
                         .toRule(),
 
                     ExposeIndexRule.INSTANCE,