diff --git a/src/storage/buffer/file_worker/hnsw_file_worker.cpp b/src/storage/buffer/file_worker/hnsw_file_worker.cpp index 2a5e66e0de..c5c68c1013 100644 --- a/src/storage/buffer/file_worker/hnsw_file_worker.cpp +++ b/src/storage/buffer/file_worker/hnsw_file_worker.cpp @@ -140,7 +140,11 @@ void HnswFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) { } else { using IndexT = std::decay_t; if constexpr (IndexT::kOwnMem) { - index = IndexT::Load(*file_handle_).release(); + if (from_spill) { + index = IndexT::Load(*file_handle_).release(); + } else { + index = IndexT::LoadFromPtr(*file_handle_, file_size).release(); + } } else { UnrecoverableError("Invalid index type."); } diff --git a/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/data_store.cppm index 2284d9d3ea..692e50da26 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 @@ -215,6 +215,24 @@ public: return ret; } + 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 chunk_size = 1; + while (chunk_size < cur_vec_num) { + chunk_size <<= 1; + } + This ret = This(chunk_size, 1 /*max_chunk_n*/, std::move(vec_store_meta), std::move(graph_store_meta)); + ret.cur_vec_num_ = cur_vec_num; + + SizeT mem_usage = 0; + ret.inners_[0] = Inner::LoadFromPtr(ptr, cur_vec_num, chunk_size, ret.vec_store_meta_, ret.graph_store_meta_, mem_usage); + ret.mem_usage_.store(mem_usage); + return ret; + } + void SetGraph(GraphStoreMeta &&graph_meta, Vector> &&graph_inners) { this->graph_store_meta_ = std::move(graph_meta); for (SizeT i = 0; i < graph_inners.size(); ++i) { @@ -615,6 +633,20 @@ public: return ret; } + static This LoadFromPtr(const char *&ptr, + SizeT cur_vec_num, + SizeT chunk_size, + VecStoreMeta &vec_store_meta, + GraphStoreMeta &graph_store_meta, + SizeT &mem_usage) { + auto vec_store_inner = VecStoreInner::LoadFromPtr(ptr, cur_vec_num, chunk_size, vec_store_meta, mem_usage); + auto graph_store_inner = GraphStoreInner::LoadFromPtr(ptr, cur_vec_num, chunk_size, graph_store_meta, mem_usage); + This ret(chunk_size, std::move(vec_store_inner), std::move(graph_store_inner)); + std::memcpy(ret.labels_.get(), ptr, sizeof(LabelType) * cur_vec_num); + ptr += sizeof(LabelType) * cur_vec_num; + return ret; + } + // vec store template Iterator> Pair AddVec(Iterator &&query_iter, VertexType start_idx, SizeT remain_num, const VecStoreMeta &meta, SizeT &mem_usage) { diff --git a/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/graph_store.cppm index 60c4b0c8b1..67439a2303 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 @@ -383,6 +383,32 @@ public: return graph_store; } + static GraphStoreInner LoadFromPtr(const char *&ptr, SizeT cur_vertex_n, SizeT max_vertex, const GraphStoreMeta &meta, SizeT &mem_usage) { + SizeT layer_sum = ReadBufAdv(ptr); + GraphStoreInner graph_store(max_vertex, meta, cur_vertex_n); + const char *graph = ptr; + ptr += cur_vertex_n * meta.level0_size(); + std::memcpy(graph_store.graph_.get(), graph, cur_vertex_n * meta.level0_size()); + + auto loaded_layers = MakeUnique(layer_sum * meta.levelx_size()); + char *loaded_layers_p = loaded_layers.get(); + for (VertexType vertex_i = 0; vertex_i < (VertexType)cur_vertex_n; ++vertex_i) { + VertexL0 *v = graph_store.GetLevel0(vertex_i, meta); + if (v->layer_n_) { + std::memcpy(loaded_layers_p, ptr, meta.levelx_size() * v->layer_n_); + v->layers_p_ = loaded_layers_p; + loaded_layers_p += meta.levelx_size() * v->layer_n_; + ptr += meta.levelx_size() * v->layer_n_; + } else { + v->layers_p_ = nullptr; + } + } + graph_store.loaded_layers_ = std::move(loaded_layers); + + mem_usage += max_vertex * meta.level0_size() + layer_sum * meta.levelx_size(); + return graph_store; + } + void AddVertex(VertexType vertex_i, i32 layer_n, const GraphStoreMeta &meta, SizeT &mem_usage) { VertexL0 *v = GetLevel0(vertex_i, meta); v->neighbor_n_ = 0; diff --git a/src/storage/knn_index/knn_hnsw/data_store/lvq_vec_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/lvq_vec_store.cppm index af9bd18f48..c0aeabcb9f 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 @@ -230,6 +230,18 @@ public: return meta; } + static This LoadFromPtr(const char *&ptr) { + SizeT dim = ReadBufAdv(ptr); + This meta(dim); + std::memcpy(meta.mean_.get(), ptr, sizeof(MeanType) * dim); + ptr += sizeof(MeanType) * dim; + if constexpr (!std::is_same_v>) { + std::memcpy(&meta.global_cache_, ptr, sizeof(GlobalCacheType)); + ptr += sizeof(GlobalCacheType); + } + return meta; + } + template Iterator> void Optimize(Iterator &&query_iter, const Vector> &inners, SizeT &mem_usage) { auto new_mean = MakeUnique(this->dim_); @@ -378,6 +390,14 @@ public: return ret; } + static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { + This ret(max_vec_num, meta); + std::memcpy(ret.ptr_.get(), ptr, cur_vec_num * meta.compress_data_size()); + ptr += cur_vec_num * meta.compress_data_size(); + mem_usage += max_vec_num * meta.compress_data_size(); + return ret; + } + void SetVec(SizeT idx, const DataType *vec, const Meta &meta, SizeT &mem_usage) { meta.CompressTo(vec, GetVecMut(idx, meta)); } private: diff --git a/src/storage/knn_index/knn_hnsw/data_store/plain_vec_store.cppm b/src/storage/knn_index/knn_hnsw/data_store/plain_vec_store.cppm index cda6e3918a..2aed9e35b6 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 @@ -152,6 +152,14 @@ public: return ret; } + static This LoadFromPtr(const char *&ptr, SizeT cur_vec_num, SizeT max_vec_num, const Meta &meta, SizeT &mem_usage) { + This ret(max_vec_num, meta); + std::memcpy(ret.ptr_.get(), ptr, sizeof(DataType) * cur_vec_num * meta.dim()); + ptr += sizeof(DataType) * cur_vec_num * meta.dim(); + mem_usage += sizeof(DataType) * max_vec_num * meta.dim(); + return ret; + } + void SetVec(SizeT idx, const DataType *vec, const Meta &meta, SizeT &mem_usage) { Copy(vec, vec + meta.dim(), GetVecMut(idx, meta)); } private: diff --git a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm index cd7f610fe7..40045fc9fd 100644 --- a/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm +++ b/src/storage/knn_index/knn_hnsw/hnsw_alg.cppm @@ -484,6 +484,20 @@ public: return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); } + static UniquePtr LoadFromPtr(LocalFileHandle &file_handle, SizeT size) { + auto buffer = MakeUnique(size); + file_handle.Read(buffer.get(), size); + const char *ptr = buffer.get(); + SizeT M = ReadBufAdv(ptr); + SizeT ef_construction = ReadBufAdv(ptr); + auto data_store = DataStore::LoadFromPtr(ptr); + Distance distance(data_store.dim()); + if (SizeT diff = ptr - buffer.get(); diff != size) { + UnrecoverableError("LoadFromPtr failed"); + } + return MakeUnique(M, ef_construction, std::move(data_store), std::move(distance)); + } + UniquePtr> CompressToLVQ() && { if constexpr (std::is_same_v) { return MakeUnique(std::move(*this)); diff --git a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp index 2d9d7262ca..417ed6acaa 100644 --- a/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp +++ b/src/unit_test/storage/knnindex/knn_hnsw/test_hnsw.cpp @@ -178,6 +178,33 @@ class HnswAlgTest : public BaseTest { test_func(hnsw_index); +#ifdef USE_MMAP + VirtualStore::MunmapFile(filepath); +#endif + } + { + SizeT file_size = VirtualStore::GetFileSize(filepath); +#define USE_MMAP +#ifdef USE_MMAP + unsigned char *data_ptr = nullptr; + int ret = VirtualStore::MmapFile(filepath, data_ptr, file_size); + if (ret < 0) { + UnrecoverableError("mmap failed"); + } + const char *ptr = reinterpret_cast(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 = Hnsw::LoadFromPtr(ptr, file_size); + + test_func(hnsw_index); + #ifdef USE_MMAP VirtualStore::MunmapFile(filepath); #endif diff --git a/test/sql/dql/knn/embedding/test_knn_hnsw_l2_lvq.slt b/test/sql/dql/knn/embedding/test_knn_hnsw_l2_lvq.slt index 41ba44497b..3921d46b42 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 895f6efb56..20112d82ce 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);