Add a method to batch search for the MetadataIndexNode offset of devices

java/tsfile/src/main/java/org/apache/tsfile/file/metadata/MetadataIndexNode.java

@@ -33,10 +33,12 @@
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 public class MetadataIndexNode {
 
+  private static final double LOG2 = Math.log(2);
   protected static final TSFileConfig config = TSFileDescriptor.getInstance().getConfig();
   protected final List<IMetadataIndexEntry> children;
   protected long endOffset;
@@ -127,6 +129,10 @@ public static MetadataIndexNode deserializeFrom(
 
   public Pair<IMetadataIndexEntry, Long> getChildIndexEntry(Comparable key, boolean exactSearch) {
     int index = binarySearchInChildren(key, exactSearch);
+    return getChildIndexEntry(index);
+  }
+
+  private Pair<IMetadataIndexEntry, Long> getChildIndexEntry(int index) {
     if (index == -1) {
       return null;
     }
@@ -165,6 +171,105 @@ int binarySearchInChildren(Comparable key, boolean exactSearch) {
     }
   }
 
+  public List<Pair<IMetadataIndexEntry, Long>> getChildIndexEntries(
+      List<? extends Comparable> keys, boolean exactSearch) {
+    int[] indexArr =
+        keys.size() >= children.size()
+                || (keys.size() * Math.log(children.size()) / LOG2)
+                    > (keys.size() + children.size())
+            ? mergeSearchInChildren(keys, exactSearch)
+            : binarySearchInChildren(keys, exactSearch);
+    List<Pair<IMetadataIndexEntry, Long>> pairs = new ArrayList<>();
+    int previousIndex = -1;
+    Pair<IMetadataIndexEntry, Long> previousPair = null;
+    for (int idx : indexArr) {
+      if (previousIndex == idx) {
+        pairs.add(previousPair);
+      } else {
+        Pair<IMetadataIndexEntry, Long> current = getChildIndexEntry(idx);
+        pairs.add(current);
+        previousIndex = idx;
+        previousPair = current;
+      }
+    }
+    return pairs;
+  }
+
+  int[] binarySearchInChildren(List<? extends Comparable> keys, boolean exactSearch) {
+    int[] results = new int[keys.size()];
+    Arrays.fill(results, -1);
+    int currentLow = 0;
+    int high = children.size() - 1;
+
+    for (int i = 0; i < keys.size(); i++) {
+      Comparable key = keys.get(i);
+      if (currentLow > high) {
+        Arrays.fill(results, i, keys.size(), exactSearch ? -1 : currentLow - 1);
+        return results;
+      }
+
+      int foundIndex = -1;
+      int start = currentLow;
+      int end = high;
+
+      while (start <= end) {
+        int mid = (start + end) >>> 1;
+        IMetadataIndexEntry midVal = children.get(mid);
+        int cmp = midVal.getCompareKey().compareTo(key);
+
+        if (cmp < 0) {
+          start = mid + 1;
+        } else if (cmp > 0) {
+          end = mid - 1;
+        } else {
+          foundIndex = mid;
+          break;
+        }
+      }
+
+      if (foundIndex >= 0) {
+        results[i] = foundIndex;
+        currentLow = foundIndex + 1;
+      } else {
+        if (exactSearch) {
+          results[i] = -1;
+        } else {
+          int insertPos = start - 1;
+          results[i] = insertPos;
+          currentLow = start;
+        }
+      }
+    }
+    return results;
+  }
+
+  int[] mergeSearchInChildren(List<? extends Comparable> keys, boolean exactSearch) {
+    int[] results = new int[keys.size()];
+    int i = 0;
+    int j = 0;
+    while (i < keys.size() && j < children.size()) {
+      Comparable currentKey = keys.get(i);
+      Comparable currentChild = children.get(j).getCompareKey();
+      int cmp = currentKey.compareTo(currentChild);
+      if (cmp == 0) {
+        results[i] = j;
+        i++;
+        j++;
+      } else if (cmp > 0) {
+        j++;
+      } else {
+        if (exactSearch) {
+          results[i] = -1;
+        } else {
+          results[i] = j - 1;
+        }
+        i++;
+      }
+    }
+    Arrays.fill(results, i, keys.size(), exactSearch ? -1 : children.size() - 1);
+    return results;
+  }
+
   public boolean isDeviceLevel() {
     return this.nodeType == MetadataIndexNodeType.INTERNAL_DEVICE
         || this.nodeType == MetadataIndexNodeType.LEAF_DEVICE;

java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java

@@ -67,6 +67,8 @@
 import org.apache.tsfile.read.common.Path;
 import org.apache.tsfile.read.controller.CachedChunkLoaderImpl;
 import org.apache.tsfile.read.controller.MetadataQuerierByFileImpl;
+import org.apache.tsfile.read.query.DeviceMetadataIndexEntriesQueryResult;
+import org.apache.tsfile.read.query.DeviceMetadataIndexNodeOffsetsQueryContext;
 import org.apache.tsfile.read.reader.TsFileInput;
 import org.apache.tsfile.read.reader.page.PageReader;
 import org.apache.tsfile.read.reader.page.TimePageReader;
@@ -535,6 +537,123 @@ private Map<String, TimeseriesMetadata> readDeviceMetadataFromDisk(IDeviceID dev
     return deviceMetadata;
   }
 
+  /**
+   * Find the offset of MetadataIndexNode corresponding to every device to avoid repeated reading of
+   * internal MetadataIndexNode
+   *
+   * @param table table name. ("" or null) for cases with different tables
+   * @param sortedDevices devices should be sorted
+   * @param ioSizeRecorder can be null
+   * @return Each element of the outer array corresponds to the device at this index. The inner
+   *     array size is 2, the first element is the start offset, and the second is the end offset
+   * @throws IOException io error
+   */
+  public DeviceMetadataIndexEntriesQueryResult getDeviceMetadataIndexNodeOffsets(
+      String table, List<IDeviceID> sortedDevices, LongConsumer ioSizeRecorder) throws IOException {
+    DeviceMetadataIndexNodeOffsetsQueryContext context =
+        new DeviceMetadataIndexNodeOffsetsQueryContext(sortedDevices.size(), getAllMetadataSize());
+    if (sortedDevices.isEmpty()) {
+      return context.compact();
+    }
+    readFileMetadata();
+    // for many table
+    if (table == null || table.isEmpty()) {
+      int tableStartIndex = 0;
+      String previousTable = null;
+      for (int i = 0; i < sortedDevices.size(); i++) {
+        String currentTable = sortedDevices.get(i).getTableName();
+        if (i == 0) {
+          previousTable = currentTable;
+          continue;
+        }
+        if (!previousTable.equals(currentTable)) {
+          getDeviceMetadataIndexNodeOffsets(
+              context, sortedDevices, tableStartIndex, i, previousTable, ioSizeRecorder);
+          tableStartIndex = i;
+          previousTable = currentTable;
+        }
+      }
+      // for last table
+      getDeviceMetadataIndexNodeOffsets(
+          context,
+          sortedDevices,
+          tableStartIndex,
+          sortedDevices.size(),
+          previousTable,
+          ioSizeRecorder);
+    } else {
+      getDeviceMetadataIndexNodeOffsets(
+          context, sortedDevices, 0, sortedDevices.size(), table, ioSizeRecorder);
+    }
+    return context.compact();
+  }
+
+  private void getDeviceMetadataIndexNodeOffsets(
+      DeviceMetadataIndexNodeOffsetsQueryContext context,
+      List<IDeviceID> devices,
+      int deviceStartIdx,
+      int deviceEndIdx,
+      String table,
+      LongConsumer ioSizeRecorder)
+      throws IOException {
+    MetadataIndexNode tableMetadataIndexNode = getTableRootNode(table);
+    if (tableMetadataIndexNode == null) {
+      return;
+    }
+    getDeviceMetadataIndexNodeOffsets(
+        context, devices, deviceStartIdx, deviceEndIdx, tableMetadataIndexNode, ioSizeRecorder);
+  }
+
+  private void getDeviceMetadataIndexNodeOffsets(
+      DeviceMetadataIndexNodeOffsetsQueryContext context,
+      List<IDeviceID> devices,
+      int deviceStartIdx,
+      int deviceEndIdx,
+      MetadataIndexNode startNode,
+      LongConsumer ioSizeRecorder)
+      throws IOException {
+    MetadataIndexNodeType metadataIndexNodeType = startNode.getNodeType();
+    boolean exactSearch = metadataIndexNodeType == MetadataIndexNodeType.LEAF_DEVICE;
+    List<Pair<IMetadataIndexEntry, Long>> entries =
+        startNode.getChildIndexEntries(devices.subList(deviceStartIdx, deviceEndIdx), exactSearch);
+    Iterator<Pair<IMetadataIndexEntry, Long>> metadataIndexEntriesIterator = entries.iterator();
+    int startIdxOfChild = deviceStartIdx;
+    Pair<IMetadataIndexEntry, Long> previousPair = null;
+    for (int i = deviceStartIdx; i < deviceEndIdx; i++) {
+      Pair<IMetadataIndexEntry, Long> pair = metadataIndexEntriesIterator.next();
+      if (exactSearch) {
+        if (pair != null) {
+          context.addDeviceMetadataIndexNodeOffset(i, pair.getLeft().getOffset(), pair.getRight());
+        }
+        continue;
+      }
+      if (previousPair == null) {
+        previousPair = pair;
+        continue;
+      }
+      if (previousPair == pair) {
+        continue;
+      }
+      IMetadataIndexEntry entry = previousPair.getLeft();
+      ByteBuffer buffer = readData(entry.getOffset(), previousPair.getRight(), ioSizeRecorder);
+      MetadataIndexNode lastNode =
+          MetadataIndexNode.deserializeFrom(buffer, true, deserializeConfig);
+      getDeviceMetadataIndexNodeOffsets(
+          context, devices, startIdxOfChild, i, lastNode, ioSizeRecorder);
+      previousPair = pair;
+      startIdxOfChild = i;
+    }
+    if (exactSearch || previousPair == null) {
+      return;
+    }
+    // for last entry
+    IMetadataIndexEntry entry = previousPair.getLeft();
+    ByteBuffer buffer = readData(entry.getOffset(), previousPair.getRight(), ioSizeRecorder);
+    MetadataIndexNode lastNode = MetadataIndexNode.deserializeFrom(buffer, true, deserializeConfig);
+    getDeviceMetadataIndexNodeOffsets(
+        context, devices, startIdxOfChild, deviceEndIdx, lastNode, ioSizeRecorder);
+  }
+
   public TimeseriesMetadata readTimeseriesMetadata(
       IDeviceID device, String measurement, boolean ignoreNotExists) throws IOException {
     return readTimeseriesMetadata(device, measurement, ignoreNotExists, null);
@@ -546,11 +665,31 @@ public TimeseriesMetadata readTimeseriesMetadata(
       boolean ignoreNotExistDevice,
       LongConsumer ioSizeConsumer)
       throws IOException {
+    return readTimeseriesMetadata(device, null, measurement, ignoreNotExistDevice, ioSizeConsumer);
+  }
+
+  public TimeseriesMetadata readTimeseriesMetadata(
+      IDeviceID device,
+      long[] deviceMetadataIndexNodeOffset,
+      String measurement,
+      boolean ignoreNotExistDevice,
+      LongConsumer ioSizeConsumer)
+      throws IOException {
     readFileMetadata(ioSizeConsumer);
     MetadataIndexNode deviceMetadataIndexNode =
         tsFileMetaData.getTableMetadataIndexNode(device.getTableName());
-    Pair<IMetadataIndexEntry, Long> metadataIndexPair =
-        getMetadataAndEndOffsetOfDeviceNode(deviceMetadataIndexNode, device, true, ioSizeConsumer);
+    Pair<IMetadataIndexEntry, Long> metadataIndexPair;
+    if (deviceMetadataIndexNodeOffset == null) {
+      metadataIndexPair =
+          getMetadataAndEndOffsetOfDeviceNode(
+              deviceMetadataIndexNode, device, true, ioSizeConsumer);
+    } else {
+      metadataIndexPair =
+          new Pair<>(
+              new DeviceMetadataIndexEntry(device, deviceMetadataIndexNodeOffset[0]),
+              deviceMetadataIndexNodeOffset[1]);
+    }
+
     if (metadataIndexPair == null) {
       if (ignoreNotExistDevice) {
         return null;
@@ -683,13 +822,15 @@ public ITimeSeriesMetadata readITimeseriesMetadata(Path path, boolean ignoreNotE
    */
   public List<TimeseriesMetadata> readTimeseriesMetadata(
       IDeviceID device,
+      long[] deviceMetadataIndexNodeOffset,
       String measurement,
       Set<String> allSensors,
       boolean ignoreNotExistDevice,
       LongConsumer ioSizeRecorder)
       throws IOException {
     Pair<IMetadataIndexEntry, Long> metadataIndexPair =
-        getLeafMetadataIndexPair(device, measurement, ioSizeRecorder);
+        getLeafMetadataIndexPair(
+            device, deviceMetadataIndexNodeOffset, measurement, ioSizeRecorder);
     if (metadataIndexPair == null) {
       if (ignoreNotExistDevice) {
         return Collections.emptyList();
@@ -744,17 +885,28 @@ public List<TimeseriesMetadata> readTimeseriesMetadata(
 
   /* Get leaf MetadataIndexPair which contains path */
   private Pair<IMetadataIndexEntry, Long> getLeafMetadataIndexPair(
-      IDeviceID device, String measurement, LongConsumer ioSizeRecorder) throws IOException {
+      IDeviceID device,
+      long[] deviceMetadataIndexNodeOffset,
+      String measurement,
+      LongConsumer ioSizeRecorder)
+      throws IOException {
     readFileMetadata(ioSizeRecorder);
     MetadataIndexNode deviceMetadataIndexNode =
         tsFileMetaData.getTableMetadataIndexNode(device.getTableName());
-    Pair<IMetadataIndexEntry, Long> metadataIndexPair =
-        getMetadataAndEndOffsetOfDeviceNode(deviceMetadataIndexNode, device, true, ioSizeRecorder);
-    if (metadataIndexPair == null) {
-      return null;
+    Pair<IMetadataIndexEntry, Long> metadataIndexPair = null;
+    ByteBuffer buffer;
+    if (deviceMetadataIndexNodeOffset == null) {
+      metadataIndexPair =
+          getMetadataAndEndOffsetOfDeviceNode(
+              deviceMetadataIndexNode, device, true, ioSizeRecorder);
+      if (metadataIndexPair == null) {
+        return null;
+      }
+      buffer =
+          readData(metadataIndexPair.left.getOffset(), metadataIndexPair.right, ioSizeRecorder);
+    } else {
+      buffer = readData(deviceMetadataIndexNodeOffset[0], deviceMetadataIndexNodeOffset[1]);
     }
-    ByteBuffer buffer =
-        readData(metadataIndexPair.left.getOffset(), metadataIndexPair.right, ioSizeRecorder);
     MetadataIndexNode metadataIndexNode = deviceMetadataIndexNode;
     if (!metadataIndexNode.getNodeType().equals(MetadataIndexNodeType.LEAF_MEASUREMENT)) {
       try {