Add optimize in mmap hnsw

infiniflow · Jan 6, 2025 · 9bd9bdf · 9bd9bdf
1 parent c38e6eb
commit 9bd9bdf
Show file tree

Hide file tree

Showing 9 changed files with 136 additions and 5 deletions.
diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp
@@ -140,7 +140,11 @@ void HnswFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
             } else {
                 using IndexT = std::decay_t<decltype(*index)>;
                 if constexpr (IndexT::kOwnMem) {
-                    index = IndexT::Load(*file_handle_).release();
+                    if (from_spill) {
+                        index = IndexT::Load(*file_handle_).release();
+                    } else {
+                        index = IndexT::LoadFromPtr(*file_handle_, file_size).release();
+                    }
                 } else {
                     UnrecoverableError("Invalid index type.");
                 }

diff --git a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm
@@ -215,6 +215,24 @@ public:
         return ret;
     }
 
+    static This LoadFromPtr(const char *&ptr) {
+        SizeT cur_vec_num = ReadBufAdv<SizeT>(ptr);
+        VecStoreMeta vec_store_meta = VecStoreMeta::LoadFromPtr(ptr);
+        GraphStoreMeta graph_store_meta = GraphStoreMeta::LoadFromPtr(ptr);
+
+        SizeT chunk_size = 1;
+        while (chunk_size < cur_vec_num) {
+            chunk_size <<= 1;
+        }
+        This ret = This(chunk_size, 1 /*max_chunk_n*/, std::move(vec_store_meta), std::move(graph_store_meta));
+        ret.cur_vec_num_ = cur_vec_num;
+
+        SizeT mem_usage = 0;
+        ret.inners_[0] = Inner::LoadFromPtr(ptr, cur_vec_num, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage);
+        ret.mem_usage_.store(mem_usage);
+        return ret;
+    }
+
     void SetGraph(GraphStoreMeta &&graph_meta, Vector<GraphStoreInner<OwnMem>> &&graph_inners) {
         this->graph_store_meta_ = std::move(graph_meta);
         for (SizeT i = 0; i < graph_inners.size(); ++i) {
@@ -615,6 +633,20 @@ public:
         return ret;
     }
 
+    static This LoadFromPtr(const char *&ptr,
+                            SizeT cur_vec_num,
+                            SizeT chunk_size,
+                            VecStoreMeta &vec_store_meta,
+                            GraphStoreMeta &graph_store_meta,
+                            SizeT &mem_usage) {
+        auto vec_store_inner = VecStoreInner::LoadFromPtr(ptr, cur_vec_num, chunk_size, vec_store_meta, mem_usage);
+        auto graph_store_inner = GraphStoreInner::LoadFromPtr(ptr, cur_vec_num, chunk_size, graph_store_meta, mem_usage);
+        This ret(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner));
+        std::memcpy(ret.labels_.get(), ptr, sizeof(LabelType) * cur_vec_num);
+        ptr += sizeof(LabelType) * cur_vec_num;
+        return ret;
+    }
+
     // vec store
     template <DataIteratorConcept<QueryVecType, LabelType> Iterator>
     Pair<SizeT, bool> AddVec(Iterator &&query_iter, VertexType start_idx, SizeT remain_num, const VecStoreMeta &meta, SizeT &mem_usage) {

diff --git a/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm
@@ -383,6 +383,32 @@ public:
         return graph_store;
     }
 
+    static GraphStoreInner LoadFromPtr(const char *&ptr, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) {
+        SizeT layer_sum = ReadBufAdv<SizeT>(ptr);
+        GraphStoreInner graph_store(max_vertex, meta, cur_vertex_n);
+        const char *graph = ptr;
+        ptr += cur_vertex_n * meta.level0_size();
+        std::memcpy(graph_store.graph_.get(), graph, cur_vertex_n * meta.level0_size());
+
+        auto loaded_layers = MakeUnique<char[]>(layer_sum * meta.levelx_size());
+        char *loaded_layers_p = loaded_layers.get();
+        for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) {
+            VertexL0 *v = graph_store.GetLevel0(vertex_i, meta);
+            if (v->layer_n_) {
+                std::memcpy(loaded_layers_p, ptr, meta.levelx_size() * v->layer_n_);
+                v->layers_p_ = loaded_layers_p;
+                loaded_layers_p += meta.levelx_size() * v->layer_n_;
+                ptr += meta.levelx_size() * v->layer_n_;
+            } else {
+                v->layers_p_ = nullptr;
+            }
+        }
+        graph_store.loaded_layers_ = std::move(loaded_layers);
+
+        mem_usage += max_vertex * meta.level0_size() + layer_sum * meta.levelx_size();
+        return graph_store;
+    }
+
     void AddVertex(VertexType vertex_i, i32 layer_n, const GraphStoreMeta &meta, SizeT &mem_usage) {
         VertexL0 *v = GetLevel0(vertex_i, meta);
         v->neighbor_n_ = 0;

diff --git a/src/storage/knn_index/knn_hnsw/data_store/lvq_vec_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/lvq_vec_store.cppm
@@ -230,6 +230,18 @@ public:
         return meta;
     }
 
+    static This LoadFromPtr(const char *&ptr) {
+        SizeT dim = ReadBufAdv<SizeT>(ptr);
+        This meta(dim);
+        std::memcpy(meta.mean_.get(), ptr, sizeof(MeanType) * dim);
+        ptr += sizeof(MeanType) * dim;
+        if constexpr (!std::is_same_v<GlobalCacheType, Tuple<>>) {
+            std::memcpy(&meta.global_cache_, ptr, sizeof(GlobalCacheType));
+            ptr += sizeof(GlobalCacheType);
+        }
+        return meta;
+    }
+
     template <typename LabelType, DataIteratorConcept<const DataType *, LabelType> Iterator>
     void Optimize(Iterator &&query_iter, const Vector<Pair<Inner *, SizeT>> &inners, SizeT &mem_usage) {
         auto new_mean = MakeUnique<MeanType[]>(this->dim_);
@@ -378,6 +390,14 @@ public:
         return ret;
     }
 
+    static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) {
+        This ret(max_vec_num, meta);
+        std::memcpy(ret.ptr_.get(), ptr, cur_vec_num * meta.compress_data_size());
+        ptr += cur_vec_num * meta.compress_data_size();
+        mem_usage += max_vec_num * meta.compress_data_size();
+        return ret;
+    }
+
     void SetVec(SizeT idx, const DataType *vec, const Meta &meta, SizeT &mem_usage) { meta.CompressTo(vec, GetVecMut(idx, meta)); }
 
 private:

diff --git a/src/storage/knn_index/knn_hnsw/data_store/plain_vec_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/plain_vec_store.cppm
@@ -152,6 +152,14 @@ public:
         return ret;
     }
 
+    static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) {
+        This ret(max_vec_num, meta);
+        std::memcpy(ret.ptr_.get(), ptr, sizeof(DataType) * cur_vec_num * meta.dim());
+        ptr += sizeof(DataType) * cur_vec_num * meta.dim();
+        mem_usage += sizeof(DataType) * max_vec_num * meta.dim();
+        return ret;
+    }
+
     void SetVec(SizeT idx, const DataType *vec, const Meta &meta, SizeT &mem_usage) { Copy(vec, vec + meta.dim(), GetVecMut(idx, meta)); }
 
 private:

diff --git a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm
@@ -484,6 +484,20 @@ public:
         return MakeUnique<This>(M, ef_construction, std::move(data_store), std::move(distance));
     }
 
+    static UniquePtr<This> LoadFromPtr(LocalFileHandle &file_handle, SizeT size) {
+        auto buffer = MakeUnique<char []>(size);
+        file_handle.Read(buffer.get(), size);
+        const char *ptr = buffer.get();
+        SizeT M = ReadBufAdv<SizeT>(ptr);
+        SizeT ef_construction = ReadBufAdv<SizeT>(ptr);
+        auto data_store = DataStore::LoadFromPtr(ptr);
+        Distance distance(data_store.dim());
+        if (SizeT diff = ptr - buffer.get(); diff != size) {
+            UnrecoverableError("LoadFromPtr failed");
+        }
+        return MakeUnique<This>(M, ef_construction, std::move(data_store), std::move(distance));
+    }
+
     UniquePtr<KnnHnsw<CompressVecStoreType, LabelType>> CompressToLVQ() && {
         if constexpr (std::is_same_v<VecStoreType, CompressVecStoreType>) {
             return MakeUnique<This>(std::move(*this));

diff --git a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp
@@ -178,6 +178,33 @@ class HnswAlgTest : public BaseTest {
 
             test_func(hnsw_index);
 
+#ifdef USE_MMAP
+            VirtualStore::MunmapFile(filepath);
+#endif
+        }
+        {
+            SizeT file_size = VirtualStore::GetFileSize(filepath);
+#define USE_MMAP
+#ifdef USE_MMAP
+            unsigned char *data_ptr = nullptr;
+            int ret = VirtualStore::MmapFile(filepath, data_ptr, file_size);
+            if (ret < 0) {
+                UnrecoverableError("mmap failed");
+            }
+            const char *ptr = reinterpret_cast<const char *>(data_ptr);
+#else
+            auto [file_handle, status] = VirtualStore::Open(filepath, FileAccessMode::kRead);
+            if (!status.ok()) {
+                UnrecoverableError(status.message());
+            }
+            auto buffer = MakeUnique<char[]>(file_size);
+            file_handle->Read(buffer.get(), file_size);
+            const char *ptr = buffer.get();
+#endif
+            auto hnsw_index = Hnsw::LoadFromPtr(ptr, file_size);
+
+            test_func(hnsw_index);
+
 #ifdef USE_MMAP
             VirtualStore::MunmapFile(filepath);
 #endif

diff --git a/test/sql/dql/knn/embedding/test_knn_hnsw_l2_lvq.slt b/test/sql/dql/knn/embedding/test_knn_hnsw_l2_lvq.slt
@@ -34,8 +34,8 @@ SELECT c1 FROM test_knn_hnsw_l2 SEARCH MATCH VECTOR (c2, [0.3, 0.3, 0.2, 0.2], '
 8
 6
 
-# statement ok
-# OPTIMIZE idx1 ON test_knn_hnsw_l2 WITH (lvq_avg);
+statement ok
+OPTIMIZE idx1 ON test_knn_hnsw_l2 WITH (lvq_avg);
 
 query I
 SELECT c1 FROM test_knn_hnsw_l2 SEARCH MATCH VECTOR (c2, [0.3, 0.3, 0.2, 0.2], 'float', 'l2', 3) WITH (ef = 6, rerank);

diff --git a/test/sql/dql/knn/embedding/test_knn_hnsw_l2_lvq2.slt b/test/sql/dql/knn/embedding/test_knn_hnsw_l2_lvq2.slt
@@ -23,8 +23,8 @@ SELECT c1 FROM test_knn_hnsw_l2 SEARCH MATCH VECTOR (c2, [0.3, 0.3, 0.2, 0.2], '
 6
 4
 
-# statement ok
-# OPTIMIZE idx1 ON test_knn_hnsw_l2 WITH (compress_to_lvq);
+statement ok
+OPTIMIZE idx1 ON test_knn_hnsw_l2 WITH (compress_to_lvq);
 
 query I
 SELECT c1 FROM test_knn_hnsw_l2 SEARCH MATCH VECTOR (c2, [0.3, 0.3, 0.2, 0.2], 'float', 'l2', 3) WITH (ef = 4, rerank);