diff --git a/include/vsag/constants.h b/include/vsag/constants.h index cbc68825..eb6cbe7c 100644 --- a/include/vsag/constants.h +++ b/include/vsag/constants.h @@ -59,6 +59,12 @@ extern const char* const DISKANN_PARAMETER_USE_REFERENCE; extern const char* const DISKANN_PARAMETER_USE_OPQ; extern const char* const DISKANN_PARAMETER_USE_ASYNC_IO; extern const char* const DISKANN_PARAMETER_USE_BSA; +extern const char* const DISKANN_PARAMETER_GRAPH_TYPE; +extern const char* const DISKANN_PARAMETER_ALPHA; +extern const char* const DISKANN_PARAMETER_TURN; +extern const char* const DISKANN_PARAMETER_SAMPLE_RATE; +extern const char* const DISKANN_GRAPH_TYPE_VAMANA; +extern const char* const DISKANN_GRAPH_TYPE_ODESCENT; extern const char* const DISKANN_PARAMETER_BEAM_SEARCH; extern const char* const DISKANN_PARAMETER_IO_LIMIT; diff --git a/src/allocator_wrapper.h b/src/allocator_wrapper.h index 5479aae8..cfdb38df 100644 --- a/src/allocator_wrapper.h +++ b/src/allocator_wrapper.h @@ -42,6 +42,11 @@ class AllocatorWrapper { return allocator_ == other.allocator_; } + bool + operator!=(const AllocatorWrapper& other) const noexcept { + return allocator_ != other.allocator_; + } + inline pointer allocate(size_type n, const_void_pointer hint = 0) { return static_cast(allocator_->Allocate(n * sizeof(value_type))); diff --git a/src/constants.cpp b/src/constants.cpp index d791cc52..2bfa2a11 100644 --- a/src/constants.cpp +++ b/src/constants.cpp @@ -66,6 +66,13 @@ const char* const DISKANN_PARAMETER_BEAM_SEARCH = "beam_search"; const char* const DISKANN_PARAMETER_IO_LIMIT = "io_limit"; const char* const DISKANN_PARAMETER_EF_SEARCH = "ef_search"; const char* const DISKANN_PARAMETER_REORDER = "use_reorder"; +const char* const DISKANN_PARAMETER_GRAPH_TYPE = "graph_type"; +const char* const DISKANN_PARAMETER_ALPHA = "alpha"; +const char* const DISKANN_PARAMETER_TURN = "turn"; +const char* const DISKANN_PARAMETER_SAMPLE_RATE = "sample_rate"; + +const char* const DISKANN_GRAPH_TYPE_VAMANA = "vamana"; +const char* const DISKANN_GRAPH_TYPE_ODESCENT = "odescent"; const char* const HNSW_PARAMETER_EF_RUNTIME = "ef_search"; const char* const HNSW_PARAMETER_M = "max_degree"; diff --git a/src/impl/odescent_graph_builder.cpp b/src/impl/odescent_graph_builder.cpp new file mode 100644 index 00000000..3f4f50e1 --- /dev/null +++ b/src/impl/odescent_graph_builder.cpp @@ -0,0 +1,377 @@ + +// Copyright 2024-present the vsag project +// +// 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 +// +// http://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. + +#include "odescent_graph_builder.h" + +#include + +namespace vsag { + +class LCG { +public: + LCG() { + auto now = std::chrono::steady_clock::now(); + auto timestamp = + std::chrono::duration_cast(now.time_since_epoch()).count(); + current_ = static_cast(timestamp); + } + + float + NextFloat() { + current_ = (A * current_ + C) % M; + return static_cast(current_) / static_cast(M); + } + +private: + unsigned int current_; + static const uint32_t A = 1664525; + static const uint32_t C = 1013904223; + static const uint32_t M = 4294967295; // 2^32 - 1 +}; + +bool +Odescent::Build(const DatasetPtr& dataset) { + if (is_build_) { + return false; + } + is_build_ = true; + dim_ = dataset->GetDim(); + data_num_ = dataset->GetNumElements(); + data_ = dataset->GetFloat32Vectors(); + min_in_degree_ = std::min(min_in_degree_, data_num_ - 1); + Vector(data_num_, allocator_).swap(points_lock_); + Vector> old_neighbors(allocator_); + Vector> new_neighbors(allocator_); + old_neighbors.resize(data_num_, UnorderedSet(allocator_)); + new_neighbors.resize(data_num_, UnorderedSet(allocator_)); + for (int i = 0; i < data_num_; ++i) { + old_neighbors[i].reserve(max_degree_); + new_neighbors[i].reserve(max_degree_); + } + init_graph(); + { + for (int i = 0; i < turn_; ++i) { + sample_candidates(old_neighbors, new_neighbors, sample_rate_); + update_neighbors(old_neighbors, new_neighbors); + repair_no_in_edge(); + } + if (pruning_) { + prune_graph(); + add_reverse_edges(); + } + } + return true; +} + +void +Odescent::SaveGraph(std::stringstream& out) { + size_t file_offset = 0; // we will use this if we want + out.seekp(file_offset, out.beg); + size_t index_size = 24; + uint32_t max_degree = 0; + out.write((char*)&index_size, sizeof(uint64_t)); + out.write((char*)&max_degree, sizeof(uint32_t)); + uint32_t ep_u32 = 0; + size_t num_frozen = 0; + out.write((char*)&ep_u32, sizeof(uint32_t)); + out.write((char*)&num_frozen, sizeof(size_t)); + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. + auto final_graph = GetGraph(); + for (uint32_t i = 0; i < data_num_; i++) { + uint32_t gk = (uint32_t)final_graph[i].size(); + out.write((char*)&gk, sizeof(uint32_t)); + out.write((char*)final_graph[i].data(), gk * sizeof(uint32_t)); + max_degree = + final_graph[i].size() > max_degree ? (uint32_t)final_graph[i].size() : max_degree; + index_size += (size_t)(sizeof(uint32_t) * (gk + 1)); + } + out.seekp(file_offset, out.beg); + out.write((char*)&index_size, sizeof(uint64_t)); + out.write((char*)&max_degree, sizeof(uint32_t)); +} + +Vector> +Odescent::GetGraph() { + Vector> extract_graph(allocator_); + extract_graph.resize(data_num_, Vector(allocator_)); + for (int i = 0; i < data_num_; ++i) { + extract_graph[i].resize(graph[i].neighbors.size()); + for (int j = 0; j < graph[i].neighbors.size(); ++j) { + extract_graph[i][j] = graph[i].neighbors[j].id; + } + } + + return extract_graph; +} + +void +Odescent::init_graph() { + graph.resize(data_num_, Linklist(allocator_)); + + auto task = [&, this](int64_t start, int64_t end) { + std::random_device rd; + std::uniform_int_distribution k_generate(0, data_num_ - 1); + std::mt19937 rng(rd()); + for (int i = start; i < end; ++i) { + UnorderedSet ids_set(allocator_); + ids_set.insert(i); + graph[i].neighbors.reserve(max_degree_); + int max_neighbors = std::min(data_num_ - 1, max_degree_); + for (int j = 0; j < max_neighbors; ++j) { + uint32_t id = i; + if (data_num_ - 1 < max_degree_) { + id = (i + j + 1) % data_num_; + } else { + while (ids_set.find(id) != ids_set.end()) { + id = k_generate(rng); + } + } + ids_set.insert(id); + auto dist = get_distance(i, id); + graph[i].neighbors.emplace_back(id, dist); + graph[i].greast_neighbor_distance = + std::max(graph[i].greast_neighbor_distance, dist); + } + } + }; + parallelize_task(task); +} + +void +Odescent::update_neighbors(Vector>& old_neighbors, + Vector>& new_neighbors) { + Vector> futures(allocator_); + auto task = [&, this](int64_t start, int64_t end) { + for (int i = start; i < end; ++i) { + Vector new_neighbors_candidates(allocator_); + for (uint32_t node_id : new_neighbors[i]) { + for (int k = 0; k < new_neighbors_candidates.size(); ++k) { + auto neighbor_id = new_neighbors_candidates[k]; + float dist = get_distance(node_id, neighbor_id); + if (dist < graph[node_id].greast_neighbor_distance) { + std::lock_guard lock(points_lock_[node_id]); + graph[node_id].neighbors.emplace_back(neighbor_id, dist); + } + if (dist < graph[neighbor_id].greast_neighbor_distance) { + std::lock_guard lock(points_lock_[neighbor_id]); + graph[neighbor_id].neighbors.emplace_back(node_id, dist); + } + } + new_neighbors_candidates.push_back(node_id); + + for (uint32_t neighbor_id : old_neighbors[i]) { + if (node_id == neighbor_id) { + continue; + } + float dist = get_distance(neighbor_id, node_id); + if (dist < graph[node_id].greast_neighbor_distance) { + std::lock_guard lock(points_lock_[node_id]); + graph[node_id].neighbors.emplace_back(neighbor_id, dist); + } + if (dist < graph[neighbor_id].greast_neighbor_distance) { + std::lock_guard lock(points_lock_[neighbor_id]); + graph[neighbor_id].neighbors.emplace_back(node_id, dist); + } + } + } + old_neighbors[i].clear(); + new_neighbors[i].clear(); + } + }; + parallelize_task(task); + + auto resize_task = [&, this](int64_t start, int64_t end) { + for (uint32_t i = start; i < end; ++i) { + auto& neighbors = graph[i].neighbors; + std::sort(neighbors.begin(), neighbors.end()); + neighbors.erase(std::unique(neighbors.begin(), neighbors.end()), neighbors.end()); + if (neighbors.size() > max_degree_) { + neighbors.resize(max_degree_); + } + graph[i].greast_neighbor_distance = neighbors.back().distance; + } + }; + parallelize_task(resize_task); +} + +void +Odescent::add_reverse_edges() { + Vector reverse_graph(allocator_); + reverse_graph.resize(data_num_, Linklist(allocator_)); + for (int i = 0; i < data_num_; ++i) { + reverse_graph[i].neighbors.reserve(max_degree_); + } + for (int i = 0; i < data_num_; ++i) { + for (const auto& node : graph[i].neighbors) { + reverse_graph[node.id].neighbors.emplace_back(i, node.distance); + } + } + + auto task = [&, this](int64_t start, int64_t end) { + for (int i = start; i < end; ++i) { + auto& neighbors = graph[i].neighbors; + neighbors.insert(neighbors.end(), + reverse_graph[i].neighbors.begin(), + reverse_graph[i].neighbors.end()); + std::sort(neighbors.begin(), neighbors.end()); + neighbors.erase(std::unique(neighbors.begin(), neighbors.end()), neighbors.end()); + if (neighbors.size() > max_degree_) { + neighbors.resize(max_degree_); + } + } + }; + parallelize_task(task); +} + +void +Odescent::sample_candidates(Vector>& old_neighbors, + Vector>& new_neighbors, + float sample_rate) { + auto task = [&, this](int64_t start, int64_t end) { + LCG r; + for (int i = start; i < end; ++i) { + auto& neighbors = graph[i].neighbors; + for (int j = 0; j < neighbors.size(); ++j) { + float current_state = r.NextFloat(); + if (current_state < sample_rate) { + if (neighbors[j].old) { + { + std::lock_guard lock(points_lock_[i]); + old_neighbors[i].insert(neighbors[j].id); + } + { + std::lock_guard inner_lock(points_lock_[neighbors[j].id]); + old_neighbors[neighbors[j].id].insert(i); + } + } else { + { + std::lock_guard lock(points_lock_[i]); + new_neighbors[i].insert(neighbors[j].id); + } + { + std::lock_guard inner_lock(points_lock_[neighbors[j].id]); + new_neighbors[neighbors[j].id].insert(i); + } + neighbors[j].old = true; + } + } + } + } + }; + parallelize_task(task); +} + +void +Odescent::repair_no_in_edge() { + Vector in_edges_count(data_num_, 0, allocator_); + for (int i = 0; i < data_num_; ++i) { + for (auto& neigbor : graph[i].neighbors) { + in_edges_count[neigbor.id]++; + } + } + + Vector replace_pos(data_num_, std::min(data_num_ - 1, max_degree_) - 1, allocator_); + for (int i = 0; i < data_num_; ++i) { + auto& link = graph[i].neighbors; + int need_replace_loc = 0; + while (in_edges_count[i] < min_in_degree_ && need_replace_loc < max_degree_) { + uint32_t need_replace_id = link[need_replace_loc].id; + bool has_connect = false; + for (auto& neigbor : graph[need_replace_id].neighbors) { + if (neigbor.id == i) { + has_connect = true; + break; + } + } + if (replace_pos[need_replace_id] > 0 && not has_connect) { + auto& replace_node = graph[need_replace_id].neighbors[replace_pos[need_replace_id]]; + auto replace_id = replace_node.id; + if (in_edges_count[replace_id] > min_in_degree_) { + in_edges_count[replace_id]--; + replace_node.id = i; + replace_node.distance = link[need_replace_loc].distance; + in_edges_count[i]++; + } + replace_pos[need_replace_id]--; + } + need_replace_loc++; + } + } +} + +void +Odescent::prune_graph() { + Vector in_edges_count(data_num_, 0, allocator_); + for (int i = 0; i < data_num_; ++i) { + for (int j = 0; j < graph[i].neighbors.size(); ++j) { + in_edges_count[graph[i].neighbors[j].id]++; + } + } + + auto task = [&, this](int64_t start, int64_t end) { + for (int loc = start; loc < end; ++loc) { + auto& neighbors = graph[loc].neighbors; + std::sort(neighbors.begin(), neighbors.end()); + neighbors.erase(std::unique(neighbors.begin(), neighbors.end()), neighbors.end()); + Vector candidates(allocator_); + candidates.reserve(max_degree_); + for (int i = 0; i < neighbors.size(); ++i) { + bool flag = true; + int cur_in_edge = 0; + { + std::lock_guard lock(points_lock_[neighbors[i].id]); + cur_in_edge = in_edges_count[neighbors[i].id]; + } + if (cur_in_edge > min_in_degree_) { + for (int j = 0; j < candidates.size(); ++j) { + if (get_distance(neighbors[i].id, candidates[j].id) * alpha_ < + neighbors[i].distance) { + flag = false; + { + std::lock_guard lock(points_lock_[neighbors[i].id]); + in_edges_count[neighbors[i].id]--; + } + break; + } + } + } + if (flag) { + candidates.push_back(neighbors[i]); + } + } + neighbors.swap(candidates); + if (neighbors.size() > max_degree_) { + neighbors.resize(max_degree_); + } + } + }; + parallelize_task(task); +} + +void +Odescent::parallelize_task(std::function task) { + Vector> futures(allocator_); + for (int64_t i = 0; i < data_num_; i += block_size_) { + int end = std::min(i + block_size_, data_num_); + futures.push_back(thread_pool_->GeneralEnqueue(task, i, end)); + } + for (auto& future : futures) { + future.get(); + } +} + +} // namespace vsag \ No newline at end of file diff --git a/src/impl/odescent_graph_builder.h b/src/impl/odescent_graph_builder.h new file mode 100644 index 00000000..0e941a65 --- /dev/null +++ b/src/impl/odescent_graph_builder.h @@ -0,0 +1,156 @@ + +// Copyright 2024-present the vsag project +// +// 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 +// +// http://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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "data_cell/flatten_datacell.h" +#include "logger.h" +#include "safe_allocator.h" +#include "simd/simd.h" +#include "utils.h" +#include "vsag/dataset.h" + +namespace vsag { + +struct Node { + bool old = false; + uint32_t id; + float distance; + + Node(uint32_t id, float distance) { + this->id = id; + this->distance = distance; + } + + Node(uint32_t id, float distance, bool old) { + this->id = id; + this->distance = distance; + this->old = old; + } + Node() { + } + + bool + operator<(const Node& other) const { + if (distance != other.distance) { + return distance < other.distance; + } + return old && not other.old; + } + + bool + operator==(const Node& other) const { + return id == other.id; + } +}; + +struct Linklist { + Vector neighbors; + float greast_neighbor_distance; + Linklist(Allocator* allocator) + : neighbors(allocator), greast_neighbor_distance(std::numeric_limits::max()) { + } +}; + +class Odescent { +public: + Odescent(int64_t max_degree, + float alpha, + int64_t turn, + float sample_rate, + FlattenInterfacePtr flatten_interface, + Allocator* allocator, + SafeThreadPool* thread_pool, + bool pruning = true) + : max_degree_(max_degree), + alpha_(alpha), + turn_(turn), + sample_rate_(sample_rate), + flatten_interface_(flatten_interface), + pruning_(pruning), + allocator_(allocator), + graph(allocator), + points_lock_(allocator), + thread_pool_(thread_pool) { + } + + bool + Build(const DatasetPtr& dataset); + + void + SaveGraph(std::stringstream& out); + + Vector> + GetGraph(); + +private: + inline float + get_distance(uint32_t loc1, uint32_t loc2) { + return flatten_interface_->ComputePairVectors(loc1, loc2); + } + + void + init_graph(); + + void + update_neighbors(Vector>& old_neigbors, + Vector>& new_neigbors); + + void + add_reverse_edges(); + + void + sample_candidates(Vector>& old_neigbors, + Vector>& new_neigbors, + float sample_rate); + + void + repair_no_in_edge(); + + void + prune_graph(); + +private: + void + parallelize_task(std::function task); + + size_t dim_; + int64_t data_num_; + int64_t is_build_ = false; + const float* data_; + + int64_t max_degree_; + float alpha_; + int64_t turn_; + Vector graph; + int64_t min_in_degree_ = 1; + int64_t block_size_{10000}; + Vector points_lock_; + SafeThreadPool* thread_pool_; + + bool pruning_{true}; + float sample_rate_{0.3}; + Allocator* allocator_; + + FlattenInterfacePtr flatten_interface_; +}; + +} // namespace vsag diff --git a/src/impl/odescent_graph_builder_test.cpp b/src/impl/odescent_graph_builder_test.cpp new file mode 100644 index 00000000..f8643d35 --- /dev/null +++ b/src/impl/odescent_graph_builder_test.cpp @@ -0,0 +1,117 @@ + +// Copyright 2024-present the vsag project +// +// 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 +// +// http://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. + +#include "odescent_graph_builder.h" + +#include +#include +#include +#include +#include + +#include "data_cell/flatten_interface.h" +#include "safe_allocator.h" +#include "typing.h" + +const static std::string FLATTEN_PARAM = R"( + { + "io_type": "memory", + "io_params": {}, + "codes_type": "flatten_codes", + "quantization_type": "fp32", + "quantization_params": {} + } +)"; + +size_t +calculate_overlap(const std::vector& vec1, vsag::Vector& vec2, int K) { + int size1 = std::min(K, static_cast(vec1.size())); + int size2 = std::min(K, static_cast(vec2.size())); + + std::vector top_k_vec1(vec1.begin(), vec1.begin() + size1); + std::vector top_k_vec2(vec2.begin(), vec2.begin() + size2); + + std::sort(top_k_vec1.rbegin(), top_k_vec1.rend()); + std::sort(top_k_vec2.rbegin(), top_k_vec2.rend()); + + std::set set1(top_k_vec1.begin(), top_k_vec1.end()); + std::set set2(top_k_vec2.begin(), top_k_vec2.end()); + + std::set intersection; + std::set_intersection(set1.begin(), + set1.end(), + set2.begin(), + set2.end(), + std::inserter(intersection, intersection.begin())); + return intersection.size(); +} + +TEST_CASE("build nndescent", "[ut][nndescent]") { + int64_t num_vectors = 2000; + size_t dim = 128; + int64_t max_degree = 32; + + auto vectors = new float[dim * num_vectors]; + + std::mt19937 rng; + rng.seed(47); + std::uniform_real_distribution<> distrib_real; + for (int64_t i = 0; i < dim * num_vectors; ++i) { + vectors[i] = distrib_real(rng); + } + + std::vector>> ground_truths(num_vectors); + vsag::IndexCommonParam param; + param.dim_ = dim; + param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + param.thread_pool_ = vsag::SafeThreadPool::FactoryDefaultThreadPool(); + vsag::JsonType type = vsag::JsonType::parse(FLATTEN_PARAM); + vsag::FlattenInterfacePtr flatten_interface_ptr = + vsag::FlattenInterface::MakeInstance(type, param); + flatten_interface_ptr->Train(vectors, num_vectors); + flatten_interface_ptr->BatchInsertVector(vectors, num_vectors); + + vsag::DatasetPtr dataset = vsag::Dataset::Make(); + dataset->NumElements(num_vectors)->Float32Vectors(vectors)->Dim(dim)->Owner(true); + vsag::Odescent graph(max_degree, + 1, + 30, + 0.3, + flatten_interface_ptr, + param.allocator_.get(), + param.thread_pool_.get(), + false); + graph.Build(dataset); + + auto extract_graph = graph.GetGraph(); + + float hit_edge_count = 0; + for (int i = 0; i < num_vectors; ++i) { + for (int j = 0; j < num_vectors; ++j) { + if (i != j) { + ground_truths[i].emplace_back(flatten_interface_ptr->ComputePairVectors(i, j), j); + } + } + std::sort(ground_truths[i].begin(), ground_truths[i].end()); + std::vector truths_edges(max_degree); + for (int j = 0; j < max_degree; ++j) { + truths_edges[j] = ground_truths[i][j].second; + } + hit_edge_count += calculate_overlap(truths_edges, extract_graph[i], max_degree); + } + REQUIRE(hit_edge_count / (num_vectors * max_degree) > 0.95); +} \ No newline at end of file diff --git a/src/index/diskann.cpp b/src/index/diskann.cpp index ed2c9149..dedb841f 100644 --- a/src/index/diskann.cpp +++ b/src/index/diskann.cpp @@ -27,6 +27,8 @@ #include #include +#include "data_cell/flatten_datacell.h" +#include "impl/odescent_graph_builder.h" #include "vsag/constants.h" #include "vsag/errors.h" #include "vsag/expected.hpp" @@ -49,6 +51,16 @@ const static std::string BUILD_CURRENT_ROUND = "round"; const static std::string BUILD_NODES = "builded_nodes"; const static std::string BUILD_FAILED_LOC = "failed_loc"; +const static std::string FLATTEN_PARAM = R"( + { + "io_type": "memory", + "io_params": {}, + "codes_type": "flatten_codes", + "quantization_type": "sq8", + "quantization_params": {} + } +)"; + template Binary to_binary(T& value) { @@ -145,7 +157,8 @@ DiskANN::DiskANN(DiskannParameters& diskann_params, const IndexCommonParam& inde use_opq_(diskann_params.use_opq), use_bsa_(diskann_params.use_bsa), use_async_io_(diskann_params.use_async_io), - index_common_param_(index_common_param) { + diskann_params_(diskann_params), + common_param_(index_common_param) { if (not use_async_io_) { pool_ = index_common_param_.thread_pool_; } @@ -212,7 +225,28 @@ DiskANN::build(const DatasetPtr& base) { auto data_num = base->GetNumElements(); std::vector failed_locs; - { + if (diskann_params_.graph_type == DISKANN_GRAPH_TYPE_ODESCENT) { + SlowTaskTimer t("odescent build full (graph)"); + JsonType type = JsonType::parse(FLATTEN_PARAM); + FlattenInterfacePtr flatten_interface_ptr = + FlattenInterface::MakeInstance(type, common_param_); + flatten_interface_ptr->Train(vectors, data_num); + flatten_interface_ptr->BatchInsertVector(vectors, data_num); + vsag::Odescent graph(2 * R_, + diskann_params_.alpha, + diskann_params_.turn, + diskann_params_.sample_rate, + flatten_interface_ptr, + common_param_.allocator_.get(), + common_param_.thread_pool_.get()); + graph.Build(base); + graph.SaveGraph(graph_stream_); + int data_num_int32 = data_num; + int data_dim_int32 = data_dim; + tag_stream_.write((char*)&data_num_int32, sizeof(data_num_int32)); + tag_stream_.write((char*)&data_dim_int32, sizeof(data_dim_int32)); + tag_stream_.write((char*)ids, data_num * sizeof(ids)); + } else if (diskann_params_.graph_type == DISKANN_GRAPH_TYPE_VAMANA) { SlowTaskTimer t("diskann build full (graph)"); // build graph build_index_ = std::make_shared>( diff --git a/src/index/diskann.h b/src/index/diskann.h index 08cac3e6..290cf192 100644 --- a/src/index/diskann.h +++ b/src/index/diskann.h @@ -243,6 +243,9 @@ class DiskANN : public Index { mutable std::shared_mutex rw_mutex_; + IndexCommonParam common_param_; + DiskannParameters diskann_params_; + private: // Request Statistics mutable std::mutex stats_mutex_; std::shared_ptr pool_; diff --git a/src/index/diskann_zparameters.cpp b/src/index/diskann_zparameters.cpp index 0f89bebc..996a14e7 100644 --- a/src/index/diskann_zparameters.cpp +++ b/src/index/diskann_zparameters.cpp @@ -58,16 +58,6 @@ DiskannParameters::FromJson(JsonType& diskann_param_obj, IndexCommonParam index_ CHECK_ARGUMENT((5 <= obj.max_degree) and (obj.max_degree <= 128), fmt::format("max_degree({}) must in range[5, 128]", obj.max_degree)); - // set obj.ef_construction - CHECK_ARGUMENT( - diskann_param_obj.contains(DISKANN_PARAMETER_L), - fmt::format("parameters[{}] must contains {}", INDEX_DISKANN, DISKANN_PARAMETER_L)); - obj.ef_construction = diskann_param_obj[DISKANN_PARAMETER_L]; - CHECK_ARGUMENT((obj.max_degree <= obj.ef_construction) and (obj.ef_construction <= 1000), - fmt::format("ef_construction({}) must in range[$max_degree({}), 64]", - obj.ef_construction, - obj.max_degree)); - // set obj.pq_dims CHECK_ARGUMENT( diskann_param_obj.contains(DISKANN_PARAMETER_DISK_PQ_DIMS), @@ -105,6 +95,41 @@ DiskannParameters::FromJson(JsonType& diskann_param_obj, IndexCommonParam index_ obj.use_async_io = diskann_param_obj[DISKANN_PARAMETER_USE_ASYNC_IO]; } + // set obj.graph_type + if (diskann_param_obj.contains(DISKANN_PARAMETER_GRAPH_TYPE)) { + obj.graph_type = diskann_param_obj[DISKANN_PARAMETER_GRAPH_TYPE]; + } + + if (obj.graph_type == DISKANN_GRAPH_TYPE_VAMANA) { + // set obj.ef_construction + CHECK_ARGUMENT( + diskann_param_obj.contains(DISKANN_PARAMETER_L), + fmt::format("parameters[{}] must contains {}", INDEX_DISKANN, DISKANN_PARAMETER_L)); + obj.ef_construction = diskann_param_obj[DISKANN_PARAMETER_L]; + CHECK_ARGUMENT((obj.max_degree <= obj.ef_construction) and (obj.ef_construction <= 1000), + fmt::format("ef_construction({}) must in range[$max_degree({}), 64]", + obj.ef_construction, + obj.max_degree)); + } else if (obj.graph_type == DISKANN_GRAPH_TYPE_ODESCENT) { + // set obj.alpha + if (diskann_param_obj.contains(DISKANN_PARAMETER_ALPHA)) { + obj.alpha = diskann_param_obj[DISKANN_PARAMETER_ALPHA]; + } + // set obj.turn + if (diskann_param_obj.contains(DISKANN_PARAMETER_TURN)) { + obj.turn = diskann_param_obj[DISKANN_PARAMETER_TURN]; + } + // set obj.sample_rate + if (diskann_param_obj.contains(DISKANN_PARAMETER_SAMPLE_RATE)) { + obj.sample_rate = diskann_param_obj[DISKANN_PARAMETER_SAMPLE_RATE]; + } + } else { + throw std::invalid_argument(fmt::format("parameters[{}] must in [{}, {}], now is {}", + DISKANN_PARAMETER_GRAPH_TYPE, + DISKANN_GRAPH_TYPE_VAMANA, + DISKANN_GRAPH_TYPE_ODESCENT, + obj.graph_type)); + } return obj; } diff --git a/src/index/diskann_zparameters.h b/src/index/diskann_zparameters.h index 2e5f76a4..4aa3b556 100644 --- a/src/index/diskann_zparameters.h +++ b/src/index/diskann_zparameters.h @@ -44,6 +44,12 @@ struct DiskannParameters { bool use_bsa = false; bool use_async_io = false; + // use new construction method + std::string graph_type = "vamana"; + float alpha = 1.2; + int64_t turn = 40; + float sample_rate = 0.3; + private: DiskannParameters() = default; }; diff --git a/src/index/hgraph_zparameters.cpp b/src/index/hgraph_zparameters.cpp index 3171a28b..3c14e174 100644 --- a/src/index/hgraph_zparameters.cpp +++ b/src/index/hgraph_zparameters.cpp @@ -20,7 +20,6 @@ #include "fmt/format-inl.h" #include "inner_string_params.h" #include "vsag/constants.h" - namespace vsag { static const std::unordered_map> EXTERNAL_MAPPING = { diff --git a/src/resource.cpp b/src/resource.cpp index 5df82396..cb6ba41e 100644 --- a/src/resource.cpp +++ b/src/resource.cpp @@ -32,7 +32,7 @@ Resource::Resource(Allocator* allocator, ThreadPool* thread_pool) { if (thread_pool != nullptr) { this->thread_pool = std::make_shared(thread_pool, false); } else { - this->allocator = SafeAllocator::FactoryDefaultAllocator(); + this->thread_pool = SafeThreadPool::FactoryDefaultThreadPool(); } } diff --git a/tests/test_diskann.cpp b/tests/test_diskann.cpp index 6f65787b..20b2e4ba 100644 --- a/tests/test_diskann.cpp +++ b/tests/test_diskann.cpp @@ -65,6 +65,47 @@ TEST_CASE_METHOD(fixtures::DiskANNTestIndex, "diskann build test", "[ft][index][ } } +TEST_CASE_METHOD(fixtures::DiskANNTestIndex, "diskann build and search", "[ft][index][diskann]") { + vsag::Options::Instance().logger()->SetLevel(vsag::Logger::Level::kDEBUG); + auto dims = fixtures::get_common_used_dims(1); + auto metric_type = GENERATE("l2", "ip"); + auto graph_type = GENERATE("vamana", "odescent"); + const std::string name = "diskann"; + + constexpr auto build_parameter_json = R"( + {{ + "dtype": "float32", + "metric_type": "{}", + "dim": {}, + "diskann": {{ + "max_degree": 16, + "ef_construction": 200, + "graph_type": "{}", + "pq_dims": 32, + "pq_sample_rate": 1 + }} + }} + )"; + constexpr auto search_param = R"( + { + "diskann": { + "ef_search": 200, + "io_limit": 200, + "beam_search": 4, + "use_reorder": true + } + } + )"; + auto count = 1000; + for (auto dim : dims) { + auto param = fmt::format(build_parameter_json, metric_type, dim, graph_type); + auto index = TestFactory(name, param, true); + auto dataset = pool.GetDatasetAndCreate(dim, count, metric_type); + TestBuildIndex(index, dataset, true); + TestKnnSearch(index, dataset, search_param, 0.99, true); + } +} + TEST_CASE("DiskAnn Float Recall", "[ft][diskann]") { int dim = 128; // Dimension of the elements int max_elements = 1000; // Maximum number of elements, should be known beforehand diff --git a/tools/test_performance.cpp b/tools/test_performance.cpp index 783cb310..5851890b 100644 --- a/tools/test_performance.cpp +++ b/tools/test_performance.cpp @@ -98,7 +98,8 @@ class PerfTools { const std::string& build_parameters) { spdlog::debug("index_name: " + index_name); spdlog::debug("build_parameters: " + build_parameters); - auto index = Factory::CreateIndex(index_name, build_parameters).value(); + Engine e; + auto index = e.CreateIndex(index_name, build_parameters).value(); spdlog::debug("dataset_path: " + dataset_path); auto eval_dataset = EvalDataset::Load(dataset_path);