diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/AlignedSinglePageWholeChunkReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/AlignedSinglePageWholeChunkReader.java index 9c3b3514c830e..8cf21b144d9d2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/AlignedSinglePageWholeChunkReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/AlignedSinglePageWholeChunkReader.java @@ -219,44 +219,120 @@ public static long calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData( public static List calculatePageEstimatedMemoryUsageInBytesWithBatchDataList( final Chunk timeChunk, final List valueChunkList) throws IOException { - final ByteBuffer timeChunkDataBuffer = timeChunk.getData().duplicate(); - final List 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 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 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 pageUncompressedSizeInBytesList = new ArrayList<>(); + private final List pageRowCountList = new ArrayList<>(); + private List 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 pageEstimatedMemoryUsageInBytesList = + calculatePageEstimatedMemoryUsageInBytesList(valueChunk, false); + return pageEstimatedMemoryUsageInBytesList.isEmpty() + ? 0 + : pageEstimatedMemoryUsageInBytesList.get(0); + } + + void addValueChunk(final Chunk valueChunk) throws IOException { + pageEstimatedMemoryUsageInBytesList = + calculatePageEstimatedMemoryUsageInBytesList(valueChunk, true); + } + + List getPageEstimatedMemoryUsageInBytesList() { + return pageEstimatedMemoryUsageInBytesList; + } + + private List calculatePageEstimatedMemoryUsageInBytesList() { + final List 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 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 newPageUncompressedSizeInBytesList = + new ArrayList<>(pageUncompressedSizeInBytesList); + final List 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/container/scan/SinglePageWholeChunkReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/SinglePageWholeChunkReader.java index 1f741ddfc70fd..c2d4630badd42 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/SinglePageWholeChunkReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/SinglePageWholeChunkReader.java @@ -180,6 +180,7 @@ static int getPageRowCount(final PageHeader pageHeader, final Chunk chunk) throw 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(pageHeader, chunkDataBuffer, chunk.getHeader()); final Decoder decoder = @@ -206,25 +207,42 @@ static long estimatePageMemoryUsageInBytesWithBatchData( final long pageUncompressedSizeInBytes, final int rowCount, final List 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 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; @@ -263,18 +281,21 @@ private static long estimateVectorValueMemoryUsageInBytes( final List 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) { switch (dataType) { case BOOLEAN: diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/TsFileInsertionScanDataContainer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/TsFileInsertionScanDataContainer.java index 6572c92eb9a70..fff7576e95f7b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/TsFileInsertionScanDataContainer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/TsFileInsertionScanDataContainer.java @@ -67,7 +67,6 @@ import java.io.File; import java.io.IOException; -import java.nio.ByteBuffer; import java.time.LocalDate; import java.util.ArrayList; import java.util.Collections; @@ -792,11 +791,12 @@ private boolean filterChunk( private boolean useNextPendingAlignedChunk(final byte marker) throws IOException { while (!pendingAlignedChunkGroups.isEmpty()) { - final PendingAlignedChunkGroup pendingAlignedChunkGroup = pendingAlignedChunkGroups.remove(0); - pendingAlignedChunkSize = - Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize); + final PendingAlignedChunkGroup pendingAlignedChunkGroup = pendingAlignedChunkGroups.get(0); if (pendingAlignedChunkGroup.valueChunkList.isEmpty()) { + pendingAlignedChunkGroups.remove(0); + pendingAlignedChunkSize = + Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize); continue; } @@ -806,36 +806,42 @@ private boolean useNextPendingAlignedChunk(final byte marker) throws IOException valueChunk.getData().rewind(); } - currentMeasurements.clear(); - currentMeasurements.addAll(pendingAlignedChunkGroup.measurements); - modsInfos.clear(); - modsInfos.addAll(pendingAlignedChunkGroup.modsInfos); - - currentIsMultiPage = isMultiPageList.get(pendingAlignedChunkGroup.timeChunkIndex); - if (!currentIsMultiPage) { + final boolean nextIsMultiPage = isMultiPageList.get(pendingAlignedChunkGroup.timeChunkIndex); + if (!nextIsMultiPage) { resizePageDataMemoryIfNeeded( AlignedSinglePageWholeChunkReader.calculatePageEstimatedMemoryUsageInBytes( timeChunk, pendingAlignedChunkGroup.valueChunkList)); } final List pageEstimatedMemoryUsageInBytesList = - currentIsMultiPage - ? AlignedSinglePageWholeChunkReader - .calculatePageEstimatedMemoryUsageInBytesWithBatchDataList( - timeChunk, pendingAlignedChunkGroup.valueChunkList) + nextIsMultiPage + ? pendingAlignedChunkGroup.pageMemoryUsageEstimator + .getPageEstimatedMemoryUsageInBytesList() : Collections.emptyList(); final long maxPageEstimatedMemoryUsageInBytes = pageEstimatedMemoryUsageInBytesList.isEmpty() ? 0 : pageEstimatedMemoryUsageInBytesList.get(0); resizePageDataMemoryIfNeeded(maxPageEstimatedMemoryUsageInBytes); - chunkReader = - currentIsMultiPage + final IChunkReader nextChunkReader = + nextIsMultiPage ? new MemoryControlledChunkReader( new AlignedChunkReader( timeChunk, pendingAlignedChunkGroup.valueChunkList, filter), pageEstimatedMemoryUsageInBytesList) : new AlignedSinglePageWholeChunkReader( timeChunk, pendingAlignedChunkGroup.valueChunkList); + + // Publish the state transition only after memory reservations and reader construction + // succeed. + pendingAlignedChunkGroups.remove(0); + pendingAlignedChunkSize = + Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize); + currentMeasurements.clear(); + currentMeasurements.addAll(pendingAlignedChunkGroup.measurements); + modsInfos.clear(); + modsInfos.addAll(pendingAlignedChunkGroup.modsInfos); + currentIsMultiPage = nextIsMultiPage; + chunkReader = nextChunkReader; currentIsAligned = true; if (marker != Byte.MIN_VALUE) { lastMarker = marker; @@ -906,6 +912,7 @@ private void cacheAlignedValueChunk(final CachedAlignedValueChunk valueChunk) th resizeChunkMemoryBlockIfFirstValueChunkExceedsLimit(pendingAlignedChunkGroup, valueChunk); resizePageDataMemoryBlockIfFirstValueChunkExceedsLimit(pendingAlignedChunkGroup, valueChunk); + pendingAlignedChunkGroup.pageMemoryUsageEstimator.addValueChunk(valueChunk.chunk); pendingAlignedChunkGroup.valueChunkList.add(valueChunk.chunk); pendingAlignedChunkGroup.chunkSize += valueChunk.valueChunkSize; pendingAlignedChunkSize += valueChunk.valueChunkSize; @@ -919,7 +926,8 @@ private void cacheAlignedValueChunk(final CachedAlignedValueChunk valueChunk) th currentDevice, Collections.singletonList(measurementID), currentModifications)); } - private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int timeChunkIndex) { + private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int timeChunkIndex) + throws IOException { final PendingAlignedChunkGroup pendingAlignedChunkGroup = findPendingAlignedChunkGroup(timeChunkIndex); if (Objects.nonNull(pendingAlignedChunkGroup)) { @@ -929,7 +937,8 @@ private PendingAlignedChunkGroup getOrCreatePendingAlignedChunkGroup(final int t 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) { @@ -1003,49 +1012,12 @@ private long calculateMaxAlignedPageMemorySizeWithBatchData( final PendingAlignedChunkGroup pendingAlignedChunkGroup, final CachedAlignedValueChunk valueChunk) throws IOException { - final List 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 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) { @@ -1102,14 +1074,20 @@ private Statistics findNonAlignedChunkStatistics( private static class PendingAlignedChunkGroup { private final int timeChunkIndex; + private final AlignedSinglePageWholeChunkReader.PageMemoryUsageEstimator + pageMemoryUsageEstimator; private final List valueChunkList = new ArrayList<>(); private final List measurements = new ArrayList<>(); private final List 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/container/scan/AlignedPageMemoryUsageEstimatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/AlignedPageMemoryUsageEstimatorTest.java new file mode 100644 index 0000000000000..48d90037db97c --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/container/scan/AlignedPageMemoryUsageEstimatorTest.java @@ -0,0 +1,129 @@ +/* + * 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.container.scan; + +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.AlignedChunkMetadata; +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.read.common.Path; +import org.apache.tsfile.write.TsFileWriter; +import org.apache.tsfile.write.record.Tablet; +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 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); + } + } + tablet.rowSize = 8; + + try (final TsFileWriter writer = new TsFileWriter(tsFile)) { + writer.registerAlignedTimeseries(new Path("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 alignedChunkMetadataList = + reader.getAlignedChunkMetadata(deviceID); + Assert.assertEquals(1, alignedChunkMetadataList.size()); + final AlignedChunkMetadata alignedChunkMetadata = alignedChunkMetadataList.get(0); + final Chunk timeChunk = + reader.readMemChunk((ChunkMetadata) alignedChunkMetadata.getTimeChunkMetadata()); + Assert.assertEquals(MetaMarker.CHUNK_HEADER, timeChunk.getHeader().getChunkType() & 0x3F); + + final List 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 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 appendAndCopy(final List chunks, final Chunk chunk) { + final List copy = new ArrayList<>(chunks); + copy.add(chunk); + return copy; + } +}