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 6a98ca7c4828..3664a845d4f8 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
@@ -37,9 +37,8 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
@@ -56,6 +55,8 @@ public static class MemoryReservationFuture<V> extends AbstractFuture<V> {
     private final String fragmentInstanceId;
     private final String planNodeId;
     private final long bytesToReserve;
+    // Assigned once by the pool before publication. Zero denotes an untracked future.
+    private long reservationId;
 
     /**
      * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the
@@ -173,6 +174,8 @@ 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 final Map<Long, MemoryReservationFuture<Void>> pendingReservations =
+        new ConcurrentHashMap<>();
     // Enabled by local planning before drivers start; request limits are supplied by handles.
     private volatile boolean sharedBudget;
 
@@ -185,8 +188,11 @@ private AtomicLong getUsedBytes(MemorySide side) {
   private final Map<String, Map<String, FragmentMemoryReservation>> queryMemoryReservations =
       new ConcurrentHashMap<>();
 
-  private final Queue<MemoryReservationFuture<Void>> memoryReservationFutures =
-      new ConcurrentLinkedQueue<>();
+  // Sequence keys retain FIFO wakeup order and allow removal without scanning unrelated requests.
+  // Only blocked reservations use these indexes; successful admission keeps its existing fast path.
+  private final Map<Long, MemoryReservationFuture<Void>> memoryReservationFutures =
+      new ConcurrentSkipListMap<>();
+  private final AtomicLong nextReservationId = new AtomicLong();
 
   public MemoryPool(String id, MemoryManager memoryManager, long maxBytesPerFragmentInstance) {
     this.id = Validate.notNull(id, DataNodeQueryMessages.EXCEPTION_ID_CAN_NOT_BE_NULL_DOT_BDD2AD7D);
@@ -302,13 +308,10 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap(
       // Pending requests must not outlive their FI registration. Cancel before removing the
       // account so another query's free() cannot retry a request against a missing FI. Do not
       // hold the registration lock while notifying cancellation listeners.
-      Iterator<MemoryReservationFuture<Void>> pending = memoryReservationFutures.iterator();
-      while (pending.hasNext()) {
-        MemoryReservationFuture<Void> future = pending.next();
-        if (queryId.equals(future.getQueryId())
-            && fragmentInstanceId.equals(future.getFragmentInstanceId())) {
+      if (fragmentRelatedMemory != null) {
+        for (MemoryReservationFuture<Void> future :
+            fragmentRelatedMemory.pendingReservations.values()) {
           tryCancel(future);
-          pending.remove();
         }
       }
       synchronized (queryMemoryReservations) {
@@ -392,11 +395,21 @@ public MemoryReservationResult reserveWithPriority(
       MemoryReservationFuture<Void> result =
           MemoryReservationFuture.create(
               queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve);
-      memoryReservationFutures.add(result);
+      FragmentMemoryReservation fragment =
+          getFragmentReservation(queryId, fragmentInstanceId, planNodeId);
+      result.reservationId = nextReservationId.incrementAndGet();
+      fragment.pendingReservations.put(result.reservationId, result);
+      memoryReservationFutures.put(result.reservationId, 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);
+      try {
+        if (tryCompleteReservation(result)) {
+          removePendingReservation(result);
+        }
+      } catch (RuntimeException | Error e) {
+        // A concurrently removed registration must not leave a published request behind.
+        tryCancel(result);
+        throw e;
       }
       return new MemoryReservationResult(result, false, bytesToReserve);
     }
@@ -440,6 +453,9 @@ public synchronized long tryCancel(ListenableFuture<Void> future) {
     synchronized (future) {
       // If the future is not a MemoryReservationFuture, it must have been completed.
       if (future.isDone()) {
+        if (future instanceof MemoryReservationFuture) {
+          removePendingReservation((MemoryReservationFuture<Void>) future);
+        }
         return 0L;
       }
       Validate.isTrue(
@@ -447,6 +463,7 @@ public synchronized long tryCancel(ListenableFuture<Void> future) {
           DataNodeQueryMessages.EXCEPTION_INVALID_FUTURE_TYPE_14507EF5
               + future.getClass().getSimpleName());
       future.cancel(true);
+      removePendingReservation((MemoryReservationFuture<Void>) future);
     }
     return ((MemoryReservationFuture<Void>) future).getBytesToReserve();
   }
@@ -481,15 +498,26 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l
     if (memoryReservationFutures.isEmpty()) {
       return;
     }
-    Iterator<MemoryReservationFuture<Void>> iterator = memoryReservationFutures.iterator();
+    Iterator<MemoryReservationFuture<Void>> iterator = memoryReservationFutures.values().iterator();
     while (iterator.hasNext()) {
       MemoryReservationFuture<Void> future = iterator.next();
       if (tryCompleteReservation(future)) {
-        iterator.remove();
+        removePendingReservation(future);
       }
     }
   }
 
+  private void removePendingReservation(MemoryReservationFuture<Void> future) {
+    memoryReservationFutures.remove(future.reservationId, future);
+    Map<String, FragmentMemoryReservation> fragments =
+        queryMemoryReservations.get(future.getQueryId());
+    FragmentMemoryReservation fragment =
+        fragments == null ? null : fragments.get(future.getFragmentInstanceId());
+    if (fragment != null) {
+      fragment.pendingReservations.remove(future.reservationId, future);
+    }
+  }
+
   private boolean tryCompleteReservation(MemoryReservationFuture<Void> future) {
     synchronized (future) {
       if (future.isDone()) {
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 9e5cd2e0ea00..055bc4b19f98 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
@@ -324,7 +324,8 @@ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) {
         sharedBudget ? inputMaxBytes : 0, sharedBudget ? outputMaxBytes : 0);
     if (sharedBudget) {
       memoryPool.registerFragmentInstanceSharedBudget(
-          instanceContext.getId().getQueryId().getId(), instanceContext.getId().getFullId());
+          instanceContext.getId().getQueryId().getId(),
+          instanceContext.getId().getFragmentInstanceId());
       collectOperators.forEach(CollectOperator::enableReadyInputConsumption);
       LOGGER.debug(
           DataNodeQueryMessages.LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109,
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 25957b97897d..d2b901d7f235 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.AbstractCollection;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -41,6 +48,7 @@
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.awaitility.Awaitility.await;
 
@@ -259,6 +267,203 @@ public void testDeregisterDoesNotCancelOtherFragmentPendingReservation() {
     Assert.assertEquals(0, pool.getReservedBytes());
   }
 
+  @Test
+  public void testOwnerCancellationAndDeregistrationDoNotScanUnrelatedWaiters() throws Exception {
+    assertDeregistrationDoesNotScanUnrelatedWaiters(true);
+  }
+
+  @Test
+  public void testDeregistrationCancelsOnlyIndexedWaiters() throws Exception {
+    assertDeregistrationDoesNotScanUnrelatedWaiters(false);
+  }
+
+  private void assertDeregistrationDoesNotScanUnrelatedWaiters(boolean ownerCancels)
+      throws Exception {
+    AtomicInteger visits = countGlobalPendingVisits();
+    pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512);
+    pool.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "p", MemorySide.INPUT);
+    pool.tryReserve("holder", "fi", "p", 512, 512);
+    int count = 1024;
+    List<MemoryReservationResult> victims = new ArrayList<>();
+    List<MemoryReservationResult> unrelated = new ArrayList<>();
+    for (int i = 0; i < count; i++) {
+      String fi = Integer.toString(i);
+      pool.registerPlanNodeIdToQueryMemoryMap("victim", fi, "p", MemorySide.INPUT);
+      pool.registerPlanNodeIdToQueryMemoryMap("unrelated", fi, "p", MemorySide.INPUT);
+      victims.add(pool.reserveWithPriority("victim", fi, "p", 8, 512, false));
+      unrelated.add(pool.reserveWithPriority("unrelated", fi, "p", 8, 512, false));
+    }
+    for (int i = 0; i < count; i++) {
+      if (ownerCancels) {
+        Assert.assertEquals(8, pool.tryCancel(victims.get(i).getFuture()));
+      }
+      pool.deRegisterFragmentInstanceFromQueryMemoryMap("victim", Integer.toString(i), true);
+      Assert.assertTrue(victims.get(i).getFuture().isCancelled());
+      Assert.assertFalse(unrelated.get(i).getFuture().isDone());
+    }
+    // An operation-count assertion avoids timing noise and detects a global scan for every FI.
+    Assert.assertEquals(0, visits.get());
+    Assert.assertEquals(count, pool.getMemoryReservationSize());
+    for (int i = 0; i < count; i++) {
+      Assert.assertEquals(8, pool.tryCancel(unrelated.get(i).getFuture()));
+      pool.deRegisterFragmentInstanceFromQueryMemoryMap("unrelated", Integer.toString(i), true);
+    }
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512);
+    pool.free("holder", "fi", "p", 512);
+    pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true);
+    pool.deRegisterFragmentInstanceFromQueryMemoryMap("holder", "fi", true);
+    Assert.assertEquals(0, pool.getMemoryReservationSize());
+    Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+    Assert.assertEquals(0, pool.getReservedBytes());
+  }
+
+  @SuppressWarnings("unchecked")
+  private AtomicInteger countGlobalPendingVisits() throws Exception {
+    AtomicInteger visits = new AtomicInteger();
+    Field field = MemoryPool.class.getDeclaredField("memoryReservationFutures");
+    field.setAccessible(true);
+    Object pending = field.get(pool);
+    if (pending instanceof Queue) {
+      // Keep this branch so the regression also runs against the former queue implementation.
+      Queue<MemoryPool.MemoryReservationFuture<Void>> counted =
+          new ConcurrentLinkedQueue<MemoryPool.MemoryReservationFuture<Void>>() {
+            @Override
+            public Iterator<MemoryPool.MemoryReservationFuture<Void>> iterator() {
+              return countVisits(super.iterator(), visits);
+            }
+          };
+      counted.addAll((Queue<MemoryPool.MemoryReservationFuture<Void>>) pending);
+      field.set(pool, counted);
+    } else {
+      Map<Long, MemoryPool.MemoryReservationFuture<Void>> counted =
+          new ConcurrentSkipListMap<Long, MemoryPool.MemoryReservationFuture<Void>>() {
+            @Override
+            public Collection<MemoryPool.MemoryReservationFuture<Void>> values() {
+              Collection<MemoryPool.MemoryReservationFuture<Void>> values = super.values();
+              return new AbstractCollection<MemoryPool.MemoryReservationFuture<Void>>() {
+                @Override
+                public Iterator<MemoryPool.MemoryReservationFuture<Void>> iterator() {
+                  return countVisits(values.iterator(), visits);
+                }
+
+                @Override
+                public int size() {
+                  return values.size();
+                }
+              };
+            }
+          };
+      counted.putAll((Map<Long, MemoryPool.MemoryReservationFuture<Void>>) pending);
+      field.set(pool, counted);
+    }
+    return visits;
+  }
+
+  private static <T> Iterator<T> countVisits(Iterator<T> delegate, AtomicInteger visits) {
+    return new Iterator<T>() {
+      @Override
+      public boolean hasNext() {
+        return delegate.hasNext();
+      }
+
+      @Override
+      public T next() {
+        visits.incrementAndGet();
+        return delegate.next();
+      }
+
+      @Override
+      public void remove() {
+        delegate.remove();
+      }
+    };
+  }
+
+  @Test
+  public void testCancellationKeepsFifoWakeupOrder() {
+    pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512);
+    pool.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "p", MemorySide.INPUT);
+    pool.tryReserve("holder", "fi", "p", 512, 512);
+    List<MemoryReservationResult> requests = new ArrayList<>();
+    for (int i = 0; i < 3; i++) {
+      String fi = Integer.toString(i);
+      pool.registerPlanNodeIdToQueryMemoryMap("waiting", fi, "p", MemorySide.INPUT);
+      requests.add(pool.reserveWithPriority("waiting", fi, "p", 64, 512, false));
+    }
+    Assert.assertEquals(64, pool.tryCancel(requests.get(1).getFuture()));
+    Assert.assertEquals(0, pool.tryCancel(requests.get(1).getFuture()));
+    Assert.assertEquals(2, pool.getMemoryReservationSize());
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 64);
+    Assert.assertTrue(requests.get(0).getFuture().isDone());
+    Assert.assertFalse(requests.get(2).getFuture().isDone());
+    pool.free("waiting", "0", "p", 64);
+    Assert.assertTrue(requests.get(2).getFuture().isDone());
+    pool.free("waiting", "2", "p", 64);
+    pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 448);
+    pool.free("holder", "fi", "p", 512);
+    for (int i = 0; i < 3; i++) {
+      pool.deRegisterFragmentInstanceFromQueryMemoryMap("waiting", Integer.toString(i), true);
+    }
+    Assert.assertEquals(0, pool.getMemoryReservationSize());
+    Assert.assertEquals(0, pool.getReservedBytes());
+  }
+
+  @Test
+  public void testConcurrentCancellationAndGrantRemoveAllPendingIndexes() throws Exception {
+    ExecutorService executor = Executors.newFixedThreadPool(2);
+    try {
+      for (int round = 0; round < 32; round++) {
+        MemoryPool concurrent = new MemoryPool("race", new MemoryManager(1024), 512);
+        concurrent.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "a", MemorySide.INPUT);
+        concurrent.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "b", MemorySide.INPUT);
+        concurrent.tryReserve("holder", "fi", "a", 512, 512);
+        concurrent.tryReserve("holder", "fi", "b", 512, 512);
+        List<MemoryReservationResult> requests = new ArrayList<>();
+        for (int i = 0; i < 16; i++) {
+          String fi = Integer.toString(i);
+          concurrent.registerPlanNodeIdToQueryMemoryMap("waiting", fi, "p", MemorySide.INPUT);
+          requests.add(concurrent.reserveWithPriority("waiting", fi, "p", 16, 512, false));
+        }
+        CountDownLatch start = new CountDownLatch(1);
+        Future<?> grant =
+            executor.submit(
+                () -> {
+                  start.await();
+                  concurrent.free("holder", "fi", "a", 512);
+                  return null;
+                });
+        Future<?> cancel =
+            executor.submit(
+                () -> {
+                  start.await();
+                  for (MemoryReservationResult request : requests) {
+                    concurrent.tryCancel(request.getFuture());
+                  }
+                  return null;
+                });
+        start.countDown();
+        grant.get(10, TimeUnit.SECONDS);
+        cancel.get(10, TimeUnit.SECONDS);
+        for (int i = 0; i < requests.size(); i++) {
+          Assert.assertTrue(requests.get(i).getFuture().isDone());
+          if (!requests.get(i).getFuture().isCancelled()) {
+            concurrent.free("waiting", Integer.toString(i), "p", 16);
+          }
+          concurrent.deRegisterFragmentInstanceFromQueryMemoryMap(
+              "waiting", Integer.toString(i), true);
+        }
+        concurrent.free("holder", "fi", "b", 512);
+        concurrent.deRegisterFragmentInstanceFromQueryMemoryMap("holder", "fi", true);
+        Assert.assertEquals(0, concurrent.getMemoryReservationSize());
+        Assert.assertEquals(0, concurrent.getQueryMemoryReservationSize());
+        Assert.assertEquals(0, concurrent.getReservedBytes());
+      }
+    } finally {
+      executor.shutdownNow();
+      Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+    }
+  }
+
   @Test
   public void testTryReserve() {
 
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
index 49623ab9e45e..2b3c847e89e4 100644
--- 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
@@ -24,6 +24,8 @@
 import org.apache.iotdb.calc.execution.operator.process.CollectOperator;
 import org.apache.iotdb.calc.execution.operator.process.TableMergeSortOperator;
 import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.memory.MemoryManager;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
@@ -35,6 +37,7 @@
 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.MPPDataExchangeManager.SourceHandleListener;
 import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle;
 import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle;
 import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle;
@@ -43,11 +46,13 @@
 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.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.db.queryengine.execution.operator.source.ExchangeOperator;
 import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
+import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
 
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.util.concurrent.SettableFuture;
@@ -55,6 +60,7 @@
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.read.common.block.TsBlock;
 import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.read.common.block.column.TsBlockSerde;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -63,6 +69,8 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -188,7 +196,7 @@ private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit
       assertEquals(0, instance.getInputMaxBytes());
       assertEquals(0, instance.getOutputMaxBytes());
     }
-    String fi = context.getFragmentInstanceId().getFullId();
+    String fi = context.getFragmentInstanceId().getFragmentInstanceId();
     pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a" + side, side);
     pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b" + side, side);
     assertTrue(
@@ -206,6 +214,56 @@ private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit
     pool.free("budget", fi, "b" + side, limit);
   }
 
+  @Test
+  public void testRealSourceHandlesShareCanonicalFragmentAccountAndCleanup() throws Exception {
+    FragmentInstanceId id = context.getFragmentInstanceId();
+    String fi = FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(id.toThrift());
+    LocalMemoryManager manager = mock(LocalMemoryManager.class);
+    when(manager.getQueryPool()).thenReturn(pool);
+    SourceHandle first = realSourceHandle(id, "first", manager);
+    SourceHandle second = realSourceHandle(id, "second", manager);
+    for (SourceHandle source : Arrays.asList(first, second)) {
+      ExchangeOperator exchange = mock(ExchangeOperator.class);
+      when(exchange.getSourceHandle()).thenReturn(source);
+      context.addExchangeOperator(exchange);
+    }
+    try {
+      context.setExchangeMemoryLimits(collect(), pool);
+      assertEquals(1024, context.getInstanceContext().getInputMaxBytes());
+      first.isBlocked();
+      second.isBlocked();
+      first.updatePendingDataBlockInfo(0, List.of(1024L));
+      second.updatePendingDataBlockInfo(0, List.of(1024L));
+      assertEquals(1024, pool.getFragmentMemoryReservedBytes("budget", fi, MemorySide.INPUT));
+      assertEquals(0, pool.getPlanNodeMemoryReservedBytes("budget", fi, "second"));
+      assertEquals(1, pool.getMemoryReservationSize());
+      first.abort();
+      assertEquals(1024, pool.getPlanNodeMemoryReservedBytes("budget", fi, "second"));
+    } finally {
+      first.abort();
+      second.abort();
+      pool.deRegisterFragmentInstanceFromQueryMemoryMap("budget", fi, true);
+    }
+    assertEquals(0, pool.getMemoryReservationSize());
+    assertEquals(0, pool.getReservedBytes());
+    assertEquals(0, pool.getQueryMemoryReservationSize());
+  }
+
+  private static SourceHandle realSourceHandle(
+      FragmentInstanceId id, String planNode, LocalMemoryManager manager) {
+    return new SourceHandle(
+        new TEndPoint("remote", 12345),
+        new TFragmentInstanceId(id.getQueryId().getId(), 1, "0"),
+        id.toThrift(),
+        planNode,
+        0,
+        manager,
+        mock(ExecutorService.class),
+        new TsBlockSerde(),
+        mock(SourceHandleListener.class),
+        mock(IClientManager.class));
+  }
+
   @Test
   public void testPureCollectUsesSharedSideBudgets() {
     ISourceHandle source = addSource(context);
@@ -515,7 +573,7 @@ public void testMergeSortGetsBothHeadsDespitePrefetchOnOneInput() throws Excepti
     context.setExchangeMemoryLimits(
         node(PlanNodeType.TABLE_MERGESORT_NODE, collect(), node(PlanNodeType.TABLE_EXCHANGE_NODE)),
         pool);
-    String fi = context.getFragmentInstanceId().getFullId();
+    String fi = context.getFragmentInstanceId().getFragmentInstanceId();
     pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", MemorySide.INPUT);
     pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", MemorySide.INPUT);
     for (int i = 0; i < 4; i++) {
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java
index ac32c831c43e..31c562199fd6 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java
@@ -130,7 +130,7 @@ public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() th
       verify(secondHandle).setMaxBytesCanReserve(1024);
 
       String query = instance.getId().getQueryId().getId();
-      String fi = instance.getId().getFullId();
+      String fi = instance.getId().getFragmentInstanceId();
       pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "first", MemorySide.INPUT);
       pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "second", MemorySide.INPUT);
       assertTrue(
@@ -144,7 +144,7 @@ public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() th
       collect.close();
     } finally {
       pool.deRegisterFragmentInstanceFromQueryMemoryMap(
-          instance.getId().getQueryId().getId(), instance.getId().getFullId(), true);
+          instance.getId().getQueryId().getId(), instance.getId().getFragmentInstanceId(), true);
       long reserved = freeBefore - planner.getFreeMemoryForOperators();
       if (reserved > 0) {
         planner.releaseToFreeMemoryForOperators(reserved);