diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java index 970b337b9ca59..a56423ec2f8d2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java @@ -35,10 +35,14 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; import java.util.Map; +import java.util.NavigableSet; import java.util.Objects; -import java.util.PriorityQueue; +import java.util.Set; +import java.util.TreeSet; /** * Buffers chunks and deletions of one source TsFile per target region partition, managing @@ -63,9 +67,13 @@ interface DispatchCallback { private final Map partition2OffsetCalculator = new HashMap<>(); - /** Max-heap tracking (partitionKey, bufferedSize) for memory budget eviction. */ - private final PriorityQueue largestPieceRegions = - new PriorityQueue<>((a, b) -> Long.compare(b.dataSize(), a.dataSize())); + /** One ordered entry per non-empty partition, so tracking stays bounded as chunks stream in. */ + private final NavigableSet largestPieceRegions = + new TreeSet<>( + Comparator.comparingLong(PieceSize::dataSize) + .reversed() + .thenComparing(entry -> entry.key().regionId()) + .thenComparingLong(entry -> entry.key().timePartitionStart())); PieceDispatcher( final LoadSingleTsFileNode singleTsFileNode, @@ -116,23 +124,29 @@ void offerChunk(final ChunkData chunkData, final TRegionReplicaSet replicaSet) partition2ReplicaSetAndNode.computeIfAbsent( key, k -> new Pair<>(replicaSet, createPieceNode(k))); + largestPieceRegions.remove(new PieceSize(key, targetPair.getRight().getDataSize())); targetPair.getRight().addTsFileData(chunkData); - trackHeapEntry(key, targetPair.getRight().getDataSize()); + trackPieceSize(key, targetPair.getRight().getDataSize()); } - /** Broadcasts deletion data across every active partition buffer. */ - void addDeletionToAll(final DeletionData deletionData) { + /** Buffers one copy per region; its receiver applies the deletion to all staged partitions. */ + void addDeletionToAllRegions(final DeletionData deletionData) { if (deletionData == null || partition2ReplicaSetAndNode.isEmpty()) { return; } final long deletionSize = deletionData.getDataSize(); + final Set regions = new HashSet<>(); for (final Map.Entry> entry : partition2ReplicaSetAndNode.entrySet()) { + if (!regions.add(entry.getKey().regionId())) { + continue; + } memoryBuffer.add(deletionSize); final LoadTsFilePieceNode pieceNode = entry.getValue().getRight(); + largestPieceRegions.remove(new PieceSize(entry.getKey(), pieceNode.getDataSize())); pieceNode.addTsFileData(deletionData); - trackHeapEntry(entry.getKey(), pieceNode.getDataSize()); + trackPieceSize(entry.getKey(), pieceNode.getDataSize()); } } @@ -143,7 +157,7 @@ void addDeletionToAll(final DeletionData deletionData) { /** Progressively dispatches largest buffered pieces until memory usage is back within budget. */ boolean dispatchLargestUntilMemoryEnough() throws LoadFileException { while (!memoryBuffer.isMemoryEnough()) { - final PartitionKey key = pollLargestValidPartition(); + final PartitionKey key = pollLargestPartition(); if (key == null) { break; } @@ -158,6 +172,7 @@ boolean dispatchLargestUntilMemoryEnough() throws LoadFileException { final long pieceSize = pieceNode.getDataSize(); if (!dispatchCallback.dispatch(pieceNode, pair.getLeft())) { + trackPieceSize(key, pieceSize); return false; } @@ -192,10 +207,12 @@ boolean flushAll() throws LoadFileException { } // ------------------------------------------------------------------------- - // Helpers & Eviction Heap Management + // Helpers & Eviction Tracking // ------------------------------------------------------------------------- private void replacePieceNode(final PartitionKey key, final TRegionReplicaSet replicaSet) { + largestPieceRegions.remove( + new PieceSize(key, partition2ReplicaSetAndNode.get(key).getRight().getDataSize())); partition2ReplicaSetAndNode.put(key, new Pair<>(replicaSet, createPieceNode(key))); } @@ -208,25 +225,15 @@ private LoadTsFilePieceNode createPieceNode(final PartitionKey key) { pieceIndex); } - private void trackHeapEntry(final PartitionKey key, final long size) { - largestPieceRegions.offer(new HeapEntry(key, size)); + private void trackPieceSize(final PartitionKey key, final long size) { + if (size > 0) { + largestPieceRegions.add(new PieceSize(key, size)); + } } - private PartitionKey pollLargestValidPartition() { - while (!largestPieceRegions.isEmpty()) { - final HeapEntry entry = largestPieceRegions.poll(); - final Pair pair = - partition2ReplicaSetAndNode.get(entry.key()); - if (pair == null) { - continue; - } - - final long actualSize = pair.getRight().getDataSize(); - if (entry.dataSize() == actualSize && actualSize > 0) { - return entry.key(); - } - } - return null; + private PartitionKey pollLargestPartition() { + final PieceSize entry = largestPieceRegions.pollFirst(); + return entry == null ? null : entry.key(); } void clear() { @@ -242,5 +249,5 @@ void clear() { private record PartitionKey(TConsensusGroupId regionId, long timePartitionStart) {} - private record HeapEntry(PartitionKey key, long dataSize) {} + private record PieceSize(PartitionKey key, long dataSize) {} } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java index 0481d0ce42d9f..fcb704431f398 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java @@ -49,6 +49,7 @@ public class TsFileSplitConsumer implements TsFileSplitter.TsFileDataConsumer { private final Consumer progressIndexCallback; private final List nonDirectionalChunkData = new ArrayList<>(); + private boolean isProcessingDeletions; public TsFileSplitConsumer( final LoadSingleTsFileNode singleTsFileNode, @@ -111,10 +112,17 @@ private boolean addOrSendChunkData(final ChunkData chunkData) throws LoadFileExc } private boolean addOrSendDeletionData(final DeletionData deletionData) throws LoadFileException { - // Route pending chunks first to ensure deletions never precede the chunks they apply to + // The receiver broadcasts each deletion to all staged partitions of its region. Stage all + // chunks first, then send each deletion once per region, including partitions already evicted. routeChunkData(); - dispatcher.addDeletionToAll(deletionData); - return true; + if (!isProcessingDeletions) { + if (!dispatcher.flushAll()) { + return false; + } + isProcessingDeletions = true; + } + dispatcher.addDeletionToAllRegions(deletionData); + return dispatcher.dispatchLargestUntilMemoryEnough(); } // ------------------------------------------------------------------------- @@ -157,6 +165,7 @@ boolean sendAllTsFileData() throws LoadFileException { /** Drops all unrouted chunks, clears partition buffers, and releases cached memory accounting. */ void clear() { nonDirectionalChunkData.clear(); + isProcessingDeletions = false; dispatcher.clear(); memoryBuffer.clear(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java index d3d9efb1345a5..d72fd02e54b88 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java @@ -45,6 +45,7 @@ import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; +import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; @@ -54,6 +55,7 @@ import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.read.TimeValuePair; +import org.apache.tsfile.read.TsFileSequenceReader; import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.Pair; import org.apache.tsfile.utils.RamUsageEstimator; @@ -71,10 +73,12 @@ import java.nio.file.StandardOpenOption; import java.util.ArrayList; import java.util.HashMap; +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.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; @@ -466,6 +470,37 @@ private void writeDeletion(final DeletionData deletionData) throws IOException { // Consensus Lifecycle: PREPARE & LOAD // ------------------------------------------------------------------------- + private void validateStagedFiles( + final Map timePartition2ExpectedDataEnd) + throws IOException, LoadFileException { + final Set partitions = new HashSet<>(timePartition2Writer.keySet()); + if (timePartition2ExpectedDataEnd != null) { + partitions.addAll(timePartition2ExpectedDataEnd.keySet()); + } + for (final TTimePartitionSlot partition : partitions) { + final TsFilePrecalculatedChunkWriter writer = timePartition2Writer.get(partition); + final File stagedFile = + writer == null + ? new File(taskDir, stagedFileName(partition) + TsFileConstant.TSFILE_SUFFIX) + : writer.getFile(); + final LoadTsFileProgress progress = timePartition2Progress.get(partition); + final long fileLength = stagedFile.length(); + final long expectedDataEnd = expectedDataEndOf(timePartition2ExpectedDataEnd, partition); + if (writer == null + || progress == null + || !progress.isReady(fileLength) + || (expectedDataEnd >= 0L && fileLength < expectedDataEnd)) { + throw new LoadFileException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19, + stagedFile.getAbsolutePath(), + uuidDirName(stagedFile), + Math.max(progress == null ? 0L : progress.getTotalLength(), expectedDataEnd), + fileLength)); + } + } + } + /** * Seals modification files, validates chunk continuity, and seals TsFiles in the PREPARE round. */ @@ -477,12 +512,16 @@ void prepare( checkNotClosed(); touch(); - // 1. Close modification files first to freeze deletion views + // Validate every announced partition before sealing any file. A missing partition or a file + // without durable chunk records must not turn into a successful partial import. + validateStagedFiles(timePartition2ExpectedDataEnd); + + // Close modification files to freeze deletion views. for (final ModificationFile modFile : timePartition2ModificationFile.values()) { modFile.close(); } - // 2. Validate continuity and seal writer metadata zone + // Seal writer metadata zones only after every partition has passed validation. for (final Map.Entry entry : timePartition2Writer.entrySet()) { final TsFilePrecalculatedChunkWriter writer = entry.getValue(); @@ -490,23 +529,6 @@ void prepare( continue; } - final LoadTsFileProgress progress = timePartition2Progress.get(entry.getKey()); - if (progress != null && progress.exists() && progress.getTotalLength() > 0) { - final long fileLength = writer.getFile().length(); - final long expectedDataEnd = - expectedDataEndOf(timePartition2ExpectedDataEnd, entry.getKey()); - if (!progress.isReady(fileLength) - || (expectedDataEnd >= 0L && fileLength < expectedDataEnd)) { - throw new LoadFileException( - String.format( - StorageEngineMessages - .EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19, - writer.getFile().getAbsolutePath(), - uuidDirName(writer.getFile()), - Math.max(progress.getTotalLength(), expectedDataEnd), - fileLength)); - } - } // Rebuild from the staged devices so recovered writers also retain the table schemas that // Pipe needs to enumerate table data in the sealed file. if (PathUtils.isTableModelDatabase(dataRegion.getDatabaseName())) { @@ -721,7 +743,21 @@ TsFileWriterManager recoverFromDisk() throws IOException { final TTimePartitionSlot timePartitionSlot = new TTimePartitionSlot(parseTimePartitionStart(tsFile.getName())); final long resumeOffset = progress.getTotalLength(); - if (tsFile.length() > resumeOffset) { + final TsFileResource resource = new TsFileResource(tsFile); + boolean sealed = false; + if (resource.resourceFileExists()) { + try (final TsFileSequenceReader reader = + new TsFileSequenceReader(tsFile.getAbsolutePath(), false)) { + sealed = TSFileConfig.MAGIC_STRING.equals(reader.readTailMagic()); + } + if (sealed) { + resource.deserialize(); + resource.setStatus(TsFileResourceStatus.NORMAL); + } + } + // A resource sidecar and a completed footer record a successful PREPARE. Its coordinator + // may send COMMIT immediately after recovery, so preserve both the sealed bytes and marks. + if (!sealed && tsFile.length() > resumeOffset) { try (final FileChannel truncatingChannel = FileChannel.open(tsFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE)) { truncatingChannel.truncate(resumeOffset); @@ -745,15 +781,15 @@ TsFileWriterManager recoverFromDisk() throws IOException { try { channel = FileChannel.open(tsFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE); - channel.position(resumeOffset); + channel.position(sealed ? tsFile.length() : resumeOffset); final TsFilePrecalculatedChunkWriter writer = - new TsFilePrecalculatedChunkWriter(tsFile, channel); + new TsFilePrecalculatedChunkWriter(tsFile, channel, sealed); writer.setRejectHoles(LoadConsensusCompat.usesIoTConsensusV2()); writer.restoreChunkMetadata(restoreChunkMetadata(records)); publishedLengths.put(tsFile, tsFile.length()); timePartition2Writer.put(timePartitionSlot, writer); - timePartition2Resource.put(timePartitionSlot, new TsFileResource(tsFile)); + timePartition2Resource.put(timePartitionSlot, resource); timePartition2Progress.put(timePartitionSlot, progress); final File modificationFile = ModificationFile.getExclusiveMods(tsFile); @@ -928,11 +964,15 @@ private static void deleteIfExistsWithRetry(final Path path) throws IOException /** The name a staged file of this region and time partition has. */ private String stagedFileName(final TTimePartitionSlot timePartitionSlot) { + return stagedFilePrefix() + timePartitionSlot.getStartTime(); + } + + private String stagedFilePrefix() { return String.join( IoTDBConstant.FILE_NAME_SEPARATOR, dataRegion.getDatabaseName(), dataRegion.getDataRegionIdString(), - Long.toString(timePartitionSlot.getStartTime())); + ""); } private static String uuidDirName(final File tsFile) { @@ -940,10 +980,12 @@ private static String uuidDirName(final File tsFile) { return parent == null ? tsFile.getName() : parent.getName(); } - private static long parseTimePartitionStart(final String tsFileName) { - final String baseName = tsFileName.substring(0, tsFileName.lastIndexOf('.')); - final int lastDash = baseName.lastIndexOf(IoTDBConstant.FILE_NAME_SEPARATOR); - return Long.parseLong(baseName.substring(lastDash + 1)); + private long parseTimePartitionStart(final String tsFileName) { + // The partition may be negative: its sign is part of the number, not a name separator. + return Long.parseLong( + tsFileName.substring( + stagedFilePrefix().length(), + tsFileName.length() - TsFileConstant.TSFILE_SUFFIX.length())); } /** The bytes this piece added to each staged file, as the refs a replica reads them back with. */ diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java index 16109c64de2db..9b17746f942a7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java @@ -24,6 +24,7 @@ import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; @@ -455,7 +456,12 @@ private void getChunkMetadata( private void handleModification(List deletions) throws LoadFileException { for (final ModEntry mod : deletions) { - consumer.apply(new DeletionData(mod)); + if (!consumer.apply(new DeletionData(mod))) { + throw new IllegalStateException( + String.format( + DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S, + tsFile.getPath())); + } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java index b48b7d14cb17d..7e375d9fe7042 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java @@ -109,12 +109,18 @@ public TsFilePrecalculatedChunkWriter(final File file) throws IOException { public TsFilePrecalculatedChunkWriter(final File file, final FileChannel channel) throws IOException { + this(file, channel, false); + } + + /** Resumes a staged file, preserving a completed PREPARE when it is already sealed. */ + public TsFilePrecalculatedChunkWriter( + final File file, final FileChannel channel, final boolean sealed) throws IOException { this.file = file; this.out = new WritableFileChannelOutput( Objects.requireNonNull( channel, StorageEngineMessages.EXCEPTION_CHANNEL_CANNOT_BE_NULL_0F79E0FB)); - this.sealed = false; + this.sealed = sealed; this.maxPhysicalFileSize = channel.size(); LOGGER.info( StorageEngineMessages diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcherTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcherTest.java new file mode 100644 index 0000000000000..b94ef5aa07bf5 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcherTest.java @@ -0,0 +1,146 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.lang.reflect.Field; +import java.util.Collection; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.when; + +public class PieceDispatcherTest { + + @Test + public void testEvictionQueueStaysBoundedAcrossChunksAndFlushes() throws Exception { + final MemoryBoundedBuffer buffer = Mockito.mock(MemoryBoundedBuffer.class); + final AtomicLong usedBytes = new AtomicLong(); + final long budget = chunk(0L).getDataSize() * 5; + when(buffer.isMemoryEnough()).thenAnswer(ignored -> usedBytes.get() <= budget); + Mockito.doAnswer( + invocation -> { + usedBytes.addAndGet(-invocation.getArgument(0)); + return null; + }) + .when(buffer) + .release(anyLong()); + final AtomicInteger dispatchedChunks = new AtomicInteger(); + final PieceDispatcher dispatcher = + new PieceDispatcher( + loadNode(), + buffer, + (piece, region) -> { + dispatchedChunks.addAndGet(piece.getAllTsFileData().size()); + return true; + }); + final TRegionReplicaSet replicaSet = + new TRegionReplicaSet() + .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion, 1)); + final Field queueField = PieceDispatcher.class.getDeclaredField("largestPieceRegions"); + queueField.setAccessible(true); + + for (int i = 0; i < 128; i++) { + final NonAlignedChunkData chunk = chunk(i % 2); + usedBytes.addAndGet(chunk.getDataSize()); + dispatcher.offerChunk(chunk, replicaSet); + assertTrue(dispatcher.dispatchLargestUntilMemoryEnough()); + assertTrue( + "Eviction tracking must be bounded by the two live partition buffers", + ((Collection) queueField.get(dispatcher)).size() <= 2); + } + assertTrue(dispatcher.flushAll()); + assertEquals(128, dispatchedChunks.get()); + assertEquals(0L, usedBytes.get()); + assertTrue( + "Flushed pieces must leave no eviction entries", + ((Collection) queueField.get(dispatcher)).isEmpty()); + } + + @Test + public void testFailedEvictionCanBeRetried() throws Exception { + final MemoryBoundedBuffer buffer = Mockito.mock(MemoryBoundedBuffer.class); + final AtomicBoolean memoryEnough = new AtomicBoolean(false); + when(buffer.isMemoryEnough()).thenAnswer(ignored -> memoryEnough.get()); + Mockito.doAnswer( + ignored -> { + memoryEnough.set(true); + return null; + }) + .when(buffer) + .release(anyLong()); + final AtomicInteger attempts = new AtomicInteger(); + final PieceDispatcher dispatcher = + new PieceDispatcher(loadNode(), buffer, (piece, region) -> attempts.incrementAndGet() > 1); + final NonAlignedChunkData chunk = chunk(0L); + dispatcher.offerChunk( + chunk, + new TRegionReplicaSet() + .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion, 1))); + + assertFalse(dispatcher.dispatchLargestUntilMemoryEnough()); + assertTrue(dispatcher.dispatchLargestUntilMemoryEnough()); + assertEquals(2, attempts.get()); + Mockito.verify(buffer).release(chunk.getDataSize()); + } + + private static LoadSingleTsFileNode loadNode() { + final LoadSingleTsFileNode node = Mockito.mock(LoadSingleTsFileNode.class); + when(node.getPlanNodeId()).thenReturn(new PlanNodeId("bounded-queue")); + final TsFileResource resource = Mockito.mock(TsFileResource.class); + when(resource.getTsFile()).thenReturn(new File("bounded-queue.tsfile")); + when(node.getTsFileResource()).thenReturn(resource); + return node; + } + + private static NonAlignedChunkData chunk(final long partition) { + final NonAlignedChunkData chunk = + new NonAlignedChunkData( + new StringArrayDeviceID("root", "test", "d1"), + new ChunkHeader( + "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0), + new TTimePartitionSlot(partition)); + chunk.writeDecodePage(new long[] {partition + 1}, new Object[] {1}, 1); + chunk.endChunk(); + return chunk; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumerTest.java new file mode 100644 index 0000000000000..93bf7dc4c7bcf --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumerTest.java @@ -0,0 +1,233 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; +import org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; +import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.utils.Pair; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.Mockito; + +import java.io.File; +import java.lang.reflect.Constructor; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.when; + +public class TsFileSplitConsumerTest { + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void testDeletionThatExceedsMemoryBudgetTriggersDispatch() throws Exception { + final int memoryBudget = IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize() >> 2; + final Constructor constructor = + LoadTsFileDataCacheMemoryBlock.class.getDeclaredConstructor(long.class); + constructor.setAccessible(true); + final LoadTsFileDataCacheMemoryBlock memoryBlock = constructor.newInstance(1024 * 1024L); + + final TRegionReplicaSet replicaSet = + new TRegionReplicaSet() + .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion, 1)); + final DataPartitionBatchFetcher partitionFetcher = + Mockito.mock(DataPartitionBatchFetcher.class); + when(partitionFetcher.queryDataPartition(anyList(), anyString())) + .thenReturn(List.of(replicaSet)); + + final LoadSingleTsFileNode loadNode = Mockito.mock(LoadSingleTsFileNode.class); + when(loadNode.getPlanNodeId()).thenReturn(new PlanNodeId("load-memory-budget")); + final TsFileResource resource = Mockito.mock(TsFileResource.class); + when(resource.getTsFile()).thenReturn(new File("load-memory-budget.tsfile")); + when(loadNode.getTsFileResource()).thenReturn(resource); + + final AtomicInteger dispatchedPieces = new AtomicInteger(); + final TsFileSplitConsumer consumer = + new TsFileSplitConsumer( + loadNode, + memoryBlock, + new DataPartitionResolver(partitionFetcher, "root.test"), + null, + (pieceNode, targetReplicaSet) -> { + if (pieceNode.getAllTsFileData().stream().anyMatch(DeletionData.class::isInstance)) { + dispatchedPieces.incrementAndGet(); + } + return true; + }); + + final NonAlignedChunkData chunk = + new NonAlignedChunkData( + new StringArrayDeviceID("root", "test", "d1"), + new ChunkHeader( + "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0), + new TTimePartitionSlot(0L)); + chunk.writeDecodePage(new long[] {1L}, new Object[] {1}, 1); + chunk.endChunk(); + + final ModEntry deletion = Mockito.mock(ModEntry.class); + when(deletion.serializedSize()).thenReturn(memoryBudget + 1); + + try { + Assert.assertTrue(consumer.apply(chunk)); + Assert.assertEquals(0, dispatchedPieces.get()); + + Assert.assertTrue(consumer.apply(new DeletionData(deletion))); + + Assert.assertEquals( + "A deletion must trigger largest-first dispatch when its buffered piece exceeds budget", + 1, + dispatchedPieces.get()); + } finally { + consumer.clear(); + } + } + + @Test + public void testDeletionsAreSentOncePerRegionAfterAllPartitions() throws Exception { + final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + final String[] originalLoadDirs = config.getLoadTsFileDirs(); + final File loadDir = temporaryFolder.newFolder("load"); + config.setLoadTsFileDirs(new String[] {loadDir.getAbsolutePath()}); + final LoadTsFileManager firstRegion = manager(1); + final LoadTsFileManager secondRegion = manager(2); + final Map managers = Map.of(1, firstRegion, 2, secondRegion); + final TRegionReplicaSet firstReplicaSet = replicaSet(1); + final TRegionReplicaSet secondReplicaSet = replicaSet(2); + final DataPartitionBatchFetcher partitionFetcher = + Mockito.mock(DataPartitionBatchFetcher.class); + when(partitionFetcher.queryDataPartition(anyList(), anyString())) + .thenAnswer( + invocation -> { + final List> + slots = invocation.getArgument(0); + return slots.stream() + .map( + slot -> + slot.right.getStartTime() < 2000L ? firstReplicaSet : secondReplicaSet) + .toList(); + }); + final Constructor constructor = + LoadTsFileDataCacheMemoryBlock.class.getDeclaredConstructor(long.class); + constructor.setAccessible(true); + final TsFileSplitConsumer consumer = + new TsFileSplitConsumer( + loadNode(), + constructor.newInstance(1024 * 1024L), + new DataPartitionResolver(partitionFetcher, "root"), + null, + (piece, target) -> { + try { + managers + .get(target.getRegionId().getId()) + .writePiece("load-deletions", piece.getAllTsFileData()); + return true; + } catch (Exception e) { + throw new AssertionError(e); + } + }); + final ModEntry firstDeletion = + new TreeDeletionEntry(new MeasurementPath("root.test.d1.s1"), 0L, 10L); + final ModEntry secondDeletion = + new TreeDeletionEntry(new MeasurementPath("root.test.d1.s1"), 20L, 30L); + try { + for (long partition : new long[] {0L, 1000L, 2000L}) { + Assert.assertTrue(consumer.apply(chunk(partition))); + } + Assert.assertTrue(consumer.apply(new DeletionData(firstDeletion))); + Assert.assertTrue(consumer.apply(new DeletionData(secondDeletion))); + Assert.assertTrue(consumer.sendAllTsFileData()); + for (long partition : new long[] {0L, 1000L, 2000L}) { + final int regionId = partition < 2000L ? 1 : 2; + final File taskDir = managers.get(regionId).getActiveTaskDirs().get(0); + final File staged = + new File(taskDir, "root.test-" + regionId + "-" + partition + ".tsfile"); + Assert.assertEquals( + "Each partition must receive every deletion exactly once", + List.of(firstDeletion, secondDeletion), + ModificationFile.readAllModifications(staged, false)); + } + } finally { + consumer.clear(); + firstRegion.stop(); + secondRegion.stop(); + config.setLoadTsFileDirs(originalLoadDirs); + } + } + + private static LoadTsFileManager manager(final int regionId) { + final DataRegion region = Mockito.mock(DataRegion.class); + when(region.getDatabaseName()).thenReturn("root.test"); + when(region.getDataRegionIdString()).thenReturn(String.valueOf(regionId)); + return new LoadTsFileManager(region); + } + + private static TRegionReplicaSet replicaSet(final int regionId) { + return new TRegionReplicaSet() + .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion, regionId)); + } + + private static LoadSingleTsFileNode loadNode() { + final LoadSingleTsFileNode node = Mockito.mock(LoadSingleTsFileNode.class); + when(node.getPlanNodeId()).thenReturn(new PlanNodeId("load-deletions")); + final TsFileResource resource = Mockito.mock(TsFileResource.class); + when(resource.getTsFile()).thenReturn(new File("load-deletions.tsfile")); + when(node.getTsFileResource()).thenReturn(resource); + return node; + } + + private static NonAlignedChunkData chunk(final long partition) { + final NonAlignedChunkData chunk = + new NonAlignedChunkData( + new StringArrayDeviceID("root", "test", "d1"), + new ChunkHeader( + "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0), + new TTimePartitionSlot(partition)); + chunk.writeDecodePage(new long[] {partition + 1}, new Object[] {1}, 1); + chunk.endChunk(); + return chunk; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java index 1797c4011602c..4ccbc98aec631 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java @@ -21,7 +21,10 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.path.MeasurementPath; import org.apache.iotdb.commons.utils.TimePartitionUtils; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; +import org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry; import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.enums.ColumnCategory; @@ -61,9 +64,47 @@ import java.util.Collections; import java.util.List; import java.util.Objects; +import java.util.concurrent.atomic.AtomicInteger; public class TsFileSplitterTest { + @Test + public void testRejectedDeletionStopsSplitting() throws Exception { + final File source = Files.createTempFile("rejected-load-deletion", ".tsfile").toFile(); + final File mods = ModificationFile.getExclusiveMods(source); + try { + try (final TsFileIOWriter writer = new TsFileIOWriter(source)) { + writer.startChunkGroup(new StringArrayDeviceID("root", "sg", "d1")); + final ChunkWriterImpl chunkWriter = + new ChunkWriterImpl(new MeasurementSchema("s1", TSDataType.INT32)); + chunkWriter.write(1L, 1); + chunkWriter.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } + try (final ModificationFile modificationFile = new ModificationFile(mods, false)) { + modificationFile.write(new TreeDeletionEntry(new MeasurementPath("root.sg.d1.s1"), 0L, 1L)); + modificationFile.write(new TreeDeletionEntry(new MeasurementPath("root.sg.d1.s1"), 2L, 3L)); + } + final AtomicInteger consumedDeletions = new AtomicInteger(); + final TsFileSplitter splitter = + new TsFileSplitter( + source, + data -> { + if (data instanceof DeletionData) { + consumedDeletions.incrementAndGet(); + return false; + } + return true; + }); + Assert.assertThrows(IllegalStateException.class, splitter::splitTsFileByDataPartition); + Assert.assertEquals(1, consumedDeletions.get()); + } finally { + Files.deleteIfExists(mods.toPath()); + Files.deleteIfExists(source.toPath()); + } + } + // Verify the splitter initializes the v3 deserialize configuration for a valid v3 TsFile. @Test public void testSplitV3TsFile() throws Exception { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerRecoveryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerRecoveryTest.java new file mode 100644 index 0000000000000..9c76b75c6a172 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerRecoveryTest.java @@ -0,0 +1,254 @@ +/* + * 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.storageengine.load; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.consensus.index.impl.SimpleProgressIndex; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.queryengine.plan.scheduler.load.ChunkOffsetCalculator; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TsFileReader; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TsFileWriterManagerRecoveryTest { + + private static final String DATABASE = "root.load_recovery_test"; + private static final IDeviceID DEVICE = + new StringArrayDeviceID("root", "load_recovery_test", "d1"); + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private DataRegion dataRegion; + + @Before + public void setUp() { + dataRegion = Mockito.mock(DataRegion.class); + Mockito.when(dataRegion.getDatabaseName()).thenReturn(DATABASE); + Mockito.when(dataRegion.getDataRegionIdString()).thenReturn("0"); + } + + @Test + public void testResumeNegativeTimePartition() throws Exception { + assertTimePartitionResumes(-604_800_000L); + } + + @Test + public void testResumeMinimumTimePartition() throws Exception { + assertTimePartitionResumes(Long.MIN_VALUE); + } + + @Test + public void testCommitPreparedFileAfterRecovery() throws Exception { + final File taskDir = temporaryFolder.newFolder("prepared-task"); + final NonAlignedChunkData chunk = chunk(0L, 1); + new ChunkOffsetCalculator().assign(chunk); + final ProgressIndex progressIndex = new SimpleProgressIndex(1, 10L); + TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir); + try { + manager.writePiece(List.of(chunk)); + manager.prepare( + true, + Map.of(chunk.getTimePartitionSlot(), progressIndex), + Map.of(chunk.getTimePartitionSlot(), chunkEnd(chunk))); + final File staged = stagedFile(taskDir, 0L); + final byte[] preparedBytes = Files.readAllBytes(staged.toPath()); + manager.close(true); + + manager = new TsFileWriterManager(dataRegion, taskDir).recoverFromDisk(); + + assertArrayEquals( + "Recovery must preserve a prepared TsFile", + preparedBytes, + Files.readAllBytes(staged.toPath())); + // The coordinator already received the PREPARE acknowledgement and only sends COMMIT. + manager.loadAll(true, false); + final ArgumentCaptor imported = ArgumentCaptor.forClass(TsFileResource.class); + Mockito.verify(dataRegion) + .loadNewTsFile( + imported.capture(), + Mockito.eq(false), + Mockito.eq(true), + Mockito.eq(false), + Mockito.any()); + assertTrue(imported.getValue().isGeneratedByPipe()); + assertEquals(progressIndex, imported.getValue().getMaxProgressIndexAfterClose()); + assertPoints(staged, 0L, 1); + } finally { + manager.close(); + } + } + + private void assertTimePartitionResumes(final long partitionStart) throws Exception { + final File taskDir = temporaryFolder.newFolder("task"); + final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator(); + final NonAlignedChunkData first = chunk(partitionStart, 1); + final NonAlignedChunkData second = chunk(partitionStart, 2); + calculator.assign(first); + calculator.assign(second); + + TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir); + try { + manager.writePiece(List.of(first)); + manager.close(true); + manager = new TsFileWriterManager(dataRegion, taskDir).recoverFromDisk(); + manager.writePiece(List.of(second)); + manager.prepare( + false, Collections.emptyMap(), Map.of(second.getTimePartitionSlot(), chunkEnd(second))); + + // Both the chunk staged before recovery and the one appended afterwards must survive. + assertPoints(stagedFile(taskDir, partitionStart), partitionStart, 1, 2); + } finally { + manager.close(); + } + } + + @Test + public void testPrepareRejectsMissingPartition() throws Exception { + final File taskDir = temporaryFolder.newFolder("task"); + final TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir); + final NonAlignedChunkData first = chunk(0L, 1); + final NonAlignedChunkData missing = chunk(604_800_000L, 1); + new ChunkOffsetCalculator().assign(first); + new ChunkOffsetCalculator().assign(missing); + final Map expectedDataEnd = + Map.of( + first.getTimePartitionSlot(), chunkEnd(first), + missing.getTimePartitionSlot(), chunkEnd(missing)); + try { + manager.writePiece(List.of(first)); + final long unsealedLength = stagedFile(taskDir, 0L).length(); + assertPrepareFails(manager, expectedDataEnd, stagedFile(taskDir, 604_800_000L)); + assertEquals(unsealedLength, stagedFile(taskDir, 0L).length()); + + manager.writePiece(List.of(missing)); + manager.prepare(false, Collections.emptyMap(), expectedDataEnd); + assertPoints(stagedFile(taskDir, 0L), 0L, 1); + assertPoints(stagedFile(taskDir, 604_800_000L), 604_800_000L, 1); + } finally { + manager.close(); + } + } + + @Test + public void testPrepareRejectsPartitionWithoutRecordedChunks() throws Exception { + final File taskDir = temporaryFolder.newFolder("task"); + final TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir); + final NonAlignedChunkData complete = chunk(0L, 1); + new ChunkOffsetCalculator().assign(complete); + final Map expectedDataEnd = + Map.of(complete.getTimePartitionSlot(), chunkEnd(complete)); + try { + // A failed piece may have opened the file without recording any durable chunk. + try { + manager.writePiece(List.of(chunk(0L, 1))); + fail("A chunk without a layout must fail after opening its staged file"); + } catch (final IOException expected) { + assertTrue(stagedFile(taskDir, 0L).isFile()); + } + assertPrepareFails(manager, expectedDataEnd, stagedFile(taskDir, 0L)); + + manager.writePiece(List.of(complete)); + manager.prepare(false, Collections.emptyMap(), expectedDataEnd); + assertPoints(stagedFile(taskDir, 0L), 0L, 1); + } finally { + manager.close(); + } + } + + private static void assertPrepareFails( + final TsFileWriterManager manager, + final Map expectedDataEnd, + final File missingDataFile) + throws Exception { + try { + manager.prepare(false, Collections.emptyMap(), expectedDataEnd); + fail("PREPARE must reject a partition whose announced data has not been staged"); + } catch (final LoadFileException e) { + assertTrue(e.getMessage(), e.getMessage().contains(missingDataFile.getAbsolutePath())); + } + } + + private static NonAlignedChunkData chunk(final long partitionStart, final int value) { + final NonAlignedChunkData chunk = + new NonAlignedChunkData( + DEVICE, + new ChunkHeader( + "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0), + new TTimePartitionSlot(partitionStart)); + chunk.writeDecodePage(new long[] {partitionStart + value}, new Object[] {value}, 1); + chunk.endChunk(); + return chunk; + } + + private static long chunkEnd(final NonAlignedChunkData chunk) { + return chunk.getChunkLayout().offset() + chunk.getChunkLayout().length(); + } + + private static File stagedFile(final File taskDir, final long partitionStart) { + return new File(taskDir, DATABASE + "-0-" + partitionStart + ".tsfile"); + } + + private static void assertPoints(final File file, final long partitionStart, final int... values) + throws Exception { + try (final TsFileReader reader = new TsFileReader(file)) { + final QueryDataSet result = + reader.query(QueryExpression.create(List.of(new Path(DEVICE, "s1", false)), null)); + for (final int value : values) { + assertTrue(result.hasNext()); + final RowRecord row = result.next(); + assertEquals(partitionStart + value, row.getTimestamp()); + assertEquals(value, row.getField(0).getIntV()); + } + assertFalse(result.hasNext()); + } + } +}