Skip to content

Commit 6bc95b1

Browse files
CaideyipiJackieTien97
Caideyipi
authored and
JackieTien97
committed
Pipe: account page decode memory in scan parser (#17807)
* Pipe: account page decode memory in scan parser * Fix pipe scan parser single page row count * Fix pipe scan parser page memory test
1 parent 7128b9f commit 6bc95b1

5 files changed

Lines changed: 568 additions & 24 deletions

File tree

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

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import java.nio.ByteBuffer;
4040
import java.util.ArrayList;
4141
import java.util.List;
42+
import java.util.Objects;
4243
import java.util.function.LongConsumer;
4344

4445
/**
@@ -71,7 +72,7 @@ public AlignedSinglePageWholeChunkReader(
7172
this.timeChunkDataBuffer = timeChunk.getData();
7273
this.encryptParam = timeChunk.getEncryptParam();
7374
this.pageEstimatedMemoryUsageInBytes =
74-
calculatePageEstimatedMemoryUsageInBytes(timeChunk, valueChunkList);
75+
calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(timeChunk, valueChunkList);
7576

7677
valueChunkList.forEach(
7778
chunk -> {
@@ -216,4 +217,56 @@ public static long calculatePageEstimatedMemoryUsageInBytes(
216217

217218
return estimatedMemoryUsageInBytes;
218219
}
220+
221+
public static long calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(
222+
final Chunk timeChunk, final List<Chunk> valueChunkList) throws IOException {
223+
final List<Long> pageEstimatedMemoryUsageInBytesList =
224+
calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(timeChunk, valueChunkList);
225+
return pageEstimatedMemoryUsageInBytesList.isEmpty()
226+
? 0
227+
: pageEstimatedMemoryUsageInBytesList.get(0);
228+
}
229+
230+
public static List<Long> calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
231+
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());
234+
for (final Chunk valueChunk : valueChunkList) {
235+
valueChunkDataBufferList.add(
236+
Objects.isNull(valueChunk) ? null : valueChunk.getData().duplicate());
237+
}
238+
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+
}
256+
257+
final PageHeader valuePageHeader =
258+
SinglePageWholeChunkReader.deserializePageHeader(
259+
valueChunkDataBuffer, valueChunk.getHeader());
260+
pageUncompressedSizeInBytes += valuePageHeader.getUncompressedSize();
261+
valueDataTypeList.add(valueChunk.getHeader().getDataType());
262+
SinglePageWholeChunkReader.skipCompressedPageData(valueChunkDataBuffer, valuePageHeader);
263+
}
264+
pageEstimatedMemoryUsageInBytesList.add(
265+
SinglePageWholeChunkReader.estimatePageMemoryUsageInBytesWithBatchData(
266+
pageUncompressedSizeInBytes,
267+
SinglePageWholeChunkReader.getPageRowCount(timePageHeader, timeChunk),
268+
valueDataTypeList));
269+
}
270+
return SinglePageWholeChunkReader.toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
271+
}
219272
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan;
21+
22+
import org.apache.tsfile.read.common.BatchData;
23+
import org.apache.tsfile.read.reader.IChunkReader;
24+
import org.apache.tsfile.read.reader.IPageReader;
25+
26+
import java.io.IOException;
27+
import java.util.List;
28+
29+
class MemoryControlledChunkReader implements IChunkReader, EstimatedMemoryChunkReader {
30+
31+
private final IChunkReader delegate;
32+
private final List<Long> pageEstimatedMemoryUsageInBytesList;
33+
private int pageIndex;
34+
35+
MemoryControlledChunkReader(
36+
final IChunkReader delegate, final List<Long> pageEstimatedMemoryUsageInBytesList) {
37+
this.delegate = delegate;
38+
this.pageEstimatedMemoryUsageInBytesList = pageEstimatedMemoryUsageInBytesList;
39+
}
40+
41+
@Override
42+
public long getCurrentPageEstimatedMemoryUsageInBytes() {
43+
return pageIndex < pageEstimatedMemoryUsageInBytesList.size()
44+
? pageEstimatedMemoryUsageInBytesList.get(pageIndex)
45+
: 0;
46+
}
47+
48+
@Override
49+
public boolean hasNextSatisfiedPage() throws IOException {
50+
return delegate.hasNextSatisfiedPage();
51+
}
52+
53+
@Override
54+
public BatchData nextPageData() throws IOException {
55+
try {
56+
return delegate.nextPageData();
57+
} finally {
58+
++pageIndex;
59+
}
60+
}
61+
62+
@Override
63+
public void close() throws IOException {
64+
delegate.close();
65+
}
66+
67+
@Override
68+
public void markDataTypeModifiedAndCannotUseStatistics() {
69+
delegate.markDataTypeModifiedAndCannotUseStatistics();
70+
}
71+
72+
@Override
73+
public List<IPageReader> loadPageReaderList() throws IOException {
74+
return delegate.loadPageReaderList();
75+
}
76+
}

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

Lines changed: 190 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import org.apache.tsfile.encoding.decoder.Decoder;
2626
import org.apache.tsfile.encrypt.EncryptParameter;
2727
import org.apache.tsfile.encrypt.IDecryptor;
28+
import org.apache.tsfile.enums.TSDataType;
29+
import org.apache.tsfile.file.MetaMarker;
2830
import org.apache.tsfile.file.header.ChunkHeader;
2931
import org.apache.tsfile.file.header.PageHeader;
3032
import org.apache.tsfile.file.metadata.enums.EncryptionType;
@@ -33,10 +35,15 @@
3335
import org.apache.tsfile.read.reader.chunk.AbstractChunkReader;
3436
import org.apache.tsfile.read.reader.page.LazyLoadPageData;
3537
import org.apache.tsfile.read.reader.page.PageReader;
38+
import org.apache.tsfile.utils.RamUsageEstimator;
3639

3740
import java.io.IOException;
3841
import java.io.Serializable;
3942
import java.nio.ByteBuffer;
43+
import java.util.ArrayList;
44+
import java.util.Collections;
45+
import java.util.List;
46+
import java.util.Objects;
4047

4148
import static org.apache.tsfile.file.metadata.enums.CompressionType.UNCOMPRESSED;
4249

@@ -53,7 +60,8 @@ public SinglePageWholeChunkReader(Chunk chunk) throws IOException {
5360
this.chunkHeader = chunk.getHeader();
5461
this.chunkDataBuffer = chunk.getData();
5562
this.encryptParam = chunk.getEncryptParam();
56-
this.pageEstimatedMemoryUsageInBytes = calculatePageEstimatedMemoryUsageInBytes(chunk);
63+
this.pageEstimatedMemoryUsageInBytes =
64+
calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(chunk);
5765
initAllPageReaders();
5866
}
5967

@@ -91,11 +99,190 @@ public long getCurrentPageEstimatedMemoryUsageInBytes() {
9199
public static long calculatePageEstimatedMemoryUsageInBytes(final Chunk chunk)
92100
throws IOException {
93101
final ByteBuffer chunkDataBuffer = chunk.getData().duplicate();
94-
final PageHeader pageHeader =
95-
PageHeader.deserializeFrom(chunkDataBuffer, (Statistics<? extends Serializable>) null);
102+
final PageHeader pageHeader = deserializePageHeader(chunkDataBuffer, chunk.getHeader());
96103
return pageHeader.getUncompressedSize();
97104
}
98105

106+
public static long calculateMaxPageEstimatedMemoryUsageInBytes(final Chunk chunk)
107+
throws IOException {
108+
final ByteBuffer chunkDataBuffer = chunk.getData().duplicate();
109+
long maxPageEstimatedMemoryUsageInBytes = 0;
110+
while (chunkDataBuffer.remaining() > 0) {
111+
final PageHeader pageHeader = deserializePageHeader(chunkDataBuffer, chunk.getHeader());
112+
maxPageEstimatedMemoryUsageInBytes =
113+
Math.max(maxPageEstimatedMemoryUsageInBytes, pageHeader.getUncompressedSize());
114+
skipCompressedPageData(chunkDataBuffer, pageHeader);
115+
}
116+
return maxPageEstimatedMemoryUsageInBytes;
117+
}
118+
119+
public static long calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(final Chunk chunk)
120+
throws IOException {
121+
final List<Long> pageEstimatedMemoryUsageInBytesList =
122+
calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(chunk);
123+
return pageEstimatedMemoryUsageInBytesList.isEmpty()
124+
? 0
125+
: pageEstimatedMemoryUsageInBytesList.get(0);
126+
}
127+
128+
public static List<Long> calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
129+
final Chunk chunk) throws IOException {
130+
final ByteBuffer chunkDataBuffer = chunk.getData().duplicate();
131+
final List<Long> pageEstimatedMemoryUsageInBytesList = new ArrayList<>();
132+
while (chunkDataBuffer.remaining() > 0) {
133+
final PageHeader pageHeader = deserializePageHeader(chunkDataBuffer, chunk.getHeader());
134+
pageEstimatedMemoryUsageInBytesList.add(
135+
estimatePageMemoryUsageInBytesWithBatchData(
136+
pageHeader, chunk, Collections.singletonList(chunk.getHeader().getDataType())));
137+
skipCompressedPageData(chunkDataBuffer, pageHeader);
138+
}
139+
return toSuffixMaxList(pageEstimatedMemoryUsageInBytesList);
140+
}
141+
142+
static PageHeader deserializePageHeader(
143+
final ByteBuffer chunkDataBuffer, final ChunkHeader chunkHeader) throws IOException {
144+
return isSinglePageChunk(chunkHeader)
145+
? PageHeader.deserializeFrom(chunkDataBuffer, (Statistics<? extends Serializable>) null)
146+
: PageHeader.deserializeFrom(chunkDataBuffer, chunkHeader.getDataType());
147+
}
148+
149+
static boolean isSinglePageChunk(final ChunkHeader chunkHeader) {
150+
return (chunkHeader.getChunkType() & 0x3F) == MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER;
151+
}
152+
153+
static void skipCompressedPageData(
154+
final ByteBuffer chunkDataBuffer, final PageHeader pageHeader) {
155+
chunkDataBuffer.position(chunkDataBuffer.position() + pageHeader.getCompressedSize());
156+
}
157+
158+
static List<Long> toSuffixMaxList(final List<Long> pageEstimatedMemoryUsageInBytesList) {
159+
long suffixMaxPageEstimatedMemoryUsageInBytes = 0;
160+
for (int i = pageEstimatedMemoryUsageInBytesList.size() - 1; i >= 0; --i) {
161+
suffixMaxPageEstimatedMemoryUsageInBytes =
162+
Math.max(
163+
suffixMaxPageEstimatedMemoryUsageInBytes, pageEstimatedMemoryUsageInBytesList.get(i));
164+
pageEstimatedMemoryUsageInBytesList.set(i, suffixMaxPageEstimatedMemoryUsageInBytes);
165+
}
166+
return pageEstimatedMemoryUsageInBytesList;
167+
}
168+
169+
static long estimatePageMemoryUsageInBytesWithBatchData(
170+
final PageHeader timePageHeader,
171+
final Chunk timeChunk,
172+
final List<TSDataType> valueDataTypeList) {
173+
return estimatePageMemoryUsageInBytesWithBatchData(
174+
timePageHeader.getUncompressedSize(),
175+
getPageRowCount(timePageHeader, timeChunk),
176+
valueDataTypeList);
177+
}
178+
179+
static int getPageRowCount(final PageHeader pageHeader, final Chunk chunk) {
180+
if (isSinglePageChunk(chunk.getHeader())) {
181+
return Objects.isNull(chunk.getChunkStatistic())
182+
? 0
183+
: saturateToInt(chunk.getChunkStatistic().getCount());
184+
}
185+
return saturateToInt(pageHeader.getNumOfValues());
186+
}
187+
188+
private static int saturateToInt(final long value) {
189+
return (int) Math.min(Integer.MAX_VALUE, value);
190+
}
191+
192+
static long estimatePageMemoryUsageInBytesWithBatchData(
193+
final long pageUncompressedSizeInBytes,
194+
final int rowCount,
195+
final List<TSDataType> valueDataTypeList) {
196+
return pageUncompressedSizeInBytes
197+
+ estimateBatchDataMemoryUsageInBytes(rowCount, valueDataTypeList);
198+
}
199+
200+
private static long estimateBatchDataMemoryUsageInBytes(
201+
final int rowCount, final List<TSDataType> valueDataTypeList) {
202+
final int valueCount = valueDataTypeList.size();
203+
final long segmentCount = Math.max(1, (rowCount + 15L) / 16);
204+
long estimatedMemoryUsageInBytes = RamUsageEstimator.sizeOfLongArray(16) * segmentCount;
205+
206+
if (valueCount == 1) {
207+
estimatedMemoryUsageInBytes +=
208+
estimateSingleValueArrayMemoryUsageInBytes(rowCount, valueDataTypeList.get(0));
209+
} else if (valueCount > 1) {
210+
estimatedMemoryUsageInBytes += RamUsageEstimator.sizeOfObjectArray(16) * segmentCount;
211+
estimatedMemoryUsageInBytes +=
212+
(long) rowCount
213+
* (RamUsageEstimator.sizeOfObjectArray(valueCount)
214+
+ estimateVectorValueMemoryUsageInBytes(valueDataTypeList));
215+
}
216+
217+
return estimatedMemoryUsageInBytes;
218+
}
219+
220+
private static long estimateSingleValueArrayMemoryUsageInBytes(
221+
final int rowCount, final TSDataType dataType) {
222+
final long segmentCount = Math.max(1, (rowCount + 15L) / 16);
223+
if (Objects.isNull(dataType)) {
224+
return 0;
225+
}
226+
227+
switch (dataType) {
228+
case BOOLEAN:
229+
return RamUsageEstimator.sizeOfBooleanArray(16) * segmentCount;
230+
case INT32:
231+
case DATE:
232+
return RamUsageEstimator.sizeOfIntArray(16) * segmentCount;
233+
case INT64:
234+
case TIMESTAMP:
235+
return RamUsageEstimator.sizeOfLongArray(16) * segmentCount;
236+
case FLOAT:
237+
return RamUsageEstimator.sizeOfFloatArray(16) * segmentCount;
238+
case DOUBLE:
239+
return RamUsageEstimator.sizeOfDoubleArray(16) * segmentCount;
240+
case TEXT:
241+
case BLOB:
242+
case STRING:
243+
return RamUsageEstimator.sizeOfObjectArray(16) * segmentCount;
244+
default:
245+
return 0;
246+
}
247+
}
248+
249+
private static long estimateVectorValueMemoryUsageInBytes(
250+
final List<TSDataType> valueDataTypeList) {
251+
long estimatedMemoryUsageInBytes = 0;
252+
for (final TSDataType dataType : valueDataTypeList) {
253+
if (Objects.isNull(dataType)) {
254+
continue;
255+
}
256+
257+
estimatedMemoryUsageInBytes +=
258+
RamUsageEstimator.alignObjectSize(
259+
RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
260+
+ estimateTsPrimitiveTypeValueMemoryUsageInBytes(dataType));
261+
}
262+
return estimatedMemoryUsageInBytes;
263+
}
264+
265+
private static long estimateTsPrimitiveTypeValueMemoryUsageInBytes(final TSDataType dataType) {
266+
switch (dataType) {
267+
case BOOLEAN:
268+
return 1;
269+
case INT32:
270+
case DATE:
271+
case FLOAT:
272+
return Integer.BYTES;
273+
case INT64:
274+
case TIMESTAMP:
275+
case DOUBLE:
276+
return Long.BYTES;
277+
case TEXT:
278+
case BLOB:
279+
case STRING:
280+
return RamUsageEstimator.NUM_BYTES_OBJECT_REF;
281+
default:
282+
return 0;
283+
}
284+
}
285+
99286
/////////////////////////////////////////////////////////////////////////////////////////////////
100287
// util methods
101288
/////////////////////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)