From b975e6c1171df62488e4a3e0b2fa01933e3c308b Mon Sep 17 00:00:00 2001 From: shizy Date: Tue, 4 Aug 2026 17:11:02 +0800 Subject: [PATCH 01/16] clone partial columns of AlignedTVList during query # Conflicts: # iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java --- .../fragment/FragmentInstanceContext.java | 45 ++- .../utils/ResourceByPathUtils.java | 276 +++++++++++------- .../memtable/AbstractWritableMemChunk.java | 2 +- .../memtable/AlignedReadOnlyMemChunk.java | 4 +- .../dataregion/memtable/ReadOnlyMemChunk.java | 4 +- .../db/utils/datastructure/AlignedTVList.java | 185 ++++++++++-- .../FragmentInstanceExecutionTest.java | 91 ++++++ .../datastructure/AlignedTVListTest.java | 31 ++ 8 files changed, 497 insertions(+), 141 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 70cc46428be9a..62c581ad91a2f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/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> alignedTVListColumnAccessMap = new ConcurrentHashMap<>(); + public static FragmentInstanceContext createFragmentInstanceContext( FragmentInstanceId id, FragmentInstanceStateMachine stateMachine, @@ -218,6 +223,43 @@ 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 columnIndexList) { + Set accessedColumns = + alignedTVListColumnAccessMap.computeIfAbsent(tvList, ignored -> new HashSet<>()); + columnIndexList.stream() + .filter(Objects::nonNull) + .forEach( + index -> { + if (index >= 0) { + accessedColumns.add(index); + } + }); + } + + /** + * 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 getAccessedAlignedColumns(TVList tvList) { + Set accessedColumns = alignedTVListColumnAccessMap.get(tvList); + return accessedColumns == null + ? Collections.emptySet() + : Collections.unmodifiableSet(accessedColumns); + } + @TestOnly public static FragmentInstanceContext createFragmentInstanceContext( FragmentInstanceId id, FragmentInstanceStateMachine stateMachine) { @@ -897,12 +939,12 @@ public void releaseResourceWhenAllDriversAreClosed() { */ private void releaseTVListOwnedByQuery() { for (TVList tvList : tvListSet) { - long tvListRamSize = tvList.calculateRamSize().getRamSize(); tvList.lockQueryList(); Set 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 +1022,7 @@ public synchronized void releaseResource() { // release TVList/AlignedTVList owned by current query releaseTVListOwnedByQuery(); + alignedTVListColumnAccessMap.clear(); fileModCache = null; nonExistentModFiles = null; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index fa2f603d6facc..adf2198019a63 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -38,6 +38,7 @@ import org.apache.iotdb.db.storageengine.dataregion.modification.Modification; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.iotdb.db.utils.ModificationUtils; +import org.apache.iotdb.db.utils.datastructure.AlignedTVList; import org.apache.iotdb.db.utils.datastructure.TVList; import org.apache.tsfile.enums.TSDataType; @@ -63,9 +64,11 @@ import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Set; import static org.apache.iotdb.commons.path.AlignedPath.VECTOR_PLACEHOLDER; @@ -121,7 +124,8 @@ protected Map prepareTvListMapForQuery( QueryContext context, IWritableMemChunk memChunk, boolean isWorkMemTable, - Filter globalTimeFilter) { + Filter globalTimeFilter, + List columnIndexList) { // should copy globalTimeFilter because GroupByMonthFilter is stateful Filter copyTimeFilter = null; if (globalTimeFilter != null) { @@ -146,113 +150,157 @@ protected Map prepareTvListMapForQuery( } } - // mutable tvlist - TVList list = memChunk.getWorkingTVList(); - TVList cloneList = null; - TVList.RamInfo listRamInfo = list.calculateRamSize(); - list.lockQueryList(); - try { - if (copyTimeFilter != null - && !copyTimeFilter.satisfyStartEndTime(list.getMinTime(), list.getMaxTime())) { - return tvListQueryMap; - } + TVList.RamInfo listRamInfo = null; + + // calculateRamSize (synchronized method on TVList) was previously called before + // lockQueryList to avoid deadlock concerns. For partial clone of AlignedTVList, however + // calculateRamSize must now be called inside the lockQueryList section because it depends on + // accessing columns on the AlignedTVList. + // This is safe because the lock ordering — queryListLock must always be acquired before the + // TVList intrinsic lock (via synchronized methods like calculateRamSize, clone). So no AB-BA + // deadlock is possible. + while (true) { + // The working TVList may be replaced by another query. Always lock the actual list being used + // and verify it is still the current working list before touching query-owned fields. + TVList list = memChunk.getWorkingTVList(); + list.lockQueryList(); + try { + if (list != memChunk.getWorkingTVList()) { + continue; + } + if (copyTimeFilter != null + && !copyTimeFilter.satisfyStartEndTime(list.getMinTime(), list.getMaxTime())) { + return tvListQueryMap; + } - if (!isWorkMemTable) { - /* - * 1. Q1 queries this TVList while it is still in the working memtable and records a smaller - * visible row count. - * 2. Later writes append out-of-order rows to the same TVList, then FLUSH moves the - * memtable to the flushing list. - * 3. Q2 queries the flushing memtable. If Q2 directly reuses the original mutable TVList, - * Q2's query-side sort may reorder the indices in place. - * 4. Q1 continues to read with its old row count and the reordered indices. The converted - * value index can exceed Q1's bitmap range and cause out-of-bound access. - * - * Therefore, this flushing branch can reuse the original list only when it is already - * sorted or no active query is using it. Otherwise, Q2 should read from - * workingListForFlush. - */ - boolean canUseListDirectly = list.isSorted() || list.getQueryContextSet().isEmpty(); - LOGGER.debug( - "Flushing MemTable - add current query context to mutable TVList's query list"); - if (canUseListDirectly) { - list.getQueryContextSet().add(context); - tvListQueryMap.put(list, list.rowCount()); - } else { - TVList workingListForFlushSort = memChunk.initWorkingListForFlushIfNecessary(list, true); + if (!isWorkMemTable) { /* - * The query will read from workingListForFlushSort, but cloneForFlushSort() only clones - * times and indices. The value arrays and bitmaps are still shared with the original - * list. - * - * Therefore, this query must also hold the original list until it finishes. Adding - * context to list.getQueryContextSet() lets flush/query cleanup see that the original - * list is still in use. Adding list to context.tvListSet makes - * releaseTVListOwnedByQuery() remove this context from the original list later. + * 1. Q1 queries this TVList while it is still in the working memtable and records a smaller + * visible row count. + * 2. Later writes append out-of-order rows to the same TVList, then FLUSH moves the + * memtable to the flushing list. + * 3. Q2 queries the flushing memtable. If Q2 directly reuses the original mutable TVList, + * Q2's query-side sort may reorder the indices in place. + * 4. Q1 continues to read with its old row count and the reordered indices. The converted + * value index can exceed Q1's bitmap range and cause out-of-bound access. * - * Do not put the original list into tvListQueryMap here. The actual read path must use - * workingListForFlushSort to avoid sorting the original list in place. + * Therefore, this flushing branch can reuse the original list only when it is already + * sorted or no active query is using it. Otherwise, Q2 should read from + * workingListForFlush. */ - list.getQueryContextSet().add(context); - context.addTVListToSet(Collections.singleton(list)); - workingListForFlushSort.getQueryContextSet().add(context); - tvListQueryMap.put(workingListForFlushSort, workingListForFlushSort.rowCount()); - } - } else { - if (list.isSorted() || list.getQueryContextSet().isEmpty()) { + boolean canUseListDirectly = list.isSorted() || list.getQueryContextSet().isEmpty(); LOGGER.debug( - "Working MemTable - add current query context to mutable TVList's query list when it's sorted or no other query on it"); - list.getQueryContextSet().add(context); - tvListQueryMap.put(list, list.rowCount()); - } else { - /* - * +----------------------+ - * | MemTable | - * | | - * | +------------+ | +-----------------+ - * | | TVList |<---+--+ +---+ Previous Query | - * | +-----^------+ | | | +-----------------+ - * | | | | | - * +----------+-----------+ | | +----------------+ - * | Clone +---+---+ Current Query | - * +-----+------+ | +----------------+ - * | TVList | <---------+ - * +------------+ - */ - LOGGER.debug( - "Working MemTable - clone mutable TVList and replace old TVList in working MemTable"); - QueryContext firstQuery = list.getQueryContextSet().iterator().next(); - // reserve query memory - if (firstQuery instanceof FragmentInstanceContext) { - MemoryReservationManager memoryReservationManager = - ((FragmentInstanceContext) firstQuery).getMemoryReservationContext(); - memoryReservationManager.reserveMemoryCumulatively(listRamInfo.getRamSize()); - list.setReservedMemoryBytes(listRamInfo.getRamSize()); + "Flushing MemTable - add current query context to mutable TVList's query list"); + if (canUseListDirectly) { + list.getQueryContextSet().add(context); + tvListQueryMap.put(list, list.rowCount()); + } else { + TVList workingListForFlushSort = + memChunk.initWorkingListForFlushIfNecessary(list, true); + /* + * The query will read from workingListForFlushSort, but cloneForFlushSort() only clones + * times and indices. The value arrays and bitmaps are still shared with the original + * list. + * + * Therefore, this query must also hold the original list until it finishes. Adding + * context to list.getQueryContextSet() lets flush/query cleanup see that the original + * list is still in use. Adding list to context.tvListSet makes + * releaseTVListOwnedByQuery() remove this context from the original list later. + * + * Do not put the original list into tvListQueryMap here. The actual read path must use + * workingListForFlushSort to avoid sorting the original list in place. + */ + list.getQueryContextSet().add(context); + context.addTVListToSet(Collections.singleton(list)); + workingListForFlushSort.getQueryContextSet().add(context); + tvListQueryMap.put(workingListForFlushSort, workingListForFlushSort.rowCount()); } - list.setOwnerQuery(firstQuery); + } else { + if (list.isSorted() || list.getQueryContextSet().isEmpty()) { + LOGGER.debug( + "Working MemTable - add current query context to mutable TVList's query list when it's sorted or no other query on it"); + list.getQueryContextSet().add(context); + tvListQueryMap.put(list, list.rowCount()); + + // columnIndexList is to track column-level access for AlignedTVList. + // For TVList (primitive time series), it remains null and column tracking is not + // needed. + if (columnIndexList != null && context instanceof FragmentInstanceContext) { + ((FragmentInstanceContext) context).putAccessedColumns(list, columnIndexList); + } + } else { + /* + * +----------------------+ + * | MemTable | + * | | + * | +------------+ | +-----------------+ + * | | TVList |<---+--+ +---+ Previous Query | + * | +-----^------+ | | | +-----------------+ + * | | | | | + * +----------+-----------+ | | +----------------+ + * | Clone +---+---+ Current Query | + * +-----+------+ | +----------------+ + * | TVList | <---------+ + * +------------+ + */ + LOGGER.debug( + "Working MemTable - clone mutable TVList and replace old TVList in working MemTable"); + + Set columnsToClone = getAccessedColumnsForQuery(list); + listRamInfo = + (columnsToClone == null) + ? list.calculateRamSize() + : ((AlignedTVList) list).calculateRamSize(columnsToClone); + + // reserve query memory + QueryContext firstQuery = list.getQueryContextSet().iterator().next(); + if (firstQuery instanceof FragmentInstanceContext) { + MemoryReservationManager memoryReservationManager = + ((FragmentInstanceContext) firstQuery).getMemoryReservationContext(); + memoryReservationManager.reserveMemoryCumulatively(listRamInfo.getRamSize()); + list.setReservedMemoryBytes(listRamInfo.getRamSize()); + } + list.setOwnerQuery(firstQuery); + + // clone TVList + TVList cloneList = + (columnsToClone == null) + ? list.clone() + : ((AlignedTVList) list).clone(columnsToClone); + + cloneList.getQueryContextSet().add(context); + tvListQueryMap.put(cloneList, cloneList.rowCount()); + if (columnIndexList != null && context instanceof FragmentInstanceContext) { + ((FragmentInstanceContext) context).putAccessedColumns(cloneList, columnIndexList); + } - // clone TVList - cloneList = list.clone(); - cloneList.getQueryContextSet().add(context); - tvListQueryMap.put(cloneList, cloneList.rowCount()); + if (columnsToClone != null) { + ((AlignedTVList) list) + .moveUnclonedColumnsTo((AlignedTVList) cloneList, columnsToClone); + } + memChunk.setWorkingTVList(cloneList); + } } + } catch (MemoryNotEnoughException ex) { + if (listRamInfo != null) { + LOGGER.warn( + "Failed to reserve memory for TVList: ramSize {}, timestampsSize {}, arrayMemCost {}, rowCount {}, dataTypes {}", + listRamInfo.getRamSize(), + listRamInfo.getTimestampsSize(), + listRamInfo.getArrayMemCost(), + listRamInfo.getRowCount(), + listRamInfo.getDataTypes()); + } + throw ex; + } finally { + list.unlockQueryList(); } - } catch (MemoryNotEnoughException ex) { - LOGGER.warn( - "Failed to reserve memory for TVList: ramSize {}, timestampsSize {}, arrayMemCost {}, rowCount {}, dataTypes {}", - listRamInfo.getRamSize(), - listRamInfo.getTimestampsSize(), - listRamInfo.getArrayMemCost(), - listRamInfo.getRowCount(), - listRamInfo.getDataTypes()); - throw ex; - } finally { - list.unlockQueryList(); - } - if (cloneList != null) { - memChunk.setWorkingTVList(cloneList); + return tvListQueryMap; } - return tvListQueryMap; + } + + protected Set getAccessedColumnsForQuery(TVList tvList) { + return null; } } @@ -400,15 +448,15 @@ public ReadOnlyMemChunk getReadOnlyMemChunkFromMemTable( return null; } - // prepare AlignedTVList for query. It should clone TVList if necessary. - Map alignedTvListQueryMap = - prepareTvListMapForQuery( - context, alignedMemChunk, modsToMemtable == null, globalTimeFilter); - // column index list for the query List columnIndexList = alignedMemChunk.buildColumnIndexList(partialPath.getSchemaList()); + // prepare AlignedTVList for query. It should clone TVList if necessary. + Map alignedTvListQueryMap = + prepareTvListMapForQuery( + context, alignedMemChunk, modsToMemtable == null, globalTimeFilter, columnIndexList); + List> deletionList = null; if (modsToMemtable != null) { deletionList = @@ -419,6 +467,26 @@ public ReadOnlyMemChunk getReadOnlyMemChunkFromMemTable( context, columnIndexList, getMeasurementSchema(), alignedTvListQueryMap, deletionList); } + /** + * This method is called from prepareTvListMapForQuery with tvList.lockQueryList() held, ensuring + * thread-safe access to queryContextSet. + * + * @param tvList the TVList to get accessed columns for + * @return set of accessed column indices, or empty set if no columns are tracked + */ + @Override + protected Set getAccessedColumnsForQuery(TVList tvList) { + Set accessedColumns = new HashSet<>(); + for (QueryContext queryContext : tvList.getQueryContextSet()) { + if (!(queryContext instanceof FragmentInstanceContext)) { + return null; + } + accessedColumns.addAll( + ((FragmentInstanceContext) queryContext).getAccessedAlignedColumns(tvList)); + } + return accessedColumns; + } + public VectorMeasurementSchema getMeasurementSchema() { List measurementList = partialPath.getMeasurementList(); TSDataType[] types = new TSDataType[measurementList.size()]; @@ -571,7 +639,7 @@ public ReadOnlyMemChunk getReadOnlyMemChunkFromMemTable( memTableMap.get(deviceID).getMemChunkMap().get(partialPath.getMeasurement()); // prepare TVList for query. It should clone TVList if necessary. Map tvListQueryMap = - prepareTvListMapForQuery(context, memChunk, modsToMemtable == null, globalTimeFilter); + prepareTvListMapForQuery(context, memChunk, modsToMemtable == null, globalTimeFilter, null); List deletionList = null; if (modsToMemtable != null) { deletionList = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java index 6c773942fb72b..13fa84d039e27 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java @@ -101,7 +101,6 @@ protected void maybeReleaseTvList(TVList tvList) { } private void tryReleaseTvList(TVList tvList) { - long tvListRamSize = tvList.calculateRamSize().getRamSize(); tvList.lockQueryList(); try { if (tvList.getQueryContextSet().isEmpty()) { @@ -113,6 +112,7 @@ private void tryReleaseTvList(TVList tvList) { if (firstQuery instanceof FragmentInstanceContext) { MemoryReservationManager memoryReservationManager = ((FragmentInstanceContext) firstQuery).getMemoryReservationContext(); + long tvListRamSize = tvList.calculateRamSize().getRamSize(); memoryReservationManager.reserveMemoryCumulatively(tvListRamSize); tvList.setReservedMemoryBytes(tvListRamSize); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedReadOnlyMemChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedReadOnlyMemChunk.java index bb2ee311d3047..f54af9cfcbb4a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedReadOnlyMemChunk.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedReadOnlyMemChunk.java @@ -122,12 +122,12 @@ public void sortTvLists() { // We must update queryRowCount here, otherwise, it may be used later to build // BitMaps, causing bitmap array size mismatch and possible out of bound. entry.setValue(alignedTvList.sort()); - long alignedTvListRamSize = alignedTvList.calculateRamSize().getRamSize(); alignedTvList.lockQueryList(); try { FragmentInstanceContext ownerQuery = (FragmentInstanceContext) alignedTvList.getOwnerQuery(); if (ownerQuery != null) { + long alignedTvListRamSize = alignedTvList.calculateRamSize().getRamSize(); long deltaBytes = alignedTvListRamSize - alignedTvList.getReservedMemoryBytes(); if (deltaBytes > 0) { ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes); @@ -367,12 +367,12 @@ public IPointReader getPointReader() { int queryLength = entry.getValue(); if (!alignedTvList.isSorted() && queryLength > alignedTvList.seqRowCount()) { entry.setValue(alignedTvList.sort()); - long alignedTvListRamSize = alignedTvList.calculateRamSize().getRamSize(); alignedTvList.lockQueryList(); try { FragmentInstanceContext ownerQuery = (FragmentInstanceContext) alignedTvList.getOwnerQuery(); if (ownerQuery != null) { + long alignedTvListRamSize = alignedTvList.calculateRamSize().getRamSize(); long deltaBytes = alignedTvListRamSize - alignedTvList.getReservedMemoryBytes(); if (deltaBytes > 0) { ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/ReadOnlyMemChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/ReadOnlyMemChunk.java index c0a71bf7edcdc..223e9ebb81142 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/ReadOnlyMemChunk.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/ReadOnlyMemChunk.java @@ -136,11 +136,11 @@ public void sortTvLists() { int queryRowCount = entry.getValue(); if (!tvList.isSorted() && queryRowCount > tvList.seqRowCount()) { entry.setValue(tvList.sort()); - long tvListRamSize = tvList.calculateRamSize().getRamSize(); tvList.lockQueryList(); try { FragmentInstanceContext ownerQuery = (FragmentInstanceContext) tvList.getOwnerQuery(); if (ownerQuery != null) { + long tvListRamSize = tvList.calculateRamSize().getRamSize(); long deltaBytes = tvListRamSize - tvList.getReservedMemoryBytes(); if (deltaBytes > 0) { ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes); @@ -288,11 +288,11 @@ public IPointReader getPointReader() { int queryLength = entry.getValue(); if (!tvList.isSorted() && queryLength > tvList.seqRowCount()) { entry.setValue(tvList.sort()); - long tvListRamSize = tvList.calculateRamSize().getRamSize(); tvList.lockQueryList(); try { FragmentInstanceContext ownerQuery = (FragmentInstanceContext) tvList.getOwnerQuery(); if (ownerQuery != null) { + long tvListRamSize = tvList.calculateRamSize().getRamSize(); long deltaBytes = tvListRamSize - tvList.getReservedMemoryBytes(); if (deltaBytes > 0) { ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index f787ff64fda16..bdf0633d6a430 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -53,6 +53,7 @@ import java.util.Arrays; import java.util.List; import java.util.Objects; +import java.util.Set; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -162,34 +163,46 @@ public synchronized AlignedTVList cloneForFlushSort() { public synchronized AlignedTVList clone() { AlignedTVList cloneList = AlignedTVList.newAlignedList(new ArrayList<>(dataTypes)); cloneAs(cloneList); - System.arraycopy( - memoryBinaryChunkSize, 0, cloneList.memoryBinaryChunkSize, 0, dataTypes.size()); + cloneColumnDataTo(cloneList, null); + return cloneList; + } + + public synchronized AlignedTVList clone(Set columnsToClone) { + AlignedTVList cloneList = AlignedTVList.newAlignedList(new ArrayList<>(dataTypes)); + cloneAs(cloneList); + cloneColumnDataTo(cloneList, columnsToClone); + return cloneList; + } + + public synchronized void moveUnclonedColumnsTo( + AlignedTVList cloneList, Set columnsToClone) { + if (columnsToClone == null) { + return; + } + if (bitMaps != null && cloneList.bitMaps == null) { + for (int i = 0; i < values.size(); i++) { + if (values.get(i) != null && !columnsToClone.contains(i) && bitMaps.get(i) != null) { + throw new IllegalStateException( + "Target AlignedTVList is not ready to receive moved bitmaps"); + } + } + } for (int i = 0; i < values.size(); i++) { - // Clone value List columnValues = values.get(i); - for (Object valueArray : columnValues) { - cloneList.values.get(i).add(cloneValue(dataTypes.get(i), valueArray)); + if (columnValues == null || columnsToClone.contains(i)) { + continue; } - // Clone bitmap in columnIndex - if (bitMaps != null && bitMaps.get(i) != null) { - List columnBitMaps = bitMaps.get(i); - if (cloneList.bitMaps == null) { - cloneList.bitMaps = new ArrayList<>(dataTypes.size()); - for (int j = 0; j < dataTypes.size(); j++) { - cloneList.bitMaps.add(null); - } - } - if (cloneList.bitMaps.get(i) == null) { - List cloneColumnBitMaps = new ArrayList<>(); - for (BitMap bitMap : columnBitMaps) { - cloneColumnBitMaps.add(bitMap == null ? null : bitMap.clone()); - } - cloneList.bitMaps.set(i, cloneColumnBitMaps); - } + cloneList.values.set(i, columnValues); + values.set(i, null); + if (bitMaps != null && bitMaps.get(i) != null && cloneList.bitMaps != null) { + cloneList.bitMaps.set(i, bitMaps.get(i)); + bitMaps.set(i, null); } + memoryBinaryChunkSize[i] = 0; } - cloneList.materializedBitmapMemoryCost = materializedBitmapMemoryCost; - return cloneList; + // Column ownership changed on both lists, so refresh their per-block memory cost. + refreshArrayMemCostWithoutIndex(); + cloneList.refreshArrayMemCostWithoutIndex(); } @SuppressWarnings("squid:S3776") // Suppress high Cognitive Complexity warning @@ -203,10 +216,14 @@ public synchronized void putAlignedValue(long timestamp, Object[] value) { timestamps.get(arrayIndex)[elementIndex] = timestamp; for (int i = 0; i < values.size(); i++) { Object columnValue = value[i]; - List columnValues = values.get(i); if (columnValue == null) { markNullValue(i, arrayIndex, elementIndex); } + List columnValues = values.get(i); + if (columnValues == null) { + throw new IllegalStateException( + String.format("Missing value arrays for aligned column index %d during append", i)); + } switch (dataTypes.get(i)) { case TEXT: case BLOB: @@ -604,6 +621,13 @@ public synchronized Pair delete( * or delete wrong rows. */ public synchronized void deleteColumn(int columnIndex) { + List columnValues = values.get(columnIndex); + if (columnValues == null) { + throw new IllegalStateException( + String.format( + "Missing value arrays for aligned column index %d during delete", columnIndex)); + } + if (bitMaps == null) { List> localBitMaps = new ArrayList<>(dataTypes.size()); for (int j = 0; j < dataTypes.size(); j++) { @@ -611,9 +635,10 @@ public synchronized void deleteColumn(int columnIndex) { } bitMaps = localBitMaps; } + if (bitMaps.get(columnIndex) == null) { List columnBitMaps = new ArrayList<>(); - for (int i = 0; i < values.get(columnIndex).size(); i++) { + for (int i = 0; i < columnValues.size(); i++) { columnBitMaps.add(new BitMap(ARRAY_SIZE)); } bitMaps.set(columnIndex, columnBitMaps); @@ -670,6 +695,69 @@ protected Object cloneValue(TSDataType type, Object value) { } } + /* + * There are two clone modes: + * 1. Full clone: columnsToClone is null, meaning no column filter is applied. All columns are + * deep-cloned. + * 2. Partial clone: columnsToClone is non-null. Columns in columnsToClone are deep-cloned for the + * query that keeps using the source TVList; columns not in columnsToClone are not copied here. + * They are moved from the source TVList to cloneList later, and cloneList becomes the new + * working list in the memtable. + * + * This method only performs the allocation phase: clone requested value/bitmap arrays and prepare + * bitmap containers that will be needed by moved columns. It must not clear or move columns from + * the source TVList here. The destructive move is committed by moveUnclonedColumnsTo() only after + * cloneList is fully prepared for publication. + */ + private void cloneColumnDataTo(AlignedTVList cloneList, Set columnsToClone) { + boolean cloneAllColumns = columnsToClone == null; + System.arraycopy( + memoryBinaryChunkSize, 0, cloneList.memoryBinaryChunkSize, 0, dataTypes.size()); + boolean hasBitMapsToMove = false; + for (int i = 0; i < values.size(); i++) { + // Clone value + List columnValues = values.get(i); + if (columnValues == null) { + throw new IllegalStateException( + String.format("Missing value arrays for aligned column index %d during clone", i)); + } + boolean shouldCloneColumn = cloneAllColumns || columnsToClone.contains(i); + if (!shouldCloneColumn) { + hasBitMapsToMove |= bitMaps != null && bitMaps.get(i) != null; + continue; + } + + for (Object valueArray : columnValues) { + cloneList.values.get(i).add(cloneValue(dataTypes.get(i), valueArray)); + } + // Clone bitmap in columnIndex + if (bitMaps != null && bitMaps.get(i) != null) { + List columnBitMaps = bitMaps.get(i); + if (cloneList.bitMaps == null) { + cloneList.bitMaps = new ArrayList<>(dataTypes.size()); + for (int j = 0; j < dataTypes.size(); j++) { + cloneList.bitMaps.add(null); + } + } + if (cloneList.bitMaps.get(i) == null) { + List cloneColumnBitMaps = new ArrayList<>(); + for (BitMap bitMap : columnBitMaps) { + cloneColumnBitMaps.add(bitMap == null ? null : bitMap.clone()); + } + cloneList.bitMaps.set(i, cloneColumnBitMaps); + } + } + } + cloneList.materializedBitmapMemoryCost = materializedBitmapMemoryCost; + + if (hasBitMapsToMove && cloneList.bitMaps == null) { + cloneList.bitMaps = new ArrayList<>(dataTypes.size()); + for (int i = 0; i < dataTypes.size(); i++) { + cloneList.bitMaps.add(null); + } + } + } + @Override protected void clearValue() { for (int i = 0; i < dataTypes.size(); i++) { @@ -703,7 +791,12 @@ protected void expandValues() { indices.add((int[]) getPrimitiveArraysByType(TSDataType.INT32)); } for (int i = 0; i < dataTypes.size(); i++) { - values.get(i).add(getPrimitiveArraysByType(dataTypes.get(i))); + List columnValues = values.get(i); + if (columnValues == null) { + throw new IllegalStateException( + String.format("Missing value arrays for aligned column index %d during expand", i)); + } + columnValues.add(getPrimitiveArraysByType(dataTypes.get(i))); if (bitMaps != null && bitMaps.get(i) != null) { bitMaps.get(i).add(null); materializedBitmapMemoryCost += bitmapReferenceRamCost(); @@ -829,6 +922,10 @@ private void arrayCopy(Object[] value, int idx, int arrayIndex, int elementIndex continue; } List columnValues = values.get(i); + if (columnValues == null) { + throw new IllegalStateException( + String.format("Missing value arrays for aligned column index %d during arrayCopy", i)); + } switch (dataTypes.get(i)) { case TEXT: case BLOB: @@ -871,6 +968,14 @@ private void arrayCopy(Object[] value, int idx, int arrayIndex, int elementIndex } private BitMap getBitMap(int columnIndex, int arrayIndex) { + List columnValues = values.get(columnIndex); + if (columnValues == null) { + throw new IllegalStateException( + String.format( + "Missing value arrays for aligned column index %d during mark null value", + columnIndex)); + } + // init BitMaps if doesn't have if (bitMaps == null) { List> localBitMaps = new ArrayList<>(dataTypes.size()); @@ -883,7 +988,7 @@ private BitMap getBitMap(int columnIndex, int arrayIndex) { // if the bitmap in columnIndex is null, init the bitmap of this column from the beginning if (bitMaps.get(columnIndex) == null) { List columnBitMaps = new ArrayList<>(); - for (int i = 0; i < values.get(columnIndex).size(); i++) { + for (int i = 0; i < columnValues.size(); i++) { columnBitMaps.add(null); } bitMaps.set(columnIndex, columnBitMaps); @@ -919,6 +1024,15 @@ public synchronized RamInfo calculateRamSize() { new ArrayList<>(dataTypes)); } + public synchronized RamInfo calculateRamSize(Set columnsToClone) { + return new RamInfo( + timestamps.size(), + alignedTvListArrayMemCost(columnsToClone), + getRamSize(), + rowCount, + new ArrayList<>(dataTypes)); + } + public synchronized long getRamSize() { return (long) timestamps.size() * (arrayMemCostWithoutIndex @@ -986,12 +1100,17 @@ public static long alignedTvListArrayMemCost(TSDataType[] types) { * * @return AlignedTvListArrayMemSize */ - public long alignedTvListArrayMemCost() { + public long alignedTvListArrayMemCost(Set columnsToClone) { long size = 0; + int retainedColumnNum = 0; // value array mem size for (int column = 0; column < dataTypes.size(); column++) { + if (columnsToClone != null && !columnsToClone.contains(column)) { + continue; + } TSDataType type = dataTypes.get(column); - if (type != null) { + if (type != null && values.get(column) != null) { + retainedColumnNum++; size += (long) PrimitiveArrayManager.ARRAY_SIZE * (long) type.getDataTypeSize(); } } @@ -1004,12 +1123,16 @@ public long alignedTvListArrayMemCost() { // index array mem size size += (indices != null) ? PrimitiveArrayManager.ARRAY_SIZE * 4L : 0; // array headers mem size - size += (long) NUM_BYTES_ARRAY_HEADER * (2 + dataTypes.size()); + size += (long) NUM_BYTES_ARRAY_HEADER * (2 + retainedColumnNum); // Object references size in ArrayList - size += (long) NUM_BYTES_OBJECT_REF * (2 + dataTypes.size()); + size += (long) NUM_BYTES_OBJECT_REF * (2 + retainedColumnNum); return size; } + public long alignedTvListArrayMemCost() { + return alignedTvListArrayMemCost((Set) null); + } + /** * Get the single column array mem cost by give type. * diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java index cfc7f887dcfd3..60ce07d38cbfc 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory; import org.apache.iotdb.commons.exception.IllegalPathException; import org.apache.iotdb.commons.exception.MetadataException; +import org.apache.iotdb.commons.path.AlignedPath; import org.apache.iotdb.commons.path.MeasurementPath; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.db.conf.IoTDBDescriptor; @@ -49,6 +50,7 @@ import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.read.reader.IPointReader; +import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; import org.junit.Test; import org.mockito.Mockito; @@ -57,8 +59,11 @@ import java.io.IOException; import java.io.PrintStream; 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.Executor; import java.util.concurrent.ExecutorService; @@ -238,6 +243,72 @@ public void testTVListCloneForQuery() { } } + @Test + public void testAlignedTVListPartialColumnClone() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(1); + ExecutorService instanceNotificationExecutor = + IoTDBThreadPoolFactory.newFixedThreadPool(2, "test-aligned-partial-clone"); + + try { + // Create MemTable with AlignedPath + List schemaList = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + schemaList.add(new MeasurementSchema("sensor_" + i, TSDataType.INT64)); + } + String deviceId = "d1"; + IMemTable memTable = createMemTable(deviceId, schemaList); + + // Verify we have unsorted AlignedTVList + assertEquals(1, memTable.getMemTableMap().size()); + IWritableMemChunkGroup memChunkGroup = memTable.getMemTableMap().values().iterator().next(); + assertEquals(1, memChunkGroup.getMemChunkMap().size()); + IWritableMemChunk memChunk = memChunkGroup.getMemChunkMap().values().iterator().next(); + TVList tvList = memChunk.getWorkingTVList(); + assertFalse(tvList.isSorted()); + assertEquals(6424, tvList.calculateRamSize().getRamSize()); + assertEquals(100, tvList.rowCount()); + + // FragmentInstance Context + FragmentInstanceId id1 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 1), "1"); + FragmentInstanceStateMachine stateMachine1 = + new FragmentInstanceStateMachine(id1, instanceNotificationExecutor); + FragmentInstanceContext context1 = createFragmentInstanceContext(id1, stateMachine1); + + FragmentInstanceId id2 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 2), "2"); + FragmentInstanceStateMachine stateMachine2 = + new FragmentInstanceStateMachine(id2, instanceNotificationExecutor); + FragmentInstanceContext context2 = createFragmentInstanceContext(id2, stateMachine2); + + // Query 1: sensor_2 and sensor_0 + List measurements1 = Arrays.asList("sensor_2", "sensor_0"); + List schemas1 = Arrays.asList(schemaList.get(2), schemaList.get(0)); + AlignedPath fullPath1 = new AlignedPath(deviceId, measurements1, schemas1); + + ReadOnlyMemChunk readOnlyMemChunk1 = + memTable.query(context1, fullPath1, Long.MIN_VALUE, null, null); + Set accessedColumnsForQuery1 = context1.getAccessedAlignedColumns(tvList); + assertEquals(new HashSet<>(Arrays.asList(0, 2)), accessedColumnsForQuery1); + + // Query 2: sensor_1 and sensor_3 + List measurements2 = Arrays.asList("sensor_1", "sensor_3"); + List schemas2 = Arrays.asList(schemaList.get(1), schemaList.get(3)); + AlignedPath fullPath2 = new AlignedPath(deviceId, measurements2, schemas2); + ReadOnlyMemChunk readOnlyMemChunk2 = + memTable.query(context2, fullPath2, Long.MIN_VALUE, null, null); + + // Only cloned sensor_2 and sensor_0 exist + assertEquals(3232, tvList.calculateRamSize().getRamSize()); + assertEquals( + 1104, ((AlignedTVList) tvList).calculateRamSize(Collections.emptySet()).getRamSize()); + assertEquals(100, tvList.rowCount()); + + } catch (Exception e) { + fail(e.getMessage()); + } finally { + instanceNotificationExecutor.shutdown(); + } + } + private FragmentInstanceExecution createFragmentInstanceExecution(int id, Executor executor) throws CpuNotEnoughException { IDriverScheduler scheduler = Mockito.mock(IDriverScheduler.class); @@ -298,4 +369,24 @@ private IMemTable createMemTable(String deviceId, String measurementId) } return memTable; } + + private IMemTable createMemTable(String deviceId, List schemaList) + throws IllegalPathException { + PrimitiveMemTable memTable = new PrimitiveMemTable("root.test", "1"); + + // Insert data in reverse order to make it unsorted + int rows = 100; + for (int i = rows - 1; i >= 0; i--) { + Object[] values = new Object[5]; + for (int j = 0; j < 5; j++) { + values[j] = (long) i * 100 + j; + } + memTable.writeAlignedRow( + DeviceIDFactory.getInstance().getDeviceID(new PartialPath(deviceId)), + schemaList, + i, + values); + } + return memTable; + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java index e27dc009f637e..0698b73f44355 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java @@ -29,7 +29,9 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.Set; import static org.apache.iotdb.db.storageengine.rescon.memory.PrimitiveArrayManager.ARRAY_SIZE; import static org.apache.tsfile.utils.RamUsageEstimator.NUM_BYTES_ARRAY_HEADER; @@ -334,4 +336,33 @@ public void testCalculateChunkSize() { Assert.assertEquals(tvList.memoryBinaryChunkSize[1], 0); Assert.assertEquals(tvList.memoryBinaryChunkSize[2], 0); } + + @Test + public void testMovesUnclonedColumns() { + List dataTypes = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + dataTypes.add(TSDataType.INT64); + } + AlignedTVList tvList = AlignedTVList.newAlignedList(dataTypes); + tvList.putAlignedValue(0, new Object[] {1L, 2L, null}); + + Set columnsToClone = Collections.singleton(1); + AlignedTVList clonedTvList = tvList.clone(columnsToClone); + + Assert.assertNotNull(tvList.getValues().get(0)); + Assert.assertNotNull(tvList.getValues().get(2)); + Assert.assertEquals(1L, tvList.getLongByValueIndex(0, 0)); + Assert.assertTrue(tvList.isNullValue(0, 2)); + Assert.assertEquals(2L, clonedTvList.getLongByValueIndex(0, 1)); + + tvList.moveUnclonedColumnsTo(clonedTvList, columnsToClone); + + Assert.assertNull(tvList.getValues().get(0)); + Assert.assertNull(tvList.getValues().get(2)); + Assert.assertTrue(tvList.isNullValue(0, 0)); + Assert.assertTrue(tvList.isNullValue(0, 2)); + Assert.assertEquals(1L, clonedTvList.getLongByValueIndex(0, 0)); + Assert.assertEquals(2L, clonedTvList.getLongByValueIndex(0, 1)); + Assert.assertTrue(clonedTvList.isNullValue(0, 2)); + } } From 3c6f499ffd52091b996d0c83c6cc75687eceabb3 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Tue, 4 Aug 2026 19:10:20 +0800 Subject: [PATCH 02/16] Fix race when concurrent queries replace working TVList during partial clone Re-fetch and re-verify the working TVList under the memChunk lock after acquiring its queryListLock. The clone, column move and working-list swap now happen in the same memChunk critical section, so a concurrent query can never observe a working TVList whose uncloned columns were already moved away, and the fast paths never read a detached list. --- .../utils/ResourceByPathUtils.java | 182 ++++++++++-------- 1 file changed, 104 insertions(+), 78 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index adf2198019a63..43b007dcccab9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -160,16 +160,22 @@ protected Map prepareTvListMapForQuery( // TVList intrinsic lock (via synchronized methods like calculateRamSize, clone). So no AB-BA // deadlock is possible. while (true) { - // The working TVList may be replaced by another query. Always lock the actual list being used - // and verify it is still the current working list before touching query-owned fields. - TVList list = memChunk.getWorkingTVList(); - list.lockQueryList(); + // The working TVList may be replaced by a concurrent query via clone-and-swap + // (memChunk.setWorkingTVList(clone)). A queryListLock held on a detached candidate does + // not protect the current working TVList, so after acquiring the lock, re-verify it is + // still the current working list under the memChunk lock. If it was replaced while + // waiting for candidate's queryListLock, retry with the current one. + final TVList candidate = memChunk.getWorkingTVList(); + candidate.lockQueryList(); try { - if (list != memChunk.getWorkingTVList()) { - continue; + synchronized (memChunk) { + if (memChunk.getWorkingTVList() != candidate) { + continue; + } } + if (copyTimeFilter != null - && !copyTimeFilter.satisfyStartEndTime(list.getMinTime(), list.getMaxTime())) { + && !copyTimeFilter.satisfyStartEndTime(candidate.getMinTime(), candidate.getMaxTime())) { return tvListQueryMap; } @@ -188,15 +194,16 @@ protected Map prepareTvListMapForQuery( * sorted or no active query is using it. Otherwise, Q2 should read from * workingListForFlush. */ - boolean canUseListDirectly = list.isSorted() || list.getQueryContextSet().isEmpty(); + boolean canUseListDirectly = + candidate.isSorted() || candidate.getQueryContextSet().isEmpty(); LOGGER.debug( "Flushing MemTable - add current query context to mutable TVList's query list"); if (canUseListDirectly) { - list.getQueryContextSet().add(context); - tvListQueryMap.put(list, list.rowCount()); + candidate.getQueryContextSet().add(context); + tvListQueryMap.put(candidate, candidate.rowCount()); } else { TVList workingListForFlushSort = - memChunk.initWorkingListForFlushIfNecessary(list, true); + memChunk.initWorkingListForFlushIfNecessary(candidate, true); /* * The query will read from workingListForFlushSort, but cloneForFlushSort() only clones * times and indices. The value arrays and bitmaps are still shared with the original @@ -210,76 +217,96 @@ protected Map prepareTvListMapForQuery( * Do not put the original list into tvListQueryMap here. The actual read path must use * workingListForFlushSort to avoid sorting the original list in place. */ - list.getQueryContextSet().add(context); - context.addTVListToSet(Collections.singleton(list)); + candidate.getQueryContextSet().add(context); + context.addTVListToSet(Collections.singleton(candidate)); workingListForFlushSort.getQueryContextSet().add(context); tvListQueryMap.put(workingListForFlushSort, workingListForFlushSort.rowCount()); } - } else { - if (list.isSorted() || list.getQueryContextSet().isEmpty()) { - LOGGER.debug( - "Working MemTable - add current query context to mutable TVList's query list when it's sorted or no other query on it"); - list.getQueryContextSet().add(context); - tvListQueryMap.put(list, list.rowCount()); - - // columnIndexList is to track column-level access for AlignedTVList. - // For TVList (primitive time series), it remains null and column tracking is not - // needed. - if (columnIndexList != null && context instanceof FragmentInstanceContext) { - ((FragmentInstanceContext) context).putAccessedColumns(list, columnIndexList); - } - } else { - /* - * +----------------------+ - * | MemTable | - * | | - * | +------------+ | +-----------------+ - * | | TVList |<---+--+ +---+ Previous Query | - * | +-----^------+ | | | +-----------------+ - * | | | | | - * +----------+-----------+ | | +----------------+ - * | Clone +---+---+ Current Query | - * +-----+------+ | +----------------+ - * | TVList | <---------+ - * +------------+ - */ - LOGGER.debug( - "Working MemTable - clone mutable TVList and replace old TVList in working MemTable"); - - Set columnsToClone = getAccessedColumnsForQuery(list); - listRamInfo = - (columnsToClone == null) - ? list.calculateRamSize() - : ((AlignedTVList) list).calculateRamSize(columnsToClone); - - // reserve query memory - QueryContext firstQuery = list.getQueryContextSet().iterator().next(); - if (firstQuery instanceof FragmentInstanceContext) { - MemoryReservationManager memoryReservationManager = - ((FragmentInstanceContext) firstQuery).getMemoryReservationContext(); - memoryReservationManager.reserveMemoryCumulatively(listRamInfo.getRamSize()); - list.setReservedMemoryBytes(listRamInfo.getRamSize()); - } - list.setOwnerQuery(firstQuery); - - // clone TVList - TVList cloneList = - (columnsToClone == null) - ? list.clone() - : ((AlignedTVList) list).clone(columnsToClone); - - cloneList.getQueryContextSet().add(context); - tvListQueryMap.put(cloneList, cloneList.rowCount()); - if (columnIndexList != null && context instanceof FragmentInstanceContext) { - ((FragmentInstanceContext) context).putAccessedColumns(cloneList, columnIndexList); - } + return tvListQueryMap; + } - if (columnsToClone != null) { - ((AlignedTVList) list) - .moveUnclonedColumnsTo((AlignedTVList) cloneList, columnsToClone); - } - memChunk.setWorkingTVList(cloneList); + if (candidate.isSorted() || candidate.getQueryContextSet().isEmpty()) { + LOGGER.debug( + "Working MemTable - add current query context to mutable TVList's query list when it's sorted or no other query on it"); + candidate.getQueryContextSet().add(context); + tvListQueryMap.put(candidate, candidate.rowCount()); + + // columnIndexList is to track column-level access for AlignedTVList. + // For TVList (primitive time series), it remains null and column tracking is not needed. + if (columnIndexList != null && context instanceof FragmentInstanceContext) { + ((FragmentInstanceContext) context).putAccessedColumns(candidate, columnIndexList); + } + return tvListQueryMap; + } + + /* + * +----------------------+ + * | MemTable | + * | | + * | +------------+ | +-----------------+ + * | | TVList |<---+--+ +---+ Previous Query | + * | +-----^------+ | | | +-----------------+ + * | | | | | + * +----------+-----------+ | | +----------------+ + * | Clone +---+---+ Current Query | + * +-----+------+ | +----------------+ + * | TVList | <---------+ + * +------------+ + */ + LOGGER.debug( + "Working MemTable - clone mutable TVList and replace old TVList in working MemTable"); + + synchronized (memChunk) { + // Re-check defensively before cloning and publishing the replacement. The clone and the + // working-list swap must be done in the same memChunk critical section, so a concurrent + // query can never observe a working TVList whose columns have already been moved away. + if (memChunk.getWorkingTVList() != candidate) { + continue; } + + // calculateRamSize (synchronized method on TVList) was previously called before + // lockQueryList to avoid deadlock concerns. For partial clone of AlignedTVList, however + // calculateRamSize must now be called inside the lockQueryList section because it depends + // on accessing columns on the AlignedTVList. + // This is safe because the lock ordering - queryListLock must always be acquired before + // the TVList intrinsic lock (via synchronized methods like calculateRamSize, clone). So + // no AB-BA deadlock is possible. + Set columnsToClone = getAccessedColumnsForQuery(candidate); + listRamInfo = + (columnsToClone == null) + ? candidate.calculateRamSize() + : ((AlignedTVList) candidate).calculateRamSize(columnsToClone); + + // reserve query memory + QueryContext firstQuery = candidate.getQueryContextSet().iterator().next(); + if (firstQuery instanceof FragmentInstanceContext) { + MemoryReservationManager memoryReservationManager = + ((FragmentInstanceContext) firstQuery).getMemoryReservationContext(); + memoryReservationManager.reserveMemoryCumulatively(listRamInfo.getRamSize()); + candidate.setReservedMemoryBytes(listRamInfo.getRamSize()); + } + candidate.setOwnerQuery(firstQuery); + + // clone TVList + TVList cloneList = + (columnsToClone == null) + ? candidate.clone() + : ((AlignedTVList) candidate).clone(columnsToClone); + + cloneList.getQueryContextSet().add(context); + tvListQueryMap.put(cloneList, cloneList.rowCount()); + if (columnIndexList != null && context instanceof FragmentInstanceContext) { + ((FragmentInstanceContext) context).putAccessedColumns(cloneList, columnIndexList); + } + + // Move the uncloned columns and publish the clone in the same memChunk critical section, + // so a concurrent query never observes a working TVList whose columns were already moved. + if (columnsToClone != null) { + ((AlignedTVList) candidate) + .moveUnclonedColumnsTo((AlignedTVList) cloneList, columnsToClone); + } + memChunk.setWorkingTVList(cloneList); + return tvListQueryMap; } } catch (MemoryNotEnoughException ex) { if (listRamInfo != null) { @@ -293,9 +320,8 @@ protected Map prepareTvListMapForQuery( } throw ex; } finally { - list.unlockQueryList(); + candidate.unlockQueryList(); } - return tvListQueryMap; } } From ddb6043f83931810147b588dd2e6d0cef6e73c10 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Tue, 4 Aug 2026 19:21:34 +0800 Subject: [PATCH 03/16] Fix NPE in AlignedTVList constructor caused by partial clone mem cost calc refreshArrayMemCostWithoutIndex() was invoked before the values field was initialized, but alignedTvListArrayMemCost(Set) now dereferences values.get(column) to skip moved columns, so construction of any aligned TVList threw NPE. Initialize values before computing the array mem cost. --- .../apache/iotdb/db/utils/datastructure/AlignedTVList.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index bdf0633d6a430..1ef24ef2131f0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -97,12 +97,13 @@ public abstract class AlignedTVList extends TVList { super(); dataTypes = types; memoryBinaryChunkSize = new long[dataTypes.size()]; - refreshArrayMemCostWithoutIndex(); - values = new ArrayList<>(types.size()); for (int i = 0; i < types.size(); i++) { values.add(new ArrayList<>()); } + // arrayMemCostWithoutIndex depends on per-column value arrays, so values must be + // initialized before computing it + refreshArrayMemCostWithoutIndex(); } public static AlignedTVList newAlignedList(List dataTypes) { From 667178e113e2cadeeb4ab7141632f686d71f6db6 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Tue, 4 Aug 2026 19:26:12 +0800 Subject: [PATCH 04/16] spotless --- .../schemaengine/schemaregion/utils/ResourceByPathUtils.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index 43b007dcccab9..36fb7cb1fe5c2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -175,7 +175,8 @@ protected Map prepareTvListMapForQuery( } if (copyTimeFilter != null - && !copyTimeFilter.satisfyStartEndTime(candidate.getMinTime(), candidate.getMaxTime())) { + && !copyTimeFilter.satisfyStartEndTime( + candidate.getMinTime(), candidate.getMaxTime())) { return tvListQueryMap; } From ec4ae27a5caa4bea9ec68223eb628d12afb93948 Mon Sep 17 00:00:00 2001 From: shizy Date: Tue, 4 Aug 2026 19:27:16 +0800 Subject: [PATCH 05/16] remove one test temporarily --- .../FragmentInstanceExecutionTest.java | 66 ------------------- 1 file changed, 66 deletions(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java index 60ce07d38cbfc..149c03e96fbd1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java @@ -243,72 +243,6 @@ public void testTVListCloneForQuery() { } } - @Test - public void testAlignedTVListPartialColumnClone() { - IoTDBDescriptor.getInstance().getConfig().setDataNodeId(1); - ExecutorService instanceNotificationExecutor = - IoTDBThreadPoolFactory.newFixedThreadPool(2, "test-aligned-partial-clone"); - - try { - // Create MemTable with AlignedPath - List schemaList = new ArrayList<>(); - for (int i = 0; i < 5; i++) { - schemaList.add(new MeasurementSchema("sensor_" + i, TSDataType.INT64)); - } - String deviceId = "d1"; - IMemTable memTable = createMemTable(deviceId, schemaList); - - // Verify we have unsorted AlignedTVList - assertEquals(1, memTable.getMemTableMap().size()); - IWritableMemChunkGroup memChunkGroup = memTable.getMemTableMap().values().iterator().next(); - assertEquals(1, memChunkGroup.getMemChunkMap().size()); - IWritableMemChunk memChunk = memChunkGroup.getMemChunkMap().values().iterator().next(); - TVList tvList = memChunk.getWorkingTVList(); - assertFalse(tvList.isSorted()); - assertEquals(6424, tvList.calculateRamSize().getRamSize()); - assertEquals(100, tvList.rowCount()); - - // FragmentInstance Context - FragmentInstanceId id1 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 1), "1"); - FragmentInstanceStateMachine stateMachine1 = - new FragmentInstanceStateMachine(id1, instanceNotificationExecutor); - FragmentInstanceContext context1 = createFragmentInstanceContext(id1, stateMachine1); - - FragmentInstanceId id2 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 2), "2"); - FragmentInstanceStateMachine stateMachine2 = - new FragmentInstanceStateMachine(id2, instanceNotificationExecutor); - FragmentInstanceContext context2 = createFragmentInstanceContext(id2, stateMachine2); - - // Query 1: sensor_2 and sensor_0 - List measurements1 = Arrays.asList("sensor_2", "sensor_0"); - List schemas1 = Arrays.asList(schemaList.get(2), schemaList.get(0)); - AlignedPath fullPath1 = new AlignedPath(deviceId, measurements1, schemas1); - - ReadOnlyMemChunk readOnlyMemChunk1 = - memTable.query(context1, fullPath1, Long.MIN_VALUE, null, null); - Set accessedColumnsForQuery1 = context1.getAccessedAlignedColumns(tvList); - assertEquals(new HashSet<>(Arrays.asList(0, 2)), accessedColumnsForQuery1); - - // Query 2: sensor_1 and sensor_3 - List measurements2 = Arrays.asList("sensor_1", "sensor_3"); - List schemas2 = Arrays.asList(schemaList.get(1), schemaList.get(3)); - AlignedPath fullPath2 = new AlignedPath(deviceId, measurements2, schemas2); - ReadOnlyMemChunk readOnlyMemChunk2 = - memTable.query(context2, fullPath2, Long.MIN_VALUE, null, null); - - // Only cloned sensor_2 and sensor_0 exist - assertEquals(3232, tvList.calculateRamSize().getRamSize()); - assertEquals( - 1104, ((AlignedTVList) tvList).calculateRamSize(Collections.emptySet()).getRamSize()); - assertEquals(100, tvList.rowCount()); - - } catch (Exception e) { - fail(e.getMessage()); - } finally { - instanceNotificationExecutor.shutdown(); - } - } - private FragmentInstanceExecution createFragmentInstanceExecution(int id, Executor executor) throws CpuNotEnoughException { IDriverScheduler scheduler = Mockito.mock(IDriverScheduler.class); From 8da6498cdb620af26391a2ef279f6e86ae13fbdc Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Tue, 4 Aug 2026 19:30:13 +0800 Subject: [PATCH 06/16] spotless --- .../execution/fragment/FragmentInstanceExecutionTest.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java index 149c03e96fbd1..0f1b1c7d25318 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java @@ -22,7 +22,6 @@ import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory; import org.apache.iotdb.commons.exception.IllegalPathException; import org.apache.iotdb.commons.exception.MetadataException; -import org.apache.iotdb.commons.path.AlignedPath; import org.apache.iotdb.commons.path.MeasurementPath; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.db.conf.IoTDBDescriptor; @@ -59,11 +58,8 @@ import java.io.IOException; import java.io.PrintStream; 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.Executor; import java.util.concurrent.ExecutorService; From e421731714ad249c2d61dafb7a19c7ff97930fd0 Mon Sep 17 00:00:00 2001 From: shizy Date: Tue, 4 Aug 2026 19:55:34 +0800 Subject: [PATCH 07/16] release non-query columns during release tvlist --- .../memtable/AbstractWritableMemChunk.java | 42 +++++++++++++++++ .../db/utils/datastructure/AlignedTVList.java | 46 +++++++++++++++++++ .../datastructure/AlignedTVListTest.java | 37 +++++++++++++++ 3 files changed, 125 insertions(+) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java index 13fa84d039e27..6bb4572f18032 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java @@ -24,6 +24,7 @@ import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext; import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView; +import org.apache.iotdb.db.utils.datastructure.AlignedTVList; import org.apache.iotdb.db.utils.datastructure.BatchEncodeInfo; import org.apache.iotdb.db.utils.datastructure.TVList; @@ -35,8 +36,10 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Set; import java.util.concurrent.BlockingQueue; public abstract class AbstractWritableMemChunk implements IWritableMemChunk { @@ -100,6 +103,11 @@ protected void maybeReleaseTvList(TVList tvList) { } } + /** + * Try to release the TVList. If there are active queries, transfer memory ownership to the first + * query. For AlignedTVList, this will release non-query columns before transferring to reduce + * memory footprint. + */ private void tryReleaseTvList(TVList tvList) { tvList.lockQueryList(); try { @@ -107,6 +115,21 @@ private void tryReleaseTvList(TVList tvList) { tvList.clear(); } else { QueryContext firstQuery = tvList.getQueryContextSet().iterator().next(); + + // For AlignedTVList with active queries, release non-query columns before + // transferring memory ownership to reduce memory footprint. + if (tvList instanceof AlignedTVList) { + AlignedTVList alignedTVList = (AlignedTVList) tvList; + + // Get the union of all columns accessed by queries + Set accessedColumns = getAccessedColumnsForQuery(alignedTVList); + + if (accessedColumns != null && !accessedColumns.isEmpty()) { + // Release non-query columns to reduce memory before ownership transfer + alignedTVList.releaseNonQueryColumns(accessedColumns); + } + } + // transfer memory from write process to read process. Here it reserves read memory and // releaseFlushedMemTable will release write memory. if (firstQuery instanceof FragmentInstanceContext) { @@ -124,6 +147,25 @@ private void tryReleaseTvList(TVList tvList) { } } + /** + * Get the union of all columns accessed by active queries on this TVList. This method must be + * called with tvList.lockQueryList() held. + */ + private Set getAccessedColumnsForQuery(AlignedTVList alignedTVList) { + Set accessedColumns = new HashSet<>(); + for (QueryContext queryContext : alignedTVList.getQueryContextSet()) { + if (!(queryContext instanceof FragmentInstanceContext)) { + return null; + } + FragmentInstanceContext ctx = (FragmentInstanceContext) queryContext; + Set columns = ctx.getAccessedAlignedColumns(alignedTVList); + if (columns != null && !columns.isEmpty()) { + accessedColumns.addAll(columns); + } + } + return accessedColumns; + } + @Override public abstract void putLong(long t, long v); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 1ef24ef2131f0..2918f088a9245 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -206,6 +206,52 @@ public synchronized void moveUnclonedColumnsTo( cloneList.refreshArrayMemCostWithoutIndex(); } + /** + * Release memory for non-query columns in this TVList. This is used during memory ownership + * transfer from write process to read process to reduce memory footprint. Only columns that are + * accessed by active queries are retained; all other columns are released. + * + * @param columnsToKeep set of column indices that are accessed by queries and should be kept + */ + public synchronized void releaseNonQueryColumns(Set columnsToKeep) { + if (columnsToKeep == null || columnsToKeep.isEmpty()) { + return; + } + + for (int i = 0; i < values.size(); i++) { + // Skip columns that should be kept or are already null + if (columnsToKeep.contains(i)) { + continue; + } + + List columnValues = values.get(i); + if (columnValues == null) { + continue; + } + + // Release memory for non-query columns + for (Object dataArray : columnValues) { + PrimitiveArrayManager.release(dataArray); + } + columnValues.clear(); + memoryBinaryChunkSize[i] = 0; + + // Release bitmap memory for non-query columns + if (bitMaps != null && bitMaps.get(i) != null) { + for (BitMap bitMap : bitMaps.get(i)) { + if (bitMap != null) { + materializedBitmapMemoryCost -= bitmapRamCost(); + } + } + bitMaps.get(i).clear(); + materializedBitmapMemoryCost -= (long) bitMaps.get(i).size() * bitmapReferenceRamCost(); + } + } + + // Refresh per-block memory cost after releasing columns + refreshArrayMemCostWithoutIndex(); + } + @SuppressWarnings("squid:S3776") // Suppress high Cognitive Complexity warning @Override public synchronized void putAlignedValue(long timestamp, Object[] value) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java index 0698b73f44355..8ade90a773dca 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java @@ -30,6 +30,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Set; @@ -365,4 +366,40 @@ public void testMovesUnclonedColumns() { Assert.assertEquals(2L, clonedTvList.getLongByValueIndex(0, 1)); Assert.assertTrue(clonedTvList.isNullValue(0, 2)); } + + @Test + public void testReleaseNonQueryColumnsWithBitmaps() { + List dataTypes = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + dataTypes.add(TSDataType.INT64); + } + AlignedTVList tvList = AlignedTVList.newAlignedList(dataTypes); + for (int i = 0; i < 100; i++) { + Object[] values = new Object[3]; + values[0] = (long) i; + values[1] = null; // This will create a bitmap + values[2] = (long) (i * 100); + tvList.putAlignedValue(i, values); + } + + // Verify bitmaps were created for column 1 + Assert.assertNotNull(tvList.getBitMaps()); + Assert.assertNotNull(tvList.getBitMaps().get(1)); + + // Keep only column 0 and 2, release column 1 + Set columnsToKeep = new HashSet<>(Arrays.asList(0, 2)); + tvList.releaseNonQueryColumns(columnsToKeep); + + // Verify column 1 is released + Assert.assertTrue(tvList.getValues().get(1).isEmpty()); + Assert.assertTrue(tvList.getBitMaps().get(1).isEmpty()); + + // Verify columns 0 and 2 are intact + Assert.assertFalse(tvList.getValues().get(0).isEmpty()); + Assert.assertFalse(tvList.getValues().get(2).isEmpty()); + for (int i = 0; i < 100; i++) { + Assert.assertEquals((long) i, tvList.getLongByValueIndex(i, 0)); + Assert.assertEquals((long) (i * 100), tvList.getLongByValueIndex(i, 2)); + } + } } From bc6ca6d36a5bb17affabb4a086396a0f5ac6991d Mon Sep 17 00:00:00 2001 From: shizy Date: Tue, 4 Aug 2026 20:28:08 +0800 Subject: [PATCH 08/16] fix getAccessedColumnsForQuery --- .../utils/ResourceByPathUtils.java | 27 +--------------- .../memtable/AbstractWritableMemChunk.java | 22 +------------ .../db/utils/datastructure/AlignedTVList.java | 32 +++++++++++++++---- .../iotdb/db/utils/datastructure/TVList.java | 10 ++++++ 4 files changed, 37 insertions(+), 54 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index 36fb7cb1fe5c2..65eacf0e05847 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -64,7 +64,6 @@ import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; -import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -272,7 +271,7 @@ protected Map prepareTvListMapForQuery( // This is safe because the lock ordering - queryListLock must always be acquired before // the TVList intrinsic lock (via synchronized methods like calculateRamSize, clone). So // no AB-BA deadlock is possible. - Set columnsToClone = getAccessedColumnsForQuery(candidate); + Set columnsToClone = candidate.getAccessedColumnsForQuery(); listRamInfo = (columnsToClone == null) ? candidate.calculateRamSize() @@ -325,10 +324,6 @@ protected Map prepareTvListMapForQuery( } } } - - protected Set getAccessedColumnsForQuery(TVList tvList) { - return null; - } } class AlignedResourceByPathUtils extends ResourceByPathUtils { @@ -494,26 +489,6 @@ public ReadOnlyMemChunk getReadOnlyMemChunkFromMemTable( context, columnIndexList, getMeasurementSchema(), alignedTvListQueryMap, deletionList); } - /** - * This method is called from prepareTvListMapForQuery with tvList.lockQueryList() held, ensuring - * thread-safe access to queryContextSet. - * - * @param tvList the TVList to get accessed columns for - * @return set of accessed column indices, or empty set if no columns are tracked - */ - @Override - protected Set getAccessedColumnsForQuery(TVList tvList) { - Set accessedColumns = new HashSet<>(); - for (QueryContext queryContext : tvList.getQueryContextSet()) { - if (!(queryContext instanceof FragmentInstanceContext)) { - return null; - } - accessedColumns.addAll( - ((FragmentInstanceContext) queryContext).getAccessedAlignedColumns(tvList)); - } - return accessedColumns; - } - public VectorMeasurementSchema getMeasurementSchema() { List measurementList = partialPath.getMeasurementList(); TSDataType[] types = new TSDataType[measurementList.size()]; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java index 6bb4572f18032..0f5edd22adbaa 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java @@ -36,7 +36,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Set; @@ -122,7 +121,7 @@ private void tryReleaseTvList(TVList tvList) { AlignedTVList alignedTVList = (AlignedTVList) tvList; // Get the union of all columns accessed by queries - Set accessedColumns = getAccessedColumnsForQuery(alignedTVList); + Set accessedColumns = alignedTVList.getAccessedColumnsForQuery(); if (accessedColumns != null && !accessedColumns.isEmpty()) { // Release non-query columns to reduce memory before ownership transfer @@ -147,25 +146,6 @@ private void tryReleaseTvList(TVList tvList) { } } - /** - * Get the union of all columns accessed by active queries on this TVList. This method must be - * called with tvList.lockQueryList() held. - */ - private Set getAccessedColumnsForQuery(AlignedTVList alignedTVList) { - Set accessedColumns = new HashSet<>(); - for (QueryContext queryContext : alignedTVList.getQueryContextSet()) { - if (!(queryContext instanceof FragmentInstanceContext)) { - return null; - } - FragmentInstanceContext ctx = (FragmentInstanceContext) queryContext; - Set columns = ctx.getAccessedAlignedColumns(alignedTVList); - if (columns != null && !columns.isEmpty()) { - accessedColumns.addAll(columns); - } - } - return accessedColumns; - } - @Override public abstract void putLong(long t, long v); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 2918f088a9245..549870343e900 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -19,6 +19,8 @@ package org.apache.iotdb.db.utils.datastructure; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext; import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALWriteUtils; @@ -51,6 +53,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashSet; import java.util.List; import java.util.Objects; import java.util.Set; @@ -201,6 +204,7 @@ public synchronized void moveUnclonedColumnsTo( } memoryBinaryChunkSize[i] = 0; } + materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps); // Column ownership changed on both lists, so refresh their per-block memory cost. refreshArrayMemCostWithoutIndex(); cloneList.refreshArrayMemCostWithoutIndex(); @@ -238,16 +242,11 @@ public synchronized void releaseNonQueryColumns(Set columnsToKeep) { // Release bitmap memory for non-query columns if (bitMaps != null && bitMaps.get(i) != null) { - for (BitMap bitMap : bitMaps.get(i)) { - if (bitMap != null) { - materializedBitmapMemoryCost -= bitmapRamCost(); - } - } - bitMaps.get(i).clear(); - materializedBitmapMemoryCost -= (long) bitMaps.get(i).size() * bitmapReferenceRamCost(); + bitMaps.set(i, null); } } + materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps); // Refresh per-block memory cost after releasing columns refreshArrayMemCostWithoutIndex(); } @@ -611,6 +610,25 @@ public List getTsDataTypes() { return dataTypes; } + /** + * Get the union of all columns accessed by queries on this AlignedTVList. This method should be + * called with queryListLock held for thread safety. + * + * @return set of accessed column indices, or empty set if no columns are tracked or no queries + * are present + */ + @Override + public Set getAccessedColumnsForQuery() { + Set accessedColumns = new HashSet<>(); + for (QueryContext queryContext : getQueryContextSet()) { + if (queryContext instanceof FragmentInstanceContext) { + accessedColumns.addAll( + ((FragmentInstanceContext) queryContext).getAccessedAlignedColumns(this)); + } + } + return accessedColumns; + } + @Override /* * Must be synchronized with sort() on the same TVList instance: a query may sort diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java index a085c29e1309b..60e421c9af5be 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java @@ -809,6 +809,16 @@ public Set getQueryContextSet() { return queryContextSet; } + /** + * Get the union of all columns accessed by queries on this TVList. For non-AlignedTVList, returns + * empty set. This method should be called with queryListLock held for thread safety. + * + * @return set of accessed column indices, or empty set if no columns are tracked + */ + public Set getAccessedColumnsForQuery() { + return null; + } + public List getBitMap() { return bitMap; } From 6ab294fca9557510c1fc595c3f8e7a05bfa60613 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Tue, 4 Aug 2026 20:39:55 +0800 Subject: [PATCH 09/16] tmp fix --- .../db/utils/datastructure/AlignedTVList.java | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 2918f088a9245..e4d1a77249d78 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -145,7 +145,7 @@ public TVList getTvListByColumnIndex(List columnIndex, List alignedTvList.bitMaps = bitMaps; alignedTvList.rowCount = this.rowCount; alignedTvList.allValueColDeletedMap = getAllValueColDeletedMap(); - alignedTvList.materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps); + alignedTvList.materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps, null); return alignedTvList; } @@ -1075,18 +1075,23 @@ public synchronized RamInfo calculateRamSize(Set columnsToClone) { return new RamInfo( timestamps.size(), alignedTvListArrayMemCost(columnsToClone), - getRamSize(), + getRamSize(columnsToClone), rowCount, new ArrayList<>(dataTypes)); } public synchronized long getRamSize() { - return (long) timestamps.size() + return timestamps.size() * (arrayMemCostWithoutIndex + (indices != null ? (long) PrimitiveArrayManager.ARRAY_SIZE * Integer.BYTES : 0)) + materializedBitmapMemoryCost; } + public synchronized long getRamSize(Set columnsToClone) { + return timestamps.size() * alignedTvListArrayMemCost(columnsToClone) + + calculateBitmapRamCost(bitMaps, columnsToClone); + } + private void refreshArrayMemCostWithoutIndex() { arrayMemCostWithoutIndex = alignedTvListArrayMemCost(); if (indices != null) { @@ -1094,16 +1099,21 @@ private void refreshArrayMemCostWithoutIndex() { } } - private static long calculateBitmapRamCost(List> bitMaps) { + private static long calculateBitmapRamCost( + List> bitMaps, Set columnsToClone) { if (bitMaps == null) { return 0; } long size = 0; - for (List columnBitMaps : bitMaps) { + for (int i = 0, length = bitMaps.size(); i < length; i++) { + if (columnsToClone != null && !columnsToClone.contains(i)) { + continue; + } + List columnBitMaps = bitMaps.get(i); if (columnBitMaps == null) { continue; } - size += (long) columnBitMaps.size() * bitmapReferenceRamCost(); + size += columnBitMaps.size() * bitmapReferenceRamCost(); for (BitMap bitMap : columnBitMaps) { if (bitMap != null) { size += bitmapRamCost(); From 14aec7eba055090e51269ed6e351e882fe7a185e Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Tue, 4 Aug 2026 20:46:08 +0800 Subject: [PATCH 10/16] fix all --- .../apache/iotdb/db/utils/datastructure/AlignedTVList.java | 6 +++--- .../iotdb/db/utils/datastructure/AlignedTVListTest.java | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 0e0ce380ee1cd..3063631f93da6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -204,7 +204,7 @@ public synchronized void moveUnclonedColumnsTo( } memoryBinaryChunkSize[i] = 0; } - materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps); + materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps, columnsToClone); // Column ownership changed on both lists, so refresh their per-block memory cost. refreshArrayMemCostWithoutIndex(); cloneList.refreshArrayMemCostWithoutIndex(); @@ -237,7 +237,7 @@ public synchronized void releaseNonQueryColumns(Set columnsToKeep) { for (Object dataArray : columnValues) { PrimitiveArrayManager.release(dataArray); } - columnValues.clear(); + values.set(i, null); memoryBinaryChunkSize[i] = 0; // Release bitmap memory for non-query columns @@ -246,7 +246,7 @@ public synchronized void releaseNonQueryColumns(Set columnsToKeep) { } } - materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps); + materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps, columnsToKeep); // Refresh per-block memory cost after releasing columns refreshArrayMemCostWithoutIndex(); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java index 8ade90a773dca..6027364061cb4 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java @@ -391,8 +391,8 @@ public void testReleaseNonQueryColumnsWithBitmaps() { tvList.releaseNonQueryColumns(columnsToKeep); // Verify column 1 is released - Assert.assertTrue(tvList.getValues().get(1).isEmpty()); - Assert.assertTrue(tvList.getBitMaps().get(1).isEmpty()); + Assert.assertNull(tvList.getValues().get(1)); + Assert.assertNull(tvList.getBitMaps().get(1)); // Verify columns 0 and 2 are intact Assert.assertFalse(tvList.getValues().get(0).isEmpty()); From 256cece33ddb9a7e8ee88703c446b1dfacfa5f75 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Tue, 4 Aug 2026 21:19:58 +0800 Subject: [PATCH 11/16] fix all --- .../schemaregion/utils/ResourceByPathUtils.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index 65eacf0e05847..f7c74efa5addd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -222,6 +222,12 @@ protected Map prepareTvListMapForQuery( workingListForFlushSort.getQueryContextSet().add(context); tvListQueryMap.put(workingListForFlushSort, workingListForFlushSort.rowCount()); } + + // columnIndexList is to track column-level access for AlignedTVList. + // For TVList (primitive time series), it remains null and column tracking is not needed. + if (columnIndexList != null && context instanceof FragmentInstanceContext) { + ((FragmentInstanceContext) context).putAccessedColumns(candidate, columnIndexList); + } return tvListQueryMap; } From c4f31171b69d38f42ae5c39c29226438dccec03f Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 5 Aug 2026 11:04:55 +0800 Subject: [PATCH 12/16] Fix partial TVList clone memory accounting --- .../fragment/FragmentInstanceContext.java | 5 + .../memory/FakedMemoryReservationManager.java | 3 + .../memory/MemoryReservationManager.java | 6 + ...NotThreadSafeMemoryReservationManager.java | 16 +- .../ThreadSafeMemoryReservationManager.java | 5 + .../utils/ResourceByPathUtils.java | 96 +++++--- .../db/utils/datastructure/AlignedTVList.java | 211 ++++++++++++++++-- ...alExecutionPlannerOperatorsMemoryTest.java | 46 ++++ .../datastructure/AlignedTVListTest.java | 118 +++++++++- 9 files changed, 446 insertions(+), 60 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 62c581ad91a2f..24b233efcdf64 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -245,6 +245,11 @@ public void putAccessedColumns(TVList tvList, List columnIndexList) { }); } + /** 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 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java index 8d0c9ae5997e0..1742a8070b36d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java +++ b/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() {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java index eddec15facce2..9a9036b1d97e9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java +++ b/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. diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java index e4f211ea7640f..514a2935f6ce3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java +++ b/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) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java index 0a1c6eee4181e..71676e5b77fe9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java +++ b/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(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index f7c74efa5addd..ef3ffd6acc840 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -144,6 +144,11 @@ protected Map prepareTvListMapForQuery( "Flushing/Working MemTable - add current query context to immutable TVList's query list"); tvList.getQueryContextSet().add(context); tvListQueryMap.put(tvList, tvList.rowCount()); + // columnIndexList is to track column-level access for AlignedTVList. + // For TVList (primitive time series), it remains null and column tracking is not needed. + if (columnIndexList != null && context instanceof FragmentInstanceContext) { + ((FragmentInstanceContext) context).putAccessedColumns(tvList, columnIndexList); + } } finally { tvList.unlockQueryList(); } @@ -283,36 +288,77 @@ protected Map prepareTvListMapForQuery( ? candidate.calculateRamSize() : ((AlignedTVList) candidate).calculateRamSize(columnsToClone); - // reserve query memory QueryContext firstQuery = candidate.getQueryContextSet().iterator().next(); - if (firstQuery instanceof FragmentInstanceContext) { - MemoryReservationManager memoryReservationManager = - ((FragmentInstanceContext) firstQuery).getMemoryReservationContext(); - memoryReservationManager.reserveMemoryCumulatively(listRamInfo.getRamSize()); - candidate.setReservedMemoryBytes(listRamInfo.getRamSize()); - } - candidate.setOwnerQuery(firstQuery); + TVList cloneList = null; + AlignedTVList.PartialClonePlan partialClonePlan = null; + FragmentInstanceContext cloneContext = + columnIndexList != null && context instanceof FragmentInstanceContext + ? (FragmentInstanceContext) context + : null; + MemoryReservationManager memoryReservationManager = + firstQuery instanceof FragmentInstanceContext + ? ((FragmentInstanceContext) firstQuery).getMemoryReservationContext() + : null; + boolean reservationNeedsRollback = false; + boolean replacementPublished = false; + try { + // Reserve before allocating the clone, so this transient memory increase is still + // protected by query-memory admission control. Ownership is not published yet, and a + // later preparation failure rolls this exact reservation back immediately. + if (memoryReservationManager != null) { + memoryReservationManager.reserveMemoryCumulatively(listRamInfo.getRamSize()); + reservationNeedsRollback = true; + } - // clone TVList - TVList cloneList = - (columnsToClone == null) - ? candidate.clone() - : ((AlignedTVList) candidate).clone(columnsToClone); + // Clone and validate without changing the source list. PartialClonePlan.commit is the + // only destructive step and is allocation-free. + if (columnsToClone == null) { + cloneList = candidate.clone(); + } else { + partialClonePlan = ((AlignedTVList) candidate).preparePartialClone(columnsToClone); + cloneList = partialClonePlan.getCloneList(); + } - cloneList.getQueryContextSet().add(context); - tvListQueryMap.put(cloneList, cloneList.rowCount()); - if (columnIndexList != null && context instanceof FragmentInstanceContext) { - ((FragmentInstanceContext) context).putAccessedColumns(cloneList, columnIndexList); - } + cloneList.getQueryContextSet().add(context); + tvListQueryMap.put(cloneList, cloneList.rowCount()); + if (cloneContext != null) { + cloneContext.putAccessedColumns(cloneList, columnIndexList); + } - // Move the uncloned columns and publish the clone in the same memChunk critical section, - // so a concurrent query never observes a working TVList whose columns were already moved. - if (columnsToClone != null) { - ((AlignedTVList) candidate) - .moveUnclonedColumnsTo((AlignedTVList) cloneList, columnsToClone); + if (partialClonePlan != null) { + partialClonePlan.commit(); + } + memChunk.setWorkingTVList(cloneList); + replacementPublished = true; + + // Publish query ownership only after the replacement is fully committed. The + // candidate query-list lock prevents its owner from being released concurrently. + if (memoryReservationManager != null) { + candidate.setReservedMemoryBytes(listRamInfo.getRamSize()); + } + candidate.setOwnerQuery(firstQuery); + reservationNeedsRollback = false; + return tvListQueryMap; + } catch (RuntimeException | Error failure) { + if (reservationNeedsRollback) { + try { + memoryReservationManager.releaseMemoryImmediately(listRamInfo.getRamSize()); + } catch (RuntimeException | Error rollbackFailure) { + failure.addSuppressed(rollbackFailure); + } + } + + // Before commit, remove the only external reference installed for the unpublished + // clone. Its arrays can then be reclaimed while candidate remains the working list. + if (!replacementPublished && cloneList != null) { + cloneList.getQueryContextSet().remove(context); + tvListQueryMap.remove(cloneList); + if (cloneContext != null) { + cloneContext.removeAccessedColumns(cloneList); + } + } + throw failure; } - memChunk.setWorkingTVList(cloneList); - return tvListQueryMap; } } catch (MemoryNotEnoughException ex) { if (listRamInfo != null) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 3063631f93da6..6f35a783a61d3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -82,6 +82,56 @@ public abstract class AlignedTVList extends TVList { private long materializedBitmapMemoryCost; private long arrayMemCostWithoutIndex; + /** + * A fully prepared partial clone. All allocations and validations are completed before this plan + * is returned, so {@link #commit()} only moves already captured references and updates primitive + * accounting fields. + */ + public static final class PartialClonePlan { + private final AlignedTVList sourceList; + private final AlignedTVList cloneList; + private final List[] valueColumnsToMove; + private final List[] bitmapColumnsToMove; + private final long sourceArrayMemCostWithoutIndex; + private final long cloneArrayMemCostWithoutIndex; + private final long sourceBitmapMemoryCost; + private final long cloneBitmapMemoryCost; + + private boolean committed; + + private PartialClonePlan( + AlignedTVList sourceList, + AlignedTVList cloneList, + List[] valueColumnsToMove, + List[] bitmapColumnsToMove, + long sourceArrayMemCostWithoutIndex, + long cloneArrayMemCostWithoutIndex, + long sourceBitmapMemoryCost, + long cloneBitmapMemoryCost) { + this.sourceList = sourceList; + this.cloneList = cloneList; + this.valueColumnsToMove = valueColumnsToMove; + this.bitmapColumnsToMove = bitmapColumnsToMove; + this.sourceArrayMemCostWithoutIndex = sourceArrayMemCostWithoutIndex; + this.cloneArrayMemCostWithoutIndex = cloneArrayMemCostWithoutIndex; + this.sourceBitmapMemoryCost = sourceBitmapMemoryCost; + this.cloneBitmapMemoryCost = cloneBitmapMemoryCost; + } + + public AlignedTVList getCloneList() { + return cloneList; + } + + /** Commit the prepared ownership transfer. This method is idempotent and allocation-free. */ + public synchronized void commit() { + if (committed) { + return; + } + sourceList.commitPartialClone(this); + committed = true; + } + } + // Data type list -> list of TVList, add 1 when expanded -> primitive array of basic type // Index relation: columnIndex(dataTypeIndex) -> arrayIndex -> elementIndex protected List> values; @@ -178,36 +228,98 @@ public synchronized AlignedTVList clone(Set columnsToClone) { return cloneList; } + /** + * Prepare a partial clone without changing this TVList. The returned plan must be committed only + * after the query-memory reservation succeeds. + */ + public synchronized PartialClonePlan preparePartialClone(Set columnsToClone) { + Set retainedColumns = + new HashSet<>(Objects.requireNonNull(columnsToClone, "columnsToClone cannot be null")); + AlignedTVList cloneList = AlignedTVList.newAlignedList(new ArrayList<>(dataTypes)); + cloneAs(cloneList); + cloneColumnDataTo(cloneList, retainedColumns); + return prepareMovePlan(cloneList, retainedColumns); + } + public synchronized void moveUnclonedColumnsTo( AlignedTVList cloneList, Set columnsToClone) { if (columnsToClone == null) { return; } - if (bitMaps != null && cloneList.bitMaps == null) { - for (int i = 0; i < values.size(); i++) { - if (values.get(i) != null && !columnsToClone.contains(i) && bitMaps.get(i) != null) { + Set retainedColumns = new HashSet<>(columnsToClone); + prepareMovePlan(cloneList, retainedColumns).commit(); + } + + @SuppressWarnings("unchecked") + private PartialClonePlan prepareMovePlan(AlignedTVList cloneList, Set retainedColumns) { + Objects.requireNonNull(cloneList, "cloneList cannot be null"); + int columnCount = values.size(); + if (cloneList.values.size() != columnCount + || cloneList.memoryBinaryChunkSize.length != memoryBinaryChunkSize.length) { + throw new IllegalStateException("Target AlignedTVList has incompatible column containers"); + } + + List[] valueColumnsToMove = (List[]) new List[columnCount]; + List[] bitmapColumnsToMove = (List[]) new List[columnCount]; + for (int i = 0; i < columnCount; i++) { + if (retainedColumns.contains(i)) { + continue; + } + + List columnValues = values.get(i); + if (columnValues == null) { + throw new IllegalStateException( + String.format("Missing value arrays for aligned column index %d during move", i)); + } + if (cloneList.values.get(i) == null || !cloneList.values.get(i).isEmpty()) { + throw new IllegalStateException( + String.format("Target value column index %d is not ready for move", i)); + } + valueColumnsToMove[i] = columnValues; + + if (bitMaps != null && bitMaps.get(i) != null) { + if (cloneList.bitMaps == null + || cloneList.bitMaps.size() != bitMaps.size() + || cloneList.bitMaps.get(i) != null) { throw new IllegalStateException( - "Target AlignedTVList is not ready to receive moved bitmaps"); + String.format("Target bitmap column index %d is not ready for move", i)); } + bitmapColumnsToMove[i] = bitMaps.get(i); } } - for (int i = 0; i < values.size(); i++) { - List columnValues = values.get(i); - if (columnValues == null || columnsToClone.contains(i)) { + + return new PartialClonePlan( + this, + cloneList, + valueColumnsToMove, + bitmapColumnsToMove, + calculateArrayMemCostWithoutIndex(retainedColumns), + cloneList.calculateArrayMemCostWithoutIndex(null), + calculateBitmapRamCost(bitMaps, retainedColumns), + calculateBitmapRamCost(bitMaps, null)); + } + + private synchronized void commitPartialClone(PartialClonePlan plan) { + for (int i = 0; i < plan.valueColumnsToMove.length; i++) { + List columnValues = plan.valueColumnsToMove[i]; + if (columnValues == null) { continue; } - cloneList.values.set(i, columnValues); + + plan.cloneList.values.set(i, columnValues); values.set(i, null); - if (bitMaps != null && bitMaps.get(i) != null && cloneList.bitMaps != null) { - cloneList.bitMaps.set(i, bitMaps.get(i)); + List columnBitMaps = plan.bitmapColumnsToMove[i]; + if (columnBitMaps != null) { + plan.cloneList.bitMaps.set(i, columnBitMaps); bitMaps.set(i, null); } memoryBinaryChunkSize[i] = 0; } - materializedBitmapMemoryCost = calculateBitmapRamCost(bitMaps, columnsToClone); - // Column ownership changed on both lists, so refresh their per-block memory cost. - refreshArrayMemCostWithoutIndex(); - cloneList.refreshArrayMemCostWithoutIndex(); + + arrayMemCostWithoutIndex = plan.sourceArrayMemCostWithoutIndex; + plan.cloneList.arrayMemCostWithoutIndex = plan.cloneArrayMemCostWithoutIndex; + materializedBitmapMemoryCost = plan.sourceBitmapMemoryCost; + plan.cloneList.materializedBitmapMemoryCost = plan.cloneBitmapMemoryCost; } /** @@ -1102,19 +1214,77 @@ public synchronized long getRamSize() { return timestamps.size() * (arrayMemCostWithoutIndex + (indices != null ? (long) PrimitiveArrayManager.ARRAY_SIZE * Integer.BYTES : 0)) - + materializedBitmapMemoryCost; + + materializedBitmapMemoryCost + + calculateContainerRamCost(null); } public synchronized long getRamSize(Set columnsToClone) { return timestamps.size() * alignedTvListArrayMemCost(columnsToClone) - + calculateBitmapRamCost(bitMaps, columnsToClone); + + calculateBitmapRamCost(bitMaps, columnsToClone) + + calculateContainerRamCost(columnsToClone); + } + + private long calculateArrayMemCostWithoutIndex(Set retainedColumns) { + long arrayMemCost = alignedTvListArrayMemCost(retainedColumns); + if (indices != null) { + arrayMemCost -= (long) PrimitiveArrayManager.ARRAY_SIZE * Integer.BYTES; + } + return arrayMemCost; } private void refreshArrayMemCostWithoutIndex() { - arrayMemCostWithoutIndex = alignedTvListArrayMemCost(); + arrayMemCostWithoutIndex = calculateArrayMemCostWithoutIndex(null); + } + + /** + * Calculate the one-time container memory retained by this list. Primitive-array references in + * the time/index/value lists and bitmap lists are already charged by the per-block accounting, so + * only their list objects and backing-array headers are added here. In contrast, references in + * the outer column containers are not charged elsewhere and are counted in full. + */ + private long calculateContainerRamCost(Set retainedColumns) { + long size = 0; + + size += listRamCostWithReferences(dataTypes); + size += RamUsageEstimator.sizeOfLongArray(memoryBinaryChunkSize.length); + size += listRamCostWithoutReferences(timestamps); if (indices != null) { - arrayMemCostWithoutIndex -= (long) PrimitiveArrayManager.ARRAY_SIZE * Integer.BYTES; + size += listRamCostWithoutReferences(indices); + } + + size += listRamCostWithReferences(values); + for (int i = 0; i < values.size(); i++) { + if (retainedColumns != null && !retainedColumns.contains(i)) { + continue; + } + List columnValues = values.get(i); + if (columnValues != null) { + size += listRamCostWithoutReferences(columnValues); + } } + + if (bitMaps != null) { + size += listRamCostWithReferences(bitMaps); + for (int i = 0; i < bitMaps.size(); i++) { + if (retainedColumns != null && !retainedColumns.contains(i)) { + continue; + } + List columnBitMaps = bitMaps.get(i); + if (columnBitMaps != null) { + size += listRamCostWithoutReferences(columnBitMaps); + } + } + } + return size; + } + + private static long listRamCostWithReferences(List list) { + return RamUsageEstimator.shallowSizeOf(list) + RamUsageEstimator.sizeOfObjectArray(list.size()); + } + + private static long listRamCostWithoutReferences(List list) { + return RamUsageEstimator.shallowSizeOf(list) + + (list.isEmpty() ? 0 : RamUsageEstimator.sizeOfObjectArray(0)); } private static long calculateBitmapRamCost( @@ -1189,10 +1359,7 @@ public long alignedTvListArrayMemCost(Set columnsToClone) { size += (long) PrimitiveArrayManager.ARRAY_SIZE * (long) type.getDataTypeSize(); } } - // size is 0 when all types are null - if (size == 0) { - return size; - } + // time array mem size size += PrimitiveArrayManager.ARRAY_SIZE * 8L; // index array mem size diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java index 6d0cabb044313..9768d1177d5c1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.queryengine.plan.planner; import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.exception.MemoryNotEnoughException; import org.apache.iotdb.db.queryengine.plan.planner.memory.NotThreadSafeMemoryReservationManager; import org.junit.After; @@ -161,4 +162,49 @@ public void testMemoryReservationManagerNormalPriorityReserveAndRelease() { Assert.assertEquals(0L, manager.getReservedBytesInTotalForTest()); Assert.assertEquals(freeBefore, PLANNER.getFreeMemoryForOperators()); } + + @Test + public void testImmediateReservationRollback() { + long request = Math.min(1024L, PLANNER.getFreeMemoryForOperators()); + if (request <= 0) { + return; + } + + NotThreadSafeMemoryReservationManager manager = + new NotThreadSafeMemoryReservationManager(new QueryId("normal_query"), "test"); + long freeBefore = PLANNER.getFreeMemoryForOperators(); + + manager.reserveMemoryCumulatively(request); + manager.releaseMemoryImmediately(request); + manager.reserveMemoryImmediately(); + + Assert.assertEquals(0L, manager.getReservedBytesInTotalForTest()); + Assert.assertEquals(freeBefore, PLANNER.getFreeMemoryForOperators()); + + manager.reserveMemoryImmediately(request); + manager.releaseMemoryImmediately(request); + + Assert.assertEquals(0L, manager.getReservedBytesInTotalForTest()); + Assert.assertEquals(freeBefore, PLANNER.getFreeMemoryForOperators()); + } + + @Test + public void testFailedCumulativeReservationDoesNotRemainPending() { + long freeBefore = PLANNER.getFreeMemoryForOperators(); + long request = freeBefore + MEMORY_BATCH_THRESHOLD; + NotThreadSafeMemoryReservationManager manager = + new NotThreadSafeMemoryReservationManager(new QueryId("normal_query"), "test"); + + try { + manager.reserveMemoryCumulatively(request); + Assert.fail("Expected insufficient query memory"); + } catch (MemoryNotEnoughException expected) { + // expected + } + + // A stale pending reservation would make this retry fail again. + manager.reserveMemoryImmediately(); + Assert.assertEquals(0L, manager.getReservedBytesInTotalForTest()); + Assert.assertEquals(freeBefore, PLANNER.getFreeMemoryForOperators()); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java index 6027364061cb4..d14fb2c95f6b3 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java @@ -184,11 +184,11 @@ public void testBitmapIsAllocatedLazilyWithCompactBackingArray() { ARRAY_SIZE / Byte.SIZE + 1, firstColumnBitMaps.get(2).getByteArray().length); Assert.assertTrue(tvList.isNullValue(ARRAY_SIZE * 2 + 1, 0)); Assert.assertFalse(tvList.isNullValue(ARRAY_SIZE * 2, 0)); - Assert.assertEquals( + long primitiveArrayAndBitmapCost = 3L * tvList.alignedTvListArrayMemCost() + 3L * AlignedTVList.bitmapReferenceRamCost() - + AlignedTVList.bitmapRamCost(), - tvList.getRamSize()); + + AlignedTVList.bitmapRamCost(); + Assert.assertTrue(tvList.getRamSize() > primitiveArrayAndBitmapCost); Assert.assertEquals(tvList.getRamSize(), tvList.calculateRamSize().getRamSize()); Assert.assertEquals(tvList.getRamSize(), tvList.clone().getRamSize()); Assert.assertEquals(tvList.getRamSize(), tvList.cloneForFlushSort().getRamSize()); @@ -205,14 +205,14 @@ public void testExtendedColumnRamCostIncludesActualBitmaps() { long ramSizeBeforeExtension = tvList.getRamSize(); tvList.extendColumn(TSDataType.INT32); - Assert.assertEquals( - 2L - * (AlignedTVList.valueListArrayMemCost(TSDataType.INT32) - + AlignedTVList.bitmapReferenceRamCost() - + AlignedTVList.bitmapRamCost()), - tvList.getRamSize() - ramSizeBeforeExtension); + Assert.assertTrue( + tvList.getRamSize() - ramSizeBeforeExtension + >= 2L + * (AlignedTVList.valueListArrayMemCost(TSDataType.INT32) + + AlignedTVList.bitmapReferenceRamCost() + + AlignedTVList.bitmapRamCost())); tvList.clear(); - Assert.assertEquals(0, tvList.getRamSize()); + Assert.assertTrue(tvList.getRamSize() > 0); } @Test @@ -348,7 +348,9 @@ public void testMovesUnclonedColumns() { tvList.putAlignedValue(0, new Object[] {1L, 2L, null}); Set columnsToClone = Collections.singleton(1); - AlignedTVList clonedTvList = tvList.clone(columnsToClone); + long retainedRamSize = tvList.calculateRamSize(columnsToClone).getRamSize(); + AlignedTVList.PartialClonePlan partialClonePlan = tvList.preparePartialClone(columnsToClone); + AlignedTVList clonedTvList = partialClonePlan.getCloneList(); Assert.assertNotNull(tvList.getValues().get(0)); Assert.assertNotNull(tvList.getValues().get(2)); @@ -356,7 +358,7 @@ public void testMovesUnclonedColumns() { Assert.assertTrue(tvList.isNullValue(0, 2)); Assert.assertEquals(2L, clonedTvList.getLongByValueIndex(0, 1)); - tvList.moveUnclonedColumnsTo(clonedTvList, columnsToClone); + partialClonePlan.commit(); Assert.assertNull(tvList.getValues().get(0)); Assert.assertNull(tvList.getValues().get(2)); @@ -365,6 +367,98 @@ public void testMovesUnclonedColumns() { Assert.assertEquals(1L, clonedTvList.getLongByValueIndex(0, 0)); Assert.assertEquals(2L, clonedTvList.getLongByValueIndex(0, 1)); Assert.assertTrue(clonedTvList.isNullValue(0, 2)); + Assert.assertEquals(retainedRamSize, tvList.calculateRamSize().getRamSize()); + } + + @Test + public void testPartialRamSizeIncludesWideColumnContainers() { + int columnCount = 256; + List dataTypes = new ArrayList<>(columnCount); + Object[] values = new Object[columnCount]; + for (int i = 0; i < columnCount; i++) { + dataTypes.add(TSDataType.INT64); + values[i] = (long) i; + } + + AlignedTVList tvList = AlignedTVList.newAlignedList(dataTypes); + tvList.putAlignedValue(1, values); + Set retainedColumns = Collections.singleton(0); + long primitiveArrayCost = + (long) tvList.getTimestamps().size() * tvList.alignedTvListArrayMemCost(retainedColumns); + long retainedRamSize = tvList.calculateRamSize(retainedColumns).getRamSize(); + + // memoryBinaryChunkSize and the outer column containers remain N-wide after partial move. + Assert.assertTrue(retainedRamSize - primitiveArrayCost >= (long) columnCount * Long.BYTES); + + AlignedTVList.PartialClonePlan plan = tvList.preparePartialClone(retainedColumns); + plan.commit(); + Assert.assertEquals(retainedRamSize, tvList.calculateRamSize().getRamSize()); + } + + @Test + public void testPartialReservationMatchesCleanupCalculation() { + for (boolean createIndices : new boolean[] {false, true}) { + for (boolean retainValueColumn : new boolean[] {false, true}) { + AlignedTVList tvList = + AlignedTVList.newAlignedList( + new ArrayList<>( + Arrays.asList(TSDataType.INT64, TSDataType.INT64, TSDataType.INT64))); + for (int i = 0; i <= ARRAY_SIZE; i++) { + long time = createIndices ? ARRAY_SIZE - i : i; + tvList.putAlignedValue( + time, new Object[] {(long) i, i % 2 == 0 ? null : (long) i, (long) i}); + } + if (createIndices) { + Assert.assertFalse(tvList.isSorted()); + tvList.sort(); + Assert.assertNotNull(tvList.getIndices()); + } else { + Assert.assertNull(tvList.getIndices()); + } + + Set retainedColumns = + retainValueColumn ? Collections.singleton(1) : Collections.emptySet(); + long reservedMemoryBytes = tvList.calculateRamSize(retainedColumns).getRamSize(); + tvList.setReservedMemoryBytes(reservedMemoryBytes); + + AlignedTVList.PartialClonePlan plan = tvList.preparePartialClone(retainedColumns); + plan.commit(); + + long cleanupMemoryBytes = tvList.calculateRamSize().getRamSize(); + String scenario = + String.format( + "createIndices=%s, retainValueColumn=%s", createIndices, retainValueColumn); + Assert.assertEquals(scenario, reservedMemoryBytes, cleanupMemoryBytes); + Assert.assertEquals(scenario, tvList.getReservedMemoryBytes(), cleanupMemoryBytes); + } + } + } + + @Test + public void testMoveValidationFailureLeavesSourceUntouched() { + AlignedTVList tvList = + AlignedTVList.newAlignedList( + Arrays.asList(TSDataType.INT64, TSDataType.INT64, TSDataType.INT64)); + tvList.putAlignedValue(0, new Object[] {null, 2L, 3L}); + + Set columnsToClone = Collections.singleton(1); + AlignedTVList clonedTvList = tvList.clone(columnsToClone); + clonedTvList.bitMaps = null; + + try { + tvList.moveUnclonedColumnsTo(clonedTvList, columnsToClone); + Assert.fail("Expected move preparation to reject the incomplete target"); + } catch (IllegalStateException expected) { + // expected + } + + Assert.assertNotNull(tvList.getValues().get(0)); + Assert.assertNotNull(tvList.getValues().get(1)); + Assert.assertNotNull(tvList.getValues().get(2)); + Assert.assertNotNull(tvList.getBitMaps().get(0)); + Assert.assertTrue(tvList.isNullValue(0, 0)); + Assert.assertEquals(2L, tvList.getLongByValueIndex(0, 1)); + Assert.assertEquals(3L, tvList.getLongByValueIndex(0, 2)); } @Test From cf31a18e6be2414eda003de715efefb4df173f8c Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 5 Aug 2026 11:44:12 +0800 Subject: [PATCH 13/16] Fix concurrent temporary TVList query registration --- .../utils/ResourceByPathUtils.java | 11 +- .../utils/ResourceByPathUtilsTest.java | 109 ++++++++++++++++++ 2 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtilsTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java index ef3ffd6acc840..6a3f7f23ab6cf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java @@ -224,7 +224,16 @@ protected Map prepareTvListMapForQuery( */ candidate.getQueryContextSet().add(context); context.addTVListToSet(Collections.singleton(candidate)); - workingListForFlushSort.getQueryContextSet().add(context); + // Query preparation is serialized by candidate's query-list lock, but cleanup removes + // the context under workingListForFlushSort's own lock. Use the same lock for this add + // to avoid concurrently mutating its HashSet. The lock order here is candidate first, + // then workingListForFlushSort; cleanup never holds both locks at the same time. + workingListForFlushSort.lockQueryList(); + try { + workingListForFlushSort.getQueryContextSet().add(context); + } finally { + workingListForFlushSort.unlockQueryList(); + } tvListQueryMap.put(workingListForFlushSort, workingListForFlushSort.rowCount()); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtilsTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtilsTest.java new file mode 100644 index 0000000000000..cd1ece64f083c --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtilsTest.java @@ -0,0 +1,109 @@ +/* + * 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.schemaengine.schemaregion.utils; + +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunk; +import org.apache.iotdb.db.utils.datastructure.TVList; + +import org.apache.tsfile.enums.TSDataType; +import org.junit.Assert; +import org.junit.Test; + +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ResourceByPathUtilsTest { + + @Test + public void testFlushingQueryLocksTemporaryTVListBeforeRegistration() throws Exception { + TVList candidate = TVList.newList(TSDataType.INT64); + candidate.putLong(2, 2); + candidate.putLong(1, 1); + Assert.assertFalse(candidate.isSorted()); + + QueryContext previousQuery = new QueryContext(1, false); + candidate.lockQueryList(); + try { + candidate.getQueryContextSet().add(previousQuery); + } finally { + candidate.unlockQueryList(); + } + + TVList temporaryList = candidate.cloneForFlushSort(); + IWritableMemChunk memChunk = mock(IWritableMemChunk.class); + when(memChunk.getSortedList()).thenReturn(Collections.emptyList()); + when(memChunk.getWorkingTVList()).thenReturn(candidate); + CountDownLatch temporaryListInitialized = new CountDownLatch(1); + when(memChunk.initWorkingListForFlushIfNecessary(candidate, true)) + .thenAnswer( + ignored -> { + temporaryListInitialized.countDown(); + return temporaryList; + }); + + ResourceByPathUtils resourceByPathUtils = + ResourceByPathUtils.getResourceInstance( + new MeasurementPath("root.test.d.s", TSDataType.INT64)); + QueryContext currentQuery = new QueryContext(2, false); + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future> result = null; + try { + temporaryList.lockQueryList(); + try { + result = + executor.submit( + () -> + resourceByPathUtils.prepareTvListMapForQuery( + currentQuery, memChunk, false, null, null)); + Assert.assertTrue(temporaryListInitialized.await(3, TimeUnit.SECONDS)); + Future> blockedResult = result; + Assert.assertThrows( + TimeoutException.class, () -> blockedResult.get(200, TimeUnit.MILLISECONDS)); + } finally { + temporaryList.unlockQueryList(); + } + + Map tvListQueryMap = result.get(3, TimeUnit.SECONDS); + Assert.assertTrue(tvListQueryMap.containsKey(temporaryList)); + temporaryList.lockQueryList(); + try { + Assert.assertTrue(temporaryList.getQueryContextSet().contains(currentQuery)); + } finally { + temporaryList.unlockQueryList(); + } + } finally { + if (result != null) { + result.cancel(true); + } + executor.shutdownNow(); + executor.awaitTermination(3, TimeUnit.SECONDS); + } + } +} From cbeecc310680bbda7ef6f132f82e439b71e25b10 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 5 Aug 2026 12:03:11 +0800 Subject: [PATCH 14/16] Remove obsolete partial TVList clone API --- .../db/utils/datastructure/AlignedTVList.java | 13 ++------ .../datastructure/AlignedTVListTest.java | 31 +++++++++---------- 2 files changed, 17 insertions(+), 27 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 6f35a783a61d3..8923c57d3089c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -241,15 +241,6 @@ public synchronized PartialClonePlan preparePartialClone(Set columnsToC return prepareMovePlan(cloneList, retainedColumns); } - public synchronized void moveUnclonedColumnsTo( - AlignedTVList cloneList, Set columnsToClone) { - if (columnsToClone == null) { - return; - } - Set retainedColumns = new HashSet<>(columnsToClone); - prepareMovePlan(cloneList, retainedColumns).commit(); - } - @SuppressWarnings("unchecked") private PartialClonePlan prepareMovePlan(AlignedTVList cloneList, Set retainedColumns) { Objects.requireNonNull(cloneList, "cloneList cannot be null"); @@ -883,8 +874,8 @@ protected Object cloneValue(TSDataType type, Object value) { * * This method only performs the allocation phase: clone requested value/bitmap arrays and prepare * bitmap containers that will be needed by moved columns. It must not clear or move columns from - * the source TVList here. The destructive move is committed by moveUnclonedColumnsTo() only after - * cloneList is fully prepared for publication. + * the source TVList here. The destructive move is performed only by PartialClonePlan.commit() + * after cloneList and the ownership-transfer plan are fully prepared for publication. */ private void cloneColumnDataTo(AlignedTVList cloneList, Set columnsToClone) { boolean cloneAllColumns = columnsToClone == null; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java index d14fb2c95f6b3..f84710bf6a08b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/datastructure/AlignedTVListTest.java @@ -435,30 +435,29 @@ public void testPartialReservationMatchesCleanupCalculation() { } @Test - public void testMoveValidationFailureLeavesSourceUntouched() { + public void testPartialCloneFailureLeavesSourceUntouched() { AlignedTVList tvList = AlignedTVList.newAlignedList( Arrays.asList(TSDataType.INT64, TSDataType.INT64, TSDataType.INT64)); tvList.putAlignedValue(0, new Object[] {null, 2L, 3L}); - Set columnsToClone = Collections.singleton(1); - AlignedTVList clonedTvList = tvList.clone(columnsToClone); - clonedTvList.bitMaps = null; - - try { - tvList.moveUnclonedColumnsTo(clonedTvList, columnsToClone); - Assert.fail("Expected move preparation to reject the incomplete target"); - } catch (IllegalStateException expected) { - // expected - } + List firstColumnValues = tvList.getValues().get(0); + List secondColumnValues = tvList.getValues().get(1); + List thirdColumnValues = tvList.getValues().get(2); + List firstColumnBitMaps = tvList.getBitMaps().get(0); + Object invalidThirdColumnArray = new int[ARRAY_SIZE]; + thirdColumnValues.set(0, invalidThirdColumnArray); - Assert.assertNotNull(tvList.getValues().get(0)); - Assert.assertNotNull(tvList.getValues().get(1)); - Assert.assertNotNull(tvList.getValues().get(2)); - Assert.assertNotNull(tvList.getBitMaps().get(0)); + Set columnsToClone = new HashSet<>(Arrays.asList(0, 1, 2)); + Assert.assertThrows(ClassCastException.class, () -> tvList.preparePartialClone(columnsToClone)); + + Assert.assertSame(firstColumnValues, tvList.getValues().get(0)); + Assert.assertSame(secondColumnValues, tvList.getValues().get(1)); + Assert.assertSame(thirdColumnValues, tvList.getValues().get(2)); + Assert.assertSame(invalidThirdColumnArray, tvList.getValues().get(2).get(0)); + Assert.assertSame(firstColumnBitMaps, tvList.getBitMaps().get(0)); Assert.assertTrue(tvList.isNullValue(0, 0)); Assert.assertEquals(2L, tvList.getLongByValueIndex(0, 1)); - Assert.assertEquals(3L, tvList.getLongByValueIndex(0, 2)); } @Test From 0c2d4431cb5a7063e1be2f7fb93b068e4fac3832 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 5 Aug 2026 14:02:14 +0800 Subject: [PATCH 15/16] Fix aligned TVList RAM cost test expectations --- .../memtable/TsFileProcessorTest.java | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessorTest.java index db3b84dbf583f..0e2e533cc8715 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessorTest.java @@ -688,11 +688,11 @@ public void alignedTvListRamCostTest() // Test Tablet processor.insertTablet(genInsertTableNode(0, true), 0, 10, new TSStatus[10]); IMemTable memTable = processor.getWorkMemTable(); - Assert.assertEquals(1596552, memTable.getTVListsRamCost()); + Assert.assertEquals(1764688, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNode(100, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(1596552, memTable.getTVListsRamCost()); + Assert.assertEquals(1764688, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNode(200, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(1596552, memTable.getTVListsRamCost()); + Assert.assertEquals(1764688, memTable.getTVListsRamCost()); Assert.assertEquals(90000, memTable.getTotalPointsNum()); Assert.assertEquals(720360, memTable.memSize()); // Test records @@ -701,7 +701,7 @@ public void alignedTvListRamCostTest() record.addTuple(DataPoint.getDataPoint(dataType, measurementId, String.valueOf(i))); processor.insert(buildInsertRowNodeByTSRecord(record), new long[4]); } - Assert.assertEquals(1598168, memTable.getTVListsRamCost()); + Assert.assertEquals(1766304, memTable.getTVListsRamCost()); Assert.assertEquals(90100, memTable.getTotalPointsNum()); Assert.assertEquals(721560, memTable.memSize()); } @@ -724,21 +724,21 @@ public void alignedTvListRamCostTest2() // Test Tablet processor.insertTablet(genInsertTableNode(0, true), 0, 10, new TSStatus[10]); IMemTable memTable = processor.getWorkMemTable(); - Assert.assertEquals(1596552, memTable.getTVListsRamCost()); + Assert.assertEquals(1764688, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNodeFors3000ToS6000(0, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(3552552, memTable.getTVListsRamCost()); + Assert.assertEquals(4152728, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNode(100, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(3552552, memTable.getTVListsRamCost()); + Assert.assertEquals(4152728, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNodeFors3000ToS6000(100, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(3552552, memTable.getTVListsRamCost()); + Assert.assertEquals(4152728, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNode(200, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(3552552, memTable.getTVListsRamCost()); + Assert.assertEquals(4152728, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNodeFors3000ToS6000(200, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(3552552, memTable.getTVListsRamCost()); + Assert.assertEquals(4152728, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNode(300, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(6937104, memTable.getTVListsRamCost()); + Assert.assertEquals(7537280, memTable.getTVListsRamCost()); processor.insertTablet(genInsertTableNodeFors3000ToS6000(300, true), 0, 10, new TSStatus[10]); - Assert.assertEquals(7105104, memTable.getTVListsRamCost()); + Assert.assertEquals(7705280, memTable.getTVListsRamCost()); Assert.assertEquals(240000, memTable.getTotalPointsNum()); Assert.assertEquals(1920960, memTable.memSize()); @@ -748,14 +748,14 @@ public void alignedTvListRamCostTest2() record.addTuple(DataPoint.getDataPoint(dataType, measurementId, String.valueOf(i))); processor.insert(buildInsertRowNodeByTSRecord(record), new long[4]); } - Assert.assertEquals(7106720, memTable.getTVListsRamCost()); + Assert.assertEquals(7706896, memTable.getTVListsRamCost()); // Test records for (int i = 1; i <= 100; i++) { TSRecord record = new TSRecord(i, deviceId); record.addTuple(DataPoint.getDataPoint(dataType, "s1", String.valueOf(i))); processor.insert(buildInsertRowNodeByTSRecord(record), new long[4]); } - Assert.assertEquals(7108336, memTable.getTVListsRamCost()); + Assert.assertEquals(7708512, memTable.getTVListsRamCost()); Assert.assertEquals(240200, memTable.getTotalPointsNum()); Assert.assertEquals(1923360, memTable.memSize()); } From b2f7b2eaf78f55301f90b6e6e861273a337a112c Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 5 Aug 2026 15:03:21 +0800 Subject: [PATCH 16/16] Remove unused partial clone overload --- .../apache/iotdb/db/utils/datastructure/AlignedTVList.java | 7 ------- 1 file changed, 7 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java index 8923c57d3089c..d319b9c9a0d46 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java @@ -221,13 +221,6 @@ public synchronized AlignedTVList clone() { return cloneList; } - public synchronized AlignedTVList clone(Set columnsToClone) { - AlignedTVList cloneList = AlignedTVList.newAlignedList(new ArrayList<>(dataTypes)); - cloneAs(cloneList); - cloneColumnDataTo(cloneList, columnsToClone); - return cloneList; - } - /** * Prepare a partial clone without changing this TVList. The returned plan must be committed only * after the query-memory reservation succeeds.