[To dev/1.3] clone partial columns of aligned tvlist for query

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java

@@ -69,8 +69,10 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
@@ -162,6 +164,9 @@ public class FragmentInstanceContext extends QueryContext {
   private long closedUnseqFileNum = 0;
   private boolean highestPriority = false;
 
+  // accessed value columns on each referenced AlignedTVList.
+  private final Map<TVList, Set<Integer>> alignedTVListColumnAccessMap = new ConcurrentHashMap<>();
+
   public static FragmentInstanceContext createFragmentInstanceContext(
       FragmentInstanceId id,
       FragmentInstanceStateMachine stateMachine,
@@ -218,6 +223,48 @@ public void setQueryDataSourceType(QueryDataSourceType queryDataSourceType) {
     this.queryDataSourceType = queryDataSourceType;
   }
 
+  /**
+   * Record columns of the AlignedTVList accessed by the query. This method is called from
+   * prepareTvListMapForQuery with tvList.lockQueryList() held. Even though the HashSet inside
+   * alignedTVListColumnAccessMap is not thread-safe, the calling pattern guarantees thread safety
+   * without requiring additional synchronization.
+   *
+   * @param tvList the TVList being accessed
+   * @param columnIndexList list of column indices being accessed
+   */
+  public void putAccessedColumns(TVList tvList, List<Integer> columnIndexList) {
+    Set<Integer> accessedColumns =
+        alignedTVListColumnAccessMap.computeIfAbsent(tvList, ignored -> new HashSet<>());
+    columnIndexList.stream()
+        .filter(Objects::nonNull)
+        .forEach(
+            index -> {
+              if (index >= 0) {
+                accessedColumns.add(index);
+              }
+            });
+  }
+
+  /** Remove column-access metadata for an unpublished TVList when clone preparation fails. */
+  public void removeAccessedColumns(TVList tvList) {
+    alignedTVListColumnAccessMap.remove(tvList);
+  }
+
+  /**
+   * Get columns of the AlignedTVList accessed by the query. This method is called from
+   * prepareTvListMapForQuery with tvList.lockQueryList() held, ensuring that no other thread can
+   * change accessed columns for the same TVList concurrently.
+   *
+   * @param tvList the TVList being accessed
+   * @return set of column indices being accessed
+   */
+  public Set<Integer> getAccessedAlignedColumns(TVList tvList) {
+    Set<Integer> accessedColumns = alignedTVListColumnAccessMap.get(tvList);
+    return accessedColumns == null
+        ? Collections.emptySet()
+        : Collections.unmodifiableSet(accessedColumns);
+  }
+
   @TestOnly
   public static FragmentInstanceContext createFragmentInstanceContext(
       FragmentInstanceId id, FragmentInstanceStateMachine stateMachine) {
@@ -897,12 +944,12 @@ public void releaseResourceWhenAllDriversAreClosed() {
    */
   private void releaseTVListOwnedByQuery() {
     for (TVList tvList : tvListSet) {
-      long tvListRamSize = tvList.calculateRamSize().getRamSize();
       tvList.lockQueryList();
       Set<QueryContext> queryContextSet = tvList.getQueryContextSet();
       try {
         queryContextSet.remove(this);
         if (tvList.getOwnerQuery() == this) {
+          long tvListRamSize = tvList.calculateRamSize().getRamSize();
           if (tvList.getReservedMemoryBytes() != tvListRamSize) {
             LOGGER.warn(
                 "Release TVList owned by query: allocate size {}, release size {}",
@@ -980,6 +1027,7 @@ public synchronized void releaseResource() {
 
     // release TVList/AlignedTVList owned by current query
     releaseTVListOwnedByQuery();
+    alignedTVListColumnAccessMap.clear();
 
     fileModCache = null;
     nonExistentModFiles = null;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java

@@ -32,6 +32,9 @@ public void reserveMemoryImmediately() {}
   @Override
   public void releaseMemoryCumulatively(long size) {}
 
+  @Override
+  public void releaseMemoryImmediately(long size) {}
+
   @Override
   public void releaseAllReservedMemory() {}
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java

@@ -40,6 +40,12 @@ public interface MemoryReservationManager {
    */
   void releaseMemoryCumulatively(final long size);
 
+  /**
+   * Release the given size immediately. This is used to roll back a reservation when the operation
+   * protected by that reservation fails before ownership is published.
+   */
+  void releaseMemoryImmediately(final long size);
+
   /**
    * Release all reserved memory immediately. Make sure this method is called when the lifecycle of
    * this manager ends, Or the memory to be released in the batch may not be released correctly.

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java

@@ -79,7 +79,14 @@ public long getFallbackBytesInTotalForTest() {
   public void reserveMemoryCumulatively(final long size) {
     bytesToBeReserved += size;
     if (bytesToBeReserved >= MEMORY_BATCH_THRESHOLD) {
-      reserveMemoryImmediately();
+      try {
+        reserveMemoryImmediately();
+      } catch (RuntimeException | Error failure) {
+        // reserveMemoryImmediately can fail only while asking the planner for memory, before it
+        // updates this manager's counters. Keep the caller-visible reservation operation atomic.
+        bytesToBeReserved -= size;
+        throw failure;
+      }
     }
   }
 
@@ -127,6 +134,13 @@ public void releaseMemoryCumulatively(final long size) {
     }
   }
 
+  @Override
+  public void releaseMemoryImmediately(final long size) {
+    if (size > 0) {
+      releaseBytesImmediately(size);
+    }
+  }
+
   private void releaseBytesImmediately(final long size) {
     long poolBytes = deductReleaseAccounting(size);
     if (poolBytes > 0) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java

@@ -51,6 +51,11 @@ public synchronized void releaseMemoryCumulatively(long size) {
     super.releaseMemoryCumulatively(size);
   }
 
+  @Override
+  public synchronized void releaseMemoryImmediately(long size) {
+    super.releaseMemoryImmediately(size);
+  }
+
   @Override
   public synchronized void releaseAllReservedMemory() {
     super.releaseAllReservedMemory();