enhance: cache reusable format reader metadata (#552)

This change makes repeated reads cheaper by keeping reusable
format-reader metadata around for the lifetime of a Reader. The main
need here is to avoid reopening the same file metadata over and over
when different read paths, projections, predicates, chunk reads, or take
calls touch the same column group files.

The implementation adds a typed metadata cache for Parquet, Iceberg,
Lance, and Vortex readers. Each format now separates the stable file
metadata loading step from the stateful reader creation step, so schema,
row group layout, footer data, Lance fragment metadata, Vortex file
handles, and Iceberg positional delete information can be reused while
each actual reader still gets its own projection and predicate state. At
a high level, the cache now sits behind ReaderImpl and is shared by all
column-group read paths. The important part is that we cache only
immutable metadata, then create a fresh reader from that metadata
whenever the caller brings a different projection or predicate.

```
  ReaderImpl
    owns MetadataCache
      owns FormatReaderMetadataCaches<Parquet, Iceberg, Lance, Vortex>
        owns FormatReaderMetadataCache<ReaderT>
          entries_[cache_key] = immutable FormatReaderMetadata<Payload>
          in_flight_loads_[cache_key] = singleflight loader state

  read flow
    metadata = cache.get_or_open(key, load_metadata)
    reader = FormatReader::create_from_metadata(metadata, file, projection, predicate)
```

The cache also tracenhance: cache reusable format reader metadata

This change makes repeated reads cheaper by keeping reusable
format-reader metadata around for the lifetime of a Reader. The main
need here is to avoid reopening the same file metadata over and over
when different read paths, projections, predicates, chunk reads, or take
calls touch the same column group files.

The implementation adds a typed metadata cache for Parquet, Iceberg,
Lance, and Vortex readers. Each format now separates the stable file
metadata loading step from the stateful reader creation step, so schema,
row group layout, footer data, Lance fragment metadata, Vortex file
handles, and Iceberg positional delete information can be reused while
each actual reader still gets its own projection aks in-flight loads per
key, so concurrent requests for the same file wait on the same metadata
load instead of doing duplicate work. It is wired through ReaderImpl,
ColumnGroupReader, and ColumnGroupLazyReader, so batch reads, chunk
readers, and lazy take reads all use the same metadata source while
keeping the caller-facing reader behavior unchanged.

---------

Signed-off-by: jiaqizho <jiaqi.zhou@zilliz.com>

Author: jiaqizho

cpp/ffi_exports.map

@@ -57,6 +57,7 @@
     loon_properties_reader_record_batch_max_rows;
     loon_properties_reader_record_batch_max_size;
     loon_properties_reader_logical_chunk_rows;
+    loon_properties_reader_metadata_cache_enable;
     loon_properties_reader_parquet_prebuffer_hole_size_limit;
     loon_properties_reader_parquet_prebuffer_range_size_limit;
     loon_properties_reader_vortex_split_row_indices;

cpp/ffi_exports_mac.map

@@ -55,6 +55,7 @@ _loon_properties_writer_vortex_enable_statistics
 _loon_properties_reader_record_batch_max_rows
 _loon_properties_reader_record_batch_max_size
 _loon_properties_reader_logical_chunk_rows
+_loon_properties_reader_metadata_cache_enable
 _loon_properties_reader_parquet_prebuffer_hole_size_limit
 _loon_properties_reader_parquet_prebuffer_range_size_limit
 _loon_properties_reader_vortex_split_row_indices

cpp/include/milvus-storage/ffi_c.h

@@ -125,6 +125,7 @@ FFI_EXPORT extern const char* loon_properties_writer_vortex_enable_statistics;
 FFI_EXPORT extern const char* loon_properties_reader_record_batch_max_rows;
 FFI_EXPORT extern const char* loon_properties_reader_record_batch_max_size;
 FFI_EXPORT extern const char* loon_properties_reader_logical_chunk_rows;
+FFI_EXPORT extern const char* loon_properties_reader_metadata_cache_enable;
 FFI_EXPORT extern const char* loon_properties_reader_parquet_prebuffer_hole_size_limit;
 FFI_EXPORT extern const char* loon_properties_reader_parquet_prebuffer_range_size_limit;
 FFI_EXPORT extern const char* loon_properties_reader_vortex_split_row_indices;
@@ -561,7 +562,10 @@ FFI_EXPORT LoonFFIResult loon_reader_new(const LoonColumnGroups* column_groups,
 
 /**
  * @brief Sets a key retriever callback for dynamic key retrieval
- * use to the KMS(key management system) integration
+ * use to the KMS(key management system) integration.
+ *
+ * This is a setup-only API and is not thread-safe with read operations. Call it
+ * before creating record batch readers, chunk readers, or calling take.
  */
 FFI_EXPORT void loon_reader_set_keyretriever(LoonReaderHandle reader,
                                              const char* (*key_retriever)(const char* metadata));

cpp/include/milvus-storage/format/column_group_lazy_reader.h

@@ -19,6 +19,7 @@
 
 #include "milvus-storage/column_groups.h"
 #include "milvus-storage/properties.h"
+#include "milvus-storage/format/format_reader_cache.h"
 #include "milvus-storage/format/format_reader.h"
 #include "milvus-storage/thread_pool.h"
 
@@ -31,7 +32,7 @@ class ColumnGroupLazyReader {
   /**
    * @brief Take a table from the column group
    *
-   * Thread-safe: each call clones the FormatReader, safe for concurrent use on the same object.
+   * Thread-safe: each call opens independent FormatReaders from reusable metadata.
    *
    * @param row_indices the row indices to take, MUST be uniqued and sorted
    * @return arrow::Result<std::shared_ptr<arrow::Table>>
@@ -44,7 +45,8 @@ class ColumnGroupLazyReader {
       const std::shared_ptr<milvus_storage::api::ColumnGroup>& column_group,
       const milvus_storage::api::Properties& properties,
       const std::vector<std::string>& needed_columns,
-      const std::function<std::string(const std::string&)>& key_retriever);
+      const std::function<std::string(const std::string&)>& key_retriever,
+      const milvus_storage::MetadataCache& cache = milvus_storage::MetadataCache());
 };
 
-};  // namespace milvus_storage::api
+};  // namespace milvus_storage::api

cpp/include/milvus-storage/format/column_group_reader.h

@@ -18,6 +18,7 @@
 
 #include "milvus-storage/column_groups.h"
 #include "milvus-storage/properties.h"
+#include "milvus-storage/format/format_reader_cache.h"
 #include "milvus-storage/format/format_reader.h"
 #include "milvus-storage/thread_pool.h"
 
@@ -34,7 +35,7 @@ class ColumnGroupReader {
   // NOT thread-safe: concurrent calls on the same object may race on the underlying FormatReader.
   virtual arrow::Result<std::shared_ptr<arrow::RecordBatch>> get_chunk(int64_t chunk_index) = 0;
 
-  // Thread-safe: each call clones the FormatReader, safe for concurrent use on the same object.
+  // Thread-safe: each call opens an independent FormatReader from reusable metadata.
   virtual arrow::Result<std::vector<std::shared_ptr<arrow::RecordBatch>>> get_chunks(
       const std::vector<int64_t>& chunk_indices, size_t parallelism = 1) = 0;
 
@@ -59,7 +60,8 @@ class ColumnGroupReader {
       const std::vector<std::string>& needed_columns,
       const milvus_storage::api::Properties& properties,
       const std::function<std::string(const std::string&)>& key_retriever,
-      const std::string& predicate = "");
+      const std::string& predicate = "",
+      const milvus_storage::MetadataCache& cache = milvus_storage::MetadataCache());
 };
 
 }  // namespace milvus_storage::api

cpp/include/milvus-storage/format/format_reader.h

@@ -14,7 +14,12 @@
 
 #pragma once
 
+#include <concepts>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
 #include <memory>
+#include <string>
 #include <vector>
 
 #include <arrow/status.h>
@@ -27,6 +32,8 @@
 
 namespace milvus_storage {
 
+using KeyRetriever = std::function<std::string(const std::string&)>;
+
 struct RowGroupInfo {
   public:
   size_t start_offset;
@@ -36,6 +43,16 @@ struct RowGroupInfo {
   std::string ToString() const;
 };
 
+template <typename Payload>
+struct FormatReaderMetadata {
+  std::string cache_key;
+  std::string path;
+  std::shared_ptr<arrow::Schema> file_schema;
+  std::vector<RowGroupInfo> row_group_infos;
+  uint64_t cache_size = 0;
+  Payload payload;
+};
+
 /**
  * FormatReader is a reader to read the format file.
  * It exists both blocking and streaming read interfaces.
@@ -49,6 +66,12 @@ struct RowGroupInfo {
  */
 class FormatReader {
   public:
+  template <typename ReaderT>
+  using MetaTrait = typename ReaderT::MetaTrait;
+
+  template <typename ReaderT>
+  using MetadataPtr = typename MetaTrait<ReaderT>::MetadataPtr;
+
   virtual ~FormatReader() = default;
 
   // open the format reader, usage to initialize the reader
@@ -88,6 +111,26 @@ class FormatReader {
   // set a predicate string for filtering (default no-op for formats that don't support it)
   virtual void set_predicate(const std::string& /*predicate*/) {}
 
+  // Load reusable file metadata without applying read-time state such as
+  // projection or predicate. The returned metadata is safe to share through
+  // MetadataCache and later reuse to create independent readers.
+  template <typename ReaderT>
+  static arrow::Result<MetadataPtr<ReaderT>> load_metadata(const api::ColumnGroupFile& file,
+                                                           const api::Properties& properties,
+                                                           const KeyRetriever& key_retriever);
+
+  // Create a new stateful reader from cached metadata. The file carries
+  // manifest-owned values such as file_size and footer_size; read_schema,
+  // needed_columns, and predicate are applied here so callers can create
+  // independent readers with different projections or filters from the same
+  // cached metadata.
+  template <typename ReaderT>
+  static arrow::Result<std::shared_ptr<ReaderT>> create_from_metadata(MetadataPtr<ReaderT> metadata,
+                                                                      const api::ColumnGroupFile& file,
+                                                                      const std::shared_ptr<arrow::Schema>& read_schema,
+                                                                      const std::vector<std::string>& needed_columns,
+                                                                      const std::string& predicate);
+
   // create format reader
   static arrow::Result<std::shared_ptr<FormatReader>> create(
       const std::shared_ptr<arrow::Schema>& read_schema,
@@ -99,4 +142,57 @@ class FormatReader {
 
 };  // class FormatReader
 
+template <typename ReaderT>
+concept FormatReaderWithMetadata =
+    std::derived_from<ReaderT, FormatReader> && requires(const api::ColumnGroupFile& file,
+                                                         const api::Properties& properties,
+                                                         const KeyRetriever& key_retriever,
+                                                         typename ReaderT::MetaTrait::MetadataPtr metadata,
+                                                         const api::ColumnGroupFile& metadata_file,
+                                                         const std::shared_ptr<arrow::Schema>& read_schema,
+                                                         const std::vector<std::string>& needed_columns,
+                                                         const std::string& predicate) {
+      typename ReaderT::MetaTrait::Payload;
+      typename ReaderT::MetaTrait::Metadata;
+      typename ReaderT::MetaTrait::MetadataPtr;
+
+      requires std::same_as<typename ReaderT::MetaTrait::Metadata,
+                            FormatReaderMetadata<typename ReaderT::MetaTrait::Payload>>;
+      requires std::same_as<typename ReaderT::MetaTrait::MetadataPtr,
+                            std::shared_ptr<const typename ReaderT::MetaTrait::Metadata>>;
+
+      { ReaderT::MetaTrait::cache_key(file) } -> std::convertible_to<std::string>;
+      {
+        ReaderT::MetaTrait::load_metadata(file, properties, key_retriever)
+      } -> std::same_as<arrow::Result<typename ReaderT::MetaTrait::MetadataPtr>>;
+      {
+        ReaderT::MetaTrait::create_from_metadata(metadata, metadata_file, read_schema, needed_columns, predicate)
+      } -> std::same_as<arrow::Result<std::shared_ptr<ReaderT>>>;
+      { metadata->row_group_infos } -> std::same_as<const std::vector<RowGroupInfo>&>;
+      { metadata->file_schema } -> std::same_as<const std::shared_ptr<arrow::Schema>&>;
+    };
+
+template <typename ReaderT>
+arrow::Result<FormatReader::MetadataPtr<ReaderT>> FormatReader::load_metadata(const api::ColumnGroupFile& file,
+                                                                              const api::Properties& properties,
+                                                                              const KeyRetriever& key_retriever) {
+  static_assert(FormatReaderWithMetadata<ReaderT>,
+                "ReaderT must derive from FormatReader and define MetaTrait with Payload, Metadata, MetadataPtr, "
+                "cache_key, load_metadata, and create_from_metadata.");
+  return ReaderT::MetaTrait::load_metadata(file, properties, key_retriever);
+}
+
+template <typename ReaderT>
+arrow::Result<std::shared_ptr<ReaderT>> FormatReader::create_from_metadata(
+    MetadataPtr<ReaderT> metadata,
+    const api::ColumnGroupFile& file,
+    const std::shared_ptr<arrow::Schema>& read_schema,
+    const std::vector<std::string>& needed_columns,
+    const std::string& predicate) {
+  static_assert(FormatReaderWithMetadata<ReaderT>,
+                "ReaderT must derive from FormatReader and define MetaTrait with Payload, Metadata, MetadataPtr, "
+                "cache_key, load_metadata, and create_from_metadata.");
+  return ReaderT::MetaTrait::create_from_metadata(metadata, file, read_schema, needed_columns, predicate);
+}
+
 }  // namespace milvus_storage

cpp/include/milvus-storage/format/format_reader_cache.h

@@ -0,0 +1,165 @@
+// Copyright 2023 Zilliz
+//
+// Licensed 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.
+
+#pragma once
+
+#include <concepts>
+#include <condition_variable>
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <typeindex>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
+
+#include <arrow/result.h>
+#include <arrow/status.h>
+
+#include "milvus-storage/format/format_reader.h"
+
+namespace milvus_storage {
+
+namespace iceberg {
+class IcebergFormatReader;
+}  // namespace iceberg
+
+namespace lance {
+class LanceTableReader;
+}  // namespace lance
+
+namespace parquet {
+class ParquetFormatReader;
+}  // namespace parquet
+
+namespace vortex {
+class VortexFormatReader;
+}  // namespace vortex
+
+// Thread-safe metadata cache for one concrete FormatReader type.
+// Cached metadata is immutable and can be reused to create independent
+// stateful readers with different projections or predicates.
+template <typename ReaderT>
+class FormatReaderMetadataCache final {
+  static_assert(FormatReaderWithMetadata<ReaderT>,
+                "ReaderT must derive from FormatReader and define MetaTrait with Payload, Metadata, MetadataPtr, "
+                "cache_key, load_metadata, and create_from_metadata.");
+
+  public:
+  using ReaderType = ReaderT;
+  using Trait = typename FormatReader::template MetaTrait<ReaderT>;
+  using MetadataPtr = typename Trait::MetadataPtr;
+  using MetadataLoader = std::function<arrow::Result<MetadataPtr>()>;
+
+  std::optional<MetadataPtr> get(const std::string& key) const;
+
+  arrow::Status add(std::string key, MetadataPtr metadata);
+
+  arrow::Result<MetadataPtr> get_or_open(const std::string& key, const MetadataLoader& load_fn);
+
+  private:
+  struct Entry {
+    MetadataPtr metadata;
+  };
+
+  // Per-key singleflight state. The first cache miss creates this marker and
+  // runs load_fn outside mutex_; waiters for the same key block on cv while
+  // unrelated keys can still load concurrently.
+  struct InFlightLoad {
+    std::condition_variable cv;
+    bool done = false;
+    arrow::Status status = arrow::Status::OK();
+    MetadataPtr metadata;
+  };
+
+  mutable std::mutex mutex_;
+  std::unordered_map<std::string, Entry> entries_;
+  std::unordered_map<std::string, std::shared_ptr<InFlightLoad>> in_flight_loads_;
+};
+
+// Owns one typed metadata cache for each ReaderT in the template list.
+// Callers still retrieve caches statically with get<ReaderT>(); this class only
+// groups the per-format caches into one value that can be embedded elsewhere.
+template <typename... ReaderTs>
+class FormatReaderMetadataCaches final {
+  public:
+  FormatReaderMetadataCaches() : caches_(std::make_shared<FormatReaderMetadataCache<ReaderTs>>()...) {}
+
+  template <typename ReaderT>
+  [[nodiscard]] std::shared_ptr<FormatReaderMetadataCache<ReaderT>> get() const {
+    static_assert((std::same_as<ReaderT, ReaderTs> || ...), "ReaderT must be a supported metadata cache reader type");
+    return std::get<std::shared_ptr<FormatReaderMetadataCache<ReaderT>>>(caches_);
+  }
+
+  private:
+  std::tuple<std::shared_ptr<FormatReaderMetadataCache<ReaderTs>>...> caches_;
+};
+
+// Public cache handle carried by ReaderImpl and passed down to column-group
+// readers. Concrete reader headers are intentionally not included here, so
+// installed consumers can include public reader headers without private bridge
+// headers from the source tree.
+class MetadataCache final {
+  public:
+  explicit MetadataCache(bool enabled = true);
+
+  [[nodiscard]] bool enabled() const { return enabled_; }
+
+  template <typename ReaderT>
+  [[nodiscard]] std::shared_ptr<FormatReaderMetadataCache<ReaderT>> get() const {
+    static_assert(FormatReaderWithMetadata<ReaderT>,
+                  "ReaderT must derive from FormatReader and define MetaTrait with Payload, Metadata, MetadataPtr, "
+                  "cache_key, load_metadata, and create_from_metadata.");
+
+    std::lock_guard<std::mutex> lock(state_->mutex);
+    auto [it, inserted] = state_->caches.try_emplace(std::type_index(typeid(ReaderT)));
+    if (inserted || !it->second) {
+      it->second = std::make_shared<FormatReaderMetadataCache<ReaderT>>();
+    }
+    return std::static_pointer_cast<FormatReaderMetadataCache<ReaderT>>(it->second);
+  }
+
+  template <typename Visitor>
+  auto dispatch(const std::string& format, Visitor&& visitor) const {
+    using ReturnT = decltype(std::forward<Visitor>(visitor)(get<parquet::ParquetFormatReader>()));
+
+    if (format == LOON_FORMAT_PARQUET) {
+      return std::forward<Visitor>(visitor)(get<parquet::ParquetFormatReader>());
+    }
+    if (format == LOON_FORMAT_VORTEX) {
+      return std::forward<Visitor>(visitor)(get<vortex::VortexFormatReader>());
+    }
+    if (format == LOON_FORMAT_LANCE_TABLE) {
+      return std::forward<Visitor>(visitor)(get<lance::LanceTableReader>());
+    }
+    if (format == LOON_FORMAT_ICEBERG_TABLE) {
+      return std::forward<Visitor>(visitor)(get<iceberg::IcebergFormatReader>());
+    }
+
+    return ReturnT(arrow::Status::Invalid("Unknown column group format: ", format));
+  }
+
+  private:
+  struct State {
+    mutable std::mutex mutex;
+    std::unordered_map<std::type_index, std::shared_ptr<void>> caches;
+  };
+
+  bool enabled_ = true;
+  std::shared_ptr<State> state_ = std::make_shared<State>();
+};
+
+}  // namespace milvus_storage