IGNITE-28896 SQL Calcite: Support BigDecimal for FETCH, OFFSET and LIMIT

modules/calcite/src/main/java/org/apache/calcite/adapter/enumerable/FetchOffsetRoundingPolicy.java

@@ -0,0 +1,36 @@
+/*
+ * 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.calcite.adapter.enumerable;
+
+import java.math.BigDecimal;
+
+/** Defines how enumerable FETCH and OFFSET values are rounded. */
+// TODO https://issues.apache.org/jira/browse/CALCITE-7624
+//  Check and remove after update to Calcite 1.43.
+public interface FetchOffsetRoundingPolicy {
+    /** Default policy that preserves the value produced by validation or parameter binding. */
+    FetchOffsetRoundingPolicy NONE = value -> value;
+
+    /**
+     * Rounds a FETCH or OFFSET value.
+     *
+     * @param value Value to round.
+     * @return Rounded value.
+     */
+    BigDecimal round(BigDecimal value);
+}

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

@@ -17,6 +17,8 @@
 
 package org.apache.ignite.internal.processors.query.calcite.exec;
 
+import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
@@ -28,6 +30,7 @@
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import com.google.common.collect.ImmutableList;
+import org.apache.calcite.adapter.enumerable.FetchOffsetRoundingPolicy;
 import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelNode;
@@ -126,6 +129,7 @@
 import org.apache.ignite.internal.processors.query.calcite.trait.TraitUtils;
 import org.apache.ignite.internal.processors.query.calcite.type.IgniteTypeFactory;
 import org.apache.ignite.internal.processors.query.calcite.util.Commons;
+import org.apache.ignite.internal.processors.query.calcite.util.IgniteMath;
 import org.apache.ignite.internal.processors.query.calcite.util.RexUtils;
 import org.apache.ignite.internal.util.typedef.F;
 import org.jetbrains.annotations.Nullable;
@@ -144,6 +148,9 @@ public class LogicalRelImplementor<Row> implements IgniteRelVisitor<Node<Row>> {
     /** */
     private final ExecutionContext<Row> ctx;
 
+    /** FETCH/OFFSET rounding policy used by the Calcite engine. */
+    private final FetchOffsetRoundingPolicy fetchOffsetRoundingPolicy;
+
     /** */
     private final AffinityService affSrvc;
 
@@ -176,6 +183,7 @@ public LogicalRelImplementor(
         this.ctx = ctx;
 
         expressionFactory = ctx.expressionFactory();
+        fetchOffsetRoundingPolicy = getFetchOffsetRoundingPolicy(ctx);
     }
 
     /** {@inheritDoc} */
@@ -568,7 +576,7 @@ private boolean hasExchange(RelNode rel) {
                 rowType,
                 idxBndRel.first() ? cmp : cmp.reversed(),
                 null,
-                () -> 1
+                () -> BigDecimal.ONE
             );
 
             sortNode.register(scanNode);
@@ -630,8 +638,8 @@ 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<BigDecimal> offset = fetchOffsetSupplier(rel.offset(), "OFFSET");
+        Supplier<BigDecimal> fetch = fetchOffsetSupplier(rel.fetch(), "FETCH");
 
         LimitNode<Row> node = new LimitNode<>(ctx, rel.getRowType(), offset, fetch);
 
@@ -646,8 +654,8 @@ private boolean hasExchange(RelNode rel) {
     @Override public Node<Row> visit(IgniteSort 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<BigDecimal> offset = fetchOffsetSupplier(rel.offset, "OFFSET");
+        Supplier<BigDecimal> fetch = fetchOffsetSupplier(rel.fetch, "FETCH");
 
         SortNode<Row> node = new SortNode<>(ctx, rel.getRowType(), expressionFactory.comparator(collation), offset,
             fetch);
@@ -1050,4 +1058,34 @@ private ScanStorageNode<Row> createStorageScan(
             otherColMapping
         );
     }
+
+    /** */
+    private static FetchOffsetRoundingPolicy getFetchOffsetRoundingPolicy(ExecutionContext<?> ctx) {
+        FetchOffsetRoundingPolicy roundingPlc = ctx.unwrap(FetchOffsetRoundingPolicy.class);
+
+        return roundingPlc == null ? FetchOffsetRoundingPolicy.NONE : roundingPlc;
+    }
+
+    /** Converts a FETCH/OFFSET expression to the row count expected by Ignite execution nodes. */
+    private @Nullable Supplier<BigDecimal> fetchOffsetSupplier(@Nullable RexNode node, String kind) {
+        if (node == null)
+            return null;
+
+        Supplier<Number> val = expressionFactory.execute(node);
+
+        return () -> fetchOffsetValue(val.get(), kind);
+    }
+
+    /** Converts a FETCH/OFFSET runtime value to a row count. */
+    private BigDecimal fetchOffsetValue(Object val, String kind) {
+        if (!(val instanceof Number))
+            throw new IllegalArgumentException(kind + " must be a number");
+
+        BigDecimal decimal = IgniteMath.convertToBigDecimal((Number)val);
+
+        if (decimal.signum() < 0)
+            throw new IllegalArgumentException(kind + " must not be negative");
+
+        return fetchOffsetRoundingPolicy.round(decimal).setScale(0, RoundingMode.CEILING);
+    }
 }

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

@@ -17,6 +17,7 @@
 
 package org.apache.ignite.internal.processors.query.calcite.exec.rel;
 
+import java.math.BigDecimal;
 import java.util.function.Supplier;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
@@ -26,19 +27,16 @@
 /** Offset, fetch|limit support node. */
 public class LimitNode<Row> extends AbstractNode<Row> implements SingleNode<Row>, Downstream<Row> {
     /** Offset if its present, otherwise 0. */
-    private final int offset;
+    private final BigDecimal offset;
 
-    /** Fetch if its present, otherwise 0. */
-    private final int fetch;
+    /** Fetch if its present, otherwise all rows. */
+    private final @Nullable BigDecimal fetch;
 
     /** Already processed (pushed to upstream) rows count. */
-    private int rowsProcessed;
+    private BigDecimal rowsProcessed = BigDecimal.ZERO;
 
-    /** Fetch can be unset, in this case we need all rows. */
-    private @Nullable Supplier<Integer> fetchNode;
-
-    /** Waiting results counter. */
-    private int waiting;
+    /** Number of upstream rows that remain to be processed for the current downstream request. */
+    private BigDecimal waiting = BigDecimal.ZERO;
 
     /**
      * Constructor.
@@ -49,75 +47,76 @@ public class LimitNode<Row> extends AbstractNode<Row> implements SingleNode<Row>
     public LimitNode(
         ExecutionContext<Row> ctx,
         RelDataType rowType,
-        Supplier<Integer> offsetNode,
-        Supplier<Integer> fetchNode
+        @Nullable Supplier<BigDecimal> offsetNode,
+        @Nullable Supplier<BigDecimal> fetchNode
     ) {
         super(ctx, rowType);
 
-        offset = offsetNode == null ? 0 : offsetNode.get();
-        fetch = fetchNode == null ? 0 : fetchNode.get();
-        this.fetchNode = fetchNode;
+        offset = offsetNode == null ? BigDecimal.ZERO : offsetNode.get();
+        fetch = fetchNode == null ? null : fetchNode.get();
     }
 
     /** {@inheritDoc} */
     @Override public void request(int rowsCnt) throws Exception {
         assert !F.isEmpty(sources()) && sources().size() == 1;
-        assert rowsCnt > 0;
+        assert rowsCnt > 0 : rowsCnt;
 
         if (fetchNone()) {
             end();
 
             return;
         }
 
-        if (offset > 0 && rowsProcessed == 0)
-            rowsCnt = offset + rowsCnt;
-
-        waiting = rowsCnt;
+        waiting = BigDecimal.valueOf(rowsCnt);
 
-        if (fetch > 0)
-            rowsCnt = Math.min(rowsCnt, (fetch + offset) - rowsProcessed);
+        if (rowsProcessed.compareTo(offset) < 0)
+            waiting = waiting.add(offset.subtract(rowsProcessed));
 
-        checkState();
-
-        source().request(rowsCnt);
+        requestNextBatch();
     }
 
     /** {@inheritDoc} */
     @Override public void push(Row row) throws Exception {
-        if (waiting == -1)
+        if (waiting.signum() < 0)
             return;
 
-        ++rowsProcessed;
+        rowsProcessed = rowsProcessed.add(BigDecimal.ONE);
 
-        --waiting;
+        waiting = waiting.subtract(BigDecimal.ONE);
 
         checkState();
 
-        if (rowsProcessed > offset) {
-            if (fetchNode == null || (fetchNode != null && rowsProcessed <= fetch + offset))
-                downstream().push(row);
+        boolean endAfterRow = fetchNone() && waiting.signum() > 0;
+        boolean reqNextAfterRow = !endAfterRow && waiting.signum() > 0
+            && rowsProcessed.remainder(BigDecimal.valueOf(IN_BUFFER_SIZE)).signum() == 0;
+
+        if (rowsProcessed.compareTo(offset) > 0
+            && (fetch == null || rowsProcessed.compareTo(offset.add(fetch)) <= 0)) {
+            downstream().push(row);
         }
 
-        if (fetch > 0 && rowsProcessed == fetch + offset && waiting > 0)
+        if (endAfterRow)
             end();
+        else if (reqNextAfterRow)
+            requestNextBatch();
     }
 
     /** {@inheritDoc} */
     @Override public void end() throws Exception {
-        if (waiting == -1)
+        if (waiting.signum() < 0)
             return;
 
         assert downstream() != null;
 
-        waiting = -1;
+        waiting = BigDecimal.ONE.negate();
 
         downstream().end();
     }
 
     /** {@inheritDoc} */
     @Override protected void rewindInternal() {
-        rowsProcessed = 0;
+        rowsProcessed = BigDecimal.ZERO;
+        waiting = BigDecimal.ZERO;
     }
 
     /** {@inheritDoc} */
@@ -130,6 +129,27 @@ public LimitNode(
 
     /** {@code True} if requested 0 results, or all already processed. */
     private boolean fetchNone() {
-        return (fetchNode != null && fetch == 0) || (fetch > 0 && rowsProcessed == fetch + offset);
+        return fetch != null && (fetch.signum() == 0 || rowsProcessed.compareTo(offset.add(fetch)) >= 0);
+    }
+
+    /** Requests the next upstream batch, taking large decimal offsets into account. */
+    private void requestNextBatch() throws Exception {
+        BigDecimal bufSize = BigDecimal.valueOf(IN_BUFFER_SIZE);
+        BigDecimal rowsAvailable = waiting;
+
+        if (fetch != null)
+            rowsAvailable = rowsAvailable.min(offset.add(fetch).subtract(rowsProcessed));
+
+        BigDecimal rowsToReq = bufSize.subtract(rowsProcessed.remainder(bufSize)).min(rowsAvailable);
+
+        if (rowsToReq.signum() == 0) {
+            end();
+
+            return;
+        }
+
+        checkState();
+
+        source().request(rowsToReq.intValueExact());
     }
 }