feat: [Tiered] Make load list work as warmup hint (#42490)

Related to #42489
See also #41435

This PR's main target is to make partial load field list work as caching
layer warmup policy hint. If user specify load field list, the fields
not included in the list shall use `disabled` warmup policy and be able
to lazily loaded if any read op uses them.

The major changes are listed here:
- Pass load list to segcore and creating collection&schema
- Add util functions to check field shall be proactively loaded
- Adapt storage v2 column group, which may lead to hint fail if columns
share same group

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/core/src/common/Schema.cpp

@@ -91,7 +91,7 @@ Schema::ConvertToArrowSchema() const {
 }
 
 std::unique_ptr<std::vector<FieldMeta>>
-Schema::absent_fields(Schema& old_schema) const {
+Schema::AbsentFields(Schema& old_schema) const {
     std::vector<FieldMeta> result;
     for (const auto& [field_id, field_meta] : fields_) {
         auto it = old_schema.fields_.find(field_id);

internal/core/src/common/Schema.h

@@ -16,10 +16,12 @@
 
 #pragma once
 
+#include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -27,6 +29,7 @@
 #include "FieldMeta.h"
 #include "boost/stacktrace/frame.hpp"
 #include "boost/stacktrace/stacktrace_fwd.hpp"
+#include "common/Types.h"
 #include "pb/schema.pb.h"
 #include "log/Log.h"
 #include "Consts.h"
@@ -238,7 +241,7 @@ class Schema {
     }
 
     const std::unordered_map<FieldId, FieldMeta>
-    get_field_metas(std::vector<FieldId> field_ids) {
+    get_field_metas(const std::vector<FieldId>& field_ids) {
         std::unordered_map<FieldId, FieldMeta> field_metas;
         for (const auto& field_id : field_ids) {
             field_metas.emplace(field_id, operator[](field_id));
@@ -272,6 +275,28 @@ class Schema {
     const ArrowSchemaPtr
     ConvertToArrowSchema() const;
 
+    void
+    UpdateLoadFields(const std::vector<int64_t>& field_ids) {
+        load_fields_.clear();
+        for (auto field_id : field_ids) {
+            load_fields_.emplace(field_id);
+        }
+    }
+
+    bool
+    ShallLoadField(FieldId field_id) {
+        return load_fields_.empty() || load_fields_.count(field_id) > 0;
+    }
+
+    std::vector<int64_t>
+    load_fields() {
+        auto fields = std::vector<int64_t>();
+        for (auto field_id : field_ids_) {
+            fields.emplace_back(field_id.get());
+        }
+        return std::move(fields);
+    }
+
  public:
     static std::shared_ptr<Schema>
     ParseFrom(const milvus::proto::schema::CollectionSchema& schema_proto);
@@ -290,7 +315,7 @@ class Schema {
     }
 
     std::unique_ptr<std::vector<FieldMeta>>
-    absent_fields(Schema& old_schema) const;
+    AbsentFields(Schema& old_schema) const;
 
  private:
     int64_t debug_id = START_USER_FIELDID;
@@ -306,6 +331,10 @@ class Schema {
     std::optional<FieldId> primary_field_id_opt_;
     std::optional<FieldId> dynamic_field_id_opt_;
 
+    // field partial load list
+    // work as hint now
+    std::unordered_set<FieldId> load_fields_;
+
     // schema_version_, currently marked with update timestamp
     uint64_t schema_version_;
 };

internal/core/src/mmap/ChunkedColumnInterface.h

@@ -28,7 +28,8 @@ class ChunkedColumnInterface {
 
     // Default implementation does nothing.
     virtual void
-    ManualEvictCache() const {}
+    ManualEvictCache() const {
+    }
 
     // Get raw data pointer of a specific chunk
     virtual cachinglayer::PinWrapper<const char*>

internal/core/src/mmap/Types.h

@@ -31,10 +31,12 @@ struct FieldDataInfo {
 
     FieldDataInfo(int64_t field_id,
                   size_t row_count,
-                  std::string mmap_dir_path = "")
+                  std::string mmap_dir_path = "",
+                  bool in_list = false)
         : field_id(field_id),
           row_count(row_count),
-          mmap_dir_path(std::move(mmap_dir_path)) {
+          mmap_dir_path(std::move(mmap_dir_path)),
+          in_load_list(in_list) {
         arrow_reader_channel = std::make_shared<ArrowReaderChannel>();
     }
 
@@ -42,5 +44,6 @@ struct FieldDataInfo {
     size_t row_count;
     std::string mmap_dir_path;
     std::shared_ptr<ArrowReaderChannel> arrow_reader_channel;
+    bool in_load_list = false;
 };
 }  // namespace milvus

internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp

@@ -269,12 +269,22 @@ ChunkedSegmentSealedImpl::load_column_group_data_internal(
             metadata->GetGroupFieldIDList().GetFieldIDList(
                 column_group_id.get());
         std::vector<FieldId> milvus_field_ids;
+
+        // if multiple fields share same column group
+        // hint for not loading certain field shall not be working for now
+        // warmup will be disabled only when all columns are not in load list
+        bool merged_in_load_list = false;
         for (int i = 0; i < field_id_list.size(); ++i) {
             milvus_field_ids.emplace_back(field_id_list.Get(i));
+            merged_in_load_list =
+                merged_in_load_list ||
+                schema_->ShallLoadField(FieldId(field_id_list.Get(i)));
         }
 
-        auto column_group_info = FieldDataInfo(
-            column_group_id.get(), num_rows, load_info.mmap_dir_path);
+        auto column_group_info = FieldDataInfo(column_group_id.get(),
+                                               num_rows,
+                                               load_info.mmap_dir_path,
+                                               merged_in_load_list);
         LOG_INFO("segment {} loads column group {} with num_rows {}",
                  this->get_segment_id(),
                  column_group_id.get(),
@@ -321,8 +331,10 @@ ChunkedSegmentSealedImpl::load_field_data_internal(
 
         auto field_id = FieldId(id);
 
-        auto field_data_info =
-            FieldDataInfo(field_id.get(), num_rows, load_info.mmap_dir_path);
+        auto field_data_info = FieldDataInfo(field_id.get(),
+                                             num_rows,
+                                             load_info.mmap_dir_path,
+                                             schema_->ShallLoadField(field_id));
         LOG_INFO("segment {} loads field {} with num_rows {}, sorted by pk {}",
                  this->get_segment_id(),
                  field_id.get(),
@@ -1872,7 +1884,7 @@ ChunkedSegmentSealedImpl::Reopen(SchemaPtr sch) {
     index_ready_bitset_.resize(sch->size());
     binlog_index_bitset_.resize(sch->size());
 
-    auto absent_fields = sch->absent_fields(*schema_);
+    auto absent_fields = sch->AbsentFields(*schema_);
     for (const auto& field_meta : *absent_fields) {
         // vector field is not supported to be "added field", thus if a vector
         // field is absent, it means for some reason we want to skip loading this

internal/core/src/segcore/Collection.cpp

@@ -79,8 +79,14 @@ Collection::parse_schema(const void* schema_proto_blob,
 
     AssertInfo(suc, "parse schema proto failed");
 
+    auto old_schema = schema_;
+
     schema_ = Schema::ParseFrom(collection_schema);
     schema_->set_schema_version(version);
+
+    if (old_schema) {
+        schema_->UpdateLoadFields(old_schema->load_fields());
+    }
 }
 
 }  // namespace milvus::segcore

internal/core/src/segcore/SegmentGrowingImpl.cpp

@@ -1197,7 +1197,7 @@ void
 SegmentGrowingImpl::Reopen(SchemaPtr sch) {
     std::unique_lock lck(mutex_);
 
-    auto absent_fields = sch->absent_fields(*schema_);
+    auto absent_fields = sch->AbsentFields(*schema_);
 
     for (const auto& field_meta : *absent_fields) {
         fill_empty_field(field_meta);

internal/core/src/segcore/Utils.cpp

@@ -978,8 +978,12 @@ upper_bound(const ConcurrentVector<Timestamp>& timestamps,
 
 // Get the globally configured cache warmup policy for the given content type.
 CacheWarmupPolicy
-getCacheWarmupPolicy(bool is_vector, bool is_index) {
+getCacheWarmupPolicy(bool is_vector, bool is_index, bool in_load_list) {
     auto& manager = milvus::cachinglayer::Manager::GetInstance();
+    // if field not in load list(hint), disable warmup
+    if (!in_load_list) {
+        return CacheWarmupPolicy::CacheWarmupPolicy_Disable;
+    }
     if (is_index) {
         return is_vector ? manager.getVectorIndexCacheWarmupPolicy()
                          : manager.getScalarIndexCacheWarmupPolicy();

internal/core/src/segcore/Utils.h

@@ -131,6 +131,6 @@ upper_bound(const ConcurrentVector<Timestamp>& timestamps,
             Timestamp value);
 
 CacheWarmupPolicy
-getCacheWarmupPolicy(bool is_vector, bool is_index);
+getCacheWarmupPolicy(bool is_vector, bool is_index, bool in_load_list = true);
 
 }  // namespace milvus::segcore

internal/core/src/segcore/collection_c.cpp

@@ -49,6 +49,21 @@ UpdateSchema(CCollection collection,
     }
 }
 
+CStatus
+UpdateLoadFields(CCollection collection,
+                 const int64_t* field_ids,
+                 const int64_t length) {
+    try {
+        auto col = static_cast<milvus::segcore::Collection*>(collection);
+
+        col->get_schema()->UpdateLoadFields(
+            std::vector<int64_t>(field_ids, field_ids + length));
+        return milvus::SuccessCStatus();
+    } catch (std::exception& e) {
+        return milvus::FailureCStatus(&e);
+    }
+}
+
 CStatus
 SetIndexMeta(CCollection collection,
              const void* proto_blob,

internal/core/src/segcore/collection_c.h

@@ -31,6 +31,11 @@ UpdateSchema(CCollection collection,
              const int64_t length,
              const uint64_t version);
 
+CStatus
+UpdateLoadFields(CCollection collection,
+                 const int64_t* field_ids,
+                 const int64_t length);
+
 CStatus
 SetIndexMeta(CCollection collection,
              const void* proto_blob,

internal/core/src/segcore/storagev1translator/ChunkTranslator.cpp

@@ -44,7 +44,8 @@ ChunkTranslator::ChunkTranslator(
                      : milvus::cachinglayer::StorageType::MEMORY,
             milvus::segcore::getCacheWarmupPolicy(
                 IsVectorDataType(field_meta.get_data_type()),
-                /* is_index */ false),
+                /* is_index */ false,
+                /* in_load_list*/ field_data_info.in_load_list),
             /* support_eviction */ false) {
     AssertInfo(!SystemProperty::Instance().IsSystem(FieldId(field_id_)),
                "ChunkTranslator not supported for system field");
@@ -68,7 +69,7 @@ ChunkTranslator::load_chunk(milvus::cachinglayer::cid_t cid) {
     pool.Submit(LoadArrowReaderFromRemote,
                 std::vector<std::string>{files_and_rows_[cid].first},
                 channel);
-    LOG_DEBUG("segment {} submits load field {} chunk {} task to thread pool",
+    LOG_INFO("segment {} submits load field {} chunk {} task to thread pool",
              segment_id_,
              field_id_,
              cid);

internal/core/src/segcore/storagev1translator/DefaultValueChunkTranslator.cpp

@@ -30,7 +30,8 @@ DefaultValueChunkTranslator::DefaultValueChunkTranslator(
                      : milvus::cachinglayer::StorageType::MEMORY,
             milvus::segcore::getCacheWarmupPolicy(
                 IsVectorDataType(field_meta.get_data_type()),
-                /* is_index */ false),
+                /* is_index */ false,
+                /* in_load_list, set to false to reduce memory usage */ false),
             /* support_eviction */ false) {
     meta_.num_rows_until_chunk_.push_back(0);
     meta_.num_rows_until_chunk_.push_back(field_data_info.row_count);

internal/core/src/segcore/storagev1translator/InterimSealedIndexTranslator.cpp

@@ -62,7 +62,12 @@ InterimSealedIndexTranslator::get_cells(
                                               vec_data_](size_t id) {
             const void* data;
             int64_t data_id = id;
-            field_raw_data_ptr->BulkValueAt([&data, &data_id](const char* value, size_t i) {data = static_cast<const void*>(value);}, &data_id, 1);                       
+            field_raw_data_ptr->BulkValueAt(
+                [&data, &data_id](const char* value, size_t i) {
+                    data = static_cast<const void*>(value);
+                },
+                &data_id,
+                1);
             return data;
         };
 
@@ -71,28 +76,29 @@ InterimSealedIndexTranslator::get_cells(
                 index_type_,
                 metric_type_,
                 knowhere::Version::GetCurrentVersion().VersionNumber(),
-                view_data, false);
+                view_data,
+                false);
         } else if (vec_data_type_ == DataType::VECTOR_FLOAT16) {
-            vec_index =
-                std::make_unique<index::VectorMemIndex<knowhere::fp16>>(
+            vec_index = std::make_unique<index::VectorMemIndex<knowhere::fp16>>(
                 index_type_,
                 metric_type_,
                 knowhere::Version::GetCurrentVersion().VersionNumber(),
-                    view_data, false);
-        } else if (vec_data_type_ ==
-                    DataType::VECTOR_BFLOAT16) {
-            vec_index =
-                std::make_unique<index::VectorMemIndex<knowhere::bf16>>(
+                view_data,
+                false);
+        } else if (vec_data_type_ == DataType::VECTOR_BFLOAT16) {
+            vec_index = std::make_unique<index::VectorMemIndex<knowhere::bf16>>(
                 index_type_,
                 metric_type_,
                 knowhere::Version::GetCurrentVersion().VersionNumber(),
-                    view_data, false);
+                view_data,
+                false);
         }
     } else {
         vec_index = std::make_unique<index::VectorMemIndex<float>>(
             index_type_,
             metric_type_,
-            knowhere::Version::GetCurrentVersion().VersionNumber(), false);
+            knowhere::Version::GetCurrentVersion().VersionNumber(),
+            false);
     }
 
     auto num_chunk = vec_data_->num_chunks();

internal/querynodev2/segments/collection.go

@@ -287,6 +287,7 @@ func NewCollection(collectionID int64, schema *schemapb.CollectionSchema, indexM
 	isGpuIndex := false
 	req := &segcore.CreateCCollectionRequest{
 		Schema:        loadSchema,
+		LoadFieldList: loadFieldIDs.Collect(),
 	}
 	if indexMeta != nil && len(indexMeta.GetIndexMetas()) > 0 && indexMeta.GetMaxIndexRowCount() > 0 {
 		req.IndexMeta = indexMeta

internal/util/segcore/collection.go

@@ -24,6 +24,7 @@ type CreateCCollectionRequest struct {
 	CollectionID  int64
 	Schema        *schemapb.CollectionSchema
 	IndexMeta     *segcorepb.CollectionIndexMeta
+	LoadFieldList []int64
 }
 
 // CreateCCollection creates a CCollection from a CreateCCollectionRequest.
@@ -51,6 +52,14 @@ func CreateCCollection(req *CreateCCollectionRequest) (*CCollection, error) {
 			return nil, err
 		}
 	}
+	if req.LoadFieldList != nil {
+		status = C.UpdateLoadFields(ptr, (*C.int64_t)(unsafe.Pointer(&req.LoadFieldList[0])),
+			C.int64_t(len(req.LoadFieldList)))
+		if err := ConsumeCStatusIntoError(&status); err != nil {
+			C.DeleteCollection(ptr)
+			return nil, err
+		}
+	}
 	return &CCollection{
 		collectionID: req.CollectionID,
 		ptr:          ptr,