From 156b566df082393fa7fcae00f06272d6105f04c5 Mon Sep 17 00:00:00 2001
From: ColinLee <shuolin_l@163.com>
Date: Wed, 14 Oct 2026 12:11:34 +0800
Subject: [PATCH 1/3] Document exchange fan-in implementation scope

---
 ...2026-10-14-exchange-fanin-budget-design.md | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md

diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
new file mode 100644
index 0000000000000..b2bc07d1ddbb0
--- /dev/null
+++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
@@ -0,0 +1,71 @@
+<!--
+
+    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.
+
+-->
+
+# Exchange 无序扇入与输入输出预算
+
+基线：Apache master `aec1ad6139cb4dd1450725f53418cf40c26efb4e`。
+依据：[技术设计文档](https://timechor.feishu.cn/docx/EpDFdmS94ovMCWxluEOc3WGDn8d)，修订 861。
+
+## 已确认范围
+
+实现 FI 输入输出共享账目、Collect/MappingCollect 任意就绪消费，以及由无序汇集引起的计划属性修正。
+聚合使用 HashAggregation：只清空实际输入已失去连续性的 preGroupedSymbols；有效分支内的流式聚合保留。
+MergeSort、Join、ASOF、CrossJoin 只作依赖验证，不修改其消费算法，不新增进度保障机制。
+TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍需匹配实际输入。
+
+## 额度与账目
+
+F 为配置的单 FI 交换预算。INPUT=F/2，OUTPUT=F-INPUT；同 FI 的所有 pipeline 继承相同固定值。
+远端 SourceHandle 计入接收 FI 的 INPUT；SinkChannel 和 SharedTsBlockQueue 计入生产 FI 的 OUTPUT。
+本地共享队列只计费一次，创建先后不改变所有者。
+
+MemoryPool 维护 query→FI→planNode 明细，注册时固定 side。
+FI 两侧使用 AtomicLong；planNode 字节明细通过 compute 增减。
+申请、回滚、释放均对应修改全局池、FI 单侧与 planNode 三层账目。
+额度判断使用本次增加后的单侧总数；稳定快照中的总数等于同侧明细之和。
+保留现有最高优先级查询的零计费兜底与交换组件的收发、等待、释放机制。
+
+## Collect 生命周期
+
+AnyChildBlocked 缓存直接 child.isBlocked 的 Future，每个 pending 周期只有一个 listener。
+新轮次重定向现有 listener；发布后检查完成状态，覆盖注册与完成交错。
+回调只记录就绪或错误，在锁外完成选择 Future；所有 child 方法由 Driver 线程调用。
+重复 isBlocked 与输出背压保留选择；取消对外等待不向 child 传播。
+
+一次 next 消耗一次选择。null/空块刷新等待，不表示结束。
+hasNext=false 后检查 isFinished，确认结束才关闭 child 并移除记录。
+MappingCollect 按实际 childId 映射。关闭唤醒待完成选择，忽略迟到回调，逐一尝试关闭全部输入。
+保留内存 R 为各存活 child 保留内存之和；峰值为 max(R,max(R-r_i+p_i))；返回大小为 max(b_i)。
+
+## 计划属性
+
+多路 Collect/Union 打断分支连续性和排序；单路可透传。
+Filter 与符号透传 Project 不恢复属性；有效 Sort/MergeSort 可以重新建立属性。
+在 SINGLE/PARTIAL/FINAL 实际 child 组装后校验聚合，包括 DISTINCT/mask、普通分阶段聚合及 Union 下推。
+Group 在分区连续性失效时改为完整 Sort。复用 master 的 StreamSort 前缀校验和 Window/Fill 修复。
+SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除恢复排序。
+
+## 验证
+
+覆盖共享额度并发、INPUT/OUTPUT 隔离、所有者、奇数 F、超大申请、回滚释放取消及稳定账目一致性；
+覆盖 100/1000 child、A 阻塞 B 就绪、listener 复用、注册竞争、输出背压、空块、结束、异常取消、关闭和映射；
+覆盖文档 SQL 对应的聚合、Group/Window、StreamSort、Union 与 Fill 计划及实际结果。
+运行 Spotless、英文/中文全 reactor test-compile、针对性单测，以及本分支新增或修改的 IT。

From ff59338f89ab272120dcc58afef6b184878c1164 Mon Sep 17 00:00:00 2001
From: ColinLee <shuolin_l@163.com>
Date: Wed, 14 Oct 2026 12:47:59 +0800
Subject: [PATCH 2/3] Share fragment exchange budgets and consume ready Collect
 inputs

---
 ...2026-10-14-exchange-fanin-budget-design.md |  28 ++
 .../it/query/recent/IoTDBCollectFanInIT.java  | 276 +++++++++++
 .../operator/process/AnyChildBlocked.java     | 242 ++++++++++
 .../operator/process/CollectOperator.java     | 136 ++++--
 .../process/MappingCollectOperator.java       |  56 ++-
 .../operator/process/CollectOperatorTest.java | 442 ++++++++++++++++++
 .../iotdb/db/i18n/DataNodeQueryMessages.java  |   2 +
 .../iotdb/db/i18n/DataNodeQueryMessages.java  |   2 +
 .../exchange/SharedTsBlockQueue.java          |   9 +-
 .../exchange/sink/ShuffleSinkHandle.java      |   4 +-
 .../execution/exchange/sink/SinkChannel.java  |   9 +-
 .../exchange/source/SourceHandle.java         |   9 +-
 .../execution/memory/MemoryPool.java          | 195 ++++++--
 .../execution/memory/MemorySide.java          |  31 ++
 .../DataNodeTableOperatorGenerator.java       |   4 +-
 .../planner/LocalExecutionPlanContext.java    |  46 +-
 .../plan/planner/LocalExecutionPlanner.java   |   4 +-
 .../plan/planner/OperatorTreeGenerator.java   |   6 +-
 .../TableDistributedPlanGenerator.java        | 122 ++++-
 .../optimizations/ParallelizeGrouping.java    |  16 +
 .../optimizations/SortElimination.java        |  25 +
 .../exchange/MPPDataExchangeManagerTest.java  |  53 +++
 .../execution/exchange/SourceHandleTest.java  |   3 +-
 .../execution/memory/MemoryPoolTest.java      | 115 ++++-
 .../LocalExecutionPlanContextBudgetTest.java  |  78 ++++
 .../planner/CollectInputPropertyTest.java     | 149 ++++++
 .../ParallelizeGroupingTest.java              |  54 +++
 27 files changed, 1968 insertions(+), 148 deletions(-)
 create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java
 create mode 100644 iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java
 create mode 100644 iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java
 create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java
 create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java
 create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java

diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
index b2bc07d1ddbb0..a817039cab4b9 100644
--- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
+++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
@@ -69,3 +69,31 @@ SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除
 覆盖 100/1000 child、A 阻塞 B 就绪、listener 复用、注册竞争、输出背压、空块、结束、异常取消、关闭和映射；
 覆盖文档 SQL 对应的聚合、Group/Window、StreamSort、Union 与 Fill 计划及实际结果。
 运行 Spotless、英文/中文全 reactor test-compile、针对性单测，以及本分支新增或修改的 IT。
+
+## 实现验证记录（2026-10-14）
+
+第一版实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。
+集成验证包含普通与 DISTINCT 聚合、Union 分阶段聚合、嵌套窗口、Fill、排序/TopK、Join 依赖、树模型以及树表视图。
+
+构建先执行以下干净构建，生成当前分支的节点包并安装本地依赖：
+
+```sh
+mvn -o clean install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests
+```
+
+全 reactor 中文编译与英文编译/回归均通过；英文 test 阶段包含完整 test-compile：
+
+```sh
+mvn -o test-compile -DskipTests -P with-zh-locale
+mvn -o test -Dtest=CollectOperatorTest,MemoryPoolTest,LocalExecutionPlanContextBudgetTest,CollectInputPropertyTest,SharedTsBlockQueueTest,SourceHandleTest,SinkChannelTest,LocalSourceHandleTest,LocalSinkChannelTest,MPPDataExchangeManagerTest,ParallelizeGroupingTest,TableFunctionSortTest -Dsurefire.failIfNoSpecifiedTests=false
+```
+
+两个集成环境均只运行本分支新增的测试类：
+
+```sh
+mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableSimpleIT -P with-integration-tests
+mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableClusterIT -P with-integration-tests
+```
+
+使用离线模式是为了避免远端 SNAPSHOT ZIP 下载停滞；节点执行的是当前分支重新打包的代码。
+Spotless、Checkstyle、新增文件许可证及中英文消息键名/占位符一致性检查通过。
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java
new file mode 100644
index 0000000000000..adc72b6046e50
--- /dev/null
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java
@@ -0,0 +1,276 @@
+/*
+ * 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.relational.it.query.recent;
+
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.TableClusterIT;
+import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.Statement;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({TableLocalStandaloneIT.class, TableClusterIT.class})
+public class IoTDBCollectFanInIT {
+  private static final String DATABASE = "collect_fan_in";
+  private static final int DEVICES = 24;
+  private static final int ROWS = 8;
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    EnvFactory.getEnv()
+        .getConfig()
+        .getCommonConfig()
+        .setMaxTsBlockLineNumber(2)
+        .setSortBufferSize(128 * 1024)
+        .setDataRegionGroupExtensionPolicy("CUSTOM")
+        .setDataRegionPerDataNode(3);
+    EnvFactory.getEnv().initClusterEnvironment();
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("CREATE DATABASE " + DATABASE);
+      statement.execute("USE " + DATABASE);
+      for (String table : new String[] {"readings1", "readings2"}) {
+        statement.execute(
+            "CREATE TABLE " + table + " (device STRING TAG, zone STRING TAG, value INT32 FIELD)");
+        for (int device = 0; device < DEVICES; device++) {
+          for (int value = 1; value <= ROWS; value++) {
+            statement.execute(
+                "INSERT INTO "
+                    + table
+                    + " (time,device,zone,value) VALUES ("
+                    + value
+                    + ",'"
+                    + device(device)
+                    + "',"
+                    + (zone(device) == null ? "NULL" : "'" + zone(device) + "'")
+                    + ","
+                    + value
+                    + ")");
+          }
+        }
+      }
+      statement.execute("FLUSH");
+    }
+  }
+
+  @Test
+  public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("CREATE DATABASE root.collect_tree");
+      for (int device = 0; device < DEVICES; device++) {
+        String path = "root.collect_tree.z" + device % 3 + "." + device(device);
+        statement.execute("CREATE ALIGNED TIMESERIES " + path + "(value INT32)");
+        for (int value = 1; value <= ROWS; value++) {
+          statement.execute(
+              "INSERT INTO " + path + "(time,value) ALIGNED VALUES (" + value + "," + value + ")");
+        }
+      }
+      try (ResultSet result =
+          statement.executeQuery("SELECT value FROM root.collect_tree.*.* ALIGN BY DEVICE")) {
+        int rows = 0;
+        while (result.next()) {
+          assertEquals(result.getLong("Time"), result.getLong("value"));
+          rows++;
+        }
+        assertEquals(DEVICES * ROWS, rows);
+      }
+    }
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      statement.execute(
+          "CREATE VIEW tree_readings(zone STRING TAG, device STRING TAG, value INT32 FIELD) AS root.collect_tree.**");
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT device, SUM(value) FROM ("
+                  + "SELECT device, value FROM tree_readings UNION ALL SELECT device, value FROM tree_readings"
+                  + ") u GROUP BY device ORDER BY device")) {
+        for (int device = 0; device < DEVICES; device++) {
+          assertTrue(result.next());
+          assertEquals(device(device), result.getString(1));
+          assertEquals(72.0, result.getDouble(2), 0.0);
+        }
+        assertFalse(result.next());
+      }
+    }
+  }
+
+  @AfterClass
+  public static void tearDown() {
+    EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  private static String device(int id) {
+    return "d" + (id < 10 ? "0" : "") + id;
+  }
+
+  private static String zone(int id) {
+    return id % 3 == 0 ? null : "z" + id % 3;
+  }
+
+  @Test
+  public void testGroupedAggregationAfterGlobalWindow() throws Exception {
+    checkGroups(
+        "SELECT device, zone, value FROM (SELECT device, zone, value, "
+            + "COUNT(*) OVER () AS n FROM readings1) t WHERE n > 0",
+        1);
+  }
+
+  @Test
+  public void testUnionPartialFinalAndDistinctMask() throws Exception {
+    String input =
+        "SELECT device, zone, value FROM readings1 UNION ALL "
+            + "SELECT device, zone, value FROM readings2";
+    checkGroups(input, 2);
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      // Without DISTINCT, this query exercises branch PARTIAL and post-Collect FINAL aggregation.
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT device, zone, SUM(value) FROM ("
+                  + input
+                  + ") u GROUP BY device, zone ORDER BY device")) {
+        for (int device = 0; device < DEVICES; device++) {
+          assertTrue(result.next());
+          assertEquals(device(device), result.getString(1));
+          assertEquals(zone(device), result.getString(2));
+          assertEquals(72.0, result.getDouble(3), 0.0);
+        }
+        assertFalse(result.next());
+      }
+    }
+  }
+
+  @Test
+  public void testOrderedFanInAndJoinDependencies() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      for (String limit : new String[] {"", " LIMIT 13"}) {
+        try (ResultSet result =
+            statement.executeQuery(
+                "SELECT device, value FROM (SELECT device, value FROM readings1 UNION ALL "
+                    + "SELECT device, value FROM readings2) u ORDER BY device, value"
+                    + limit)) {
+          int expectedRows = limit.isEmpty() ? DEVICES * ROWS * 2 : 13;
+          for (int row = 0; row < expectedRows; row++) {
+            assertTrue(result.next());
+            assertEquals(device(row / (ROWS * 2)), result.getString(1));
+            assertEquals(row / 2 % ROWS + 1, result.getInt(2));
+          }
+          assertFalse(result.next());
+        }
+      }
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT COUNT(*) FROM readings1 a JOIN readings2 b "
+                  + "ON a.device = b.device AND a.time = b.time WHERE a.value = b.value")) {
+        assertTrue(result.next());
+        assertEquals(DEVICES * ROWS, result.getLong(1));
+        assertFalse(result.next());
+      }
+    }
+  }
+
+  private void checkGroups(String input, int copies) throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT device, zone, SUM(value), "
+                  + "COUNT(*), COUNT(DISTINCT value) FROM ("
+                  + input
+                  + ") u GROUP BY device, zone ORDER BY device")) {
+        for (int device = 0; device < DEVICES; device++) {
+          assertTrue(result.next());
+          assertEquals(device(device), result.getString(1));
+          assertEquals(zone(device), result.getString(2));
+          assertEquals(36.0 * copies, result.getDouble(3), 0.0);
+          assertEquals((long) ROWS * copies, result.getLong(4));
+          assertEquals(ROWS, result.getLong(5));
+        }
+        assertFalse(result.next());
+      }
+    }
+  }
+
+  @Test
+  public void testNestedWindowAfterCollectRestoresCompletePartitions() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT device, zone, value, "
+                  + "SUM(value) OVER (PARTITION BY device, zone ORDER BY value) AS running_sum FROM ("
+                  + "SELECT device, zone, value, COUNT(*) OVER () AS n FROM readings1) t "
+                  + "WHERE n > 0 ORDER BY device, value")) {
+        for (int device = 0; device < DEVICES; device++) {
+          for (int value = 1; value <= ROWS; value++) {
+            assertTrue(result.next());
+            assertEquals(device(device), result.getString(1));
+            assertEquals(zone(device), result.getString(2));
+            assertEquals(value, result.getInt(3));
+            assertEquals((double) value * (value + 1) / 2, result.getDouble(4), 0.0);
+          }
+        }
+        assertFalse(result.next());
+      }
+    }
+  }
+
+  @Test
+  public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT device, zone, value, COUNT(*) OVER () AS n "
+                  + "FROM readings1 FILL METHOD CONSTANT 0 ORDER BY device, value DESC")) {
+        for (int device = 0; device < DEVICES; device++) {
+          for (int value = ROWS; value > 0; value--) {
+            assertTrue(result.next());
+            assertEquals(device(device), result.getString(1));
+            assertEquals(value, result.getInt(3));
+            assertEquals((long) DEVICES * ROWS, result.getLong(4));
+          }
+        }
+        assertFalse(result.next());
+      }
+    }
+  }
+}
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java
new file mode 100644
index 0000000000000..c0397bf1ff5a4
--- /dev/null
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java
@@ -0,0 +1,242 @@
+/*
+ * 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.calc.execution.operator.process;
+
+import org.apache.iotdb.calc.execution.operator.Operator;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import org.apache.tsfile.utils.Accountable;
+import org.apache.tsfile.utils.RamUsageEstimator;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+
+import static com.google.common.util.concurrent.Futures.getDone;
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+
+/**
+ * Observes direct children without calling operators from completion callbacks. The driver owns
+ * selection/invalidation; the lock protects listener records shared with asynchronous completions.
+ */
+final class AnyChildBlocked implements Accountable {
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AnyChildBlocked.class);
+  private static final long LISTENER_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(ChildFutureListener.class);
+  private static final long ROUND_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(WaitRound.class)
+          + RamUsageEstimator.shallowSizeOfInstance(SettableFuture.class);
+
+  private final List<Operator> children;
+  private final Object lock = new Object();
+  private final ChildFutureListener[] listeners;
+  private WaitRound currentRound;
+  private Throwable firstFailure;
+  private boolean closed;
+  private int nextChild;
+
+  AnyChildBlocked(List<Operator> children) {
+    this.children = children;
+    this.listeners = new ChildFutureListener[children.size()];
+  }
+
+  ListenableFuture<Integer> getAnyChild() {
+    WaitRound round;
+    synchronized (lock) {
+      if (currentRound != null) {
+        return currentRound.selected;
+      }
+      round = new WaitRound();
+      currentRound = round;
+      for (ChildFutureListener listener : listeners) {
+        if (listener != null) {
+          listener.target = round;
+        }
+      }
+    }
+
+    // Activate every live direct child, even if an earlier one has already become ready. Only
+    // invalidated records call child.isBlocked(); a long-pending future keeps its one listener.
+    for (int offset = 0; offset < children.size(); offset++) {
+      int childId = (nextChild + offset) % children.size();
+      Operator child;
+      ChildFutureListener listener;
+      synchronized (lock) {
+        if (closed) {
+          break;
+        }
+        child = children.get(childId);
+        listener = listeners[childId];
+      }
+      if (child == null) {
+        continue;
+      }
+      if (listener == null) {
+        ListenableFuture<?> observedFuture = child.isBlocked();
+        synchronized (lock) {
+          if (closed) {
+            break;
+          }
+          listener = new ChildFutureListener(childId, observedFuture, round);
+          listeners[childId] = listener;
+        }
+        // The target is published before registration. addListener also handles a future that
+        // finishes between isBlocked and registration, including inline completion.
+        observedFuture.addListener(listener, directExecutor());
+      } else if (listener.observedFuture.isDone()) {
+        // A completed record may have notified a previous round. Recheck it for this round.
+        listener.run();
+      }
+    }
+
+    Throwable failure;
+    boolean isClosed;
+    synchronized (lock) {
+      failure = firstFailure;
+      isClosed = closed;
+    }
+    if (isClosed || children.isEmpty()) {
+      round.selected.set(-1);
+    } else if (failure != null) {
+      round.selected.setException(failure);
+    }
+    return round.selected;
+  }
+
+  void checkFailure() throws Exception {
+    Throwable failure;
+    synchronized (lock) {
+      failure = firstFailure;
+    }
+    if (failure instanceof Exception) {
+      throw (Exception) failure;
+    }
+    if (failure instanceof Error) {
+      throw (Error) failure;
+    }
+    if (failure != null) {
+      throw new RuntimeException(failure);
+    }
+  }
+
+  void invalidateChild(int childId) {
+    synchronized (lock) {
+      listeners[childId] = null;
+      if (currentRound != null) {
+        currentRound.valid = false;
+        currentRound = null;
+      }
+      nextChild = (childId + 1) % children.size();
+    }
+  }
+
+  void finishChild(int childId) {
+    invalidateChild(childId);
+  }
+
+  void close() {
+    WaitRound round;
+    synchronized (lock) {
+      closed = true;
+      round = currentRound;
+      currentRound = null;
+      if (round != null) {
+        round.valid = false;
+      }
+      for (ChildFutureListener listener : listeners) {
+        if (listener != null) {
+          listener.target = null;
+        }
+      }
+      Arrays.fill(listeners, null);
+    }
+    if (round != null) {
+      round.selected.set(-1);
+    }
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    synchronized (lock) {
+      long size =
+          INSTANCE_SIZE
+              + RamUsageEstimator.shallowSizeOfInstance(Object.class)
+              + RamUsageEstimator.shallowSizeOf(listeners);
+      for (ChildFutureListener listener : listeners) {
+        if (listener != null) {
+          size += LISTENER_SIZE;
+        }
+      }
+      return size + (currentRound == null ? 0 : ROUND_SIZE);
+    }
+  }
+
+  private static final class WaitRound {
+    private final SettableFuture<Integer> selected = SettableFuture.create();
+    private boolean valid = true;
+  }
+
+  private final class ChildFutureListener implements Runnable {
+    private final int childId;
+    private final ListenableFuture<?> observedFuture;
+    private WaitRound target;
+
+    private ChildFutureListener(int childId, ListenableFuture<?> observedFuture, WaitRound target) {
+      this.childId = childId;
+      this.observedFuture = observedFuture;
+      this.target = target;
+    }
+
+    @Override
+    public void run() {
+      Throwable failure = null;
+      try {
+        getDone(observedFuture);
+      } catch (ExecutionException e) {
+        failure = e.getCause();
+      } catch (CancellationException e) {
+        failure = e;
+      }
+      WaitRound round;
+      synchronized (lock) {
+        if (closed || listeners[childId] != this) {
+          return;
+        }
+        if (failure != null && firstFailure == null) {
+          firstFailure = failure;
+        }
+        round = target;
+        if (round == null || !round.valid || round != currentRound) {
+          return;
+        }
+        failure = firstFailure;
+      }
+      // Completing a selection can invoke downstream listeners synchronously. Do not hold lock.
+      if (failure == null) {
+        round.selected.set(childId);
+      } else {
+        round.selected.setException(failure);
+      }
+    }
+  }
+}
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java
index 921d273d04e8a..2cc5deec59ea2 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java
@@ -30,60 +30,102 @@
 
 import java.util.List;
 
+import static com.google.common.util.concurrent.Futures.getDone;
+import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
+
 public class CollectOperator implements ProcessOperator {
   private static final long INSTANCE_SIZE =
       RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class);
 
   private final CommonOperatorContext operatorContext;
   protected final List<Operator> children;
-  private boolean inited = false;
-
-  protected int currentIndex;
+  private final AnyChildBlocked anyChildBlocked;
+  private ListenableFuture<Integer> selectionFuture;
+  private ListenableFuture<Integer> exposedFuture;
+  private int remainingChildren;
+  private boolean closed;
 
   public CollectOperator(CommonOperatorContext operatorContext, List<Operator> children) {
     this.operatorContext = operatorContext;
     this.children = children;
-    this.currentIndex = 0;
+    this.anyChildBlocked = new AnyChildBlocked(children);
+    this.remainingChildren = children.size();
   }
 
   @Override
-  public boolean hasNext() throws Exception {
-    return currentIndex < children.size();
+  public boolean hasNext() {
+    return !closed && remainingChildren > 0;
   }
 
   @Override
   public TsBlock next() throws Exception {
-    if (children.get(currentIndex).hasNextWithTimer()) {
-      return children.get(currentIndex).nextWithTimer();
-    } else {
-      closeCurrentChild(currentIndex);
-      currentIndex++;
+    anyChildBlocked.checkFailure();
+    if (!hasNext()) {
+      return null;
+    }
+    if (selectionFuture == null) {
+      isBlocked();
+    }
+    anyChildBlocked.checkFailure();
+    if (!selectionFuture.isDone()) {
+      return null;
+    }
+    int childId = getDone(selectionFuture);
+    // Consume this selection before invoking any operator method. Another next must acquire a
+    // new readiness observation even when this advancement returns null or an empty block.
+    selectionFuture = null;
+    exposedFuture = null;
+    if (childId < 0 || closed || children.get(childId) == null) {
       return null;
     }
+    Operator child = children.get(childId);
+    try {
+      if (child.hasNextWithTimer()) {
+        TsBlock block = child.nextWithTimer();
+        return block == null ? null : transformBlock(childId, block);
+      }
+      if (child.isFinished()) {
+        closeCurrentChild(childId);
+      }
+      return null;
+    } finally {
+      anyChildBlocked.invalidateChild(childId);
+    }
+  }
+
+  protected TsBlock transformBlock(int childId, TsBlock block) {
+    return block;
   }
 
   protected void closeCurrentChild(int index) throws Exception {
-    children.get(index).close();
+    Operator child = children.get(index);
+    if (child == null) {
+      return;
+    }
+    // Retire the record even if closing its subtree throws.
+    anyChildBlocked.finishChild(index);
     children.set(index, null);
+    remainingChildren--;
+    child.close();
   }
 
   @Override
   public ListenableFuture<?> isBlocked() {
-    if (!inited) {
-      inited = true;
-      for (Operator child : children) {
-        child.isBlocked();
-      }
-    }
-    if (currentIndex >= children.size()) {
+    if (!hasNext()) {
       return NOT_BLOCKED;
     }
-    return children.get(currentIndex).isBlocked();
+    if (selectionFuture == null) {
+      selectionFuture = anyChildBlocked.getAnyChild();
+    }
+    if (exposedFuture == null || exposedFuture.isCancelled()) {
+      exposedFuture = nonCancellationPropagating(selectionFuture);
+    }
+    return exposedFuture;
   }
 
   @Override
-  public boolean isFinished() throws Exception {
-    return currentIndex >= children.size();
+  public boolean isFinished() {
+    return !hasNext();
   }
 
   @Override
@@ -93,35 +135,67 @@ public CommonOperatorContext getOperatorContext() {
 
   @Override
   public void close() throws Exception {
-    for (int i = currentIndex, n = children.size(); i < n; i++) {
-      Operator currentChild = children.get(i);
-      if (currentChild != null) {
+    if (closed) {
+      return;
+    }
+    closed = true;
+    anyChildBlocked.close();
+    selectionFuture = null;
+    exposedFuture = null;
+    Exception firstException = null;
+    for (int i = 0; i < children.size(); i++) {
+      try {
         closeCurrentChild(i);
+      } catch (Exception e) {
+        if (firstException == null) {
+          firstException = e;
+        } else if (firstException != e) {
+          firstException.addSuppressed(e);
+        }
       }
     }
+    if (firstException != null) {
+      throw firstException;
+    }
   }
 
   @Override
   public long calculateMaxPeekMemory() {
-    long maxPeekMemory = 0;
+    long retained = calculateRetainedSizeAfterCallingNext();
+    long peak = retained;
     for (Operator child : children) {
-      maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter());
+      if (child != null) {
+        peak =
+            Math.max(
+                peak,
+                retained
+                    - child.calculateRetainedSizeAfterCallingNext()
+                    + child.calculateMaxPeekMemoryWithCounter());
+      }
     }
-    return maxPeekMemory;
+    return peak;
   }
 
   @Override
   public long calculateMaxReturnSize() {
     long maxReturnSize = 0;
     for (Operator child : children) {
-      maxReturnSize = Math.max(maxReturnSize, child.calculateMaxReturnSize());
+      if (child != null) {
+        maxReturnSize = Math.max(maxReturnSize, child.calculateMaxReturnSize());
+      }
     }
     return maxReturnSize;
   }
 
   @Override
   public long calculateRetainedSizeAfterCallingNext() {
-    return 0L;
+    long retained = 0;
+    for (Operator child : children) {
+      if (child != null) {
+        retained += child.calculateRetainedSizeAfterCallingNext();
+      }
+    }
+    return retained;
   }
 
   @TestOnly
@@ -132,7 +206,9 @@ public List<Operator> getChildren() {
   @Override
   public long ramBytesUsed() {
     return INSTANCE_SIZE
+        + anyChildBlocked.ramBytesUsed()
         + children.stream()
+            .filter(java.util.Objects::nonNull)
             .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject)
             .sum()
         + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext);
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java
index 8bc1c18e9e66a..02850bf0caf16 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java
@@ -45,37 +45,43 @@ public MappingCollectOperator(
       List<List<Integer>> mappings) {
     super(operatorContext, children);
     this.mappings = mappings;
-    outputColumnsCount = mappings.get(0).size();
+    outputColumnsCount = mappings.isEmpty() ? 0 : mappings.get(0).size();
   }
 
   @Override
-  public TsBlock next() throws Exception {
-    if (children.get(currentIndex).hasNextWithTimer()) {
-      TsBlock tsBlock = children.get(currentIndex).nextWithTimer();
-      if (tsBlock == null) {
-        return null;
-      } else {
-        Column[] columns = new Column[outputColumnsCount];
-        List<Integer> mapping = mappings.get(currentIndex);
-        for (int i = 0; i < columns.length; i++) {
-          columns[i] = tsBlock.getColumn(mapping.get(i));
-        }
-        return TsBlock.wrapBlocksWithoutCopy(
-            tsBlock.getPositionCount(),
-            new RunLengthEncodedColumn(
-                CommonOperatorUtils.TIME_COLUMN_TEMPLATE, tsBlock.getPositionCount()),
-            columns);
-      }
-    } else {
-      closeCurrentChild(currentIndex);
-      currentIndex++;
-      return null;
+  protected TsBlock transformBlock(int childId, TsBlock tsBlock) {
+    Column[] columns = new Column[outputColumnsCount];
+    List<Integer> mapping = mappings.get(childId);
+    for (int i = 0; i < columns.length; i++) {
+      columns[i] = tsBlock.getColumn(mapping.get(i));
     }
+    return TsBlock.wrapBlocksWithoutCopy(
+        tsBlock.getPositionCount(),
+        new RunLengthEncodedColumn(
+            CommonOperatorUtils.TIME_COLUMN_TEMPLATE, tsBlock.getPositionCount()),
+        columns);
   }
 
+  @Override
   protected void closeCurrentChild(int index) throws Exception {
-    children.get(index).close();
-    children.set(index, null);
-    mappings.set(index, null);
+    try {
+      super.closeCurrentChild(index);
+    } finally {
+      mappings.set(index, null);
+    }
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    long size =
+        super.ramBytesUsed()
+            + INSTANCE_SIZE
+            - RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class);
+    for (List<Integer> mapping : mappings) {
+      if (mapping != null) {
+        size += RamUsageEstimator.shallowSizeOf(mapping.toArray());
+      }
+    }
+    return size;
   }
 }
diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java
new file mode 100644
index 0000000000000..abf7a82acf3d9
--- /dev/null
+++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java
@@ -0,0 +1,442 @@
+/*
+ * 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.calc.execution.operator.process;
+
+import org.apache.iotdb.calc.execution.operator.CommonOperatorContext;
+import org.apache.iotdb.calc.execution.operator.Operator;
+import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+
+import com.google.common.util.concurrent.ForwardingListenableFuture;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+public class CollectOperatorTest {
+  private static TsBlock block(int... values) {
+    TsBlockBuilder builder =
+        new TsBlockBuilder(Collections.nCopies(values.length, TSDataType.INT32));
+    builder.getTimeColumnBuilder().writeLong(0);
+    for (int i = 0; i < values.length; i++) {
+      builder.getColumnBuilder(i).writeInt(values[i]);
+    }
+    builder.declarePosition();
+    return builder.build();
+  }
+
+  private static CollectOperator collect(Operator... children) {
+    return new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children)));
+  }
+
+  @Test
+  public void testReadyChildAndStableSelectionDuringBackpressure() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    b.ready();
+    CollectOperator collect = collect(a, b);
+    ListenableFuture<?> selected = collect.isBlocked();
+    assertTrue(selected.isDone());
+    a.ready();
+    assertSame(selected, collect.isBlocked());
+    assertSame(b.output, collect.next());
+    assertEquals(0, a.advances);
+    assertEquals(1, b.advances);
+    assertTrue(collect.isBlocked().isDone());
+    assertSame(a.output, collect.next());
+    collect.close();
+  }
+
+  @Test
+  public void testThousandChildrenAndOneListenerPerPendingPeriod() throws Exception {
+    for (int count : new int[] {100, 1000}) {
+      List<Operator> children = new ArrayList<>();
+      for (int i = 0; i < count; i++) {
+        children.add(new Child(block(i)));
+      }
+      CollectOperator collect = new CollectOperator(new TestContext(), children);
+      Child ready = (Child) children.get(count - 1);
+      for (int round = 0; round < 5; round++) {
+        ListenableFuture<?> wait = collect.isBlocked();
+        assertFalse(wait.isDone());
+        for (int repeat = 0; repeat < 10; repeat++) {
+          assertSame(wait, collect.isBlocked());
+        }
+        ready.ready();
+        wait.get(5, TimeUnit.SECONDS);
+        assertSame(ready.output, collect.next());
+        ready.reset();
+      }
+      for (int i = 0; i < count - 1; i++) {
+        Child child = (Child) children.get(i);
+        assertEquals(1, child.waitCalls);
+        assertEquals(1, child.wait.registrations.get());
+      }
+      collect.close();
+    }
+  }
+
+  @Test
+  public void testPendingListenerRetargetedAndStaleCompletionIgnored() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    CollectOperator collect = collect(a, b);
+    ListenableFuture<?> first = collect.isBlocked();
+    b.ready();
+    assertSame(b.output, collect.next());
+    Runnable stale = b.wait.listener;
+    b.reset();
+    ListenableFuture<?> second = collect.isBlocked();
+    assertFalse(second.isDone());
+    stale.run();
+    assertFalse(second.isDone());
+    a.ready();
+    second.get(5, TimeUnit.SECONDS);
+    assertEquals(1, a.wait.registrations.get());
+    assertSame(a.output, collect.next());
+    assertTrue(first.isDone());
+    collect.close();
+  }
+
+  @Test
+  public void testNullEmptyAndUnfinishedChildRefreshReadiness() throws Exception {
+    Child child = new Child(null);
+    child.ready();
+    CollectOperator collect = collect(child);
+    assertTrue(collect.isBlocked().isDone());
+    assertNull(collect.next());
+    child.reset();
+    assertFalse(collect.isBlocked().isDone());
+    child.output = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32)).build();
+    child.ready();
+    assertTrue(collect.next().isEmpty());
+    child.reset();
+    child.hasNext = false;
+    child.ready();
+    assertNull(collect.next());
+    assertFalse(collect.isFinished());
+    assertEquals(0, child.closes);
+    child.finished = true;
+    assertNull(collect.next());
+    assertTrue(collect.isFinished());
+    assertEquals(1, child.closes);
+    assertTrue(collect.isBlocked().isDone());
+  }
+
+  @Test
+  public void testCancellationDoesNotPropagateToChild() throws Exception {
+    Child child = new Child(block(1));
+    CollectOperator collect = collect(child);
+    ListenableFuture<?> exposed = collect.isBlocked();
+    assertTrue(exposed.cancel(true));
+    assertFalse(child.wait.isCancelled());
+    ListenableFuture<?> replacement = collect.isBlocked();
+    assertFalse(replacement.isCancelled());
+    child.ready();
+    replacement.get(5, TimeUnit.SECONDS);
+    assertSame(child.output, collect.next());
+    collect.close();
+  }
+
+  @Test
+  public void testFailureOfUnselectedChildIsNotLost() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    a.ready();
+    CollectOperator collect = collect(a, b);
+    assertTrue(collect.isBlocked().isDone());
+    IOException expected = new IOException("child failure");
+    b.wait.delegate.setException(expected);
+    try {
+      collect.next();
+      fail("must propagate the failed child even when another child was selected");
+    } catch (IOException e) {
+      assertSame(expected, e);
+    } finally {
+      collect.close();
+    }
+    assertEquals(0, a.advances);
+  }
+
+  @Test
+  public void testChildCancellationPropagates() throws Exception {
+    Child child = new Child(block(1));
+    CollectOperator collect = collect(child);
+    ListenableFuture<?> wait = collect.isBlocked();
+    child.wait.cancel(false);
+    assertTrue(wait.isDone());
+    try {
+      collect.next();
+      fail("must propagate child cancellation");
+    } catch (CancellationException expected) {
+      assertEquals(0, child.advances);
+    } finally {
+      collect.close();
+    }
+  }
+
+  @Test
+  public void testCompletionDuringRegistrationAndOnAnotherThread() throws Exception {
+    Child a = new Child(block(1));
+    a.wait.beforeRegistration = a::ready;
+    CollectOperator collect = collect(a);
+    assertTrue(collect.isBlocked().isDone());
+    assertSame(a.output, collect.next());
+    a.reset();
+    ListenableFuture<?> pending = collect.isBlocked();
+    Thread completion = new Thread(a::ready);
+    completion.start();
+    pending.get(5, TimeUnit.SECONDS);
+    completion.join();
+    assertSame(a.output, collect.next());
+    collect.close();
+  }
+
+  @Test
+  public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    a.closeFailure = new IOException("close failure");
+    CollectOperator collect = collect(a, b);
+    ListenableFuture<?> pending = collect.isBlocked();
+    try {
+      collect.close();
+      fail("must propagate close failure");
+    } catch (IOException e) {
+      assertSame(a.closeFailure, e);
+    }
+    pending.get(5, TimeUnit.SECONDS);
+    assertEquals(1, a.closes);
+    assertEquals(1, b.closes);
+    a.ready();
+    b.ready();
+    assertTrue(collect.isFinished());
+    assertNull(collect.next());
+    collect.close();
+    assertEquals(1, b.closes);
+  }
+
+  @Test
+  public void testMappingUsesActualChildIdentity() throws Exception {
+    Child a = new Child(block(10, 11));
+    Child b = new Child(block(20, 21));
+    b.ready();
+    List<List<Integer>> mappings = new ArrayList<>();
+    mappings.add(Collections.singletonList(0));
+    mappings.add(Collections.singletonList(1));
+    MappingCollectOperator collect =
+        new MappingCollectOperator(
+            new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings);
+    assertTrue(collect.isBlocked().isDone());
+    assertEquals(21, collect.next().getColumn(0).getInt(0));
+    a.ready();
+    assertEquals(10, collect.next().getColumn(0).getInt(0));
+    collect.close();
+    assertNull(mappings.get(0));
+    assertNull(mappings.get(1));
+  }
+
+  @Test
+  public void testMemoryIncludesAllLiveRetainedStates() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    a.retained = 30;
+    a.peak = 100;
+    a.returnSize = 10;
+    b.retained = 50;
+    b.peak = 80;
+    b.returnSize = 20;
+    CollectOperator collect = collect(a, b);
+    assertEquals(80, collect.calculateRetainedSizeAfterCallingNext());
+    assertEquals(150, collect.calculateMaxPeekMemory());
+    assertEquals(20, collect.calculateMaxReturnSize());
+    assertTrue(collect.ramBytesUsed() > 0);
+    a.hasNext = false;
+    a.finished = true;
+    a.ready();
+    collect.next();
+    assertEquals(50, collect.calculateRetainedSizeAfterCallingNext());
+    assertEquals(80, collect.calculateMaxPeekMemory());
+    collect.close();
+    assertEquals(0, collect.calculateMaxPeekMemory());
+    assertEquals(0, collect.calculateMaxReturnSize());
+    assertEquals(0, collect.calculateRetainedSizeAfterCallingNext());
+  }
+
+  private static final class CountingFuture
+      extends ForwardingListenableFuture.SimpleForwardingListenableFuture<Void> {
+    private final SettableFuture<Void> delegate;
+    private final AtomicInteger registrations = new AtomicInteger();
+    private Runnable beforeRegistration;
+    private Runnable listener;
+
+    private CountingFuture() {
+      this(SettableFuture.create());
+    }
+
+    private CountingFuture(SettableFuture<Void> delegate) {
+      super(delegate);
+      this.delegate = delegate;
+    }
+
+    @Override
+    public void addListener(Runnable listener, Executor executor) {
+      this.listener = listener;
+      registrations.incrementAndGet();
+      if (beforeRegistration != null) {
+        beforeRegistration.run();
+      }
+      super.addListener(listener, executor);
+    }
+  }
+
+  private static final class TestContext extends CommonOperatorContext {
+    private TestContext() {
+      super(0, new PlanNodeId("collect"), "test");
+    }
+
+    @Override
+    public MemoryReservationManager getMemoryReservationContext() {
+      return null;
+    }
+
+    @Override
+    public int getFragmentId() {
+      return 0;
+    }
+
+    @Override
+    public int getPipelineId() {
+      return 0;
+    }
+
+    @Override
+    public long ramBytesUsed() {
+      return 0;
+    }
+  }
+
+  private static final class Child implements Operator {
+    private final Thread driver = Thread.currentThread();
+    private final CommonOperatorContext context = new TestContext();
+    private CountingFuture wait = new CountingFuture();
+    private TsBlock output;
+    private boolean hasNext = true;
+    private boolean finished;
+    private int advances;
+    private int waitCalls;
+    private int closes;
+    private IOException closeFailure;
+    private long retained, peak, returnSize;
+
+    private Child(TsBlock output) {
+      this.output = output;
+    }
+
+    private void ready() {
+      wait.delegate.set(null);
+    }
+
+    private void reset() {
+      wait = new CountingFuture();
+    }
+
+    private void checkDriver() {
+      assertSame(driver, Thread.currentThread());
+    }
+
+    @Override
+    public CommonOperatorContext getOperatorContext() {
+      return context;
+    }
+
+    @Override
+    public ListenableFuture<?> isBlocked() {
+      checkDriver();
+      waitCalls++;
+      return wait;
+    }
+
+    @Override
+    public boolean hasNext() {
+      checkDriver();
+      return hasNext;
+    }
+
+    @Override
+    public TsBlock next() {
+      checkDriver();
+      advances++;
+      return output;
+    }
+
+    @Override
+    public boolean isFinished() {
+      checkDriver();
+      return finished;
+    }
+
+    @Override
+    public void close() throws Exception {
+      checkDriver();
+      closes++;
+      if (closeFailure != null) throw closeFailure;
+    }
+
+    @Override
+    public long calculateMaxPeekMemory() {
+      return peak;
+    }
+
+    @Override
+    public long calculateMaxReturnSize() {
+      return returnSize;
+    }
+
+    @Override
+    public long calculateRetainedSizeAfterCallingNext() {
+      return retained;
+    }
+
+    @Override
+    public long ramBytesUsed() {
+      return 0;
+    }
+  }
+}
diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
index 73b27b0622304..c0df17a3122ff 100644
--- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
+++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
@@ -3940,4 +3940,6 @@ private DataNodeQueryMessages() {}
   public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB =
       "Only InMemoryDeviceEntryDataSet supports get inline device entries";
 
+  public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "Plan node %s is already registered on memory side %s, not %s";
+  public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "Memory reservation is not registered for query %s, fragment %s, plan node %s";
 }
diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
index ed1fefe479af7..b3b20f477c714 100644
--- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
+++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
@@ -4696,4 +4696,6 @@ private DataNodeQueryMessages() {}
   public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB =
       "只有 InMemoryDeviceEntryDataSet 支持获取内存中的设备条目";
 
+  public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "PlanNode %s 已注册到内存侧 %s，不能注册到 %s";
+  public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "查询 %s、Fragment %s、PlanNode %s 尚未注册内存账目";
 }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java
index 5812d4e37cdd3..cb8c681e4e9da 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java
@@ -27,6 +27,7 @@
 import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle;
 import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager;
 import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult;
+import org.apache.iotdb.db.queryengine.execution.memory.MemorySide;
 import org.apache.iotdb.db.utils.CommonUtils;
 import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
 
@@ -85,7 +86,8 @@ public class SharedTsBlockQueue {
   private LocalSinkChannel sinkChannel;
 
   private long maxBytesCanReserve =
-      IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+      MemorySide.OUTPUT.getMaxBytes(
+          IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance());
   private final boolean isHighestPriority;
 
   private volatile Throwable abortedCause = null;
@@ -265,7 +267,10 @@ public ListenableFuture<Void> add(TsBlock tsBlock) {
       localMemoryManager
           .getQueryPool()
           .registerPlanNodeIdToQueryMemoryMap(
-              localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId);
+              localFragmentInstanceId.queryId,
+              fullFragmentInstanceId,
+              localPlanNodeId,
+              MemorySide.OUTPUT);
       alreadyRegistered = true;
     }
     MemoryReservationResult reserveResult =
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
index e54062158a3f1..1de3ccb897df7 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
@@ -22,6 +22,7 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
 import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager;
+import org.apache.iotdb.db.queryengine.execution.memory.MemorySide;
 import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet;
 import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
 
@@ -76,7 +77,8 @@ public class ShuffleSinkHandle implements ISinkHandle {
 
   /** max bytes this ShuffleSinkHandle can reserve. */
   private long maxBytesCanReserve =
-      IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+      MemorySide.OUTPUT.getMaxBytes(
+          IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance());
 
   private static final long INSTANCE_SIZE =
       RamUsageEstimator.shallowSizeOfInstance(ShuffleSinkHandle.class)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
index d24a1552a5b22..0dec1cbf07f26 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
@@ -30,6 +30,7 @@
 import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SinkListener;
 import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager;
 import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult;
+import org.apache.iotdb.db.queryengine.execution.memory.MemorySide;
 import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet;
 import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet;
 import org.apache.iotdb.db.utils.SetThreadName;
@@ -131,7 +132,8 @@ private TsBlockInfo(TsBlock tsBlock, ByteBuffer serializedTsBlock, long tsBlockS
 
   /** max bytes this SinkChannel can reserve. */
   private long maxBytesCanReserve =
-      IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+      MemorySide.OUTPUT.getMaxBytes(
+          IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance());
 
   private final boolean isHighestPriority;
 
@@ -235,7 +237,10 @@ public SinkChannel(
     localMemoryManager
         .getQueryPool()
         .registerPlanNodeIdToQueryMemoryMap(
-            localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId);
+            localFragmentInstanceId.queryId,
+            fullFragmentInstanceId,
+            localPlanNodeId,
+            MemorySide.OUTPUT);
   }
 
   @Override
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
index aed3f9533fb1d..b1dc2781ed658 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
@@ -31,6 +31,7 @@
 import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SourceHandleListener;
 import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager;
 import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult;
+import org.apache.iotdb.db.queryengine.execution.memory.MemorySide;
 import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet;
 import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet;
 import org.apache.iotdb.db.utils.SetThreadName;
@@ -117,7 +118,8 @@ public class SourceHandle implements ISourceHandle {
 
   /** max bytes this SourceHandle can reserve. */
   private long maxBytesCanReserve =
-      IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+      MemorySide.INPUT.getMaxBytes(
+          IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance());
 
   /**
    * this is set to true after calling isBlocked() at least once which indicates that this
@@ -223,7 +225,10 @@ public SourceHandle(
     localMemoryManager
         .getQueryPool()
         .registerPlanNodeIdToQueryMemoryMap(
-            localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId);
+            localFragmentInstanceId.queryId,
+            fullFragmentInstanceId,
+            localPlanNodeId,
+            MemorySide.INPUT);
   }
 
   @Override
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
index b2d35568fa31c..56c9974764a5f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
@@ -40,6 +40,7 @@
 import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
@@ -57,8 +58,8 @@ public static class MemoryReservationFuture<V> extends AbstractFuture<V> {
     private final long bytesToReserve;
 
     /**
-     * MemoryReservationFuture is created when SinkHandle or SourceHandle tries to reserve memory
-     * from pool. This field is max Bytes that SinkHandle or SourceHandle can reserve.
+     * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the
+     * pool. This limit is shared by all plan nodes registered on the same side of that fragment.
      */
     private final long maxBytesCanReserve;
 
@@ -154,8 +155,32 @@ public long getReservedBytes() {
   private final IMemoryBlock memoryBlock;
   private final long maxBytesPerFragmentInstance;
 
-  /** queryId -> fragmentInstanceId -> planNodeId -> bytesReserved. */
-  private final Map<String, Map<String, Map<String, Long>>> queryMemoryReservations =
+  private static class PlanNodeReservation {
+    private final MemorySide side;
+    private volatile long bytesReserved;
+
+    private PlanNodeReservation(MemorySide side) {
+      this.side = side;
+    }
+
+    @Override
+    public String toString() {
+      return side + ":" + bytesReserved;
+    }
+  }
+
+  private static class FragmentMemoryReservation {
+    private final Map<String, PlanNodeReservation> planNodes = new ConcurrentHashMap<>();
+    private final AtomicLong inputUsedBytes = new AtomicLong();
+    private final AtomicLong outputUsedBytes = new AtomicLong();
+
+    private AtomicLong getUsedBytes(MemorySide side) {
+      return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes;
+    }
+  }
+
+  /** queryId -> fragmentInstanceId -> FI side totals and planNode details. */
+  private final Map<String, Map<String, FragmentMemoryReservation>> queryMemoryReservations =
       new ConcurrentHashMap<>();
 
   private final Queue<MemoryReservationFuture<Void>> memoryReservationFutures =
@@ -204,12 +229,27 @@ public int getMemoryReservationSize() {
    * planNodeId to queryMemoryReservationsMap first.
    */
   public void registerPlanNodeIdToQueryMemoryMap(
-      String queryId, String fragmentInstanceId, String planNodeId) {
+      String queryId, String fragmentInstanceId, String planNodeId, MemorySide side) {
+    java.util.Objects.requireNonNull(side);
     synchronized (queryMemoryReservations) {
       queryMemoryReservations
           .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>())
-          .computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>())
-          .putIfAbsent(planNodeId, 0L);
+          .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation())
+          .planNodes
+          .compute(
+              planNodeId,
+              (key, reservation) -> {
+                if (reservation != null && reservation.side != side) {
+                  throw new IllegalArgumentException(
+                      String.format(
+                          DataNodeQueryMessages
+                              .EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625,
+                          planNodeId,
+                          reservation.side,
+                          side));
+                }
+                return reservation == null ? new PlanNodeReservation(side) : reservation;
+              });
     }
   }
 
@@ -224,15 +264,19 @@ public void registerPlanNodeIdToQueryMemoryMap(
    */
   public void deRegisterFragmentInstanceFromQueryMemoryMap(
       String queryId, String fragmentInstanceId, boolean forceDeregister) {
-    Map<String, Map<String, Long>> queryRelatedMemory = queryMemoryReservations.get(queryId);
+    Map<String, FragmentMemoryReservation> queryRelatedMemory =
+        queryMemoryReservations.get(queryId);
     if (queryRelatedMemory != null) {
-      Map<String, Long> fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
+      FragmentMemoryReservation fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
       boolean hasPotentialMemoryLeak = false;
       // fragmentRelatedMemory could be null if the FI has not reserved any memory(For example,
       // next() of root operator returns no data)
       if (fragmentRelatedMemory != null) {
         hasPotentialMemoryLeak =
-            fragmentRelatedMemory.values().stream().anyMatch(value -> value != 0L);
+            fragmentRelatedMemory.inputUsedBytes.get() != 0L
+                || fragmentRelatedMemory.outputUsedBytes.get() != 0L
+                || fragmentRelatedMemory.planNodes.values().stream()
+                    .anyMatch(value -> value.bytesReserved != 0L);
       }
       if (!forceDeregister && hasPotentialMemoryLeak) {
         // If hasPotentialMemoryLeak is true, it means that LocalSourceChannel/LocalSourceHandles
@@ -252,9 +296,9 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap(
       }
       if (hasPotentialMemoryLeak) {
         // hasPotentialMemoryLeak means that fragmentRelatedMemory is not null
-        List<Map.Entry<String, Long>> invalidEntryList =
-            fragmentRelatedMemory.entrySet().stream()
-                .filter(entry -> entry.getValue() != 0L)
+        List<Map.Entry<String, PlanNodeReservation>> invalidEntryList =
+            fragmentRelatedMemory.planNodes.entrySet().stream()
+                .filter(entry -> entry.getValue().bytesReserved != 0L)
                 .collect(Collectors.toList());
         throw new MemoryLeakException(
             String.format(
@@ -390,25 +434,26 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l
     Validate.notNull(queryId, DataNodeQueryMessages.EXCEPTION_QUERYID_CAN_NOT_BE_NULL_DOT_16639DBE);
     Validate.isTrue(bytes > 0L);
 
-    try {
-      queryMemoryReservations
-          .get(queryId)
-          .get(fragmentInstanceId)
-          .computeIfPresent(
-              planNodeId,
-              (k, reservedMemory) -> {
-                if (reservedMemory < bytes) {
-                  throw new IllegalArgumentException(
-                      DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED);
-                }
-                return reservedMemory - bytes;
-              });
-    } catch (NullPointerException e) {
+    Map<String, FragmentMemoryReservation> fragments = queryMemoryReservations.get(queryId);
+    FragmentMemoryReservation fragment =
+        fragments == null ? null : fragments.get(fragmentInstanceId);
+    if (fragment == null) {
       throw new IllegalArgumentException(
           DataNodeQueryMessages
-              .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2,
-          e);
+              .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2);
     }
+    fragment.planNodes.compute(
+        planNodeId,
+        (key, reservation) -> {
+          requireReservation(reservation, queryId, fragmentInstanceId, planNodeId);
+          if (reservation.bytesReserved < bytes) {
+            throw new IllegalArgumentException(
+                DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED);
+          }
+          reservation.bytesReserved -= bytes;
+          fragment.getUsedBytes(reservation.side).addAndGet(-bytes);
+          return reservation;
+        });
 
     memoryBlock.release(bytes);
 
@@ -440,16 +485,34 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l
   }
 
   public long getQueryMemoryReservedBytes(String queryId) {
-    if (!queryMemoryReservations.containsKey(queryId)) {
+    Map<String, FragmentMemoryReservation> fragments = queryMemoryReservations.get(queryId);
+    if (fragments == null) {
       return 0L;
     }
     long sum = 0;
-    for (Map<String, Long> map : queryMemoryReservations.get(queryId).values()) {
-      sum = sum + map.values().stream().reduce(0L, Long::sum);
+    for (FragmentMemoryReservation fragment : fragments.values()) {
+      sum += fragment.inputUsedBytes.get() + fragment.outputUsedBytes.get();
     }
     return sum;
   }
 
+  public long getFragmentMemoryReservedBytes(
+      String queryId, String fragmentInstanceId, MemorySide side) {
+    Map<String, FragmentMemoryReservation> fragments = queryMemoryReservations.get(queryId);
+    FragmentMemoryReservation fragment =
+        fragments == null ? null : fragments.get(fragmentInstanceId);
+    return fragment == null ? 0L : fragment.getUsedBytes(side).get();
+  }
+
+  public long getPlanNodeMemoryReservedBytes(
+      String queryId, String fragmentInstanceId, String planNodeId) {
+    Map<String, FragmentMemoryReservation> fragments = queryMemoryReservations.get(queryId);
+    FragmentMemoryReservation fragment =
+        fragments == null ? null : fragments.get(fragmentInstanceId);
+    PlanNodeReservation reservation = fragment == null ? null : fragment.planNodes.get(planNodeId);
+    return reservation == null ? 0L : reservation.bytesReserved;
+  }
+
   public long getReservedBytes() {
     return memoryBlock.getUsedMemoryInBytes();
   }
@@ -460,22 +523,68 @@ public boolean tryReserve(
       String planNodeId,
       long bytesToReserve,
       long maxBytesCanReserve) {
+    FragmentMemoryReservation fragment =
+        getFragmentReservation(queryId, fragmentInstanceId, planNodeId);
+    PlanNodeReservation reservation =
+        requireReservation(
+            fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId);
     long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve);
-    long queryRemainingBytes =
-        maxBytesCanReserve
-            - queryMemoryReservations
-                .get(queryId)
-                .get(fragmentInstanceId)
-                .merge(planNodeId, bytesToReserve, Long::sum);
-    return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0;
+    long sideUsedBytes = fragment.getUsedBytes(reservation.side).addAndGet(bytesToReserve);
+    fragment.planNodes.compute(
+        planNodeId,
+        (key, value) -> {
+          value.bytesReserved += bytesToReserve;
+          return value;
+        });
+    return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes()
+        && sideUsedBytes <= maxBytesCanReserve;
   }
 
   private void rollbackReserve(
       String queryId, String fragmentInstanceId, String planNodeId, long bytesToReserve) {
-    queryMemoryReservations
-        .get(queryId)
-        .get(fragmentInstanceId)
-        .merge(planNodeId, -bytesToReserve, Long::sum);
+    FragmentMemoryReservation fragment =
+        getFragmentReservation(queryId, fragmentInstanceId, planNodeId);
+    fragment.planNodes.compute(
+        planNodeId,
+        (key, reservation) -> {
+          reservation.bytesReserved -= bytesToReserve;
+          fragment.getUsedBytes(reservation.side).addAndGet(-bytesToReserve);
+          return reservation;
+        });
     memoryBlock.release(bytesToReserve);
   }
+
+  private FragmentMemoryReservation getFragmentReservation(
+      String queryId, String fragmentInstanceId, String planNodeId) {
+    Map<String, FragmentMemoryReservation> fragments = queryMemoryReservations.get(queryId);
+    FragmentMemoryReservation fragment =
+        fragments == null ? null : fragments.get(fragmentInstanceId);
+    if (fragment == null) {
+      throw new IllegalArgumentException(
+          String.format(
+              DataNodeQueryMessages
+                  .EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F,
+              queryId,
+              fragmentInstanceId,
+              planNodeId));
+    }
+    return fragment;
+  }
+
+  private PlanNodeReservation requireReservation(
+      PlanNodeReservation reservation,
+      String queryId,
+      String fragmentInstanceId,
+      String planNodeId) {
+    if (reservation == null) {
+      throw new IllegalArgumentException(
+          String.format(
+              DataNodeQueryMessages
+                  .EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F,
+              queryId,
+              fragmentInstanceId,
+              planNodeId));
+    }
+    return reservation;
+  }
 }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java
new file mode 100644
index 0000000000000..a59ef20d26848
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java
@@ -0,0 +1,31 @@
+/*
+ * 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.queryengine.execution.memory;
+
+/** The exchange budget charged by a registered plan node. */
+public enum MemorySide {
+  INPUT,
+  OUTPUT;
+
+  public long getMaxBytes(long fragmentExchangeMaxBytes) {
+    long inputMaxBytes = fragmentExchangeMaxBytes / 2;
+    return this == INPUT ? inputMaxBytes : fragmentExchangeMaxBytes - inputMaxBytes;
+  }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java
index cb2f6f02acb57..fd539be7467e2 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java
@@ -343,7 +343,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte
       Operator child = node.getChildren().get(0).accept(this, context);
       List<Operator> children = new ArrayList<>(1);
       children.add(child);
-      sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve());
+      sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes());
       context.getDriverContext().setSink(sinkHandle);
       return new IdentitySinkOperator(
           operatorContext, children, downStreamChannelIndex, sinkHandle);
@@ -385,7 +385,7 @@ public Operator visitTableExchange(ExchangeNode node, LocalExecutionPlanContext
     if (!isSameNode) {
       context.addExchangeSumNum(1);
     }
-    sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve());
+    sourceHandle.setMaxBytesCanReserve(context.getInputMaxBytes());
     ExchangeOperator exchangeOperator =
         new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId());
     context.addExchangeOperator(exchangeOperator);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java
index 7a5f52749aa4a..bfed28ccc4ca4 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java
@@ -32,8 +32,10 @@
 import org.apache.iotdb.db.queryengine.execution.driver.DriverContext;
 import org.apache.iotdb.db.queryengine.execution.driver.SchemaDriverContext;
 import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
+import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle;
 import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
+import org.apache.iotdb.db.queryengine.execution.memory.MemorySide;
 import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator;
 import org.apache.iotdb.db.queryengine.plan.analyze.TemplatedInfo;
 import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider;
@@ -47,8 +49,6 @@
 import org.apache.tsfile.read.filter.basic.Filter;
 import org.apache.tsfile.utils.Binary;
 import org.apache.tsfile.utils.Pair;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nullable;
 
@@ -70,7 +70,6 @@
 // Attention: We should use thread-safe data structure for members that are shared by all pipelines
 public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext {
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(LocalExecutionPlanContext.class);
   // Save operators in this pipeline, a new one will be created when creating another pipeline
   private final DriverContext driverContext;
   private final AtomicInteger nextOperatorId;
@@ -83,6 +82,8 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext
   private List<PipelineDriverFactory> pipelineDriverFactories;
   private List<ExchangeOperator> exchangeOperatorList = new ArrayList<>();
   private int exchangeSumNum = 0;
+  private final long inputMaxBytes;
+  private final long outputMaxBytes;
 
   private List<TSDataType> cachedDataTypes;
 
@@ -110,6 +111,10 @@ public LocalExecutionPlanContext(
       TypeProvider typeProvider,
       FragmentInstanceContext instanceContext,
       DataNodeQueryContext dataNodeQueryContext) {
+    long fragmentExchangeMaxBytes =
+        IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+    this.inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes);
+    this.outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes);
     this.typeProvider = typeProvider;
     this.allSensorsMap = new ConcurrentHashMap<>();
     this.nextOperatorId = new AtomicInteger(0);
@@ -121,6 +126,8 @@ public LocalExecutionPlanContext(
 
   // For creating subContext, differ from parent context mainly in driver context
   public LocalExecutionPlanContext(LocalExecutionPlanContext parentContext) {
+    this.inputMaxBytes = parentContext.inputMaxBytes;
+    this.outputMaxBytes = parentContext.outputMaxBytes;
     this.nextOperatorId = parentContext.nextOperatorId;
     this.typeProvider = parentContext.typeProvider;
     this.allSensorsMap = parentContext.allSensorsMap;
@@ -142,6 +149,10 @@ public LocalExecutionPlanContext(
       final TypeProvider typeProvider,
       final FragmentInstanceContext instanceContext,
       final ISchemaRegion schemaRegion) {
+    long fragmentExchangeMaxBytes =
+        IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+    this.inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes);
+    this.outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes);
     this.allSensorsMap = new ConcurrentHashMap<>();
     this.typeProvider = typeProvider;
     this.nextOperatorId = new AtomicInteger(0);
@@ -247,10 +258,12 @@ public void setExchangeSumNum(int exchangeSumNum) {
     this.exchangeSumNum = exchangeSumNum;
   }
 
-  public long getMaxBytesOneHandleCanReserve() {
-    long maxBytesPerFI =
-        IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
-    return exchangeSumNum == 0 ? maxBytesPerFI : maxBytesPerFI / exchangeSumNum;
+  public long getInputMaxBytes() {
+    return inputMaxBytes;
+  }
+
+  public long getOutputMaxBytes() {
+    return outputMaxBytes;
   }
 
   public void addExchangeSumNum(int addValue) {
@@ -261,16 +274,16 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) {
     this.exchangeOperatorList.add(exchangeOperator);
   }
 
-  public void setMaxBytesOneHandleCanReserve() {
-    long maxBytesOneHandleCanReserve = getMaxBytesOneHandleCanReserve();
-    LOGGER.debug(
-        DataNodeQueryMessages
-            .MAXBYTESONEHANDLECANRESERVE_FOR_EXCHANGEOPERATOR_IS_ARG_EXCHANGESUMNUM_IS_ARG,
-        maxBytesOneHandleCanReserve,
-        exchangeSumNum);
+  public void setExchangeMemoryLimits() {
     exchangeOperatorList.forEach(
-        exchangeOperator ->
-            exchangeOperator.getSourceHandle().setMaxBytesCanReserve(maxBytesOneHandleCanReserve));
+        exchangeOperator -> {
+          exchangeOperator
+              .getSourceHandle()
+              .setMaxBytesCanReserve(
+                  exchangeOperator.getSourceHandle() instanceof PipelineSourceHandle
+                      ? outputMaxBytes
+                      : inputMaxBytes);
+        });
   }
 
   public Set<String> getAllSensors(IDeviceID deviceId, String sensorId) {
@@ -303,6 +316,7 @@ public void setISink(ISink sink) {
     checkArgument(
         driverContext.getSink() == null,
         DataNodeQueryMessages.EXCEPTION_THERE_MUST_BE_AT_MOST_ONE_SINKNODE_A965AFE7);
+    sink.setMaxBytesCanReserve(outputMaxBytes);
     driverContext.setSink(sink);
   }
 
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java
index feb6b0cb74afd..1a85cb148c4f4 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java
@@ -133,7 +133,7 @@ public List<PipelineDriverFactory> plan(
     context.getTimePartitions().ifPresent(instanceContext::setTimePartitions);
 
     // set maxBytes one SourceHandle can reserve after visiting the whole tree
-    context.setMaxBytesOneHandleCanReserve();
+    context.setExchangeMemoryLimits();
 
     return context.getPipelineDriverFactories();
   }
@@ -166,7 +166,7 @@ public List<PipelineDriverFactory> plan(
     context.addPipelineDriverFactory(root, context.getDriverContext(), 0);
 
     // set maxBytes one SourceHandle can reserve after visiting the whole tree
-    context.setMaxBytesOneHandleCanReserve();
+    context.setExchangeMemoryLimits();
 
     return context.getPipelineDriverFactories();
   }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java
index b6a57db5445d9..d643ec03752e1 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java
@@ -2577,7 +2577,7 @@ public Operator visitExchange(ExchangeNode node, LocalExecutionPlanContext conte
     if (!isSameNode) {
       context.addExchangeSumNum(1);
     }
-    sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve());
+    sourceHandle.setMaxBytesCanReserve(context.getInputMaxBytes());
     ExchangeOperator exchangeOperator =
         new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId());
     context.addExchangeOperator(exchangeOperator);
@@ -2616,7 +2616,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte
             node.getPlanNodeId().getId(),
             context.getInstanceContext());
     List<Operator> children = dealWithConsumeChildrenOneByOneNode(node, context);
-    sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve());
+    sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes());
     context.getDriverContext().setSink(sinkHandle);
 
     return new IdentitySinkOperator(operatorContext, children, downStreamChannelIndex, sinkHandle);
@@ -2658,7 +2658,7 @@ public Operator visitShuffleSink(ShuffleSinkNode node, LocalExecutionPlanContext
             localInstanceId.toThrift(),
             node.getPlanNodeId().getId(),
             context.getInstanceContext());
-    sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve());
+    sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes());
     context.getDriverContext().setSink(sinkHandle);
 
     return new ShuffleHelperOperator(operatorContext, children, downStreamChannelIndex, sinkHandle);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java
index 66f6da892f3e8..d7461431824e2 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java
@@ -630,11 +630,27 @@ public List<PlanNode> visitTopK(TopKNode node, PlanContext context) {
       if (canTopKEliminated(node.getOrderingScheme(), node.getCount(), childrenNodes.get(0))) {
         return childrenNodes;
       }
-      node.setChildren(Collections.singletonList(childrenNodes.get(0)));
-      return Collections.singletonList(node);
+      TopKNode physical =
+          new TopKNode(
+              node.getPlanNodeId(),
+              childrenNodes,
+              node.getOrderingScheme(),
+              node.getCount(),
+              node.getOutputSymbols(),
+              canSortEliminated(
+                  node.getOrderingScheme(),
+                  nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId())));
+      physical.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId());
+      return Collections.singletonList(physical);
     }
 
-    TopKNode newTopKNode = (TopKNode) node.clone();
+    TopKNode newTopKNode =
+        new TopKNode(
+            node.getPlanNodeId(),
+            node.getOrderingScheme(),
+            node.getCount(),
+            node.getOutputSymbols(),
+            true);
     // root TopK is not a runtime filter producer.
     newTopKNode.setTopKRuntimeFilterSourceId(null);
     for (PlanNode child : childrenNodes) {
@@ -649,7 +665,8 @@ public List<PlanNode> visitTopK(TopKNode node, PlanContext context) {
                 node.getOrderingScheme(),
                 node.getCount(),
                 node.getOutputSymbols(),
-                node.isChildrenDataInOrder());
+                canSortEliminated(
+                    node.getOrderingScheme(), nodeOrderingMap.get(child.getPlanNodeId())));
         regionTopK.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId());
         newChild = regionTopK;
       }
@@ -689,6 +706,12 @@ public List<PlanNode> visitGroup(GroupNode node, PlanContext context) {
     for (PlanNode child : childrenNodes) {
       if (canSortEliminated(node.getOrderingScheme(), nodeOrderingMap.get(child.getPlanNodeId()))) {
         result.add(child);
+      } else if (groupingInvalidatedByCollect(
+          child, node.getOrderingScheme().getOrderBy().subList(0, node.getPartitionKeyCount()))) {
+        SortNode sort =
+            new SortNode(queryId.genPlanNodeId(), child, node.getOrderingScheme(), false, false);
+        nodeOrderingMap.put(sort.getPlanNodeId(), sort.getOrderingScheme());
+        result.add(sort);
       } else {
         GroupNode subGroupNode =
             new GroupNode(
@@ -1100,7 +1123,7 @@ public List<PlanNode> visitExternalTsFileAggregationScan(
     if (needFinalAggregation) {
       OrderingScheme childOrdering = nodeOrderingMap.get(result.get(0).getPlanNodeId());
       finalAggregation.setChild(mergeChildrenViaCollectOrMergeSort(childOrdering, result));
-      return Collections.singletonList(finalAggregation);
+      return Collections.singletonList(validateAggregationGrouping(finalAggregation));
     }
     return result;
   }
@@ -2045,7 +2068,7 @@ public List<PlanNode> visitAggregation(AggregationNode node, PlanContext context
     }
 
     List<PlanNode> childrenNodes = node.getChild().accept(this, context);
-    OrderingScheme childOrdering = nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId());
+    OrderingScheme childOrdering = commonOrdering(childrenNodes);
     if (node.isStreamable()) {
       // Child has Ordering, we need to check if it is the Ordering we expected
       if (childOrdering != null) {
@@ -2118,7 +2141,7 @@ public List<PlanNode> visitAggregation(AggregationNode node, PlanContext context
                       if (physicalAggregation.isStreamable() && childOrdering != null) {
                         nodeOrderingMap.put(planNodeId, expectedOrderingSchema);
                       }
-                      return aggregationNode;
+                      return validateAggregationGrouping(aggregationNode);
                     })
                 .collect(Collectors.toList());
 
@@ -2128,10 +2151,10 @@ public List<PlanNode> visitAggregation(AggregationNode node, PlanContext context
             new CollectNode(queryId.genPlanNodeId(), aggregationNodes.get(0).getOutputSymbols());
         collectNode.setChildren(aggregationNodes);
         splitResult.left.setChild(collectNode);
-        return Collections.singletonList(splitResult.left);
+        return Collections.singletonList(validateAggregationGrouping(splitResult.left));
       }
 
-      return Collections.singletonList(physicalAggregation);
+      return Collections.singletonList(validateAggregationGrouping(physicalAggregation));
     }
 
     // We cannot do multi-stage Aggregate if any aggregation-function is distinct.
@@ -2144,7 +2167,9 @@ public List<PlanNode> visitAggregation(AggregationNode node, PlanContext context
               nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes);
       node.setChild(physicalChild);
       return Collections.singletonList(
-          withRateFunctionInputOrdering(node, nodeOrderingMap.get(physicalChild.getPlanNodeId())));
+          validateAggregationGrouping(
+              withRateFunctionInputOrdering(
+                  node, nodeOrderingMap.get(physicalChild.getPlanNodeId()))));
     }
     Pair<AggregationNode, AggregationNode> splitResult = split(node, symbolAllocator, queryId);
     AggregationNode intermediate = splitResult.right;
@@ -2167,13 +2192,73 @@ public List<PlanNode> visitAggregation(AggregationNode node, PlanContext context
                   if (node.isStreamable() && childOrdering != null) {
                     nodeOrderingMap.put(planNodeId, expectedOrderingSchema);
                   }
-                  return aggregationNode;
+                  return validateAggregationGrouping(aggregationNode);
                 })
             .collect(Collectors.toList());
     splitResult.left.setChild(
         mergeChildrenViaCollectOrMergeSort(
             nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes));
-    return Collections.singletonList(splitResult.left);
+    return Collections.singletonList(validateAggregationGrouping(splitResult.left));
+  }
+
+  private AggregationNode validateAggregationGrouping(AggregationNode node) {
+    if (!node.getPreGroupedSymbols().isEmpty()
+        && groupingInvalidatedByCollect(node.getChild(), node.getPreGroupedSymbols())) {
+      node.setPreGroupedSymbols(ImmutableList.of());
+      nodeOrderingMap.remove(node.getPlanNodeId());
+    }
+    return node;
+  }
+
+  /** Whether an unordered fan-in invalidates a previously established grouping property. */
+  private boolean groupingInvalidatedByCollect(PlanNode node, List<Symbol> grouping) {
+    if (grouping.isEmpty()) {
+      return false;
+    }
+    if ((node instanceof CollectNode || node instanceof UnionNode)
+        && node.getChildren().size() > 1) {
+      return true;
+    }
+    OrderingScheme ordering = nodeOrderingMap.get(node.getPlanNodeId());
+    if (ordering != null && prefixMatched(ordering, grouping)) {
+      return false;
+    }
+    // A completed aggregation emits each full grouping key once, even when it hashes its input.
+    if (node instanceof AggregationNode) {
+      AggregationNode aggregation = (AggregationNode) node;
+      if (!aggregation.getStep().isOutputPartial()
+          && aggregation.getGroupingSetCount() == 1
+          && grouping.containsAll(aggregation.getGroupingKeys())) {
+        return false;
+      }
+    }
+    if (node instanceof ProjectNode) {
+      ProjectNode project = (ProjectNode) node;
+      List<Symbol> inputs = new ArrayList<>();
+      for (Symbol symbol : grouping) {
+        Expression expression = project.getAssignments().get(symbol);
+        if (!(expression instanceof SymbolReference)) {
+          // A computed key cannot inherit a branch-local continuity proof.
+          return containsUnorderedCollect(project.getChild());
+        }
+        inputs.add(Symbol.from(expression));
+      }
+      return groupingInvalidatedByCollect(project.getChild(), inputs);
+    }
+    for (PlanNode child : node.getChildren()) {
+      if (groupingInvalidatedByCollect(child, grouping)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private boolean containsUnorderedCollect(PlanNode node) {
+    if ((node instanceof CollectNode || node instanceof UnionNode)
+        && node.getChildren().size() > 1) {
+      return true;
+    }
+    return node.getChildren().stream().anyMatch(this::containsUnorderedCollect);
   }
 
   private static AggregationNode withRateFunctionInputOrdering(
@@ -2301,7 +2386,8 @@ public List<PlanNode> visitAggregationTableScan(
         throw new IllegalStateException(
             DataNodeQueryMessages.LIST_PLANNODE_SIZE_SHOULD_1_BUT_NOW_IS);
       }
-      resultTableScanNodeList = Collections.singletonList(finalAggregation);
+      resultTableScanNodeList =
+          Collections.singletonList(validateAggregationGrouping(finalAggregation));
     }
 
     return resultTableScanNodeList;
@@ -2460,7 +2546,7 @@ private List<PlanNode> constructSpilledAggregationTableScanByRegionReplicaSet(
             mergeChildrenViaCollectOrMergeSort(
                 nodeOrderingMap.get(result.get(0).getPlanNodeId()), result));
       }
-      return Collections.singletonList(finalAggregation);
+      return Collections.singletonList(validateAggregationGrouping(finalAggregation));
     }
     return result;
   }
@@ -2677,7 +2763,8 @@ public List<PlanNode> visitAggregationTreeDeviceViewScan(
         throw new IllegalStateException(
             DataNodeQueryMessages.LIST_PLANNODE_SIZE_SHOULD_1_BUT_NOW_IS);
       }
-      resultTableScanNodeList = Collections.singletonList(finalAggregation);
+      resultTableScanNodeList =
+          Collections.singletonList(validateAggregationGrouping(finalAggregation));
     }
 
     return resultTableScanNodeList;
@@ -2863,7 +2950,7 @@ private List<PlanNode> constructSpilledAggregationTreeDeviceViewScanByRegionRepl
             mergeChildrenViaCollectOrMergeSort(
                 nodeOrderingMap.get(result.get(0).getPlanNodeId()), result));
       }
-      return Collections.singletonList(finalAggregation);
+      return Collections.singletonList(validateAggregationGrouping(finalAggregation));
     }
     return result;
   }
@@ -4007,6 +4094,9 @@ public NodeDistribution getNodeDistribution(PlanNodeId nodeId) {
     public void clearExpectedOrderingScheme() {
       expectedOrderingScheme = null;
       hasSortProperty = false;
+      // Grouping pushdown depends on this ordering. A boundary such as a global window
+      // must not pass an outer Group's partitioning request to its scans.
+      pushDownGrouping = false;
     }
 
     public void setExpectedOrderingScheme(OrderingScheme expectedOrderingScheme) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java
index 66c4eff6e3e03..ac3fc65fe6683 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java
@@ -34,6 +34,7 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TableFunctionProcessorNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
 import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode;
@@ -101,6 +102,21 @@ public PlanNode visitPlan(PlanNode node, Context context) {
       return newNode;
     }
 
+    @Override
+    public PlanNode visitUnion(UnionNode node, Context context) {
+      if (node.getChildren().size() <= 1) {
+        return visitPlan(node, context);
+      }
+      if (!context.canSkip()) {
+        context.canOptimized = TO_SORT;
+      }
+      // Neither branch may establish a property for the interleaved Union output.
+      return node.replaceChildren(
+          node.getChildren().stream()
+              .map(child -> child.accept(this, new Context(null, 0)))
+              .collect(Collectors.toList()));
+    }
+
     /** We need to make sure: context#partitionKey can match the prefix of childOrderSchema */
     private void checkPrefixMatch(Context context, List<Symbol> childOrder) {
       if (context.canSkip()) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java
index 64ef1b8c2adc6..c35f66128052e 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java
@@ -22,6 +22,7 @@
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.OrderingScheme;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.FillNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GapFillNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.PatternRecognitionNode;
@@ -31,6 +32,7 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TableFunctionProcessorNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKRankingNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PruneTableScanColumns;
@@ -38,6 +40,7 @@
 
 import com.google.common.collect.ImmutableSet;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
@@ -75,6 +78,28 @@ public PlanNode visitPlan(PlanNode node, Context context) {
       return newNode;
     }
 
+    @Override
+    public PlanNode visitCollect(CollectNode node, Context context) {
+      return visitUnorderedFanIn(node, context);
+    }
+
+    @Override
+    public PlanNode visitUnion(UnionNode node, Context context) {
+      return visitUnorderedFanIn(node, context);
+    }
+
+    private PlanNode visitUnorderedFanIn(PlanNode node, Context context) {
+      if (node.getChildren().size() <= 1) {
+        return visitPlan(node, context);
+      }
+      List<PlanNode> children = new ArrayList<>();
+      for (PlanNode child : node.getChildren()) {
+        children.add(child.accept(this, new Context()));
+      }
+      context.addCannotEliminateSort(true);
+      return node.replaceChildren(children);
+    }
+
     @Override
     public PlanNode visitProject(ProjectNode node, Context context) {
       Context newContext = new Context();
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
index ce49104248123..a5573159a94e5 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
@@ -40,9 +40,11 @@
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine;
 import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager;
 import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool;
+import org.apache.iotdb.db.queryengine.execution.memory.MemorySide;
 import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
 import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent;
 
+import org.apache.tsfile.read.common.block.TsBlock;
 import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
@@ -180,6 +182,13 @@ public void testCreateLocalSinkHandle() {
     Assert.assertEquals(
         ((LocalSinkChannel) shuffleSinkHandle.getChannel(0)).getSharedTsBlockQueue(),
         ((LocalSourceHandle) localSourceHandle).getSharedTsBlockQueue());
+    assertProducerOwnsQueue(
+        spyMemoryPool,
+        shuffleSinkHandle,
+        localSourceHandle,
+        localFragmentInstanceId,
+        remoteFragmentInstanceId,
+        localPlanNodeId);
   }
 
   @Test
@@ -237,5 +246,49 @@ public void testCreateLocalSourceHandle() {
     Assert.assertEquals(
         ((LocalSinkChannel) shuffleSinkHandle.getChannel(0)).getSharedTsBlockQueue(),
         ((LocalSourceHandle) localSourceHandle).getSharedTsBlockQueue());
+    assertProducerOwnsQueue(
+        spyMemoryPool,
+        shuffleSinkHandle,
+        localSourceHandle,
+        localFragmentInstanceId,
+        remoteFragmentInstanceId,
+        localPlanNodeId);
+  }
+
+  private static void assertProducerOwnsQueue(
+      MemoryPool pool,
+      ISinkHandle sink,
+      ISourceHandle source,
+      TFragmentInstanceId producer,
+      TFragmentInstanceId consumer,
+      String planNodeId) {
+    SharedTsBlockQueue queue = ((LocalSourceHandle) source).getSharedTsBlockQueue();
+    sink.setMaxBytesCanReserve(512);
+    source.setMaxBytesCanReserve(128);
+    Assert.assertEquals(512, queue.getMaxBytesCanReserve());
+    String producerId =
+        FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(producer);
+    String consumerId =
+        FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(consumer);
+    TsBlock block = Utils.createMockTsBlock(256);
+    synchronized (queue) {
+      Assert.assertTrue(queue.add(block).isDone());
+      Assert.assertEquals(256, pool.getReservedBytes());
+      Assert.assertEquals(
+          256, pool.getPlanNodeMemoryReservedBytes(producer.queryId, producerId, planNodeId));
+      Assert.assertEquals(
+          256,
+          pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.OUTPUT));
+      Assert.assertEquals(
+          0, pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.INPUT));
+      Assert.assertEquals(
+          0, pool.getFragmentMemoryReservedBytes(consumer.queryId, consumerId, MemorySide.INPUT));
+      Assert.assertEquals(
+          0, pool.getFragmentMemoryReservedBytes(consumer.queryId, consumerId, MemorySide.OUTPUT));
+      Assert.assertSame(block, queue.remove());
+      Assert.assertEquals(0, pool.getReservedBytes());
+      Assert.assertEquals(
+          0, pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.OUTPUT));
+    }
   }
 }
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java
index 2d9e5683cfdfb..3cfebb1ba295c 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java
@@ -66,9 +66,10 @@ public class SourceHandleTest {
   public static void beforeClass() {
     maxBytesPerFI =
         IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+    // INPUT receives half of the fragment budget, allowing five buffered blocks.
     IoTDBDescriptor.getInstance()
         .getMemoryConfig()
-        .setMaxBytesPerFragmentInstance(5 * MOCK_TSBLOCK_SIZE);
+        .setMaxBytesPerFragmentInstance(10 * MOCK_TSBLOCK_SIZE);
   }
 
   @AfterClass
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
index 5f2285e4ca910..19fc0f17d3367 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
@@ -50,7 +50,8 @@ public class MemoryPoolTest {
   public void before() {
     MemoryManager memoryManager = Mockito.spy(new MemoryManager(1024L));
     pool = new MemoryPool("test", memoryManager, 512L);
-    pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID);
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT);
   }
 
   @Test
@@ -87,7 +88,7 @@ public void testDeregisterDoesNotRemoveReregisteredQueryMemory() throws Exceptio
         // Replace the query map while the delayed cleanup still refers to the old one.
         pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true);
         Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
-        pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result", PLAN_NODE_ID);
+        pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result", PLAN_NODE_ID, MemorySide.INPUT);
         Assert.assertTrue(
             pool.tryReserveForTest(QUERY_ID, "result", PLAN_NODE_ID, 256L, Long.MAX_VALUE));
       }
@@ -123,7 +124,7 @@ public void testDeregisterMissingFragmentPreservesReservedMemory() {
 
   @Test
   public void testDeregisterPreservesOtherFragments() {
-    pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID);
+    pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID, MemorySide.INPUT);
     Assert.assertTrue(
         pool.tryReserveForTest(QUERY_ID, "other", PLAN_NODE_ID, 256L, Long.MAX_VALUE));
 
@@ -518,4 +519,112 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() {
     Assert.assertTrue(r.getFuture().isDone());
     Assert.assertEquals(256L, pool.getReservedBytes());
   }
+
+  @Test
+  public void testSharedInputBudgetAndIndependentOutput() {
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT);
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT);
+    Assert.assertTrue(
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 256, false)
+            .isReserveSuccess());
+    MemoryReservationResult pending =
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256, false);
+    Assert.assertFalse(pending.isReserveSuccess());
+    Assert.assertFalse(pending.getFuture().isDone());
+    Assert.assertEquals(
+        200, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
+    Assert.assertEquals(
+        0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2"));
+    Assert.assertTrue(
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 256, false)
+            .isReserveSuccess());
+    Assert.assertEquals(456, pool.getReservedBytes());
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200);
+    Assert.assertTrue(pending.getFuture().isDone());
+    Assert.assertEquals(
+        100, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
+    Assert.assertEquals(
+        256,
+        pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.OUTPUT));
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100);
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256);
+    Assert.assertEquals(0, pool.getReservedBytes());
+    Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID));
+  }
+
+  @Test
+  public void testSideRegistrationCannotChange() {
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT);
+    try {
+      pool.registerPlanNodeIdToQueryMemoryMap(
+          QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.OUTPUT);
+      Assert.fail("conflicting side must be rejected");
+    } catch (IllegalArgumentException expected) {
+      Assert.assertEquals(0, pool.getReservedBytes());
+    }
+  }
+
+  @Test
+  public void testOddBudgetAndSingleRequestBoundary() {
+    Assert.assertEquals(256, MemorySide.INPUT.getMaxBytes(513));
+    Assert.assertEquals(257, MemorySide.OUTPUT.getMaxBytes(513));
+    try {
+      pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, 256, false);
+      Assert.fail("request exceeds the input budget");
+    } catch (IllegalArgumentException expected) {
+      Assert.assertEquals(0, pool.getReservedBytes());
+      Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID));
+    }
+  }
+
+  @Test
+  public void testConcurrentPlanNodesShareOneSideCounter() throws Exception {
+    java.util.concurrent.ExecutorService executor =
+        java.util.concurrent.Executors.newFixedThreadPool(16);
+    java.util.concurrent.CountDownLatch start = new java.util.concurrent.CountDownLatch(1);
+    java.util.concurrent.CountDownLatch attempted = new java.util.concurrent.CountDownLatch(16);
+    java.util.concurrent.CountDownLatch release = new java.util.concurrent.CountDownLatch(1);
+    java.util.List<java.util.concurrent.Future<?>> tasks = new java.util.ArrayList<>();
+    try {
+      for (int i = 0; i < 16; i++) {
+        String id = "pipeline-" + i;
+        pool.registerPlanNodeIdToQueryMemoryMap(
+            QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT);
+        tasks.add(
+            executor.submit(
+                () -> {
+                  start.await();
+                  boolean reserved =
+                      pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128);
+                  attempted.countDown();
+                  release.await();
+                  if (reserved) pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32);
+                  return null;
+                }));
+      }
+      start.countDown();
+      Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS));
+      long details = 0;
+      for (int i = 0; i < 16; i++) {
+        details +=
+            pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "pipeline-" + i);
+      }
+      Assert.assertTrue(details > 0 && details <= 128);
+      Assert.assertEquals(
+          details,
+          pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
+      Assert.assertEquals(details, pool.getReservedBytes());
+      release.countDown();
+      for (java.util.concurrent.Future<?> task : tasks) task.get(10, TimeUnit.SECONDS);
+      Assert.assertEquals(0, pool.getReservedBytes());
+      Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID));
+    } finally {
+      start.countDown();
+      release.countDown();
+      executor.shutdownNow();
+    }
+  }
 }
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java
new file mode 100644
index 0000000000000..68bff035fe2c3
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java
@@ -0,0 +1,78 @@
+/*
+ * 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.queryengine.plan.planner;
+
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
+import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
+import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle;
+import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle;
+import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
+import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator;
+import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class LocalExecutionPlanContextBudgetTest {
+  @Test
+  public void testFixedBudgetsInheritedByPipelinesAndWiredBySide() {
+    long previous =
+        IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
+    try {
+      IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(513);
+      FragmentInstanceContext instance = mock(FragmentInstanceContext.class);
+      when(instance.getId())
+          .thenReturn(new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0"));
+      LocalExecutionPlanContext context =
+          new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null);
+      context.setExchangeSumNum(500);
+      assertEquals(256, context.getInputMaxBytes());
+      assertEquals(257, context.getOutputMaxBytes());
+      IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024);
+      LocalExecutionPlanContext child = context.createSubContext();
+      assertEquals(256, child.getInputMaxBytes());
+      assertEquals(257, child.getOutputMaxBytes());
+      ISink sink = mock(ISink.class);
+      child.setISink(sink);
+      verify(sink).setMaxBytesCanReserve(257);
+      ISourceHandle remote = mock(ISourceHandle.class);
+      PipelineSourceHandle pipeline = mock(PipelineSourceHandle.class);
+      ExchangeOperator remoteOperator = mock(ExchangeOperator.class);
+      ExchangeOperator pipelineOperator = mock(ExchangeOperator.class);
+      when(remoteOperator.getSourceHandle()).thenReturn(remote);
+      when(pipelineOperator.getSourceHandle()).thenReturn(pipeline);
+      context.addExchangeOperator(remoteOperator);
+      child.addExchangeOperator(pipelineOperator);
+      context.setExchangeMemoryLimits();
+      verify(remote).setMaxBytesCanReserve(256);
+      verify(pipeline).setMaxBytesCanReserve(257);
+    } finally {
+      IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previous);
+    }
+  }
+}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java
new file mode 100644
index 0000000000000..3021363139c72
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java
@@ -0,0 +1,149 @@
+/*
+ * 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.queryengine.plan.relational.planner;
+
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.AggregationNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class CollectInputPropertyTest {
+  private static PlanNode plan(String sql) {
+    PlanTester tester = new PlanTester();
+    tester.createPlan(sql);
+    return tester.getFragmentPlan(0);
+  }
+
+  private static <T> List<T> nodes(PlanNode node, Class<T> type) {
+    List<T> result = new ArrayList<>();
+    if (type.isInstance(node)) result.add(type.cast(node));
+    for (PlanNode child : node.getChildren()) result.addAll(nodes(child, type));
+    return result;
+  }
+
+  @Test
+  public void testAggregationAboveGlobalWindowHashesInterleavedGroups() {
+    for (String aggregate : new String[] {"SUM(s1)", "COUNT(DISTINCT s1)"}) {
+      PlanNode root =
+          plan(
+              "SELECT tag1, tag2, tag3, "
+                  + aggregate
+                  + " FROM ("
+                  + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n "
+                  + "FROM table1 WHERE tag1 != 'shenzhen') t WHERE n > 0 GROUP BY tag1, tag2, tag3");
+      List<AggregationNode> aggregations = nodes(root, AggregationNode.class);
+      assertFalse(aggregations.isEmpty());
+      assertFalse(nodes(root, CollectNode.class).isEmpty());
+      for (AggregationNode aggregation : aggregations) {
+        if (!nodes(aggregation.getChild(), CollectNode.class).isEmpty()) {
+          assertTrue(aggregation.getPreGroupedSymbols().isEmpty());
+          assertFalse(aggregation.isStreamable());
+        }
+      }
+      assertTrue(
+          "Hash aggregation should not add a grouping sort", nodes(root, SortNode.class).isEmpty());
+    }
+  }
+
+  @Test
+  public void testPartialPreGroupingAlsoFallsBackToHash() {
+    PlanNode root =
+        plan(
+            "SELECT tag1, tag2, tag3, s1, COUNT(*) FROM ("
+                + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n "
+                + "FROM table1 WHERE tag1 != 'shenzhen') t WHERE n > 0 GROUP BY tag1, tag2, tag3, s1");
+    List<AggregationNode> aggregations = nodes(root, AggregationNode.class);
+    assertFalse(aggregations.isEmpty());
+    assertFalse(nodes(root, CollectNode.class).isEmpty());
+    assertTrue(aggregations.get(0).getPreGroupedSymbols().isEmpty());
+  }
+
+  @Test
+  public void testNestedWindowRestoresPartitionContiguityWithSort() {
+    PlanNode root =
+        plan(
+            "SELECT tag1, tag2, tag3, SUM(s1) OVER ("
+                + "PARTITION BY tag1, tag2, tag3 ORDER BY s1) AS running_sum "
+                + "FROM (SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1) t WHERE n > 0");
+    assertFalse(nodes(root, CollectNode.class).isEmpty());
+    assertTrue(nodes(root, GroupNode.class).isEmpty());
+    List<WindowNode> windows = nodes(root, WindowNode.class);
+    assertEquals(2, windows.size());
+    WindowNode outer =
+        windows.stream()
+            .filter(w -> !w.getSpecification().getPartitionBy().isEmpty())
+            .findFirst()
+            .get();
+    assertFalse(nodes(outer.getChild(), SortNode.class).isEmpty());
+  }
+
+  @Test
+  public void testStreamSortAfterGlobalWindowUsesFullSort() {
+    PlanNode root =
+        plan(
+            "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n "
+                + "FROM table1 ORDER BY tag1, tag2, tag3, s1");
+    assertFalse(nodes(root, CollectNode.class).isEmpty());
+    assertFalse(nodes(root, SortNode.class).isEmpty());
+    assertTrue(nodes(root, StreamSortNode.class).isEmpty());
+  }
+
+  @Test
+  public void testConstantFillDoesNotRemoveRestoredSort() {
+    PlanNode root =
+        plan(
+            "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1 "
+                + "FILL METHOD CONSTANT 0 ORDER BY tag1, tag2, tag3");
+    assertFalse(nodes(root, CollectNode.class).isEmpty());
+    assertFalse(nodes(root, SortNode.class).isEmpty());
+    assertTrue(nodes(root, StreamSortNode.class).isEmpty());
+  }
+
+  @Test
+  public void testUnionFinalAggregationDoesNotInheritBranchGrouping() {
+    PlanNode root =
+        plan(
+            "SELECT tag1, tag2, tag3, SUM(s1) FROM ("
+                + "SELECT tag1, tag2, tag3, s1 FROM table1 UNION ALL "
+                + "SELECT tag1, tag2, tag3, s1 FROM table2) u GROUP BY tag1, tag2, tag3");
+    List<AggregationNode> aggregations = nodes(root, AggregationNode.class);
+    assertFalse(aggregations.isEmpty());
+    boolean foundFinalCollect = false;
+    for (AggregationNode aggregation : aggregations) {
+      if (aggregation.getChild() instanceof CollectNode) {
+        foundFinalCollect = true;
+        assertTrue(aggregation.getPreGroupedSymbols().isEmpty());
+      }
+    }
+    assertTrue(foundFinalCollect);
+  }
+}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java
index e2326688112b2..b5b69f7663f03 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java
@@ -26,6 +26,9 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.OrderingScheme;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.ValuesNode;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.execution.warnings.WarningCollector;
@@ -36,6 +39,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExternalTsFileScanNode;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import org.junit.Test;
 
@@ -108,4 +112,54 @@ region, new ColumnSchema("region", STRING, false, TAG),
       resource.closeByQueryExecution();
     }
   }
+
+  @Test
+  public void testUnionDoesNotInheritFirstBranchGrouping() {
+    Symbol key = new Symbol("key");
+    OrderingScheme ordering =
+        new OrderingScheme(ImmutableList.of(key), ImmutableMap.of(key, ASC_NULLS_LAST));
+    PlanNode left =
+        new SortNode(
+            new PlanNodeId("leftSort"),
+            new ValuesNode(new PlanNodeId("left"), ImmutableList.of(key), ImmutableList.of()),
+            ordering,
+            false,
+            false);
+    PlanNode right =
+        new SortNode(
+            new PlanNodeId("rightSort"),
+            new ValuesNode(new PlanNodeId("right"), ImmutableList.of(key), ImmutableList.of()),
+            ordering,
+            false,
+            false);
+    UnionNode union =
+        new UnionNode(
+            new PlanNodeId("union"),
+            ImmutableList.of(left, right),
+            ImmutableListMultimap.of(key, key, key, key),
+            ImmutableList.of(key));
+    assertTrue(
+        optimizeGroup(new GroupNode(new PlanNodeId("group"), union, ordering, 1))
+            instanceof SortNode);
+    // An explicit sort above the fan-in restores the grouping property.
+    SortNode restored = new SortNode(new PlanNodeId("restored"), union, ordering, false, false);
+    assertTrue(
+        optimizeGroup(new GroupNode(new PlanNodeId("group2"), restored, ordering, 1))
+            instanceof GroupNode);
+  }
+
+  private PlanNode optimizeGroup(GroupNode group) {
+    return new ParallelizeGrouping()
+        .optimize(
+            group,
+            new PlanOptimizer.Context(
+                SESSION_INFO,
+                new Analysis(null, Collections.emptyMap()),
+                TEST_MATADATA,
+                new MPPQueryContext(new QueryId("group_union")),
+                new SymbolAllocator(),
+                new QueryId("group_union"),
+                WarningCollector.NOOP,
+                PlanOptimizersStatsCollector.createPlanOptimizersStatsCollector()));
+  }
 }

From 9e54ec5c95cafe7ed2ea5c01779bf894f5498e27 Mon Sep 17 00:00:00 2001
From: ColinLee <shuolin_l@163.com>
Date: Wed, 14 Oct 2026 14:53:56 +0800
Subject: [PATCH 3/3] Fix exchange reservation wakeup race and strengthen
 fan-in tests

---
 ...2026-10-14-exchange-fanin-budget-design.md |  22 ++-
 .../it/query/recent/IoTDBCollectFanInIT.java  |  23 +++
 .../operator/process/CollectOperatorTest.java | 102 ++++++++++++
 .../execution/memory/MemoryPool.java          |  51 +++---
 .../execution/memory/MemoryPoolTest.java      | 154 +++++++++++++++++-
 .../planner/CollectInputPropertyTest.java     |  28 ++++
 6 files changed, 353 insertions(+), 27 deletions(-)

diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
index a817039cab4b9..e191622940dc7 100644
--- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
+++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md
@@ -72,7 +72,7 @@ SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除
 
 ## 实现验证记录（2026-10-14）
 
-第一版实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。
+上述约定范围的实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。
 集成验证包含普通与 DISTINCT 聚合、Union 分阶段聚合、嵌套窗口、Fill、排序/TopK、Join 依赖、树模型以及树表视图。
 
 构建先执行以下干净构建，生成当前分支的节点包并安装本地依赖：
@@ -97,3 +97,23 @@ mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfa
 
 使用离线模式是为了避免远端 SNAPSHOT ZIP 下载停滞；节点执行的是当前分支重新打包的代码。
 Spotless、Checkstyle、新增文件许可证及中英文消息键名/占位符一致性检查通过。
+
+## 自查与补测（2026-10-14）
+
+自查复现了 MemoryPool 的等待申请漏唤醒：一次申请额度不足后，旧申请先释放全部额度，
+新申请才进入等待队列；释放线程看不到新申请，导致空闲额度下仍无限等待。
+回归测试用 latch 固定这一时序，在修复前稳定失败。
+修复在申请入队后再次检查额度，并与 free/tryCancel 共用该 future 的锁，避免重复计费。
+
+新增 8 项单元测试：等待申请发布竞态、取消不占用共享额度、16 路阻塞申请反复获得额度后账目归零；
+空 Collect、全部阻塞时异常唤醒、并发就绪后每路恰好消费一次；
+Collect 后 TopK 不误用有序输入，以及有效分支保留流式聚合。
+增强关闭异常测试，检查其他输入仍被关闭且后续异常被保留。
+新增一项集成测试，验证全局窗口后按字段倒序的 TopK 跨设备结果及完整窗口计数。
+
+补测后的 118 项定向单测通过（14 个测试类，零失败、零跳过）；全 reactor 中英文编译通过。
+在前述英文回归命令中加入 `OrderingPropertyPropagationTest,TopKRuntimeFilterDistributedPlanTest`，
+同时验证有效排序优化与 TopK runtime filter 不受影响。
+重新执行 `mvn -o install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests`
+后运行前述两条 verify 命令；IoTDBCollectFanInIT 的 7 项测试在单节点和三 DataNode 集群中分别全部通过。
+Spotless、Checkstyle、许可证与 `git diff --check` 检查通过。
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java
index adc72b6046e50..8667f01ddfd70 100644
--- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java
@@ -273,4 +273,27 @@ public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception {
       }
     }
   }
+
+  @Test
+  public void testTopKAfterGlobalWindowConsumesUnorderedInput() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE);
+      int limit = DEVICES + 5;
+      try (ResultSet result =
+          statement.executeQuery(
+              "SELECT device, zone, value, COUNT(*) OVER () AS n FROM readings1 "
+                  + "ORDER BY value DESC, device LIMIT "
+                  + limit)) {
+        for (int row = 0; row < limit; row++) {
+          assertTrue(result.next());
+          assertEquals(device(row % DEVICES), result.getString(1));
+          assertEquals(zone(row % DEVICES), result.getString(2));
+          assertEquals(ROWS - row / DEVICES, result.getInt(3));
+          assertEquals((long) DEVICES * ROWS, result.getLong(4));
+        }
+        assertFalse(result.next());
+      }
+    }
+  }
 }
diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java
index abf7a82acf3d9..b446f9c999b95 100644
--- a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java
+++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java
@@ -36,9 +36,16 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -233,6 +240,7 @@ public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Excep
     Child a = new Child(block(1));
     Child b = new Child(block(2));
     a.closeFailure = new IOException("close failure");
+    b.closeFailure = new IOException("second close failure");
     CollectOperator collect = collect(a, b);
     ListenableFuture<?> pending = collect.isBlocked();
     try {
@@ -240,6 +248,8 @@ public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Excep
       fail("must propagate close failure");
     } catch (IOException e) {
       assertSame(a.closeFailure, e);
+      assertEquals(1, e.getSuppressed().length);
+      assertSame(b.closeFailure, e.getSuppressed()[0]);
     }
     pending.get(5, TimeUnit.SECONDS);
     assertEquals(1, a.closes);
@@ -299,6 +309,98 @@ public void testMemoryIncludesAllLiveRetainedStates() throws Exception {
     assertEquals(0, collect.calculateRetainedSizeAfterCallingNext());
   }
 
+  @Test
+  public void testEmptyCollectIsAlreadyFinished() throws Exception {
+    CollectOperator collect = collect();
+    assertFalse(collect.hasNext());
+    assertTrue(collect.isFinished());
+    assertTrue(collect.isBlocked().isDone());
+    assertNull(collect.next());
+    assertEquals(0, collect.calculateMaxPeekMemory());
+    assertEquals(0, collect.calculateMaxReturnSize());
+    collect.close();
+    collect.close();
+  }
+
+  @Test
+  public void testFailureWakesAllPendingChildren() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    CollectOperator collect = collect(a, b);
+    ListenableFuture<?> pending = collect.isBlocked();
+    assertFalse(pending.isDone());
+    IOException expected = new IOException("blocked child failure");
+    b.wait.delegate.setException(expected);
+    try {
+      pending.get(5, TimeUnit.SECONDS);
+      fail("child failure must wake the driver");
+    } catch (ExecutionException e) {
+      assertSame(expected, e.getCause());
+    }
+    try {
+      collect.next();
+      fail("child failure must propagate without advancing other children");
+    } catch (IOException e) {
+      assertSame(expected, e);
+    } finally {
+      collect.close();
+    }
+    assertEquals(0, a.advances);
+    assertEquals(0, b.advances);
+    assertEquals(1, a.closes);
+    assertEquals(1, b.closes);
+  }
+
+  @Test
+  public void testConcurrentCompletionsConsumeEveryChildExactlyOnce() throws Exception {
+    int count = 16;
+    List<Child> children = new ArrayList<>();
+    for (int i = 0; i < count; i++) children.add(new Child(block(i)));
+    CollectOperator collect = new CollectOperator(new TestContext(), new ArrayList<>(children));
+    ListenableFuture<?> pending = collect.isBlocked();
+    assertFalse(pending.isDone());
+    ExecutorService executor = Executors.newFixedThreadPool(count);
+    CountDownLatch start = new CountDownLatch(1);
+    List<Future<?>> completions = new ArrayList<>();
+    try {
+      for (Child child : children) {
+        completions.add(
+            executor.submit(
+                () -> {
+                  start.await();
+                  child.ready();
+                  return null;
+                }));
+      }
+      start.countDown();
+      pending.get(5, TimeUnit.SECONDS);
+      for (Future<?> completion : completions) completion.get(5, TimeUnit.SECONDS);
+      Set<Integer> consumed = new HashSet<>();
+      for (int round = 0; round < 2 * count && collect.hasNext(); round++) {
+        assertTrue(collect.isBlocked().isDone());
+        TsBlock result = collect.next();
+        if (result != null) {
+          int childId = result.getColumn(0).getInt(0);
+          assertTrue("child emitted more than once", consumed.add(childId));
+          children.get(childId).hasNext = false;
+          children.get(childId).finished = true;
+        }
+      }
+      assertEquals(count, consumed.size());
+      assertTrue(collect.isFinished());
+      assertNull(collect.next());
+      for (Child child : children) {
+        assertEquals(1, child.advances);
+        assertEquals(1, child.closes);
+      }
+    } finally {
+      start.countDown();
+      executor.shutdownNow();
+      assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+      collect.close();
+    }
+  }
+
   private static final class CountingFuture
       extends ForwardingListenableFuture.SimpleForwardingListenableFuture<Void> {
     private final SettableFuture<Void> delegate;
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
index 56c9974764a5f..07eced179d94e 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
@@ -365,10 +365,15 @@ public MemoryReservationResult reserveWithPriority(
           DataNodeQueryMessages.BLOCKED_RESERVE_REQUEST_ARG_BYTES_MEMORY_FOR_PLANNODEIDARG,
           bytesToReserve,
           planNodeId);
-      ListenableFuture<Void> result =
+      MemoryReservationFuture<Void> result =
           MemoryReservationFuture.create(
               queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve);
-      memoryReservationFutures.add((MemoryReservationFuture<Void>) result);
+      memoryReservationFutures.add(result);
+      // A concurrent free may have missed this request before it was published. Retry after
+      // publication, under the same future lock used by free and tryCancel to avoid charging twice.
+      if (tryCompleteReservation(result)) {
+        memoryReservationFutures.remove(result);
+      }
       return new MemoryReservationResult(result, false, bytesToReserve);
     }
   }
@@ -463,27 +468,35 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l
     Iterator<MemoryReservationFuture<Void>> iterator = memoryReservationFutures.iterator();
     while (iterator.hasNext()) {
       MemoryReservationFuture<Void> future = iterator.next();
-      synchronized (future) {
-        if (future.isCancelled() || future.isDone()) {
-          iterator.remove();
-          continue;
-        }
-        long bytesToReserve = future.getBytesToReserve();
-        String curQueryId = future.getQueryId();
-        String curFragmentInstanceId = future.getFragmentInstanceId();
-        String curPlanNodeId = future.getPlanNodeId();
-        long maxBytesCanReserve = future.getMaxBytesCanReserve();
-        if (tryReserve(
-            curQueryId, curFragmentInstanceId, curPlanNodeId, bytesToReserve, maxBytesCanReserve)) {
-          future.set(null);
-          iterator.remove();
-        } else {
-          rollbackReserve(curQueryId, curFragmentInstanceId, curPlanNodeId, bytesToReserve);
-        }
+      if (tryCompleteReservation(future)) {
+        iterator.remove();
       }
     }
   }
 
+  private boolean tryCompleteReservation(MemoryReservationFuture<Void> future) {
+    synchronized (future) {
+      if (future.isDone()) {
+        return true;
+      }
+      long bytesToReserve = future.getBytesToReserve();
+      String queryId = future.getQueryId();
+      String fragmentInstanceId = future.getFragmentInstanceId();
+      String planNodeId = future.getPlanNodeId();
+      if (tryReserve(
+          queryId,
+          fragmentInstanceId,
+          planNodeId,
+          bytesToReserve,
+          future.getMaxBytesCanReserve())) {
+        future.set(null);
+        return true;
+      }
+      rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve);
+      return false;
+    }
+  }
+
   public long getQueryMemoryReservedBytes(String queryId) {
     Map<String, FragmentMemoryReservation> fragments = queryMemoryReservations.get(queryId);
     if (fragments == null) {
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
index 19fc0f17d3367..0254fa1ac33c3 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
@@ -32,8 +32,15 @@
 import java.lang.management.ManagementFactory;
 import java.lang.management.ThreadInfo;
 import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.awaitility.Awaitility.await;
 
@@ -582,12 +589,11 @@ public void testOddBudgetAndSingleRequestBoundary() {
 
   @Test
   public void testConcurrentPlanNodesShareOneSideCounter() throws Exception {
-    java.util.concurrent.ExecutorService executor =
-        java.util.concurrent.Executors.newFixedThreadPool(16);
-    java.util.concurrent.CountDownLatch start = new java.util.concurrent.CountDownLatch(1);
-    java.util.concurrent.CountDownLatch attempted = new java.util.concurrent.CountDownLatch(16);
-    java.util.concurrent.CountDownLatch release = new java.util.concurrent.CountDownLatch(1);
-    java.util.List<java.util.concurrent.Future<?>> tasks = new java.util.ArrayList<>();
+    ExecutorService executor = Executors.newFixedThreadPool(16);
+    CountDownLatch start = new CountDownLatch(1);
+    CountDownLatch attempted = new CountDownLatch(16);
+    CountDownLatch release = new CountDownLatch(1);
+    List<Future<?>> tasks = new ArrayList<>();
     try {
       for (int i = 0; i < 16; i++) {
         String id = "pipeline-" + i;
@@ -618,7 +624,7 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception {
           pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
       Assert.assertEquals(details, pool.getReservedBytes());
       release.countDown();
-      for (java.util.concurrent.Future<?> task : tasks) task.get(10, TimeUnit.SECONDS);
+      for (Future<?> task : tasks) task.get(10, TimeUnit.SECONDS);
       Assert.assertEquals(0, pool.getReservedBytes());
       Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID));
     } finally {
@@ -627,4 +633,138 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception {
       executor.shutdownNow();
     }
   }
+
+  @Test
+  public void testReleaseBeforePendingReservationIsPublished() throws Exception {
+    pool = Mockito.spy(pool);
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT);
+    Assert.assertTrue(
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false)
+            .isReserveSuccess());
+    CountDownLatch failedAttempt = new CountDownLatch(1);
+    CountDownLatch released = new CountDownLatch(1);
+    AtomicBoolean firstAttempt = new AtomicBoolean(true);
+    Mockito.doAnswer(
+            invocation -> {
+              boolean granted = (boolean) invocation.callRealMethod();
+              if (!granted && firstAttempt.compareAndSet(true, false)) {
+                failedAttempt.countDown();
+                Assert.assertTrue(released.await(10, TimeUnit.SECONDS));
+              }
+              return granted;
+            })
+        .when(pool)
+        .tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256, 256);
+    FutureTask<MemoryReservationResult> reservation =
+        new FutureTask<>(
+            () ->
+                pool.reserveWithPriority(
+                    QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256, 256, false));
+    Thread reserver = new Thread(reservation, "reserve-shared-input");
+    reserver.start();
+    try {
+      Assert.assertTrue(failedAttempt.await(10, TimeUnit.SECONDS));
+      pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256);
+      released.countDown();
+      MemoryReservationResult result = reservation.get(10, TimeUnit.SECONDS);
+      Assert.assertTrue(
+          "Releasing the last owner must wake the pending request", result.getFuture().isDone());
+      result.getFuture().get(10, TimeUnit.SECONDS);
+      Assert.assertEquals(256, pool.getReservedBytes());
+      Assert.assertEquals(
+          256,
+          pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
+      Assert.assertEquals(
+          256, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2"));
+      pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256);
+      Assert.assertEquals(0, pool.getReservedBytes());
+    } finally {
+      released.countDown();
+      reserver.join(10000);
+      Assert.assertFalse(reserver.isAlive());
+    }
+  }
+
+  @Test
+  public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Exception {
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", MemorySide.INPUT);
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", MemorySide.INPUT);
+    pool.registerPlanNodeIdToQueryMemoryMap(
+        QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT);
+    pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false);
+    MemoryReservationResult cancelled =
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", 128, 256, false);
+    MemoryReservationResult waiting =
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", 256, 256, false);
+    Assert.assertFalse(cancelled.getFuture().isDone());
+    Assert.assertFalse(waiting.getFuture().isDone());
+    Assert.assertTrue(
+        pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 128, 256, false)
+            .isReserveSuccess());
+    Assert.assertEquals(128, pool.tryCancel(cancelled.getFuture()));
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256);
+    waiting.getFuture().get(10, TimeUnit.SECONDS);
+    Assert.assertTrue(cancelled.getFuture().isCancelled());
+    Assert.assertEquals(
+        0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled"));
+    Assert.assertEquals(
+        256, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
+    Assert.assertEquals(
+        128,
+        pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.OUTPUT));
+    Assert.assertEquals(384, pool.getReservedBytes());
+    Assert.assertEquals(0, pool.getMemoryReservationSize());
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", 256);
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 128);
+    Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID));
+    Assert.assertEquals(0, pool.getReservedBytes());
+  }
+
+  @Test
+  public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() throws Exception {
+    int count = 16;
+    ExecutorService executor = Executors.newFixedThreadPool(count);
+    CountDownLatch attempted = new CountDownLatch(count);
+    List<Future<?>> tasks = new ArrayList<>();
+    pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 128, 128, false);
+    try {
+      for (int i = 0; i < count; i++) {
+        String id = "waiting-" + i;
+        pool.registerPlanNodeIdToQueryMemoryMap(
+            QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT);
+        tasks.add(
+            executor.submit(
+                () -> {
+                  for (int round = 0; round < 50; round++) {
+                    MemoryReservationResult result =
+                        pool.reserveWithPriority(
+                            QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128, false);
+                    if (round == 0) attempted.countDown();
+                    result.getFuture().get(10, TimeUnit.SECONDS);
+                    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, result.getReservedBytes());
+                  }
+                  return null;
+                }));
+      }
+      Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS));
+      Assert.assertEquals(count, pool.getMemoryReservationSize());
+      pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 128);
+      for (Future<?> task : tasks) task.get(30, TimeUnit.SECONDS);
+      for (int i = 0; i < count; i++) {
+        Assert.assertEquals(
+            0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting-" + i));
+      }
+      Assert.assertEquals(
+          0, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT));
+      Assert.assertEquals(0, pool.getReservedBytes());
+      Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID));
+      Assert.assertEquals(0, pool.getMemoryReservationSize());
+    } finally {
+      executor.shutdownNow();
+      Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+    }
+  }
 }
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java
index 3021363139c72..5bae986bf4762 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java
@@ -25,6 +25,7 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode;
 
 import org.junit.Test;
@@ -146,4 +147,31 @@ public void testUnionFinalAggregationDoesNotInheritBranchGrouping() {
     }
     assertTrue(foundFinalCollect);
   }
+
+  @Test
+  public void testTopKAboveGlobalWindowDoesNotAssumeOrderedInput() {
+    PlanNode root =
+        plan(
+            "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1 "
+                + "ORDER BY s1 DESC, tag1, tag2, tag3 LIMIT 5");
+    List<TopKNode> topKs = nodes(root, TopKNode.class);
+    assertFalse(topKs.isEmpty());
+    boolean foundUnorderedInput = false;
+    for (TopKNode topK : topKs) {
+      if (!nodes(topK.getChildren().get(0), CollectNode.class).isEmpty()) {
+        foundUnorderedInput = true;
+        assertFalse(topK.isChildrenDataInOrder());
+      }
+    }
+    assertTrue(foundUnorderedInput);
+  }
+
+  @Test
+  public void testOrderedScanBranchesKeepStreamingAggregation() {
+    PlanNode root =
+        plan("SELECT tag1, tag2, tag3, SUM(s1 + 1) FROM table1 GROUP BY tag1, tag2, tag3");
+    List<AggregationNode> aggregations = nodes(root, AggregationNode.class);
+    assertFalse(aggregations.isEmpty());
+    assertTrue(aggregations.stream().anyMatch(AggregationNode::isStreamable));
+  }
 }