IGNITE-28925 SQL Calcite: Fix unexpected downstream().end() call for execution nodes

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

@@ -38,7 +38,7 @@
  */
 public abstract class AbstractNode<Row> implements Node<Row> {
     /** */
-    protected static final int IN_BUFFER_SIZE = IgniteSystemProperties.getInteger(IGNITE_CALCITE_EXEC_IN_BUFFER_SIZE, 512);
+    public static final int IN_BUFFER_SIZE = IgniteSystemProperties.getInteger(IGNITE_CALCITE_EXEC_IN_BUFFER_SIZE, 512);
 
     /** */
     protected static final int MODIFY_BATCH_SIZE = IgniteSystemProperties.getInteger(IGNITE_CALCITE_EXEC_MODIFY_BATCH_SIZE, 100);

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

@@ -114,6 +114,11 @@ public static <Row> CollectNode<Row> createCountCollector(ExecutionContext<Row>
 
         if (waiting == 0)
             source().request(waiting = IN_BUFFER_SIZE);
+        else if (waiting < 0) {
+            requested = 0;
+
+            downstream().end();
+        }
     }
 
     /** {@inheritDoc} */
@@ -138,6 +143,7 @@ public static <Row> CollectNode<Row> createCountCollector(ExecutionContext<Row>
     @Override public void end() throws Exception {
         assert downstream() != null;
         assert waiting > 0;
+        assert requested > 0;
 
         checkState();
 
@@ -146,10 +152,13 @@ public static <Row> CollectNode<Row> createCountCollector(ExecutionContext<Row>
         if (isClosed())
             return;
 
+        requested--;
+
+        downstream().push(collector.get());
+
         if (requested > 0) {
             requested = 0;
 
-            downstream().push(collector.get());
             downstream().end();
         }
     }

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

@@ -282,6 +282,8 @@ private void onRequest() throws Exception {
                 break;
 
             case END:
+                requested = 0;
+
                 downstream().end();
                 break;
 
@@ -338,8 +340,11 @@ private void onEndLeft() throws Exception {
 
             state = State.END;
 
-            if (requested > 0)
+            if (requested > 0) {
+                requested = 0;
+
                 downstream().end();
+            }
         }
         else {
             prepareCorrelations();
@@ -469,8 +474,9 @@ private void join() throws Exception {
 
                 state = State.END;
 
-                if (requested > 0)
-                    downstream().end();
+                requested = 0;
+
+                downstream().end();
 
                 return;
             }

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

@@ -101,8 +101,11 @@ public SortAggregateNode(
 
             source().request(IN_BUFFER_SIZE);
         }
-        else if (waiting < 0)
+        else if (waiting < 0) {
+            requested = 0;
+
             downstream().end();
+        }
     }
 
     /** {@inheritDoc} */
@@ -163,8 +166,11 @@ else if (waiting < 0)
             doPush();
         }
 
-        if (requested > 0)
+        if (requested > 0) {
+            requested = 0;
+
             downstream().end();
+        }
 
         grp = null;
         prevRow = null;

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

@@ -205,10 +205,11 @@ private void flush() throws Exception {
             }
 
             if (reversed == null ? rows.isEmpty() : reversed.isEmpty()) {
-                if (requested > 0)
-                    downstream().end();
+                if (requested > 0) {
+                    requested = 0;
 
-                requested = 0;
+                    downstream().end();
+                }
             }
         }
         finally {

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/CollectIntegrationTest.java

@@ -0,0 +1,52 @@
+/*
+ * 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.integration;
+
+import org.apache.ignite.internal.processors.query.calcite.exec.rel.AbstractNode;
+import org.junit.Test;
+
+/**
+ * Integration test for collect node.
+ */
+public class CollectIntegrationTest extends AbstractBasicIntegrationTest {
+    /**
+     * Tests that collect node correctly handles the case when downstream requests
+     * limited number of rows, where collect must push one row and then
+     * properly terminate downstream.
+     */
+    @Test
+    public void testRequestLimitedRowsCountFromCollect() {
+        sql("CREATE TABLE t(a INT)");
+
+        sql("INSERT INTO t (a) VALUES (?)", 0);
+
+        String sql = "SELECT /*+ CNL_JOIN */ ARRAY(SELECT a FROM t) FROM t LIMIT 1";
+
+        assertQuery(sql).resultSize(1).check();
+
+        /**
+         * The data source size of (buffer size + 1) is used to ensure that multiple batches are needed
+         * on right hand of CNLJ to process all input rows, in this case left hand is not requested
+         * immediately after endLeft() call.
+         */
+        for (int i = 1; i < AbstractNode.IN_BUFFER_SIZE + 1; i++)
+            sql("INSERT INTO t (a) VALUES (?)", i);
+
+        assertQuery(sql).resultSize(1).check();
+    }
+}

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/SortAggregateIntegrationTest.java

@@ -28,6 +28,7 @@
 import org.apache.ignite.cache.QueryIndexType;
 import org.apache.ignite.cache.affinity.rendezvous.RendezvousAffinityFunction;
 import org.apache.ignite.configuration.IgniteConfiguration;
+import org.apache.ignite.internal.processors.query.calcite.exec.rel.AbstractNode;
 import org.apache.ignite.internal.util.typedef.F;
 import org.junit.Test;
 
@@ -155,6 +156,31 @@ public void testNullsReordering() {
             .check();
     }
 
+    /**
+     * Tests that sort aggregate node correctly handles the case when input data
+     * ends exactly when the requested number of rows is satisfied.
+     */
+    @Test
+    public void testRequestRowsAfterInputEnds() {
+        /**
+         * With input rows count equals to the buffer size, the last row completes both
+         * the input data and the requested count in the same cycle. This triggers
+         * a synchronous request() call from within push() to fill the buffer, and
+         * the node must properly handle the termination on the subsequent request()
+         * call rather than on end().
+         */
+        int bufSize = AbstractNode.IN_BUFFER_SIZE;
+
+        sql("CREATE TABLE t0(a INTEGER PRIMARY KEY, b INTEGER) WITH template=replicated," + atomicity());
+
+        for (int i = 0; i < bufSize; i++)
+            sql("INSERT INTO t0 VALUES (?, ?)", i, i);
+
+        assertQuery("SELECT t1.a FROM t0 AS t1 JOIN (SELECT a, count(a) FROM t0 GROUP BY a) AS t2 ON t1.a = t2.a")
+            .resultSize(bufSize)
+            .check();
+    }
+
     /**
      * @param c Cache.
      * @param rows Rows count.

modules/calcite/src/test/java/org/apache/ignite/testsuites/IntegrationTestSuite.java

@@ -32,6 +32,7 @@
 import org.apache.ignite.internal.processors.query.calcite.integration.CalciteBasicSecondaryIndexIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.CalciteErrorHandlilngIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.CalcitePlanningDumpTest;
+import org.apache.ignite.internal.processors.query.calcite.integration.CollectIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.CorrelatesIntegrationTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.DataTypesTest;
 import org.apache.ignite.internal.processors.query.calcite.integration.DateTimeTest;
@@ -187,6 +188,7 @@
     SelectByKeyFieldTest.class,
     WindowIntegrationTest.class,
     CalciteMessageUnmarshalThreadIntegrationTest.class,
+    CollectIntegrationTest.class,
 })
 public class IntegrationTestSuite {
 }