[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/TestUtils.java

@@ -65,6 +65,7 @@
 import static org.apache.iotdb.itbase.constant.TestConstant.TIMESTAMP_STR;
 import static org.apache.iotdb.itbase.env.BaseEnv.TABLE_SQL_DIALECT;
 import static org.apache.iotdb.itbase.env.BaseEnv.TREE_SQL_DIALECT;
+import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets;
 import static org.awaitility.Awaitility.await;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -807,6 +808,15 @@ public static void assertResultSetEqual(
       final String expectedHeader,
       final Collection<String> expectedResult,
       final boolean printResult) {
+    if (expectedResult instanceof Set) {
+      List<ResultSet> nodeResults = getNodeResultSets(actualResultSet);
+      if (nodeResults.size() != 1 || nodeResults.get(0) != actualResultSet) {
+        for (ResultSet nodeResult : nodeResults) {
+          assertResultSetEqual(nodeResult, expectedHeader, expectedResult, printResult);
+        }
+        return;
+      }
+    }
     try {
       final ResultSetMetaData resultSetMetaData = actualResultSet.getMetaData();
       final StringBuilder header = new StringBuilder();
@@ -868,27 +878,36 @@ public static void assertResultSetEqual(
       String expectedHeader,
       Set<String> expectedRetSet,
       Consumer consumer) {
-    try {
-      ResultSetMetaData resultSetMetaData = actualResultSet.getMetaData();
+    assertResultSetEqual(actualResultSet, expectedHeader, expectedRetSet, false);
+  }
+
+  /**
+   * Checks unordered rows on every node, preserving duplicate multiplicity.
+   *
+   * <p>The caller is responsible for ensuring that the query's final result has no ordering or
+   * arbitrary slice semantics that require deterministic row order.
+   */
+  public static void assertResultSetEqualUnordered(
+      ResultSet resultSet, String expectedHeader, List<String> expectedRows) throws SQLException {
+    List<String> sortedExpected = new ArrayList<>(expectedRows);
+    Collections.sort(sortedExpected);
+    for (ResultSet nodeResult : getNodeResultSets(resultSet)) {
+      ResultSetMetaData metadata = nodeResult.getMetaData();
       StringBuilder header = new StringBuilder();
-      for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) {
-        header.append(resultSetMetaData.getColumnName(i)).append(",");
+      for (int i = 1; i <= metadata.getColumnCount(); i++) {
+        header.append(metadata.getColumnName(i)).append(",");
       }
       assertEquals(expectedHeader, header.toString());
-
-      Set<String> actualRetSet = new HashSet<>();
-
-      while (actualResultSet.next()) {
-        StringBuilder builder = new StringBuilder();
-        for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) {
-          builder.append(actualResultSet.getString(i)).append(",");
+      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(",");
         }
-        actualRetSet.add(builder.toString());
+        actualRows.add(row.toString());
       }
-      assertStringSetEqual(expectedRetSet, actualRetSet);
-    } catch (Exception e) {
-      e.printStackTrace();
-      Assert.fail(String.valueOf(e));
+      Collections.sort(actualRows);
+      assertEquals(sortedExpected, actualRows);
     }
   }
 

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

@@ -44,6 +44,7 @@
 
 import java.sql.Connection;
 import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Arrays;
@@ -54,6 +55,8 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
+import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets;
+
 @RunWith(IoTDBTestRunner.class)
 @Category({MultiClusterIT3.class})
 public class IoTDBShowReceiversTripleIT {
@@ -214,39 +217,44 @@ 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);
+        final ResultSet clusterResult = statement.executeQuery(sql)) {
+      for (ResultSet resultSet : getNodeResultSets(clusterResult)) {
+        boolean hasFirstPipe = false;
+        boolean hasSecondPipe = false;
+        final Set<Integer> receiverDataNodeIds = new HashSet<>();
+        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);
+          }
         }
-        if (receiverDataNodeId == secondReceiverDataNodeId
-            && pipeIds.contains(secondPipeName + "@")
-            && getInt(resultSet, "ConnectionCount", "connection_count") >= 1
-            && getInt(resultSet, "PipeCount", "pipe_count") >= 1) {
-          hasSecondPipe = true;
-          receiverDataNodeIds.add(receiverDataNodeId);
+        if (!hasFirstPipe || !hasSecondPipe || receiverDataNodeIds.size() < 2) {
+          return false;
         }
       }
     }
-    return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2;
+    return true;
   }
 
   private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception {
@@ -268,20 +276,27 @@ private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception
   private static String getString(
       final ResultSet resultSet, final String treeColumnName, final String tableColumnName)
       throws SQLException {
-    try {
-      return resultSet.getString(treeColumnName);
-    } catch (final SQLException ignored) {
-      return resultSet.getString(tableColumnName);
-    }
+    return resultSet.getString(getColumnIndex(resultSet, treeColumnName, tableColumnName));
   }
 
   private static int getInt(
       final ResultSet resultSet, final String treeColumnName, final String tableColumnName)
       throws SQLException {
-    try {
-      return resultSet.getInt(treeColumnName);
-    } catch (final SQLException ignored) {
-      return resultSet.getInt(tableColumnName);
+    return resultSet.getInt(getColumnIndex(resultSet, treeColumnName, tableColumnName));
+  }
+
+  private static int getColumnIndex(
+      final ResultSet resultSet, final String treeColumnName, final String tableColumnName)
+      throws SQLException {
+    // SHOW RECEIVERS and information_schema.receivers expose different column names.
+    final ResultSetMetaData metadata = resultSet.getMetaData();
+    for (int column = 1; column <= metadata.getColumnCount(); column++) {
+      final String name = metadata.getColumnName(column);
+      if (treeColumnName.equals(name) || tableColumnName.equals(name)) {
+        return column;
+      }
     }
+    throw new AssertionError(
+        "Missing receiver column: " + treeColumnName + " / " + tableColumnName);
   }
 }

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

@@ -37,6 +37,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -269,20 +270,22 @@ private void testSelectAll() {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
         Statement statement = connection.createStatement()) {
       statement.execute("use \"test\"");
-      Map<String, Long> lastTimeMap = new HashMap<>();
-      try (ResultSet resultSet = statement.executeQuery(selectSql)) {
-        int cnt = 0;
-        while (resultSet.next()) {
-          String tag = resultSet.getString("tag1") + "_" + resultSet.getString("tag2");
-          long before = lastTimeMap.getOrDefault(tag, -1L);
-          long cur = resultSet.getTimestamp("time").getTime();
-          if (cur <= before) {
-            fail("time order wrong!");
+      try (ResultSet clusterResult = statement.executeQuery(selectSql)) {
+        for (ResultSet resultSet : getNodeResultSets(clusterResult)) {
+          Map<String, Long> lastTimeMap = new HashMap<>();
+          int cnt = 0;
+          while (resultSet.next()) {
+            String tag = resultSet.getString("tag1") + "_" + resultSet.getString("tag2");
+            long before = lastTimeMap.getOrDefault(tag, -1L);
+            long cur = resultSet.getTimestamp("time").getTime();
+            if (cur <= before) {
+              fail("time order wrong!");
+            }
+            lastTimeMap.put(tag, cur);
+            cnt++;
           }
-          lastTimeMap.put(tag, cur);
-          cnt++;
+          assertEquals(16030, cnt);
         }
-        assertEquals(16030, cnt);
       }
     } catch (Exception e) {
       e.printStackTrace();

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

@@ -41,6 +41,7 @@
 
 import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
+import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets;
 import static org.junit.Assert.fail;
 
 @RunWith(IoTDBTestRunner.class)
@@ -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 =
+      try (ResultSet clusterResult =
           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));
-        }
-      }
-    }
+        for (ResultSet resultSet : getNodeResultSets(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

@@ -49,6 +49,7 @@
 import java.util.List;
 import java.util.Locale;
 
+import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -269,24 +270,32 @@ public void testPartialColumn() throws SQLException {
 
       // It is allowed to insert without tags
       statement.execute("INSERT INTO vehicle1(time, s1, s2) SELECT time, s1, s2 FROM vehicle0");
-      ResultSet resultSet = statement.executeQuery("SELECT deviceId, manufacturer FROM vehicle1");
-      while (resultSet.next()) {
-        for (int i = 1; i <= 2; i++) {
-          assertNull(resultSet.getString(i));
+      try (ResultSet clusterResult =
+          statement.executeQuery("SELECT deviceId, manufacturer FROM vehicle1")) {
+        for (ResultSet resultSet : getNodeResultSets(clusterResult)) {
+          while (resultSet.next()) {
+            for (int i = 1; i <= 2; i++) {
+              assertNull(resultSet.getString(i));
+            }
+          }
         }
       }
 
       // insert partial tag & fields
       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");
-      while (resultSet.next()) {
-        for (int i = 1; i <= 4; i++) {
-          if (i % 2 == 1) {
-            assertNull(resultSet.getString(i));
-          } else {
-            assertNotNull(resultSet.getString(i));
+      try (ResultSet clusterResult =
+          statement.executeQuery("SELECT deviceId, manufacturer, s1, s2 FROM vehicle1")) {
+        for (ResultSet resultSet : getNodeResultSets(clusterResult)) {
+          while (resultSet.next()) {
+            for (int i = 1; i <= 4; i++) {
+              if (i % 2 == 1) {
+                assertNull(resultSet.getString(i));
+              } else {
+                assertNotNull(resultSet.getString(i));
+              }
+            }
           }
         }
       }