Add readChunksInDevice method to TsFileSequenceReader (#1765)
* add readChunksInDevice method to TsFileSequenceReader
* add a test for TsFileSequenceReader#readChunksInDevice(String)
* close the reader after test readChunksInDevice
diff --git a/tsfile/src/main/java/org/apache/iotdb/tsfile/read/TsFileSequenceReader.java b/tsfile/src/main/java/org/apache/iotdb/tsfile/read/TsFileSequenceReader.java
index 776fb08..2d3313f 100644
--- a/tsfile/src/main/java/org/apache/iotdb/tsfile/read/TsFileSequenceReader.java
+++ b/tsfile/src/main/java/org/apache/iotdb/tsfile/read/TsFileSequenceReader.java
@@ -686,6 +686,34 @@
}
/**
+ * read all Chunks of given device.
+ * <p>
+ * note that this method loads all the chunks into memory, so it needs to be invoked carefully.
+ *
+ * @param device name
+ * @return measurement -> chunks list
+ */
+ public Map<String, List<Chunk>> readChunksInDevice(String device) throws IOException {
+ List<ChunkMetadata> chunkMetadataList = new ArrayList<>();
+ Map<String, List<ChunkMetadata>> chunkMetadataInDevice = readChunkMetadataInDevice(device);
+ for (List<ChunkMetadata> chunkMetadataListInDevice : chunkMetadataInDevice.values()) {
+ chunkMetadataList.addAll(chunkMetadataListInDevice);
+ }
+
+ Map<String, List<Chunk>> chunksInDevice = new HashMap<>();
+ chunkMetadataList.sort(Comparator.comparing(ChunkMetadata::getOffsetOfChunkHeader));
+ for (ChunkMetadata chunkMetadata : chunkMetadataList) {
+ Chunk chunk = readMemChunk(chunkMetadata);
+ String measurement = chunk.getHeader().getMeasurementID();
+ if (!chunksInDevice.containsKey(measurement)) {
+ chunksInDevice.put(measurement, new ArrayList<>());
+ }
+ chunksInDevice.get(measurement).add(chunk);
+ }
+ return chunksInDevice;
+ }
+
+ /**
* not thread safe.
*
* @param type given tsfile data type
diff --git a/tsfile/src/test/java/org/apache/iotdb/tsfile/read/TsFileSequenceReaderTest.java b/tsfile/src/test/java/org/apache/iotdb/tsfile/read/TsFileSequenceReaderTest.java
index b4c5671..f8c80ec 100644
--- a/tsfile/src/test/java/org/apache/iotdb/tsfile/read/TsFileSequenceReaderTest.java
+++ b/tsfile/src/test/java/org/apache/iotdb/tsfile/read/TsFileSequenceReaderTest.java
@@ -22,16 +22,23 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.stream.Collectors;
import org.apache.iotdb.tsfile.common.conf.TSFileConfig;
import org.apache.iotdb.tsfile.file.MetaMarker;
import org.apache.iotdb.tsfile.file.footer.ChunkGroupFooter;
import org.apache.iotdb.tsfile.file.header.ChunkHeader;
import org.apache.iotdb.tsfile.file.header.PageHeader;
+import org.apache.iotdb.tsfile.file.metadata.ChunkMetadata;
+import org.apache.iotdb.tsfile.read.common.Chunk;
import org.apache.iotdb.tsfile.utils.FileGenerator;
import org.apache.iotdb.tsfile.utils.Pair;
import org.junit.After;
+import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
@@ -107,4 +114,35 @@
*/
reader.close();
}
-}
\ No newline at end of file
+
+ @Test
+ public void readChunksInDevice() throws IOException {
+ TsFileSequenceReader reader = new TsFileSequenceReader(FILE_PATH);
+
+ for (String device : reader.getAllDevices()) {
+ Map<String, Set<Chunk>> expectedChunksInDevice = new HashMap<>();
+ for (Entry<String, List<ChunkMetadata>> entry : reader.readChunkMetadataInDevice(device)
+ .entrySet()) {
+ expectedChunksInDevice.putIfAbsent(entry.getKey(), new HashSet<>());
+ for (ChunkMetadata chunkMetadata : entry.getValue()) {
+ expectedChunksInDevice.get(entry.getKey()).add(reader.readMemChunk(chunkMetadata));
+ }
+ }
+
+ Map<String, List<Chunk>> actualChunksInDevice = reader.readChunksInDevice(device);
+
+ for (Entry<String, Set<Chunk>> entry : expectedChunksInDevice.entrySet()) {
+ Set<String> expectedChunkStrings = entry.getValue().stream()
+ .map(chunk -> chunk.getHeader().toString()).collect(Collectors.toSet());
+
+ Assert.assertTrue(actualChunksInDevice.containsKey(entry.getKey()));
+ Set<String> actualChunkStrings = actualChunksInDevice.get(entry.getKey()).stream()
+ .map(chunk -> chunk.getHeader().toString()).collect(Collectors.toSet());
+
+ Assert.assertEquals(expectedChunkStrings, actualChunkStrings);
+ }
+ }
+
+ reader.close();
+ }
+}
