add cache table schema map option (#369)

* add cache table schema map option

* modify method name

* fix ut

java/tsfile/src/main/java/org/apache/tsfile/compatibility/CompatibilityUtils.java

@@ -42,6 +42,8 @@ private CompatibilityUtils() {
         org.apache.tsfile.common.conf.TSFileConfig.VERSION_NUMBER_V3;
     v3DeserializeConfig.tsFileMetadataBufferDeserializer =
         CompatibilityUtils::deserializeTsFileMetadataFromV3;
+    v3DeserializeConfig.cacheTableSchemaMapTsFileMetadataBufferDeserializer =
+        CompatibilityUtils::deserializeTsFileMetadataFromV3;
     v3DeserializeConfig.deviceIDBufferDeserializer =
         ((buffer, context) -> {
           final PlainDeviceID deviceID = PlainDeviceID.deserialize(buffer);

java/tsfile/src/main/java/org/apache/tsfile/compatibility/DeserializeConfig.java

@@ -36,7 +36,9 @@
 public class DeserializeConfig {
   public byte versionNumber = org.apache.tsfile.common.conf.TSFileConfig.VERSION_NUMBER;
   public BufferDeserializer<TsFileMetadata> tsFileMetadataBufferDeserializer =
-      TsFileMetadata::deserializeFrom;
+      TsFileMetadata::deserializeWithoutCacheTableSchemaMap;
+  public BufferDeserializer<TsFileMetadata> cacheTableSchemaMapTsFileMetadataBufferDeserializer =
+      TsFileMetadata::deserializeAndCacheTableSchemaMap;
 
   public BufferDeserializer<MetadataIndexNode> deviceMetadataIndexNodeBufferDeserializer =
       (buffer, context) -> MetadataIndexNode.deserializeFrom(buffer, true, context);

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

@@ -46,6 +46,7 @@ public class TsFileMetadata {
   // List of <name, offset, childMetadataIndexType>
   private Map<String, MetadataIndexNode> tableMetadataIndexNodeMap;
   private Map<String, TableSchema> tableSchemaMap;
+  private boolean hasTableSchemaMapCache;
   private Map<String, String> tsFileProperties;
 
   // offset of MetaMarker.SEPARATOR
@@ -57,13 +58,24 @@ public class TsFileMetadata {
 
   private String encryptType;
 
+  public static TsFileMetadata deserializeAndCacheTableSchemaMap(
+      ByteBuffer buffer, DeserializeConfig context) {
+    return deserializeFrom(buffer, context, true);
+  }
+
+  public static TsFileMetadata deserializeWithoutCacheTableSchemaMap(
+      ByteBuffer buffer, DeserializeConfig context) {
+    return deserializeFrom(buffer, context, false);
+  }
+
   /**
    * deserialize data from the buffer.
    *
    * @param buffer -buffer use to deserialize
    * @return -an instance of TsFileMetaData
    */
-  public static TsFileMetadata deserializeFrom(ByteBuffer buffer, DeserializeConfig context) {
+  public static TsFileMetadata deserializeFrom(
+      ByteBuffer buffer, DeserializeConfig context, boolean needTableSchemaMap) {
     TsFileMetadata fileMetaData = new TsFileMetadata();
 
     int startPos = buffer.position();
@@ -84,10 +96,13 @@ public static TsFileMetadata deserializeFrom(ByteBuffer buffer, DeserializeConfi
     for (int i = 0; i < tableSchemaNum; i++) {
       String tableName = ReadWriteIOUtils.readVarIntString(buffer);
       TableSchema tableSchema = context.tableSchemaBufferDeserializer.deserialize(buffer, context);
-      tableSchema.setTableName(tableName);
-      tableSchemaMap.put(tableName, tableSchema);
+      if (needTableSchemaMap) {
+        tableSchema.setTableName(tableName);
+        tableSchemaMap.put(tableName, tableSchema);
+      }
     }
     fileMetaData.setTableSchemaMap(tableSchemaMap);
+    fileMetaData.hasTableSchemaMapCache = needTableSchemaMap;
 
     // metaOffset
     long metaOffset = ReadWriteIOUtils.readLong(buffer);
@@ -267,6 +282,7 @@ public void setTableMetadataIndexNodeMap(
 
   public void setTableSchemaMap(Map<String, TableSchema> tableSchemaMap) {
     this.tableSchemaMap = tableSchemaMap;
+    this.hasTableSchemaMapCache = true;
   }
 
   public Map<String, MetadataIndexNode> getTableMetadataIndexNodeMap() {
@@ -281,6 +297,10 @@ public MetadataIndexNode getTableMetadataIndexNode(String tableName) {
     return metadataIndexNode;
   }
 
+  public boolean hasTableSchemaMapCache() {
+    return hasTableSchemaMapCache;
+  }
+
   public Map<String, TableSchema> getTableSchemaMap() {
     return tableSchemaMap;
   }

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

@@ -22,6 +22,7 @@
 import org.apache.tsfile.common.conf.TSFileConfig;
 import org.apache.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.common.constant.TsFileConstant;
+import org.apache.tsfile.compatibility.BufferDeserializer;
 import org.apache.tsfile.compatibility.CompatibilityUtils;
 import org.apache.tsfile.compatibility.DeserializeConfig;
 import org.apache.tsfile.compress.IUnCompressor;
@@ -52,6 +53,7 @@
 import org.apache.tsfile.file.metadata.MeasurementMetadataIndexEntry;
 import org.apache.tsfile.file.metadata.MetadataIndexNode;
 import org.apache.tsfile.file.metadata.TableDeviceMetadata;
+import org.apache.tsfile.file.metadata.TableSchema;
 import org.apache.tsfile.file.metadata.TimeseriesMetadata;
 import org.apache.tsfile.file.metadata.TsFileMetadata;
 import org.apache.tsfile.file.metadata.enums.CompressionType;
@@ -134,6 +136,7 @@ public class TsFileSequenceReader implements AutoCloseable {
   private byte fileVersion;
 
   private DeserializeConfig deserializeConfig = new DeserializeConfig();
+  private volatile boolean cacheTableSchemaMap = false;
 
   /**
    * Create a file reader of the given file. The reader will read the tail of the file to get the
@@ -285,6 +288,10 @@ private void checkFileVersion() throws FileVersionTooOldException {
     }
   }
 
+  public void setEnableCacheTableSchemaMap() {
+    this.cacheTableSchemaMap = true;
+  }
+
   public void loadMetadataSize() throws IOException {
     loadMetadataSize(null);
   }
@@ -394,9 +401,7 @@ public TsFileMetadata readFileMetadata(LongConsumer ioSizeRecorder) throws IOExc
       if (tsFileMetaData == null) {
         synchronized (this) {
           if (tsFileMetaData == null) {
-            tsFileMetaData =
-                deserializeConfig.tsFileMetadataBufferDeserializer.deserialize(
-                    readData(fileMetadataPos, fileMetadataSize, ioSizeRecorder), deserializeConfig);
+            tsFileMetaData = forceReadFileMetadata(cacheTableSchemaMap, ioSizeRecorder);
           }
         }
       }
@@ -409,6 +414,34 @@ public TsFileMetadata readFileMetadata(LongConsumer ioSizeRecorder) throws IOExc
     return tsFileMetaData;
   }
 
+  public Map<String, TableSchema> getTableSchemaMap() throws IOException {
+    return getTableSchemaMap(null);
+  }
+
+  public Map<String, TableSchema> getTableSchemaMap(LongConsumer ioSizeRecorder)
+      throws IOException {
+    if (tsFileMetaData != null && tsFileMetaData.hasTableSchemaMapCache()) {
+      return tsFileMetaData.getTableSchemaMap();
+    }
+    TsFileMetadata tempTsFileMetadata = forceReadFileMetadata(true, ioSizeRecorder);
+    if (cacheTableSchemaMap) {
+      synchronized (this) {
+        this.tsFileMetaData = tempTsFileMetadata;
+      }
+    }
+    return tempTsFileMetadata.getTableSchemaMap();
+  }
+
+  private TsFileMetadata forceReadFileMetadata(
+      boolean needTableSchemaMap, LongConsumer ioSizeRecorder) throws IOException {
+    ByteBuffer buffer = readData(fileMetadataPos, fileMetadataSize, ioSizeRecorder);
+    BufferDeserializer<TsFileMetadata> deserializer =
+        needTableSchemaMap
+            ? deserializeConfig.cacheTableSchemaMapTsFileMetadataBufferDeserializer
+            : deserializeConfig.tsFileMetadataBufferDeserializer;
+    return deserializer.deserialize(buffer, deserializeConfig);
+  }
+
   /**
    * This function does not modify the position of the file reader.
    *

java/tsfile/src/main/java/org/apache/tsfile/read/controller/IMetadataQuerier.java

@@ -24,6 +24,7 @@
 import org.apache.tsfile.file.metadata.IChunkMetadata;
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.file.metadata.MetadataIndexNode;
+import org.apache.tsfile.file.metadata.TableSchema;
 import org.apache.tsfile.file.metadata.TsFileMetadata;
 import org.apache.tsfile.read.common.Path;
 import org.apache.tsfile.read.common.TimeRange;
@@ -55,6 +56,8 @@ List<List<IChunkMetadata>> getChunkMetadataLists(
 
   TsFileMetadata getWholeFileMetadata();
 
+  Map<String, TableSchema> getTableSchemaMap();
+
   /**
    * this will load all chunk metadata of given paths into cache.
    *

java/tsfile/src/main/java/org/apache/tsfile/read/controller/MetadataQuerierByFileImpl.java

@@ -27,6 +27,7 @@
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.file.metadata.ITimeSeriesMetadata;
 import org.apache.tsfile.file.metadata.MetadataIndexNode;
+import org.apache.tsfile.file.metadata.TableSchema;
 import org.apache.tsfile.file.metadata.TimeseriesMetadata;
 import org.apache.tsfile.file.metadata.TsFileMetadata;
 import org.apache.tsfile.read.TsFileSequenceReader;
@@ -60,11 +61,14 @@ public class MetadataQuerierByFileImpl implements IMetadataQuerier {
   private LRUCache<Pair<IDeviceID, String>, List<IChunkMetadata>> deviceIdChunkMetadataCache;
 
   private TsFileSequenceReader tsFileReader;
+  private Map<String, TableSchema> tableSchemaMap;
 
   /** Constructor of MetadataQuerierByFileImpl. */
   public MetadataQuerierByFileImpl(TsFileSequenceReader tsFileReader) throws IOException {
     this.tsFileReader = tsFileReader;
+    this.tsFileReader.setEnableCacheTableSchemaMap();
     this.fileMetaData = tsFileReader.readFileMetadata();
+    this.tableSchemaMap = tsFileReader.getTableSchemaMap();
     deviceIdChunkMetadataCache =
         new LRUCache<Pair<IDeviceID, String>, List<IChunkMetadata>>(CACHED_ENTRY_NUMBER) {
           @Override
@@ -127,6 +131,11 @@ public TsFileMetadata getWholeFileMetadata() {
     return fileMetaData;
   }
 
+  @Override
+  public Map<String, TableSchema> getTableSchemaMap() {
+    return tableSchemaMap;
+  }
+
   @Override
   @SuppressWarnings("squid:S3776") // Suppress high Cognitive Complexity warning
   public void loadChunkMetaDatas(List<Path> paths) throws IOException {

java/tsfile/src/main/java/org/apache/tsfile/read/query/executor/TableQueryExecutor.java

@@ -76,7 +76,7 @@ public TsBlockReader query(
       throws ReadProcessException {
     TsFileMetadata fileMetadata = metadataQuerier.getWholeFileMetadata();
     MetadataIndexNode tableRoot = fileMetadata.getTableMetadataIndexNode(tableName);
-    TableSchema tableSchema = fileMetadata.getTableSchemaMap().get(tableName);
+    TableSchema tableSchema = metadataQuerier.getTableSchemaMap().get(tableName);
     if (tableRoot == null || tableSchema == null) {
       return new EmptyTsBlockReader();
     }

java/tsfile/src/main/java/org/apache/tsfile/read/v4/DeviceTableModelReader.java

@@ -25,7 +25,6 @@
 import org.apache.tsfile.exception.write.NoMeasurementException;
 import org.apache.tsfile.exception.write.NoTableException;
 import org.apache.tsfile.file.metadata.TableSchema;
-import org.apache.tsfile.file.metadata.TsFileMetadata;
 import org.apache.tsfile.read.TsFileSequenceReader;
 import org.apache.tsfile.read.controller.CachedChunkLoaderImpl;
 import org.apache.tsfile.read.controller.IChunkLoader;
@@ -57,6 +56,7 @@ public class DeviceTableModelReader implements ITsFileReader {
 
   public DeviceTableModelReader(File file) throws IOException {
     this.fileReader = new TsFileSequenceReader(file.getPath());
+    this.fileReader.setEnableCacheTableSchemaMap();
     this.metadataQuerier = new MetadataQuerierByFileImpl(fileReader);
     this.chunkLoader = new CachedChunkLoaderImpl(fileReader);
     this.queryExecutor =
@@ -66,23 +66,21 @@ public DeviceTableModelReader(File file) throws IOException {
 
   @TsFileApi
   public List<TableSchema> getAllTableSchema() throws IOException {
-    Map<String, TableSchema> tableSchemaMap = fileReader.readFileMetadata().getTableSchemaMap();
+    Map<String, TableSchema> tableSchemaMap = fileReader.getTableSchemaMap();
     return new ArrayList<>(tableSchemaMap.values());
   }
 
   @TsFileApi
   public Optional<TableSchema> getTableSchemas(String tableName) throws IOException {
-    TsFileMetadata tsFileMetadata = fileReader.readFileMetadata();
-    Map<String, TableSchema> tableSchemaMap = tsFileMetadata.getTableSchemaMap();
+    Map<String, TableSchema> tableSchemaMap = fileReader.getTableSchemaMap();
     return Optional.ofNullable(tableSchemaMap.get(tableName.toLowerCase()));
   }
 
   @TsFileApi
   public ResultSet query(String tableName, List<String> columnNames, long startTime, long endTime)
       throws IOException, NoTableException, NoMeasurementException, ReadProcessException {
     String lowerCaseTableName = tableName.toLowerCase();
-    TsFileMetadata tsFileMetadata = fileReader.readFileMetadata();
-    TableSchema tableSchema = tsFileMetadata.getTableSchemaMap().get(lowerCaseTableName);
+    TableSchema tableSchema = fileReader.getTableSchemaMap().get(lowerCaseTableName);
     if (tableSchema == null) {
       throw new NoTableException(tableName);
     }

java/tsfile/src/main/java/org/apache/tsfile/utils/TsFileSketchTool.java

@@ -202,13 +202,10 @@ private void printTsFileMetadata(TsFileMetadata tsFileMetaData) {
       }
 
       // table schema
-      printlnBoth(
-          pw,
-          String.format("%20s", pos)
-              + "|\tTableSchemaCnt="
-              + tsFileMetaData.getTableSchemaMap().size());
+      Map<String, TableSchema> tableSchemaMap = reader.getTableSchemaMap();
+      printlnBoth(pw, String.format("%20s", pos) + "|\tTableSchemaCnt=" + tableSchemaMap.size());
       pos += Integer.BYTES;
-      for (Entry<String, TableSchema> entry : tsFileMetaData.getTableSchemaMap().entrySet()) {
+      for (Entry<String, TableSchema> entry : tableSchemaMap.entrySet()) {
         final String tableName = entry.getKey();
         final TableSchema tableSchema = entry.getValue();
 

java/tsfile/src/test/java/org/apache/tsfile/file/metadata/TsFileMetadataTest.java

@@ -72,7 +72,7 @@ private TsFileMetadata deSerialized() {
       ByteBuffer buffer = ByteBuffer.allocate((int) channel.size());
       channel.read(buffer);
       buffer.rewind();
-      metaData = TsFileMetadata.deserializeFrom(buffer, deserializeConfig);
+      metaData = TsFileMetadata.deserializeAndCacheTableSchemaMap(buffer, deserializeConfig);
       return metaData;
     } catch (IOException e) {
       e.printStackTrace();

java/tsfile/src/test/java/org/apache/tsfile/read/TsFileSequenceReaderTest.java

@@ -32,6 +32,7 @@
 import org.apache.tsfile.file.metadata.IChunkMetadata;
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.file.metadata.IDeviceID.Factory;
+import org.apache.tsfile.file.metadata.TableSchema;
 import org.apache.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.tsfile.read.common.Path;
 import org.apache.tsfile.utils.BloomFilter;
@@ -40,10 +41,13 @@
 import org.apache.tsfile.utils.TsFileGeneratorUtils;
 import org.apache.tsfile.write.TsFileWriter;
 import org.apache.tsfile.write.record.TSRecord;
+import org.apache.tsfile.write.record.Tablet;
 import org.apache.tsfile.write.record.datapoint.DoubleDataPoint;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.apache.tsfile.write.schema.Schema;
+import org.apache.tsfile.write.v4.ITsFileWriter;
+import org.apache.tsfile.write.v4.TsFileWriterBuilder;
 
 import org.junit.After;
 import org.junit.Assert;
@@ -53,7 +57,9 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.file.Files;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -225,4 +231,35 @@ public void testReadEmptyPageInSelfCheck() throws IOException, WriteProcessExcep
           reader.selfCheck(new Schema(), new ArrayList<>(), false));
     }
   }
+
+  @Test
+  public void testGetTableSchemaMap() throws IOException, WriteProcessException {
+    File file = new File(FILE_PATH);
+    try {
+      tsFile.close();
+      Files.deleteIfExists(file.toPath());
+    } catch (IOException ignored) {
+    }
+    TableSchema tableSchema =
+        new TableSchema(
+            "t1",
+            Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT32)),
+            Collections.singletonList(Tablet.ColumnCategory.FIELD));
+    try (ITsFileWriter writer =
+        new TsFileWriterBuilder().tableSchema(tableSchema).file(file).build()) {
+      Tablet tablet =
+          new Tablet(Collections.singletonList("s1"), Collections.singletonList(TSDataType.INT32));
+      tablet.addTimestamp(0, 1);
+      tablet.addValue("s1", 0, 1);
+      writer.write(tablet);
+    }
+    try (TsFileSequenceReader reader = new TsFileSequenceReader(FILE_PATH)) {
+      Assert.assertFalse(reader.readFileMetadata().hasTableSchemaMapCache());
+      Assert.assertEquals(1, reader.getTableSchemaMap().size());
+      Assert.assertFalse(reader.readFileMetadata().hasTableSchemaMapCache());
+      reader.setEnableCacheTableSchemaMap();
+      Assert.assertEquals(1, reader.getTableSchemaMap().size());
+      Assert.assertTrue(reader.readFileMetadata().hasTableSchemaMapCache());
+    }
+  }
 }

java/tsfile/src/test/java/org/apache/tsfile/write/TsFileWriteApiTest.java

@@ -923,7 +923,7 @@ public void writeTableTsFileWithUpperCaseColumns() throws IOException, WriteProc
       writer.writeTable(tablet);
     }
     try (TsFileSequenceReader reader = new TsFileSequenceReader(f.getPath())) {
-      Map<String, TableSchema> tableSchemaMap = reader.readFileMetadata().getTableSchemaMap();
+      Map<String, TableSchema> tableSchemaMap = reader.getTableSchemaMap();
       TableSchema tableSchemaInTsFile = tableSchemaMap.get("table1");
       Assert.assertNotNull(tableSchemaInTsFile);
       for (IMeasurementSchema columnSchema : tableSchemaInTsFile.getColumnSchemas()) {