[Query] Share table fragment exchange budgets and consume ready Collect inputs

integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java

@@ -40,6 +40,7 @@
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.util.Calendar;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -68,6 +69,17 @@ public ClusterTestResultSet(
     this.env = env;
   }
 
+  /**
+   * Returns each node's result for assertions that do not require identical row order or plan text.
+   * Every returned result must be checked. The caller must close the original result set and must
+   * not mix per-node iteration with iteration through the cluster wrapper.
+   */
+  public static List<ResultSet> getNodeResultSets(ResultSet resultSet) {
+    return resultSet instanceof ClusterTestResultSet
+        ? Collections.unmodifiableList(((ClusterTestResultSet) resultSet).resultSets)
+        : Collections.singletonList(resultSet);
+  }
+
   @Override
   public boolean next() throws SQLException {
     RequestDelegate<Boolean> delegate = createRequestDelegate();

integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java

@@ -0,0 +1,120 @@
+/*
+ * 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.iotdb.db.it.utils;
+
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QueryBody;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QuerySpecification;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.TableSubquery;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Explain;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser;
+import org.apache.iotdb.itbase.runtime.ClusterTestResultSet;
+
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.time.ZoneOffset;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/** Explicit SQL-aware assertions for queries affected by ready-input Collect consumption. */
+public final class CollectQueryAssertions {
+  private CollectQueryAssertions() {}
+
+  /** Only queries without final ordering or an arbitrary top-level slice may use this path. */
+  public static List<ResultSet> getUnorderedResultSets(String sql, ResultSet resultSet) {
+    org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement statement =
+        new SqlParser().createStatement(sql, ZoneOffset.UTC, null);
+    assertTrue("Unordered assertions require a SELECT query: " + sql, statement instanceof Query);
+    Query query = (Query) statement;
+    while (true) {
+      assertFalse(
+          "Final ORDER BY requires ordered assertions: " + sql, query.getOrderBy().isPresent());
+      assertFalse(
+          "LIMIT requires deterministic result selection: " + sql, query.getLimit().isPresent());
+      assertFalse(
+          "OFFSET requires deterministic result selection: " + sql, query.getOffset().isPresent());
+      QueryBody body = query.getQueryBody();
+      if (body instanceof QuerySpecification) {
+        QuerySpecification specification = (QuerySpecification) body;
+        assertFalse(
+            "Final ORDER BY requires ordered assertions: " + sql,
+            specification.getOrderBy().isPresent());
+        assertFalse(
+            "LIMIT requires deterministic result selection: " + sql,
+            specification.getLimit().isPresent());
+        assertFalse(
+            "OFFSET requires deterministic result selection: " + sql,
+            specification.getOffset().isPresent());
+      }
+      if (!(body instanceof TableSubquery)) {
+        break;
+      }
+      // Parentheses around the whole query retain its ordering; nested FROM/UNION branches do not.
+      query = ((TableSubquery) body).getQuery();
+    }
+    return nodeResults(resultSet);
+  }
+
+  /** EXPLAIN text may contain node-specific details; each node must satisfy the plan assertions. */
+  public static List<ResultSet> getExplainResultSets(String sql, ResultSet resultSet) {
+    assertTrue(
+        "Plan assertions require EXPLAIN: " + sql,
+        new SqlParser().createStatement(sql, ZoneOffset.UTC, null) instanceof Explain);
+    return nodeResults(resultSet);
+  }
+
+  private static List<ResultSet> nodeResults(ResultSet resultSet) {
+    return resultSet instanceof ClusterTestResultSet
+        ? ClusterTestResultSet.getNodeResultSets(resultSet)
+        : Collections.singletonList(resultSet);
+  }
+
+  /** Compare every node independently while retaining duplicate multiplicity. */
+  public static void assertUnorderedQueryResult(
+      String sql, ResultSet resultSet, String expectedHeader, List<String> expectedRows)
+      throws SQLException {
+    List<String> sortedExpected = new ArrayList<>(expectedRows);
+    Collections.sort(sortedExpected);
+    for (ResultSet nodeResult : getUnorderedResultSets(sql, resultSet)) {
+      ResultSetMetaData metadata = nodeResult.getMetaData();
+      StringBuilder header = new StringBuilder();
+      for (int i = 1; i <= metadata.getColumnCount(); i++) {
+        header.append(metadata.getColumnName(i)).append(",");
+      }
+      assertEquals(expectedHeader, header.toString());
+      List<String> actualRows = new ArrayList<>();
+      while (nodeResult.next()) {
+        StringBuilder row = new StringBuilder();
+        for (int i = 1; i <= metadata.getColumnCount(); i++) {
+          row.append(nodeResult.getString(i)).append(",");
+        }
+        actualRows.add(row.toString());
+      }
+      Collections.sort(actualRows);
+      assertEquals(sortedExpected, actualRows);
+    }
+  }
+}

integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java

@@ -366,7 +366,11 @@ public static String generateHeaderResults() {
   }
 
   public static String getQuerySql(final String table) {
-    return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from " + table;
+    // ClusterTestResultSet compares replicas row by row. All tags plus time form a unique key,
+    // so require the same row order even when Collect consumes inputs in readiness order.
+    return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from "
+        + table
+        + " order by s0,s1,s2,s3,time";
   }
 
   public static String getQueryCountSql(final String table) {

integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java

@@ -54,6 +54,8 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
+import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets;
+
 @RunWith(IoTDBTestRunner.class)
 @Category({MultiClusterIT3.class})
 public class IoTDBShowReceiversTripleIT {
@@ -214,37 +216,66 @@ private boolean hasReceiversOnDifferentDataNodes(
       final String secondPipeName,
       final int secondReceiverDataNodeId)
       throws SQLException {
-    boolean hasFirstPipe = false;
-    boolean hasSecondPipe = false;
-    final Set<Integer> receiverDataNodeIds = new HashSet<>();
     try (final Connection connection =
             receiverEnv.getConnection(
                 SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect);
         final Statement statement = connection.createStatement();
         final ResultSet resultSet = statement.executeQuery(sql)) {
-      while (resultSet.next()) {
-        final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids");
-        if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type"))
-            || !"thrift".equals(getString(resultSet, "Protocol", "protocol"))
-            || pipeIds == null) {
-          continue;
-        }
-        final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id");
-        if (receiverDataNodeId == firstReceiverDataNodeId
-            && pipeIds.contains(firstPipeName + "@")
-            && getInt(resultSet, "ConnectionCount", "connection_count") >= 1
-            && getInt(resultSet, "PipeCount", "pipe_count") >= 1) {
-          hasFirstPipe = true;
-          receiverDataNodeIds.add(receiverDataNodeId);
-        }
-        if (receiverDataNodeId == secondReceiverDataNodeId
-            && pipeIds.contains(secondPipeName + "@")
-            && getInt(resultSet, "ConnectionCount", "connection_count") >= 1
-            && getInt(resultSet, "PipeCount", "pipe_count") >= 1) {
-          hasSecondPipe = true;
-          receiverDataNodeIds.add(receiverDataNodeId);
+      // Each node must see both senders, but unordered table queries may return different row
+      // orders.
+      for (final ResultSet nodeResult :
+          BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect)
+              ? getUnorderedResultSets(sql, resultSet)
+              : Collections.singletonList(resultSet)) {
+        if (!hasExpectedReceivers(
+            nodeResult,
+            sqlDialect,
+            firstPipeName,
+            firstReceiverDataNodeId,
+            secondPipeName,
+            secondReceiverDataNodeId)) {
+          return false;
         }
       }
+      return true;
+    }
+  }
+
+  private static boolean hasExpectedReceivers(
+      final ResultSet resultSet,
+      final String sqlDialect,
+      final String firstPipeName,
+      final int firstReceiverDataNodeId,
+      final String secondPipeName,
+      final int secondReceiverDataNodeId)
+      throws SQLException {
+    boolean hasFirstPipe = false;
+    boolean hasSecondPipe = false;
+    final Set<Integer> receiverDataNodeIds = new HashSet<>();
+    while (resultSet.next()) {
+      final String pipeIds = getString(resultSet, sqlDialect, "PipeIDs", "pipe_ids");
+      if (!"DataNode"
+              .equals(getString(resultSet, sqlDialect, "ReceiverNodeType", "receiver_node_type"))
+          || !"thrift".equals(getString(resultSet, sqlDialect, "Protocol", "protocol"))
+          || pipeIds == null) {
+        continue;
+      }
+      final int receiverDataNodeId =
+          getInt(resultSet, sqlDialect, "ReceiverNodeId", "receiver_node_id");
+      if (receiverDataNodeId == firstReceiverDataNodeId
+          && pipeIds.contains(firstPipeName + "@")
+          && getInt(resultSet, sqlDialect, "ConnectionCount", "connection_count") >= 1
+          && getInt(resultSet, sqlDialect, "PipeCount", "pipe_count") >= 1) {
+        hasFirstPipe = true;
+        receiverDataNodeIds.add(receiverDataNodeId);
+      }
+      if (receiverDataNodeId == secondReceiverDataNodeId
+          && pipeIds.contains(secondPipeName + "@")
+          && getInt(resultSet, sqlDialect, "ConnectionCount", "connection_count") >= 1
+          && getInt(resultSet, sqlDialect, "PipeCount", "pipe_count") >= 1) {
+        hasSecondPipe = true;
+        receiverDataNodeIds.add(receiverDataNodeId);
+      }
     }
     return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2;
   }
@@ -266,22 +297,22 @@ private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception
   }
 
   private static String getString(
-      final ResultSet resultSet, final String treeColumnName, final String tableColumnName)
+      final ResultSet resultSet,
+      final String sqlDialect,
+      final String treeColumnName,
+      final String tableColumnName)
       throws SQLException {
-    try {
-      return resultSet.getString(treeColumnName);
-    } catch (final SQLException ignored) {
-      return resultSet.getString(tableColumnName);
-    }
+    return resultSet.getString(
+        BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) ? tableColumnName : treeColumnName);
   }
 
   private static int getInt(
-      final ResultSet resultSet, final String treeColumnName, final String tableColumnName)
+      final ResultSet resultSet,
+      final String sqlDialect,
+      final String treeColumnName,
+      final String tableColumnName)
       throws SQLException {
-    try {
-      return resultSet.getInt(treeColumnName);
-    } catch (final SQLException ignored) {
-      return resultSet.getInt(tableColumnName);
-    }
+    return resultSet.getInt(
+        BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) ? tableColumnName : treeColumnName);
   }
 }

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java

@@ -264,7 +264,7 @@ public void testMultiDeviceQueryAndDelete() {
   }
 
   private void testSelectAll() {
-    String selectSql = "select * from t";
+    String selectSql = "select * from t ORDER BY tag1, tag2, time";
 
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
         Statement statement = connection.createStatement()) {

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java

@@ -39,6 +39,7 @@
 import java.util.HashSet;
 import java.util.Set;
 
+import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
 import static org.junit.Assert.fail;
@@ -200,37 +201,39 @@ public void testLargeResultIsSplitWithoutDataLoss() throws Exception {
         "large_result",
         "org.apache.iotdb.db.query.udf.example.relational.LargeResultTableFunction");
 
-    Set<String> returnedRows = new HashSet<>();
     try (Connection connection = EnvFactory.getEnv().getTableConnection();
         Statement statement = connection.createStatement()) {
       statement.execute("USE " + DATABASE_NAME);
-      try (ResultSet resultSet =
-          statement.executeQuery(
-              "SELECT * FROM large_result(vehicle, "
-                  + LARGE_RESULT_REPEAT_COUNT
-                  + ", "
-                  + LARGE_RESULT_PAYLOAD_SIZE
-                  + ")")) {
-        while (resultSet.next()) {
-          int repeatIndex = resultSet.getInt("repeat_index");
-          long time = resultSet.getLong("time");
-          Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT);
-          Assert.assertEquals(
-              repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE),
-              resultSet.getString("payload"));
-          assertPassThroughColumns(resultSet, time);
-          Assert.assertTrue(
-              "Duplicate result row for time " + time + " and repeat index " + repeatIndex,
-              returnedRows.add(time + ":" + repeatIndex));
-        }
-      }
-    }
+      String sql =
+          "SELECT * FROM large_result(vehicle, "
+              + LARGE_RESULT_REPEAT_COUNT
+              + ", "
+              + LARGE_RESULT_PAYLOAD_SIZE
+              + ")";
+      try (ResultSet clusterResult = statement.executeQuery(sql)) {
+        for (ResultSet resultSet : getUnorderedResultSets(sql, clusterResult)) {
+          Set<String> returnedRows = new HashSet<>();
+          while (resultSet.next()) {
+            int repeatIndex = resultSet.getInt("repeat_index");
+            long time = resultSet.getLong("time");
+            Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT);
+            Assert.assertEquals(
+                repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE),
+                resultSet.getString("payload"));
+            assertPassThroughColumns(resultSet, time);
+            Assert.assertTrue(
+                "Duplicate result row for time " + time + " and repeat index " + repeatIndex,
+                returnedRows.add(time + ":" + repeatIndex));
+          }
 
-    long[] inputTimes = new long[] {1, 2, 3, 5};
-    Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size());
-    for (long time : inputTimes) {
-      for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) {
-        Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex));
+          long[] inputTimes = new long[] {1, 2, 3, 5};
+          Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size());
+          for (long time : inputTimes) {
+            for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) {
+              Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex));
+            }
+          }
+        }
       }
     }
   }

integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java

@@ -279,7 +279,9 @@ public void testPartialColumn() throws SQLException {
       statement.execute("DELETE FROM vehicle1");
       statement.execute(
           "INSERT INTO vehicle1(time, s2, manufacturer) SELECT time, s2, manufacturer FROM vehicle0");
-      resultSet = statement.executeQuery("SELECT deviceId, manufacturer, s1, s2 FROM vehicle1");
+      resultSet =
+          statement.executeQuery(
+              "SELECT deviceId, manufacturer, s1, s2 FROM vehicle1 ORDER BY manufacturer, s2");
       while (resultSet.next()) {
         for (int i = 1; i <= 4; i++) {
           if (i % 2 == 1) {