You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ji...@apache.org on 2026/10/20 04:26:54 UTC
(iotdb) branch master updated: [Pipe] Avoid quadratic aligned page memory estimation (#18917)
This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 8a8a6da57be [Pipe] Avoid quadratic aligned page memory estimation (#18917)
8a8a6da57be is described below
commit 8a8a6da57beda6589ffc955aab7432b9976c1985
Author: Caideyipi <87...@users.noreply.github.com>
AuthorDate: Tue Oct 20 12:26:48 2026 +0800
[Pipe] Avoid quadratic aligned page memory estimation (#18917)
---
.../scan/AlignedSinglePageWholeChunkReader.java | 140 ++++++++++++++++-----
.../parser/scan/SinglePageWholeChunkReader.java | 49 +++++---
.../scan/TsFileInsertionEventScanParser.java | 68 +++-------
.../scan/AlignedPageMemoryUsageEstimatorTest.java | 130 +++++++++++++++++++
4 files changed, 292 insertions(+), 95 deletions(-)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java
index f24b9d60264..c1ea72da8cf 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java
@@ -229,44 +229,120 @@ public class AlignedSinglePageWholeChunkReader extends AbstractChunkReader
public static List<Long> calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
final Chunk timeChunk, final List<Chunk> valueChunkList) throws IOException {
- final ByteBuffer timeChunkDataBuffer = timeChunk.getData().duplicate();
- final List<ByteBuffer> valueChunkDataBufferList = new ArrayList<>(valueChunkList.size());
+ final PageMemoryUsageEstimator estimator = new PageMemoryUsageEstimator(timeChunk);
for (final Chunk valueChunk : valueChunkList) {
- valueChunkDataBufferList.add(
- Objects.isNull(valueChunk) ? null : valueChunk.getData().duplicate());
+ estimator.addValueChunk(valueChunk);
}
+ return estimator.getPageEstimatedMemoryUsageInBytesList();
+ }
- final List<Long> pageEstimatedMemoryUsageInBytesList = new ArrayList<>();
- while (timeChunkDataBuffer.remaining() > 0) {
- long pageUncompressedSizeInBytes = 0;
- final PageHeader timePageHeader =
- SinglePageWholeChunkReader.deserializePageHeader(
- timeChunkDataBuffer, timeChunk.getHeader());
- pageUncompressedSizeInBytes += timePageHeader.getUncompressedSize();
- SinglePageWholeChunkReader.skipCompressedPageData(timeChunkDataBuffer, timePageHeader);
-
- final List<TSDataType> valueDataTypeList = new ArrayList<>(valueChunkList.size());
- for (int i = 0; i < valueChunkList.size(); ++i) {
- final Chunk valueChunk = valueChunkList.get(i);
- final ByteBuffer valueChunkDataBuffer = valueChunkDataBufferList.get(i);
- if (Objects.isNull(valueChunk) || Objects.isNull(valueChunkDataBuffer)) {
- valueDataTypeList.add(null);
- continue;
- }
+ static final class PageMemoryUsageEstimator {
- final PageHeader valuePageHeader =
+ private final List<Long> pageUncompressedSizeInBytesList = new ArrayList<>();
+ private final List<Integer> pageRowCountList = new ArrayList<>();
+ private List<Long> pageEstimatedMemoryUsageInBytesList = new ArrayList<>();
+
+ private int valueCount;
+ private TSDataType singleValueDataType;
+ private long vectorValueMemoryUsageInBytes;
+
+ PageMemoryUsageEstimator(final Chunk timeChunk) throws IOException {
+ final ByteBuffer timeChunkDataBuffer = timeChunk.getData().duplicate();
+ timeChunkDataBuffer.rewind();
+ while (timeChunkDataBuffer.remaining() > 0) {
+ final PageHeader timePageHeader =
SinglePageWholeChunkReader.deserializePageHeader(
- valueChunkDataBuffer, valueChunk.getHeader());
- pageUncompressedSizeInBytes += valuePageHeader.getUncompressedSize();
- valueDataTypeList.add(valueChunk.getHeader().getDataType());
- SinglePageWholeChunkReader.skipCompressedPageData(valueChunkDataBuffer, valuePageHeader);
+ timeChunkDataBuffer, timeChunk.getHeader());
+ pageUncompressedSizeInBytesList.add((long) timePageHeader.getUncompressedSize());
+ pageRowCountList.add(SinglePageWholeChunkReader.getPageRowCount(timePageHeader, timeChunk));
+ SinglePageWholeChunkReader.skipCompressedPageData(timeChunkDataBuffer, timePageHeader);
+ }
+ pageEstimatedMemoryUsageInBytesList = calculatePageEstimatedMemoryUsageInBytesList();
+ }
+
+ long calculateMaxPageEstimatedMemoryUsageInBytesWithValueChunk(final Chunk valueChunk)
+ throws IOException {
+ final List<Long> pageEstimatedMemoryUsageInBytesList =
+ calculatePageEstimatedMemoryUsageInBytesList(valueChunk, false);
+ return pageEstimatedMemoryUsageInBytesList.isEmpty()
+ ? 0
+ : pageEstimatedMemoryUsageInBytesList.get(0);
+ }
+
+ void addValueChunk(final Chunk valueChunk) throws IOException {
+ pageEstimatedMemoryUsageInBytesList =
+ calculatePageEstimatedMemoryUsageInBytesList(valueChunk, true);
+ }
+
+ List<Long> getPageEstimatedMemoryUsageInBytesList() {
+ return pageEstimatedMemoryUsageInBytesList;
+ }
+
+ private List<Long> calculatePageEstimatedMemoryUsageInBytesList() {
+ final List<Long> pageEstimatedMemoryUsageInBytesList =
+ new ArrayList<>(pageUncompressedSizeInBytesList.size());
+ for (int i = 0; i < pageUncompressedSizeInBytesList.size(); ++i) {
+ pageEstimatedMemoryUsageInBytesList.add(
+ SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
+ pageUncompressedSizeInBytesList.get(i),
+ pageRowCountList.get(i),
+ valueCount,
+ singleValueDataType,
+ vectorValueMemoryUsageInBytes));
}
- pageEstimatedMemoryUsageInBytesList.add(
- SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
- pageUncompressedSizeInBytes,
- SinglePageWholeChunkReader.getPageRowCount(timePageHeader, timeChunk),
- valueDataTypeList));
+ return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
+ }
+
+ private List<Long> calculatePageEstimatedMemoryUsageInBytesList(
+ final Chunk valueChunk, final boolean commit) throws IOException {
+ final ByteBuffer valueChunkDataBuffer;
+ final TSDataType valueDataType;
+ if (Objects.isNull(valueChunk)) {
+ valueChunkDataBuffer = null;
+ valueDataType = null;
+ } else {
+ valueChunkDataBuffer = valueChunk.getData().duplicate();
+ valueChunkDataBuffer.rewind();
+ valueDataType = valueChunk.getHeader().getDataType();
+ }
+
+ final int newValueCount = valueCount + 1;
+ final TSDataType newSingleValueDataType = valueCount == 0 ? valueDataType : null;
+ final long newVectorValueMemoryUsageInBytes =
+ vectorValueMemoryUsageInBytes
+ + SinglePageWholeChunkReader.estimateVectorValueMemoryUsageInBytes(valueDataType);
+ final List<Long> newPageUncompressedSizeInBytesList =
+ new ArrayList<>(pageUncompressedSizeInBytesList);
+ final List<Long> pageEstimatedMemoryUsageInBytesList =
+ new ArrayList<>(pageUncompressedSizeInBytesList.size());
+
+ for (int i = 0; i < newPageUncompressedSizeInBytesList.size(); ++i) {
+ if (Objects.nonNull(valueChunkDataBuffer)) {
+ final PageHeader valuePageHeader =
+ SinglePageWholeChunkReader.deserializePageHeader(
+ valueChunkDataBuffer, valueChunk.getHeader());
+ newPageUncompressedSizeInBytesList.set(
+ i, newPageUncompressedSizeInBytesList.get(i) + valuePageHeader.getUncompressedSize());
+ SinglePageWholeChunkReader.skipCompressedPageData(valueChunkDataBuffer, valuePageHeader);
+ }
+ pageEstimatedMemoryUsageInBytesList.add(
+ SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
+ newPageUncompressedSizeInBytesList.get(i),
+ pageRowCountList.get(i),
+ newValueCount,
+ newSingleValueDataType,
+ newVectorValueMemoryUsageInBytes));
+ }
+
+ if (commit) {
+ pageUncompressedSizeInBytesList.clear();
+ pageUncompressedSizeInBytesList.addAll(newPageUncompressedSizeInBytesList);
+ valueCount = newValueCount;
+ singleValueDataType = newSingleValueDataType;
+ vectorValueMemoryUsageInBytes = newVectorValueMemoryUsageInBytes;
+ }
+
+ return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
}
- return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java
index 4fdb6f4e4a5..7bf4889b20e 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java
@@ -191,6 +191,7 @@ public class SinglePageWholeChunkReader extends AbstractChunkReader
private static int countSinglePageTimeValues(final Chunk chunk) throws IOException {
final ByteBuffer chunkDataBuffer = chunk.getData().duplicate();
+ chunkDataBuffer.rewind();
final PageHeader pageHeader = deserializePageHeader(chunkDataBuffer, chunk.getHeader());
final ByteBuffer pageData =
deserializePageData(
@@ -222,25 +223,42 @@ public class SinglePageWholeChunkReader extends AbstractChunkReader
final long pageUncompressedSizeInBytes,
final int rowCount,
final List<TSDataType> valueDataTypeList) {
+ final int valueCount = valueDataTypeList.size();
+ return estimatePageMemoryUsageInBytesWithBatchData(
+ pageUncompressedSizeInBytes,
+ rowCount,
+ valueCount,
+ valueCount == 1 ? valueDataTypeList.get(0) : null,
+ estimateVectorValueMemoryUsageInBytes(valueDataTypeList));
+ }
+
+ static long estimatePageMemoryUsageInBytesWithBatchData(
+ final long pageUncompressedSizeInBytes,
+ final int rowCount,
+ final int valueCount,
+ final TSDataType singleValueDataType,
+ final long vectorValueMemoryUsageInBytes) {
return pageUncompressedSizeInBytes
- + estimateBatchDataMemoryUsageInBytes(rowCount, valueDataTypeList);
+ + estimateBatchDataMemoryUsageInBytes(
+ rowCount, valueCount, singleValueDataType, vectorValueMemoryUsageInBytes);
}
private static long estimateBatchDataMemoryUsageInBytes(
- final int rowCount, final List<TSDataType> valueDataTypeList) {
- final int valueCount = valueDataTypeList.size();
+ final int rowCount,
+ final int valueCount,
+ final TSDataType singleValueDataType,
+ final long vectorValueMemoryUsageInBytes) {
final long segmentCount = Math.max(1, (rowCount + 15L) / 16);
long estimatedMemoryUsageInBytes = RamUsageEstimator.sizeOfLongArray(16) * segmentCount;
if (valueCount == 1) {
estimatedMemoryUsageInBytes +=
- estimateSingleValueArrayMemoryUsageInBytes(rowCount, valueDataTypeList.get(0));
+ estimateSingleValueArrayMemoryUsageInBytes(rowCount, singleValueDataType);
} else if (valueCount > 1) {
estimatedMemoryUsageInBytes += RamUsageEstimator.sizeOfObjectArray(16) * segmentCount;
estimatedMemoryUsageInBytes +=
(long) rowCount
- * (RamUsageEstimator.sizeOfObjectArray(valueCount)
- + estimateVectorValueMemoryUsageInBytes(valueDataTypeList));
+ * (RamUsageEstimator.sizeOfObjectArray(valueCount) + vectorValueMemoryUsageInBytes);
}
return estimatedMemoryUsageInBytes;
@@ -260,18 +278,21 @@ public class SinglePageWholeChunkReader extends AbstractChunkReader
final List<TSDataType> valueDataTypeList) {
long estimatedMemoryUsageInBytes = 0;
for (final TSDataType dataType : valueDataTypeList) {
- if (Objects.isNull(dataType)) {
- continue;
- }
-
- estimatedMemoryUsageInBytes +=
- RamUsageEstimator.alignObjectSize(
- RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
- + estimateTsPrimitiveTypeValueMemoryUsageInBytes(dataType));
+ estimatedMemoryUsageInBytes += estimateVectorValueMemoryUsageInBytes(dataType);
}
return estimatedMemoryUsageInBytes;
}
+ static long estimateVectorValueMemoryUsageInBytes(final TSDataType dataType) {
+ if (Objects.isNull(dataType)) {
+ return 0;
+ }
+
+ return RamUsageEstimator.alignObjectSize(
+ RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
+ + estimateTsPrimitiveTypeValueMemoryUsageInBytes(dataType));
+ }
+
private static long estimateTsPrimitiveTypeValueMemoryUsageInBytes(final TSDataType dataType) {
return Type.fromTsDataType(dataType).estimateValueSize();
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java
index 2ce36709f65..c368b0087a6 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java
@@ -74,7 +74,6 @@ import org.apache.tsfile.write.schema.MeasurementSchema;
import java.io.File;
import java.io.IOException;
-import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -893,9 +892,8 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
}
final List<Long> pageEstimatedMemoryUsageInBytesList =
nextIsMultiPage
- ? AlignedSinglePageWholeChunkReader
- .calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
- timeChunk, pendingAlignedChunkGroup.valueChunkList)
+ ? pendingAlignedChunkGroup.pageMemoryUsageEstimator
+ .getPageEstimatedMemoryUsageInBytesList()
: Collections.emptyList();
final long maxPageEstimatedMemoryUsageInBytes =
pageEstimatedMemoryUsageInBytesList.isEmpty()
@@ -996,6 +994,7 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
resizeChunkMemoryBlockIfFirstValueChunkExceedsLimit(pendingAlignedChunkGroup, valueChunk);
resizePageDataMemoryBlockIfFirstValueChunkExceedsLimit(pendingAlignedChunkGroup, valueChunk);
+ pendingAlignedChunkGroup.pageMemoryUsageEstimator.addValueChunk(valueChunk.chunk);
pendingAlignedChunkGroup.valueChunkList.add(valueChunk.chunk);
pendingAlignedChunkGroup.chunkSize += valueChunk.valueChunkSize;
pendingAlignedChunkSize += valueChunk.valueChunkSize;
@@ -1013,7 +1012,8 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
currentDevice, Collections.singletonList(measurementID), currentModifications, false));
}
- private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int timeChunkIndex) {
+ private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int timeChunkIndex)
+ throws IOException {
final PendingAlignedChunkGroup pendingAlignedChunkGroup =
findPendingAlignedChunkGroup(timeChunkIndex);
if (Objects.nonNull(pendingAlignedChunkGroup)) {
@@ -1023,7 +1023,8 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
final PendingAlignedChunkGroup newPendingAlignedChunkGroup =
new PendingAlignedChunkGroup(
timeChunkIndex,
- PipeMemoryWeightUtil.calculateChunkRamBytesUsed(timeChunkList.get(timeChunkIndex)));
+ PipeMemoryWeightUtil.calculateChunkRamBytesUsed(timeChunkList.get(timeChunkIndex)),
+ timeChunkList.get(timeChunkIndex));
pendingAlignedChunkSize += newPendingAlignedChunkGroup.chunkSize;
for (int i = 0; i < pendingAlignedChunkGroups.size(); ++i) {
@@ -1099,49 +1100,12 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
final PendingAlignedChunkGroup pendingAlignedChunkGroup,
final CachedAlignedValueChunk valueChunk)
throws IOException {
- final List<Chunk> valueChunkList =
- new ArrayList<>(
- (Objects.isNull(pendingAlignedChunkGroup)
- ? 0
- : pendingAlignedChunkGroup.valueChunkList.size())
- + 1);
- if (Objects.nonNull(pendingAlignedChunkGroup)) {
- valueChunkList.addAll(pendingAlignedChunkGroup.valueChunkList);
- }
- valueChunkList.add(valueChunk.chunk);
-
final Chunk timeChunk = timeChunkList.get(timeChunkIndex);
- final int timeChunkDataPosition = timeChunk.getData().position();
- final List<Integer> valueChunkDataPositions = new ArrayList<>(valueChunkList.size());
- for (final Chunk chunk : valueChunkList) {
- valueChunkDataPositions.add(Objects.isNull(chunk) ? 0 : chunk.getData().position());
- }
-
- rewindChunkData(timeChunk);
- valueChunkList.forEach(this::rewindChunkData);
- try {
- return AlignedSinglePageWholeChunkReader
- .calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(timeChunk, valueChunkList);
- } finally {
- timeChunk.getData().position(timeChunkDataPosition);
- for (int i = 0; i < valueChunkList.size(); ++i) {
- final Chunk chunk = valueChunkList.get(i);
- if (Objects.nonNull(chunk)) {
- chunk.getData().position(valueChunkDataPositions.get(i));
- }
- }
- }
- }
-
- private void rewindChunkData(final Chunk chunk) {
- if (Objects.isNull(chunk)) {
- return;
- }
-
- final ByteBuffer chunkData = chunk.getData();
- if (Objects.nonNull(chunkData)) {
- chunkData.rewind();
- }
+ final AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator estimator =
+ Objects.isNull(pendingAlignedChunkGroup)
+ ? new AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator(timeChunk)
+ : pendingAlignedChunkGroup.pageMemoryUsageEstimator;
+ return estimator.calculateMaxPageEstimatedMemoryUsageInBytesWithValueChunk(valueChunk.chunk);
}
private boolean isSinglePageValueChunk(final ChunkHeader chunkHeader) {
@@ -1207,14 +1171,20 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
private static class PendingAlignedChunkGroup {
private final int timeChunkIndex;
+ private final AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator
+ pageMemoryUsageEstimator;
private final List<Chunk> valueChunkList = new ArrayList<>();
private final List<IMeasurementSchema> measurements = new ArrayList<>();
private final List<ModsOperationUtil.ModsInfo> modsInfos = new ArrayList<>();
private long chunkSize;
- private PendingAlignedChunkGroup(final int timeChunkIndex, final long timeChunkSize) {
+ private PendingAlignedChunkGroup(
+ final int timeChunkIndex, final long timeChunkSize, final Chunk timeChunk)
+ throws IOException {
this.timeChunkIndex = timeChunkIndex;
this.chunkSize = timeChunkSize;
+ this.pageMemoryUsageEstimator =
+ new AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator(timeChunk);
}
}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedPageMemoryUsageEstimatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedPageMemoryUsageEstimatorTest.java
new file mode 100644
index 00000000000..22ba15742fe
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedPageMemoryUsageEstimatorTest.java
@@ -0,0 +1,130 @@
+/*
+ * 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.pipe.event.common.tsfile.parser.scan;
+
+import org.apache.iotdb.commons.path.PartialPath;
+
+import org.apache.tsfile.common.conf.TSFileDescriptor;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.MetaMarker;
+import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata;
+import org.apache.tsfile.file.metadata.ChunkMetadata;
+import org.apache.tsfile.file.metadata.IChunkMetadata;
+import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.read.TsFileSequenceReader;
+import org.apache.tsfile.read.common.Chunk;
+import org.apache.tsfile.write.TsFileWriter;
+import org.apache.tsfile.write.record.Tablet;
+import org.apache.tsfile.write.schema.IMeasurementSchema;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+
+public class AlignedPageMemoryUsageEstimatorTest {
+
+ @Test
+ public void testIncrementalEstimatesMatchFullEstimates() throws Exception {
+ final int originalMaxNumberOfPointsInPage =
+ TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage();
+ final File tsFile = new File("aligned-page-memory-estimator.tsfile");
+
+ try {
+ TSFileDescriptor.getInstance().getConfig().setMaxNumberOfPointsInPage(2);
+ if (tsFile.exists()) {
+ Assert.assertTrue(tsFile.delete());
+ }
+
+ final List<IMeasurementSchema> schemas = new ArrayList<>();
+ for (int i = 0; i < 4; ++i) {
+ schemas.add(
+ new MeasurementSchema(
+ "s" + i, TSDataType.INT64, TSEncoding.PLAIN, CompressionType.LZ4));
+ }
+ final Tablet tablet = new Tablet("root.sg.d", schemas, 8);
+ for (int row = 0; row < 8; ++row) {
+ tablet.addTimestamp(row, row);
+ for (int measurement = 0; measurement < schemas.size(); ++measurement) {
+ tablet.addValue("s" + measurement, row, (long) row + measurement);
+ }
+ }
+
+ try (final TsFileWriter writer = new TsFileWriter(tsFile)) {
+ writer.registerAlignedTimeseries(new PartialPath("root.sg.d"), schemas);
+ writer.writeAligned(tablet);
+ }
+
+ try (final TsFileSequenceReader reader = new TsFileSequenceReader(tsFile.getAbsolutePath())) {
+ final IDeviceID deviceID = reader.getDeviceMeasurementsMap().keySet().iterator().next();
+ final List<AbstractAlignedChunkMetadata> alignedChunkMetadataList =
+ reader.getAlignedChunkMetadata(deviceID, true);
+ Assert.assertEquals(1, alignedChunkMetadataList.size());
+ final AbstractAlignedChunkMetadata alignedChunkMetadata = alignedChunkMetadataList.get(0);
+ final Chunk timeChunk =
+ reader.readMemChunk((ChunkMetadata) alignedChunkMetadata.getTimeChunkMetadata());
+ Assert.assertEquals(MetaMarker.CHUNK_HEADER, timeChunk.getHeader().getChunkType() & 0x3F);
+
+ final List<Chunk> prefix = new ArrayList<>();
+ final AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator estimator =
+ new AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator(timeChunk);
+ for (final IChunkMetadata valueChunkMetadata :
+ alignedChunkMetadata.getValueChunkMetadataList()) {
+ final Chunk valueChunk = reader.readMemChunk((ChunkMetadata) valueChunkMetadata);
+ Assert.assertEquals(
+ MetaMarker.CHUNK_HEADER, valueChunk.getHeader().getChunkType() & 0x3F);
+
+ final long expectedWithCandidate =
+ AlignedSinglePageWholeChunkReader
+ .calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(
+ timeChunk, appendAndCopy(prefix, valueChunk));
+ Assert.assertEquals(
+ expectedWithCandidate,
+ estimator.calculateMaxPageEstimatedMemoryUsageInBytesWithValueChunk(valueChunk));
+
+ estimator.addValueChunk(valueChunk);
+ prefix.add(valueChunk);
+ final List<Long> expectedPageMemoryUsage =
+ AlignedSinglePageWholeChunkReader
+ .calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(timeChunk, prefix);
+ Assert.assertEquals(
+ expectedPageMemoryUsage, estimator.getPageEstimatedMemoryUsageInBytesList());
+ }
+ }
+ } finally {
+ TSFileDescriptor.getInstance()
+ .getConfig()
+ .setMaxNumberOfPointsInPage(originalMaxNumberOfPointsInPage);
+ if (tsFile.exists()) {
+ Assert.assertTrue(tsFile.delete());
+ }
+ }
+ }
+
+ private List<Chunk> appendAndCopy(final List<Chunk> chunks, final Chunk chunk) {
+ final List<Chunk> copy = new ArrayList<>(chunks);
+ copy.add(chunk);
+ return copy;
+ }
+}