improve sparse vector index mmap: to mmap almost everything (#928)

Signed-off-by: Buqian Zheng <zhengbuqian@gmail.com>
zilliztech · Nov 11, 2024 · 504011c · 504011c
1 parent 8aa9f94
commit 504011c
Show file tree

Hide file tree

Showing 5 changed files with 400 additions and 147 deletions.
diff --git a/include/knowhere/sparse_utils.h b/include/knowhere/sparse_utils.h
@@ -274,4 +274,82 @@ class MaxMinHeap {
     std::vector<SparseIdVal<T>> pool_;
 };  // class MaxMinHeap
 
+// A std::vector like container but uses fixed size free memory(typically from
+// mmap) as backing store and can only be appended at the end.
+//
+// Must be initialized with a valid pointer to memory when used. The memory must be
+// valid during the lifetime of this object. After initialization, GrowableVectorView will
+// have space for mmap_element_count_ elements, none of which are initialized.
+//
+// Currently only used in sparse InvertedIndex. Move to other places if needed.
+template <typename T>
+class GrowableVectorView {
+ public:
+    using value_type = T;
+    using size_type = size_t;
+
+    GrowableVectorView() = default;
+
+    void
+    initialize(void* data, size_type byte_size) {
+        if (byte_size % sizeof(T) != 0) {
+            throw std::invalid_argument("GrowableVectorView byte_size must be a multiple of element size");
+        }
+        mmap_data_ = data;
+        mmap_byte_size_ = byte_size;
+        mmap_element_count_ = 0;
+    }
+
+    size_type
+    capacity() const {
+        return mmap_byte_size_ / sizeof(T);
+    }
+
+    size_type
+    size() const {
+        return mmap_element_count_;
+    }
+
+    template <typename... Args>
+    T&
+    emplace_back(Args&&... args) {
+        if (size() == capacity()) {
+            throw std::out_of_range("emplace_back on a full GrowableVectorView");
+        }
+        auto* elem = reinterpret_cast<T*>(mmap_data_) + mmap_element_count_++;
+        return *new (elem) T(std::forward<Args>(args)...);
+    }
+
+    T&
+    operator[](size_type i) {
+        return reinterpret_cast<T*>(mmap_data_)[i];
+    }
+
+    const T&
+    operator[](size_type i) const {
+        return reinterpret_cast<const T*>(mmap_data_)[i];
+    }
+
+    T&
+    at(size_type i) {
+        if (i >= mmap_element_count_) {
+            throw std::out_of_range("GrowableVectorView index out of range");
+        }
+        return reinterpret_cast<T*>(mmap_data_)[i];
+    }
+
+    const T&
+    at(size_type i) const {
+        if (i >= mmap_element_count_) {
+            throw std::out_of_range("GrowableVectorView index out of range");
+        }
+        return reinterpret_cast<const T*>(mmap_data_)[i];
+    }
+
+ private:
+    void* mmap_data_ = nullptr;
+    size_type mmap_byte_size_ = 0;
+    size_type mmap_element_count_ = 0;
+};
+
 }  // namespace knowhere::sparse
diff --git a/src/index/sparse/sparse_index_node.cc b/src/index/sparse/sparse_index_node.cc
@@ -53,7 +53,7 @@ class SparseInvertedIndexNode : public IndexNode {
             return Status::invalid_metric_type;
         }
         auto drop_ratio_build = cfg.drop_ratio_build.value_or(0.0f);
-        auto index_or = CreateIndex(cfg);
+        auto index_or = CreateIndex</*mmapped=*/false>(cfg);
         if (!index_or.has_value()) {
             return index_or.error();
         }
@@ -206,12 +206,12 @@ class SparseInvertedIndexNode : public IndexNode {
             return Status::invalid_binary_set;
         }
         MemoryIOReader reader(binary->data.get(), binary->size);
-        auto index_or = CreateIndex(cfg);
+        auto index_or = CreateIndex</*mmapped=*/false>(cfg);
         if (!index_or.has_value()) {
             return index_or.error();
         }
         index_ = index_or.value();
-        return index_->Load(reader, false);
+        return index_->Load(reader);
     }
 
     Status
@@ -237,13 +237,13 @@ class SparseInvertedIndexNode : public IndexNode {
         if (madvise(map_, map_size_, MADV_RANDOM) != 0) {
             LOG_KNOWHERE_WARNING_ << "Failed to madvise file: " << strerror(errno);
         }
-        auto index_or = CreateIndex(cfg);
+        auto index_or = CreateIndex</*mmapped=*/true>(cfg);
         if (!index_or.has_value()) {
             return index_or.error();
         }
         index_ = index_or.value();
         MemoryIOReader map_reader((uint8_t*)map_, map_size_);
-        return index_->Load(map_reader, true);
+        return index_->Load(map_reader, map_flags);
     }
 
     static std::unique_ptr<BaseConfig>
@@ -278,10 +278,11 @@ class SparseInvertedIndexNode : public IndexNode {
     }
 
  private:
+    template <bool mmapped>
     expected<sparse::BaseInvertedIndex<T>*>
     CreateIndex(const SparseInvertedIndexConfig& cfg) const {
         if (IsMetricType(cfg.metric_type.value(), metric::BM25)) {
-            auto idx = new sparse::InvertedIndex<T, use_wand, true>();
+            auto idx = new sparse::InvertedIndex<T, use_wand, true, mmapped>();
             if (!cfg.bm25_k1.has_value() || !cfg.bm25_b.has_value() || !cfg.bm25_avgdl.has_value()) {
                 return expected<sparse::BaseInvertedIndex<T>*>::Err(
                     Status::invalid_args, "BM25 parameters k1, b, and avgdl must be set when building/loading");
@@ -293,7 +294,7 @@ class SparseInvertedIndexNode : public IndexNode {
             idx->SetBM25Params(k1, b, avgdl, max_score_ratio);
             return idx;
         } else {
-            return new sparse::InvertedIndex<T, use_wand, false>();
+            return new sparse::InvertedIndex<T, use_wand, false, mmapped>();
         }
     }