From 68df975876d75240ad8545c90e256eddf6be05b2 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Mon, 23 Dec 2024 19:28:26 +0800 Subject: [PATCH 01/12] Add own mem template in hnsw store. --- .../buffer/file_worker/hnsw_file_worker.cpp | 37 ++ .../buffer/file_worker/hnsw_file_worker.cppm | 4 + .../knn_hnsw/data_store/data_store.cppm | 528 +++++++++++------- .../knn_hnsw/data_store/data_store_util.cppm | 52 ++ .../knn_hnsw/data_store/graph_store.cppm | 209 ++++--- .../knn_hnsw/data_store/lvq_vec_store.cppm | 290 ++++++---- .../knn_hnsw/data_store/plain_vec_store.cppm | 87 ++- .../knn_hnsw/data_store/sparse_vec_store.cppm | 1 + .../knn_hnsw/data_store/vec_store_type.cppm | 63 ++- .../knn_index/knn_hnsw/dist_func_cos.cppm | 7 +- .../knn_index/knn_hnsw/dist_func_ip.cppm | 7 +- .../knn_index/knn_hnsw/dist_func_l2.cppm | 7 +- 12 files changed, 843 insertions(+), 449 deletions(-) create mode 100644 src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index ece759fd4e..4970359cc9 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -140,4 +140,41 @@ void HnswFileWorker::ReadFromFileImpl(SizeT file_size) { *hnsw_index); } +bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) { + if (mmap_data_ != nullptr) { + UnrecoverableError("Mmap data is already allocated."); + } + mmap_data_ = reinterpret_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get()))); + auto *hnsw_index = reinterpret_cast(mmap_data_); + std::visit( + [&](auto &&index) { + using T = std::decay_t; + if constexpr (std::is_same_v) { + UnrecoverableError("Invalid index type."); + } else { + // using IndexT = std::decay_t; + // index = IndexT::LoadFromPtr(static_cast(ptr), size).release(); + } + }, + *hnsw_index); + return true; +} + +void HnswFileWorker::FreeFromMmapImpl() { + if (mmap_data_ == nullptr) { + UnrecoverableError("Mmap data is not allocated."); + } + auto *hnsw_index = reinterpret_cast(mmap_data_); + std::visit( + [&](auto &&index) { + using T = std::decay_t; + if constexpr (!std::is_same_v) { + delete index; + } + }, + *hnsw_index); + delete hnsw_index; + mmap_data_ = nullptr; +} + } // namespace infinity \ No newline at end of file diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cppm b/src/storage/buffer/file_worker/hnsw_file_worker.cppm index fe5b8191fa..31952b88a0 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cppm +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cppm @@ -56,6 +56,10 @@ protected: void ReadFromFileImpl(SizeT file_size) override; + bool ReadFromMmapImpl(const void *ptr, SizeT size) override; + + void FreeFromMmapImpl() override; + private: SizeT index_size_{}; }; 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 index f6df8fefa6..1bfcd18a00 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm @@ -26,33 +26,34 @@ import local_file_handle; import vec_store_type; import graph_store; import infinity_exception; +import serialize; +import data_store_util; namespace infinity { -template +template class DataStoreInner; export template class DataStoreChunkIter; -template +template class DataStoreIter; #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wunused-variable" -export template -class DataStore { +export template +class DataStoreBase { public: - using This = DataStore; + using This = DataStoreBase; using DataType = typename VecStoreT::DataType; using QueryVecType = typename VecStoreT::QueryVecType; - using Inner = DataStoreInner; - using VecStoreMeta = typename VecStoreT::Meta; - using VecStoreInner = typename VecStoreT::Inner; + using Inner = DataStoreInner; + using VecStoreMeta = typename VecStoreT::template Meta; friend class DataStoreChunkIter; - friend class DataStoreIter; + friend class DataStoreIter; public: template @@ -61,24 +62,13 @@ public: template struct has_compress_type> : std::true_type {}; -private: - DataStore(SizeT chunk_size, SizeT max_chunk_n, VecStoreMeta &&vec_store_meta, GraphStoreMeta &&graph_store_meta) - : chunk_size_(chunk_size), max_chunk_n_(max_chunk_n), vec_store_meta_(std::move(vec_store_meta)), - graph_store_meta_(std::move(graph_store_meta)) { - assert(chunk_size > 0); - assert((chunk_size & (chunk_size - 1)) == 0); - chunk_shift_ = __builtin_ctzll(chunk_size); - inners_ = MakeUnique(max_chunk_n); - } - -public: - DataStore() : chunk_size_(0), max_chunk_n_(0), chunk_shift_(0), cur_vec_num_(0) {} - DataStore(This &&other) + DataStoreBase() : chunk_size_(0), max_chunk_n_(0), chunk_shift_(0), cur_vec_num_(0) {} + DataStoreBase(This &&other) : chunk_size_(std::exchange(other.chunk_size_, 0)), max_chunk_n_(std::exchange(other.max_chunk_n_, 0)), chunk_shift_(std::exchange(other.chunk_shift_, 0)), cur_vec_num_(other.cur_vec_num_.exchange(0)), vec_store_meta_(std::move(other.vec_store_meta_)), graph_store_meta_(std::move(other.graph_store_meta_)), inners_(std::exchange(other.inners_, nullptr)), mem_usage_(other.mem_usage_.exchange(0)) {} - DataStore &operator=(This &&other) { + DataStoreBase &operator=(This &&other) { if (this != &other) { chunk_size_ = std::exchange(other.chunk_size_, 0); max_chunk_n_ = std::exchange(other.max_chunk_n_, 0); @@ -91,7 +81,7 @@ public: } return *this; } - ~DataStore() { + ~DataStoreBase() { if (!inners_) { return; } @@ -103,23 +93,7 @@ public: } } - static This Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT Mmax0, SizeT Mmax) { - bool normalize = false; - if constexpr (has_compress_type::value) { - normalize = std::is_same_v::Meta>; - } - VecStoreMeta vec_store_meta = VecStoreMeta::Make(dim, normalize); - GraphStoreMeta graph_store_meta = GraphStoreMeta::Make(Mmax0, Mmax); - This ret(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); - ret.cur_vec_num_ = 0; - - SizeT mem_usage = 0; - ret.inners_[0] = Inner::Make(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 &&graph_inners) { + void SetGraph(GraphStoreMeta &&graph_meta, Vector> &&graph_inners) { graph_store_meta_ = std::move(graph_meta); for (SizeT i = 0; i < graph_inners.size(); ++i) { inners_[i].SetGraphStoreInner(std::move(graph_inners[i])); @@ -159,97 +133,11 @@ public: } } - static This Load(LocalFileHandle &file_handle, SizeT max_chunk_n = 0) { - SizeT chunk_size; - file_handle.Read(&chunk_size, sizeof(chunk_size)); - SizeT max_chunk_n1; - file_handle.Read(&max_chunk_n1, sizeof(max_chunk_n1)); - if (max_chunk_n == 0) { - max_chunk_n = max_chunk_n1; - } - assert(max_chunk_n >= max_chunk_n1); - - SizeT cur_vec_num; - file_handle.Read(&cur_vec_num, sizeof(cur_vec_num)); - VecStoreMeta vec_store_meta = VecStoreMeta::Load(file_handle); - GraphStoreMeta graph_store_meta = GraphStoreMeta::Load(file_handle); - - This ret = This(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); - ret.cur_vec_num_ = cur_vec_num; - - SizeT mem_usage = 0; - auto [chunk_num, last_chunk_size] = ret.ChunkInfo(cur_vec_num); - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size : last_chunk_size; - ret.inners_[i] = Inner::Load(file_handle, cur_chunk_size, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); - } - ret.mem_usage_.store(mem_usage); - return ret; - } - - template Iterator> - Pair AddVec(Iterator &&query_iter) { - SizeT mem_usage = 0; - SizeT cur_vec_num = this->cur_vec_num(); - SizeT start_idx = cur_vec_num; - auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); - while (true) { - SizeT remain_size = chunk_size_ - last_chunk_size; - auto [insert_n, used_up] = inners_[chunk_num - 1].AddVec(std::move(query_iter), last_chunk_size, remain_size, vec_store_meta_, mem_usage); - cur_vec_num += insert_n; - last_chunk_size += insert_n; - if (cur_vec_num == max_chunk_n_ * chunk_size_) { - break; - } - if (last_chunk_size == chunk_size_) { - inners_[chunk_num++] = Inner::Make(chunk_size_, vec_store_meta_, graph_store_meta_, mem_usage); - last_chunk_size = 0; - } - if (used_up) { - break; - } - } - cur_vec_num_.store(cur_vec_num); - mem_usage_.fetch_add(mem_usage); - return {start_idx, cur_vec_num}; - } - - template Iterator> - Pair OptAddVec(Iterator &&query_iter) { - if constexpr (VecStoreT::HasOptimize) { - SizeT mem_usage = 0; - SizeT cur_vec_num = this->cur_vec_num(); - auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); - if (chunk_num > 0) { - Vector> vec_inners; - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; - vec_inners.emplace_back(inners_[i].vec_store_inner(), chunk_size); - } - Iterator query_iter_copy = query_iter; - vec_store_meta_.template Optimize(std::move(query_iter_copy), vec_inners, mem_usage); - } - mem_usage_.fetch_add(mem_usage); - } - return AddVec(std::move(query_iter)); - } - - void Optimize() { - if constexpr (!VecStoreT::HasOptimize) { - return; - } - DenseVectorIter empty_iter(nullptr, dim(), 0); - AddVec(std::move(empty_iter)); - } - typename VecStoreT::StoreType GetVec(SizeT vec_i) const { const auto &[inner, idx] = GetInner(vec_i); return inner.GetVec(idx, vec_store_meta_); } - template - DataStore CompressToLVQ() &&; - typename VecStoreT::QueryType MakeQuery(QueryVecType query) const { return vec_store_meta_.MakeQuery(query); } void PrefetchVec(SizeT vec_i) const { @@ -262,21 +150,10 @@ public: SizeT dim() const { return vec_store_meta_.dim(); } // Graph store - void AddVertex(VertexType vec_i, i32 layer_n) { - auto [inner, idx] = GetInner(vec_i); - SizeT mem_usage = 0; - inner.AddVertex(idx, layer_n, graph_store_meta_, mem_usage); - mem_usage_.fetch_add(mem_usage); - } - Pair GetNeighbors(VertexType vertex_i, i32 layer_i) const { const auto &[inner, idx] = GetInner(vertex_i); return inner.GetNeighbors(idx, layer_i, graph_store_meta_); } - Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i) { - auto [inner, idx] = GetInner(vertex_i); - return inner.GetNeighborsMut(idx, layer_i, graph_store_meta_); - } Pair GetEnterPoint() const { return graph_store_meta_.GetEnterPoint(); } @@ -305,7 +182,7 @@ public: SizeT mem_usage() const { return mem_usage_.load(); } -private: +protected: Pair GetInner(SizeT vec_i) { return {inners_[vec_i >> chunk_shift_], vec_i & (chunk_size_ - 1)}; } Pair GetInner(SizeT vec_i) const { return {inners_[vec_i >> chunk_shift_], vec_i & (chunk_size_ - 1)}; } @@ -318,7 +195,7 @@ private: return {chunk_num, last_chunk_size}; }; -private: +protected: SizeT chunk_size_; SizeT max_chunk_n_; SizeT chunk_shift_; @@ -371,33 +248,206 @@ public: } }; -#pragma clang diagnostic pop -//----------------------------------------------- Inner ----------------------------------------------- - -template -class DataStoreInner { +export template +class DataStore : public DataStoreBase { public: - using This = DataStoreInner; + using This = DataStore; + using Base = DataStoreBase; using DataType = typename VecStoreT::DataType; using QueryVecType = typename VecStoreT::QueryVecType; - using VecStoreInner = typename VecStoreT::Inner; - using VecStoreMeta = typename VecStoreT::Meta; + using Inner = DataStoreInner; + using VecStoreMeta = typename VecStoreT::template Meta; + using VecStoreInner = typename VecStoreT::template Inner; + +private: + DataStore(SizeT chunk_size, SizeT max_chunk_n, VecStoreMeta &&vec_store_meta, GraphStoreMeta &&graph_store_meta) { + this->chunk_size_ = chunk_size; + this->max_chunk_n_ = max_chunk_n; + this->vec_store_meta_ = std::move(vec_store_meta); + this->graph_store_meta_ = std::move(graph_store_meta); + assert(chunk_size > 0); + assert((chunk_size & (chunk_size - 1)) == 0); + this->chunk_shift_ = __builtin_ctzll(chunk_size); + this->inners_ = MakeUnique(max_chunk_n); + } + +public: + static This Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT Mmax0, SizeT Mmax) { + bool normalize = false; + if constexpr (Base::template has_compress_type::value) { + normalize = std::is_same_v::template Meta>; + } + VecStoreMeta vec_store_meta = VecStoreMeta::Make(dim, normalize); + GraphStoreMeta graph_store_meta = GraphStoreMeta::Make(Mmax0, Mmax); + This ret(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); + ret.cur_vec_num_ = 0; + + SizeT mem_usage = 0; + ret.inners_[0] = Inner::Make(chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); + ret.mem_usage_.store(mem_usage); + return ret; + } + + static This Load(LocalFileHandle &file_handle, SizeT max_chunk_n = 0) { + SizeT chunk_size; + file_handle.Read(&chunk_size, sizeof(chunk_size)); + SizeT max_chunk_n1; + file_handle.Read(&max_chunk_n1, sizeof(max_chunk_n1)); + if (max_chunk_n == 0) { + max_chunk_n = max_chunk_n1; + } + assert(max_chunk_n >= max_chunk_n1); - friend class DataStoreIter; + SizeT cur_vec_num; + file_handle.Read(&cur_vec_num, sizeof(cur_vec_num)); + VecStoreMeta vec_store_meta = VecStoreMeta::Load(file_handle); + GraphStoreMeta graph_store_meta = GraphStoreMeta::Load(file_handle); + + This ret = This(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); + ret.cur_vec_num_ = cur_vec_num; + + SizeT mem_usage = 0; + auto [chunk_num, last_chunk_size] = ret.ChunkInfo(cur_vec_num); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size : last_chunk_size; + ret.inners_[i] = Inner::Load(file_handle, cur_chunk_size, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); + } + ret.mem_usage_.store(mem_usage); + return ret; + } + + template Iterator> + Pair AddVec(Iterator &&query_iter) { + SizeT mem_usage = 0; + SizeT cur_vec_num = this->cur_vec_num(); + SizeT start_idx = cur_vec_num; + auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); + while (true) { + SizeT remain_size = this->chunk_size_ - last_chunk_size; + auto [insert_n, used_up] = + this->inners_[chunk_num - 1].AddVec(std::move(query_iter), last_chunk_size, remain_size, this->vec_store_meta_, mem_usage); + cur_vec_num += insert_n; + last_chunk_size += insert_n; + if (cur_vec_num == this->max_chunk_n_ * this->chunk_size_) { + break; + } + if (last_chunk_size == this->chunk_size_) { + this->inners_[chunk_num++] = Inner::Make(this->chunk_size_, this->vec_store_meta_, this->graph_store_meta_, mem_usage); + last_chunk_size = 0; + } + if (used_up) { + break; + } + } + this->cur_vec_num_.store(cur_vec_num); + this->mem_usage_.fetch_add(mem_usage); + return {start_idx, cur_vec_num}; + } + + template Iterator> + Pair OptAddVec(Iterator &&query_iter) { + if constexpr (VecStoreT::HasOptimize) { + SizeT mem_usage = 0; + SizeT cur_vec_num = this->cur_vec_num(); + auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); + if (chunk_num > 0) { + Vector> vec_inners; + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? this->chunk_size_ : last_chunk_size; + vec_inners.emplace_back(this->inners_[i].vec_store_inner(), chunk_size); + } + Iterator query_iter_copy = query_iter; + this->vec_store_meta_.template Optimize(std::move(query_iter_copy), vec_inners, mem_usage); + } + this->mem_usage_.fetch_add(mem_usage); + } + return AddVec(std::move(query_iter)); + } + + void Optimize() { + if constexpr (!VecStoreT::HasOptimize) { + return; + } + DenseVectorIter empty_iter(nullptr, this->dim(), 0); + AddVec(std::move(empty_iter)); + } + + void AddVertex(VertexType vec_i, i32 layer_n) { + auto [inner, idx] = this->GetInner(vec_i); + SizeT mem_usage = 0; + inner.AddVertex(idx, layer_n, this->graph_store_meta_, mem_usage); + this->mem_usage_.fetch_add(mem_usage); + } + Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i) { + auto [inner, idx] = this->GetInner(vertex_i); + return inner.GetNeighborsMut(idx, layer_i, this->graph_store_meta_); + } + + template + DataStore CompressToLVQ() &&; +}; + +export template +class DataStore : public DataStoreBase { +public: + using This = DataStore; + using VecStoreMeta = typename VecStoreT::template Meta; + using Base = DataStoreBase; + using Inner = DataStoreInner; private: - DataStoreInner(SizeT chunk_size, VecStoreInner vec_store_inner, GraphStoreInner graph_store_inner) - : vec_store_inner_(std::move(vec_store_inner)), graph_store_inner_(std::move(graph_store_inner)), - labels_(MakeUnique(chunk_size)), vertex_mutex_(MakeUnique(chunk_size)) {} + DataStore(SizeT chunk_size, SizeT max_chunk_n, VecStoreMeta vec_store_meta, GraphStoreMeta graph_store_meta, SizeT cur_vec_num) { + this->chunk_size_ = chunk_size; + this->max_chunk_n_ = max_chunk_n; + this->vec_store_meta_ = std::move(vec_store_meta); + this->graph_store_meta_ = std::move(graph_store_meta); + assert(chunk_size > 0); + assert((chunk_size & (chunk_size - 1)) == 0); + this->chunk_shift_ = __builtin_ctzll(chunk_size); + this->inners_ = MakeUnique(max_chunk_n); + this->cur_vec_num_.store(cur_vec_num); + } public: - DataStoreInner() = default; + static This LoadFromPtr(const char *&ptr, SizeT max_chunk_n = 0) { + SizeT chunk_size = ReadBufAdv(ptr); + SizeT max_chunk_n1 = ReadBufAdv(ptr); + if (max_chunk_n == 0) { + max_chunk_n = max_chunk_n1; + } + assert(max_chunk_n >= max_chunk_n1); - static This Make(SizeT chunk_size, VecStoreMeta &vec_store_meta, GraphStoreMeta &graph_store_meta, SizeT &mem_usage) { - auto vec_store_inner = VecStoreInner::Make(chunk_size, vec_store_meta, mem_usage); - auto graph_store_inner = GraphStoreInner::Make(chunk_size, graph_store_meta, mem_usage); - return This(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner)); + SizeT cur_vec_num = ReadBufAdv(ptr); + VecStoreMeta vec_store_meta = VecStoreMeta::LoadFromPtr(ptr); + GraphStoreMeta graph_store_meta = GraphStoreMeta::LoadFromPtr(ptr); + + This ret = This(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta), cur_vec_num); + + auto [chunk_num, last_chunk_size] = ret.ChunkInfo(cur_vec_num); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size : last_chunk_size; + ret.inners_[i] = Inner::LoadFromPtr(ptr, cur_chunk_size, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_); + } + return ret; } +}; + +#pragma clang diagnostic pop +//----------------------------------------------- Inner ----------------------------------------------- + +template +class DataStoreInnerBase { +public: + using This = DataStoreInner; + using DataType = typename VecStoreT::DataType; + using VecStoreInner = typename VecStoreT::template Inner; + using VecStoreMeta = typename VecStoreT::template Meta; + using GraphStoreInner = GraphStoreInner; + + friend class DataStoreIter; + +public: + DataStoreInnerBase() = default; void Free(SizeT cur_vec_num, const GraphStoreMeta &graph_store_meta) { graph_store_inner_.Free(cur_vec_num, graph_store_meta); } @@ -415,53 +465,15 @@ public: file_handle.Append(labels_.get(), sizeof(LabelType) * cur_vec_num); } - static This Load(LocalFileHandle &file_handle, - SizeT cur_vec_num, - SizeT chunk_size, - VecStoreMeta &vec_store_meta, - GraphStoreMeta &graph_store_meta, - SizeT &mem_usage) { - auto vec_store_inner = VecStoreInner::Load(file_handle, cur_vec_num, chunk_size, vec_store_meta, mem_usage); - auto graph_store_iner = GraphStoreInner::Load(file_handle, cur_vec_num, chunk_size, graph_store_meta, mem_usage); - This ret(chunk_size, std::move(vec_store_inner), std::move(graph_store_iner)); - file_handle.Read(ret.labels_.get(), sizeof(LabelType) * cur_vec_num); - return ret; - } - // vec store - template Iterator> - Pair AddVec(Iterator &&query_iter, VertexType start_idx, SizeT remain_num, const VecStoreMeta &meta, SizeT &mem_usage) { - SizeT insert_n = 0; - bool used_up = false; - while (insert_n < remain_num) { - if (auto ret = query_iter.Next(); ret) { - auto &[vec, label] = *ret; - vec_store_inner_.SetVec(start_idx + insert_n, vec, meta, mem_usage); - labels_[start_idx + insert_n] = label; - ++insert_n; - } else { - used_up = true; - break; - } - } - return {insert_n, used_up}; - } - typename VecStoreT::StoreType GetVec(VertexType vec_i, const VecStoreMeta &meta) const { return vec_store_inner_.GetVec(vec_i, meta); } void PrefetchVec(VertexType vec_i, const VecStoreMeta &meta) const { vec_store_inner_.Prefetch(vec_i, meta); } // graph store - void AddVertex(VertexType vec_i, i32 layer_n, const GraphStoreMeta &meta, SizeT &mem_usage) { - graph_store_inner_.AddVertex(vec_i, layer_n, meta, mem_usage); - } - Pair GetNeighbors(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) const { return graph_store_inner_.GetNeighbors(vertex_i, layer_i, meta); } - Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { - return graph_store_inner_.GetNeighborsMut(vertex_i, layer_i, meta); - } LabelType GetLabel(VertexType vec_i) const { return labels_[vec_i]; } @@ -477,9 +489,7 @@ public: protected: VecStoreInner vec_store_inner_; GraphStoreInner graph_store_inner_; - UniquePtr labels_; - -private: + ArrayPtr labels_; mutable UniquePtr vertex_mutex_; public: @@ -499,6 +509,98 @@ public: void DumpGraph(std::ostream &os, SizeT chunk_size, const GraphStoreMeta &meta) const { graph_store_inner_.Dump(os, chunk_size, meta); } }; +template +class DataStoreInner : public DataStoreInnerBase { +private: + using This = DataStoreInner; + using VecStoreInner = typename VecStoreT::template Inner; + using VecStoreMeta = typename VecStoreT::template Meta; + using GraphStoreInner = GraphStoreInner; + using QueryVecType = typename VecStoreT::QueryVecType; + + DataStoreInner(SizeT chunk_size, VecStoreInner vec_store_inner, GraphStoreInner graph_store_inner) { + this->vec_store_inner_ = std::move(vec_store_inner); + this->graph_store_inner_ = std::move(graph_store_inner); + this->labels_ = MakeUnique(chunk_size); + this->vertex_mutex_ = MakeUnique(chunk_size); + } + +public: + DataStoreInner() = default; + static This Make(SizeT chunk_size, VecStoreMeta &vec_store_meta, GraphStoreMeta &graph_store_meta, SizeT &mem_usage) { + auto vec_store_inner = VecStoreInner::Make(chunk_size, vec_store_meta, mem_usage); + auto graph_store_inner = GraphStoreInner::Make(chunk_size, graph_store_meta, mem_usage); + return This(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner)); + } + + static This Load(LocalFileHandle &file_handle, + SizeT cur_vec_num, + SizeT chunk_size, + VecStoreMeta &vec_store_meta, + GraphStoreMeta &graph_store_meta, + SizeT &mem_usage) { + auto vec_store_inner = VecStoreInner::Load(file_handle, cur_vec_num, chunk_size, vec_store_meta, mem_usage); + auto graph_store_iner = GraphStoreInner::Load(file_handle, cur_vec_num, chunk_size, graph_store_meta, mem_usage); + This ret(chunk_size, std::move(vec_store_inner), std::move(graph_store_iner)); + file_handle.Read(ret.labels_.get(), sizeof(LabelType) * cur_vec_num); + return ret; + } + + // vec store + template Iterator> + Pair AddVec(Iterator &&query_iter, VertexType start_idx, SizeT remain_num, const VecStoreMeta &meta, SizeT &mem_usage) { + SizeT insert_n = 0; + bool used_up = false; + while (insert_n < remain_num) { + if (auto ret = query_iter.Next(); ret) { + auto &[vec, label] = *ret; + this->vec_store_inner_.SetVec(start_idx + insert_n, vec, meta, mem_usage); + this->labels_[start_idx + insert_n] = label; + ++insert_n; + } else { + used_up = true; + break; + } + } + return {insert_n, used_up}; + } + + // graph store + void AddVertex(VertexType vec_i, i32 layer_n, const GraphStoreMeta &meta, SizeT &mem_usage) { + this->graph_store_inner_.AddVertex(vec_i, layer_n, meta, mem_usage); + } + Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { + return this->graph_store_inner_.GetNeighborsMut(vertex_i, layer_i, meta); + } +}; + +template +class DataStoreInner : public DataStoreInnerBase { +public: + using This = DataStoreInner; + using VecStoreInner = typename VecStoreT::Inner; + using VecStoreMeta = typename VecStoreT::Meta; + using GraphStoreInner = GraphStoreInner; + +private: + DataStoreInner(SizeT chunk_size, VecStoreInner vec_store_inner, GraphStoreInner graph_store_inner, LabelType *labels) { + this->vec_store_inner_ = std::move(vec_store_inner); + this->graph_store_inner_ = std::move(graph_store_inner); + this->labels_ = labels; + this->vertex_mutex_ = MakeUnique(chunk_size); + } + +public: + static This + LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT chunk_size, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) { + auto vec_store_inner = VecStoreInner::LoadFromPtr(ptr, cur_vec_num, vec_store_meta); + auto graph_store_inner = GraphStoreInner::LoadFromPtr(ptr, cur_vec_num, chunk_size, graph_store_meta); + auto *labels = reinterpret_cast(ptr); + ptr += sizeof(LabelType) * cur_vec_num; + return This(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner), labels); + } +}; + template class DataStoreChunkIter { public: @@ -526,11 +628,11 @@ private: SizeT last_chunk_size_; }; -template +template class DataStoreInnerIter { public: - using VecMeta = typename VecStoreT::Meta; - using Inner = DataStoreInner; + using VecMeta = typename VecStoreT::template Meta; + using Inner = DataStoreInner; using StoreType = typename VecStoreT::StoreType; DataStoreInnerIter(const VecMeta *vec_meta, const Inner *inner, SizeT max_vec_num) @@ -553,11 +655,11 @@ private: SizeT cur_idx_; }; -template +template class DataStoreIter { public: using StoreType = typename VecStoreT::StoreType; - using InnerIter = DataStoreInnerIter; + using InnerIter = DataStoreInnerIter; DataStoreIter(const DataStore *data_store) : data_store_iter_(data_store), inner_iter_(None) {} @@ -584,20 +686,24 @@ private: Optional inner_iter_; }; -template +template template -DataStore DataStore::CompressToLVQ() && { +DataStore DataStore::CompressToLVQ() && { if constexpr (std::is_same_v) { return std::move(*this); } else { - const auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num()); - Vector graph_inners; + const auto [chunk_num, last_chunk_size] = this->ChunkInfo(this->cur_vec_num()); + Vector> graph_inners; for (SizeT i = 0; i < chunk_num; ++i) { - graph_inners.emplace_back(std::move(*inners_[i].graph_store_inner())); + graph_inners.emplace_back(std::move(*this->inners_[i].graph_store_inner())); } - auto ret = DataStore::Make(chunk_size_, max_chunk_n_, vec_store_meta_.dim(), Mmax0(), Mmax()); - ret.OptAddVec(DataStoreIter(this)); - ret.SetGraph(std::move(graph_store_meta_), std::move(graph_inners)); + auto ret = DataStore::Make(this->chunk_size_, + this->max_chunk_n_, + this->vec_store_meta_.dim(), + this->Mmax0(), + this->Mmax()); + ret.OptAddVec(DataStoreIter(this)); + ret.SetGraph(std::move(this->graph_store_meta_), std::move(graph_inners)); this->inners_ = nullptr; return ret; } diff --git a/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm b/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm new file mode 100644 index 0000000000..d4111b3949 --- /dev/null +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm @@ -0,0 +1,52 @@ +// Copyright(C) 2023 InfiniFlow, Inc. All rights reserved. +// +// 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 +// +// https://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. + +module; + +export module data_store_util; + +import stl; + +namespace infinity { + +export template +class ArrayPtr { +public: + ArrayPtr() = default; + ArrayPtr(UniquePtr ptr) : ptr_(std::move(ptr)) {} + + T &operator[](SizeT idx) { return ptr_[idx]; } + const T &operator[](SizeT idx) const { return ptr_[idx]; } + + T *get() const { return ptr_.get(); } + +private: + UniquePtr ptr_; +}; + +export template +class ArrayPtr { +public: + ArrayPtr() = default; + ArrayPtr(const T *ptr) : ptr_(ptr) {} + + const T &operator[](SizeT idx) const { return ptr_[idx]; } + + const T *get() const { return ptr_; } + +private: + const T *ptr_ = nullptr; +}; + +} // namespace infinity 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 index 0c0c48a689..b2d22278af 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm @@ -22,6 +22,8 @@ export module graph_store; import stl; import hnsw_common; import local_file_handle; +import data_store_util; +import serialize; namespace infinity { @@ -91,6 +93,17 @@ public: return meta; } + static GraphStoreMeta LoadFromPtr(const char *&ptr) { + SizeT Mmax0 = ReadBufAdv(ptr); + SizeT Mmax = ReadBufAdv(ptr); + GraphStoreMeta meta(Mmax0, Mmax); + i32 max_layer = ReadBufAdv(ptr); + VertexType enterpoint = ReadBufAdv(ptr); + meta.max_layer_ = max_layer; + meta.enterpoint_ = enterpoint; + return meta; + } + SizeT Mmax0() const { return Mmax0_; } SizeT Mmax() const { return Mmax_; } SizeT level0_size() const { return level0_size_; } @@ -133,13 +146,10 @@ public: } }; -export class GraphStoreInner { -private: - GraphStoreInner(SizeT max_vertex, const GraphStoreMeta &meta, SizeT loaded_vertex_n) - : graph_(MakeUnique(max_vertex * meta.level0_size())), loaded_vertex_n_(loaded_vertex_n) {} - +template +class GraphStoreInnerBase { public: - GraphStoreInner() = default; + GraphStoreInnerBase() = default; void Free(SizeT current_vertex_num, const GraphStoreMeta &meta) { for (VertexType vertex_i = loaded_vertex_n_; vertex_i < VertexType(current_vertex_num); ++vertex_i) { @@ -147,13 +157,6 @@ public: } } - static GraphStoreInner Make(SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { - GraphStoreInner graph_store(max_vertex, meta, 0); - std::fill(graph_store.graph_.get(), graph_store.graph_.get() + max_vertex * meta.level0_size(), 0); - mem_usage += max_vertex * meta.level0_size(); - return graph_store; - } - SizeT GetSizeInBytes(SizeT cur_vertex_n, const GraphStoreMeta &meta) const { SizeT size = 0; for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { @@ -182,49 +185,6 @@ public: } } - static GraphStoreInner Load(LocalFileHandle &file_handle, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { - assert(cur_vertex_n <= max_vertex); - - SizeT layer_sum; - file_handle.Read(&layer_sum, sizeof(layer_sum)); - - GraphStoreInner graph_store(max_vertex, meta, cur_vertex_n); - file_handle.Read(graph_store.graph_.get(), cur_vertex_n * meta.level0_size()); - - auto loaded_layers = MakeUnique(meta.levelx_size() * layer_sum); - 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_) { - file_handle.Read(loaded_layers_p, meta.levelx_size() * v->layer_n_); - v->layers_p_ = loaded_layers_p; - loaded_layers_p += 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; - v->layer_n_ = layer_n; - if (layer_n) { - v->layers_p_ = new char[meta.levelx_size() * layer_n]; - mem_usage += meta.levelx_size() * layer_n; - - for (i32 layer_i = 1; layer_i <= layer_n; ++layer_i) { - GetLevelX(v->layers_p_, layer_i, meta)->neighbor_n_ = 0; - } - } else { - v->layers_p_ = nullptr; - } - } - Pair GetNeighbors(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) const { const VertexL0 *v = GetLevel0(vertex_i, meta); if (layer_i == 0) { @@ -233,36 +193,20 @@ public: const VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); return {vx->neighbors_, vx->neighbor_n_}; } - Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { - VertexL0 *v = GetLevel0(vertex_i, meta); - if (layer_i == 0) { - return {v->neighbors_, &v->neighbor_n_}; - } - VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); - return {vx->neighbors_, &vx->neighbor_n_}; - } -private: +protected: const VertexL0 *GetLevel0(VertexType vertex_i, const GraphStoreMeta &meta) const { return reinterpret_cast(graph_.get() + vertex_i * meta.level0_size()); } - VertexL0 *GetLevel0(VertexType vertex_i, const GraphStoreMeta &meta) { - return reinterpret_cast(graph_.get() + vertex_i * meta.level0_size()); - } const VertexLX *GetLevelX(const char *layer_p, i32 layer_i, const GraphStoreMeta &meta) const { assert(layer_i > 0); return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); } - VertexLX *GetLevelX(char *layer_p, i32 layer_i, const GraphStoreMeta &meta) { - assert(layer_i > 0); - return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); - } -private: - UniquePtr graph_; +protected: + ArrayPtr graph_; SizeT loaded_vertex_n_; - UniquePtr loaded_layers_; //---------------------------------------------- Following is the tmp debug function. ---------------------------------------------- @@ -336,4 +280,119 @@ public: } }; +export template +class GraphStoreInner : public GraphStoreInnerBase { +private: + GraphStoreInner(SizeT max_vertex, const GraphStoreMeta &meta, SizeT loaded_vertex_n) { + this->graph_ = MakeUnique(max_vertex * meta.level0_size()); + this->loaded_vertex_n_ = loaded_vertex_n; + } + +public: + GraphStoreInner() = default; + + static GraphStoreInner Make(SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { + GraphStoreInner graph_store(max_vertex, meta, 0); + std::fill(graph_store.graph_.get(), graph_store.graph_.get() + max_vertex * meta.level0_size(), 0); + mem_usage += max_vertex * meta.level0_size(); + return graph_store; + } + + static GraphStoreInner Load(LocalFileHandle &file_handle, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { + assert(cur_vertex_n <= max_vertex); + + SizeT layer_sum; + file_handle.Read(&layer_sum, sizeof(layer_sum)); + + GraphStoreInner graph_store(max_vertex, meta, cur_vertex_n); + file_handle.Read(graph_store.graph_.get(), cur_vertex_n * meta.level0_size()); + + auto loaded_layers = MakeUnique(meta.levelx_size() * layer_sum); + 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_) { + file_handle.Read(loaded_layers_p, meta.levelx_size() * v->layer_n_); + v->layers_p_ = loaded_layers_p; + loaded_layers_p += 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; + v->layer_n_ = layer_n; + if (layer_n) { + v->layers_p_ = new char[meta.levelx_size() * layer_n]; + mem_usage += meta.levelx_size() * layer_n; + + for (i32 layer_i = 1; layer_i <= layer_n; ++layer_i) { + GetLevelX(v->layers_p_, layer_i, meta)->neighbor_n_ = 0; + } + } else { + v->layers_p_ = nullptr; + } + } + + Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { + VertexL0 *v = GetLevel0(vertex_i, meta); + if (layer_i == 0) { + return {v->neighbors_, &v->neighbor_n_}; + } + VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); + return {vx->neighbors_, &vx->neighbor_n_}; + } + +private: + VertexL0 *GetLevel0(VertexType vertex_i, const GraphStoreMeta &meta) { + return reinterpret_cast(this->graph_.get() + vertex_i * meta.level0_size()); + } + VertexLX *GetLevelX(char *layer_p, i32 layer_i, const GraphStoreMeta &meta) { + assert(layer_i > 0); + return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); + } + +private: + ArrayPtr loaded_layers_; +}; + +export template <> +class GraphStoreInner : public GraphStoreInnerBase { +public: + GraphStoreInner() = default; + + GraphStoreInner(SizeT loaded_vertex_n, const char *ptr) { + this->graph_ = ptr; + this->loaded_vertex_n_ = loaded_vertex_n; + } + + static GraphStoreInner LoadFromPtr(const char *&ptr, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta) { + assert(cur_vertex_n <= max_vertex); + + [[maybe_unused]] SizeT layer_sum = ReadBufAdv(ptr); + + GraphStoreInner graph_store(cur_vertex_n, ptr); + ptr += cur_vertex_n * meta.level0_size(); + + for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { + auto *v = const_cast(graph_store.GetLevel0(vertex_i, meta)); + if (v->layer_n_) { + v->layers_p_ = const_cast(ptr); + ptr += meta.levelx_size() * v->layer_n_; + } else { + v->layers_p_ = nullptr; + } + } + + return graph_store; + } +}; + } // namespace infinity 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 index 54e0f6dd6b..e9e942ed0c 100644 --- 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 @@ -28,6 +28,8 @@ export module lvq_vec_store; import stl; import local_file_handle; import hnsw_common; +import serialize; +import data_store_util; namespace infinity { @@ -39,22 +41,14 @@ struct LVQData { CompressType compress_vec_[]; }; -export template +export template class LVQVecStoreInner; export template -class LVQVecStoreMeta { +class LVQVecStoreMetaType { public: - // Compress type must be i8 temporarily - static_assert(std::is_same() || std::is_same()); - constexpr static SizeT max_bucket_idx_ = std::numeric_limits::max() - std::numeric_limits::min(); // 255 for i8 - - using This = LVQVecStoreMeta; - using Inner = LVQVecStoreInner; using LocalCacheType = LVQCache::LocalCacheType; - using GlobalCacheType = LVQCache::GlobalCacheType; using LVQData = LVQData; - using StoreType = const LVQData *; struct LVQQuery { UniquePtr inner_; operator const LVQData *() const { return inner_.get(); } @@ -63,22 +57,32 @@ public: LVQQuery(LVQQuery &&other) = default; ~LVQQuery() { delete[] reinterpret_cast(inner_.release()); } }; + + using StoreType = const LVQData *; using QueryType = LVQQuery; using DistanceType = f32; +}; -private: - LVQVecStoreMeta(SizeT dim) : dim_(dim), compress_data_size_(sizeof(LVQData) + sizeof(CompressType) * dim) { - mean_ = MakeUnique(dim); - std::fill(mean_.get(), mean_.get() + dim, 0); - global_cache_ = LVQCache::MakeGlobalCache(mean_.get(), dim); - } +template +class LVQVecStoreMetaBase { +public: + // Compress type must be i8 temporarily + static_assert(std::is_same() || std::is_same()); + constexpr static SizeT max_bucket_idx_ = std::numeric_limits::max() - std::numeric_limits::min(); // 255 for i8 + + using This = LVQVecStoreMetaBase; + using Inner = LVQVecStoreInner; + using LocalCacheType = LVQCache::LocalCacheType; + using GlobalCacheType = LVQCache::GlobalCacheType; + using LVQData = LVQVecStoreMetaType::LVQData; + using LVQQuery = LVQVecStoreMetaType::LVQQuery; public: - LVQVecStoreMeta() : dim_(0), compress_data_size_(0), normalize_(false) {} - LVQVecStoreMeta(This &&other) + LVQVecStoreMetaBase() : dim_(0), compress_data_size_(0), normalize_(false) {} + LVQVecStoreMetaBase(This &&other) : dim_(std::exchange(other.dim_, 0)), compress_data_size_(std::exchange(other.compress_data_size_, 0)), mean_(std::move(other.mean_)), global_cache_(std::exchange(other.global_cache_, GlobalCacheType())), normalize_(other.normalize_) {} - LVQVecStoreMeta &operator=(This &&other) { + LVQVecStoreMetaBase &operator=(This &&other) { if (this != &other) { dim_ = std::exchange(other.dim_, 0); compress_data_size_ = std::exchange(other.compress_data_size_, 0); @@ -89,13 +93,6 @@ public: return *this; } - static This Make(SizeT dim) { return This(dim); } - static This Make(SizeT dim, bool normalize) { - This ret(dim); - ret.normalize_ = normalize; - return ret; - } - SizeT GetSizeInBytes() const { return sizeof(dim_) + sizeof(MeanType) * dim_ + sizeof(GlobalCacheType); } void Save(LocalFileHandle &file_handle) const { @@ -104,15 +101,6 @@ public: file_handle.Append(&global_cache_, sizeof(GlobalCacheType)); } - static This Load(LocalFileHandle &file_handle) { - SizeT dim; - file_handle.Read(&dim, sizeof(dim)); - This meta(dim); - file_handle.Read(meta.mean_.get(), sizeof(MeanType) * dim); - file_handle.Read(&meta.global_cache_, sizeof(GlobalCacheType)); - return meta; - } - LVQQuery MakeQuery(const DataType *vec) const { LVQQuery query(compress_data_size_); CompressTo(vec, query.inner_.get()); @@ -162,45 +150,6 @@ public: dest->local_cache_ = LVQCache::MakeLocalCache(compress, scale, dim_, mean_.get()); } - template Iterator> - void Optimize(Iterator &&query_iter, const Vector> &inners, SizeT &mem_usage) { - auto new_mean = MakeUnique(dim_); - auto temp_decompress = MakeUnique(dim_); - SizeT cur_vec_num = 0; - for (const auto [inner, size] : inners) { - for (SizeT i = 0; i < size; ++i) { - DecompressTo(inner->GetVec(i, *this), temp_decompress.get()); - for (SizeT j = 0; j < dim_; ++j) { - new_mean[j] += temp_decompress[j]; - } - } - cur_vec_num += size; - } - while (true) { - if (auto ret = query_iter.Next(); ret) { - auto &[vec, _] = *ret; - for (SizeT i = 0; i < dim_; ++i) { - new_mean[i] += vec[i]; - } - ++cur_vec_num; - } else { - break; - } - } - for (SizeT i = 0; i < dim_; ++i) { - new_mean[i] /= cur_vec_num; - } - swap(new_mean, mean_); - - for (auto [inner, size] : inners) { - for (SizeT i = 0; i < size; ++i) { - DecompressByMeanTo(inner->GetVec(i, *this), new_mean.get(), temp_decompress.get()); - inner->SetVec(i, temp_decompress.get(), *this, mem_usage); - } - } - global_cache_ = LVQCache::MakeGlobalCache(mean_.get(), dim_); - } - SizeT dim() const { return dim_; } SizeT compress_data_size() const { return compress_data_size_; } @@ -209,7 +158,7 @@ public: // for unit test const MeanType *mean() const { return mean_.get(); } -private: +protected: void DecompressByMeanTo(const LVQData *src, const MeanType *mean, DataType *dest) const { const CompressType *compress = src->compress_vec_; DataType scale = src->scale_; @@ -221,11 +170,11 @@ private: void DecompressTo(const LVQData *src, DataType *dest) const { DecompressByMeanTo(src, mean_.get(), dest); }; -private: +protected: SizeT dim_; SizeT compress_data_size_; - UniquePtr mean_; + ArrayPtr mean_; GlobalCacheType global_cache_; bool normalize_{false}; @@ -242,42 +191,126 @@ public: } }; -export template -class LVQVecStoreInner { -public: - using This = LVQVecStoreInner; - using Meta = LVQVecStoreMeta; - // Decompress: Q = scale * C + bias + Mean +export template +class LVQVecStoreMeta : public LVQVecStoreMetaBase { + using This = LVQVecStoreMeta; + using Inner = LVQVecStoreInner; using LocalCacheType = LVQCache::LocalCacheType; using LVQData = LVQData; + using GlobalCacheType = LVQCache::GlobalCacheType; private: - LVQVecStoreInner(SizeT max_vec_num, const Meta &meta) : ptr_(MakeUnique(max_vec_num * meta.compress_data_size())) {} + LVQVecStoreMeta(SizeT dim) { + this->dim_ = dim; + this->compress_data_size_ = sizeof(LVQData) + sizeof(CompressType) * dim; + this->mean_ = MakeUnique(dim); + std::fill(this->mean_.get(), this->mean_.get() + dim, 0); + this->global_cache_ = LVQCache::MakeGlobalCache(this->mean_.get(), dim); + } public: - LVQVecStoreInner() = default; - - static This Make(SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { - auto ret = This(max_vec_num, meta); - mem_usage += max_vec_num * meta.compress_data_size(); + LVQVecStoreMeta() = default; + static This Make(SizeT dim) { return This(dim); } + static This Make(SizeT dim, bool normalize) { + This ret(dim); + ret.normalize_ = normalize; return ret; } - SizeT GetSizeInBytes(SizeT cur_vec_num, const Meta &meta) const { return cur_vec_num * meta.compress_data_size(); } + static This Load(LocalFileHandle &file_handle) { + SizeT dim; + file_handle.Read(&dim, sizeof(dim)); + This meta(dim); + file_handle.Read(meta.mean_.get(), sizeof(MeanType) * dim); + file_handle.Read(&meta.global_cache_, sizeof(GlobalCacheType)); + return meta; + } - void Save(LocalFileHandle &file_handle, SizeT cur_vec_num, const Meta &meta) const { - file_handle.Append(ptr_.get(), cur_vec_num * meta.compress_data_size()); + template Iterator> + void Optimize(Iterator &&query_iter, const Vector> &inners, SizeT &mem_usage) { + auto new_mean = MakeUnique(this->dim_); + auto temp_decompress = MakeUnique(this->dim_); + SizeT cur_vec_num = 0; + for (const auto [inner, size] : inners) { + for (SizeT i = 0; i < size; ++i) { + this->DecompressTo(inner->GetVec(i, *this), temp_decompress.get()); + for (SizeT j = 0; j < this->dim_; ++j) { + new_mean[j] += temp_decompress[j]; + } + } + cur_vec_num += size; + } + while (true) { + if (auto ret = query_iter.Next(); ret) { + auto &[vec, _] = *ret; + for (SizeT i = 0; i < this->dim_; ++i) { + new_mean[i] += vec[i]; + } + ++cur_vec_num; + } else { + break; + } + } + for (SizeT i = 0; i < this->dim_; ++i) { + new_mean[i] /= cur_vec_num; + } + this->mean_ = std::move(new_mean); + + for (auto [inner, size] : inners) { + for (SizeT i = 0; i < size; ++i) { + this->DecompressByMeanTo(inner->GetVec(i, *this), new_mean.get(), temp_decompress.get()); + inner->SetVec(i, temp_decompress.get(), *this, mem_usage); + } + } + this->global_cache_ = LVQCache::MakeGlobalCache(this->mean_.get(), this->dim_); } +}; - static This Load(LocalFileHandle &file_handle, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { - assert(cur_vec_num <= max_vec_num); - This ret(max_vec_num, meta); - file_handle.Read(ret.ptr_.get(), cur_vec_num * meta.compress_data_size()); - mem_usage += max_vec_num * meta.compress_data_size(); - return ret; +export template +class LVQVecStoreMeta : public LVQVecStoreMetaBase { + using This = LVQVecStoreMeta; + using LocalCacheType = LVQCache::LocalCacheType; + using LVQData = LVQData; + using GlobalCacheType = LVQCache::GlobalCacheType; + +private: + LVQVecStoreMeta(SizeT dim, MeanType *mean, GlobalCacheType global_cache) { + this->dim_ = dim; + this->compress_data_size_ = sizeof(LVQData) + sizeof(CompressType) * dim; + this->mean_ = mean; + this->global_cache_ = global_cache; } - void SetVec(SizeT idx, const DataType *vec, const Meta &meta, SizeT &mem_usage) { meta.CompressTo(vec, GetVecMut(idx, meta)); } +public: + LVQVecStoreMeta() = default; + + static This LoadFromPtr(const char *&ptr) { + SizeT dim = ReadBufAdv(ptr); + auto *mean = reinterpret_cast(const_cast(ptr)); + ptr += sizeof(MeanType) * dim; + GlobalCacheType global_cache = ReadBufAdv(ptr); + This meta(dim, mean, global_cache); + return meta; + } +}; + +template +class LVQVecStoreInnerBase { +public: + using This = LVQVecStoreInner; + using Meta = LVQVecStoreMetaBase; + // Decompress: Q = scale * C + bias + Mean + using LocalCacheType = LVQCache::LocalCacheType; + using LVQData = LVQData; + +public: + LVQVecStoreInnerBase() = default; + + SizeT GetSizeInBytes(SizeT cur_vec_num, const Meta &meta) const { return cur_vec_num * meta.compress_data_size(); } + + void Save(LocalFileHandle &file_handle, SizeT cur_vec_num, const Meta &meta) const { + file_handle.Append(ptr_.get(), cur_vec_num * meta.compress_data_size()); + } const LVQData *GetVec(SizeT idx, const Meta &meta) const { return reinterpret_cast(ptr_.get() + idx * meta.compress_data_size()); @@ -285,11 +318,8 @@ public: void Prefetch(VertexType vec_i, const Meta &meta) const { _mm_prefetch(reinterpret_cast(GetVec(vec_i, meta)), _MM_HINT_T0); } -private: - LVQData *GetVecMut(SizeT idx, const Meta &meta) { return reinterpret_cast(ptr_.get() + idx * meta.compress_data_size()); } - -private: - UniquePtr ptr_; +protected: + ArrayPtr ptr_; public: void Dump(std::ostream &os, SizeT offset, SizeT chunk_size, const Meta &meta) const { @@ -307,4 +337,56 @@ public: } }; +export template +class LVQVecStoreInner : public LVQVecStoreInnerBase { + using This = LVQVecStoreInner; + using Meta = LVQVecStoreMetaBase; + using LocalCacheType = LVQCache::LocalCacheType; + using LVQData = LVQData; + +private: + LVQVecStoreInner(SizeT max_vec_num, const Meta &meta) { this->ptr_ = MakeUnique(max_vec_num * meta.compress_data_size()); } + +public: + LVQVecStoreInner() = default; + + static This Make(SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { + auto ret = This(max_vec_num, meta); + mem_usage += max_vec_num * meta.compress_data_size(); + return ret; + } + + static This Load(LocalFileHandle &file_handle, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { + assert(cur_vec_num <= max_vec_num); + This ret(max_vec_num, meta); + file_handle.Read(ret.ptr_.get(), 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: + LVQData *GetVecMut(SizeT idx, const Meta &meta) { return reinterpret_cast(this->ptr_.get() + idx * meta.compress_data_size()); } +}; + +export template +class LVQVecStoreInner : public LVQVecStoreInnerBase { + using This = LVQVecStoreInner; + using Meta = LVQVecStoreMetaBase; + +private: + LVQVecStoreInner(const char *ptr) { this->ptr_ = ptr; } + +public: + LVQVecStoreInner() = default; + + static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, const Meta &meta) { + const char *p = ptr; + This ret(p); + ptr += cur_vec_num * meta.compress_data_size(); + return ret; + } +}; + } // namespace infinity \ No newline at end of file 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 index e5d56b4724..808dc319b4 100644 --- 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 @@ -27,6 +27,8 @@ export module plain_vec_store; import stl; import local_file_handle; import hnsw_common; +import serialize; +import data_store_util; namespace infinity { @@ -64,6 +66,11 @@ public: return This(dim); } + static This LoadFromPtr(const char *&ptr) { + SizeT dim = ReadBufAdv(ptr); + return This(dim); + } + QueryType MakeQuery(const DataType *vec) const { return vec; } SizeT dim() const { return dim_; } @@ -75,22 +82,14 @@ public: void Dump(std::ostream &os) const { os << "[CONST] dim: " << dim_ << std::endl; } }; -export template -class PlainVecStoreInner { +template +class PlainVecStoreInnerBase { public: - using This = PlainVecStoreInner; + using This = PlainVecStoreInnerBase; using Meta = PlainVecStoreMeta; -private: - PlainVecStoreInner(SizeT max_vec_num, const Meta &meta) : ptr_(MakeUnique(max_vec_num * meta.dim())) {} - public: - PlainVecStoreInner() = default; - - static This Make(SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { - mem_usage += sizeof(DataType) * max_vec_num * meta.dim(); - return This(max_vec_num, meta); - } + PlainVecStoreInnerBase() = default; SizeT GetSizeInBytes(SizeT cur_vec_num, const Meta &meta) const { return sizeof(DataType) * cur_vec_num * meta.dim(); } @@ -98,25 +97,12 @@ public: file_handle.Append(ptr_.get(), sizeof(DataType) * cur_vec_num * meta.dim()); } - static This Load(LocalFileHandle &file_handle, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { - assert(cur_vec_num <= max_vec_num); - This ret(max_vec_num, meta); - file_handle.Read(ret.ptr_.get(), 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)); } - const DataType *GetVec(SizeT idx, const Meta &meta) const { return ptr_.get() + idx * meta.dim(); } void Prefetch(VertexType vec_i, const Meta &meta) const { _mm_prefetch(reinterpret_cast(GetVec(vec_i, meta)), _MM_HINT_T0); } -private: - DataType *GetVecMut(SizeT idx, const Meta &meta) { return ptr_.get() + idx * meta.dim(); } - -private: - UniquePtr ptr_; +protected: + ArrayPtr ptr_; public: void Dump(std::ostream &os, SizeT offset, SizeT chunk_size, const Meta &meta) const { @@ -131,4 +117,51 @@ public: } }; +export template +class PlainVecStoreInner : public PlainVecStoreInnerBase { + using This = PlainVecStoreInner; + using Meta = PlainVecStoreMeta; + +protected: + PlainVecStoreInner(SizeT max_vec_num, const Meta &meta) { this->ptr_ = MakeUnique(max_vec_num * meta.dim()); } + +public: + PlainVecStoreInner() = default; + + static This Make(SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { + mem_usage += sizeof(DataType) * max_vec_num * meta.dim(); + return This(max_vec_num, meta); + } + + static This Load(LocalFileHandle &file_handle, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { + assert(cur_vec_num <= max_vec_num); + This ret(max_vec_num, meta); + file_handle.Read(ret.ptr_.get(), 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: + DataType *GetVecMut(SizeT idx, const Meta &meta) { return this->ptr_.get() + idx * meta.dim(); } +}; + +export template +class PlainVecStoreInner : public PlainVecStoreInnerBase { + using This = PlainVecStoreInner; + using Meta = PlainVecStoreMeta; + +protected: + PlainVecStoreInner(const DataType *ptr) { this->ptr_ = ptr; } + +public: + static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta) { + assert(cur_vec_num <= max_vec_num); + const auto *p = reinterpret_cast(ptr); // fixme + ptr += sizeof(DataType) * cur_vec_num * meta.dim(); + return This(p); + } +}; + } // namespace infinity \ No newline at end of file diff --git a/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm index bc4c184fa8..def76f99a0 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm @@ -43,6 +43,7 @@ private: SparseVecStoreMeta(SizeT max_dim) : max_dim_(max_dim) {} public: + SparseVecStoreMeta() = default; static This Make(SizeT max_dim) { return This(max_dim); } static This Make(SizeT max_dim, bool) { return This(max_dim); } diff --git a/src/storage/knn_index/knn_hnsw/data_store/vec_store_type.cppm b/src/storage/knn_index/knn_hnsw/data_store/vec_store_type.cppm index cfe6d986c3..61eeff5a93 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/vec_store_type.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/vec_store_type.cppm @@ -42,11 +42,13 @@ class PlainCosVecStoreType { public: using DataType = DataT; using CompressType = void; + template using Meta = PlainVecStoreMeta; - using Inner = PlainVecStoreInner; + template + using Inner = PlainVecStoreInner; using QueryVecType = const DataType *; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using StoreType = typename Meta::StoreType; + using QueryType = typename Meta::QueryType; using Distance = PlainCosDist; static constexpr bool HasOptimize = false; @@ -62,11 +64,13 @@ class PlainL2VecStoreType { public: using DataType = DataT; using CompressType = void; + template using Meta = PlainVecStoreMeta; - using Inner = PlainVecStoreInner; + template + using Inner = PlainVecStoreInner; using QueryVecType = const DataType *; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using StoreType = typename Meta::StoreType; + using QueryType = typename Meta::QueryType; using Distance = PlainL2Dist; static constexpr bool HasOptimize = false; @@ -82,11 +86,13 @@ class PlainIPVecStoreType { public: using DataType = DataT; using CompressType = void; + template using Meta = PlainVecStoreMeta; - using Inner = PlainVecStoreInner; + template + using Inner = PlainVecStoreInner; using QueryVecType = const DataType *; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using StoreType = typename Meta::StoreType; + using QueryType = typename Meta::QueryType; using Distance = PlainIPDist; static constexpr bool HasOptimize = false; @@ -102,11 +108,13 @@ class SparseIPVecStoreType { public: using DataType = DataT; using CompressType = void; + template using Meta = SparseVecStoreMeta; + template using Inner = SparseVecStoreInner; using QueryVecType = SparseVecRef; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using StoreType = typename Meta::StoreType; + using QueryType = typename Meta::QueryType; using Distance = SparseIPDist; static constexpr bool HasOptimize = false; @@ -122,11 +130,14 @@ class LVQCosVecStoreType { public: using DataType = DataT; using CompressType = CompressT; - using Meta = LVQVecStoreMeta>; - using Inner = LVQVecStoreInner>; + template + using Meta = LVQVecStoreMeta, OwnMem>; + template + using Inner = LVQVecStoreInner, OwnMem>; using QueryVecType = const DataType *; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using MetaType = LVQVecStoreMetaType>; + using StoreType = typename MetaType::StoreType; + using QueryType = typename MetaType::QueryType; using Distance = LVQCosDist; static constexpr bool HasOptimize = true; @@ -142,11 +153,14 @@ class LVQL2VecStoreType { public: using DataType = DataT; using CompressType = CompressT; - using Meta = LVQVecStoreMeta>; - using Inner = LVQVecStoreInner>; + template + using Meta = LVQVecStoreMeta, OwnMem>; + template + using Inner = LVQVecStoreInner, OwnMem>; using QueryVecType = const DataType *; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using MetaType = LVQVecStoreMetaType>; + using StoreType = typename MetaType::StoreType; + using QueryType = typename MetaType::QueryType; using Distance = LVQL2Dist; static constexpr bool HasOptimize = true; @@ -162,11 +176,14 @@ class LVQIPVecStoreType { public: using DataType = DataT; using CompressType = CompressT; - using Meta = LVQVecStoreMeta>; - using Inner = LVQVecStoreInner>; + template + using Meta = LVQVecStoreMeta, OwnMem>; + template + using Inner = LVQVecStoreInner, OwnMem>; using QueryVecType = const DataType *; - using StoreType = typename Meta::StoreType; - using QueryType = typename Meta::QueryType; + using MetaType = LVQVecStoreMetaType>; + using StoreType = typename MetaType::StoreType; + using QueryType = typename MetaType::QueryType; using Distance = LVQIPDist; static constexpr bool HasOptimize = true; diff --git a/src/storage/knn_index/knn_hnsw/dist_func_cos.cppm b/src/storage/knn_index/knn_hnsw/dist_func_cos.cppm index 2ea84c48e3..ca209a12c1 100644 --- a/src/storage/knn_index/knn_hnsw/dist_func_cos.cppm +++ b/src/storage/knn_index/knn_hnsw/dist_func_cos.cppm @@ -114,9 +114,9 @@ export template class LVQCosDist { public: using This = LVQCosDist; - using VecStoreMeta = LVQVecStoreMeta>; - using StoreType = typename VecStoreMeta::StoreType; - using DistanceType = typename VecStoreMeta::DistanceType; + using VecStoreMetaType = LVQVecStoreMetaType>; + using StoreType = typename VecStoreMetaType::StoreType; + using DistanceType = typename VecStoreMetaType::DistanceType; private: using SIMDFuncType = i32 (*)(const CompressType *, const CompressType *, SizeT); @@ -147,6 +147,7 @@ public: } } + template DataType operator()(const StoreType &v1, const StoreType &v2, const VecStoreMeta &vec_store_meta) const { SizeT dim = vec_store_meta.dim(); i32 c1c2_ip = SIMDFunc(v1->compress_vec_, v2->compress_vec_, dim); diff --git a/src/storage/knn_index/knn_hnsw/dist_func_ip.cppm b/src/storage/knn_index/knn_hnsw/dist_func_ip.cppm index 99f7e9360a..8230c3b36b 100644 --- a/src/storage/knn_index/knn_hnsw/dist_func_ip.cppm +++ b/src/storage/knn_index/knn_hnsw/dist_func_ip.cppm @@ -129,9 +129,9 @@ export template class LVQIPDist { public: using This = LVQIPDist; - using VecStoreMeta = LVQVecStoreMeta>; - using StoreType = typename VecStoreMeta::StoreType; - using DistanceType = typename VecStoreMeta::DistanceType; + using VecStoreMetaType = LVQVecStoreMetaType>; + using StoreType = typename VecStoreMetaType::StoreType; + using DistanceType = typename VecStoreMetaType::DistanceType; private: using SIMDFuncType = i32 (*)(const CompressType *, const CompressType *, SizeT); @@ -162,6 +162,7 @@ public: } } + template DataType operator()(const StoreType &v1, const StoreType &v2, const VecStoreMeta &vec_store_meta) const { SizeT dim = vec_store_meta.dim(); i32 c1c2_ip = SIMDFunc(v1->compress_vec_, v2->compress_vec_, dim); diff --git a/src/storage/knn_index/knn_hnsw/dist_func_l2.cppm b/src/storage/knn_index/knn_hnsw/dist_func_l2.cppm index f4859b43e7..cda83eb9bc 100644 --- a/src/storage/knn_index/knn_hnsw/dist_func_l2.cppm +++ b/src/storage/knn_index/knn_hnsw/dist_func_l2.cppm @@ -120,9 +120,9 @@ export template class LVQL2Dist { public: using This = LVQL2Dist; - using VecStoreMeta = LVQVecStoreMeta>; - using StoreType = typename VecStoreMeta::StoreType; - using DistanceType = typename VecStoreMeta::DistanceType; + using VecStoreMetaType = LVQVecStoreMetaType>; + using StoreType = typename VecStoreMetaType::StoreType; + using DistanceType = typename VecStoreMetaType::DistanceType; private: using SIMDFuncType = i32 (*)(const CompressType *, const CompressType *, SizeT); @@ -153,6 +153,7 @@ public: } } + template DataType operator()(const StoreType &v1, const StoreType &v2, const VecStoreMeta &vec_store_meta) const { SizeT dim = vec_store_meta.dim(); i32 c1c2_ip = SIMDFunc(v1->compress_vec_, v2->compress_vec_, dim); From 9064cf189f3872c4eb55fb3214f88e3887db6aa3 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Tue, 24 Dec 2024 16:00:19 +0800 Subject: [PATCH 02/12] Refactor hnsw with mmap. --- src/parser/type/serialize.h | 16 ++ .../knn_hnsw/data_store/data_store.cppm | 81 +++++----- .../knn_hnsw/data_store/data_store_util.cppm | 22 +++ .../knn_hnsw/data_store/graph_store.cppm | 137 ++++++++++------- .../knn_hnsw/data_store/lvq_vec_store.cppm | 8 +- .../knn_hnsw/data_store/plain_vec_store.cppm | 4 +- .../knn_hnsw/dist_func_sparse_ip.cppm | 1 + src/storage/knn_index/knn_hnsw/hnsw_alg.cppm | 138 +++++++++++------- .../storage/knnindex/knn_hnsw/test_hnsw.cpp | 93 +++++++++++- 9 files changed, 354 insertions(+), 146 deletions(-) diff --git a/src/parser/type/serialize.h b/src/parser/type/serialize.h index c218f7d8ea..9af9d125e9 100644 --- a/src/parser/type/serialize.h +++ b/src/parser/type/serialize.h @@ -57,6 +57,16 @@ inline std::string ReadBuf(const char *buf) { return str; } +template <> +inline std::tuple<> ReadBuf>(const char *buf) { + return {}; +} + +template <> +inline std::tuple<> ReadBufAdv>(const char *&buf) { + return {}; +} + template <> inline std::string ReadBufAdv(const char *&buf) { int32_t size = ReadBufAdv(buf); @@ -93,6 +103,12 @@ inline void WriteBufAdv(char *&buf, const std::string &value) { buf += len; } +template <> +inline void WriteBuf>(char *const buf, const std::tuple<> &) {} + +template <> +inline void WriteBufAdv>(char *&buf, const std::tuple<> &) {} + template inline void WriteBufVecAdv(char *&buf, const T *data, size_t size) { static_assert(std::is_standard_layout_v, "T must be POD"); 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 index 1bfcd18a00..cff059c810 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm @@ -34,7 +34,7 @@ namespace infinity { template class DataStoreInner; -export template +export template class DataStoreChunkIter; template @@ -52,7 +52,7 @@ public: using Inner = DataStoreInner; using VecStoreMeta = typename VecStoreT::template Meta; - friend class DataStoreChunkIter; + friend class DataStoreChunkIter; friend class DataStoreIter; public: @@ -117,22 +117,6 @@ public: return size; } - void Save(LocalFileHandle &file_handle) const { - SizeT cur_vec_num = this->cur_vec_num(); - auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); - - file_handle.Append(&chunk_size_, sizeof(chunk_size_)); - file_handle.Append(&max_chunk_n_, sizeof(max_chunk_n_)); - - file_handle.Append(&cur_vec_num, sizeof(cur_vec_num)); - vec_store_meta_.Save(file_handle); - graph_store_meta_.Save(file_handle); - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; - inners_[i].Save(file_handle, chunk_size, vec_store_meta_, graph_store_meta_); - } - } - typename VecStoreT::StoreType GetVec(SizeT vec_i) const { const auto &[inner, idx] = GetInner(vec_i); return inner.GetVec(idx, vec_store_meta_); @@ -272,6 +256,8 @@ private: } public: + DataStore() = default; + static This Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT Mmax0, SizeT Mmax) { bool normalize = false; if constexpr (Base::template has_compress_type::value) { @@ -288,6 +274,22 @@ public: return ret; } + void Save(LocalFileHandle &file_handle) { + SizeT cur_vec_num = this->cur_vec_num(); + auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); + + file_handle.Append(&this->chunk_size_, sizeof(this->chunk_size_)); + file_handle.Append(&this->max_chunk_n_, sizeof(this->max_chunk_n_)); + + file_handle.Append(&cur_vec_num, sizeof(cur_vec_num)); + this->vec_store_meta_.Save(file_handle); + this->graph_store_meta_.Save(file_handle, cur_vec_num); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? this->chunk_size_ : last_chunk_size; + this->inners_[i].Save(file_handle, chunk_size, this->vec_store_meta_, this->graph_store_meta_); + } + } + static This Load(LocalFileHandle &file_handle, SizeT max_chunk_n = 0) { SizeT chunk_size; file_handle.Read(&chunk_size, sizeof(chunk_size)); @@ -409,6 +411,7 @@ private: } public: + DataStore() = default; static This LoadFromPtr(const char *&ptr, SizeT max_chunk_n = 0) { SizeT chunk_size = ReadBufAdv(ptr); SizeT max_chunk_n1 = ReadBufAdv(ptr); @@ -459,12 +462,6 @@ public: return size; } - void Save(LocalFileHandle &file_handle, SizeT cur_vec_num, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) const { - vec_store_inner_.Save(file_handle, cur_vec_num, vec_store_meta); - graph_store_inner_.Save(file_handle, cur_vec_num, graph_store_meta); - file_handle.Append(labels_.get(), sizeof(LabelType) * cur_vec_num); - } - // vec store typename VecStoreT::StoreType GetVec(VertexType vec_i, const VecStoreMeta &meta) const { return vec_store_inner_.GetVec(vec_i, meta); } @@ -533,6 +530,12 @@ public: return This(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner)); } + void Save(LocalFileHandle &file_handle, SizeT cur_vec_num, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) { + this->vec_store_inner_.Save(file_handle, cur_vec_num, vec_store_meta); + this->graph_store_inner_.Save(file_handle, cur_vec_num, graph_store_meta); + file_handle.Append(this->labels_.get(), sizeof(LabelType) * cur_vec_num); + } + static This Load(LocalFileHandle &file_handle, SizeT cur_vec_num, SizeT chunk_size, @@ -578,12 +581,12 @@ template class DataStoreInner : public DataStoreInnerBase { public: using This = DataStoreInner; - using VecStoreInner = typename VecStoreT::Inner; - using VecStoreMeta = typename VecStoreT::Meta; + using VecStoreInner = typename VecStoreT::template Inner; + using VecStoreMeta = typename VecStoreT::template Meta; using GraphStoreInner = GraphStoreInner; private: - DataStoreInner(SizeT chunk_size, VecStoreInner vec_store_inner, GraphStoreInner graph_store_inner, LabelType *labels) { + DataStoreInner(SizeT chunk_size, VecStoreInner vec_store_inner, GraphStoreInner graph_store_inner, const LabelType *labels) { this->vec_store_inner_ = std::move(vec_store_inner); this->graph_store_inner_ = std::move(graph_store_inner); this->labels_ = labels; @@ -591,6 +594,8 @@ private: } public: + DataStoreInner() = default; + static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT chunk_size, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) { auto vec_store_inner = VecStoreInner::LoadFromPtr(ptr, cur_vec_num, vec_store_meta); @@ -601,12 +606,12 @@ public: } }; -template +template class DataStoreChunkIter { public: - using Inner = typename DataStore::Inner; + using Inner = typename DataStore::Inner; - DataStoreChunkIter(const DataStore *data_store) : data_store_(data_store) { + DataStoreChunkIter(const DataStore *data_store) : data_store_(data_store) { std::tie(chunk_num_, last_chunk_size_) = data_store_->ChunkInfo(data_store_->cur_vec_num()); } @@ -620,7 +625,7 @@ public: return ret; } - const DataStore *data_store_; + const DataStore *data_store_; private: SizeT cur_chunk_i_ = 0; @@ -661,7 +666,7 @@ public: using StoreType = typename VecStoreT::StoreType; using InnerIter = DataStoreInnerIter; - DataStoreIter(const DataStore *data_store) : data_store_iter_(data_store), inner_iter_(None) {} + DataStoreIter(const DataStore *data_store) : data_store_iter_(data_store), inner_iter_(None) {} Optional> Next() { if (!inner_iter_.has_value()) { @@ -682,7 +687,7 @@ public: } private: - DataStoreChunkIter data_store_iter_; + DataStoreChunkIter data_store_iter_; Optional inner_iter_; }; @@ -697,11 +702,11 @@ DataStore DataStoreinners_[i].graph_store_inner())); } - auto ret = DataStore::Make(this->chunk_size_, - this->max_chunk_n_, - this->vec_store_meta_.dim(), - this->Mmax0(), - this->Mmax()); + auto ret = DataStore::Make(this->chunk_size_, + this->max_chunk_n_, + this->vec_store_meta_.dim(), + this->Mmax0(), + this->Mmax()); ret.OptAddVec(DataStoreIter(this)); ret.SetGraph(std::move(this->graph_store_meta_), std::move(graph_inners)); this->inners_ = nullptr; diff --git a/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm b/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm index d4111b3949..c6ee0d7867 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm @@ -49,4 +49,26 @@ private: const T *ptr_ = nullptr; }; +export template +class PPtr { +public: + PPtr() = default; + void set(char *ptr) { ptr_ = ptr; } + char *get() const { return ptr_; } + +private: + char *ptr_; +}; + +export template <> +class PPtr { +public: + PPtr() = default; + void set(const char *ptr) { ptr_ = ptr; } + const char *get() const { return ptr_; } + +private: + const char *ptr_ = nullptr; +}; + } // namespace infinity 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 index b2d22278af..183735dc8f 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm @@ -42,19 +42,20 @@ export class GraphStoreMeta { private: GraphStoreMeta(SizeT Mmax0, SizeT Mmax) : Mmax0_(Mmax0), Mmax_(Mmax), level0_size_(sizeof(VertexL0) + sizeof(VertexType) * Mmax0), - levelx_size_(sizeof(VertexLX) + sizeof(VertexType) * Mmax) {} + levelx_size_(sizeof(VertexLX) + sizeof(VertexType) * Mmax), last_save_num_(0) {} public: - GraphStoreMeta() : Mmax0_(0), Mmax_(0), level0_size_(0), levelx_size_(0), max_layer_(-1), enterpoint_(-1) {} + GraphStoreMeta() : Mmax0_(0), Mmax_(0), level0_size_(0), levelx_size_(0), last_save_num_(0), max_layer_(-1), enterpoint_(-1) {} GraphStoreMeta(GraphStoreMeta &&other) : Mmax0_(std::exchange(other.Mmax0_, 0)), Mmax_(std::exchange(other.Mmax_, 0)), level0_size_(std::exchange(other.level0_size_, 0)), - levelx_size_(std::exchange(other.levelx_size_, 0)), max_layer_(std::exchange(other.max_layer_, -1)), - enterpoint_(std::exchange(other.enterpoint_, -1)) {} + levelx_size_(std::exchange(other.levelx_size_, 0)), last_save_num_(std::exchange(other.last_save_num_, 0)), + max_layer_(std::exchange(other.max_layer_, -1)), enterpoint_(std::exchange(other.enterpoint_, -1)) {} GraphStoreMeta &operator=(GraphStoreMeta &&other) { Mmax0_ = std::exchange(other.Mmax0_, 0); Mmax_ = std::exchange(other.Mmax_, 0); level0_size_ = std::exchange(other.level0_size_, 0); levelx_size_ = std::exchange(other.levelx_size_, 0); + last_save_num_ = std::exchange(other.last_save_num_, 0); max_layer_ = std::exchange(other.max_layer_, -1); enterpoint_ = std::exchange(other.enterpoint_, -1); return *this; @@ -70,9 +71,11 @@ public: SizeT GetSizeInBytes() const { return sizeof(Mmax0_) + sizeof(Mmax_) + sizeof(max_layer_) + sizeof(enterpoint_); } - void Save(LocalFileHandle &file_handle) const { + void Save(LocalFileHandle &file_handle, SizeT cur_vec_num) { file_handle.Append(&Mmax0_, sizeof(Mmax0_)); file_handle.Append(&Mmax_, sizeof(Mmax_)); + last_save_num_ = cur_vec_num; + file_handle.Append(&last_save_num_, sizeof(last_save_num_)); file_handle.Append(&max_layer_, sizeof(max_layer_)); file_handle.Append(&enterpoint_, sizeof(enterpoint_)); @@ -82,8 +85,12 @@ public: SizeT Mmax0, Mmax; file_handle.Read(&Mmax0, sizeof(Mmax0)); file_handle.Read(&Mmax, sizeof(Mmax)); + SizeT last_save_num; + file_handle.Read(&last_save_num, sizeof(last_save_num)); GraphStoreMeta meta(Mmax0, Mmax); + meta.last_save_num_ = last_save_num; + i32 max_layer; VertexType enterpoint; file_handle.Read(&max_layer, sizeof(max_layer)); @@ -96,7 +103,9 @@ public: static GraphStoreMeta LoadFromPtr(const char *&ptr) { SizeT Mmax0 = ReadBufAdv(ptr); SizeT Mmax = ReadBufAdv(ptr); + SizeT last_save_num = ReadBufAdv(ptr); GraphStoreMeta meta(Mmax0, Mmax); + meta.last_save_num_ = last_save_num; i32 max_layer = ReadBufAdv(ptr); VertexType enterpoint = ReadBufAdv(ptr); meta.max_layer_ = max_layer; @@ -108,6 +117,7 @@ public: SizeT Mmax() const { return Mmax_; } SizeT level0_size() const { return level0_size_; } SizeT levelx_size() const { return levelx_size_; } + SizeT last_save_num() const { return last_save_num_; } Pair GetEnterPoint() const { std::unique_lock lck(mtx_); @@ -132,6 +142,7 @@ private: SizeT Mmax_; SizeT level0_size_; SizeT levelx_size_; + SizeT last_save_num_; mutable std::mutex mtx_; i32 max_layer_; @@ -163,34 +174,19 @@ public: const VertexL0 *v = GetLevel0(vertex_i, meta); size += sizeof(v->layer_n_) + sizeof(v->neighbor_n_) + sizeof(VertexType) * v->neighbor_n_; for (i32 layer_i = 1; layer_i <= v->layer_n_; ++layer_i) { - const VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); + const VertexLX *vx = GetLevelX(v->layers_p_, vertex_i, layer_i, meta); size += sizeof(vx->neighbor_n_) + sizeof(VertexType) * vx->neighbor_n_; } } return size; } - void Save(LocalFileHandle &file_handle, SizeT cur_vertex_n, const GraphStoreMeta &meta) const { - SizeT layer_sum = 0; - for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { - layer_sum += GetLevel0(vertex_i, meta)->layer_n_; - } - file_handle.Append(&layer_sum, sizeof(layer_sum)); - file_handle.Append(graph_.get(), cur_vertex_n * meta.level0_size()); - for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { - const VertexL0 *v = GetLevel0(vertex_i, meta); - if (v->layer_n_) { - file_handle.Append(v->layers_p_, meta.levelx_size() * v->layer_n_); - } - } - } - Pair GetNeighbors(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) const { const VertexL0 *v = GetLevel0(vertex_i, meta); if (layer_i == 0) { return {v->neighbors_, v->neighbor_n_}; } - const VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); + const VertexLX *vx = GetLevelX(v->layers_p_, vertex_i, layer_i, meta); return {vx->neighbors_, vx->neighbor_n_}; } @@ -199,13 +195,24 @@ protected: return reinterpret_cast(graph_.get() + vertex_i * meta.level0_size()); } - const VertexLX *GetLevelX(const char *layer_p, i32 layer_i, const GraphStoreMeta &meta) const { + const VertexLX *GetLevelX(const char *layer_p, VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) const { assert(layer_i > 0); - return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); + if constexpr (OwnMem) { + if (SizeT(vertex_i) >= meta.last_save_num()) { + return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); + } + SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); + return reinterpret_cast(layer_start_.get() + offset); + } else { + assert(SizeT(vertex_i) < meta.last_save_num()); + SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); + return reinterpret_cast(layer_start_.get() + offset); + } } protected: ArrayPtr graph_; + PPtr layer_start_; SizeT loaded_vertex_n_; //---------------------------------------------- Following is the tmp debug function. ---------------------------------------------- @@ -227,7 +234,7 @@ public: assert(neighbor_idx != out_vertex_i); } for (int layer_i = 1; layer_i <= v->layer_n_; ++layer_i) { - const VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); + const VertexLX *vx = GetLevelX(v->layers_p_, vertex_i, layer_i, meta); for (int i = 0; i < vx->neighbor_n_; ++i) { VertexType neighbor_idx = vx->neighbors_[i]; assert(neighbor_idx < (VertexType)cur_vec_num && neighbor_idx >= 0); @@ -267,7 +274,7 @@ public: neighbors = v->neighbors_; neighbor_n = v->neighbor_n_; } else { - const VertexLX *vx = GetLevelX(v->layers_p_, layer, meta); + const VertexLX *vx = GetLevelX(v->layers_p_, vertex_i, layer, meta); neighbors = vx->neighbors_; neighbor_n = vx->neighbor_n_; } @@ -283,6 +290,7 @@ public: export template class GraphStoreInner : public GraphStoreInnerBase { private: + using Base = GraphStoreInnerBase; GraphStoreInner(SizeT max_vertex, const GraphStoreMeta &meta, SizeT loaded_vertex_n) { this->graph_ = MakeUnique(max_vertex * meta.level0_size()); this->loaded_vertex_n_ = loaded_vertex_n; @@ -298,6 +306,32 @@ public: return graph_store; } + void Save(LocalFileHandle &file_handle, SizeT cur_vertex_n, const GraphStoreMeta &meta) { + SizeT layer_sum = 0; + Vector layers_ptrs; + for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { + VertexL0 *v = GetLevel0(vertex_i, meta); + if (!v->layer_n_) { + continue; + } + layers_ptrs.emplace_back(v->layers_p_); + SizeT offset = layer_sum * meta.levelx_size(); + v->layers_p_ = reinterpret_cast(offset); + layer_sum += v->layer_n_; + } + file_handle.Append(&layer_sum, sizeof(layer_sum)); + file_handle.Append(this->graph_.get(), cur_vertex_n * meta.level0_size()); + for (VertexType vertex_i = 0, ptr_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { + VertexL0 *v = GetLevel0(vertex_i, meta); + if (v->layer_n_) { + char *ptr = layers_ptrs[ptr_i++]; + file_handle.Append(ptr, meta.levelx_size() * v->layer_n_); + } + } + layers_ptrs_.insert(layers_ptrs_.end(), layers_ptrs.begin(), layers_ptrs.end()); + this->loaded_vertex_n_ = cur_vertex_n; + } + static GraphStoreInner Load(LocalFileHandle &file_handle, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { assert(cur_vertex_n <= max_vertex); @@ -308,23 +342,23 @@ public: file_handle.Read(graph_store.graph_.get(), cur_vertex_n * meta.level0_size()); auto loaded_layers = MakeUnique(meta.levelx_size() * layer_sum); - 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_) { - file_handle.Read(loaded_layers_p, meta.levelx_size() * v->layer_n_); - v->layers_p_ = loaded_layers_p; - loaded_layers_p += meta.levelx_size() * v->layer_n_; - } else { - v->layers_p_ = nullptr; - } - } + file_handle.Read(loaded_layers.get(), meta.levelx_size() * layer_sum); + + graph_store.layer_start_.set(loaded_layers.get()); graph_store.loaded_layers_ = std::move(loaded_layers); mem_usage += max_vertex * meta.level0_size() + layer_sum * meta.levelx_size(); return graph_store; } + void Free(SizeT current_vertex_num, const GraphStoreMeta &meta) { + for (const char *layer_p : layers_ptrs_) { + delete[] layer_p; + } + layers_ptrs_.clear(); + Base::Free(current_vertex_num, meta); + } + void AddVertex(VertexType vertex_i, i32 layer_n, const GraphStoreMeta &meta, SizeT &mem_usage) { VertexL0 *v = GetLevel0(vertex_i, meta); v->neighbor_n_ = 0; @@ -334,7 +368,8 @@ public: mem_usage += meta.levelx_size() * layer_n; for (i32 layer_i = 1; layer_i <= layer_n; ++layer_i) { - GetLevelX(v->layers_p_, layer_i, meta)->neighbor_n_ = 0; + VertexLX *vx = GetLevelX(v->layers_p_, vertex_i, layer_i, meta); + vx->neighbor_n_ = 0; } } else { v->layers_p_ = nullptr; @@ -346,7 +381,7 @@ public: if (layer_i == 0) { return {v->neighbors_, &v->neighbor_n_}; } - VertexLX *vx = GetLevelX(v->layers_p_, layer_i, meta); + VertexLX *vx = GetLevelX(v->layers_p_, vertex_i, layer_i, meta); return {vx->neighbors_, &vx->neighbor_n_}; } @@ -354,13 +389,18 @@ private: VertexL0 *GetLevel0(VertexType vertex_i, const GraphStoreMeta &meta) { return reinterpret_cast(this->graph_.get() + vertex_i * meta.level0_size()); } - VertexLX *GetLevelX(char *layer_p, i32 layer_i, const GraphStoreMeta &meta) { + VertexLX *GetLevelX(char *layer_p, VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { assert(layer_i > 0); - return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); + if (SizeT(vertex_i) >= meta.last_save_num()) { + return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); + } + SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); + return reinterpret_cast(this->layer_start_.get() + offset); } private: ArrayPtr loaded_layers_; + Vector layers_ptrs_; }; export template <> @@ -376,20 +416,11 @@ public: static GraphStoreInner LoadFromPtr(const char *&ptr, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta) { assert(cur_vertex_n <= max_vertex); - [[maybe_unused]] SizeT layer_sum = ReadBufAdv(ptr); + SizeT layer_sum = ReadBufAdv(ptr); GraphStoreInner graph_store(cur_vertex_n, ptr); - ptr += cur_vertex_n * meta.level0_size(); - - for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { - auto *v = const_cast(graph_store.GetLevel0(vertex_i, meta)); - if (v->layer_n_) { - v->layers_p_ = const_cast(ptr); - ptr += meta.levelx_size() * v->layer_n_; - } else { - v->layers_p_ = nullptr; - } - } + graph_store.layer_start_.set(ptr + cur_vertex_n * meta.level0_size()); + ptr += cur_vertex_n * meta.level0_size() + layer_sum * meta.levelx_size(); return graph_store; } 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 index e9e942ed0c..4a090f2682 100644 --- 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 @@ -98,7 +98,9 @@ public: void Save(LocalFileHandle &file_handle) const { file_handle.Append(&dim_, sizeof(dim_)); file_handle.Append(mean_.get(), sizeof(MeanType) * dim_); - file_handle.Append(&global_cache_, sizeof(GlobalCacheType)); + if constexpr (!std::same_as>) { + file_handle.Append(&global_cache_, sizeof(GlobalCacheType)); + } } LVQQuery MakeQuery(const DataType *vec) const { @@ -222,7 +224,9 @@ public: file_handle.Read(&dim, sizeof(dim)); This meta(dim); file_handle.Read(meta.mean_.get(), sizeof(MeanType) * dim); - file_handle.Read(&meta.global_cache_, sizeof(GlobalCacheType)); + if constexpr (!std::is_same_v>) { + file_handle.Read(&meta.global_cache_, sizeof(GlobalCacheType)); + } return meta; } 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 index 808dc319b4..cbede4b320 100644 --- 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 @@ -156,8 +156,8 @@ protected: PlainVecStoreInner(const DataType *ptr) { this->ptr_ = ptr; } public: - static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta) { - assert(cur_vec_num <= max_vec_num); + PlainVecStoreInner() = default; + static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, const Meta &meta) { const auto *p = reinterpret_cast(ptr); // fixme ptr += sizeof(DataType) * cur_vec_num * meta.dim(); return This(p); diff --git a/src/storage/knn_index/knn_hnsw/dist_func_sparse_ip.cppm b/src/storage/knn_index/knn_hnsw/dist_func_sparse_ip.cppm index df4af4ef84..f02ce27e98 100644 --- a/src/storage/knn_index/knn_hnsw/dist_func_sparse_ip.cppm +++ b/src/storage/knn_index/knn_hnsw/dist_func_sparse_ip.cppm @@ -31,6 +31,7 @@ public: using DistanceType = typename VecStoreMeta::DistanceType; public: + SparseIPDist() = default; SparseIPDist(SizeT dim) {} DataType operator()(const SparseVecRef &v1, const SparseVecRef &v2, const VecStoreMeta &vec_store_meta) const { diff --git a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm index bb200b037d..80f649a619 100644 --- a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm +++ b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm @@ -28,6 +28,7 @@ import logical_type; import hnsw_common; import data_store; import third_party; +import serialize; // Fixme: some variable has implicit type conversion. // Fixme: some variable has confusing name. @@ -43,15 +44,15 @@ export struct KnnSearchOption { LogicalType column_logical_type_ = LogicalType::kEmbedding; }; -export template -class KnnHnsw { +export template +class KnnHnswBase { public: - using This = KnnHnsw; + using This = KnnHnswBase; using DataType = typename VecStoreType::DataType; using QueryVecType = typename VecStoreType::QueryVecType; using StoreType = typename VecStoreType::StoreType; using QueryType = typename VecStoreType::QueryType; - using DataStore = DataStore; + using DataStore = DataStore; using Distance = typename VecStoreType::Distance; using DistanceType = typename Distance::DistanceType; @@ -63,18 +64,11 @@ public: constexpr static int prefetch_offset_ = 0; constexpr static int prefetch_step_ = 2; - using CompressVecStoreType = decltype(VecStoreType::template ToLVQ()); - - // private: - KnnHnsw(SizeT M, SizeT ef_construction, DataStore data_store, Distance distance) - : M_(M), ef_construction_(std::max(M_, ef_construction)), mult_(1 / std::log(1.0 * M_)), data_store_(std::move(data_store)), - distance_(std::move(distance)) {} - static Pair GetMmax(SizeT M) { return {2 * M, M}; } public: - KnnHnsw() : M_(0), ef_construction_(0), mult_(0) {} - KnnHnsw(This &&other) + KnnHnswBase() : M_(0), ef_construction_(0), mult_(0) {} + KnnHnswBase(This &&other) : M_(std::exchange(other.M_, 0)), ef_construction_(std::exchange(other.ef_construction_, 0)), mult_(std::exchange(other.mult_, 0.0)), data_store_(std::move(other.data_store_)), distance_(std::move(other.distance_)) {} This &operator=(This &&other) { @@ -87,14 +81,6 @@ public: } return *this; } - ~KnnHnsw() = default; - - static UniquePtr Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT M, SizeT ef_construction) { - auto [Mmax0, Mmax] = This::GetMmax(M); - auto data_store = DataStore::Make(chunk_size, max_chunk_n, dim, Mmax0, Mmax); - Distance distance(data_store.dim()); - return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); - } SizeT GetSizeInBytes() const { return sizeof(M_) + sizeof(ef_construction_) + data_store_.GetSizeInBytes(); } @@ -104,19 +90,7 @@ public: data_store_.Save(file_handle); } - static UniquePtr Load(LocalFileHandle &file_handle) { - SizeT M; - file_handle.Read(&M, sizeof(M)); - SizeT ef_construction; - file_handle.Read(&ef_construction, sizeof(ef_construction)); - - auto data_store = DataStore::Load(file_handle); - Distance distance(data_store.dim()); - - return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); - } - -private: +protected: // >= 0 i32 GenerateRandomLayer() { static thread_local std::mt19937 generator; @@ -395,20 +369,6 @@ public: } } - UniquePtr> CompressToLVQ() && { - if constexpr (std::is_same_v) { - return MakeUnique(std::move(*this)); - } else { - using CompressedDistance = typename CompressVecStoreType::Distance; - CompressedDistance distance = std::move(distance_).ToLVQDistance(data_store_.dim()); - auto compressed_datastore = std::move(data_store_).template CompressToLVQ(); - return MakeUnique>(M_, - ef_construction_, - std::move(compressed_datastore), - std::move(distance)); - } - } - template Filter = NoneType, bool WithLock = true> Tuple, UniquePtr> KnnSearch(const QueryVecType &q, SizeT k, const Filter &filter, const KnnSearchOption &option = {}) const { @@ -459,7 +419,7 @@ public: SizeT mem_usage() const { return data_store_.mem_usage(); } -private: +protected: SizeT M_; SizeT ef_construction_; @@ -481,4 +441,84 @@ public: } }; +export template +class KnnHnsw : public KnnHnswBase { +public: + using This = KnnHnsw; + using DataStore = DataStore; + using Distance = typename VecStoreType::Distance; + using CompressVecStoreType = decltype(VecStoreType::template ToLVQ()); + + KnnHnsw(SizeT M, SizeT ef_construction, DataStore data_store, Distance distance) { + this->M_ = M; + this->ef_construction_ = std::max(M, ef_construction); + this->mult_ = 1 / std::log(1.0 * M); + this->data_store_ = std::move(data_store); + this->distance_ = std::move(distance); + } + +public: + static UniquePtr Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT M, SizeT ef_construction) { + auto [Mmax0, Mmax] = This::GetMmax(M); + auto data_store = DataStore::Make(chunk_size, max_chunk_n, dim, Mmax0, Mmax); + Distance distance(data_store.dim()); + return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); + } + + static UniquePtr Load(LocalFileHandle &file_handle) { + SizeT M; + file_handle.Read(&M, sizeof(M)); + SizeT ef_construction; + file_handle.Read(&ef_construction, sizeof(ef_construction)); + + auto data_store = DataStore::Load(file_handle); + Distance distance(data_store.dim()); + + return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); + } + + UniquePtr> CompressToLVQ() && { + if constexpr (std::is_same_v) { + return MakeUnique(std::move(*this)); + } else { + using CompressedDistance = typename CompressVecStoreType::Distance; + CompressedDistance distance = std::move(this->distance_).ToLVQDistance(this->data_store_.dim()); + auto compressed_datastore = std::move(this->data_store_).template CompressToLVQ(); + return MakeUnique>(this->M_, + this->ef_construction_, + std::move(compressed_datastore), + std::move(distance)); + } + } +}; + +export template +class KnnHnsw : public KnnHnswBase { +public: + using This = KnnHnsw; + using DataStore = DataStore; + using Distance = typename VecStoreType::Distance; + + KnnHnsw(SizeT M, SizeT ef_construction, DataStore data_store, Distance distance) { + this->M_ = M; + this->ef_construction_ = std::max(M, ef_construction); + this->mult_ = 1 / std::log(1.0 * M); + this->data_store_ = std::move(data_store); + this->distance_ = std::move(distance); + } + +public: + static UniquePtr LoadFromPtr(const char *&ptr, SizeT size) { + const char *ptr_end = ptr + size; + SizeT M = ReadBufAdv(ptr); + SizeT ef_construction = ReadBufAdv(ptr); + auto data_store = DataStore::LoadFromPtr(ptr); + Distance distance(data_store.dim()); + if (SizeT diff = ptr_end - ptr; diff != 0) { + UnrecoverableError("LoadFromPtr failed"); + } + return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); + } +}; + } // namespace infinity \ No newline at end of file diff --git a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp index bb54c9b14b..7ce3a30531 100644 --- a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp +++ b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp @@ -81,6 +81,7 @@ class HnswAlgTest : public BaseTest { EXPECT_GE(correct_rate, 0.95); }; + String filepath = save_dir_ + "/test_hnsw.bin"; { auto hnsw_index = Hnsw::Make(chunk_size, max_chunk_n, dim, M, ef_construction); auto iter = DenseVectorIter(data.get(), dim, element_size); @@ -88,7 +89,7 @@ class HnswAlgTest : public BaseTest { test_func(hnsw_index); - auto [file_handle, status] = VirtualStore::Open(save_dir_ + "/test_hnsw.bin", FileAccessMode::kWrite); + auto [file_handle, status] = VirtualStore::Open(filepath, FileAccessMode::kWrite); if (!status.ok()) { UnrecoverableError(status.message()); } @@ -96,7 +97,7 @@ class HnswAlgTest : public BaseTest { } { - auto [file_handle, status] = VirtualStore::Open(save_dir_ + "/test_hnsw.bin", FileAccessMode::kRead); + auto [file_handle, status] = VirtualStore::Open(filepath, FileAccessMode::kRead); if (!status.ok()) { UnrecoverableError(status.message()); } @@ -107,6 +108,82 @@ class HnswAlgTest : public BaseTest { } } + template + void TestLoad() { + int dim = 16; + int M = 8; + int ef_construction = 200; + int chunk_size = 16; + int max_chunk_n = 10; + int element_size = max_chunk_n * chunk_size; + + std::mt19937 rng; + rng.seed(0); + std::uniform_real_distribution distrib_real; + + auto data = MakeUnique(dim * element_size); + for (int i = 0; i < dim * element_size; ++i) { + data[i] = distrib_real(rng); + } + + auto test_func = [&](auto &hnsw_index) { + hnsw_index->Check(); + + KnnSearchOption search_option{.ef_ = 10}; + int correct = 0; + for (int i = 0; i < element_size; ++i) { + const float *query = data.get() + i * dim; + auto result = hnsw_index->KnnSearchSorted(query, 1, search_option); + if (result[0].second == (LabelT)i) { + ++correct; + } + } + float correct_rate = float(correct) / element_size; + // std::printf("correct rage: %f\n", correct_rate); + EXPECT_GE(correct_rate, 0.95); + }; + + String filepath = save_dir_ + "/test_hnsw.bin"; + { + auto hnsw_index = Hnsw::Make(chunk_size, max_chunk_n, dim, M, ef_construction); + auto iter = DenseVectorIter(data.get(), dim, element_size); + hnsw_index->InsertVecs(std::move(iter)); + + auto [file_handle, status] = VirtualStore::Open(filepath, FileAccessMode::kWrite); + if (!status.ok()) { + UnrecoverableError(status.message()); + } + hnsw_index->Save(*file_handle); + } + { + 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(data_ptr); +#else + auto [file_handle, status] = VirtualStore::Open(filepath, FileAccessMode::kRead); + if (!status.ok()) { + UnrecoverableError(status.message()); + } + auto buffer = MakeUnique(file_size); + file_handle->Read(buffer.get(), file_size); + const char *ptr = buffer.get(); +#endif + auto hnsw_index = LoadHnsw::LoadFromPtr(ptr, file_size); + + test_func(hnsw_index); + +#ifdef USE_MMAP + VirtualStore::MunmapFile(filepath); +#endif + } + } + template void TestCompress() { int dim = 16; @@ -318,3 +395,15 @@ TEST_F(HnswAlgTest, test6) { using CompressedHnsw = KnnHnsw, LabelT>; TestCompress(); } + +TEST_F(HnswAlgTest, test7) { + using Hnsw = KnnHnsw, LabelT>; + using HnswLoad = KnnHnsw, LabelT, false>; + TestLoad(); +} + +TEST_F(HnswAlgTest, test8) { + using Hnsw = KnnHnsw, LabelT>; + using HnswLoad = KnnHnsw, LabelT, false>; + TestLoad(); +} From da100b90f49b2501f7b4eda159121e56ecafb6ea Mon Sep 17 00:00:00 2001 From: shenyushi Date: Tue, 24 Dec 2024 20:36:42 +0800 Subject: [PATCH 03/12] Fix hnsw save to const. --- .../knn_index/knn_hnsw/abstract_hnsw.cpp | 137 +++++++++--------- .../knn_index/knn_hnsw/abstract_hnsw.cppm | 102 +++++++++---- .../knn_hnsw/data_store/data_store.cppm | 44 +++--- .../knn_hnsw/data_store/graph_store.cppm | 71 +++++---- src/storage/knn_index/knn_hnsw/hnsw_alg.cppm | 4 +- 5 files changed, 200 insertions(+), 158 deletions(-) diff --git a/src/storage/knn_index/knn_hnsw/abstract_hnsw.cpp b/src/storage/knn_index/knn_hnsw/abstract_hnsw.cpp index 5a6f4d2882..ff6cc6ef6b 100644 --- a/src/storage/knn_index/knn_hnsw/abstract_hnsw.cpp +++ b/src/storage/knn_index/knn_hnsw/abstract_hnsw.cpp @@ -77,32 +77,16 @@ HnswIndexInMem::HnswIndexInMem(RowID begin_row_id, using T = std::decay_t; if constexpr (!std::is_same_v) { using IndexT = std::decay_t; - hnsw_ = IndexT::Make(chunk_size, max_chunk_num, dim, M, ef_construction).release(); + if constexpr (IndexT::kOwnMem) { + index = IndexT::Make(chunk_size, max_chunk_num, dim, M, ef_construction).release(); + } else { + UnrecoverableError("HnswIndexInMem::HnswIndexInMem: index does not own memory"); + } } }, hnsw_); } -AbstractHnsw HnswIndexInMem::InitAbstractIndex(const IndexBase *index_base, const ColumnDef *column_def) { - const auto *index_hnsw = static_cast(index_base); - const auto *embedding_info = static_cast(column_def->type()->type_info().get()); - - switch (embedding_info->Type()) { - case EmbeddingDataType::kElemFloat: { - return InitAbstractIndex(index_hnsw); - } - case EmbeddingDataType::kElemUInt8: { - return InitAbstractIndex(index_hnsw); - } - case EmbeddingDataType::kElemInt8: { - return InitAbstractIndex(index_hnsw); - } - default: { - return nullptr; - } - } -} - HnswIndexInMem::~HnswIndexInMem() { SizeT mem_usage = 0; std::visit( @@ -146,27 +130,33 @@ void HnswIndexInMem::InsertVecs(SizeT block_offset, std::visit( [&](auto &&index) { using T = std::decay_t; - if constexpr (!std::is_same_v) { + if constexpr (std::is_same_v) { + return; + } else { using IndexT = std::decay_t; - using DataType = typename IndexT::DataType; - SizeT mem_usage{}; - switch (const auto &column_data_type = block_column_entry->column_type(); column_data_type->type()) { - case LogicalType::kEmbedding: { - MemIndexInserterIter iter(block_offset, block_column_entry, buffer_manager, row_offset, row_count); - InsertVecs(index, std::move(iter), config, mem_usage); - break; - } - case LogicalType::kMultiVector: { - MemIndexInserterIter> iter(block_offset, block_column_entry, buffer_manager, row_offset, row_count); - InsertVecs(index, std::move(iter), config, mem_usage); - break; - } - default: { - UnrecoverableError(fmt::format("Unsupported column type for HNSW index: {}", column_data_type->ToString())); - break; + if constexpr (IndexT::kOwnMem) { + using DataType = typename IndexT::DataType; + SizeT mem_usage{}; + switch (const auto &column_data_type = block_column_entry->column_type(); column_data_type->type()) { + case LogicalType::kEmbedding: { + MemIndexInserterIter iter(block_offset, block_column_entry, buffer_manager, row_offset, row_count); + InsertVecs(index, std::move(iter), config, mem_usage); + break; + } + case LogicalType::kMultiVector: { + MemIndexInserterIter> iter(block_offset, block_column_entry, buffer_manager, row_offset, row_count); + InsertVecs(index, std::move(iter), config, mem_usage); + break; + } + default: { + UnrecoverableError(fmt::format("Unsupported column type for HNSW index: {}", column_data_type->ToString())); + break; + } } + this->IncreaseMemoryUsageBase(mem_usage); + } else { + UnrecoverableError("HnswIndexInMem::InsertVecs: index does not own memory"); } - this->IncreaseMemoryUsageBase(mem_usage); } }, hnsw_); @@ -183,38 +173,42 @@ void HnswIndexInMem::InsertVecs(const SegmentEntry *segment_entry, using T = std::decay_t; if constexpr (!std::is_same_v) { using IndexT = std::decay_t; - using DataType = typename IndexT::DataType; + if constexpr (!IndexT::kOwnMem) { + UnrecoverableError("HnswIndexInMem::InsertVecs: index does not own memory"); + } else { + using DataType = typename IndexT::DataType; - SizeT mem_usage{}; - switch (const auto &column_data_type = segment_entry->GetTableEntry()->GetColumnDefByID(column_id)->type(); - column_data_type->type()) { - case LogicalType::kEmbedding: { - if (check_ts) { - OneColumnIterator iter(segment_entry, buffer_mgr, column_id, begin_ts); - InsertVecs(index, std::move(iter), config, mem_usage); - } else { - OneColumnIterator iter(segment_entry, buffer_mgr, column_id, begin_ts); - InsertVecs(index, std::move(iter), config, mem_usage); + SizeT mem_usage{}; + switch (const auto &column_data_type = segment_entry->GetTableEntry()->GetColumnDefByID(column_id)->type(); + column_data_type->type()) { + case LogicalType::kEmbedding: { + if (check_ts) { + OneColumnIterator iter(segment_entry, buffer_mgr, column_id, begin_ts); + InsertVecs(index, std::move(iter), config, mem_usage); + } else { + OneColumnIterator iter(segment_entry, buffer_mgr, column_id, begin_ts); + InsertVecs(index, std::move(iter), config, mem_usage); + } + break; } - break; - } - case LogicalType::kMultiVector: { - const auto ele_size = column_data_type->type_info()->Size(); - if (check_ts) { - OneColumnIterator> iter(segment_entry, buffer_mgr, column_id, begin_ts, ele_size); - InsertVecs(index, std::move(iter), config, mem_usage); - } else { - OneColumnIterator, false> iter(segment_entry, buffer_mgr, column_id, begin_ts, ele_size); - InsertVecs(index, std::move(iter), config, mem_usage); + case LogicalType::kMultiVector: { + const auto ele_size = column_data_type->type_info()->Size(); + if (check_ts) { + OneColumnIterator> iter(segment_entry, buffer_mgr, column_id, begin_ts, ele_size); + InsertVecs(index, std::move(iter), config, mem_usage); + } else { + OneColumnIterator, false> iter(segment_entry, buffer_mgr, column_id, begin_ts, ele_size); + InsertVecs(index, std::move(iter), config, mem_usage); + } + break; + } + default: { + UnrecoverableError(fmt::format("Unsupported column type for HNSW index: {}", column_data_type->ToString())); + break; } - break; - } - default: { - UnrecoverableError(fmt::format("Unsupported column type for HNSW index: {}", column_data_type->ToString())); - break; } + this->IncreaseMemoryUsageBase(mem_usage); } - this->IncreaseMemoryUsageBase(mem_usage); } }, hnsw_); @@ -231,9 +225,14 @@ SharedPtr HnswIndexInMem::Dump(SegmentIndexEntry *segment_index if constexpr (std::is_same_v) { return; } else { - row_count = index->GetVecNum(); - index_size = index->GetSizeInBytes(); - dump_size = index->mem_usage(); + using IndexT = typename std::remove_pointer_t; + if constexpr (IndexT::kOwnMem) { + row_count = index->GetVecNum(); + index_size = index->GetSizeInBytes(); + dump_size = index->mem_usage(); + } else { + UnrecoverableError("HnswIndexInMem::Dump: index does not own memory"); + } } }, hnsw_); diff --git a/src/storage/knn_index/knn_hnsw/abstract_hnsw.cppm b/src/storage/knn_index/knn_hnsw/abstract_hnsw.cppm index dd2323501e..e069428de0 100644 --- a/src/storage/knn_index/knn_hnsw/abstract_hnsw.cppm +++ b/src/storage/knn_index/knn_hnsw/abstract_hnsw.cppm @@ -62,8 +62,20 @@ export using AbstractHnsw = std::variant, Se KnnHnsw, SegmentOffset> *, KnnHnsw, SegmentOffset> *, KnnHnsw, SegmentOffset> *, - std::nullptr_t>; + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + KnnHnsw, SegmentOffset, false> *, + std::nullptr_t>; export struct HnswIndexInMem : public BaseMemIndex { public: HnswIndexInMem() : hnsw_(nullptr) {} @@ -74,21 +86,21 @@ public: HnswIndexInMem(RowID begin_row_id, const IndexBase *index_base, const ColumnDef *column_def, SegmentIndexEntry *segment_index_entry, bool trace); private: - template + template static AbstractHnsw InitAbstractIndex(const IndexHnsw *index_hnsw) { switch (index_hnsw->encode_type_) { case HnswEncodeType::kPlain: { switch (index_hnsw->metric_type_) { case MetricType::kMetricL2: { - using HnswIndex = KnnHnsw, SegmentOffset>; + using HnswIndex = KnnHnsw, SegmentOffset, OwnMem>; return static_cast(nullptr); } case MetricType::kMetricInnerProduct: { - using HnswIndex = KnnHnsw, SegmentOffset>; + using HnswIndex = KnnHnsw, SegmentOffset, OwnMem>; return static_cast(nullptr); } case MetricType::kMetricCosine: { - using HnswIndex = KnnHnsw, SegmentOffset>; + using HnswIndex = KnnHnsw, SegmentOffset, OwnMem>; return static_cast(nullptr); } default: { @@ -102,15 +114,15 @@ private: } else { switch (index_hnsw->metric_type_) { case MetricType::kMetricL2: { - using HnswIndex = KnnHnsw, SegmentOffset>; + using HnswIndex = KnnHnsw, SegmentOffset, OwnMem>; return static_cast(nullptr); } case MetricType::kMetricInnerProduct: { - using HnswIndex = KnnHnsw, SegmentOffset>; + using HnswIndex = KnnHnsw, SegmentOffset, OwnMem>; return static_cast(nullptr); } case MetricType::kMetricCosine: { - using HnswIndex = KnnHnsw, SegmentOffset>; + using HnswIndex = KnnHnsw, SegmentOffset, OwnMem>; return static_cast(nullptr); } default: { @@ -125,6 +137,27 @@ private: } } + template + static AbstractHnsw InitAbstractIndex(const IndexBase *index_base, const ColumnDef *column_def) { + const auto *index_hnsw = static_cast(index_base); + const auto *embedding_info = static_cast(column_def->type()->type_info().get()); + + switch (embedding_info->Type()) { + case EmbeddingDataType::kElemFloat: { + return InitAbstractIndex(index_hnsw); + } + case EmbeddingDataType::kElemUInt8: { + return InitAbstractIndex(index_hnsw); + } + case EmbeddingDataType::kElemInt8: { + return InitAbstractIndex(index_hnsw); + } + default: { + return nullptr; + } + } + } + template static void InsertVecs(Index &index, Iter &&iter, const HnswInsertConfig &config, SizeT &mem_usage) { auto &thread_pool = InfinityContext::instance().GetHnswBuildThreadPool(); @@ -133,32 +166,43 @@ private: } using T = std::decay_t; if constexpr (!std::is_same_v) { - SizeT mem1 = index->mem_usage(); - auto [start, end] = index->StoreData(std::forward(iter), config); - SizeT bucket_size = std::max(kBuildBucketSize, SizeT(end - start - 1) / thread_pool.size() + 1); - SizeT bucket_n = (end - start - 1) / bucket_size + 1; - - Vector> futs; - futs.reserve(bucket_n); - for (SizeT i = 0; i < bucket_n; ++i) { - SizeT i1 = start + i * bucket_size; - SizeT i2 = std::min(i1 + bucket_size, SizeT(end)); - futs.emplace_back(thread_pool.push([&index, i1, i2](int id) { - for (SizeT j = i1; j < i2; ++j) { - index->Build(j); - } - })); - } - for (auto &fut : futs) { - fut.get(); + using IndexT = std::decay_t; + if constexpr (!IndexT::kOwnMem) { + UnrecoverableError("HnswIndexInMem::InsertVecs: index does not own memory"); + } else { + SizeT mem1 = index->mem_usage(); + auto [start, end] = index->StoreData(std::forward(iter), config); + SizeT bucket_size = std::max(kBuildBucketSize, SizeT(end - start - 1) / thread_pool.size() + 1); + SizeT bucket_n = (end - start - 1) / bucket_size + 1; + + Vector> futs; + futs.reserve(bucket_n); + for (SizeT i = 0; i < bucket_n; ++i) { + SizeT i1 = start + i * bucket_size; + SizeT i2 = std::min(i1 + bucket_size, SizeT(end)); + futs.emplace_back(thread_pool.push([&index, i1, i2](int id) { + for (SizeT j = i1; j < i2; ++j) { + index->Build(j); + } + })); + } + for (auto &fut : futs) { + fut.get(); + } + SizeT mem2 = index->mem_usage(); + mem_usage = mem2 - mem1; } - SizeT mem2 = index->mem_usage(); - mem_usage = mem2 - mem1; } } public: - static AbstractHnsw InitAbstractIndex(const IndexBase *index_base, const ColumnDef *column_def); + static AbstractHnsw InitAbstractIndex(const IndexBase *index_base, const ColumnDef *column_def, bool own_mem = true) { + if (own_mem) { + return InitAbstractIndex(index_base, column_def); + } else { + return InitAbstractIndex(index_base, column_def); + } + } HnswIndexInMem(const HnswIndexInMem &) = delete; HnswIndexInMem &operator=(const HnswIndexInMem &) = delete; 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 index cff059c810..1e0b78f929 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm @@ -93,6 +93,22 @@ public: } } + void Save(LocalFileHandle &file_handle) const { + SizeT cur_vec_num = this->cur_vec_num(); + auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); + + file_handle.Append(&this->chunk_size_, sizeof(this->chunk_size_)); + file_handle.Append(&this->max_chunk_n_, sizeof(this->max_chunk_n_)); + + file_handle.Append(&cur_vec_num, sizeof(cur_vec_num)); + this->vec_store_meta_.Save(file_handle); + this->graph_store_meta_.Save(file_handle, cur_vec_num); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? this->chunk_size_ : last_chunk_size; + this->inners_[i].Save(file_handle, chunk_size, this->vec_store_meta_, this->graph_store_meta_); + } + } + void SetGraph(GraphStoreMeta &&graph_meta, Vector> &&graph_inners) { graph_store_meta_ = std::move(graph_meta); for (SizeT i = 0; i < graph_inners.size(); ++i) { @@ -274,22 +290,6 @@ public: return ret; } - void Save(LocalFileHandle &file_handle) { - SizeT cur_vec_num = this->cur_vec_num(); - auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); - - file_handle.Append(&this->chunk_size_, sizeof(this->chunk_size_)); - file_handle.Append(&this->max_chunk_n_, sizeof(this->max_chunk_n_)); - - file_handle.Append(&cur_vec_num, sizeof(cur_vec_num)); - this->vec_store_meta_.Save(file_handle); - this->graph_store_meta_.Save(file_handle, cur_vec_num); - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT chunk_size = (i < chunk_num - 1) ? this->chunk_size_ : last_chunk_size; - this->inners_[i].Save(file_handle, chunk_size, this->vec_store_meta_, this->graph_store_meta_); - } - } - static This Load(LocalFileHandle &file_handle, SizeT max_chunk_n = 0) { SizeT chunk_size; file_handle.Read(&chunk_size, sizeof(chunk_size)); @@ -452,6 +452,12 @@ public: public: DataStoreInnerBase() = default; + void Save(LocalFileHandle &file_handle, SizeT cur_vec_num, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) const { + this->vec_store_inner_.Save(file_handle, cur_vec_num, vec_store_meta); + this->graph_store_inner_.Save(file_handle, cur_vec_num, graph_store_meta); + file_handle.Append(this->labels_.get(), sizeof(LabelType) * cur_vec_num); + } + void Free(SizeT cur_vec_num, const GraphStoreMeta &graph_store_meta) { graph_store_inner_.Free(cur_vec_num, graph_store_meta); } SizeT GetSizeInBytes(SizeT chunk_size, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) const { @@ -530,12 +536,6 @@ public: return This(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner)); } - void Save(LocalFileHandle &file_handle, SizeT cur_vec_num, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) { - this->vec_store_inner_.Save(file_handle, cur_vec_num, vec_store_meta); - this->graph_store_inner_.Save(file_handle, cur_vec_num, graph_store_meta); - file_handle.Append(this->labels_.get(), sizeof(LabelType) * cur_vec_num); - } - static This Load(LocalFileHandle &file_handle, SizeT cur_vec_num, SizeT chunk_size, 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 index 183735dc8f..c3cfd1b112 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm @@ -71,11 +71,10 @@ public: SizeT GetSizeInBytes() const { return sizeof(Mmax0_) + sizeof(Mmax_) + sizeof(max_layer_) + sizeof(enterpoint_); } - void Save(LocalFileHandle &file_handle, SizeT cur_vec_num) { + void Save(LocalFileHandle &file_handle, SizeT cur_vec_num) const { file_handle.Append(&Mmax0_, sizeof(Mmax0_)); file_handle.Append(&Mmax_, sizeof(Mmax_)); - last_save_num_ = cur_vec_num; - file_handle.Append(&last_save_num_, sizeof(last_save_num_)); + file_handle.Append(&cur_vec_num, sizeof(last_save_num_)); file_handle.Append(&max_layer_, sizeof(max_layer_)); file_handle.Append(&enterpoint_, sizeof(enterpoint_)); @@ -162,6 +161,38 @@ class GraphStoreInnerBase { public: GraphStoreInnerBase() = default; + void Save(LocalFileHandle &file_handle, SizeT cur_vertex_n, const GraphStoreMeta &meta) const { + SizeT layer_sum = 0; + Vector> layers_ptrs_off; + for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { + const VertexL0 *v = GetLevel0(vertex_i, meta); + if (!v->layer_n_) { + continue; + } + SizeT offset = layer_sum * meta.levelx_size(); + SizeT off = reinterpret_cast(&v->layers_p_) - this->graph_.get(); + layers_ptrs_off.emplace_back(off, offset); + layer_sum += v->layer_n_; + } + file_handle.Append(&layer_sum, sizeof(layer_sum)); + { + auto buffer = MakeUnique(cur_vertex_n * meta.level0_size()); + std::copy(this->graph_.get(), this->graph_.get() + cur_vertex_n * meta.level0_size(), buffer.get()); + for (auto [off, offset] : layers_ptrs_off) { + char *ptr = buffer.get() + off; + *reinterpret_cast(ptr) = offset; + } + file_handle.Append(buffer.get(), cur_vertex_n * meta.level0_size()); + } + for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { + const VertexL0 *v = GetLevel0(vertex_i, meta); + if (v->layer_n_) { + char *ptr = v->layers_p_; + file_handle.Append(ptr, meta.levelx_size() * v->layer_n_); + } + } + } + void Free(SizeT current_vertex_num, const GraphStoreMeta &meta) { for (VertexType vertex_i = loaded_vertex_n_; vertex_i < VertexType(current_vertex_num); ++vertex_i) { delete[] GetLevel0(vertex_i, meta)->layers_p_; @@ -306,32 +337,6 @@ public: return graph_store; } - void Save(LocalFileHandle &file_handle, SizeT cur_vertex_n, const GraphStoreMeta &meta) { - SizeT layer_sum = 0; - Vector layers_ptrs; - for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { - VertexL0 *v = GetLevel0(vertex_i, meta); - if (!v->layer_n_) { - continue; - } - layers_ptrs.emplace_back(v->layers_p_); - SizeT offset = layer_sum * meta.levelx_size(); - v->layers_p_ = reinterpret_cast(offset); - layer_sum += v->layer_n_; - } - file_handle.Append(&layer_sum, sizeof(layer_sum)); - file_handle.Append(this->graph_.get(), cur_vertex_n * meta.level0_size()); - for (VertexType vertex_i = 0, ptr_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { - VertexL0 *v = GetLevel0(vertex_i, meta); - if (v->layer_n_) { - char *ptr = layers_ptrs[ptr_i++]; - file_handle.Append(ptr, meta.levelx_size() * v->layer_n_); - } - } - layers_ptrs_.insert(layers_ptrs_.end(), layers_ptrs.begin(), layers_ptrs.end()); - this->loaded_vertex_n_ = cur_vertex_n; - } - static GraphStoreInner Load(LocalFileHandle &file_handle, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { assert(cur_vertex_n <= max_vertex); @@ -351,14 +356,6 @@ public: return graph_store; } - void Free(SizeT current_vertex_num, const GraphStoreMeta &meta) { - for (const char *layer_p : layers_ptrs_) { - delete[] layer_p; - } - layers_ptrs_.clear(); - Base::Free(current_vertex_num, meta); - } - 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/hnsw_alg.cppm b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm index 80f649a619..6244c7a592 100644 --- a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm +++ b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm @@ -84,7 +84,7 @@ public: SizeT GetSizeInBytes() const { return sizeof(M_) + sizeof(ef_construction_) + data_store_.GetSizeInBytes(); } - void Save(LocalFileHandle &file_handle) { + void Save(LocalFileHandle &file_handle) const { file_handle.Append(&M_, sizeof(M_)); file_handle.Append(&ef_construction_, sizeof(ef_construction_)); data_store_.Save(file_handle); @@ -448,6 +448,7 @@ public: using DataStore = DataStore; using Distance = typename VecStoreType::Distance; using CompressVecStoreType = decltype(VecStoreType::template ToLVQ()); + constexpr static bool kOwnMem = OwnMem; KnnHnsw(SizeT M, SizeT ef_construction, DataStore data_store, Distance distance) { this->M_ = M; @@ -498,6 +499,7 @@ public: using This = KnnHnsw; using DataStore = DataStore; using Distance = typename VecStoreType::Distance; + constexpr static bool kOwnMem = false; KnnHnsw(SizeT M, SizeT ef_construction, DataStore data_store, Distance distance) { this->M_ = M; From f84bf231089eda8209b81fe3e08a13086ac69f4e Mon Sep 17 00:00:00 2001 From: shenyushi Date: Tue, 24 Dec 2024 20:36:55 +0800 Subject: [PATCH 04/12] Add mmap to hnsw. --- src/storage/buffer/buffer_obj.cpp | 9 +++- src/storage/buffer/buffer_obj.cppm | 2 +- .../buffer/file_worker/hnsw_file_worker.cpp | 26 ++++++++--- src/storage/meta/entry/chunk_index_entry.cpp | 10 ++++- .../meta/entry/segment_index_entry.cpp | 43 +++++++++++-------- 5 files changed, 64 insertions(+), 26 deletions(-) diff --git a/src/storage/buffer/buffer_obj.cpp b/src/storage/buffer/buffer_obj.cpp index a97372e94d..2e9971eddb 100644 --- a/src/storage/buffer/buffer_obj.cpp +++ b/src/storage/buffer/buffer_obj.cpp @@ -67,9 +67,16 @@ void BufferObj::UpdateFileWorkerInfo(UniquePtr new_file_worker) { } } -BufferHandle BufferObj::Load() { +BufferHandle BufferObj::Load(bool no_mmap) { buffer_mgr_->AddRequestCount(); std::unique_lock locker(w_locker_); + if (type_ == BufferType::kMmap && no_mmap) { + if (rc_ > 0) { + String error_message = fmt::format("Buffer {} is mmaped, but has {} references", GetFilename(), rc_); + UnrecoverableError(error_message); + } + type_ = BufferType::kPersistent; + } if (type_ == BufferType::kMmap) { switch (status_) { case BufferStatus::kLoaded: { diff --git a/src/storage/buffer/buffer_obj.cppm b/src/storage/buffer/buffer_obj.cppm index 9ee6db9126..56a9144329 100644 --- a/src/storage/buffer/buffer_obj.cppm +++ b/src/storage/buffer/buffer_obj.cppm @@ -94,7 +94,7 @@ public: public: // called by ObjectHandle when load first time for that ObjectHandle - BufferHandle Load(); + BufferHandle Load(bool no_mmap = false); // called by BufferMgr in GC process. bool Free(); diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index 4970359cc9..8edb327f44 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -45,7 +45,7 @@ HnswFileWorker::HnswFileWorker(SharedPtr data_dir, SharedPtr file_name, SharedPtr index_base, SharedPtr column_def, - PersistenceManager* persistence_manager, + PersistenceManager *persistence_manager, SizeT index_size) : IndexFileWorker(std::move(data_dir), std::move(temp_dir), @@ -113,7 +113,12 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const if constexpr (std::is_same_v) { UnrecoverableError("Invalid index type."); } else { - index->Save(*file_handle_); + using IndexT = std::decay_t; + if constexpr (IndexT::kOwnMem) { + index->Save(*file_handle_); + } else { + UnrecoverableError("Invalid index type."); + } } }, *hnsw_index); @@ -134,7 +139,11 @@ void HnswFileWorker::ReadFromFileImpl(SizeT file_size) { UnrecoverableError("Invalid index type."); } else { using IndexT = std::decay_t; - index = IndexT::Load(*file_handle_).release(); + if constexpr (IndexT::kOwnMem) { + index = IndexT::Load(*file_handle_).release(); + } else { + UnrecoverableError("Invalid index type."); + } } }, *hnsw_index); @@ -144,7 +153,7 @@ bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) { if (mmap_data_ != nullptr) { UnrecoverableError("Mmap data is already allocated."); } - mmap_data_ = reinterpret_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get()))); + mmap_data_ = reinterpret_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get(), false))); auto *hnsw_index = reinterpret_cast(mmap_data_); std::visit( [&](auto &&index) { @@ -152,8 +161,13 @@ bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) { if constexpr (std::is_same_v) { UnrecoverableError("Invalid index type."); } else { - // using IndexT = std::decay_t; - // index = IndexT::LoadFromPtr(static_cast(ptr), size).release(); + using IndexT = std::decay_t; + if constexpr (!IndexT::kOwnMem) { + const auto *p = static_cast(ptr); + index = IndexT::LoadFromPtr(p, size).release(); + } else { + UnrecoverableError("Invalid index type."); + } } }, *hnsw_index); diff --git a/src/storage/meta/entry/chunk_index_entry.cpp b/src/storage/meta/entry/chunk_index_entry.cpp index 6c12c170e5..bd62a3dfc6 100644 --- a/src/storage/meta/entry/chunk_index_entry.cpp +++ b/src/storage/meta/entry/chunk_index_entry.cpp @@ -275,7 +275,9 @@ SharedPtr ChunkIndexEntry::NewReplayChunkIndexEntry(ChunkID chu index_base, column_def, buffer_mgr->persistence_manager()); - chunk_index_entry->buffer_obj_ = buffer_mgr->GetBufferObject(std::move(file_worker)); + BufferObj *buffer_obj = buffer_mgr->GetBufferObject(std::move(file_worker)); + buffer_obj->ToMmap(); + chunk_index_entry->buffer_obj_ = buffer_obj; break; } case IndexType::kFullText: { @@ -450,7 +452,13 @@ void ChunkIndexEntry::SaveIndexFile() { if (buffer_obj_ == nullptr) { return; } + if (buffer_obj_->type() == BufferType::kMmap) { + return; + } buffer_obj_->Save(); + if (segment_index_entry_->table_index_entry()->index_base()->index_type_ == IndexType::kHnsw) { + buffer_obj_->ToMmap(); + } } void ChunkIndexEntry::DeprecateChunk(TxnTimeStamp commit_ts) { diff --git a/src/storage/meta/entry/segment_index_entry.cpp b/src/storage/meta/entry/segment_index_entry.cpp index 1e5f2b48e9..803d393d0b 100644 --- a/src/storage/meta/entry/segment_index_entry.cpp +++ b/src/storage/meta/entry/segment_index_entry.cpp @@ -753,18 +753,23 @@ void SegmentIndexEntry::OptIndex(IndexBase *index_base, if constexpr (std::is_same_v) { UnrecoverableError("Invalid index type."); } else { - using HnswIndexDataType = typename std::remove_pointer_t::DataType; - if (params->compress_to_lvq) { - if constexpr (IsAnyOf) { - UnrecoverableError("Invalid index type."); - } else { - auto *p = std::move(*index).CompressToLVQ().release(); - delete index; - *abstract_hnsw = p; + using IndexT = typename std::remove_pointer_t; + if constexpr (IndexT::kOwnMem) { + using HnswIndexDataType = typename std::remove_pointer_t::DataType; + if (params->compress_to_lvq) { + if constexpr (IsAnyOf) { + UnrecoverableError("Invalid index type."); + } else { + auto *p = std::move(*index).CompressToLVQ().release(); + delete index; + *abstract_hnsw = p; + } } - } - if (params->lvq_avg) { - index->Optimize(); + if (params->lvq_avg) { + index->Optimize(); + } + } else { + UnrecoverableError("Invalid index type."); } } }, @@ -773,7 +778,7 @@ void SegmentIndexEntry::OptIndex(IndexBase *index_base, const auto [chunk_index_entries, memory_index_entry] = this->GetHnswIndexSnapshot(); for (const auto &chunk_index_entry : chunk_index_entries) { - BufferHandle buffer_handle = chunk_index_entry->GetIndex(); + BufferHandle buffer_handle = chunk_index_entry->GetBufferObj()->Load(true /*no mmap*/); auto *abstract_hnsw = reinterpret_cast(buffer_handle.GetDataMut()); optimize_index(abstract_hnsw); chunk_index_entry->SaveIndexFile(); @@ -913,11 +918,15 @@ ChunkIndexEntry *SegmentIndexEntry::RebuildChunkIndexEntries(TxnTableStore *txn_ return; } else { using IndexT = std::decay_t; - using DataType = typename IndexT::DataType; - CappedOneColumnIterator iter(segment_entry, buffer_mgr, column_def->id(), begin_ts, row_count); - HnswInsertConfig insert_config; - insert_config.optimize_ = true; - index->InsertVecs(std::move(iter), insert_config); + if constexpr (IndexT::kOwnMem) { + using DataType = typename IndexT::DataType; + CappedOneColumnIterator iter(segment_entry, buffer_mgr, column_def->id(), begin_ts, row_count); + HnswInsertConfig insert_config; + insert_config.optimize_ = true; + index->InsertVecs(std::move(iter), insert_config); + } else { + UnrecoverableError("Invalid index type."); + } } }, abstract_hnsw); From 2a5a20ea7cfda635db3bc6c426db6c23864e78d9 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Wed, 25 Dec 2024 10:01:53 +0800 Subject: [PATCH 05/12] Fix unit test. --- src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm | 2 ++ src/storage/knn_index/knn_hnsw/data_store/lvq_vec_store.cppm | 2 +- src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm b/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm index c6ee0d7867..e53b61525e 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store_util.cppm @@ -31,6 +31,8 @@ public: T *get() const { return ptr_.get(); } + UniquePtr exchange(UniquePtr ptr) { return std::exchange(ptr_, std::move(ptr)); } + private: UniquePtr ptr_; }; 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 index 4a090f2682..d820768cd0 100644 --- 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 @@ -258,7 +258,7 @@ public: for (SizeT i = 0; i < this->dim_; ++i) { new_mean[i] /= cur_vec_num; } - this->mean_ = std::move(new_mean); + new_mean = this->mean_.exchange(std::move(new_mean)); // for (auto [inner, size] : inners) { for (SizeT i = 0; i < size; ++i) { diff --git a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp index 7ce3a30531..4ef86209fc 100644 --- a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp +++ b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp @@ -113,7 +113,7 @@ class HnswAlgTest : public BaseTest { int dim = 16; int M = 8; int ef_construction = 200; - int chunk_size = 16; + int chunk_size = 128; int max_chunk_n = 10; int element_size = max_chunk_n * chunk_size; From ad2b530f7d07bdf0b5f95db6531530e15f50fb65 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Thu, 26 Dec 2024 15:50:47 +0800 Subject: [PATCH 06/12] Load hnsw from mmap. --- .../buffer/file_worker/file_worker.cpp | 7 +- .../buffer/file_worker/file_worker.cppm | 4 +- .../buffer/file_worker/hnsw_file_worker.cpp | 65 +-- .../buffer/file_worker/hnsw_file_worker.cppm | 5 +- src/storage/io/virtual_store.cpp | 2 +- .../knn_hnsw/data_store/data_store.cppm | 538 ++++++++++-------- .../knn_hnsw/data_store/graph_store.cppm | 145 +++-- .../knn_hnsw/data_store/lvq_vec_store.cppm | 14 +- .../knn_hnsw/data_store/plain_vec_store.cppm | 11 + .../knn_hnsw/data_store/sparse_vec_store.cppm | 4 +- src/storage/knn_index/knn_hnsw/hnsw_alg.cppm | 27 +- src/storage/meta/entry/chunk_index_entry.cpp | 4 - .../storage/knnindex/knn_hnsw/test_hnsw.cpp | 2 +- .../knn/embedding/test_knn_hnsw_l2_lvq.slt | 4 +- .../knn/embedding/test_knn_hnsw_l2_lvq2.slt | 4 +- 15 files changed, 476 insertions(+), 360 deletions(-) diff --git a/src/storage/buffer/file_worker/file_worker.cpp b/src/storage/buffer/file_worker/file_worker.cpp index 041a8fff7e..d3da5b1ddc 100644 --- a/src/storage/buffer/file_worker/file_worker.cpp +++ b/src/storage/buffer/file_worker/file_worker.cpp @@ -15,6 +15,8 @@ module; #include +#include +#include module file_worker; @@ -223,7 +225,10 @@ void FileWorker::Mmap() { auto [defer_fn, read_path] = GetFilePathInner(false); bool use_object_cache = persistence_manager_ != nullptr; SizeT file_size = VirtualStore::GetFileSize(read_path); - VirtualStore::MmapFile(read_path, mmap_addr_, file_size); + int ret = VirtualStore::MmapFile(read_path, mmap_addr_, file_size); + if (ret < 0) { + UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); + } if (use_object_cache) { const void *ptr = mmap_addr_ + obj_addr_.part_offset_; this->ReadFromMmapImpl(ptr, obj_addr_.part_size_); diff --git a/src/storage/buffer/file_worker/file_worker.cppm b/src/storage/buffer/file_worker/file_worker.cppm index 1929a92bd7..7823eb7254 100644 --- a/src/storage/buffer/file_worker/file_worker.cppm +++ b/src/storage/buffer/file_worker/file_worker.cppm @@ -79,11 +79,11 @@ protected: virtual void ReadFromFileImpl(SizeT file_size) = 0; + Pair>>, String> GetFilePathInner(bool spill); + private: String ChooseFileDir(bool spill) const; - Pair>>, String> GetFilePathInner(bool spill); - public: const SharedPtr data_dir_{}; const SharedPtr temp_dir_{}; diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index 8edb327f44..e6c3366cab 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -82,7 +82,7 @@ void HnswFileWorker::AllocateInMemory() { } void HnswFileWorker::FreeInMemory() { - if (!data_) { + if (!data_ || !hnsw_mem_) { String error_message = "FreeInMemory: Data is not allocated."; UnrecoverableError(error_message); } @@ -99,6 +99,10 @@ void HnswFileWorker::FreeInMemory() { *p); delete p; data_ = nullptr; + + auto [defer_fn, read_path] = GetFilePathInner(false); + VirtualStore::MunmapFile(read_path); + hnsw_mem_ = nullptr; } bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) { @@ -115,7 +119,7 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const } else { using IndexT = std::decay_t; if constexpr (IndexT::kOwnMem) { - index->Save(*file_handle_); + index->SaveToPtr(*file_handle_); } else { UnrecoverableError("Invalid index type."); } @@ -126,35 +130,22 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const return true; } -void HnswFileWorker::ReadFromFileImpl(SizeT file_size) { +void HnswFileWorker::ReadFromFileImpl(SizeT fsize) { if (data_ != nullptr) { UnrecoverableError("Data is already allocated."); } - data_ = static_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get()))); - auto *hnsw_index = reinterpret_cast(data_); - std::visit( - [&](auto &&index) { - using T = std::decay_t; - if constexpr (std::is_same_v) { - UnrecoverableError("Invalid index type."); - } else { - using IndexT = std::decay_t; - if constexpr (IndexT::kOwnMem) { - index = IndexT::Load(*file_handle_).release(); - } else { - UnrecoverableError("Invalid index type."); - } - } - }, - *hnsw_index); -} -bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) { - if (mmap_data_ != nullptr) { - UnrecoverableError("Mmap data is already allocated."); + auto [defer_fn, read_path] = GetFilePathInner(false); + SizeT file_size = VirtualStore::GetFileSize(read_path); + u8 *mmap_addr = nullptr; + int ret = VirtualStore::MmapFile(read_path, mmap_addr, file_size); + hnsw_mem_ = mmap_addr; + if (ret < 0) { + UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); } - mmap_data_ = reinterpret_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get(), false))); - auto *hnsw_index = reinterpret_cast(mmap_data_); + + data_ = static_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get(), false /*own_mem*/))); + auto *hnsw_index = reinterpret_cast(data_); std::visit( [&](auto &&index) { using T = std::decay_t; @@ -163,32 +154,14 @@ bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) { } else { using IndexT = std::decay_t; if constexpr (!IndexT::kOwnMem) { - const auto *p = static_cast(ptr); - index = IndexT::LoadFromPtr(p, size).release(); + const auto *mmap_addr = static_cast(hnsw_mem_); + index = IndexT::LoadFromPtr(mmap_addr, file_size).release(); } else { UnrecoverableError("Invalid index type."); } } }, *hnsw_index); - return true; -} - -void HnswFileWorker::FreeFromMmapImpl() { - if (mmap_data_ == nullptr) { - UnrecoverableError("Mmap data is not allocated."); - } - auto *hnsw_index = reinterpret_cast(mmap_data_); - std::visit( - [&](auto &&index) { - using T = std::decay_t; - if constexpr (!std::is_same_v) { - delete index; - } - }, - *hnsw_index); - delete hnsw_index; - mmap_data_ = nullptr; } } // namespace infinity \ No newline at end of file diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cppm b/src/storage/buffer/file_worker/hnsw_file_worker.cppm index 31952b88a0..a37acda9f3 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cppm +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cppm @@ -56,12 +56,9 @@ protected: void ReadFromFileImpl(SizeT file_size) override; - bool ReadFromMmapImpl(const void *ptr, SizeT size) override; - - void FreeFromMmapImpl() override; - private: SizeT index_size_{}; + void *hnsw_mem_{}; }; } // namespace infinity \ No newline at end of file diff --git a/src/storage/io/virtual_store.cpp b/src/storage/io/virtual_store.cpp index 03e12c702a..4e522b39ed 100644 --- a/src/storage/io/virtual_store.cpp +++ b/src/storage/io/virtual_store.cpp @@ -409,9 +409,9 @@ i32 VirtualStore::MmapFile(const String &file_path, u8 *&data_ptr, SizeT &data_l return -1; i32 f = open(file_path.c_str(), O_RDONLY); void *tmpd = mmap(NULL, len_f, PROT_READ, MAP_SHARED, f, 0); + close(f); if (tmpd == MAP_FAILED) return -1; - close(f); i32 rc = madvise(tmpd, len_f, MADV_NORMAL 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 index 1e0b78f929..35e4d7ce63 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm @@ -34,10 +34,10 @@ namespace infinity { template class DataStoreInner; -export template +export template class DataStoreChunkIter; -template +template class DataStoreIter; #pragma clang diagnostic push @@ -47,14 +47,9 @@ export template class DataStoreBase { public: using This = DataStoreBase; - using DataType = typename VecStoreT::DataType; using QueryVecType = typename VecStoreT::QueryVecType; - using Inner = DataStoreInner; using VecStoreMeta = typename VecStoreT::template Meta; - friend class DataStoreChunkIter; - friend class DataStoreIter; - public: template struct has_compress_type : std::false_type {}; @@ -62,26 +57,82 @@ public: template struct has_compress_type> : std::true_type {}; - DataStoreBase() : chunk_size_(0), max_chunk_n_(0), chunk_shift_(0), cur_vec_num_(0) {} - DataStoreBase(This &&other) - : chunk_size_(std::exchange(other.chunk_size_, 0)), max_chunk_n_(std::exchange(other.max_chunk_n_, 0)), - chunk_shift_(std::exchange(other.chunk_shift_, 0)), cur_vec_num_(other.cur_vec_num_.exchange(0)), - vec_store_meta_(std::move(other.vec_store_meta_)), graph_store_meta_(std::move(other.graph_store_meta_)), - inners_(std::exchange(other.inners_, nullptr)), mem_usage_(other.mem_usage_.exchange(0)) {} + DataStoreBase() = default; + DataStoreBase(VecStoreMeta &&vec_store_meta, GraphStoreMeta &&graph_store_meta) + : vec_store_meta_(std::move(vec_store_meta)), graph_store_meta_(std::move(graph_store_meta)) {} + DataStoreBase(This &&other) : vec_store_meta_(std::move(other.vec_store_meta_)), graph_store_meta_(std::move(other.graph_store_meta_)) {} DataStoreBase &operator=(This &&other) { if (this != &other) { + vec_store_meta_ = std::move(other.vec_store_meta_); + graph_store_meta_ = std::move(other.graph_store_meta_); + } + return *this; + } + ~DataStoreBase() = default; + + typename VecStoreT::QueryType MakeQuery(QueryVecType query) const { return vec_store_meta_.MakeQuery(query); } + + const VecStoreMeta &vec_store_meta() const { return vec_store_meta_; } + + SizeT dim() const { return vec_store_meta_.dim(); } + + // Graph store + Pair GetEnterPoint() const { return graph_store_meta_.GetEnterPoint(); } + + SizeT Mmax0() const { return graph_store_meta_.Mmax0(); } + SizeT Mmax() const { return graph_store_meta_.Mmax(); } + +protected: + VecStoreMeta vec_store_meta_; + GraphStoreMeta graph_store_meta_; +}; + +export template +class DataStore : public DataStoreBase { +public: + using This = DataStore; + using Base = DataStoreBase; + using DataType = typename VecStoreT::DataType; + using QueryVecType = typename VecStoreT::QueryVecType; + using Inner = DataStoreInner; + using VecStoreMeta = typename VecStoreT::template Meta; + using VecStoreInner = typename VecStoreT::template Inner; + + friend class DataStoreChunkIter; + friend class DataStoreIter; + +private: + DataStore(SizeT chunk_size, SizeT max_chunk_n, VecStoreMeta &&vec_store_meta, GraphStoreMeta &&graph_store_meta) + : Base(std::move(vec_store_meta), std::move(graph_store_meta)), chunk_size_(chunk_size), max_chunk_n_(max_chunk_n), + chunk_shift_(__builtin_ctzll(chunk_size)), inners_(MakeUnique(max_chunk_n)), mem_usage_(0) { + assert(chunk_size > 0); + assert((chunk_size & (chunk_size - 1)) == 0); + cur_vec_num_ = 0; + } + +public: + DataStore() = default; + DataStore(DataStore &&other) : Base(std::move(other)) { + chunk_size_ = std::exchange(other.chunk_size_, 0); + max_chunk_n_ = std::exchange(other.max_chunk_n_, 0); + chunk_shift_ = std::exchange(other.chunk_shift_, 0); + cur_vec_num_ = other.cur_vec_num_.exchange(0); + inners_ = std::exchange(other.inners_, nullptr); + mem_usage_ = other.mem_usage_.exchange(0); + } + DataStore &operator=(DataStore &&other) { + if (this != &other) { + Base::operator=(std::move(other)); chunk_size_ = std::exchange(other.chunk_size_, 0); max_chunk_n_ = std::exchange(other.max_chunk_n_, 0); chunk_shift_ = std::exchange(other.chunk_shift_, 0); cur_vec_num_ = other.cur_vec_num_.exchange(0); - vec_store_meta_ = std::move(other.vec_store_meta_); - graph_store_meta_ = std::move(other.graph_store_meta_); inners_ = std::exchange(other.inners_, nullptr); mem_usage_ = other.mem_usage_.exchange(0); } return *this; } - ~DataStoreBase() { + ~DataStore() { if (!inners_) { return; } @@ -89,28 +140,83 @@ public: auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); for (SizeT i = 0; i < chunk_num; ++i) { SizeT chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; - inners_[i].Free(chunk_size, graph_store_meta_); + inners_[i].Free(chunk_size, this->graph_store_meta_); } } + static This Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT Mmax0, SizeT Mmax) { + bool normalize = false; + if constexpr (Base::template has_compress_type::value) { + normalize = std::is_same_v::template Meta>; + } + VecStoreMeta vec_store_meta = VecStoreMeta::Make(dim, normalize); + GraphStoreMeta graph_store_meta = GraphStoreMeta::Make(Mmax0, Mmax); + This ret(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); + ret.cur_vec_num_ = 0; + + SizeT mem_usage = 0; + ret.inners_[0] = Inner::Make(chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); + ret.mem_usage_.store(mem_usage); + return ret; + } + void Save(LocalFileHandle &file_handle) const { SizeT cur_vec_num = this->cur_vec_num(); - auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); + auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); - file_handle.Append(&this->chunk_size_, sizeof(this->chunk_size_)); - file_handle.Append(&this->max_chunk_n_, sizeof(this->max_chunk_n_)); + file_handle.Append(&chunk_size_, sizeof(chunk_size_)); + file_handle.Append(&max_chunk_n_, sizeof(max_chunk_n_)); file_handle.Append(&cur_vec_num, sizeof(cur_vec_num)); this->vec_store_meta_.Save(file_handle); this->graph_store_meta_.Save(file_handle, cur_vec_num); for (SizeT i = 0; i < chunk_num; ++i) { - SizeT chunk_size = (i < chunk_num - 1) ? this->chunk_size_ : last_chunk_size; - this->inners_[i].Save(file_handle, chunk_size, this->vec_store_meta_, this->graph_store_meta_); + SizeT chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; + inners_[i].Save(file_handle, chunk_size, this->vec_store_meta_, this->graph_store_meta_); + } + } + + void SaveToPtr(LocalFileHandle &file_handle) const { + SizeT cur_vec_num = this->cur_vec_num(); + + file_handle.Append(&cur_vec_num, sizeof(cur_vec_num)); + this->vec_store_meta_.Save(file_handle); + this->graph_store_meta_.Save(file_handle, cur_vec_num); + + auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); + Inner::SaveToPtr(file_handle, inners_.get(), this->vec_store_meta_, this->graph_store_meta_, chunk_size_, chunk_num, last_chunk_size); + } + + static This Load(LocalFileHandle &file_handle, SizeT max_chunk_n = 0) { + SizeT chunk_size; + file_handle.Read(&chunk_size, sizeof(chunk_size)); + SizeT max_chunk_n1; + file_handle.Read(&max_chunk_n1, sizeof(max_chunk_n1)); + if (max_chunk_n == 0) { + max_chunk_n = max_chunk_n1; + } + assert(max_chunk_n >= max_chunk_n1); + + SizeT cur_vec_num; + file_handle.Read(&cur_vec_num, sizeof(cur_vec_num)); + VecStoreMeta vec_store_meta = VecStoreMeta::Load(file_handle); + GraphStoreMeta graph_store_meta = GraphStoreMeta::Load(file_handle); + + This ret = This(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); + ret.cur_vec_num_ = cur_vec_num; + + SizeT mem_usage = 0; + auto [chunk_num, last_chunk_size] = ret.ChunkInfo(cur_vec_num); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size : last_chunk_size; + ret.inners_[i] = Inner::Load(file_handle, cur_chunk_size, 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> &&graph_inners) { - graph_store_meta_ = std::move(graph_meta); + this->graph_store_meta_ = std::move(graph_meta); for (SizeT i = 0; i < graph_inners.size(); ++i) { inners_[i].SetGraphStoreInner(std::move(graph_inners[i])); } @@ -124,43 +230,101 @@ public: size += sizeof(chunk_size_); size += sizeof(max_chunk_n_); size += sizeof(cur_vec_num_); - size += vec_store_meta_.GetSizeInBytes(); - size += graph_store_meta_.GetSizeInBytes(); + size += this->vec_store_meta_.GetSizeInBytes(); + size += this->graph_store_meta_.GetSizeInBytes(); for (SizeT i = 0; i < chunk_num; ++i) { SizeT chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; - size += inners_[i].GetSizeInBytes(chunk_size, vec_store_meta_, graph_store_meta_); + size += inners_[i].GetSizeInBytes(chunk_size, this->vec_store_meta_, this->graph_store_meta_); } return size; } - typename VecStoreT::StoreType GetVec(SizeT vec_i) const { - const auto &[inner, idx] = GetInner(vec_i); - return inner.GetVec(idx, vec_store_meta_); + // Vec store + template Iterator> + Pair AddVec(Iterator &&query_iter) { + SizeT mem_usage = 0; + SizeT cur_vec_num = this->cur_vec_num(); + SizeT start_idx = cur_vec_num; + auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); + while (true) { + SizeT remain_size = chunk_size_ - last_chunk_size; + auto [insert_n, used_up] = + inners_[chunk_num - 1].AddVec(std::move(query_iter), last_chunk_size, remain_size, this->vec_store_meta_, mem_usage); + cur_vec_num += insert_n; + last_chunk_size += insert_n; + if (cur_vec_num == max_chunk_n_ * chunk_size_) { + break; + } + if (last_chunk_size == chunk_size_) { + inners_[chunk_num++] = Inner::Make(chunk_size_, this->vec_store_meta_, this->graph_store_meta_, mem_usage); + last_chunk_size = 0; + } + if (used_up) { + break; + } + } + cur_vec_num_.store(cur_vec_num); + mem_usage_.fetch_add(mem_usage); + return {start_idx, cur_vec_num}; } - typename VecStoreT::QueryType MakeQuery(QueryVecType query) const { return vec_store_meta_.MakeQuery(query); } + template Iterator> + Pair OptAddVec(Iterator &&query_iter) { + if constexpr (VecStoreT::HasOptimize) { + SizeT mem_usage = 0; + SizeT cur_vec_num = this->cur_vec_num(); + auto [chunk_num, last_chunk_size] = ChunkInfo(cur_vec_num); + if (chunk_num > 0) { + Vector> vec_inners; + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; + vec_inners.emplace_back(inners_[i].vec_store_inner(), chunk_size); + } + Iterator query_iter_copy = query_iter; + this->vec_store_meta_.template Optimize(std::move(query_iter_copy), vec_inners, mem_usage); + } + mem_usage_.fetch_add(mem_usage); + } + return AddVec(std::move(query_iter)); + } + + void Optimize() { + if constexpr (!VecStoreT::HasOptimize) { + return; + } + DenseVectorIter empty_iter(nullptr, this->dim(), 0); + AddVec(std::move(empty_iter)); + } void PrefetchVec(SizeT vec_i) const { const auto &[inner, idx] = GetInner(vec_i); - inner.PrefetchVec(idx, vec_store_meta_); + inner.PrefetchVec(idx, this->vec_store_meta_); } - const VecStoreMeta &vec_store_meta() const { return vec_store_meta_; } - - SizeT dim() const { return vec_store_meta_.dim(); } + typename VecStoreT::StoreType GetVec(SizeT vec_i) const { + const auto &[inner, idx] = GetInner(vec_i); + return inner.GetVec(idx, this->vec_store_meta_); + } // Graph store - Pair GetNeighbors(VertexType vertex_i, i32 layer_i) const { - const auto &[inner, idx] = GetInner(vertex_i); - return inner.GetNeighbors(idx, layer_i, graph_store_meta_); + void AddVertex(VertexType vec_i, i32 layer_n) { + auto [inner, idx] = GetInner(vec_i); + SizeT mem_usage = 0; + inner.AddVertex(idx, layer_n, this->graph_store_meta_, mem_usage); + mem_usage_.fetch_add(mem_usage); } - Pair GetEnterPoint() const { return graph_store_meta_.GetEnterPoint(); } + Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i) { + auto [inner, idx] = GetInner(vertex_i); + return inner.GetNeighborsMut(idx, layer_i, this->graph_store_meta_); + } - Pair TryUpdateEnterPoint(i32 layer, VertexType vertex_i) { return graph_store_meta_.TryUpdateEnterPoint(layer, vertex_i); } + Pair GetNeighbors(VertexType vertex_i, i32 layer_i) const { + const auto &[inner, idx] = GetInner(vertex_i); + return inner.GetNeighbors(idx, layer_i, this->graph_store_meta_); + } - SizeT Mmax0() const { return graph_store_meta_.Mmax0(); } - SizeT Mmax() const { return graph_store_meta_.Mmax(); } + Pair TryUpdateEnterPoint(i32 layer, VertexType vertex_i) { return this->graph_store_meta_.TryUpdateEnterPoint(layer, vertex_i); } // other LabelType GetLabel(SizeT vec_i) const { @@ -182,7 +346,10 @@ public: SizeT mem_usage() const { return mem_usage_.load(); } -protected: + template + DataStore CompressToLVQ() &&; + +private: Pair GetInner(SizeT vec_i) { return {inners_[vec_i >> chunk_shift_], vec_i & (chunk_size_ - 1)}; } Pair GetInner(SizeT vec_i) const { return {inners_[vec_i >> chunk_shift_], vec_i & (chunk_size_ - 1)}; } @@ -195,14 +362,12 @@ protected: return {chunk_num, last_chunk_size}; }; -protected: +private: SizeT chunk_size_; SizeT max_chunk_n_; SizeT chunk_shift_; Atomic cur_vec_num_; - VecStoreMeta vec_store_meta_; - GraphStoreMeta graph_store_meta_; UniquePtr inners_; Atomic mem_usage_ = 0; @@ -216,10 +381,10 @@ public: for (i = 0; i < chunk_num; ++i) { i32 max_l1 = -1; SizeT chunk_size = i < chunk_num - 1 ? chunk_size_ : last_chunk_size; - inners_[i].Check(chunk_size, graph_store_meta_, i * chunk_size_, cur_vec_num, max_l1); + inners_[i].Check(chunk_size, this->graph_store_meta_, i * chunk_size_, cur_vec_num, max_l1); max_l = std::max(max_l, max_l1); } - auto [max_layer, ep] = GetEnterPoint(); + auto [max_layer, ep] = this->GetEnterPoint(); if (max_l != max_layer) { UnrecoverableError("max_l != max_layer"); } @@ -232,206 +397,97 @@ public: os << "[CONST] chunk_size: " << chunk_size_ << ", max_chunk_n: " << max_chunk_n_ << ", chunk_shift: " << chunk_shift_ << std::endl; os << "cur_vec_num: " << cur_vec_num << std::endl; - vec_store_meta_.Dump(os); + this->vec_store_meta_.Dump(os); for (SizeT i = 0; i < chunk_num; ++i) { os << "chunk " << i << std::endl; SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; - inners_[i].DumpVec(os, i * chunk_size_, cur_chunk_size, vec_store_meta_); + inners_[i].DumpVec(os, i * chunk_size_, cur_chunk_size, this->vec_store_meta_); } - graph_store_meta_.Dump(os); + this->graph_store_meta_.Dump(os); for (SizeT i = 0; i < chunk_num; ++i) { os << "chunk " << i << std::endl; SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size_ : last_chunk_size; - inners_[i].DumpGraph(os, cur_chunk_size, graph_store_meta_); + inners_[i].DumpGraph(os, cur_chunk_size, this->graph_store_meta_); } } }; -export template -class DataStore : public DataStoreBase { +export template +class DataStore : public DataStoreBase { public: - using This = DataStore; - using Base = DataStoreBase; - using DataType = typename VecStoreT::DataType; - using QueryVecType = typename VecStoreT::QueryVecType; - using Inner = DataStoreInner; - using VecStoreMeta = typename VecStoreT::template Meta; - using VecStoreInner = typename VecStoreT::template Inner; + using This = DataStore; + using VecStoreMeta = typename VecStoreT::template Meta; + using Base = DataStoreBase; + using Inner = DataStoreInner; private: - DataStore(SizeT chunk_size, SizeT max_chunk_n, VecStoreMeta &&vec_store_meta, GraphStoreMeta &&graph_store_meta) { - this->chunk_size_ = chunk_size; - this->max_chunk_n_ = max_chunk_n; - this->vec_store_meta_ = std::move(vec_store_meta); - this->graph_store_meta_ = std::move(graph_store_meta); - assert(chunk_size > 0); - assert((chunk_size & (chunk_size - 1)) == 0); - this->chunk_shift_ = __builtin_ctzll(chunk_size); - this->inners_ = MakeUnique(max_chunk_n); - } + DataStore(SizeT cur_vec_num, VecStoreMeta vec_store_meta, GraphStoreMeta graph_store_meta) + : Base(std::move(vec_store_meta), std::move(graph_store_meta)), cur_vec_num_(cur_vec_num) {} public: DataStore() = default; - - static This Make(SizeT chunk_size, SizeT max_chunk_n, SizeT dim, SizeT Mmax0, SizeT Mmax) { - bool normalize = false; - if constexpr (Base::template has_compress_type::value) { - normalize = std::is_same_v::template Meta>; + DataStore(DataStore &&other) + : Base(std::move(other)), inner_(std::move(other.inner_)), ptr_start_(other.ptr_start_), cur_vec_num_(other.cur_vec_num_) {} + DataStore &operator=(DataStore &&other) { + if (this != &other) { + Base::operator=(std::move(other)); + inner_ = std::move(other.inner_); + ptr_start_ = other.ptr_start_; + cur_vec_num_ = other.cur_vec_num_; } - VecStoreMeta vec_store_meta = VecStoreMeta::Make(dim, normalize); - GraphStoreMeta graph_store_meta = GraphStoreMeta::Make(Mmax0, Mmax); - This ret(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); - ret.cur_vec_num_ = 0; - - SizeT mem_usage = 0; - ret.inners_[0] = Inner::Make(chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); - ret.mem_usage_.store(mem_usage); - return ret; + return *this; } + ~DataStore() = default; - static This Load(LocalFileHandle &file_handle, SizeT max_chunk_n = 0) { - SizeT chunk_size; - file_handle.Read(&chunk_size, sizeof(chunk_size)); - SizeT max_chunk_n1; - file_handle.Read(&max_chunk_n1, sizeof(max_chunk_n1)); - if (max_chunk_n == 0) { - max_chunk_n = max_chunk_n1; - } - assert(max_chunk_n >= max_chunk_n1); - - SizeT cur_vec_num; - file_handle.Read(&cur_vec_num, sizeof(cur_vec_num)); - VecStoreMeta vec_store_meta = VecStoreMeta::Load(file_handle); - GraphStoreMeta graph_store_meta = GraphStoreMeta::Load(file_handle); - - This ret = This(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta)); - ret.cur_vec_num_ = cur_vec_num; + static This LoadFromPtr(const char *&ptr) { + SizeT cur_vec_num = ReadBufAdv(ptr); + VecStoreMeta vec_store_meta = VecStoreMeta::LoadFromPtr(ptr); + GraphStoreMeta graph_store_meta = GraphStoreMeta::LoadFromPtr(ptr); - SizeT mem_usage = 0; - auto [chunk_num, last_chunk_size] = ret.ChunkInfo(cur_vec_num); - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size : last_chunk_size; - ret.inners_[i] = Inner::Load(file_handle, cur_chunk_size, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); - } - ret.mem_usage_.store(mem_usage); + This ret = This(cur_vec_num, std::move(vec_store_meta), std::move(graph_store_meta)); + ret.inner_ = Inner::LoadFromPtr(ptr, cur_vec_num, cur_vec_num, ret.vec_store_meta_, ret.graph_store_meta_); return ret; } - template Iterator> - Pair AddVec(Iterator &&query_iter) { - SizeT mem_usage = 0; - SizeT cur_vec_num = this->cur_vec_num(); - SizeT start_idx = cur_vec_num; - auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); - while (true) { - SizeT remain_size = this->chunk_size_ - last_chunk_size; - auto [insert_n, used_up] = - this->inners_[chunk_num - 1].AddVec(std::move(query_iter), last_chunk_size, remain_size, this->vec_store_meta_, mem_usage); - cur_vec_num += insert_n; - last_chunk_size += insert_n; - if (cur_vec_num == this->max_chunk_n_ * this->chunk_size_) { - break; - } - if (last_chunk_size == this->chunk_size_) { - this->inners_[chunk_num++] = Inner::Make(this->chunk_size_, this->vec_store_meta_, this->graph_store_meta_, mem_usage); - last_chunk_size = 0; - } - if (used_up) { - break; - } - } - this->cur_vec_num_.store(cur_vec_num); - this->mem_usage_.fetch_add(mem_usage); - return {start_idx, cur_vec_num}; - } + typename VecStoreT::StoreType GetVec(SizeT vec_i) const { return inner_.GetVec(vec_i, this->vec_store_meta_); } - template Iterator> - Pair OptAddVec(Iterator &&query_iter) { - if constexpr (VecStoreT::HasOptimize) { - SizeT mem_usage = 0; - SizeT cur_vec_num = this->cur_vec_num(); - auto [chunk_num, last_chunk_size] = this->ChunkInfo(cur_vec_num); - if (chunk_num > 0) { - Vector> vec_inners; - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT chunk_size = (i < chunk_num - 1) ? this->chunk_size_ : last_chunk_size; - vec_inners.emplace_back(this->inners_[i].vec_store_inner(), chunk_size); - } - Iterator query_iter_copy = query_iter; - this->vec_store_meta_.template Optimize(std::move(query_iter_copy), vec_inners, mem_usage); - } - this->mem_usage_.fetch_add(mem_usage); - } - return AddVec(std::move(query_iter)); - } + void PrefetchVec(SizeT vec_i) const { inner_.PrefetchVec(vec_i, this->vec_store_meta_); } - void Optimize() { - if constexpr (!VecStoreT::HasOptimize) { - return; - } - DenseVectorIter empty_iter(nullptr, this->dim(), 0); - AddVec(std::move(empty_iter)); + Pair GetNeighbors(VertexType vertex_i, i32 layer_i) const { + return inner_.GetNeighbors(vertex_i, layer_i, this->graph_store_meta_); } - void AddVertex(VertexType vec_i, i32 layer_n) { - auto [inner, idx] = this->GetInner(vec_i); - SizeT mem_usage = 0; - inner.AddVertex(idx, layer_n, this->graph_store_meta_, mem_usage); - this->mem_usage_.fetch_add(mem_usage); - } - Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i) { - auto [inner, idx] = this->GetInner(vertex_i); - return inner.GetNeighborsMut(idx, layer_i, this->graph_store_meta_); - } + LabelType GetLabel(SizeT vec_i) const { return inner_.GetLabel(vec_i); } - template - DataStore CompressToLVQ() &&; -}; + SizeT cur_vec_num() const { return cur_vec_num_; } -export template -class DataStore : public DataStoreBase { -public: - using This = DataStore; - using VecStoreMeta = typename VecStoreT::template Meta; - using Base = DataStoreBase; - using Inner = DataStoreInner; + SizeT mem_usage() const { return 0; } + + const char *ptr_start() const { return ptr_start_; } + void set_ptr_start(const char *ptr_start) { ptr_start_ = ptr_start; } private: - DataStore(SizeT chunk_size, SizeT max_chunk_n, VecStoreMeta vec_store_meta, GraphStoreMeta graph_store_meta, SizeT cur_vec_num) { - this->chunk_size_ = chunk_size; - this->max_chunk_n_ = max_chunk_n; - this->vec_store_meta_ = std::move(vec_store_meta); - this->graph_store_meta_ = std::move(graph_store_meta); - assert(chunk_size > 0); - assert((chunk_size & (chunk_size - 1)) == 0); - this->chunk_shift_ = __builtin_ctzll(chunk_size); - this->inners_ = MakeUnique(max_chunk_n); - this->cur_vec_num_.store(cur_vec_num); - } + Inner inner_; + const char *ptr_start_ = nullptr; + SizeT cur_vec_num_ = 0; public: - DataStore() = default; - static This LoadFromPtr(const char *&ptr, SizeT max_chunk_n = 0) { - SizeT chunk_size = ReadBufAdv(ptr); - SizeT max_chunk_n1 = ReadBufAdv(ptr); - if (max_chunk_n == 0) { - max_chunk_n = max_chunk_n1; + void Check() const { + i32 max_l = -1; + inner_.Check(cur_vec_num_, this->graph_store_meta_, 0, cur_vec_num_, max_l); + auto [max_layer, ep] = this->GetEnterPoint(); + if (max_l != max_layer) { + UnrecoverableError("max_l != max_layer"); } - assert(max_chunk_n >= max_chunk_n1); - - SizeT cur_vec_num = ReadBufAdv(ptr); - VecStoreMeta vec_store_meta = VecStoreMeta::LoadFromPtr(ptr); - GraphStoreMeta graph_store_meta = GraphStoreMeta::LoadFromPtr(ptr); - - This ret = This(chunk_size, max_chunk_n, std::move(vec_store_meta), std::move(graph_store_meta), cur_vec_num); + } - auto [chunk_num, last_chunk_size] = ret.ChunkInfo(cur_vec_num); - for (SizeT i = 0; i < chunk_num; ++i) { - SizeT cur_chunk_size = (i < chunk_num - 1) ? chunk_size : last_chunk_size; - ret.inners_[i] = Inner::LoadFromPtr(ptr, cur_chunk_size, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_); - } - return ret; + void Dump() const { + std::cout << "[CONST] cur_vec_num: " << cur_vec_num_ << std::endl; + this->vec_store_meta_.Dump(); + inner_.DumpVec(std::cout, 0, cur_vec_num_, this->vec_store_meta_); + this->graph_store_meta_.Dump(); + inner_.DumpGraph(std::cout, cur_vec_num_, this->graph_store_meta_); } }; @@ -447,7 +503,7 @@ public: using VecStoreMeta = typename VecStoreT::template Meta; using GraphStoreInner = GraphStoreInner; - friend class DataStoreIter; + friend class DataStoreIter; public: DataStoreInnerBase() = default; @@ -458,6 +514,27 @@ public: file_handle.Append(this->labels_.get(), sizeof(LabelType) * cur_vec_num); } + static void SaveToPtr(LocalFileHandle &file_handle, + const This *inners, + const VecStoreMeta &vec_store_meta, + const GraphStoreMeta &graph_store_meta, + SizeT ck_size, + SizeT chunk_num, + SizeT last_chunk_size) { + Vector vec_store_inners; + Vector graph_store_inners; + for (SizeT i = 0; i < chunk_num; ++i) { + vec_store_inners.emplace_back(&inners[i].vec_store_inner_); + graph_store_inners.emplace_back(&inners[i].graph_store_inner_); + } + VecStoreInner::SaveToPtr(file_handle, vec_store_inners, vec_store_meta, ck_size, chunk_num, last_chunk_size); + GraphStoreInner::SaveToPtr(file_handle, graph_store_inners, graph_store_meta, ck_size, chunk_num, last_chunk_size); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? ck_size : last_chunk_size; + file_handle.Append(inners[i].labels_.get(), sizeof(LabelType) * chunk_size); + } + } + void Free(SizeT cur_vec_num, const GraphStoreMeta &graph_store_meta) { graph_store_inner_.Free(cur_vec_num, graph_store_meta); } SizeT GetSizeInBytes(SizeT chunk_size, const VecStoreMeta &vec_store_meta, const GraphStoreMeta &graph_store_meta) const { @@ -480,10 +557,6 @@ public: LabelType GetLabel(VertexType vec_i) const { return labels_[vec_i]; } - std::shared_lock SharedLock(VertexType vec_i) const { return std::shared_lock(vertex_mutex_[vec_i]); } - - std::unique_lock UniqueLock(VertexType vec_i) { return std::unique_lock(vertex_mutex_[vec_i]); } - VecStoreInner *vec_store_inner() { return &vec_store_inner_; } GraphStoreInner *graph_store_inner() { return &graph_store_inner_; } @@ -493,7 +566,6 @@ protected: VecStoreInner vec_store_inner_; GraphStoreInner graph_store_inner_; ArrayPtr labels_; - mutable UniquePtr vertex_mutex_; public: void Check(SizeT chunk_size, const GraphStoreMeta &meta, VertexType vertex_i_offset, SizeT cur_vec_num, i32 &max_l) const { @@ -525,7 +597,7 @@ private: this->vec_store_inner_ = std::move(vec_store_inner); this->graph_store_inner_ = std::move(graph_store_inner); this->labels_ = MakeUnique(chunk_size); - this->vertex_mutex_ = MakeUnique(chunk_size); + vertex_mutex_ = MakeUnique(chunk_size); } public: @@ -575,6 +647,13 @@ public: Pair GetNeighborsMut(VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { return this->graph_store_inner_.GetNeighborsMut(vertex_i, layer_i, meta); } + + std::shared_lock SharedLock(VertexType vec_i) const { return std::shared_lock(vertex_mutex_[vec_i]); } + + std::unique_lock UniqueLock(VertexType vec_i) { return std::unique_lock(vertex_mutex_[vec_i]); } + +private: + mutable UniquePtr vertex_mutex_; }; template @@ -590,7 +669,6 @@ private: this->vec_store_inner_ = std::move(vec_store_inner); this->graph_store_inner_ = std::move(graph_store_inner); this->labels_ = labels; - this->vertex_mutex_ = MakeUnique(chunk_size); } public: @@ -606,12 +684,12 @@ public: } }; -template +template class DataStoreChunkIter { public: - using Inner = typename DataStore::Inner; + using Inner = typename DataStore::Inner; - DataStoreChunkIter(const DataStore *data_store) : data_store_(data_store) { + DataStoreChunkIter(const DataStore *data_store) : data_store_(data_store) { std::tie(chunk_num_, last_chunk_size_) = data_store_->ChunkInfo(data_store_->cur_vec_num()); } @@ -625,7 +703,7 @@ public: return ret; } - const DataStore *data_store_; + const DataStore *data_store_; private: SizeT cur_chunk_i_ = 0; @@ -633,11 +711,11 @@ private: SizeT last_chunk_size_; }; -template +template class DataStoreInnerIter { public: - using VecMeta = typename VecStoreT::template Meta; - using Inner = DataStoreInner; + using VecMeta = typename VecStoreT::template Meta; + using Inner = DataStoreInner; using StoreType = typename VecStoreT::StoreType; DataStoreInnerIter(const VecMeta *vec_meta, const Inner *inner, SizeT max_vec_num) @@ -660,13 +738,13 @@ private: SizeT cur_idx_; }; -template +template class DataStoreIter { public: using StoreType = typename VecStoreT::StoreType; - using InnerIter = DataStoreInnerIter; + using InnerIter = DataStoreInnerIter; - DataStoreIter(const DataStore *data_store) : data_store_iter_(data_store), inner_iter_(None) {} + DataStoreIter(const DataStore *data_store) : data_store_iter_(data_store), inner_iter_(None) {} Optional> Next() { if (!inner_iter_.has_value()) { @@ -687,7 +765,7 @@ public: } private: - DataStoreChunkIter data_store_iter_; + DataStoreChunkIter data_store_iter_; Optional inner_iter_; }; @@ -707,7 +785,7 @@ DataStore DataStorevec_store_meta_.dim(), this->Mmax0(), this->Mmax()); - ret.OptAddVec(DataStoreIter(this)); + ret.OptAddVec(DataStoreIter(this)); ret.SetGraph(std::move(this->graph_store_meta_), std::move(graph_inners)); this->inners_ = nullptr; return ret; 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 index c3cfd1b112..60c4b0c8b1 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm @@ -42,20 +42,19 @@ export class GraphStoreMeta { private: GraphStoreMeta(SizeT Mmax0, SizeT Mmax) : Mmax0_(Mmax0), Mmax_(Mmax), level0_size_(sizeof(VertexL0) + sizeof(VertexType) * Mmax0), - levelx_size_(sizeof(VertexLX) + sizeof(VertexType) * Mmax), last_save_num_(0) {} + levelx_size_(sizeof(VertexLX) + sizeof(VertexType) * Mmax) {} public: - GraphStoreMeta() : Mmax0_(0), Mmax_(0), level0_size_(0), levelx_size_(0), last_save_num_(0), max_layer_(-1), enterpoint_(-1) {} + GraphStoreMeta() : Mmax0_(0), Mmax_(0), level0_size_(0), levelx_size_(0), max_layer_(-1), enterpoint_(-1) {} GraphStoreMeta(GraphStoreMeta &&other) : Mmax0_(std::exchange(other.Mmax0_, 0)), Mmax_(std::exchange(other.Mmax_, 0)), level0_size_(std::exchange(other.level0_size_, 0)), - levelx_size_(std::exchange(other.levelx_size_, 0)), last_save_num_(std::exchange(other.last_save_num_, 0)), - max_layer_(std::exchange(other.max_layer_, -1)), enterpoint_(std::exchange(other.enterpoint_, -1)) {} + levelx_size_(std::exchange(other.levelx_size_, 0)), max_layer_(std::exchange(other.max_layer_, -1)), + enterpoint_(std::exchange(other.enterpoint_, -1)) {} GraphStoreMeta &operator=(GraphStoreMeta &&other) { Mmax0_ = std::exchange(other.Mmax0_, 0); Mmax_ = std::exchange(other.Mmax_, 0); level0_size_ = std::exchange(other.level0_size_, 0); levelx_size_ = std::exchange(other.levelx_size_, 0); - last_save_num_ = std::exchange(other.last_save_num_, 0); max_layer_ = std::exchange(other.max_layer_, -1); enterpoint_ = std::exchange(other.enterpoint_, -1); return *this; @@ -74,7 +73,6 @@ public: void Save(LocalFileHandle &file_handle, SizeT cur_vec_num) const { file_handle.Append(&Mmax0_, sizeof(Mmax0_)); file_handle.Append(&Mmax_, sizeof(Mmax_)); - file_handle.Append(&cur_vec_num, sizeof(last_save_num_)); file_handle.Append(&max_layer_, sizeof(max_layer_)); file_handle.Append(&enterpoint_, sizeof(enterpoint_)); @@ -84,12 +82,8 @@ public: SizeT Mmax0, Mmax; file_handle.Read(&Mmax0, sizeof(Mmax0)); file_handle.Read(&Mmax, sizeof(Mmax)); - SizeT last_save_num; - file_handle.Read(&last_save_num, sizeof(last_save_num)); GraphStoreMeta meta(Mmax0, Mmax); - meta.last_save_num_ = last_save_num; - i32 max_layer; VertexType enterpoint; file_handle.Read(&max_layer, sizeof(max_layer)); @@ -102,9 +96,7 @@ public: static GraphStoreMeta LoadFromPtr(const char *&ptr) { SizeT Mmax0 = ReadBufAdv(ptr); SizeT Mmax = ReadBufAdv(ptr); - SizeT last_save_num = ReadBufAdv(ptr); GraphStoreMeta meta(Mmax0, Mmax); - meta.last_save_num_ = last_save_num; i32 max_layer = ReadBufAdv(ptr); VertexType enterpoint = ReadBufAdv(ptr); meta.max_layer_ = max_layer; @@ -116,7 +108,6 @@ public: SizeT Mmax() const { return Mmax_; } SizeT level0_size() const { return level0_size_; } SizeT levelx_size() const { return levelx_size_; } - SizeT last_save_num() const { return last_save_num_; } Pair GetEnterPoint() const { std::unique_lock lck(mtx_); @@ -141,7 +132,6 @@ private: SizeT Mmax_; SizeT level0_size_; SizeT levelx_size_; - SizeT last_save_num_; mutable std::mutex mtx_; i32 max_layer_; @@ -159,43 +149,71 @@ public: template class GraphStoreInnerBase { public: + using This = GraphStoreInnerBase; + GraphStoreInnerBase() = default; void Save(LocalFileHandle &file_handle, SizeT cur_vertex_n, const GraphStoreMeta &meta) const { SizeT layer_sum = 0; - Vector> layers_ptrs_off; for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { - const VertexL0 *v = GetLevel0(vertex_i, meta); - if (!v->layer_n_) { - continue; - } - SizeT offset = layer_sum * meta.levelx_size(); - SizeT off = reinterpret_cast(&v->layers_p_) - this->graph_.get(); - layers_ptrs_off.emplace_back(off, offset); - layer_sum += v->layer_n_; + layer_sum += GetLevel0(vertex_i, meta)->layer_n_; } file_handle.Append(&layer_sum, sizeof(layer_sum)); - { - auto buffer = MakeUnique(cur_vertex_n * meta.level0_size()); - std::copy(this->graph_.get(), this->graph_.get() + cur_vertex_n * meta.level0_size(), buffer.get()); - for (auto [off, offset] : layers_ptrs_off) { - char *ptr = buffer.get() + off; - *reinterpret_cast(ptr) = offset; - } - file_handle.Append(buffer.get(), cur_vertex_n * meta.level0_size()); - } + file_handle.Append(graph_.get(), cur_vertex_n * meta.level0_size()); for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { const VertexL0 *v = GetLevel0(vertex_i, meta); if (v->layer_n_) { - char *ptr = v->layers_p_; - file_handle.Append(ptr, meta.levelx_size() * v->layer_n_); + file_handle.Append(v->layers_p_, meta.levelx_size() * v->layer_n_); } } } - void Free(SizeT current_vertex_num, const GraphStoreMeta &meta) { - for (VertexType vertex_i = loaded_vertex_n_; vertex_i < VertexType(current_vertex_num); ++vertex_i) { - delete[] GetLevel0(vertex_i, meta)->layers_p_; + static void SaveToPtr(LocalFileHandle &file_handle, + const Vector &inners, + const GraphStoreMeta &meta, + SizeT ck_size, + SizeT chunk_num, + SizeT last_chunk_size) { + SizeT layer_sum = 0; + Vector>> layers_ptrs_off_vec; + for (SizeT i = 0; i < chunk_num; ++i) { + Vector> layers_ptrs_off; + SizeT chunk_size = (i < chunk_num - 1) ? ck_size : last_chunk_size; + const auto &inner = inners[i]; + for (VertexType vertex_i = 0; vertex_i < (VertexType)chunk_size; ++vertex_i) { + const VertexL0 *v = inner->GetLevel0(vertex_i, meta); + if (!v->layer_n_) { + continue; + } + SizeT offset = layer_sum * meta.levelx_size(); + SizeT ptr_off = reinterpret_cast(&v->layers_p_) - inner->graph_.get(); + layers_ptrs_off.emplace_back(ptr_off, offset); + layer_sum += v->layer_n_; + } + layers_ptrs_off_vec.emplace_back(std::move(layers_ptrs_off)); + } + file_handle.Append(&layer_sum, sizeof(layer_sum)); + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? ck_size : last_chunk_size; + const auto &inner = inners[i]; + auto buffer = MakeUnique(chunk_size * meta.level0_size()); + std::copy(inner->graph_.get(), inner->graph_.get() + chunk_size * meta.level0_size(), buffer.get()); + for (const auto &[ptr_off, offset] : layers_ptrs_off_vec[i]) { + char *ptr = buffer.get() + ptr_off; + *reinterpret_cast(ptr) = offset; + } + file_handle.Append(buffer.get(), chunk_size * meta.level0_size()); + } + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? ck_size : last_chunk_size; + const auto &inner = inners[i]; + for (VertexType vertex_i = 0; vertex_i < (VertexType)chunk_size; ++vertex_i) { + const VertexL0 *v = inner->GetLevel0(vertex_i, meta); + if (v->layer_n_) { + char *ptr = v->layers_p_; + file_handle.Append(ptr, meta.levelx_size() * v->layer_n_); + } + } } } @@ -229,14 +247,9 @@ protected: const VertexLX *GetLevelX(const char *layer_p, VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) const { assert(layer_i > 0); if constexpr (OwnMem) { - if (SizeT(vertex_i) >= meta.last_save_num()) { - return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); - } - SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); - return reinterpret_cast(layer_start_.get() + offset); + return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); } else { - assert(SizeT(vertex_i) < meta.last_save_num()); - SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); + SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); return reinterpret_cast(layer_start_.get() + offset); } } @@ -244,7 +257,6 @@ protected: protected: ArrayPtr graph_; PPtr layer_start_; - SizeT loaded_vertex_n_; //---------------------------------------------- Following is the tmp debug function. ---------------------------------------------- @@ -320,16 +332,23 @@ public: export template class GraphStoreInner : public GraphStoreInnerBase { -private: +public: using Base = GraphStoreInnerBase; - GraphStoreInner(SizeT max_vertex, const GraphStoreMeta &meta, SizeT loaded_vertex_n) { + +private: + GraphStoreInner(SizeT max_vertex, const GraphStoreMeta &meta, SizeT loaded_vertex_n) : loaded_vertex_n_(loaded_vertex_n) { this->graph_ = MakeUnique(max_vertex * meta.level0_size()); - this->loaded_vertex_n_ = loaded_vertex_n; } public: GraphStoreInner() = default; + void Free(SizeT current_vertex_num, const GraphStoreMeta &meta) { + for (VertexType vertex_i = loaded_vertex_n_; vertex_i < VertexType(current_vertex_num); ++vertex_i) { + delete[] GetLevel0(vertex_i, meta)->layers_p_; + } + } + static GraphStoreInner Make(SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { GraphStoreInner graph_store(max_vertex, meta, 0); std::fill(graph_store.graph_.get(), graph_store.graph_.get() + max_vertex * meta.level0_size(), 0); @@ -347,9 +366,17 @@ public: file_handle.Read(graph_store.graph_.get(), cur_vertex_n * meta.level0_size()); auto loaded_layers = MakeUnique(meta.levelx_size() * layer_sum); - file_handle.Read(loaded_layers.get(), meta.levelx_size() * layer_sum); - - graph_store.layer_start_.set(loaded_layers.get()); + 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_) { + file_handle.Read(loaded_layers_p, meta.levelx_size() * v->layer_n_); + v->layers_p_ = loaded_layers_p; + loaded_layers_p += 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(); @@ -388,34 +415,28 @@ private: } VertexLX *GetLevelX(char *layer_p, VertexType vertex_i, i32 layer_i, const GraphStoreMeta &meta) { assert(layer_i > 0); - if (SizeT(vertex_i) >= meta.last_save_num()) { - return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); - } - SizeT offset = reinterpret_cast(layer_p) + (layer_i - 1) * meta.levelx_size(); - return reinterpret_cast(this->layer_start_.get() + offset); + return reinterpret_cast(layer_p + (layer_i - 1) * meta.levelx_size()); } private: ArrayPtr loaded_layers_; - Vector layers_ptrs_; + SizeT loaded_vertex_n_; }; export template <> class GraphStoreInner : public GraphStoreInnerBase { public: + using Base = GraphStoreInnerBase; GraphStoreInner() = default; - GraphStoreInner(SizeT loaded_vertex_n, const char *ptr) { - this->graph_ = ptr; - this->loaded_vertex_n_ = loaded_vertex_n; - } + GraphStoreInner(const char *ptr) { this->graph_ = ptr; } static GraphStoreInner LoadFromPtr(const char *&ptr, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta) { assert(cur_vertex_n <= max_vertex); SizeT layer_sum = ReadBufAdv(ptr); - GraphStoreInner graph_store(cur_vertex_n, ptr); + GraphStoreInner graph_store(ptr); graph_store.layer_start_.set(ptr + cur_vertex_n * meta.level0_size()); ptr += cur_vertex_n * meta.level0_size() + layer_sum * meta.levelx_size(); 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 index d820768cd0..af9bd18f48 100644 --- 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 @@ -301,7 +301,7 @@ public: template class LVQVecStoreInnerBase { public: - using This = LVQVecStoreInner; + using This = LVQVecStoreInnerBase; using Meta = LVQVecStoreMetaBase; // Decompress: Q = scale * C + bias + Mean using LocalCacheType = LVQCache::LocalCacheType; @@ -316,6 +316,14 @@ public: file_handle.Append(ptr_.get(), cur_vec_num * meta.compress_data_size()); } + static void + SaveToPtr(LocalFileHandle &file_handle, const Vector &inners, const Meta &meta, SizeT ck_size, SizeT chunk_num, SizeT last_chunk_size) { + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? ck_size : last_chunk_size; + file_handle.Append(inners[i]->ptr_.get(), chunk_size * meta.compress_data_size()); + } + } + const LVQData *GetVec(SizeT idx, const Meta &meta) const { return reinterpret_cast(ptr_.get() + idx * meta.compress_data_size()); } @@ -343,10 +351,12 @@ public: export template class LVQVecStoreInner : public LVQVecStoreInnerBase { +public: using This = LVQVecStoreInner; using Meta = LVQVecStoreMetaBase; using LocalCacheType = LVQCache::LocalCacheType; using LVQData = LVQData; + using Base = LVQVecStoreInnerBase; private: LVQVecStoreInner(SizeT max_vec_num, const Meta &meta) { this->ptr_ = MakeUnique(max_vec_num * meta.compress_data_size()); } @@ -376,8 +386,10 @@ private: export template class LVQVecStoreInner : public LVQVecStoreInnerBase { +public: using This = LVQVecStoreInner; using Meta = LVQVecStoreMetaBase; + using Base = LVQVecStoreInnerBase; private: LVQVecStoreInner(const char *ptr) { this->ptr_ = ptr; } 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 index cbede4b320..cda6e3918a 100644 --- 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 @@ -87,6 +87,7 @@ class PlainVecStoreInnerBase { public: using This = PlainVecStoreInnerBase; using Meta = PlainVecStoreMeta; + using Base = PlainVecStoreInnerBase; public: PlainVecStoreInnerBase() = default; @@ -97,6 +98,14 @@ public: file_handle.Append(ptr_.get(), sizeof(DataType) * cur_vec_num * meta.dim()); } + static void + SaveToPtr(LocalFileHandle &file_handle, const Vector &inners, const Meta &meta, SizeT ck_size, SizeT chunk_num, SizeT last_chunk_size) { + for (SizeT i = 0; i < chunk_num; ++i) { + SizeT chunk_size = (i < chunk_num - 1) ? ck_size : last_chunk_size; + file_handle.Append(inners[i]->ptr_.get(), sizeof(DataType) * chunk_size * meta.dim()); + } + } + const DataType *GetVec(SizeT idx, const Meta &meta) const { return ptr_.get() + idx * meta.dim(); } void Prefetch(VertexType vec_i, const Meta &meta) const { _mm_prefetch(reinterpret_cast(GetVec(vec_i, meta)), _MM_HINT_T0); } @@ -119,8 +128,10 @@ public: export template class PlainVecStoreInner : public PlainVecStoreInnerBase { +public: using This = PlainVecStoreInner; using Meta = PlainVecStoreMeta; + using Base = PlainVecStoreInnerBase; protected: PlainVecStoreInner(SizeT max_vec_num, const Meta &meta) { this->ptr_ = MakeUnique(max_vec_num * meta.dim()); } diff --git a/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm index def76f99a0..f2194ff5cb 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/sparse_vec_store.cppm @@ -150,8 +150,8 @@ public: void Prefetch(SizeT idx, const Meta &meta) const { const SparseVecEle &vec = vecs_[idx]; - _mm_prefetch((const char*)vec.indices_.get(), _MM_HINT_T0); - _mm_prefetch((const char*)vec.data_.get(), _MM_HINT_T0); + _mm_prefetch((const char *)vec.indices_.get(), _MM_HINT_T0); + _mm_prefetch((const char *)vec.data_.get(), _MM_HINT_T0); } private: diff --git a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm index 6244c7a592..333a94a243 100644 --- a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm +++ b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm @@ -90,6 +90,12 @@ public: data_store_.Save(file_handle); } + void SaveToPtr(LocalFileHandle &file_handle) const { + file_handle.Append(&M_, sizeof(M_)); + file_handle.Append(&ef_construction_, sizeof(ef_construction_)); + data_store_.SaveToPtr(file_handle); + } + protected: // >= 0 i32 GenerateRandomLayer() { @@ -162,7 +168,7 @@ protected: } std::shared_lock lock; - if constexpr (WithLock) { + if constexpr (WithLock && OwnMem) { lock = data_store_.SharedLock(c_idx); } @@ -202,7 +208,7 @@ protected: check = false; std::shared_lock lock; - if constexpr (WithLock) { + if constexpr (WithLock && OwnMem) { lock = data_store_.SharedLock(cur_p); } @@ -508,9 +514,17 @@ public: this->data_store_ = std::move(data_store); this->distance_ = std::move(distance); } + KnnHnsw(This &&other) : KnnHnswBase(std::move(other)) {} + KnnHnsw &operator=(This &&other) { + if (this != &other) { + KnnHnswBase::operator=(std::move(other)); + } + return *this; + } public: static UniquePtr LoadFromPtr(const char *&ptr, SizeT size) { + const char *ptr_start = ptr; const char *ptr_end = ptr + size; SizeT M = ReadBufAdv(ptr); SizeT ef_construction = ReadBufAdv(ptr); @@ -519,8 +533,17 @@ public: if (SizeT diff = ptr_end - ptr; diff != 0) { UnrecoverableError("LoadFromPtr failed"); } + data_store.set_ptr_start(ptr_start); return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); } + + bool LoadAgain(const char *&ptr, SizeT size) { + if (this->data_store_.ptr_start() == ptr) { + return false; + } + *this = std::move(*LoadFromPtr(ptr, size)); + return true; + } }; } // namespace infinity \ No newline at end of file diff --git a/src/storage/meta/entry/chunk_index_entry.cpp b/src/storage/meta/entry/chunk_index_entry.cpp index bd62a3dfc6..8e4ebdb23b 100644 --- a/src/storage/meta/entry/chunk_index_entry.cpp +++ b/src/storage/meta/entry/chunk_index_entry.cpp @@ -276,7 +276,6 @@ SharedPtr ChunkIndexEntry::NewReplayChunkIndexEntry(ChunkID chu column_def, buffer_mgr->persistence_manager()); BufferObj *buffer_obj = buffer_mgr->GetBufferObject(std::move(file_worker)); - buffer_obj->ToMmap(); chunk_index_entry->buffer_obj_ = buffer_obj; break; } @@ -456,9 +455,6 @@ void ChunkIndexEntry::SaveIndexFile() { return; } buffer_obj_->Save(); - if (segment_index_entry_->table_index_entry()->index_base()->index_type_ == IndexType::kHnsw) { - buffer_obj_->ToMmap(); - } } void ChunkIndexEntry::DeprecateChunk(TxnTimeStamp commit_ts) { diff --git a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp index 4ef86209fc..2d9d7262ca 100644 --- a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp +++ b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp @@ -153,7 +153,7 @@ class HnswAlgTest : public BaseTest { if (!status.ok()) { UnrecoverableError(status.message()); } - hnsw_index->Save(*file_handle); + hnsw_index->SaveToPtr(*file_handle); } { SizeT file_size = VirtualStore::GetFileSize(filepath); 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 index 3921d46b42..41ba44497b 100644 --- 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 index 20112d82ce..895f6efb56 100644 --- 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); From 95cbf8f68b8893f5bb3684e50b42c365fae7ab5d Mon Sep 17 00:00:00 2001 From: shenyushi Date: Thu, 26 Dec 2024 16:58:28 +0800 Subject: [PATCH 07/12] fix. --- src/storage/buffer/file_worker/hnsw_file_worker.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index e6c3366cab..b6dde8bcd2 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -82,7 +82,7 @@ void HnswFileWorker::AllocateInMemory() { } void HnswFileWorker::FreeInMemory() { - if (!data_ || !hnsw_mem_) { + if (!data_ && !hnsw_mem_) { String error_message = "FreeInMemory: Data is not allocated."; UnrecoverableError(error_message); } From 075c72230891a5c5d8c7c4841410316134c65c56 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Thu, 26 Dec 2024 16:52:05 +0800 Subject: [PATCH 08/12] mmap --- .../buffer/file_worker/hnsw_file_worker.cpp | 65 +++++++++++++------ .../buffer/file_worker/hnsw_file_worker.cppm | 5 +- src/storage/meta/entry/chunk_index_entry.cpp | 1 + 3 files changed, 51 insertions(+), 20 deletions(-) diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index b6dde8bcd2..8edb327f44 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -82,7 +82,7 @@ void HnswFileWorker::AllocateInMemory() { } void HnswFileWorker::FreeInMemory() { - if (!data_ && !hnsw_mem_) { + if (!data_) { String error_message = "FreeInMemory: Data is not allocated."; UnrecoverableError(error_message); } @@ -99,10 +99,6 @@ void HnswFileWorker::FreeInMemory() { *p); delete p; data_ = nullptr; - - auto [defer_fn, read_path] = GetFilePathInner(false); - VirtualStore::MunmapFile(read_path); - hnsw_mem_ = nullptr; } bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) { @@ -119,7 +115,7 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const } else { using IndexT = std::decay_t; if constexpr (IndexT::kOwnMem) { - index->SaveToPtr(*file_handle_); + index->Save(*file_handle_); } else { UnrecoverableError("Invalid index type."); } @@ -130,22 +126,35 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const return true; } -void HnswFileWorker::ReadFromFileImpl(SizeT fsize) { +void HnswFileWorker::ReadFromFileImpl(SizeT file_size) { if (data_ != nullptr) { UnrecoverableError("Data is already allocated."); } + data_ = static_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get()))); + auto *hnsw_index = reinterpret_cast(data_); + std::visit( + [&](auto &&index) { + using T = std::decay_t; + if constexpr (std::is_same_v) { + UnrecoverableError("Invalid index type."); + } else { + using IndexT = std::decay_t; + if constexpr (IndexT::kOwnMem) { + index = IndexT::Load(*file_handle_).release(); + } else { + UnrecoverableError("Invalid index type."); + } + } + }, + *hnsw_index); +} - auto [defer_fn, read_path] = GetFilePathInner(false); - SizeT file_size = VirtualStore::GetFileSize(read_path); - u8 *mmap_addr = nullptr; - int ret = VirtualStore::MmapFile(read_path, mmap_addr, file_size); - hnsw_mem_ = mmap_addr; - if (ret < 0) { - UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); +bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) { + if (mmap_data_ != nullptr) { + UnrecoverableError("Mmap data is already allocated."); } - - data_ = static_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get(), false /*own_mem*/))); - auto *hnsw_index = reinterpret_cast(data_); + mmap_data_ = reinterpret_cast(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get(), false))); + auto *hnsw_index = reinterpret_cast(mmap_data_); std::visit( [&](auto &&index) { using T = std::decay_t; @@ -154,14 +163,32 @@ void HnswFileWorker::ReadFromFileImpl(SizeT fsize) { } else { using IndexT = std::decay_t; if constexpr (!IndexT::kOwnMem) { - const auto *mmap_addr = static_cast(hnsw_mem_); - index = IndexT::LoadFromPtr(mmap_addr, file_size).release(); + const auto *p = static_cast(ptr); + index = IndexT::LoadFromPtr(p, size).release(); } else { UnrecoverableError("Invalid index type."); } } }, *hnsw_index); + return true; +} + +void HnswFileWorker::FreeFromMmapImpl() { + if (mmap_data_ == nullptr) { + UnrecoverableError("Mmap data is not allocated."); + } + auto *hnsw_index = reinterpret_cast(mmap_data_); + std::visit( + [&](auto &&index) { + using T = std::decay_t; + if constexpr (!std::is_same_v) { + delete index; + } + }, + *hnsw_index); + delete hnsw_index; + mmap_data_ = nullptr; } } // namespace infinity \ No newline at end of file diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cppm b/src/storage/buffer/file_worker/hnsw_file_worker.cppm index a37acda9f3..31952b88a0 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cppm +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cppm @@ -56,9 +56,12 @@ protected: void ReadFromFileImpl(SizeT file_size) override; + bool ReadFromMmapImpl(const void *ptr, SizeT size) override; + + void FreeFromMmapImpl() override; + private: SizeT index_size_{}; - void *hnsw_mem_{}; }; } // namespace infinity \ No newline at end of file diff --git a/src/storage/meta/entry/chunk_index_entry.cpp b/src/storage/meta/entry/chunk_index_entry.cpp index 8e4ebdb23b..e8c1ff93aa 100644 --- a/src/storage/meta/entry/chunk_index_entry.cpp +++ b/src/storage/meta/entry/chunk_index_entry.cpp @@ -276,6 +276,7 @@ SharedPtr ChunkIndexEntry::NewReplayChunkIndexEntry(ChunkID chu column_def, buffer_mgr->persistence_manager()); BufferObj *buffer_obj = buffer_mgr->GetBufferObject(std::move(file_worker)); + buffer_obj->ToMmap(); chunk_index_entry->buffer_obj_ = buffer_obj; break; } From 3dfcceaadac7237f6b3450266901860073905809 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Fri, 27 Dec 2024 15:13:12 +0800 Subject: [PATCH 09/12] Not unmmap util cleanup. --- src/storage/buffer/buffer_obj.cpp | 7 ++++++- .../buffer/file_worker/file_worker.cpp | 21 ++++++++++++++++--- .../buffer/file_worker/file_worker.cppm | 2 ++ .../buffer/file_worker/hnsw_file_worker.cpp | 2 +- 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/src/storage/buffer/buffer_obj.cpp b/src/storage/buffer/buffer_obj.cpp index 2e9971eddb..a4f0e14401 100644 --- a/src/storage/buffer/buffer_obj.cpp +++ b/src/storage/buffer/buffer_obj.cpp @@ -227,6 +227,11 @@ void BufferObj::PickForCleanup() { LOG_INFO(fmt::format("BufferObj::PickForCleanup: obj_rc_ is {}, buffer: {}", obj_rc_, GetFilename())); return; } + if (type_ == BufferType::kMmap) { + file_worker_->Munmap(); + status_ = BufferStatus::kFreed; + return; + } switch (status_) { // when insert data into table with index, the index buffer_obj // will remain BufferStatus::kNew, so we should allow this situation @@ -336,7 +341,7 @@ void BufferObj::UnloadInner() { status_ = BufferStatus::kFreed; type_ = BufferType::kMmap; } else if (type_ == BufferType::kMmap) { - file_worker_->Munmap(); + file_worker_->MmapNotNeed(); status_ = BufferStatus::kFreed; } else { buffer_mgr_->PushGCQueue(this); diff --git a/src/storage/buffer/file_worker/file_worker.cpp b/src/storage/buffer/file_worker/file_worker.cpp index d3da5b1ddc..0eba628b67 100644 --- a/src/storage/buffer/file_worker/file_worker.cpp +++ b/src/storage/buffer/file_worker/file_worker.cpp @@ -14,9 +14,10 @@ module; -#include +#include #include -#include +#include +#include module file_worker; @@ -220,7 +221,7 @@ void FileWorker::CleanupTempFile() const { void FileWorker::Mmap() { if (mmap_addr_ != nullptr || mmap_data_ != nullptr) { - UnrecoverableError("Mmap already exists"); + return; } auto [defer_fn, read_path] = GetFilePathInner(false); bool use_object_cache = persistence_manager_ != nullptr; @@ -229,6 +230,10 @@ void FileWorker::Mmap() { if (ret < 0) { UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); } + // ret = madvise(mmap_addr_, file_size, MADV_DONTNEED); + // if (ret < 0) { + // UnrecoverableError(fmt::format("Madvise file {} failed. {}", read_path, strerror(errno))); + // } if (use_object_cache) { const void *ptr = mmap_addr_ + obj_addr_.part_offset_; this->ReadFromMmapImpl(ptr, obj_addr_.part_size_); @@ -246,4 +251,14 @@ void FileWorker::Munmap() { mmap_data_ = nullptr; } +void FileWorker::MmapNotNeed() { + // this->Munmap(); + // auto [defer_fn, read_path] = GetFilePathInner(false); + // SizeT file_size = VirtualStore::GetFileSize(read_path); + // int ret = madvise(mmap_addr_, file_size, MADV_PAGEOUT); + // if (ret < 0) { + // UnrecoverableError(fmt::format("Madvise file {} failed. {}", read_path, strerror(errno))); + // } +} + } // namespace infinity \ No newline at end of file diff --git a/src/storage/buffer/file_worker/file_worker.cppm b/src/storage/buffer/file_worker/file_worker.cppm index 7823eb7254..747210c771 100644 --- a/src/storage/buffer/file_worker/file_worker.cppm +++ b/src/storage/buffer/file_worker/file_worker.cppm @@ -103,6 +103,8 @@ public: void Munmap(); + void MmapNotNeed(); + protected: virtual bool ReadFromMmapImpl([[maybe_unused]] const void *ptr, [[maybe_unused]] SizeT size) { UnrecoverableError("Not implemented"); diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index 8edb327f44..e1adc257d6 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -115,7 +115,7 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const } else { using IndexT = std::decay_t; if constexpr (IndexT::kOwnMem) { - index->Save(*file_handle_); + index->SaveToPtr(*file_handle_); } else { UnrecoverableError("Invalid index type."); } From c9316e25b13ebda299b964db797d44107b361318 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Fri, 27 Dec 2024 15:44:29 +0800 Subject: [PATCH 10/12] Fix. --- src/storage/buffer/file_worker/file_worker.cpp | 5 ++++- .../knn_index/knn_hnsw/data_store/data_store.cppm | 8 +------- src/storage/knn_index/knn_hnsw/hnsw_alg.cppm | 10 ---------- 3 files changed, 5 insertions(+), 18 deletions(-) diff --git a/src/storage/buffer/file_worker/file_worker.cpp b/src/storage/buffer/file_worker/file_worker.cpp index 197f3a5101..ca592e5599 100644 --- a/src/storage/buffer/file_worker/file_worker.cpp +++ b/src/storage/buffer/file_worker/file_worker.cpp @@ -210,7 +210,7 @@ void FileWorker::CleanupTempFile() const { } void FileWorker::Mmap() { - if (mmap_addr_ != nullptr || mmap_data_ != nullptr) { + if (mmap_addr_ != nullptr) { return; } auto [defer_fn, read_path] = GetFilePathInner(false); @@ -234,6 +234,9 @@ void FileWorker::Mmap() { } void FileWorker::Munmap() { + if (mmap_addr_ == nullptr) { + return; + } auto [defer_fn, read_path] = GetFilePathInner(false); this->FreeFromMmapImpl(); VirtualStore::MunmapFile(read_path); 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 index 35e4d7ce63..2284d9d3ea 100644 --- a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm +++ b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm @@ -427,13 +427,11 @@ private: public: DataStore() = default; - DataStore(DataStore &&other) - : Base(std::move(other)), inner_(std::move(other.inner_)), ptr_start_(other.ptr_start_), cur_vec_num_(other.cur_vec_num_) {} + DataStore(DataStore &&other) : Base(std::move(other)), inner_(std::move(other.inner_)), cur_vec_num_(other.cur_vec_num_) {} DataStore &operator=(DataStore &&other) { if (this != &other) { Base::operator=(std::move(other)); inner_ = std::move(other.inner_); - ptr_start_ = other.ptr_start_; cur_vec_num_ = other.cur_vec_num_; } return *this; @@ -464,12 +462,8 @@ public: SizeT mem_usage() const { return 0; } - const char *ptr_start() const { return ptr_start_; } - void set_ptr_start(const char *ptr_start) { ptr_start_ = ptr_start; } - private: Inner inner_; - const char *ptr_start_ = nullptr; SizeT cur_vec_num_ = 0; public: diff --git a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm index 333a94a243..cd7f610fe7 100644 --- a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm +++ b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm @@ -524,7 +524,6 @@ public: public: static UniquePtr LoadFromPtr(const char *&ptr, SizeT size) { - const char *ptr_start = ptr; const char *ptr_end = ptr + size; SizeT M = ReadBufAdv(ptr); SizeT ef_construction = ReadBufAdv(ptr); @@ -533,17 +532,8 @@ public: if (SizeT diff = ptr_end - ptr; diff != 0) { UnrecoverableError("LoadFromPtr failed"); } - data_store.set_ptr_start(ptr_start); return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); } - - bool LoadAgain(const char *&ptr, SizeT size) { - if (this->data_store_.ptr_start() == ptr) { - return false; - } - *this = std::move(*LoadFromPtr(ptr, size)); - return true; - } }; } // namespace infinity \ No newline at end of file From abe72daf52812531fc21fbd98680b5f120e11033 Mon Sep 17 00:00:00 2001 From: shenyushi Date: Fri, 27 Dec 2024 16:18:16 +0800 Subject: [PATCH 11/12] Fix. --- src/storage/buffer/buffer_obj.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/storage/buffer/buffer_obj.cpp b/src/storage/buffer/buffer_obj.cpp index 4ff639897d..df1392abc1 100644 --- a/src/storage/buffer/buffer_obj.cpp +++ b/src/storage/buffer/buffer_obj.cpp @@ -229,7 +229,8 @@ void BufferObj::PickForCleanup() { } if (type_ == BufferType::kMmap) { file_worker_->Munmap(); - status_ = BufferStatus::kFreed; + buffer_mgr_->AddToCleanList(this, false /*do_free*/); + status_ = BufferStatus::kClean; return; } switch (status_) { From 02f98ad292fb1b7de90b2bc41301ff8a9de31acc Mon Sep 17 00:00:00 2001 From: shenyushi Date: Fri, 27 Dec 2024 20:18:07 +0800 Subject: [PATCH 12/12] Fix. --- .../buffer/file_worker/file_worker.cpp | 43 ++++++++----------- src/storage/io/virtual_store.cpp | 24 +++++++++++ src/storage/io/virtual_store.cppm | 3 ++ src/storage/meta/entry/chunk_index_entry.cpp | 3 ++ 4 files changed, 49 insertions(+), 24 deletions(-) diff --git a/src/storage/buffer/file_worker/file_worker.cpp b/src/storage/buffer/file_worker/file_worker.cpp index ca592e5599..00922bf0e6 100644 --- a/src/storage/buffer/file_worker/file_worker.cpp +++ b/src/storage/buffer/file_worker/file_worker.cpp @@ -215,21 +215,19 @@ void FileWorker::Mmap() { } auto [defer_fn, read_path] = GetFilePathInner(false); bool use_object_cache = persistence_manager_ != nullptr; - SizeT file_size = VirtualStore::GetFileSize(read_path); - int ret = VirtualStore::MmapFile(read_path, mmap_addr_, file_size); - if (ret < 0) { - UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); - } - // ret = madvise(mmap_addr_, file_size, MADV_DONTNEED); - // if (ret < 0) { - // UnrecoverableError(fmt::format("Madvise file {} failed. {}", read_path, strerror(errno))); - // } if (use_object_cache) { - const void *ptr = mmap_addr_ + obj_addr_.part_offset_; - this->ReadFromMmapImpl(ptr, obj_addr_.part_size_); + int ret = VirtualStore::MmapFilePart(read_path, obj_addr_.part_offset_, obj_addr_.part_size_, mmap_addr_); + if (ret < 0) { + UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); + } + this->ReadFromMmapImpl(mmap_addr_, obj_addr_.part_size_); } else { - const void *ptr = mmap_addr_; - this->ReadFromMmapImpl(ptr, file_size); + SizeT file_size = VirtualStore::GetFileSize(read_path); + int ret = VirtualStore::MmapFile(read_path, mmap_addr_, file_size); + if (ret < 0) { + UnrecoverableError(fmt::format("Mmap file {} failed. {}", read_path, strerror(errno))); + } + this->ReadFromMmapImpl(mmap_addr_, file_size); } } @@ -237,21 +235,18 @@ void FileWorker::Munmap() { if (mmap_addr_ == nullptr) { return; } - auto [defer_fn, read_path] = GetFilePathInner(false); this->FreeFromMmapImpl(); - VirtualStore::MunmapFile(read_path); + auto [defer_fn, read_path] = GetFilePathInner(false); + bool use_object_cache = persistence_manager_ != nullptr; + if (use_object_cache) { + VirtualStore::MunmapFilePart(mmap_addr_, obj_addr_.part_offset_, obj_addr_.part_size_); + } else { + VirtualStore::MunmapFile(read_path); + } mmap_addr_ = nullptr; mmap_data_ = nullptr; } -void FileWorker::MmapNotNeed() { - // this->Munmap(); - // auto [defer_fn, read_path] = GetFilePathInner(false); - // SizeT file_size = VirtualStore::GetFileSize(read_path); - // int ret = madvise(mmap_addr_, file_size, MADV_PAGEOUT); - // if (ret < 0) { - // UnrecoverableError(fmt::format("Madvise file {} failed. {}", read_path, strerror(errno))); - // } -} +void FileWorker::MmapNotNeed() {} } // namespace infinity \ No newline at end of file diff --git a/src/storage/io/virtual_store.cpp b/src/storage/io/virtual_store.cpp index 4e522b39ed..0267eca758 100644 --- a/src/storage/io/virtual_store.cpp +++ b/src/storage/io/virtual_store.cpp @@ -446,6 +446,30 @@ i32 VirtualStore::MunmapFile(const String &file_path) { return 0; } +i32 VirtualStore::MmapFilePart(const String &file_path, SizeT offset, SizeT length, u8 *&data_ptr) { + SizeT align_offset = offset / getpagesize() * getpagesize(); + SizeT diff = offset - align_offset; + SizeT align_length = length + diff; + int fd = open(file_path.c_str(), O_RDONLY); + if (fd < 0) { + return -1; + } + void *ret = mmap(NULL, align_length, PROT_READ, MAP_SHARED, fd, align_offset); + close(fd); + if (ret == MAP_FAILED) { + return -1; + } + data_ptr = reinterpret_cast(ret) + diff; + return 0; +} + +i32 VirtualStore::MunmapFilePart(u8 *data_ptr, SizeT offset, SizeT length) { + SizeT align_offset = offset / getpagesize() * getpagesize(); + u8 *aligned_ptr = data_ptr - offset + align_offset; + SizeT align_length = length + data_ptr - aligned_ptr; + return munmap(aligned_ptr, align_length); +} + // Remote storage StorageType VirtualStore::storage_type_ = StorageType::kInvalid; String VirtualStore::bucket_ = "infinity"; diff --git a/src/storage/io/virtual_store.cppm b/src/storage/io/virtual_store.cppm index 643a14fa6b..1a516b1c38 100644 --- a/src/storage/io/virtual_store.cppm +++ b/src/storage/io/virtual_store.cppm @@ -72,6 +72,9 @@ public: static i32 MmapFile(const String &file_path, u8 *&data_ptr, SizeT &data_len); static i32 MunmapFile(const String &file_path); + static i32 MmapFilePart(const String &file_path, SizeT offset, SizeT length, u8 *&data_ptr); + static i32 MunmapFilePart(u8 *data_ptr, SizeT offset, SizeT length); + static Status InitRemoteStore(StorageType storage_type = StorageType::kMinio, const String &URL = "http://localhost:9000", bool HTTPS = false, diff --git a/src/storage/meta/entry/chunk_index_entry.cpp b/src/storage/meta/entry/chunk_index_entry.cpp index e8c1ff93aa..bd62a3dfc6 100644 --- a/src/storage/meta/entry/chunk_index_entry.cpp +++ b/src/storage/meta/entry/chunk_index_entry.cpp @@ -456,6 +456,9 @@ void ChunkIndexEntry::SaveIndexFile() { return; } buffer_obj_->Save(); + if (segment_index_entry_->table_index_entry()->index_base()->index_type_ == IndexType::kHnsw) { + buffer_obj_->ToMmap(); + } } void ChunkIndexEntry::DeprecateChunk(TxnTimeStamp commit_ts) {