feat: LRU cache implementation (milvus-io#32567)

issue: milvus-io#32783
This pr is the implementation of lru cache on branch lru-dev.

Signed-off-by: sunby <sunbingyi1992@gmail.com>
Co-authored-by: chyezh <chyezh@outlook.com>
Co-authored-by: MrPresent-Han <chun.han@zilliz.com>
Co-authored-by: Ted Xu <ted.xu@zilliz.com>
Co-authored-by: jaime <yun.zhang@zilliz.com>
Co-authored-by: wayblink <anyang.wang@zilliz.com>

Author: 6 people

configs/milvus.yaml

@@ -367,6 +367,7 @@ queryNode:
       maxQueueLength: 16 # Maximum length of task queue in flowgraph
       maxParallelism: 1024 # Maximum number of tasks executed in parallel in the flowgraph
   enableSegmentPrune: false # use partition prune function on shard delegator
+  useStreamComputing: false
   ip:  # if not specified, use the first unicastable address
   port: 21123
   grpc:

internal/core/src/index/VectorDiskIndex.cpp

@@ -502,6 +502,11 @@ VectorDiskAnnIndex<T>::update_load_json(const Config& config) {
         }
     }
 
+    if (config.contains(kMmapFilepath)) {
+        load_config.erase(kMmapFilepath);
+        load_config[kEnableMmap] = true;
+    }
+
     return load_config;
 }
 

internal/core/src/index/VectorMemIndex.cpp

@@ -18,7 +18,6 @@
 
 #include <unistd.h>
 #include <cmath>
-#include <cstdint>
 #include <cstring>
 #include <filesystem>
 #include <memory>
@@ -33,7 +32,6 @@
 
 #include "index/Index.h"
 #include "index/IndexInfo.h"
-#include "index/Meta.h"
 #include "index/Utils.h"
 #include "common/EasyAssert.h"
 #include "config/ConfigKnowhere.h"
@@ -44,16 +42,15 @@
 #include "common/FieldData.h"
 #include "common/File.h"
 #include "common/Slice.h"
-#include "common/Tracer.h"
 #include "common/RangeSearchHelper.h"
 #include "common/Utils.h"
 #include "log/Log.h"
-#include "mmap/Types.h"
 #include "storage/DataCodec.h"
 #include "storage/MemFileManagerImpl.h"
 #include "storage/ThreadPools.h"
 #include "storage/space.h"
 #include "storage/Util.h"
+#include "storage/prometheus_client.h"
 
 namespace milvus::index {
 
@@ -733,7 +730,8 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
     }
 
     LOG_INFO("load with slice meta: {}", !slice_meta_filepath.empty());
-
+    std::chrono::duration<double> load_duration_sum;
+    std::chrono::duration<double> write_disk_duration_sum;
     if (!slice_meta_filepath
              .empty()) {  // load with the slice meta info, then we can load batch by batch
         std::string index_file_prefix = slice_meta_filepath.substr(
@@ -751,15 +749,20 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
             std::string prefix = item[NAME];
             int slice_num = item[SLICE_NUM];
             auto total_len = static_cast<size_t>(item[TOTAL_LEN]);
-
             auto HandleBatch = [&](int index) {
+                auto start_load2_mem = std::chrono::system_clock::now();
                 auto batch_data = file_manager_->LoadIndexToMemory(batch);
+                load_duration_sum +=
+                    (std::chrono::system_clock::now() - start_load2_mem);
                 for (int j = index - batch.size() + 1; j <= index; j++) {
                     std::string file_name = GenSlicedFileName(prefix, j);
                     AssertInfo(batch_data.find(file_name) != batch_data.end(),
                                "lost index slice data");
                     auto data = batch_data[file_name];
+                    auto start_write_file = std::chrono::system_clock::now();
                     auto written = file.Write(data->Data(), data->Size());
+                    write_disk_duration_sum +=
+                        (std::chrono::system_clock::now() - start_write_file);
                     AssertInfo(
                         written == data->Size(),
                         fmt::format("failed to write index data to disk {}: {}",
@@ -784,24 +787,46 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
             }
         }
     } else {
+        //1. load files into memory
+        auto start_load_files2_mem = std::chrono::system_clock::now();
         auto result = file_manager_->LoadIndexToMemory(std::vector<std::string>(
             pending_index_files.begin(), pending_index_files.end()));
+        load_duration_sum +=
+            (std::chrono::system_clock::now() - start_load_files2_mem);
+        //2. write data into files
+        auto start_write_file = std::chrono::system_clock::now();
         for (auto& [_, index_data] : result) {
             file.Write(index_data->Data(), index_data->Size());
         }
-    }
+        write_disk_duration_sum +=
+            (std::chrono::system_clock::now() - start_write_file);
+    }
+    milvus::storage::internal_storage_download_duration.Observe(
+        std::chrono::duration_cast<std::chrono::milliseconds>(load_duration_sum)
+            .count());
+    milvus::storage::internal_storage_write_disk_duration.Observe(
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            write_disk_duration_sum)
+            .count());
     file.Close();
 
     LOG_INFO("load index into Knowhere...");
     auto conf = config;
     conf.erase(kMmapFilepath);
     conf[kEnableMmap] = true;
+    auto start_deserialize = std::chrono::system_clock::now();
     auto stat = index_.DeserializeFromFile(filepath.value(), conf);
+    auto deserialize_duration =
+        std::chrono::system_clock::now() - start_deserialize;
     if (stat != knowhere::Status::success) {
         PanicInfo(ErrorCode::UnexpectedError,
                   "failed to Deserialize index: {}",
                   KnowhereStatusString(stat));
     }
+    milvus::storage::internal_storage_deserialize_duration.Observe(
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            deserialize_duration)
+            .count());
 
     auto dim = index_.Dim();
     this->SetDim(index_.Dim());
@@ -811,7 +836,18 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
                "failed to unlink mmap index file {}: {}",
                filepath.value(),
                strerror(errno));
-    LOG_INFO("load vector index done");
+    LOG_INFO(
+        "load vector index done, mmap_file_path:{}, download_duration:{}, "
+        "write_files_duration:{}, deserialize_duration:{}",
+        filepath.value(),
+        std::chrono::duration_cast<std::chrono::milliseconds>(load_duration_sum)
+            .count(),
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            write_disk_duration_sum)
+            .count(),
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            deserialize_duration)
+            .count());
 }
 
 template <typename T>

internal/core/src/segcore/CMakeLists.txt

@@ -25,6 +25,7 @@ set(SEGCORE_FILES
         SegmentSealedImpl.cpp
         FieldIndexing.cpp
         Reduce.cpp
+        StreamReduce.cpp
         metrics_c.cpp
         plan_c.cpp
         reduce_c.cpp
@@ -36,7 +37,8 @@ set(SEGCORE_FILES
         segcore_init_c.cpp
         TimestampIndex.cpp
         Utils.cpp
-        ConcurrentVector.cpp)
+        ConcurrentVector.cpp
+        ReduceUtils.cpp)
 add_library(milvus_segcore SHARED ${SEGCORE_FILES})
 
 target_link_libraries(milvus_segcore milvus_query milvus_bitset milvus_exec ${OpenMP_CXX_FLAGS} milvus-storage)

internal/core/src/segcore/ConcurrentVector.h

@@ -41,35 +41,36 @@ class ThreadSafeVector {
     template <typename... Args>
     void
     emplace_to_at_least(int64_t size, Args... args) {
+        std::lock_guard lck(mutex_);
         if (size <= size_) {
             return;
         }
-        std::lock_guard lck(mutex_);
         while (vec_.size() < size) {
             vec_.emplace_back(std::forward<Args...>(args...));
             ++size_;
         }
     }
     const Type&
     operator[](int64_t index) const {
+        std::shared_lock lck(mutex_);
         AssertInfo(index < size_,
                    fmt::format(
                        "index out of range, index={}, size_={}", index, size_));
-        std::shared_lock lck(mutex_);
         return vec_[index];
     }
 
     Type&
     operator[](int64_t index) {
+        std::shared_lock lck(mutex_);
         AssertInfo(index < size_,
                    fmt::format(
                        "index out of range, index={}, size_={}", index, size_));
-        std::shared_lock lck(mutex_);
         return vec_[index];
     }
 
     int64_t
     size() const {
+        std::lock_guard lck(mutex_);
         return size_;
     }
 
@@ -81,7 +82,7 @@ class ThreadSafeVector {
     }
 
  private:
-    std::atomic<int64_t> size_ = 0;
+    int64_t size_ = 0;
     std::deque<Type> vec_;
     mutable std::shared_mutex mutex_;
 };

internal/core/src/segcore/InsertRecord.h

@@ -598,6 +598,11 @@ struct InsertRecord {
         fields_data_.clear();
     }
 
+    bool
+    empty() const {
+        return pk2offset_->empty();
+    }
+
  public:
     ConcurrentVector<Timestamp> timestamps_;
     ConcurrentVector<idx_t> row_ids_;

internal/core/src/segcore/Reduce.cpp

@@ -19,6 +19,7 @@
 #include "Utils.h"
 #include "common/EasyAssert.h"
 #include "pkVisitor.h"
+#include "ReduceUtils.h"
 
 namespace milvus::segcore {
 
@@ -130,7 +131,6 @@ ReduceHelper::FilterInvalidSearchResult(SearchResult* search_result) {
 
 void
 ReduceHelper::FillPrimaryKey() {
-    std::vector<SearchResult*> valid_search_results;
     // get primary keys for duplicates removal
     uint32_t valid_index = 0;
     for (auto& search_result : search_results_) {
@@ -368,7 +368,7 @@ ReduceHelper::GetSearchResultDataSlice(int slice_index) {
     search_result_data->set_all_search_count(all_search_count);
 
     // `result_pairs` contains the SearchResult and result_offset info, used for filling output fields
-    std::vector<std::pair<SearchResult*, int64_t>> result_pairs(result_count);
+    std::vector<MergeBase> result_pairs(result_count);
 
     // reserve space for pks
     auto primary_field_id =
@@ -461,14 +461,14 @@ ReduceHelper::GetSearchResultDataSlice(int slice_index) {
                     group_by_values[loc] =
                         search_result->group_by_values_.value()[ki];
                 // set result offset to fill output fields data
-                result_pairs[loc] = std::make_pair(search_result, ki);
+                result_pairs[loc] = {&search_result->output_fields_data_, ki};
             }
         }
 
         // update result topKs
         search_result_data->mutable_topks()->Set(qi - nq_begin, topk_count);
     }
-    AssembleGroupByValues(search_result_data, group_by_values);
+    AssembleGroupByValues(search_result_data, group_by_values, plan_);
 
     AssertInfo(search_result_data->scores_size() == result_count,
                "wrong scores size, size = " +
@@ -498,89 +498,4 @@ ReduceHelper::GetSearchResultDataSlice(int slice_index) {
     return buffer;
 }
 
-void
-ReduceHelper::AssembleGroupByValues(
-    std::unique_ptr<milvus::proto::schema::SearchResultData>& search_result,
-    const std::vector<GroupByValueType>& group_by_vals) {
-    auto group_by_field_id = plan_->plan_node_->search_info_.group_by_field_id_;
-    if (group_by_field_id.has_value() && group_by_vals.size() > 0) {
-        auto group_by_values_field =
-            std::make_unique<milvus::proto::schema::ScalarField>();
-        auto group_by_field =
-            plan_->schema_.operator[](group_by_field_id.value());
-        DataType group_by_data_type = group_by_field.get_data_type();
-
-        int group_by_val_size = group_by_vals.size();
-        switch (group_by_data_type) {
-            case DataType::INT8: {
-                auto field_data = group_by_values_field->mutable_int_data();
-                field_data->mutable_data()->Resize(group_by_val_size, 0);
-                for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
-                    int8_t val = std::get<int8_t>(group_by_vals[idx]);
-                    field_data->mutable_data()->Set(idx, val);
-                }
-                break;
-            }
-            case DataType::INT16: {
-                auto field_data = group_by_values_field->mutable_int_data();
-                field_data->mutable_data()->Resize(group_by_val_size, 0);
-                for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
-                    int16_t val = std::get<int16_t>(group_by_vals[idx]);
-                    field_data->mutable_data()->Set(idx, val);
-                }
-                break;
-            }
-            case DataType::INT32: {
-                auto field_data = group_by_values_field->mutable_int_data();
-                field_data->mutable_data()->Resize(group_by_val_size, 0);
-                for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
-                    int32_t val = std::get<int32_t>(group_by_vals[idx]);
-                    field_data->mutable_data()->Set(idx, val);
-                }
-                break;
-            }
-            case DataType::INT64: {
-                auto field_data = group_by_values_field->mutable_long_data();
-                field_data->mutable_data()->Resize(group_by_val_size, 0);
-                for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
-                    int64_t val = std::get<int64_t>(group_by_vals[idx]);
-                    field_data->mutable_data()->Set(idx, val);
-                }
-                break;
-            }
-            case DataType::BOOL: {
-                auto field_data = group_by_values_field->mutable_bool_data();
-                field_data->mutable_data()->Resize(group_by_val_size, 0);
-                for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
-                    bool val = std::get<bool>(group_by_vals[idx]);
-                    field_data->mutable_data()->Set(idx, val);
-                }
-                break;
-            }
-            case DataType::VARCHAR: {
-                auto field_data = group_by_values_field->mutable_string_data();
-                for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
-                    std::string val =
-                        std::move(std::get<std::string>(group_by_vals[idx]));
-                    *(field_data->mutable_data()->Add()) = val;
-                }
-                break;
-            }
-            default: {
-                PanicInfo(
-                    DataTypeInvalid,
-                    fmt::format("unsupported datatype for group_by operations ",
-                                group_by_data_type));
-            }
-        }
-
-        search_result->mutable_group_by_field_value()->set_type(
-            milvus::proto::schema::DataType(group_by_data_type));
-        search_result->mutable_group_by_field_value()
-            ->mutable_scalars()
-            ->MergeFrom(*group_by_values_field.get());
-        return;
-    }
-}
-
 }  // namespace milvus::segcore