Skip to content

Commit 8a8a6da

Browse files
CaideyipiCaideyipi
authored
[Pipe] Avoid quadratic aligned page memory estimation (#18917)
1 parent f709ef5 commit 8a8a6da

4 files changed

Lines changed: 292 additions & 95 deletions

File tree

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java‎

Lines changed: 108 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -229,44 +229,120 @@ public static long calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(
229229

230230
public static List<Long> calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
231231
final Chunk timeChunk, final List<Chunk> valueChunkList) throws IOException {
232-
final ByteBuffer timeChunkDataBuffer = timeChunk.getData().duplicate();
233-
final List<ByteBuffer> valueChunkDataBufferList = new ArrayList<>(valueChunkList.size());
232+
final PageMemoryUsageEstimator estimator = new PageMemoryUsageEstimator(timeChunk);
234233
for (final Chunk valueChunk : valueChunkList) {
235-
valueChunkDataBufferList.add(
236-
Objects.isNull(valueChunk) ? null : valueChunk.getData().duplicate());
234+
estimator.addValueChunk(valueChunk);
237235
}
236+
return estimator.getPageEstimatedMemoryUsageInBytesList();
237+
}
238238

239-
final List<Long> pageEstimatedMemoryUsageInBytesList = new ArrayList<>();
240-
while (timeChunkDataBuffer.remaining() > 0) {
241-
long pageUncompressedSizeInBytes = 0;
242-
final PageHeader timePageHeader =
243-
SinglePageWholeChunkReader.deserializePageHeader(
244-
timeChunkDataBuffer, timeChunk.getHeader());
245-
pageUncompressedSizeInBytes += timePageHeader.getUncompressedSize();
246-
SinglePageWholeChunkReader.skipCompressedPageData(timeChunkDataBuffer, timePageHeader);
247-
248-
final List<TSDataType> valueDataTypeList = new ArrayList<>(valueChunkList.size());
249-
for (int i = 0; i < valueChunkList.size(); ++i) {
250-
final Chunk valueChunk = valueChunkList.get(i);
251-
final ByteBuffer valueChunkDataBuffer = valueChunkDataBufferList.get(i);
252-
if (Objects.isNull(valueChunk) || Objects.isNull(valueChunkDataBuffer)) {
253-
valueDataTypeList.add(null);
254-
continue;
255-
}
239+
static final class PageMemoryUsageEstimator {
256240

257-
final PageHeader valuePageHeader =
241+
private final List<Long> pageUncompressedSizeInBytesList = new ArrayList<>();
242+
private final List<Integer> pageRowCountList = new ArrayList<>();
243+
private List<Long> pageEstimatedMemoryUsageInBytesList = new ArrayList<>();
244+
245+
private int valueCount;
246+
private TSDataType singleValueDataType;
247+
private long vectorValueMemoryUsageInBytes;
248+
249+
PageMemoryUsageEstimator(final Chunk timeChunk) throws IOException {
250+
final ByteBuffer timeChunkDataBuffer = timeChunk.getData().duplicate();
251+
timeChunkDataBuffer.rewind();
252+
while (timeChunkDataBuffer.remaining() > 0) {
253+
final PageHeader timePageHeader =
258254
SinglePageWholeChunkReader.deserializePageHeader(
259-
valueChunkDataBuffer, valueChunk.getHeader());
260-
pageUncompressedSizeInBytes += valuePageHeader.getUncompressedSize();
261-
valueDataTypeList.add(valueChunk.getHeader().getDataType());
262-
SinglePageWholeChunkReader.skipCompressedPageData(valueChunkDataBuffer, valuePageHeader);
255+
timeChunkDataBuffer, timeChunk.getHeader());
256+
pageUncompressedSizeInBytesList.add((long) timePageHeader.getUncompressedSize());
257+
pageRowCountList.add(SinglePageWholeChunkReader.getPageRowCount(timePageHeader, timeChunk));
258+
SinglePageWholeChunkReader.skipCompressedPageData(timeChunkDataBuffer, timePageHeader);
259+
}
260+
pageEstimatedMemoryUsageInBytesList = calculatePageEstimatedMemoryUsageInBytesList();
261+
}
262+
263+
long calculateMaxPageEstimatedMemoryUsageInBytesWithValueChunk(final Chunk valueChunk)
264+
throws IOException {
265+
final List<Long> pageEstimatedMemoryUsageInBytesList =
266+
calculatePageEstimatedMemoryUsageInBytesList(valueChunk, false);
267+
return pageEstimatedMemoryUsageInBytesList.isEmpty()
268+
? 0
269+
: pageEstimatedMemoryUsageInBytesList.get(0);
270+
}
271+
272+
void addValueChunk(final Chunk valueChunk) throws IOException {
273+
pageEstimatedMemoryUsageInBytesList =
274+
calculatePageEstimatedMemoryUsageInBytesList(valueChunk, true);
275+
}
276+
277+
List<Long> getPageEstimatedMemoryUsageInBytesList() {
278+
return pageEstimatedMemoryUsageInBytesList;
279+
}
280+
281+
private List<Long> calculatePageEstimatedMemoryUsageInBytesList() {
282+
final List<Long> pageEstimatedMemoryUsageInBytesList =
283+
new ArrayList<>(pageUncompressedSizeInBytesList.size());
284+
for (int i = 0; i < pageUncompressedSizeInBytesList.size(); ++i) {
285+
pageEstimatedMemoryUsageInBytesList.add(
286+
SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
287+
pageUncompressedSizeInBytesList.get(i),
288+
pageRowCountList.get(i),
289+
valueCount,
290+
singleValueDataType,
291+
vectorValueMemoryUsageInBytes));
263292
}
264-
pageEstimatedMemoryUsageInBytesList.add(
265-
SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
266-
pageUncompressedSizeInBytes,
267-
SinglePageWholeChunkReader.getPageRowCount(timePageHeader, timeChunk),
268-
valueDataTypeList));
293+
return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
294+
}
295+
296+
private List<Long> calculatePageEstimatedMemoryUsageInBytesList(
297+
final Chunk valueChunk, final boolean commit) throws IOException {
298+
final ByteBuffer valueChunkDataBuffer;
299+
final TSDataType valueDataType;
300+
if (Objects.isNull(valueChunk)) {
301+
valueChunkDataBuffer = null;
302+
valueDataType = null;
303+
} else {
304+
valueChunkDataBuffer = valueChunk.getData().duplicate();
305+
valueChunkDataBuffer.rewind();
306+
valueDataType = valueChunk.getHeader().getDataType();
307+
}
308+
309+
final int newValueCount = valueCount + 1;
310+
final TSDataType newSingleValueDataType = valueCount == 0 ? valueDataType : null;
311+
final long newVectorValueMemoryUsageInBytes =
312+
vectorValueMemoryUsageInBytes
313+
+ SinglePageWholeChunkReader.estimateVectorValueMemoryUsageInBytes(valueDataType);
314+
final List<Long> newPageUncompressedSizeInBytesList =
315+
new ArrayList<>(pageUncompressedSizeInBytesList);
316+
final List<Long> pageEstimatedMemoryUsageInBytesList =
317+
new ArrayList<>(pageUncompressedSizeInBytesList.size());
318+
319+
for (int i = 0; i < newPageUncompressedSizeInBytesList.size(); ++i) {
320+
if (Objects.nonNull(valueChunkDataBuffer)) {
321+
final PageHeader valuePageHeader =
322+
SinglePageWholeChunkReader.deserializePageHeader(
323+
valueChunkDataBuffer, valueChunk.getHeader());
324+
newPageUncompressedSizeInBytesList.set(
325+
i, newPageUncompressedSizeInBytesList.get(i) + valuePageHeader.getUncompressedSize());
326+
SinglePageWholeChunkReader.skipCompressedPageData(valueChunkDataBuffer, valuePageHeader);
327+
}
328+
pageEstimatedMemoryUsageInBytesList.add(
329+
SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
330+
newPageUncompressedSizeInBytesList.get(i),
331+
pageRowCountList.get(i),
332+
newValueCount,
333+
newSingleValueDataType,
334+
newVectorValueMemoryUsageInBytes));
335+
}
336+
337+
if (commit) {
338+
pageUncompressedSizeInBytesList.clear();
339+
pageUncompressedSizeInBytesList.addAll(newPageUncompressedSizeInBytesList);
340+
valueCount = newValueCount;
341+
singleValueDataType = newSingleValueDataType;
342+
vectorValueMemoryUsageInBytes = newVectorValueMemoryUsageInBytes;
343+
}
344+
345+
return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
269346
}
270-
return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
271347
}
272348
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java‎

Lines changed: 35 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,7 @@ static int getPageRowCount(final PageHeader pageHeader, final Chunk chunk) throw
191191

192192
private static int countSinglePageTimeValues(final Chunk chunk) throws IOException {
193193
final ByteBuffer chunkDataBuffer = chunk.getData().duplicate();
194+
chunkDataBuffer.rewind();
194195
final PageHeader pageHeader = deserializePageHeader(chunkDataBuffer, chunk.getHeader());
195196
final ByteBuffer pageData =
196197
deserializePageData(
@@ -222,25 +223,42 @@ static long estimatePageMemoryUsageInBytesWithBatchData(
222223
final long pageUncompressedSizeInBytes,
223224
final int rowCount,
224225
final List<TSDataType> valueDataTypeList) {
226+
final int valueCount = valueDataTypeList.size();
227+
return estimatePageMemoryUsageInBytesWithBatchData(
228+
pageUncompressedSizeInBytes,
229+
rowCount,
230+
valueCount,
231+
valueCount == 1 ? valueDataTypeList.get(0) : null,
232+
estimateVectorValueMemoryUsageInBytes(valueDataTypeList));
233+
}
234+
235+
static long estimatePageMemoryUsageInBytesWithBatchData(
236+
final long pageUncompressedSizeInBytes,
237+
final int rowCount,
238+
final int valueCount,
239+
final TSDataType singleValueDataType,
240+
final long vectorValueMemoryUsageInBytes) {
225241
return pageUncompressedSizeInBytes
226-
+ estimateBatchDataMemoryUsageInBytes(rowCount, valueDataTypeList);
242+
+ estimateBatchDataMemoryUsageInBytes(
243+
rowCount, valueCount, singleValueDataType, vectorValueMemoryUsageInBytes);
227244
}
228245

229246
private static long estimateBatchDataMemoryUsageInBytes(
230-
final int rowCount, final List<TSDataType> valueDataTypeList) {
231-
final int valueCount = valueDataTypeList.size();
247+
final int rowCount,
248+
final int valueCount,
249+
final TSDataType singleValueDataType,
250+
final long vectorValueMemoryUsageInBytes) {
232251
final long segmentCount = Math.max(1, (rowCount + 15L) / 16);
233252
long estimatedMemoryUsageInBytes = RamUsageEstimator.sizeOfLongArray(16) * segmentCount;
234253

235254
if (valueCount == 1) {
236255
estimatedMemoryUsageInBytes +=
237-
estimateSingleValueArrayMemoryUsageInBytes(rowCount, valueDataTypeList.get(0));
256+
estimateSingleValueArrayMemoryUsageInBytes(rowCount, singleValueDataType);
238257
} else if (valueCount > 1) {
239258
estimatedMemoryUsageInBytes += RamUsageEstimator.sizeOfObjectArray(16) * segmentCount;
240259
estimatedMemoryUsageInBytes +=
241260
(long) rowCount
242-
* (RamUsageEstimator.sizeOfObjectArray(valueCount)
243-
+ estimateVectorValueMemoryUsageInBytes(valueDataTypeList));
261+
* (RamUsageEstimator.sizeOfObjectArray(valueCount) + vectorValueMemoryUsageInBytes);
244262
}
245263

246264
return estimatedMemoryUsageInBytes;
@@ -260,18 +278,21 @@ private static long estimateVectorValueMemoryUsageInBytes(
260278
final List<TSDataType> valueDataTypeList) {
261279
long estimatedMemoryUsageInBytes = 0;
262280
for (final TSDataType dataType : valueDataTypeList) {
263-
if (Objects.isNull(dataType)) {
264-
continue;
265-
}
266-
267-
estimatedMemoryUsageInBytes +=
268-
RamUsageEstimator.alignObjectSize(
269-
RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
270-
+ estimateTsPrimitiveTypeValueMemoryUsageInBytes(dataType));
281+
estimatedMemoryUsageInBytes += estimateVectorValueMemoryUsageInBytes(dataType);
271282
}
272283
return estimatedMemoryUsageInBytes;
273284
}
274285

286+
static long estimateVectorValueMemoryUsageInBytes(final TSDataType dataType) {
287+
if (Objects.isNull(dataType)) {
288+
return 0;
289+
}
290+
291+
return RamUsageEstimator.alignObjectSize(
292+
RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
293+
+ estimateTsPrimitiveTypeValueMemoryUsageInBytes(dataType));
294+
}
295+
275296
private static long estimateTsPrimitiveTypeValueMemoryUsageInBytes(final TSDataType dataType) {
276297
return Type.fromTsDataType(dataType).estimateValueSize();
277298
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java‎

Lines changed: 19 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,6 @@
7474

7575
import java.io.File;
7676
import java.io.IOException;
77-
import java.nio.ByteBuffer;
7877
import java.util.ArrayList;
7978
import java.util.Collections;
8079
import java.util.HashMap;
@@ -893,9 +892,8 @@ private boolean useNextPendingAlignedChunk(final byte marker) throws IOException
893892
}
894893
final List<Long> pageEstimatedMemoryUsageInBytesList =
895894
nextIsMultiPage
896-
? AlignedSinglePageWholeChunkReader
897-
.calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
898-
timeChunk, pendingAlignedChunkGroup.valueChunkList)
895+
? pendingAlignedChunkGroup.pageMemoryUsageEstimator
896+
.getPageEstimatedMemoryUsageInBytesList()
899897
: Collections.emptyList();
900898
final long maxPageEstimatedMemoryUsageInBytes =
901899
pageEstimatedMemoryUsageInBytesList.isEmpty()
@@ -996,6 +994,7 @@ private void cacheAlignedValueChunk(final CachedAlignedValueChunk valueChunk) th
996994
resizeChunkMemoryBlockIfFirstValueChunkExceedsLimit(pendingAlignedChunkGroup, valueChunk);
997995
resizePageDataMemoryBlockIfFirstValueChunkExceedsLimit(pendingAlignedChunkGroup, valueChunk);
998996

997+
pendingAlignedChunkGroup.pageMemoryUsageEstimator.addValueChunk(valueChunk.chunk);
999998
pendingAlignedChunkGroup.valueChunkList.add(valueChunk.chunk);
1000999
pendingAlignedChunkGroup.chunkSize += valueChunk.valueChunkSize;
10011000
pendingAlignedChunkSize += valueChunk.valueChunkSize;
@@ -1013,7 +1012,8 @@ private void cacheAlignedValueChunk(final CachedAlignedValueChunk valueChunk) th
10131012
currentDevice, Collections.singletonList(measurementID), currentModifications, false));
10141013
}
10151014

1016-
private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int timeChunkIndex) {
1015+
private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int timeChunkIndex)
1016+
throws IOException {
10171017
final PendingAlignedChunkGroup pendingAlignedChunkGroup =
10181018
findPendingAlignedChunkGroup(timeChunkIndex);
10191019
if (Objects.nonNull(pendingAlignedChunkGroup)) {
@@ -1023,7 +1023,8 @@ private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int t
10231023
final PendingAlignedChunkGroup newPendingAlignedChunkGroup =
10241024
new PendingAlignedChunkGroup(
10251025
timeChunkIndex,
1026-
PipeMemoryWeightUtil.calculateChunkRamBytesUsed(timeChunkList.get(timeChunkIndex)));
1026+
PipeMemoryWeightUtil.calculateChunkRamBytesUsed(timeChunkList.get(timeChunkIndex)),
1027+
timeChunkList.get(timeChunkIndex));
10271028
pendingAlignedChunkSize += newPendingAlignedChunkGroup.chunkSize;
10281029

10291030
for (int i = 0; i < pendingAlignedChunkGroups.size(); ++i) {
@@ -1099,49 +1100,12 @@ private long calculateMaxAlignedPageMemorySizeWithBatchData(
10991100
final PendingAlignedChunkGroup pendingAlignedChunkGroup,
11001101
final CachedAlignedValueChunk valueChunk)
11011102
throws IOException {
1102-
final List<Chunk> valueChunkList =
1103-
new ArrayList<>(
1104-
(Objects.isNull(pendingAlignedChunkGroup)
1105-
? 0
1106-
: pendingAlignedChunkGroup.valueChunkList.size())
1107-
+ 1);
1108-
if (Objects.nonNull(pendingAlignedChunkGroup)) {
1109-
valueChunkList.addAll(pendingAlignedChunkGroup.valueChunkList);
1110-
}
1111-
valueChunkList.add(valueChunk.chunk);
1112-
11131103
final Chunk timeChunk = timeChunkList.get(timeChunkIndex);
1114-
final int timeChunkDataPosition = timeChunk.getData().position();
1115-
final List<Integer> valueChunkDataPositions = new ArrayList<>(valueChunkList.size());
1116-
for (final Chunk chunk : valueChunkList) {
1117-
valueChunkDataPositions.add(Objects.isNull(chunk) ? 0 : chunk.getData().position());
1118-
}
1119-
1120-
rewindChunkData(timeChunk);
1121-
valueChunkList.forEach(this::rewindChunkData);
1122-
try {
1123-
return AlignedSinglePageWholeChunkReader
1124-
.calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(timeChunk, valueChunkList);
1125-
} finally {
1126-
timeChunk.getData().position(timeChunkDataPosition);
1127-
for (int i = 0; i < valueChunkList.size(); ++i) {
1128-
final Chunk chunk = valueChunkList.get(i);
1129-
if (Objects.nonNull(chunk)) {
1130-
chunk.getData().position(valueChunkDataPositions.get(i));
1131-
}
1132-
}
1133-
}
1134-
}
1135-
1136-
private void rewindChunkData(final Chunk chunk) {
1137-
if (Objects.isNull(chunk)) {
1138-
return;
1139-
}
1140-
1141-
final ByteBuffer chunkData = chunk.getData();
1142-
if (Objects.nonNull(chunkData)) {
1143-
chunkData.rewind();
1144-
}
1104+
final AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator estimator =
1105+
Objects.isNull(pendingAlignedChunkGroup)
1106+
? new AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator(timeChunk)
1107+
: pendingAlignedChunkGroup.pageMemoryUsageEstimator;
1108+
return estimator.calculateMaxPageEstimatedMemoryUsageInBytesWithValueChunk(valueChunk.chunk);
11451109
}
11461110

11471111
private boolean isSinglePageValueChunk(final ChunkHeader chunkHeader) {
@@ -1207,14 +1171,20 @@ private Statistics findNonAlignedChunkStatistics(
12071171
private static class PendingAlignedChunkGroup {
12081172

12091173
private final int timeChunkIndex;
1174+
private final AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator
1175+
pageMemoryUsageEstimator;
12101176
private final List<Chunk> valueChunkList = new ArrayList<>();
12111177
private final List<IMeasurementSchema> measurements = new ArrayList<>();
12121178
private final List<ModsOperationUtil.ModsInfo> modsInfos = new ArrayList<>();
12131179
private long chunkSize;
12141180

1215-
private PendingAlignedChunkGroup(final int timeChunkIndex, final long timeChunkSize) {
1181+
private PendingAlignedChunkGroup(
1182+
final int timeChunkIndex, final long timeChunkSize, final Chunk timeChunk)
1183+
throws IOException {
12161184
this.timeChunkIndex = timeChunkIndex;
12171185
this.chunkSize = timeChunkSize;
1186+
this.pageMemoryUsageEstimator =
1187+
new AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator(timeChunk);
12181188
}
12191189
}
12201190

0 commit comments

Comments
 (0)