diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.cc b/src/core/algorithm/hnsw/hnsw_algorithm.cc index fa553f554..50cd9b64c 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.cc +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -20,7 +20,7 @@ namespace zvec { namespace core { -HnswAlgorithm::HnswAlgorithm(HnswEntity &entity) +HnswAlgorithm::HnswAlgorithm(HnswStreamerEntityNew &entity) : entity_(entity), mt_(std::chrono::system_clock::now().time_since_epoch().count()), lock_pool_(kLockCnt) {} @@ -113,7 +113,7 @@ void HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point, auto &entity = ctx->get_entity(); HnswDistCalculator &dc = ctx->dist_calculator(); while (true) { - const Neighbors neighbors = entity.get_neighbors(level, *entry_point); + const Neighbors neighbors = entity.get_neighbors_new(level, *entry_point); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_neighbors())++; } @@ -123,7 +123,7 @@ void HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point, } std::vector neighbor_vec_blocks; - int ret = entity.get_vector(&neighbors[0], size, neighbor_vec_blocks); + int ret = entity.get_vector_new(&neighbors[0], size, neighbor_vec_blocks); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_vector())++; } @@ -208,7 +208,7 @@ void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, } candidates.pop(); - const Neighbors neighbors = entity.get_neighbors(level, main_node); + const Neighbors neighbors = entity.get_neighbors_new(level, main_node); ailego_prefetch(neighbors.data); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_neighbors())++; @@ -232,7 +232,8 @@ void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, } std::vector neighbor_vec_blocks; - int ret = entity.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks); + int ret = + entity.get_vector_new(neighbor_ids.data(), size, neighbor_vec_blocks); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_vector())++; } @@ -332,7 +333,7 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, node_id_t main_node = top->first; candidates.pop(); - const Neighbors neighbors = entity.get_neighbors(0, main_node); + const Neighbors neighbors = entity.get_neighbors_new(0, main_node); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_neighbors())++; } @@ -356,7 +357,7 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, std::vector neighbor_vec_blocks; int ret = - entity.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks); + entity.get_vector_new(neighbor_ids.data(), size, neighbor_vec_blocks); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_vector())++; } @@ -463,7 +464,7 @@ void HnswAlgorithm::reverse_update_neighbors(HnswDistCalculator &dc, uint32_t lock_idx = id & kLockMask; lock_pool_[lock_idx].lock(); - const Neighbors neighbors = entity_.get_neighbors(level, id); + const Neighbors neighbors = entity_.get_neighbors_new(level, id); size_t size = neighbors.size(); ailego_assert_with(size <= max_neighbor_cnt, "invalid neighbor size"); if (size < max_neighbor_cnt) { diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.h b/src/core/algorithm/hnsw/hnsw_algorithm.h index 886d870c6..e699477bd 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.h +++ b/src/core/algorithm/hnsw/hnsw_algorithm.h @@ -17,7 +17,7 @@ #include #include "hnsw_context.h" #include "hnsw_dist_calculator.h" -#include "hnsw_entity.h" +#include "hnsw_streamer_entity_new.h" namespace zvec { namespace core { @@ -29,7 +29,7 @@ class HnswAlgorithm { public: //! Constructor - explicit HnswAlgorithm(HnswEntity &entity); + explicit HnswAlgorithm(HnswStreamerEntityNew &entity); //! Destructor ~HnswAlgorithm() = default; @@ -116,7 +116,7 @@ class HnswAlgorithm { static constexpr uint32_t kLockCnt{1U << 8}; static constexpr uint32_t kLockMask{kLockCnt - 1U}; - HnswEntity &entity_; + HnswStreamerEntityNew &entity_; mutable std::mt19937 mt_{}; std::vector level_probas_{}; diff --git a/src/core/algorithm/hnsw/hnsw_builder.cc b/src/core/algorithm/hnsw/hnsw_builder.cc deleted file mode 100644 index da2d9faff..000000000 --- a/src/core/algorithm/hnsw/hnsw_builder.cc +++ /dev/null @@ -1,407 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_builder.h" -#include -#include -#include -#include -#include -#include -#include "hnsw_algorithm.h" -#include "hnsw_params.h" - -namespace zvec { -namespace core { - -HnswBuilder::HnswBuilder() = default; - -int HnswBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { - LOG_INFO("Begin HnswBuilder::init"); - - meta_ = meta; - auto params_copy = params; - meta_.set_builder("HnswBuilder", HnswEntity::kRevision, - std::move(params_copy)); - - size_t memory_quota = 0UL; - params.get(PARAM_HNSW_BUILDER_MEMORY_QUOTA, &memory_quota); - params.get(PARAM_HNSW_BUILDER_THREAD_COUNT, &thread_cnt_); - params.get(PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); - params.get(PARAM_HNSW_BUILDER_EFCONSTRUCTION, &ef_construction_); - params.get(PARAM_HNSW_BUILDER_CHECK_INTERVAL_SECS, &check_interval_secs_); - - params.get(PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT, &upper_max_neighbor_cnt_); - float multiplier = HnswEntity::kDefaultL0MaxNeighborCntMultiplier; - params.get(PARAM_HNSW_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, &multiplier); - l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; - scaling_factor_ = upper_max_neighbor_cnt_; - params.get(PARAM_HNSW_BUILDER_SCALING_FACTOR, &scaling_factor_); - - multiplier = HnswEntity::kDefaultNeighborPruneMultiplier; - params.get(PARAM_HNSW_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); - size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; - - if (ef_construction_ == 0) { - ef_construction_ = HnswEntity::kDefaultEfConstruction; - } - if (upper_max_neighbor_cnt_ == 0) { - upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; - } - if (upper_max_neighbor_cnt_ > kMaxNeighborCnt) { - LOG_ERROR("[%s] must be in range (0,%d]", - PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), kMaxNeighborCnt); - return IndexError_InvalidArgument; - } - if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { - LOG_ERROR("[%s]-[%d] must be <= [%s]-[%d]", - PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT.c_str(), min_neighbor_cnt_, - PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), - upper_max_neighbor_cnt_); - return IndexError_InvalidArgument; - } - if (l0_max_neighbor_cnt_ == 0) { - l0_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; - } - if (l0_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { - LOG_ERROR("L0MaxNeighborCnt must be in range (0,%d)", - HnswEntity::kMaxNeighborCnt); - return IndexError_InvalidArgument; - } - if (scaling_factor_ == 0U) { - scaling_factor_ = HnswEntity::kDefaultScalingFactor; - } - if (scaling_factor_ < 5 || scaling_factor_ > 1000) { - LOG_ERROR("[%s] must be in range [5,1000]", - PARAM_HNSW_BUILDER_SCALING_FACTOR.c_str()); - return IndexError_InvalidArgument; - } - if (thread_cnt_ == 0) { - thread_cnt_ = std::thread::hardware_concurrency(); - } - if (thread_cnt_ > std::thread::hardware_concurrency()) { - LOG_WARN("[%s] greater than cpu cores %u", - PARAM_HNSW_BUILDER_THREAD_COUNT.c_str(), - std::thread::hardware_concurrency()); - } - if (prune_cnt == 0UL) { - prune_cnt = upper_max_neighbor_cnt_; - } - - metric_ = IndexFactory::CreateMetric(meta_.metric_name()); - if (!metric_) { - LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); - return IndexError_NoExist; - } - int ret = metric_->init(meta_, meta_.metric_params()); - if (ret != 0) { - LOG_ERROR("IndexMetric init failed, ret=%d", ret); - return ret; - } - - entity_.set_vector_size(meta_.element_size()); - - entity_.set_ef_construction(ef_construction_); - entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_); - entity_.set_min_neighbor_cnt(min_neighbor_cnt_); - entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_); - entity_.set_scaling_factor(scaling_factor_); - entity_.set_memory_quota(memory_quota); - entity_.set_prune_cnt(prune_cnt); - - ret = entity_.init(); - if (ret != 0) { - return ret; - } - - alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_)); - - ret = alg_->init(); - if (ret != 0) { - return ret; - } - - state_ = BUILD_STATE_INITED; - LOG_INFO( - "End HnswBuilder::init, params: vectorSize=%u efConstruction=%u " - "l0NeighborCnt=%u upperNeighborCnt=%u scalingFactor=%u " - "memoryQuota=%zu neighborPruneCnt=%zu metricName=%s ", - meta_.element_size(), ef_construction_, l0_max_neighbor_cnt_, - upper_max_neighbor_cnt_, scaling_factor_, memory_quota, prune_cnt, - meta_.metric_name().c_str()); - - return 0; -} - -int HnswBuilder::cleanup(void) { - LOG_INFO("Begin HnswBuilder::cleanup"); - - l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt; - min_neighbor_cnt_ = 0; - upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; - ef_construction_ = HnswEntity::kDefaultEfConstruction; - scaling_factor_ = HnswEntity::kDefaultScalingFactor; - check_interval_secs_ = kDefaultLogIntervalSecs; - errcode_ = 0; - error_ = false; - entity_.cleanup(); - alg_->cleanup(); - meta_.clear(); - metric_.reset(); - stats_.clear_attributes(); - stats_.set_trained_count(0UL); - stats_.set_built_count(0UL); - stats_.set_dumped_count(0UL); - stats_.set_discarded_count(0UL); - stats_.set_trained_costtime(0UL); - stats_.set_built_costtime(0UL); - stats_.set_dumped_costtime(0UL); - state_ = BUILD_STATE_INIT; - - LOG_INFO("End HnswBuilder::cleanup"); - - return 0; -} - -int HnswBuilder::train(IndexThreads::Pointer, IndexHolder::Pointer holder) { - if (state_ != BUILD_STATE_INITED) { - LOG_ERROR("Init the builder before HnswBuilder::train"); - return IndexError_NoReady; - } - - if (!holder) { - LOG_ERROR("Input holder is nullptr while training index"); - return IndexError_InvalidArgument; - } - if (!holder->is_matched(meta_)) { - LOG_ERROR("Input holder doesn't match index meta while training index"); - return IndexError_Mismatch; - } - LOG_INFO("Begin HnswBuilder::train"); - size_t trained_cost_time = 0; - size_t trained_count = 0; - - if (metric_->support_train()) { - auto start_time = ailego::Monotime::MilliSeconds(); - auto iter = holder->create_iterator(); - if (!iter) { - LOG_ERROR("Create iterator for holder failed"); - return IndexError_Runtime; - } - while (iter->is_valid()) { - int ret = metric_->train(iter->data(), meta_.dimension()); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw build measure train failed, ret=%d", ret); - return ret; - } - iter->next(); - ++trained_count; - } - trained_cost_time = ailego::Monotime::MilliSeconds() - start_time; - } - stats_.set_trained_count(trained_count); - stats_.set_trained_costtime(trained_cost_time); - state_ = BUILD_STATE_TRAINED; - - LOG_INFO("End HnswBuilder::train"); - - return 0; -} - -int HnswBuilder::train(const IndexTrainer::Pointer & /*trainer*/) { - if (state_ != BUILD_STATE_INITED) { - LOG_ERROR("Init the builder before HnswBuilder::train"); - return IndexError_NoReady; - } - - LOG_INFO("Begin HnswBuilder::train by trainer"); - - stats_.set_trained_count(0UL); - stats_.set_trained_costtime(0UL); - state_ = BUILD_STATE_TRAINED; - - LOG_INFO("End HnswBuilder::train by trainer"); - - return 0; -} - -int HnswBuilder::build(IndexThreads::Pointer threads, - IndexHolder::Pointer holder) { - if (state_ != BUILD_STATE_TRAINED) { - LOG_ERROR("Train the index before HnswBuilder::build"); - return IndexError_NoReady; - } - - if (!holder) { - LOG_ERROR("Input holder is nullptr while building index"); - return IndexError_InvalidArgument; - } - if (!holder->is_matched(meta_)) { - LOG_ERROR("Input holder doesn't match index meta while building index"); - return IndexError_Mismatch; - } - if (!threads) { - threads = std::make_shared(thread_cnt_, false); - if (!threads) { - return IndexError_NoMemory; - } - } - - auto start_time = ailego::Monotime::MilliSeconds(); - LOG_INFO("Begin HnswBuilder::build"); - - if (holder->count() != static_cast(-1)) { - LOG_DEBUG("HnswBuilder holder documents count %lu", holder->count()); - int ret = entity_.reserve_space(holder->count()); - if (ret != 0) { - LOG_ERROR("HnswBuilde reserver space failed"); - return ret; - } - } - auto iter = holder->create_iterator(); - if (!iter) { - LOG_ERROR("Create iterator for holder failed"); - return IndexError_Runtime; - } - int ret; - error_ = false; - while (iter->is_valid()) { - level_t level = alg_->get_random_level(); - node_id_t id; - - const void *vec = iter->data(); - ret = entity_.add_vector(level, iter->key(), vec, &id); - if (ailego_unlikely(ret != 0)) { - return ret; - } - iter->next(); - } - // Holder is not needed, cleanup it. - holder.reset(); - - LOG_INFO("Finished save vector, start build graph..."); - - auto task_group = threads->make_group(); - if (!task_group) { - LOG_ERROR("Failed to create task group"); - return IndexError_Runtime; - } - - std::atomic finished{0}; - for (size_t i = 0; i < threads->count(); ++i) { - task_group->submit(ailego::Closure ::New(this, &HnswBuilder::do_build, i, - threads->count(), &finished)); - } - - while (!task_group->is_finished()) { - std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } - LOG_INFO("Built cnt %u, finished percent %.3f%%", finished.load(), - finished.load() * 100.0f / entity_.doc_cnt()); - } - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } - task_group->wait_finish(); - - stats_.set_built_count(finished.load()); - stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); - state_ = BUILD_STATE_BUILT; - - LOG_INFO("End HnswBuilder::build"); - return 0; -} - -void HnswBuilder::do_build(node_id_t idx, size_t step_size, - std::atomic *finished) { - AILEGO_DEFER([&]() { - std::lock_guard latch(mutex_); - cond_.notify_one(); - }); - HnswContext *ctx = new (std::nothrow) - HnswContext(meta_.dimension(), metric_, - std::shared_ptr(&entity_, [](HnswEntity *) {})); - if (ailego_unlikely(ctx == nullptr)) { - if (!error_.exchange(true)) { - LOG_ERROR("Failed to create context"); - errcode_ = IndexError_NoMemory; - } - return; - } - HnswContext::Pointer auto_ptr(ctx); - ctx->set_max_scan_num(entity_.doc_cnt()); - int ret = ctx->init(HnswContext::kBuilderContext); - if (ret != 0) { - if (!error_.exchange(true)) { - LOG_ERROR("Failed to init context"); - errcode_ = IndexError_Runtime; - } - return; - } - - IndexQueryMeta qmeta(meta_.data_type(), meta_.dimension()); - for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) { - ctx->reset_query(entity_.get_vector(id)); - ret = alg_->add_node(id, entity_.get_level(id), ctx); - if (ailego_unlikely(ret != 0)) { - if (!error_.exchange(true)) { - LOG_ERROR("Hnsw graph add node failed"); - errcode_ = ret; - } - return; - } - ctx->clear(); - (*finished)++; - } -} - -int HnswBuilder::dump(const IndexDumper::Pointer &dumper) { - if (state_ != BUILD_STATE_BUILT) { - LOG_INFO("Build the index before HnswBuilder::dump"); - return IndexError_NoReady; - } - - LOG_INFO("Begin HnswBuilder::dump"); - - meta_.set_searcher("HnswSearcher", HnswEntity::kRevision, ailego::Params()); - auto start_time = ailego::Monotime::MilliSeconds(); - - int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); - if (ret != 0) { - LOG_ERROR("Failed to serialize meta into dumper."); - return ret; - } - - ret = entity_.dump(dumper); - if (ret != 0) { - LOG_ERROR("HnswBuilder dump index failed"); - return ret; - } - - stats_.set_dumped_count(entity_.doc_cnt()); - stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); - - LOG_INFO("EndHnswBuilder::dump"); - return 0; -} - -INDEX_FACTORY_REGISTER_BUILDER(HnswBuilder); - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_builder.h b/src/core/algorithm/hnsw/hnsw_builder.h deleted file mode 100644 index e11452838..000000000 --- a/src/core/algorithm/hnsw/hnsw_builder.h +++ /dev/null @@ -1,93 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_algorithm.h" -#include "hnsw_builder_entity.h" - -namespace zvec { -namespace core { - -class HnswBuilder : public IndexBuilder { - public: - //! Constructor - HnswBuilder(); - - //! Initialize the builder - virtual int init(const IndexMeta &meta, - const ailego::Params ¶ms) override; - - //! Cleanup the builder - virtual int cleanup(void) override; - - //! Train the data - virtual int train(IndexThreads::Pointer, - IndexHolder::Pointer holder) override; - - //! Train the data - virtual int train(const IndexTrainer::Pointer &trainer) override; - - - //! Build the index - virtual int build(IndexThreads::Pointer threads, - IndexHolder::Pointer holder) override; - - //! Dump index into storage - virtual int dump(const IndexDumper::Pointer &dumper) override; - - //! Retrieve statistics - virtual const Stats &stats(void) const override { - return stats_; - } - - private: - void do_build(node_id_t idx, size_t step_size, - std::atomic *finished); - - constexpr static uint32_t kDefaultLogIntervalSecs = 15U; - constexpr static uint32_t kMaxNeighborCnt = 65535; - - private: - enum BUILD_STATE { - BUILD_STATE_INIT = 0, - BUILD_STATE_INITED = 1, - BUILD_STATE_TRAINED = 2, - BUILD_STATE_BUILT = 3 - }; - - HnswBuilderEntity entity_{}; - HnswAlgorithm::UPointer alg_; // impl graph algorithm - uint32_t thread_cnt_{0}; - uint32_t min_neighbor_cnt_{0}; - uint32_t upper_max_neighbor_cnt_{HnswEntity::kDefaultUpperMaxNeighborCnt}; - uint32_t l0_max_neighbor_cnt_{HnswEntity::kDefaultL0MaxNeighborCnt}; - uint32_t ef_construction_{HnswEntity::kDefaultEfConstruction}; - uint32_t scaling_factor_{HnswEntity::kDefaultScalingFactor}; - uint32_t check_interval_secs_{kDefaultLogIntervalSecs}; - - int errcode_{0}; - std::atomic_bool error_{false}; - IndexMeta meta_{}; - IndexMetric::Pointer metric_{}; - std::mutex mutex_{}; - std::condition_variable cond_{}; - Stats stats_{}; - - BUILD_STATE state_{BUILD_STATE_INIT}; -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_builder_entity.cc b/src/core/algorithm/hnsw/hnsw_builder_entity.cc deleted file mode 100644 index 472d98ee5..000000000 --- a/src/core/algorithm/hnsw/hnsw_builder_entity.cc +++ /dev/null @@ -1,198 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_builder_entity.h" -#include -#include -#include "utility/sparse_utility.h" - -namespace zvec { -namespace core { - -HnswBuilderEntity::HnswBuilderEntity() { - update_ep_and_level(kInvalidNodeId, 0U); -} - -int HnswBuilderEntity::cleanup() { - memory_quota_ = 0UL; - neighbors_size_ = 0U; - upper_neighbors_size_ = 0U; - padding_size_ = 0U; - vectors_buffer_.clear(); - keys_buffer_.clear(); - neighbors_buffer_.clear(); - upper_neighbors_buffer_.clear(); - neighbors_index_.clear(); - - vectors_buffer_.shrink_to_fit(); - keys_buffer_.shrink_to_fit(); - neighbors_buffer_.shrink_to_fit(); - upper_neighbors_buffer_.shrink_to_fit(); - neighbors_index_.shrink_to_fit(); - - this->HnswEntity::cleanup(); - - return 0; -} - -int HnswBuilderEntity::init() { - size_t size = vector_size(); - - //! aligned size to 32 - set_node_size(AlignSize(size)); - //! if node size is aligned to 1k, the build performance will downgrade - if (node_size() % 1024 == 0) { - set_node_size(AlignSize(node_size() + 1)); - } - - padding_size_ = node_size() - size; - - neighbors_size_ = neighbors_size(); - upper_neighbors_size_ = upper_neighbors_size(); - - return 0; -} - -int HnswBuilderEntity::reserve_space(size_t docs) { - if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs + - sizeof(NeighborIndex) * docs > - memory_quota_)) { - return IndexError_NoMemory; - } - - vectors_buffer_.reserve(node_size() * docs); - keys_buffer_.reserve(sizeof(key_t) * docs); - neighbors_buffer_.reserve(neighbors_size_ * docs); - neighbors_index_.reserve(docs); - - return 0; -} - -int HnswBuilderEntity::add_vector(level_t level, key_t key, const void *vec, - node_id_t *id) { - if (memory_quota_ > 0 && - (vectors_buffer_.capacity() + keys_buffer_.capacity() + - neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() + - neighbors_index_.capacity() * sizeof(NeighborIndex)) > memory_quota_) { - LOG_ERROR("Add vector failed, used memory exceed quota, cur_doc=%u", - doc_cnt()); - return IndexError_NoMemory; - } - - vectors_buffer_.append(reinterpret_cast(vec), vector_size()); - vectors_buffer_.append(padding_size_, '\0'); - keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); - - // init level 0 neighbors - neighbors_buffer_.append(neighbors_size_, '\0'); - - neighbors_index_.emplace_back(upper_neighbors_buffer_.size(), level); - - // init upper layer neighbors - for (level_t cur_level = 1; cur_level <= level; ++cur_level) { - upper_neighbors_buffer_.append(upper_neighbors_size_, '\0'); - } - - *id = (*mutable_doc_cnt())++; - - return 0; -} - -key_t HnswBuilderEntity::get_key(node_id_t id) const { - return *(reinterpret_cast(keys_buffer_.data() + - id * sizeof(key_t))); -} - -const void *HnswBuilderEntity::get_vector(node_id_t id) const { - return vectors_buffer_.data() + id * node_size(); -} - -int HnswBuilderEntity::get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const { - const void *vec = get_vector(id); - block.reset((void *)vec); - return 0; -} - -int HnswBuilderEntity::get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const { - for (uint32_t i = 0; i < count; ++i) { - vecs[i] = vectors_buffer_.data() + ids[i] * node_size(); - } - - return 0; -} - -int HnswBuilderEntity::get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const { - const void *vecs[count]; - get_vector(ids, count, vecs); - for (uint32_t i = 0; i < count; ++i) { - vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); - } - return 0; -} - -const Neighbors HnswBuilderEntity::get_neighbors(level_t level, - node_id_t id) const { - const NeighborsHeader *hd = get_neighbor_header(level, id); - return {hd->neighbor_cnt, hd->neighbors}; -} - -int HnswBuilderEntity::update_neighbors( - level_t level, node_id_t id, - const std::vector> &neighbors) { - NeighborsHeader *hd = - const_cast(get_neighbor_header(level, id)); - for (size_t i = 0; i < neighbors.size(); ++i) { - hd->neighbors[i] = neighbors[i].first; - } - hd->neighbor_cnt = neighbors.size(); - - // std::cout << "id: " << id << ", neighbour, id: "; - // for (size_t i = 0; i < neighbors.size(); ++i) { - // if (i == neighbors.size()-1) - // std::cout << neighbors[i].first << ", score:" << neighbors[i].second << - // std::endl; - // else - // std::cout << neighbors[i].first << ", score:" << neighbors[i].second << - // ", id: "; - // } - - return 0; -} - -void HnswBuilderEntity::add_neighbor(level_t level, node_id_t id, - uint32_t /*size*/, node_id_t neighbor_id) { - NeighborsHeader *hd = - const_cast(get_neighbor_header(level, id)); - hd->neighbors[hd->neighbor_cnt++] = neighbor_id; - - return; -} - -int HnswBuilderEntity::dump(const IndexDumper::Pointer &dumper) { - key_t *keys = - reinterpret_cast(const_cast(keys_buffer_.data())); - auto ret = - dump_segments(dumper, keys, [&](node_id_t id) { return get_level(id); }); - if (ailego_unlikely(ret < 0)) { - return ret; - } - - return 0; -} - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_builder_entity.h b/src/core/algorithm/hnsw/hnsw_builder_entity.h deleted file mode 100644 index 1708e3384..000000000 --- a/src/core/algorithm/hnsw/hnsw_builder_entity.h +++ /dev/null @@ -1,138 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_entity.h" - -namespace zvec { -namespace core { - -class HnswBuilderEntity : public HnswEntity { - public: - //! Add vector and key to hnsw entity, and local id will be saved to id - virtual int add_vector(level_t level, key_t key, const void *vec, - node_id_t *id) override; - - //! Get primary key of the node id - virtual key_t get_key(node_id_t id) const override; - - //! Get vector feature data by key - virtual const void *get_vector(node_id_t id) const override; - - //! Batch get vectors feature data by keys - virtual int get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const override; - - virtual int get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const override; - virtual int get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const override; - - //! Get the node id's neighbors on graph level - const NeighborsHeader *get_neighbor_header(level_t level, - node_id_t id) const { - if (level == 0) { - return reinterpret_cast( - neighbors_buffer_.data() + neighbors_size_ * id); - } else { - size_t offset = neighbors_index_[id].offset; - return reinterpret_cast( - upper_neighbors_buffer_.data() + offset + - (level - 1) * upper_neighbors_size_); - } - } - - //! Get the node id's neighbors on graph level - virtual const Neighbors get_neighbors(level_t level, - node_id_t id) const override; - - //! Replace node id in level's neighbors - virtual int update_neighbors( - level_t level, node_id_t id, - const std::vector> &neighbors) override; - - //! add a neighbor to id in graph level - virtual void add_neighbor(level_t level, node_id_t id, uint32_t size, - node_id_t neighbor_id) override; - - //! Dump the hnsw graph to dumper - virtual int dump(const IndexDumper::Pointer &dumper) override; - - //! Cleanup the entity - virtual int cleanup(void) override; - - public: - //! Constructor - HnswBuilderEntity(); - - //! Get the node graph level by id - level_t get_level(node_id_t id) const { - return neighbors_index_[id].level; - } - - //! Init builerEntity - int init(); - - //! reserve buffer space for documents - //! @param docs number of documents - int reserve_space(size_t docs); - - //! Set memory quota params - inline void set_memory_quota(size_t memory_quota) { - memory_quota_ = memory_quota; - } - - //! Get neighbors size - inline size_t neighbors_size() const { - return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); - } - - //! Get upper neighbors size - inline size_t upper_neighbors_size() const { - return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); - } - - public: - HnswBuilderEntity(const HnswBuilderEntity &) = delete; - HnswBuilderEntity &operator=(const HnswBuilderEntity &) = delete; - - private: - friend class HnswSearcherEntity; - //! class internal used only - struct NeighborIndex { - NeighborIndex(size_t off, level_t l) : offset(off), level(l) {} - uint64_t offset : 48; - uint64_t level : 16; - }; - - std::string vectors_buffer_{}; // aligned vectors - std::string keys_buffer_{}; // aligned vectors - std::string neighbors_buffer_{}; // level 0 neighbors buffer - std::string upper_neighbors_buffer_{}; // upper layer neighbors buffer - - std::string sparse_data_buffer_{}; // aligned spase data buffer - size_t sparse_data_offset_{0}; // - - // upper layer offset + level in upper_neighbors_buffer_ - std::vector neighbors_index_{}; - size_t memory_quota_{0UL}; - size_t neighbors_size_{0U}; // level 0 neighbors size - size_t upper_neighbors_size_{0U}; // level 0 neighbors size - size_t padding_size_{}; // padding size for each vector element -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_context.cc b/src/core/algorithm/hnsw/hnsw_context.cc index b930e4189..3ac975152 100644 --- a/src/core/algorithm/hnsw/hnsw_context.cc +++ b/src/core/algorithm/hnsw/hnsw_context.cc @@ -19,13 +19,13 @@ namespace zvec { namespace core { HnswContext::HnswContext(size_t dimension, const IndexMetric::Pointer &metric, - const HnswEntity::Pointer &entity) + const HnswStreamerEntityNew::Pointer &entity) : IndexContext(metric), entity_(entity), dc_(entity_.get(), metric, dimension) {} HnswContext::HnswContext(const IndexMetric::Pointer &metric, - const HnswEntity::Pointer &entity) + const HnswStreamerEntityNew::Pointer &entity) : IndexContext(metric), entity_(entity), dc_(entity_.get(), metric) {} HnswContext::~HnswContext() { @@ -201,7 +201,7 @@ int HnswContext::update(const ailego::Params ¶ms) { int HnswContext::update_context(ContextType type, const IndexMeta &meta, const IndexMetric::Pointer &metric, - const HnswEntity::Pointer &entity, + const HnswStreamerEntityNew::Pointer &entity, uint32_t magic_num) { uint32_t doc_cnt; diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h index 22bcfaadc..684eafe27 100644 --- a/src/core/algorithm/hnsw/hnsw_context.h +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -17,7 +17,7 @@ #include "utility/sparse_utility.h" #include "utility/visit_filter.h" #include "hnsw_dist_calculator.h" -#include "hnsw_entity.h" +#include "hnsw_streamer_entity_new.h" namespace zvec { namespace core { @@ -36,11 +36,11 @@ class HnswContext : public IndexContext { //! Construct HnswContext(size_t dimension, const IndexMetric::Pointer &metric, - const HnswEntity::Pointer &entity); + const HnswStreamerEntityNew::Pointer &entity); //! Construct HnswContext(const IndexMetric::Pointer &metric, - const HnswEntity::Pointer &entity); + const HnswStreamerEntityNew::Pointer &entity); //! Destructor virtual ~HnswContext(); @@ -114,9 +114,10 @@ class HnswContext : public IndexContext { //! Update context, the context may be shared by different searcher/streamer int update_context(ContextType type, const IndexMeta &meta, const IndexMetric::Pointer &metric, - const HnswEntity::Pointer &entity, uint32_t magic_num); + const HnswStreamerEntityNew::Pointer &entity, + uint32_t magic_num); - inline const HnswEntity &get_entity() const { + inline const HnswStreamerEntityNew &get_entity() const { return *entity_; } @@ -175,7 +176,7 @@ class HnswContext : public IndexContext { node_id_t id = topk_heap_[i].first; if (fetch_vector_) { results_[idx].emplace_back(entity_->get_key(id), score, id, - entity_->get_vector(id)); + entity_->get_vector_new(id)); } else { results_[idx].emplace_back(entity_->get_key(id), score, id); } @@ -238,7 +239,7 @@ class HnswContext : public IndexContext { if (fetch_vector_) { group_results_[idx][i].mutable_docs()->emplace_back( - entity_->get_key(id), score, id, entity_->get_vector(id)); + entity_->get_key(id), score, id, entity_->get_vector_new(id)); } else { group_results_[idx][i].mutable_docs()->emplace_back( entity_->get_key(id), score, id); @@ -488,7 +489,7 @@ class HnswContext : public IndexContext { constexpr static uint32_t kInvalidMgic = -1U; private: - HnswEntity::Pointer entity_; + HnswStreamerEntityNew::Pointer entity_; HnswDistCalculator dc_; IndexMetric::Pointer metric_; @@ -501,9 +502,10 @@ class HnswContext : public IndexContext { uint32_t topk_{0}; uint32_t group_topk_{0}; uint32_t filter_mode_{VisitFilter::ByteMap}; - float negative_probability_{HnswEntity::kDefaultBFNegativeProbability}; - uint32_t ef_{HnswEntity::kDefaultEf}; - float max_scan_ratio_{HnswEntity::kDefaultScanRatio}; + float negative_probability_{ + HnswStreamerEntityNew::kDefaultBFNegativeProbability}; + uint32_t ef_{HnswStreamerEntityNew::kDefaultEf}; + float max_scan_ratio_{HnswStreamerEntityNew::kDefaultScanRatio}; uint32_t magic_{0U}; std::vector results_{}; std::vector group_results_{}; diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h index 84faba40b..84bc255e7 100644 --- a/src/core/algorithm/hnsw/hnsw_dist_calculator.h +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -14,7 +14,7 @@ #pragma once #include -#include "hnsw_entity.h" +#include "hnsw_streamer_entity_new.h" namespace zvec { namespace core { @@ -33,7 +33,7 @@ class HnswDistCalculator { public: //! Constructor - HnswDistCalculator(const HnswEntity *entity, + HnswDistCalculator(const HnswStreamerEntityNew *entity, const IndexMetric::Pointer &metric, uint32_t dim) : entity_(entity), distance_(metric->distance()), @@ -43,7 +43,7 @@ class HnswDistCalculator { compare_cnt_(0) {} //! Constructor - HnswDistCalculator(const HnswEntity *entity, + HnswDistCalculator(const HnswStreamerEntityNew *entity, const IndexMetric::Pointer &metric, uint32_t dim, const void *query) : entity_(entity), @@ -54,7 +54,7 @@ class HnswDistCalculator { compare_cnt_(0) {} //! Constructor - HnswDistCalculator(const HnswEntity *entity, + HnswDistCalculator(const HnswStreamerEntityNew *entity, const IndexMetric::Pointer &metric) : entity_(entity), distance_(metric->distance()), @@ -63,14 +63,15 @@ class HnswDistCalculator { dim_(0), compare_cnt_(0) {} - void update(const HnswEntity *entity, const IndexMetric::Pointer &metric) { + void update(const HnswStreamerEntityNew *entity, + const IndexMetric::Pointer &metric) { entity_ = entity; distance_ = metric->distance(); batch_distance_ = metric->batch_distance(); } - void update(const HnswEntity *entity, const IndexMetric::Pointer &metric, - uint32_t dim) { + void update(const HnswStreamerEntityNew *entity, + const IndexMetric::Pointer &metric, uint32_t dim) { entity_ = entity; distance_ = metric->distance(); batch_distance_ = metric->batch_distance(); @@ -116,7 +117,7 @@ class HnswDistCalculator { inline dist_t dist(node_id_t id) { compare_cnt_++; - const void *feat = entity_->get_vector(id); + const void *feat = entity_->get_vector_new(id); if (ailego_unlikely(feat == nullptr)) { LOG_ERROR("Get nullptr vector, id=%u", id); error_ = true; @@ -130,8 +131,8 @@ class HnswDistCalculator { inline dist_t dist(node_id_t lhs, node_id_t rhs) { compare_cnt_++; - const void *feat = entity_->get_vector(lhs); - const void *query = entity_->get_vector(rhs); + const void *feat = entity_->get_vector_new(lhs); + const void *query = entity_->get_vector_new(rhs); if (ailego_unlikely(feat == nullptr || query == nullptr)) { LOG_ERROR("Get nullptr vector"); error_ = true; @@ -162,7 +163,7 @@ class HnswDistCalculator { inline dist_t batch_dist(node_id_t id) { compare_cnt_++; - const void *feat = entity_->get_vector(id); + const void *feat = entity_->get_vector_new(id); if (ailego_unlikely(feat == nullptr)) { LOG_ERROR("Get nullptr vector, id=%u", id); error_ = true; @@ -201,7 +202,7 @@ class HnswDistCalculator { HnswDistCalculator &operator=(const HnswDistCalculator &) = delete; private: - const HnswEntity *entity_; + const HnswStreamerEntityNew *entity_; IndexMetric::MatrixDistance distance_; IndexMetric::MatrixBatchDistance batch_distance_; diff --git a/src/core/algorithm/hnsw/hnsw_index_provider.h b/src/core/algorithm/hnsw/hnsw_index_provider.h index 4a6ccaebb..064934ce0 100644 --- a/src/core/algorithm/hnsw/hnsw_index_provider.h +++ b/src/core/algorithm/hnsw/hnsw_index_provider.h @@ -16,14 +16,15 @@ #include #include #include -#include "hnsw_entity.h" +#include "hnsw_streamer_entity_new.h" namespace zvec { namespace core { class HnswIndexProvider : public IndexProvider { public: - HnswIndexProvider(const IndexMeta &meta, const HnswEntity::Pointer &entity, + HnswIndexProvider(const IndexMeta &meta, + const HnswStreamerEntityNew::Pointer &entity, const std::string &owner) : meta_(meta), entity_(entity), owner_class_(owner) {} @@ -76,14 +77,14 @@ class HnswIndexProvider : public IndexProvider { private: class Iterator : public IndexProvider::Iterator { public: - Iterator(const HnswEntity::Pointer &entity) + Iterator(const HnswStreamerEntityNew::Pointer &entity) : entity_(entity), cur_id_(0U) {} //! Retrieve pointer of data //! NOTICE: the vec feature will be changed after iterating to next, so //! the caller need to keep a copy of it before iterator to next vector virtual const void *data(void) const override { - return entity_->get_vector(cur_id_); + return entity_->get_vector_new(cur_id_); } //! Test if the iterator is valid @@ -119,13 +120,13 @@ class HnswIndexProvider : public IndexProvider { } private: - const HnswEntity::Pointer entity_; + const HnswStreamerEntityNew::Pointer entity_; node_id_t cur_id_; }; private: const IndexMeta &meta_; - const HnswEntity::Pointer entity_; + const HnswStreamerEntityNew::Pointer entity_; const std::string owner_class_; }; diff --git a/src/core/algorithm/hnsw/hnsw_searcher.cc b/src/core/algorithm/hnsw/hnsw_searcher.cc deleted file mode 100644 index c68a146d9..000000000 --- a/src/core/algorithm/hnsw/hnsw_searcher.cc +++ /dev/null @@ -1,460 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_searcher.h" -#include "hnsw_algorithm.h" -#include "hnsw_index_provider.h" -#include "hnsw_params.h" - -namespace zvec { -namespace core { - -HnswSearcher::HnswSearcher() = default; - -HnswSearcher::~HnswSearcher() = default; - -int HnswSearcher::init(const ailego::Params &search_params) { - params_ = search_params; - params_.get(PARAM_HNSW_SEARCHER_EF, &ef_); - params_.get(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); - params_.get(PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_); - params_.get(PARAM_HNSW_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_); - params_.get(PARAM_HNSW_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE, - &neighbors_in_memory_enabled_); - params_.get(PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB, - &bf_negative_probability_); - params_.get(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD, - &bruteforce_threshold_); - params_.get(PARAM_HNSW_SEARCHER_FORCE_PADDING_RESULT_ENABLE, - &force_padding_topk_enabled_); - - if (ef_ == 0) { - ef_ = HnswEntity::kDefaultEf; - } - if (bf_negative_probability_ <= 0.0f || bf_negative_probability_ >= 1.0f) { - LOG_ERROR("[%s] must be in range (0,1)", - PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str()); - return IndexError_InvalidArgument; - } - - entity_.set_neighbors_in_memory(neighbors_in_memory_enabled_); - - state_ = STATE_INITED; - - LOG_DEBUG( - "Init params: ef=%u maxScanRatio=%f bfEnabled=%u checkCrcEnabled=%u " - "neighborsInMemoryEnabled=%u bfNagtiveProb=%f bruteForceThreshold=%u " - "forcePadding=%u", - ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_, - neighbors_in_memory_enabled_, bf_negative_probability_, - bruteforce_threshold_, force_padding_topk_enabled_); - - return 0; -} - -void HnswSearcher::print_debug_info() { - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - Neighbors neighbours = entity_.get_neighbors(0, id); - std::cout << "node: " << id << "; "; - for (uint32_t i = 0; i < neighbours.size(); ++i) { - std::cout << neighbours[i]; - - if (i == neighbours.size() - 1) { - std::cout << std::endl; - } else { - std::cout << ", "; - } - } - } -} - -int HnswSearcher::cleanup() { - LOG_INFO("Begin HnswSearcher:cleanup"); - - metric_.reset(); - meta_.clear(); - stats_.clear_attributes(); - stats_.set_loaded_count(0UL); - stats_.set_loaded_costtime(0UL); - max_scan_ratio_ = HnswEntity::kDefaultScanRatio; - max_scan_num_ = 0U; - ef_ = HnswEntity::kDefaultEf; - bf_enabled_ = false; - bf_negative_probability_ = HnswEntity::kDefaultBFNegativeProbability; - bruteforce_threshold_ = HnswEntity::kDefaultBruteForceThreshold; - check_crc_enabled_ = false; - neighbors_in_memory_enabled_ = false; - entity_.cleanup(); - state_ = STATE_INIT; - - LOG_INFO("End HnswSearcher:cleanup"); - - return 0; -} - -int HnswSearcher::load(IndexStorage::Pointer container, - IndexMetric::Pointer metric) { - if (state_ != STATE_INITED) { - LOG_ERROR("Init the searcher first before load index"); - return IndexError_Runtime; - } - - LOG_INFO("Begin HnswSearcher:load"); - - auto start_time = ailego::Monotime::MilliSeconds(); - - int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_); - if (ret != 0) { - LOG_ERROR("Failed to deserialize meta from container"); - return ret; - } - - ret = entity_.load(container, check_crc_enabled_); - if (ret != 0) { - LOG_ERROR("HnswSearcher load index failed"); - return ret; - } - - alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_)); - - if (metric) { - metric_ = metric; - } else { - metric_ = IndexFactory::CreateMetric(meta_.metric_name()); - if (!metric_) { - LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); - return IndexError_NoExist; - } - ret = metric_->init(meta_, meta_.metric_params()); - if (ret != 0) { - LOG_ERROR("IndexMetric init failed, ret=%d", ret); - return ret; - } - if (metric_->query_metric()) { - metric_ = metric_->query_metric(); - } - } - - if (!metric_->is_matched(meta_)) { - LOG_ERROR("IndexMetric not match index meta"); - return IndexError_Mismatch; - } - - max_scan_num_ = static_cast(max_scan_ratio_ * entity_.doc_cnt()); - max_scan_num_ = std::max(4096U, max_scan_num_); - - stats_.set_loaded_count(entity_.doc_cnt()); - stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); - state_ = STATE_LOADED; - magic_ = IndexContext::GenerateMagic(); - - LOG_INFO("End HnswSearcher::load"); - - return 0; -} - -int HnswSearcher::unload() { - LOG_INFO("HnswSearcher unload index"); - - meta_.clear(); - entity_.cleanup(); - metric_.reset(); - max_scan_num_ = 0; - stats_.set_loaded_count(0UL); - stats_.set_loaded_costtime(0UL); - state_ = STATE_INITED; - - return 0; -} - -int HnswSearcher::update_context(HnswContext *ctx) const { - const HnswEntity::Pointer entity = entity_.clone(); - if (!entity) { - LOG_ERROR("Failed to clone search context entity"); - return IndexError_Runtime; - } - ctx->set_max_scan_num(max_scan_num_); - ctx->set_bruteforce_threshold(bruteforce_threshold_); - - return ctx->update_context(HnswContext::kSearcherContext, meta_, metric_, - entity, magic_); -} - -int HnswSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, - uint32_t count, Context::Pointer &context) const { - if (ailego_unlikely(!query || !context)) { - LOG_ERROR("The context is not created by this searcher"); - return IndexError_Mismatch; - } - HnswContext *ctx = dynamic_cast(context.get()); - ailego_do_if_false(ctx) { - LOG_ERROR("Cast context to HnswContext failed"); - return IndexError_Cast; - } - - if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) { - return search_bf_impl(query, qmeta, count, context); - } - - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer - int ret = update_context(ctx); - if (ret != 0) { - return ret; - } - } - - ctx->clear(); - ctx->resize_results(count); - for (size_t q = 0; q < count; ++q) { - ctx->reset_query(query); - int ret = alg_->search(ctx); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher fast search failed"); - return ret; - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; - } - - return 0; -} - -int HnswSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, - uint32_t count, - Context::Pointer &context) const { - if (ailego_unlikely(!query || !context)) { - LOG_ERROR("The context is not created by this searcher"); - return IndexError_Mismatch; - } - HnswContext *ctx = dynamic_cast(context.get()); - ailego_do_if_false(ctx) { - LOG_ERROR("Cast context to HnswContext failed"); - return IndexError_Cast; - } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer - int ret = update_context(ctx); - if (ret != 0) { - return ret; - } - } - - ctx->clear(); - ctx->resize_results(count); - - if (ctx->group_by_search()) { - if (!ctx->group_by().is_valid()) { - LOG_ERROR("Invalid group-by function"); - return IndexError_InvalidArgument; - } - - std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity_.get_key(id)); - }; - - for (size_t q = 0; q < count; ++q) { - ctx->reset_query(query); - ctx->group_topk_heaps().clear(); - - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { - continue; - } - if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { - dist_t dist = ctx->dist_calculator().batch_dist(id); - - std::string group_id = group_by(id); - - auto &topk_heap = ctx->group_topk_heaps()[group_id]; - if (topk_heap.empty()) { - topk_heap.limit(ctx->group_topk()); - } - topk_heap.emplace_back(id, dist); - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } else { - for (size_t q = 0; q < count; ++q) { - ctx->reset_query(query); - ctx->topk_heap().clear(); - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { - continue; - } - if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { - dist_t dist = ctx->dist_calculator().batch_dist(id); - ctx->topk_heap().emplace(id, dist); - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; - } - - return 0; -} - -int HnswSearcher::search_bf_by_p_keys_impl( - const void *query, const std::vector> &p_keys, - const IndexQueryMeta &qmeta, uint32_t count, - Context::Pointer &context) const { - if (ailego_unlikely(!query || !context)) { - LOG_ERROR("The context is not created by this searcher"); - return IndexError_Mismatch; - } - - if (ailego_unlikely(p_keys.size() != count)) { - LOG_ERROR("The size of p_keys is not equal to count"); - return IndexError_InvalidArgument; - } - - HnswContext *ctx = dynamic_cast(context.get()); - ailego_do_if_false(ctx) { - LOG_ERROR("Cast context to HnswContext failed"); - return IndexError_Cast; - } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer - int ret = update_context(ctx); - if (ret != 0) { - return ret; - } - } - - ctx->clear(); - ctx->resize_results(count); - - if (ctx->group_by_search()) { - if (!ctx->group_by().is_valid()) { - LOG_ERROR("Invalid group-by function"); - return IndexError_InvalidArgument; - } - - std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity_.get_key(id)); - }; - - for (size_t q = 0; q < count; ++q) { - ctx->reset_query(query); - ctx->group_topk_heaps().clear(); - - for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { - uint64_t pk = p_keys[q][idx]; - if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { - node_id_t id = entity_.get_id(pk); - if (id != kInvalidNodeId) { - dist_t dist = ctx->dist_calculator().batch_dist(id); - std::string group_id = group_by(id); - - auto &topk_heap = ctx->group_topk_heaps()[group_id]; - if (topk_heap.empty()) { - topk_heap.limit(ctx->group_topk()); - } - topk_heap.emplace_back(id, dist); - } - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } else { - for (size_t q = 0; q < count; ++q) { - ctx->reset_query(query); - ctx->topk_heap().clear(); - for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { - uint64_t pk = p_keys[q][idx]; - if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { - node_id_t id = entity_.get_id(pk); - if (id != kInvalidNodeId) { - dist_t dist = ctx->dist_calculator().batch_dist(id); - ctx->topk_heap().emplace(id, dist); - } - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; - } - - return 0; -} - -IndexSearcher::Context::Pointer HnswSearcher::create_context() const { - if (ailego_unlikely(state_ != STATE_LOADED)) { - LOG_ERROR("Load the index first before create context"); - return Context::Pointer(); - } - const HnswEntity::Pointer search_ctx_entity = entity_.clone(); - if (!search_ctx_entity) { - LOG_ERROR("Failed to create search context entity"); - return Context::Pointer(); - } - HnswContext *ctx = new (std::nothrow) - HnswContext(meta_.dimension(), metric_, search_ctx_entity); - if (ailego_unlikely(ctx == nullptr)) { - LOG_ERROR("Failed to new HnswContext"); - return Context::Pointer(); - } - ctx->set_ef(ef_); - ctx->set_max_scan_num(max_scan_num_); - uint32_t filter_mode = - bf_enabled_ ? VisitFilter::BloomFilter : VisitFilter::ByteMap; - ctx->set_filter_mode(filter_mode); - ctx->set_filter_negative_probability(bf_negative_probability_); - ctx->set_magic(magic_); - ctx->set_force_padding_topk(force_padding_topk_enabled_); - ctx->set_bruteforce_threshold(bruteforce_threshold_); - if (ailego_unlikely(ctx->init(HnswContext::kSearcherContext)) != 0) { - LOG_ERROR("Init HnswContext failed"); - delete ctx; - return Context::Pointer(); - } - - return Context::Pointer(ctx); -} - -IndexProvider::Pointer HnswSearcher::create_provider(void) const { - LOG_DEBUG("HnswSearcher create provider"); - - auto entity = entity_.clone(); - if (ailego_unlikely(!entity)) { - LOG_ERROR("Clone HnswEntity failed"); - return Provider::Pointer(); - } - return Provider::Pointer( - new (std::nothrow) HnswIndexProvider(meta_, entity, "HnswSearcher")); -} - -const void *HnswSearcher::get_vector(uint64_t key) const { - return entity_.get_vector_by_key(key); -} - -INDEX_FACTORY_REGISTER_SEARCHER(HnswSearcher); - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_searcher.h b/src/core/algorithm/hnsw/hnsw_searcher.h deleted file mode 100644 index d79526df8..000000000 --- a/src/core/algorithm/hnsw/hnsw_searcher.h +++ /dev/null @@ -1,139 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_searcher_entity.h" -#include "hnsw_streamer.h" - -namespace zvec { -namespace core { - -class HnswSearcher : public IndexSearcher { - public: - using ContextPointer = IndexSearcher::Context::Pointer; - - public: - HnswSearcher(void); - ~HnswSearcher(void); - - HnswSearcher(const HnswSearcher &) = delete; - HnswSearcher &operator=(const HnswSearcher &) = delete; - - protected: - //! Initialize Searcher - virtual int init(const ailego::Params ¶ms) override; - - //! Cleanup Searcher - virtual int cleanup(void) override; - - //! Load Index from storage - virtual int load(IndexStorage::Pointer container, - IndexMetric::Pointer metric) override; - - //! Unload index from storage - virtual int unload(void) override; - - //! KNN Search - virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, - ContextPointer &context) const override { - return search_impl(query, qmeta, 1, context); - } - - //! KNN Search - virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, - uint32_t count, - ContextPointer &context) const override; - - //! Linear Search - virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, - ContextPointer &context) const override { - return search_bf_impl(query, qmeta, 1, context); - } - - //! Linear Search - virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, - uint32_t count, - ContextPointer &context) const override; - - //! Linear search by primary keys - virtual int search_bf_by_p_keys_impl( - const void *query, const std::vector> &p_keys, - const IndexQueryMeta &qmeta, ContextPointer &context) const override { - return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); - } - - //! Linear search by primary keys - virtual int search_bf_by_p_keys_impl( - const void *query, const std::vector> &p_keys, - const IndexQueryMeta &qmeta, uint32_t count, - ContextPointer &context) const override; - - //! Fetch vector by key - virtual const void *get_vector(uint64_t key) const override; - - //! Create a searcher context - virtual ContextPointer create_context() const override; - - //! Create a new iterator - virtual IndexProvider::Pointer create_provider(void) const override; - - //! Retrieve statistics - virtual const Stats &stats(void) const override { - return stats_; - } - - //! Retrieve meta of index - virtual const IndexMeta &meta(void) const override { - return meta_; - } - - //! Retrieve params of index - virtual const ailego::Params ¶ms(void) const override { - return params_; - } - - virtual void print_debug_info() override; - - private: - //! To share ctx across streamer/searcher, we need to update the context for - //! current streamer/searcher - int update_context(HnswContext *ctx) const; - - private: - enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; - - HnswSearcherEntity entity_{}; - HnswAlgorithm::UPointer alg_; // impl graph algorithm - - IndexMetric::Pointer metric_{}; - IndexMeta meta_{}; - ailego::Params params_{}; - Stats stats_; - uint32_t ef_{HnswEntity::kDefaultEf}; - uint32_t max_scan_num_{0U}; - uint32_t bruteforce_threshold_{HnswEntity::kDefaultBruteForceThreshold}; - float max_scan_ratio_{HnswEntity::kDefaultScanRatio}; - bool bf_enabled_{false}; - bool check_crc_enabled_{false}; - bool neighbors_in_memory_enabled_{false}; - bool force_padding_topk_enabled_{false}; - float bf_negative_probability_{HnswEntity::kDefaultBFNegativeProbability}; - uint32_t magic_{0U}; - - State state_{STATE_INIT}; -}; - -} // namespace core -} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_searcher_entity.cc b/src/core/algorithm/hnsw/hnsw_searcher_entity.cc deleted file mode 100644 index 6661c1db3..000000000 --- a/src/core/algorithm/hnsw/hnsw_searcher_entity.cc +++ /dev/null @@ -1,515 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_searcher_entity.h" -#include -#include "utility/sparse_utility.h" - -namespace zvec { -namespace core { - -HnswSearcherEntity::HnswSearcherEntity() {} - -int HnswSearcherEntity::cleanup(void) { - storage_.reset(); - vectors_.reset(); - keys_.reset(); - neighbors_.reset(); - neighbors_meta_.reset(); - neighbors_in_memory_enabled_ = false; - loaded_ = false; - - this->HnswEntity::cleanup(); - - return 0; -} - -key_t HnswSearcherEntity::get_key(node_id_t id) const { - const void *key; - if (ailego_unlikely(keys_->read(id * sizeof(key_t), &key, sizeof(key_t)) != - sizeof(key_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidKey; - } - return *(reinterpret_cast(key)); -} - -//! Get vector local id by key -node_id_t HnswSearcherEntity::get_id(key_t key) const { - if (ailego_unlikely(!mapping_)) { - LOG_ERROR("Index missing mapping segment"); - return kInvalidNodeId; - } - - //! Do binary search - node_id_t start = 0UL; - node_id_t end = doc_cnt(); - const void *data; - node_id_t idx = 0u; - while (start < end) { - idx = start + (end - start) / 2; - if (ailego_unlikely( - mapping_->read(idx * sizeof(node_id_t), &data, sizeof(node_id_t)) != - sizeof(node_id_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidNodeId; - } - const key_t *mkey; - node_id_t local_id = *reinterpret_cast(data); - if (ailego_unlikely(keys_->read(local_id * sizeof(key_t), - (const void **)(&mkey), - sizeof(key_t)) != sizeof(key_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidNodeId; - } - if (*mkey < key) { - start = idx + 1; - } else if (*mkey > key) { - end = idx; - } else { - return local_id; - } - } - return kInvalidNodeId; -} - -const void *HnswSearcherEntity::get_vector_by_key(key_t key) const { - node_id_t local_id = get_id(key); - if (ailego_unlikely(local_id == kInvalidNodeId)) { - return nullptr; - } - - return get_vector(local_id); -} - -const void *HnswSearcherEntity::get_vector(node_id_t id) const { - size_t read_size = vector_size(); - size_t offset = node_size() * id; - - const void *vec; - if (ailego_unlikely(vectors_->read(offset, &vec, read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed"); - return nullptr; - } - return vec; -} - -int HnswSearcherEntity::get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const { - const void *vec = get_vector(id); - block.reset((void *)vec); - return 0; -} - -const void *HnswSearcherEntity::get_vectors() const { - const void *vec; - size_t len = node_size() * doc_cnt(); - if (vectors_->read(0, &vec, len) != len) { - LOG_ERROR("Read vectors from segment failed"); - return nullptr; - } - return vec; -} - -int HnswSearcherEntity::get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const { - ailego_assert_with(count <= segment_datas_.size(), "invalid count"); - - size_t read_size = vector_size(); - - for (uint32_t i = 0; i < count; ++i) { - segment_datas_[i].offset = node_size() * ids[i]; - segment_datas_[i].length = read_size; - - ailego_assert_with(segment_datas_[i].offset < vectors_->data_size(), - "invalid offset"); - } - if (ailego_unlikely(!vectors_->read(&segment_datas_[0], count))) { - LOG_ERROR("Read vectors from segment failed"); - return IndexError_ReadData; - } - for (uint32_t i = 0; i < count; ++i) { - vecs[i] = segment_datas_[i].data; - } - - return 0; -} - -int HnswSearcherEntity::get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const { - const void *vecs[count]; - get_vector(ids, count, vecs); - for (uint32_t i = 0; i < count; ++i) { - vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); - } - return 0; -} - -const Neighbors HnswSearcherEntity::get_neighbors(level_t level, - node_id_t id) const { - if (level == 0) { - if (neighbors_in_memory_enabled_) { - auto hd = reinterpret_cast( - fixed_neighbors_.get() + neighbors_size() * id); - return {hd->neighbor_cnt, hd->neighbors}; - } - - const GraphNeighborMeta *m; - if (ailego_unlikely(neighbors_meta_->read(id * sizeof(GraphNeighborMeta), - (const void **)(&m), - sizeof(GraphNeighborMeta)) != - sizeof(GraphNeighborMeta))) { - LOG_ERROR("Read neighbors meta from segment failed"); - return {0, nullptr}; - } - - const void *data; - if (ailego_unlikely(neighbors_->read(m->offset, &data, - m->neighbor_cnt * sizeof(node_id_t)) != - m->neighbor_cnt * sizeof(node_id_t))) { - LOG_ERROR("Read neighbors from segment failed"); - return {0, nullptr}; - } - return {static_cast(m->neighbor_cnt), - reinterpret_cast(data)}; - } - - //! Read level > 0 neighbors - const HnswNeighborMeta *m; - if (ailego_unlikely(upper_neighbors_meta_->read(id * sizeof(HnswNeighborMeta), - (const void **)(&m), - sizeof(HnswNeighborMeta)) != - sizeof(HnswNeighborMeta))) { - LOG_ERROR("Read neighbors meta from segment failed"); - return {0, nullptr}; - } - - ailego_assert_with(level <= m->level, "invalid level"); - size_t offset = m->offset + (level - 1) * upper_neighbors_size(); - ailego_assert_with(offset <= upper_neighbors_->data_size(), "invalid offset"); - const void *data; - if (ailego_unlikely( - upper_neighbors_->read(offset, &data, upper_neighbors_size()) != - upper_neighbors_size())) { - LOG_ERROR("Read neighbors from segment failed"); - return {0, nullptr}; - } - - auto hd = reinterpret_cast(data); - return {hd->neighbor_cnt, hd->neighbors}; -} - -int HnswSearcherEntity::load(const IndexStorage::Pointer &container, - bool check_crc) { - storage_ = container; - - int ret = load_segments(check_crc); - if (ret != 0) { - return ret; - } - - loaded_ = true; - - LOG_INFO( - "Index info: docCnt=%u entryPoint=%u maxLevel=%d efConstruct=%zu " - "l0NeighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu " - "vectorSize=%zu nodeSize=%zu vectorSegmentSize=%zu keySegmentSize=%zu " - "neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu ", - doc_cnt(), entry_point(), cur_max_level(), ef_construction(), - l0_neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), vector_size(), - node_size(), vectors_->data_size(), keys_->data_size(), - neighbors_ == nullptr ? 0 : neighbors_->data_size(), - neighbors_meta_ == nullptr ? 0 : neighbors_meta_->data_size()); - - return 0; -} - -int HnswSearcherEntity::load_segments(bool check_crc) { - //! load header - const void *data = nullptr; - HNSWHeader hd; - auto graph_hd_segment = storage_->get(kGraphHeaderSegmentId); - if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) { - LOG_ERROR("Miss or invalid segment %s", kGraphHeaderSegmentId.c_str()); - return IndexError_InvalidFormat; - } - if (graph_hd_segment->read(0, reinterpret_cast(&data), - sizeof(hd.graph)) != sizeof(hd.graph)) { - LOG_ERROR("Read segment %s failed", kGraphHeaderSegmentId.c_str()); - return IndexError_ReadData; - } - memcpy(&hd.graph, data, sizeof(hd.graph)); - - auto hnsw_hd_segment = storage_->get(kHnswHeaderSegmentId); - if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) { - LOG_ERROR("Miss or invalid segment %s", kHnswHeaderSegmentId.c_str()); - return IndexError_InvalidFormat; - } - if (hnsw_hd_segment->read(0, reinterpret_cast(&data), - sizeof(hd.hnsw)) != sizeof(hd.hnsw)) { - LOG_ERROR("Read segment %s failed", kHnswHeaderSegmentId.c_str()); - return IndexError_ReadData; - } - memcpy(&hd.hnsw, data, sizeof(hd.hnsw)); - *mutable_header() = hd; - segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt())); - - vectors_ = storage_->get(kGraphFeaturesSegmentId); - if (!vectors_) { - LOG_ERROR("IndexStorage get segment %s failed", - kGraphFeaturesSegmentId.c_str()); - return IndexError_InvalidFormat; - } - keys_ = storage_->get(kGraphKeysSegmentId); - if (!keys_) { - LOG_ERROR("IndexStorage get segment %s failed", - kGraphKeysSegmentId.c_str()); - return IndexError_InvalidFormat; - } - - neighbors_ = storage_->get(kGraphNeighborsSegmentId); - if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) { - LOG_ERROR("IndexStorage get segment %s failed or empty", - kGraphNeighborsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - neighbors_meta_ = storage_->get(kGraphOffsetsSegmentId); - if (!neighbors_meta_ || - neighbors_meta_->data_size() < sizeof(GraphNeighborMeta) * doc_cnt()) { - LOG_ERROR("IndexStorage get segment %s failed or invalid size", - kGraphOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - upper_neighbors_ = storage_->get(kHnswNeighborsSegmentId); - if (!upper_neighbors_ || - (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) { - LOG_ERROR("IndexStorage get segment %s failed or empty", - kHnswNeighborsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - upper_neighbors_meta_ = storage_->get(kHnswOffsetsSegmentId); - if (!upper_neighbors_meta_ || upper_neighbors_meta_->data_size() < - sizeof(HnswNeighborMeta) * doc_cnt()) { - LOG_ERROR("IndexStorage get segment %s failed or invalid size", - kHnswOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - mapping_ = storage_->get(kGraphMappingSegmentId); - if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) { - LOG_ERROR("IndexStorage get segment %s failed or invalid size", - kGraphMappingSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - if (check_crc) { - std::vector segments; - segments.emplace_back(graph_hd_segment); - segments.emplace_back(hnsw_hd_segment); - segments.emplace_back(vectors_); - segments.emplace_back(keys_); - - segments.emplace_back(neighbors_); - segments.emplace_back(neighbors_meta_); - segments.emplace_back(upper_neighbors_); - segments.emplace_back(upper_neighbors_meta_); - - if (!do_crc_check(segments)) { - LOG_ERROR("Check index crc failed, the index may broken"); - return IndexError_Runtime; - } - } - - if (neighbors_in_memory_enabled_) { - int ret = load_and_flat_neighbors(); - if (ret != 0) { - return ret; - } - } - - return 0; -} - -int HnswSearcherEntity::load_and_flat_neighbors() { - fixed_neighbors_.reset( - new (std::nothrow) char[neighbors_size() * doc_cnt()]{}, - std::default_delete()); - if (!fixed_neighbors_) { - LOG_ERROR("Malloc memory failed"); - return IndexError_NoMemory; - } - - //! Get a new segemnt to release the buffer after loading neighbors - auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId); - if (!neighbors_meta) { - LOG_ERROR("IndexStorage get segment graph.offsets failed"); - return IndexError_InvalidArgument; - } - - const GraphNeighborMeta *neighbors_index = nullptr; - if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index), - neighbors_meta->data_size()) != - neighbors_meta->data_size()) { - LOG_ERROR("Read segment %s data failed", kGraphOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - const char *neighbor_data; - for (node_id_t id = 0; id < doc_cnt(); ++id) { - size_t rd_size = neighbors_index[id].neighbor_cnt * sizeof(node_id_t); - if (ailego_unlikely( - neighbors_->read(neighbors_index[id].offset, - reinterpret_cast(&neighbor_data), - rd_size) != rd_size)) { - LOG_ERROR("Read neighbors from segment failed"); - return IndexError_ReadData; - } - // copy level 0 neighbors to fixed size neighbors memory - char *dst = fixed_neighbors_.get() + neighbors_size() * id; - *reinterpret_cast(dst) = neighbors_index[id].neighbor_cnt; - memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size); - } - - return 0; -} - -int HnswSearcherEntity::get_fixed_neighbors( - std::vector *fixed_neighbors) const { - //! Get a new segemnt to release the buffer after loading neighbors - auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId); - if (!neighbors_meta) { - LOG_ERROR("IndexStorage get segment graph.offsets failed"); - return IndexError_InvalidArgument; - } - - const GraphNeighborMeta *neighbors_index = nullptr; - size_t meta_size = neighbors_meta->data_size(); - if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index), - meta_size) != meta_size) { - LOG_ERROR("Read segment %s data failed", kGraphOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - size_t fixed_neighbor_cnt = l0_neighbor_cnt(); - fixed_neighbors->resize((fixed_neighbor_cnt + 1) * doc_cnt(), kInvalidNodeId); - - size_t neighbors_cnt_offset = fixed_neighbor_cnt * doc_cnt(); - size_t total_neighbor_cnt = 0; - for (node_id_t id = 0; id < doc_cnt(); ++id) { - size_t cur_neighbor_cnt = neighbors_index[id].neighbor_cnt; - if (cur_neighbor_cnt == 0) { - (*fixed_neighbors)[neighbors_cnt_offset + id] = 0; - continue; - } - size_t rd_size = cur_neighbor_cnt * sizeof(node_id_t); - const uint32_t *neighbors; - if (neighbors_->read(neighbors_index[id].offset, - reinterpret_cast(&neighbors), - rd_size) != rd_size) { - LOG_ERROR("Read neighbors from segment failed"); - return IndexError_ReadData; - } - - // copy level 0 neighbors to fixed size neighbors memory - auto it = fixed_neighbors->begin() + id * fixed_neighbor_cnt; - std::copy(neighbors, neighbors + cur_neighbor_cnt, it); - - (*fixed_neighbors)[neighbors_cnt_offset + id] = cur_neighbor_cnt; - total_neighbor_cnt += cur_neighbor_cnt; - } - LOG_INFO("total neighbor cnt: %zu, average neighbor cnt: %zu", - total_neighbor_cnt, total_neighbor_cnt / doc_cnt()); - - return 0; -} - -bool HnswSearcherEntity::do_crc_check( - std::vector &segments) const { - constexpr size_t blk_size = 4096; - const void *data; - for (auto &segment : segments) { - size_t offset = 0; - size_t rd_size; - uint32_t crc = 0; - while (offset < segment->data_size()) { - size_t size = std::min(blk_size, segment->data_size() - offset); - if ((rd_size = segment->read(offset, &data, size)) <= 0) { - break; - } - offset += rd_size; - crc = ailego::Crc32c::Hash(data, rd_size, crc); - } - if (crc != segment->data_crc()) { - return false; - } - } - return true; -} - -const HnswEntity::Pointer HnswSearcherEntity::clone() const { - auto vectors = vectors_->clone(); - if (ailego_unlikely(!vectors)) { - LOG_ERROR("clone segment %s failed", kGraphFeaturesSegmentId.c_str()); - return HnswEntity::Pointer(); - } - auto keys = keys_->clone(); - if (ailego_unlikely(!keys)) { - LOG_ERROR("clone segment %s failed", kGraphKeysSegmentId.c_str()); - return HnswEntity::Pointer(); - } - - auto mapping = mapping_->clone(); - if (ailego_unlikely(!mapping)) { - LOG_ERROR("clone segment %s failed", kGraphMappingSegmentId.c_str()); - return HnswEntity::Pointer(); - } - - auto neighbors = neighbors_->clone(); - if (ailego_unlikely(!neighbors)) { - LOG_ERROR("clone segment %s failed", kGraphNeighborsSegmentId.c_str()); - return HnswEntity::Pointer(); - } - auto upper_neighbors = upper_neighbors_->clone(); - if (ailego_unlikely(!neighbors)) { - LOG_ERROR("clone segment %s failed", kHnswNeighborsSegmentId.c_str()); - return HnswEntity::Pointer(); - } - auto neighbors_meta = neighbors_meta_->clone(); - if (ailego_unlikely(!neighbors_meta)) { - LOG_ERROR("clone segment %s failed", kGraphOffsetsSegmentId.c_str()); - return HnswEntity::Pointer(); - } - auto upper_neighbors_meta = upper_neighbors_meta_->clone(); - if (ailego_unlikely(!upper_neighbors_meta)) { - LOG_ERROR("clone segment %s failed", kHnswOffsetsSegmentId.c_str()); - return HnswEntity::Pointer(); - } - - SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors, - upper_neighbors_meta}; - - HnswSearcherEntity *entity = new (std::nothrow) - HnswSearcherEntity(header(), vectors, keys, mapping, neighbor_group, - fixed_neighbors_, neighbors_in_memory_enabled_); - if (ailego_unlikely(!entity)) { - LOG_ERROR("HnswSearcherEntity new failed"); - } - - return HnswEntity::Pointer(entity); -} - -} // namespace core -} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw/hnsw_searcher_entity.h b/src/core/algorithm/hnsw/hnsw_searcher_entity.h deleted file mode 100644 index 342009768..000000000 --- a/src/core/algorithm/hnsw/hnsw_searcher_entity.h +++ /dev/null @@ -1,156 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_builder_entity.h" -#include "hnsw_entity.h" - -namespace zvec { -namespace core { - -class HnswSearcherEntity : public HnswEntity { - public: - using Pointer = std::shared_ptr; - using SegmentPointer = IndexStorage::Segment::Pointer; - - public: - struct SegmentGroupParam { - SegmentGroupParam(SegmentPointer neighbors_in, - SegmentPointer neighbors_meta_in, - SegmentPointer upper_neighbors_in, - SegmentPointer upper_neighbors_meta_in) - : neighbors{neighbors_in}, - neighbors_meta{neighbors_meta_in}, - upper_neighbors{upper_neighbors_in}, - upper_neighbors_meta{upper_neighbors_meta_in} {} - - SegmentPointer neighbors{nullptr}; - SegmentPointer neighbors_meta{nullptr}; - SegmentPointer upper_neighbors{nullptr}; - SegmentPointer upper_neighbors_meta{nullptr}; - }; - - //! Constructor - HnswSearcherEntity(); - - //! Make a copy of searcher entity, to support thread-safe operation. - //! The segment in container cannot be read concurrenly - virtual const HnswEntity::Pointer clone() const override; - - //! Get primary key of the node id - virtual key_t get_key(node_id_t id) const override; - - //! Get vector local id by key - node_id_t get_id(key_t key) const; - - //! Get vector feature data by key - virtual const void *get_vector_by_key(key_t key) const override; - - //! Get vector feature data by id - virtual const void *get_vector(node_id_t id) const override; - - //! Get vector feature data by id - virtual int get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const override; - - virtual int get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const override; - virtual int get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const override; - - //! Get all vectors - const void *get_vectors() const; - - //! Get the node id's neighbors on graph level - virtual const Neighbors get_neighbors(level_t level, - node_id_t id) const override; - - virtual int load(const IndexStorage::Pointer &container, - bool check_crc) override; - - int load_segments(bool check_crc); - - virtual int cleanup(void) override; - - public: - bool is_loaded() const { - return loaded_; - } - - void set_neighbors_in_memory(bool enabled) { - neighbors_in_memory_enabled_ = enabled; - } - - //! get fixed length neighbors data - int get_fixed_neighbors(std::vector *fixed_neighbors) const; - - private: - //! Constructor - HnswSearcherEntity(const HNSWHeader &hd, const SegmentPointer &vectors, - const SegmentPointer &keys, const SegmentPointer &mapping, - const SegmentGroupParam &neighbor_group, - const std::shared_ptr &fixed_neighbors, - bool neighbors_in_memory_enabled) - : HnswEntity(hd), - vectors_(vectors), - keys_(keys), - mapping_(mapping), - neighbors_(neighbor_group.neighbors), - neighbors_meta_(neighbor_group.neighbors_meta), - upper_neighbors_(neighbor_group.upper_neighbors), - upper_neighbors_meta_(neighbor_group.upper_neighbors_meta), - neighbors_in_memory_enabled_(neighbors_in_memory_enabled) { - segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()), - IndexStorage::SegmentData(0U, 0U)); - fixed_neighbors_ = fixed_neighbors; - } - - bool do_crc_check(std::vector &segments) const; - - inline size_t neighbors_size() const { - return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); - } - - inline size_t upper_neighbors_size() const { - return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); - } - - //! If neighbors_in_memory_enabled, load the level0 neighbors to memory - int load_and_flat_neighbors(void); - - public: - HnswSearcherEntity(const HnswSearcherEntity &) = delete; - HnswSearcherEntity &operator=(const HnswSearcherEntity &) = delete; - - private: - IndexStorage::Pointer storage_{}; - - SegmentPointer vectors_{}; - SegmentPointer keys_{}; - SegmentPointer mapping_{}; - - SegmentPointer neighbors_{}; - SegmentPointer neighbors_meta_{}; - SegmentPointer upper_neighbors_{}; - SegmentPointer upper_neighbors_meta_{}; - - mutable std::vector segment_datas_{}; - std::shared_ptr fixed_neighbors_{}; // level 0 fixed size neighbors - bool neighbors_in_memory_enabled_{false}; - bool loaded_{false}; -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index 5804c7d04..e08953aa5 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -35,16 +35,16 @@ HnswStreamer::~HnswStreamer() { int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { meta_ = imeta; - meta_.set_streamer("HnswStreamer", HnswEntity::kRevision, params); + meta_.set_streamer("HnswStreamer", HnswStreamerEntityNew::kRevision, params); params.get(PARAM_HNSW_STREAMER_MAX_INDEX_SIZE, &max_index_size_); params.get(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, &upper_max_neighbor_cnt_); - float multiplier = HnswEntity::kDefaultL0MaxNeighborCntMultiplier; + float multiplier = HnswStreamerEntityNew::kDefaultL0MaxNeighborCntMultiplier; params.get(PARAM_HNSW_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, &multiplier); l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; - multiplier = HnswEntity::kDefaultNeighborPruneMultiplier; + multiplier = HnswStreamerEntityNew::kDefaultNeighborPruneMultiplier; params.get(PARAM_HNSW_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; scaling_factor_ = upper_max_neighbor_cnt_; @@ -78,30 +78,30 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { return IndexError_InvalidArgument; } else if (docs_soft_limit_ == 0UL) { docs_soft_limit_ = - docs_hard_limit_ * HnswEntity::kDefaultDocsSoftLimitRatio; + docs_hard_limit_ * HnswStreamerEntityNew::kDefaultDocsSoftLimitRatio; } if (ef_ == 0U) { - ef_ = HnswEntity::kDefaultEf; + ef_ = HnswStreamerEntityNew::kDefaultEf; } if (ef_construction_ == 0U) { - ef_construction_ = HnswEntity::kDefaultEfConstruction; + ef_construction_ = HnswStreamerEntityNew::kDefaultEfConstruction; } if (upper_max_neighbor_cnt_ == 0U) { - upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; + upper_max_neighbor_cnt_ = HnswStreamerEntityNew::kDefaultUpperMaxNeighborCnt; } - if (upper_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { + if (upper_max_neighbor_cnt_ > HnswStreamerEntityNew::kMaxNeighborCnt) { LOG_ERROR("[%s] must be in range (0,%d)", PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT.c_str(), - HnswEntity::kMaxNeighborCnt); + HnswStreamerEntityNew::kMaxNeighborCnt); return IndexError_InvalidArgument; } if (l0_max_neighbor_cnt_ == 0U) { - l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt; + l0_max_neighbor_cnt_ = HnswStreamerEntityNew::kDefaultL0MaxNeighborCnt; } - if (l0_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) { + if (l0_max_neighbor_cnt_ > HnswStreamerEntityNew::kMaxNeighborCnt) { LOG_ERROR("MaxL0NeighborCnt must be in range (0,%d)", - HnswEntity::kMaxNeighborCnt); + HnswStreamerEntityNew::kMaxNeighborCnt); return IndexError_InvalidArgument; } if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { @@ -119,7 +119,7 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { } if (scaling_factor_ == 0U) { - scaling_factor_ = HnswEntity::kDefaultScalingFactor; + scaling_factor_ = HnswStreamerEntityNew::kDefaultScalingFactor; } if (scaling_factor_ < 5 || scaling_factor_ > 1000) { LOG_ERROR("[%s] must be in range [5,1000]", @@ -144,11 +144,11 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { prune_cnt = upper_max_neighbor_cnt_; } if (chunk_size_ == 0UL) { - chunk_size_ = HnswEntity::kDefaultChunkSize; + chunk_size_ = HnswStreamerEntityNew::kDefaultChunkSize; } - if (chunk_size_ > HnswEntity::kMaxChunkSize) { + if (chunk_size_ > HnswStreamerEntityNew::kMaxChunkSize) { LOG_ERROR("[%s] must be < %zu", PARAM_HNSW_STREAMER_CHUNK_SIZE.c_str(), - HnswEntity::kMaxChunkSize); + HnswStreamerEntityNew::kMaxChunkSize); return IndexError_InvalidArgument; } @@ -215,20 +215,20 @@ int HnswStreamer::cleanup(void) { } max_index_size_ = 0UL; - docs_hard_limit_ = HnswEntity::kDefaultDocsHardLimit; + docs_hard_limit_ = HnswStreamerEntityNew::kDefaultDocsHardLimit; docs_soft_limit_ = 0UL; - upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt; - l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt; - ef_ = HnswEntity::kDefaultEf; - ef_construction_ = HnswEntity::kDefaultEfConstruction; + upper_max_neighbor_cnt_ = HnswStreamerEntityNew::kDefaultUpperMaxNeighborCnt; + l0_max_neighbor_cnt_ = HnswStreamerEntityNew::kDefaultL0MaxNeighborCnt; + ef_ = HnswStreamerEntityNew::kDefaultEf; + ef_construction_ = HnswStreamerEntityNew::kDefaultEfConstruction; bf_enabled_ = false; - scaling_factor_ = HnswEntity::kDefaultScalingFactor; - bruteforce_threshold_ = HnswEntity::kDefaultBruteForceThreshold; - max_scan_limit_ = HnswEntity::kDefaultMaxScanLimit; - min_scan_limit_ = HnswEntity::kDefaultMinScanLimit; - chunk_size_ = HnswEntity::kDefaultChunkSize; - bf_negative_prob_ = HnswEntity::kDefaultBFNegativeProbability; - max_scan_ratio_ = HnswEntity::kDefaultScanRatio; + scaling_factor_ = HnswStreamerEntityNew::kDefaultScalingFactor; + bruteforce_threshold_ = HnswStreamerEntityNew::kDefaultBruteForceThreshold; + max_scan_limit_ = HnswStreamerEntityNew::kDefaultMaxScanLimit; + min_scan_limit_ = HnswStreamerEntityNew::kDefaultMinScanLimit; + chunk_size_ = HnswStreamerEntityNew::kDefaultChunkSize; + bf_negative_prob_ = HnswStreamerEntityNew::kDefaultBFNegativeProbability; + max_scan_ratio_ = HnswStreamerEntityNew::kDefaultScanRatio; state_ = STATE_INIT; check_crc_enabled_ = false; filter_same_key_ = false; @@ -342,7 +342,7 @@ int HnswStreamer::dump(const IndexDumper::Pointer &dumper) { shared_mutex_.lock(); AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); - meta_.set_searcher("HnswSearcher", HnswEntity::kRevision, ailego::Params()); + meta_.set_searcher("HnswSearcher", HnswStreamerEntityNew::kRevision, ailego::Params()); int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); if (ret != 0) { @@ -358,7 +358,7 @@ IndexStreamer::Context::Pointer HnswStreamer::create_context(void) const { return Context::Pointer(); } - HnswEntity::Pointer entity = entity_.clone(); + HnswStreamerEntityNew::Pointer entity = entity_.clone(); if (ailego_unlikely(!entity)) { LOG_ERROR("CreateContext clone init failed"); return Context::Pointer(); @@ -401,7 +401,7 @@ IndexProvider::Pointer HnswStreamer::create_provider(void) const { auto entity = entity_.clone(); if (ailego_unlikely(!entity)) { - LOG_ERROR("Clone HnswEntity failed"); + LOG_ERROR("Clone HnswStreamerEntityNew failed"); return nullptr; } return Provider::Pointer( @@ -409,7 +409,7 @@ IndexProvider::Pointer HnswStreamer::create_provider(void) const { } int HnswStreamer::update_context(HnswContext *ctx) const { - const HnswEntity::Pointer entity = entity_.clone(); + const HnswStreamerEntityNew::Pointer entity = entity_.clone(); if (!entity) { LOG_ERROR("Failed to clone search context entity"); return IndexError_Runtime; @@ -642,7 +642,7 @@ void HnswStreamer::print_debug_info() { if (entity_.get_key(id) == kInvalidKey) { continue; } - Neighbors neighbours = entity_.get_neighbors(0, id); + Neighbors neighbours = entity_.get_neighbors_new(0, id); std::cout << "node: " << id << "; "; if (neighbours.size() == 0) std::cout << std::endl; for (uint32_t i = 0; i < neighbours.size(); ++i) { diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h index b81106daf..48f377e67 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.h +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -16,7 +16,7 @@ #include #include #include "hnsw_algorithm.h" -#include "hnsw_streamer_entity.h" +#include "hnsw_streamer_entity_new.h" namespace zvec { namespace core { @@ -98,12 +98,12 @@ class HnswStreamer : public IndexStreamer { //! Fetch vector by id virtual const void *get_vector_by_id(uint32_t id) const override { - return entity_.get_vector(id); + return entity_.get_vector_new(id); } virtual int get_vector_by_id( const uint32_t id, IndexStorage::MemoryBlock &block) const override { - return entity_.get_vector(id, block); + return entity_.get_vector_new(id, block); } //! Open index from file path @@ -181,7 +181,7 @@ class HnswStreamer : public IndexStreamer { } }; - HnswStreamerEntity entity_; + HnswStreamerEntityNew entity_; HnswAlgorithm::UPointer alg_; IndexMeta meta_{}; IndexMetric::Pointer metric_{}; diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc b/src/core/algorithm/hnsw/hnsw_streamer_entity_new.cc similarity index 57% rename from src/core/algorithm/hnsw/hnsw_streamer_entity.cc rename to src/core/algorithm/hnsw/hnsw_streamer_entity_new.cc index 71e2e477b..e1875fffe 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity_new.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "hnsw_streamer_entity.h" +#include "hnsw_streamer_entity_new.h" #include // #define DEBUG_PRINT @@ -20,12 +20,77 @@ namespace zvec { namespace core { -HnswStreamerEntity::HnswStreamerEntity(IndexStreamer::Stats &stats) +const std::string HnswStreamerEntityNew::kGraphHeaderSegmentId = "graph.header"; +const std::string HnswStreamerEntityNew::kGraphFeaturesSegmentId = + "graph.features"; +const std::string HnswStreamerEntityNew::kGraphKeysSegmentId = "graph.keys"; +const std::string HnswStreamerEntityNew::kGraphNeighborsSegmentId = + "graph.neighbors"; +const std::string HnswStreamerEntityNew::kGraphOffsetsSegmentId = + "graph.offsets"; +const std::string HnswStreamerEntityNew::kGraphMappingSegmentId = + "graph.mapping"; +const std::string HnswStreamerEntityNew::kHnswHeaderSegmentId = "hnsw.header"; +const std::string HnswStreamerEntityNew::kHnswNeighborsSegmentId = + "hnsw.neighbors"; +const std::string HnswStreamerEntityNew::kHnswOffsetsSegmentId = "hnsw.offsets"; + +int64_t HnswStreamerEntityNew::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, + size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %lu, actual: %lu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return len + padding_size; +} + +int64_t HnswStreamerEntityNew::dump_header(const IndexDumper::Pointer &dumper, + const HNSWHeader &hd) const { + //! dump basic graph header. header is aligned and does not need padding + int64_t graph_hd_size = + dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size); + if (graph_hd_size < 0) { + return graph_hd_size; + } + + //! dump basic graph header. header is aligned and does not need padding + int64_t hnsw_hd_size = + dump_segment(dumper, kHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size); + if (hnsw_hd_size < 0) { + return hnsw_hd_size; + } + + return graph_hd_size + hnsw_hd_size; +} + + +HnswStreamerEntityNew::HnswStreamerEntityNew(IndexStreamer::Stats &stats) : stats_(stats) {} -HnswStreamerEntity::~HnswStreamerEntity() {} +HnswStreamerEntityNew::~HnswStreamerEntityNew() {} -int HnswStreamerEntity::init(size_t max_doc_cnt) { +int HnswStreamerEntityNew::init(size_t max_doc_cnt) { if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) { LOG_ERROR("scalingFactor=%zu is too small", scaling_factor()); return IndexError_InvalidArgument; @@ -37,7 +102,7 @@ int HnswStreamerEntity::init(size_t max_doc_cnt) { keys_map_lock_ = std::make_shared(); keys_map_ = std::make_shared>(); if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) { - LOG_ERROR("HnswStreamerEntity new object failed"); + LOG_ERROR("HnswStreamerEntityNew new object failed"); return IndexError_NoMemory; } keys_map_->set_empty_key(kInvalidKey); @@ -53,7 +118,7 @@ int HnswStreamerEntity::init(size_t max_doc_cnt) { return 0; } -int HnswStreamerEntity::cleanup() { +int HnswStreamerEntityNew::cleanup() { std::lock_guard lock(mutex_); mutable_header()->clear(); chunk_size_ = kDefaultChunkSize; @@ -77,7 +142,7 @@ int HnswStreamerEntity::cleanup() { return 0; } -int HnswStreamerEntity::update_neighbors( +int HnswStreamerEntityNew::update_neighbors( level_t level, node_id_t id, const std::vector> &neighbors) { std::vector buffer(neighbor_size_); @@ -100,8 +165,8 @@ int HnswStreamerEntity::update_neighbors( return 0; } -const Neighbors HnswStreamerEntity::get_neighbors(level_t level, - node_id_t id) const { +const Neighbors HnswStreamerEntityNew::get_neighbors(level_t level, + node_id_t id) const { Chunk *chunk = nullptr; size_t offset = 0UL; size_t neighbor_size = neighbor_size_; @@ -130,8 +195,19 @@ const Neighbors HnswStreamerEntity::get_neighbors(level_t level, return Neighbors(std::move(neighbor_block)); } +const Neighbors HnswStreamerEntityNew::get_neighbors_new(level_t level, + node_id_t id) const { + if (id) { + return get_neighbors(level, id); + } else { + const void *src = neighbors_value_ptr_->data() + id * neighbor_size_; + const NeighborsHeader *header = reinterpret_cast(src); + return Neighbors(header->neighbor_cnt, header->neighbors); + } +} + //! Get vector data by key -const void *HnswStreamerEntity::get_vector(node_id_t id) const { +const void *HnswStreamerEntityNew::get_vector(node_id_t id) const { auto loc = get_vector_chunk_loc(id); const void *vec = nullptr; ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); @@ -149,8 +225,13 @@ const void *HnswStreamerEntity::get_vector(node_id_t id) const { return vec; } -int HnswStreamerEntity::get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const { +const void *HnswStreamerEntityNew::get_vector_new(node_id_t id) const { + return vector_value_ptr_->data() + vector_size() * id; + // return get_vector(id); +} + +int HnswStreamerEntityNew::get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const { for (auto i = 0U; i < count; ++i) { auto loc = get_vector_chunk_loc(ids[i]); ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); @@ -169,8 +250,8 @@ int HnswStreamerEntity::get_vector(const node_id_t *ids, uint32_t count, return 0; } -int HnswStreamerEntity::get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const { +int HnswStreamerEntityNew::get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const { auto loc = get_vector_chunk_loc(id); ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), @@ -187,7 +268,27 @@ int HnswStreamerEntity::get_vector(const node_id_t id, return 0; } -int HnswStreamerEntity::get_vector( +int HnswStreamerEntityNew::get_vector_new( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + const void *data = vector_value_ptr_->data() + vector_size() * id; + block.reset((void *)data); + return 0; + // return get_vector(id, block); +} + +int HnswStreamerEntityNew::get_vector_new( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.reserve(count); + for (int i = 0; i < count; i++) { + const void *data = vector_value_ptr_->data() + vector_size() * ids[i]; + vec_blocks.emplace_back((void *)data); + } + return 0; + // return get_vector(ids, count, vec_blocks); +} + +int HnswStreamerEntityNew::get_vector( const node_id_t *ids, uint32_t count, std::vector &vec_blocks) const { vec_blocks.resize(count); @@ -210,7 +311,7 @@ int HnswStreamerEntity::get_vector( return 0; } -key_t HnswStreamerEntity::get_key(node_id_t id) const { +key_t HnswStreamerEntityNew::get_key(node_id_t id) const { if (use_key_info_map_) { auto loc = get_key_chunk_loc(id); IndexStorage::MemoryBlock key_block; @@ -230,8 +331,8 @@ key_t HnswStreamerEntity::get_key(node_id_t id) const { } } -void HnswStreamerEntity::add_neighbor(level_t level, node_id_t id, - uint32_t size, node_id_t neighbor_id) { +void HnswStreamerEntityNew::add_neighbor(level_t level, node_id_t id, + uint32_t size, node_id_t neighbor_id) { auto loc = get_neighbor_chunk_loc(level, id); size_t offset = loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t); @@ -252,7 +353,7 @@ void HnswStreamerEntity::add_neighbor(level_t level, node_id_t id, return; } -int HnswStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { +int HnswStreamerEntityNew::init_chunks(const Chunk::Pointer &header_chunk) { if (header_chunk->data_size() < header_size()) { LOG_ERROR("Invalid header chunk size"); return IndexError_InvalidFormat; @@ -297,8 +398,8 @@ int HnswStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { return 0; } -int HnswStreamerEntity::open(IndexStorage::Pointer stg, uint64_t max_index_size, - bool check_crc) { +int HnswStreamerEntityNew::open(IndexStorage::Pointer stg, + uint64_t max_index_size, bool check_crc) { std::lock_guard lock(mutex_); bool huge_page = stg->isHugePage(); LOG_DEBUG("huge_page: %d", (int)huge_page); @@ -378,13 +479,25 @@ int HnswStreamerEntity::open(IndexStorage::Pointer stg, uint64_t max_index_size, } } } + vector_value_ptr_ = std::make_shared(); + vector_value_ptr_->reserve(vector_size() * doc_cnt()); + for (int i = 0; i < doc_cnt(); i++) { + vector_value_ptr_->append((const char *)get_vector(i), vector_size()); + } + + neighbors_value_ptr_ = std::make_shared(); + neighbors_value_ptr_->reserve(neighbor_size_ * doc_cnt()); + for (int i = 0; i < doc_cnt(); i++) { + Neighbors neighbor = get_neighbors(0, i); + neighbors_value_ptr_->append((const char *)neighbor.neighbor_block.data(), neighbor_size_); + } stats_.set_loaded_count(doc_cnt()); return 0; } -int HnswStreamerEntity::close() { +int HnswStreamerEntityNew::close() { LOG_DEBUG("close index"); std::lock_guard lock(mutex_); @@ -399,7 +512,7 @@ int HnswStreamerEntity::close() { return broker_->close(); } -int HnswStreamerEntity::flush(uint64_t checkpoint) { +int HnswStreamerEntityNew::flush(uint64_t checkpoint) { LOG_INFO("Flush index, curDocs=%u", doc_cnt()); std::lock_guard lock(mutex_); @@ -412,7 +525,7 @@ int HnswStreamerEntity::flush(uint64_t checkpoint) { return 0; } -int HnswStreamerEntity::dump(const IndexDumper::Pointer &dumper) { +int HnswStreamerEntityNew::dump(const IndexDumper::Pointer &dumper) { LOG_INFO("Dump index, curDocs=%u", doc_cnt()); //! sort by keys, to support get_vector by key in searcher @@ -439,7 +552,7 @@ int HnswStreamerEntity::dump(const IndexDumper::Pointer &dumper) { return 0; } -int HnswStreamerEntity::check_hnsw_index(const HNSWHeader *hd) const { +int HnswStreamerEntityNew::check_hnsw_index(const HNSWHeader *hd) const { if (l0_neighbor_cnt() != hd->l0_neighbor_cnt() || upper_neighbor_cnt() != hd->upper_neighbor_cnt()) { LOG_ERROR("Param neighbor cnt: %zu:%zu mismatch index previous %zu:%zu", @@ -482,8 +595,8 @@ int HnswStreamerEntity::check_hnsw_index(const HNSWHeader *hd) const { return 0; } -int HnswStreamerEntity::add_vector(level_t level, key_t key, const void *vec, - node_id_t *id) { +int HnswStreamerEntityNew::add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) { Chunk::Pointer node_chunk; size_t chunk_offset = -1UL; @@ -555,8 +668,8 @@ int HnswStreamerEntity::add_vector(level_t level, key_t key, const void *vec, return 0; } -int HnswStreamerEntity::add_vector_with_id(level_t level, node_id_t id, - const void *vec) { +int HnswStreamerEntityNew::add_vector_with_id(level_t level, node_id_t id, + const void *vec) { Chunk::Pointer node_chunk; size_t chunk_offset = -1UL; key_t key = id; @@ -658,21 +771,22 @@ int HnswStreamerEntity::add_vector_with_id(level_t level, node_id_t id, return 0; } -void HnswStreamerEntity::update_ep_and_level(node_id_t ep, level_t level) { - HnswEntity::update_ep_and_level(ep, level); +void HnswStreamerEntityNew::update_ep_and_level(node_id_t ep, level_t level) { + base_header_.hnsw.entry_point = ep; + base_header_.hnsw.max_level = level; flush_header(); return; } -const HnswEntity::Pointer HnswStreamerEntity::clone() const { +const HnswStreamerEntityNew::Pointer HnswStreamerEntityNew::clone() const { std::vector node_chunks; node_chunks.reserve(node_chunks_.size()); for (size_t i = 0UL; i < node_chunks_.size(); ++i) { node_chunks.emplace_back(node_chunks_[i]->clone()); if (ailego_unlikely(!node_chunks[i])) { - LOG_ERROR("HnswStreamerEntity get chunk failed in clone"); - return HnswEntity::Pointer(); + LOG_ERROR("HnswStreamerEntityNew get chunk failed in clone"); + return HnswStreamerEntityNew::Pointer(); } } @@ -681,21 +795,300 @@ const HnswEntity::Pointer HnswStreamerEntity::clone() const { for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); if (ailego_unlikely(!upper_neighbor_chunks[i])) { - LOG_ERROR("HnswStreamerEntity get chunk failed in clone"); - return HnswEntity::Pointer(); + LOG_ERROR("HnswStreamerEntityNew get chunk failed in clone"); + return HnswStreamerEntityNew::Pointer(); } } - HnswStreamerEntity *entity = new (std::nothrow) HnswStreamerEntity( + HnswStreamerEntityNew *entity = new (std::nothrow) HnswStreamerEntityNew( stats_, header(), chunk_size_, node_index_mask_bits_, upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_, - std::move(node_chunks), std::move(upper_neighbor_chunks), broker_); + std::move(node_chunks), std::move(upper_neighbor_chunks), broker_, + vector_value_ptr_, neighbors_value_ptr_); if (ailego_unlikely(!entity)) { - LOG_ERROR("HnswStreamerEntity new failed"); + LOG_ERROR("HnswStreamerEntityNew new failed"); + } + return HnswStreamerEntityNew::Pointer(entity); +} + +int64_t HnswStreamerEntityNew::dump_mapping_segment( + const IndexDumper::Pointer &dumper, const key_t *keys) const { + std::vector mapping(doc_cnt()); + + std::iota(mapping.begin(), mapping.end(), 0U); + std::sort(mapping.begin(), mapping.end(), + [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; }); + + size_t size = mapping.size() * sizeof(node_id_t); + + return dump_segment(dumper, kGraphMappingSegmentId, mapping.data(), size); +} + +int64_t HnswStreamerEntityNew::dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const { + HNSWHeader dump_hd(header()); + + dump_hd.graph.node_size = AlignSize(vector_size()); + + std::vector n2o_mapping; // map new id to origin id + std::vector o2n_mapping; // map origin id to new id + if (!o2n_mapping.empty()) { + dump_hd.hnsw.entry_point = o2n_mapping[entry_point()]; + } + + //! Dump header + int64_t hd_size = dump_header(dumper, dump_hd); + if (hd_size < 0) { + return hd_size; + } + + //! Dump vectors + int64_t vecs_size = dump_vectors(dumper, n2o_mapping); + if (vecs_size < 0) { + return vecs_size; + } + + //! Dump neighbors + auto neighbors_size = + dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping); + if (neighbors_size < 0) { + return neighbors_size; + } + //! free memory + n2o_mapping = std::vector(); + o2n_mapping = std::vector(); + + //! Dump keys + size_t key_segment_size = doc_cnt() * sizeof(key_t); + int64_t keys_size = + dump_segment(dumper, kGraphKeysSegmentId, keys, key_segment_size); + if (keys_size < 0) { + return keys_size; + } + + //! Dump mapping + int64_t mapping_size = dump_mapping_segment(dumper, keys); + if (mapping_size < 0) { + return mapping_size; + } + + return hd_size + keys_size + vecs_size + neighbors_size + mapping_size; +} + +int64_t HnswStreamerEntityNew::dump_vectors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + size_t vector_dump_size = vector_size(); + + size_t padding_size = AlignSize(vector_dump_size) - vector_dump_size; + + std::vector padding(padding_size); + memset(padding.data(), 0, sizeof(char) * padding_size); + const void *data = nullptr; + uint32_t crc = 0U; + size_t vecs_size = 0UL; + + //! dump vectors + for (node_id_t id = 0; id < doc_cnt(); ++id) { + data = get_vector(reorder_mapping.empty() ? id : reorder_mapping[id]); + if (ailego_unlikely(!data)) { + return IndexError_ReadData; + } + size_t len = dumper->write(data, vector_size()); + if (len != vector_size()) { + LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, + vector_size()); + return IndexError_WriteData; + } + + crc = ailego::Crc32c::Hash(data, vector_size(), crc); + vecs_size += vector_size(); + + if (padding_size == 0) { + continue; + } + + len = dumper->write(padding.data(), padding_size); + if (len != padding_size) { + LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, padding_size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(padding.data(), padding_size, crc); + vecs_size += padding_size; + } + + int ret = dumper->append(kGraphFeaturesSegmentId, vecs_size, 0UL, crc); + if (ret != 0) { + LOG_ERROR("Dump vectors segment meta failed, ret %d", ret); + return ret; + } + + return vecs_size; +} + +int64_t HnswStreamerEntityNew::dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector graph_meta; + graph_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector mapping(l0_neighbor_cnt()); + + uint32_t min_neighbor_count = 10000; + uint32_t max_neighbor_count = 0; + size_t sum_neighbor_count = 0; + + for (node_id_t id = 0; id < doc_cnt(); ++id) { + const Neighbors neighbors = + get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(), + "invalid neighbors"); + + uint32_t neighbor_count = neighbors.size(); + if (neighbor_count < min_neighbor_count) { + min_neighbor_count = neighbor_count; + } + if (neighbor_count > max_neighbor_count) { + max_neighbor_count = neighbor_count; + } + sum_neighbor_count += neighbor_count; + + graph_meta.emplace_back(offset, neighbor_count); + size_t size = neighbors.size() * sizeof(node_id_t); + const node_id_t *data = &neighbors[0]; + if (!neighbor_mapping.empty()) { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + mapping[i] = neighbor_mapping[neighbors[i]]; + } + data = mapping.data(); + } + if (dumper->write(data, size) != size) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, size); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(data, size, crc); + offset += size; + } + + uint32_t average_neighbor_count = 0; + if (doc_cnt() > 0) { + average_neighbor_count = sum_neighbor_count / doc_cnt(); + } + LOG_INFO( + "Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] " + "average_neighbor_count[%u]", + min_neighbor_count, max_neighbor_count, average_neighbor_count); + + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; + } + ret = dumper->append(kGraphNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", + kGraphNeighborsSegmentId.c_str(), ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kGraphOffsetsSegmentId, graph_meta.data(), + graph_meta.size() * sizeof(GraphNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +int64_t HnswStreamerEntityNew::dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + std::vector hnsw_meta; + hnsw_meta.reserve(doc_cnt()); + size_t offset = 0; + uint32_t crc = 0; + std::vector buffer(upper_neighbor_cnt() + 1); + for (node_id_t id = 0; id < doc_cnt(); ++id) { + node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id]; + auto level = get_level(new_id); + if (level == 0) { + hnsw_meta.emplace_back(0U, 0U); + continue; + } + hnsw_meta.emplace_back(offset, level); + ailego_assert_with((size_t)level < kMaxGraphLayers, "invalid level"); + for (level_t cur_level = 1; cur_level <= level; ++cur_level) { + const Neighbors neighbors = get_neighbors(cur_level, new_id); + ailego_assert_with(!!neighbors.data, "invalid neighbors"); + ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level), + "invalid neighbors"); + memset(buffer.data(), 0, sizeof(node_id_t) * buffer.size()); + buffer[0] = neighbors.size(); + if (neighbor_mapping.empty()) { + memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t)); + } else { + for (node_id_t i = 0; i < neighbors.size(); ++i) { + buffer[i + 1] = neighbor_mapping[neighbors[i]]; + } + } + if (dumper->write(buffer.data(), sizeof(node_id_t) * buffer.size()) != + sizeof(node_id_t) * buffer.size()) { + LOG_ERROR("Dump graph neighbor id=%u failed, size %lu", id, + sizeof(node_id_t) * buffer.size()); + return IndexError_WriteData; + } + crc = ailego::Crc32c::Hash(buffer.data(), + sizeof(node_id_t) * buffer.size(), crc); + offset += sizeof(node_id_t) * buffer.size(); + } + } + size_t padding_size = 0; + int ret = CalcAndAddPadding(dumper, offset, &padding_size); + if (ret != 0) { + return ret; } - return HnswEntity::Pointer(entity); + + ret = dumper->append(kHnswNeighborsSegmentId, offset, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s failed, ret %d", kHnswNeighborsSegmentId.c_str(), + ret); + return ret; + } + + //! dump level 0 neighbors meta + auto len = dump_segment(dumper, kHnswOffsetsSegmentId, hnsw_meta.data(), + hnsw_meta.size() * sizeof(HnswNeighborMeta)); + if (len < 0) { + return len; + } + + return len + offset + padding_size; +} + +int HnswStreamerEntityNew::CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, + size_t *padding_size) { + *padding_size = AlignSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + + std::string padding(*padding_size, '\0'); + if (dumper->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %lu", *padding_size); + return IndexError_WriteData; + } + return 0; } + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.h b/src/core/algorithm/hnsw/hnsw_streamer_entity_new.h similarity index 61% rename from src/core/algorithm/hnsw/hnsw_streamer_entity.h rename to src/core/algorithm/hnsw/hnsw_streamer_entity_new.h index 1a01b1416..11707a029 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.h +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity_new.h @@ -28,82 +28,74 @@ namespace zvec { namespace core { -//! HnswStreamerEntity manage vector data, pkey, and node's neighbors -class HnswStreamerEntity : public HnswEntity { - public: +//! HnswStreamerEntityNew manage vector data, pkey, and node's neighbors +class HnswStreamerEntityNew { + public: // override + typedef std::shared_ptr Pointer; + //! Cleanup //! return 0 on success, or errCode in failure - virtual int cleanup() override; + int cleanup(); //! Make a copy of streamer entity, to support thread-safe operation. //! The segment in container cannot be read concurrenly - virtual const HnswEntity::Pointer clone() const override; + const HnswStreamerEntityNew::Pointer clone() const; //! Get primary key of the node id - virtual key_t get_key(node_id_t id) const override; + key_t get_key(node_id_t id) const; //! Get vector feature data by key - virtual const void *get_vector(node_id_t id) const override; + const void *get_vector(node_id_t id) const; + + const void *get_vector_new(node_id_t id) const; //! Get vectors feature data by local ids - virtual int get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const override; + int get_vector(const node_id_t *ids, uint32_t count, const void **vecs) const; - virtual int get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const override; + int get_vector(const node_id_t id, IndexStorage::MemoryBlock &block) const; - virtual int get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const override; + int get_vector(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + + int get_vector_new(const node_id_t id, + IndexStorage::MemoryBlock &block) const; + + int get_vector_new(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; //! Get the node id's neighbors on graph level //! Note: the neighbors cannot be modified, using the following //! method to get WritableNeighbors if want to - virtual const Neighbors get_neighbors(level_t level, - node_id_t id) const override; + const Neighbors get_neighbors(level_t level, node_id_t id) const; + const Neighbors get_neighbors_new(level_t level, node_id_t id) const; //! Add vector and key to hnsw entity, and local id will be saved in id - virtual int add_vector(level_t level, key_t key, const void *vec, - node_id_t *id) override; + int add_vector(level_t level, key_t key, const void *vec, node_id_t *id); //! Add vector and id to hnsw entity - virtual int add_vector_with_id(level_t level, node_id_t id, - const void *vec) override; + int add_vector_with_id(level_t level, node_id_t id, const void *vec); - virtual int update_neighbors( + int update_neighbors( level_t level, node_id_t id, - const std::vector> &neighbors) override; + const std::vector> &neighbors); //! Append neighbor_id to node id neighbors on level //! Notice: the caller must be ensure the neighbors not full - virtual void add_neighbor(level_t level, node_id_t id, uint32_t size, - node_id_t neighbor_id) override; + void add_neighbor(level_t level, node_id_t id, uint32_t size, + node_id_t neighbor_id); //! Dump index by dumper - virtual int dump(const IndexDumper::Pointer &dumper) override; - - virtual void update_ep_and_level(node_id_t ep, level_t level) override; - - void set_use_key_info_map(bool use_id_map) { - use_key_info_map_ = use_id_map; - LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); - } - - public: - //! Constructor - HnswStreamerEntity(IndexStreamer::Stats &stats); + int dump(const IndexDumper::Pointer &dumper); - //! Destructor - ~HnswStreamerEntity(); + void update_ep_and_level(node_id_t ep, level_t level); - //! Get vector feature data by key - virtual const void *get_vector_by_key(key_t key) const override { + const void *get_vector_by_key(key_t key) const { auto id = get_id(key); return id == kInvalidNodeId ? nullptr : get_vector(id); } - virtual int get_vector_by_key( - const key_t key, IndexStorage::MemoryBlock &block) const override { + int get_vector_by_key(const key_t key, + IndexStorage::MemoryBlock &block) const { auto id = get_id(key); if (id != kInvalidNodeId) { return get_vector(id, block); @@ -112,6 +104,240 @@ class HnswStreamerEntity : public HnswEntity { } } + public: // hnsw entity public + //! Get max neighbor size of graph level + inline size_t neighbor_cnt(level_t level) const { + return level == 0 ? base_header_.graph.l0_neighbor_count + : base_header_.hnsw.upper_neighbor_count; + } + + //! get max neighbor size of graph level 0 + inline size_t l0_neighbor_cnt() const { + return base_header_.graph.l0_neighbor_count; + } + + //! get min neighbor size of graph + inline size_t min_neighbor_cnt() const { + return base_header_.graph.min_neighbor_count; + } + + //! get upper neighbor size of graph level other than 0 + inline size_t upper_neighbor_cnt() const { + return base_header_.hnsw.upper_neighbor_count; + } + + //! Get current total doc of the hnsw graph + inline node_id_t *mutable_doc_cnt() { + return &base_header_.graph.doc_count; + } + + inline node_id_t doc_cnt() const { + return base_header_.graph.doc_count; + } + + //! Get hnsw graph scaling params + inline size_t scaling_factor() const { + return base_header_.hnsw.scaling_factor; + } + + //! Get prune_size + inline size_t prune_cnt() const { + return base_header_.graph.prune_neighbor_count; + } + + //! Current entity of top level graph + inline node_id_t entry_point() const { + return base_header_.hnsw.entry_point; + } + + //! Current max graph level + inline level_t cur_max_level() const { + return base_header_.hnsw.max_level; + } + + //! Retrieve index vector size + size_t vector_size() const { + return base_header_.graph.vector_size; + } + + //! Retrieve node size + size_t node_size() const { + return base_header_.graph.node_size; + } + + //! Retrieve ef constuction + size_t ef_construction() const { + return base_header_.graph.ef_construction; + } + + void set_vector_size(size_t size) { + base_header_.graph.vector_size = size; + } + + void set_prune_cnt(size_t v) { + base_header_.graph.prune_neighbor_count = v; + } + + void set_scaling_factor(size_t val) { + base_header_.hnsw.scaling_factor = val; + } + + void set_l0_neighbor_cnt(size_t cnt) { + base_header_.graph.l0_neighbor_count = cnt; + } + + void set_min_neighbor_cnt(size_t cnt) { + base_header_.graph.min_neighbor_count = cnt; + } + + void set_upper_neighbor_cnt(size_t cnt) { + base_header_.hnsw.upper_neighbor_count = cnt; + } + + void set_ef_construction(size_t ef) { + base_header_.graph.ef_construction = ef; + } + + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size); + + protected: + inline const HNSWHeader &header() const { + return base_header_; + } + + inline HNSWHeader *mutable_header() { + return &base_header_; + } + + inline size_t header_size() const { + return sizeof(base_header_); + } + + void set_node_size(size_t size) { + base_header_.graph.node_size = size; + } + + //! Dump all segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segments( + const IndexDumper::Pointer &dumper, key_t *keys, + const std::function &get_level) const; + + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + static inline size_t AlignPageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::PageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + static inline size_t AlignHugePageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + private: + //! dump mapping segment, for get_vector_by_key in provider + int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const; + + //! dump hnsw head by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_header(const IndexDumper::Pointer &dumper, + const HNSWHeader &hd) const; + + //! dump vectors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_vectors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + //! dump hnsw neighbors by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_neighbors(const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping); + if (len1 < 0) { + return len1; + } + auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping, + neighbor_mapping); + if (len2 < 0) { + return len2; + } + + return len1 + len2; + } + + //! dump segment by dumper + //! Return dump size if success, errno(<0) in failure + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + + //! Dump level 0 neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_graph_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + //! Dump upper level neighbors + //! Return dump size if success, errno(<0) in failure + int64_t dump_upper_neighbors( + const IndexDumper::Pointer &dumper, + const std::function &get_level, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + public: + const static std::string kGraphHeaderSegmentId; + const static std::string kGraphFeaturesSegmentId; + const static std::string kGraphKeysSegmentId; + const static std::string kGraphNeighborsSegmentId; + const static std::string kGraphOffsetsSegmentId; + const static std::string kGraphMappingSegmentId; + const static std::string kHnswHeaderSegmentId; + const static std::string kHnswNeighborsSegmentId; + const static std::string kHnswOffsetsSegmentId; + + constexpr static uint32_t kRevision = 0U; + constexpr static size_t kMaxGraphLayers = 15; + constexpr static uint32_t kDefaultEfConstruction = 500; + constexpr static uint32_t kDefaultEf = 500; + constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW + constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100; + constexpr static uint32_t kMaxNeighborCnt = 65535; + constexpr static float kDefaultScanRatio = 0.1f; + constexpr static uint32_t kDefaultMinScanLimit = 10000; + constexpr static uint32_t kDefaultMaxScanLimit = + std::numeric_limits::max(); + constexpr static float kDefaultBFNegativeProbability = 0.001f; + constexpr static uint32_t kDefaultScalingFactor = 50U; + constexpr static uint32_t kDefaultBruteForceThreshold = 1000U; + constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion + constexpr static float kDefaultDocsSoftLimitRatio = 0.9f; + constexpr static size_t kMaxChunkSize = 0xFFFFFFFF; + constexpr static size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL; + constexpr static size_t kDefaultMaxChunkCnt = 50000UL; + constexpr static float kDefaultNeighborPruneMultiplier = + 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier + constexpr static float kDefaultL0MaxNeighborCntMultiplier = + 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier + + public: + //! Constructor + HnswStreamerEntityNew(IndexStreamer::Stats &stats); + + //! Destructor + ~HnswStreamerEntityNew(); + + //! Get vector feature data by key + + //! Init entity int init(size_t max_doc_cnt); @@ -127,6 +353,11 @@ class HnswStreamerEntity : public HnswEntity { //! return 0 on success, or errCode in failure int close(); + void set_use_key_info_map(bool use_id_map) { + use_key_info_map_ = use_id_map; + LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); + } + //! Set meta information from entity int set_index_meta(const IndexMeta &meta) const { return IndexHelper::SerializeToStorage(meta, broker_->storage().get()); @@ -215,17 +446,19 @@ class HnswStreamerEntity : public HnswEntity { using NIHashMapPointer = std::shared_ptr; //! Private construct, only be called by clone method - HnswStreamerEntity(IndexStreamer::Stats &stats, const HNSWHeader &hd, - size_t chunk_size, uint32_t node_index_mask_bits, - uint32_t upper_neighbor_mask_bits, bool filter_same_key, - bool get_vector_enabled, - const NIHashMapPointer &upper_neighbor_index, - std::shared_ptr &keys_map_lock, - const HashMapPointer &keys_map, - bool use_key_info_map, - std::vector &&node_chunks, - std::vector &&upper_neighbor_chunks, - const ChunkBroker::Pointer &broker) + HnswStreamerEntityNew(IndexStreamer::Stats &stats, const HNSWHeader &hd, + size_t chunk_size, uint32_t node_index_mask_bits, + uint32_t upper_neighbor_mask_bits, bool filter_same_key, + bool get_vector_enabled, + const NIHashMapPointer &upper_neighbor_index, + std::shared_ptr &keys_map_lock, + const HashMapPointer &keys_map, + bool use_key_info_map, + std::vector &&node_chunks, + std::vector &&upper_neighbor_chunks, + const ChunkBroker::Pointer &broker, + std::shared_ptr vector_value_ptr, + std::shared_ptr neighbors_value_ptr) : stats_(stats), chunk_size_(chunk_size), node_index_mask_bits_(node_index_mask_bits), @@ -241,7 +474,9 @@ class HnswStreamerEntity : public HnswEntity { keys_map_(keys_map), node_chunks_(std::move(node_chunks)), upper_neighbor_chunks_(std::move(upper_neighbor_chunks)), - broker_(broker) { + broker_(broker), + vector_value_ptr_(vector_value_ptr), + neighbors_value_ptr_(neighbors_value_ptr) { *mutable_header() = hd; neighbor_size_ = neighbors_size(); @@ -472,12 +707,13 @@ class HnswStreamerEntity : public HnswEntity { } private: - HnswStreamerEntity(const HnswStreamerEntity &) = delete; - HnswStreamerEntity &operator=(const HnswStreamerEntity &) = delete; + HnswStreamerEntityNew(const HnswStreamerEntityNew &) = delete; + HnswStreamerEntityNew &operator=(const HnswStreamerEntityNew &) = delete; static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f; private: IndexStreamer::Stats &stats_; + HNSWHeader base_header_{}; HNSWHeader header_{}; std::mutex mutex_{}; size_t max_index_size_{0UL}; @@ -509,6 +745,9 @@ class HnswStreamerEntity : public HnswEntity { mutable std::vector upper_neighbor_chunks_{}; ChunkBroker::Pointer broker_{}; // chunk broker + + std::shared_ptr vector_value_ptr_{}; + std::shared_ptr neighbors_value_ptr_{}; }; } // namespace core diff --git a/src/core/utility/file_read_storage.cc b/src/core/utility/file_read_storage.cc index 8f2b79feb..f1b3f732d 100644 --- a/src/core/utility/file_read_storage.cc +++ b/src/core/utility/file_read_storage.cc @@ -289,7 +289,7 @@ class FileReadStorage : public IndexStorage { } int append(const std::string & /*id*/, size_t /*size*/) override { - return IndexError_NotImplemented; + return 0; } void refresh(uint64_t) override { diff --git a/tests/core/algorithm/hnsw/hnsw_builder_test.cc b/tests/core/algorithm/hnsw/hnsw_builder_test.cc deleted file mode 100644 index 402ae9720..000000000 --- a/tests/core/algorithm/hnsw/hnsw_builder_test.cc +++ /dev/null @@ -1,543 +0,0 @@ -// Copyright 2025-present the zvec 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 "hnsw_builder.h" -#include -#include -#include -#include -#include -#include -#include "zvec/core/framework/index_framework.h" - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-result" -#endif - -using namespace std; -using namespace zvec::ailego; - -namespace zvec { -namespace core { - -constexpr size_t static dim = 16; - -class HnswBuilderTest : public testing::Test { - protected: - void SetUp(void); - void TearDown(void); - - static std::string _dir; - static shared_ptr _index_meta_ptr; -}; - -std::string HnswBuilderTest::_dir("hnswBuilderTest"); -shared_ptr HnswBuilderTest::_index_meta_ptr; - -void HnswBuilderTest::SetUp(void) { - _index_meta_ptr.reset(new (nothrow) - IndexMeta(IndexMeta::DataType::DT_FP32, dim)); - _index_meta_ptr->set_metric("SquaredEuclidean", 0, ailego::Params()); -} - -void HnswBuilderTest::TearDown(void) { - char cmdBuf[100]; - snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str()); - system(cmdBuf); -} - -TEST_F(HnswBuilderTest, TestGeneral) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - // params.set("proxima.hnsw.builder.thread_count", 1); - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - - ASSERT_EQ(0, builder->train(holder)); - - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); - // ASSERT_GT(stats.dumped_costtime(), 0UL); - - // cleanup and rebuild - ASSERT_EQ(0, builder->cleanup()); - - auto holder2 = - make_shared>(dim); - size_t doc_cnt2 = 2000UL; - for (size_t i = 0; i < doc_cnt2; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder2->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder2)); - ASSERT_EQ(0, builder->build(holder2)); - auto dumper2 = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper2, nullptr); - ASSERT_EQ(0, dumper2->create(path)); - ASSERT_EQ(0, builder->dump(dumper2)); - ASSERT_EQ(0, dumper2->close()); - - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt2, stats.built_count()); - ASSERT_EQ(doc_cnt2, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -TEST_F(HnswBuilderTest, TestMemquota) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - params.set("proxima.hnsw.builder.memory_quota", 100000UL); - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(IndexError_NoMemory, builder->build(holder)); -} - -TEST_F(HnswBuilderTest, TestIndexThreads) { - IndexBuilder::Pointer builder1 = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder1, nullptr); - IndexBuilder::Pointer builder2 = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder2, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - std::srand(ailego::Realtime::MilliSeconds()); - auto threads = - std::make_shared(std::rand() % 4, false); - ASSERT_EQ(0, builder1->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder2->init(*_index_meta_ptr, params)); - - auto build_index1 = [&]() { - ASSERT_EQ(0, builder1->train(threads, holder)); - ASSERT_EQ(0, builder1->build(threads, holder)); - }; - auto build_index2 = [&]() { - ASSERT_EQ(0, builder2->train(threads, holder)); - ASSERT_EQ(0, builder2->build(threads, holder)); - }; - - auto t1 = std::async(std::launch::async, build_index1); - auto t2 = std::async(std::launch::async, build_index2); - t1.wait(); - t2.wait(); - - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestIndexThreads"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder1->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder2->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats1 = builder1->stats(); - ASSERT_EQ(doc_cnt, stats1.built_count()); - auto &stats2 = builder2->stats(); - ASSERT_EQ(doc_cnt, stats2.built_count()); -} - -TEST_F(HnswBuilderTest, TestCosine) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp32Converter"); - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - - ailego::Params params; - // params.set("proxima.hnsw.builder.thread_count", 1); - ASSERT_EQ(0, builder->init(index_meta, params)); - - ASSERT_EQ(0, builder->train(converted_holder)); - - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestCosine"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); - // ASSERT_GT(stats.dumped_costtime(), 0UL); - - // cleanup and rebuild - ASSERT_EQ(0, builder->cleanup()); - - auto holder2 = - make_shared>(dim); - size_t doc_cnt2 = 2000UL; - for (size_t i = 0; i < doc_cnt2; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder2->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder2)); - ASSERT_EQ(0, builder->build(holder2)); - auto dumper2 = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper2, nullptr); - ASSERT_EQ(0, dumper2->create(path)); - ASSERT_EQ(0, builder->dump(dumper2)); - ASSERT_EQ(0, dumper2->close()); - - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt2, stats.built_count()); - ASSERT_EQ(doc_cnt2, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -TEST_F(HnswBuilderTest, TestCosineFp16Converter) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp16Converter"); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - - ailego::Params params; - - // params.set("proxima.hnsw.builder.thread_count", 1); - ASSERT_EQ(0, builder->init(index_meta, params)); - - ASSERT_EQ(0, builder->train(converted_holder)); - - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestCosineFp16Converter"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); - // ASSERT_GT(stats.dumped_costtime(), 0UL); - - // cleanup and rebuild - ASSERT_EQ(0, builder->cleanup()); - - auto holder2 = - make_shared>(dim); - size_t doc_cnt2 = 2000UL; - for (size_t i = 0; i < doc_cnt2; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder2->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder2)); - ASSERT_EQ(0, builder->build(holder2)); - auto dumper2 = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper2, nullptr); - ASSERT_EQ(0, dumper2->create(path)); - ASSERT_EQ(0, builder->dump(dumper2)); - ASSERT_EQ(0, dumper2->close()); - - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt2, stats.built_count()); - ASSERT_EQ(doc_cnt2, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -TEST_F(HnswBuilderTest, TestCosineInt8Converter) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineInt8Converter"); - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - - ailego::Params params; - // params.set("proxima.hnsw.builder.thread_count", 1); - ASSERT_EQ(0, builder->init(index_meta, params)); - - ASSERT_EQ(0, builder->train(converted_holder)); - - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestCosineInt8Converter"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); - // ASSERT_GT(stats.dumped_costtime(), 0UL); - - // cleanup and rebuild - ASSERT_EQ(0, builder->cleanup()); - - auto holder2 = - make_shared>(dim); - size_t doc_cnt2 = 2000UL; - for (size_t i = 0; i < doc_cnt2; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder2->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder2)); - ASSERT_EQ(0, builder->build(holder2)); - auto dumper2 = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper2, nullptr); - ASSERT_EQ(0, dumper2->create(path)); - ASSERT_EQ(0, builder->dump(dumper2)); - ASSERT_EQ(0, dumper2->close()); - - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt2, stats.built_count()); - ASSERT_EQ(doc_cnt2, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -TEST_F(HnswBuilderTest, TestCosineInt4Converter) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineInt4Converter"); - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - - ailego::Params params; - // params.set("proxima.hnsw.builder.thread_count", 1); - ASSERT_EQ(0, builder->init(index_meta, params)); - - ASSERT_EQ(0, builder->train(converted_holder)); - - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestCosineInt4Converter"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); - // ASSERT_GT(stats.dumped_costtime(), 0UL); - - // cleanup and rebuild - ASSERT_EQ(0, builder->cleanup()); - - auto holder2 = - make_shared>(dim); - size_t doc_cnt2 = 2000UL; - for (size_t i = 0; i < doc_cnt2; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder2->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder2)); - ASSERT_EQ(0, builder->build(holder2)); - auto dumper2 = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper2, nullptr); - ASSERT_EQ(0, dumper2->create(path)); - ASSERT_EQ(0, builder->dump(dumper2)); - ASSERT_EQ(0, dumper2->close()); - - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt2, stats.built_count()); - ASSERT_EQ(doc_cnt2, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -} // namespace core -} // namespace zvec - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic pop -#endif \ No newline at end of file diff --git a/tests/core/algorithm/hnsw/hnsw_searcher_test.cpp b/tests/core/algorithm/hnsw/hnsw_searcher_test.cpp deleted file mode 100644 index d3a7004fa..000000000 --- a/tests/core/algorithm/hnsw/hnsw_searcher_test.cpp +++ /dev/null @@ -1,2775 +0,0 @@ -// Copyright 2025-present the zvec 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 -#include -#include -#include -#include -#include -#include -#include -#include "zvec/core/framework/index_builder.h" -#include "zvec/core/framework/index_factory.h" -#include "zvec/core/framework/index_meta.h" -#include "hnsw_params.h" - -using namespace std; -using namespace testing; -using namespace zvec::ailego; - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-result" -#endif - -namespace zvec { -namespace core { - -constexpr size_t static dim = 16; - -class HnswSearcherTest : public testing::Test { - protected: - void SetUp(void); - void TearDown(void); - - static std::string _dir; - static shared_ptr _index_meta_ptr; -}; - -std::string HnswSearcherTest::_dir("HnswSearcherTest/"); -shared_ptr HnswSearcherTest::_index_meta_ptr; - -void HnswSearcherTest::SetUp(void) { - _index_meta_ptr.reset(new (nothrow) - IndexMeta(IndexMeta::DataType::DT_FP32, dim)); - _index_meta_ptr->set_metric("SquaredEuclidean", 0, ailego::Params()); -} - -void HnswSearcherTest::TearDown(void) { - char cmdBuf[100]; - snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str()); - system(cmdBuf); -} - -TEST_F(HnswSearcherTest, TestRnnSearch) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(*_index_meta_ptr, ailego::Params())); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestRnnSearch"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 0.0; - } - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 50; - ctx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto &results = ctx->result(); - ASSERT_EQ(topk, results.size()); - - float radius = results[topk / 2].score(); - ctx->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - // Test Reset Threshold - ctx->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(radius, results[topk - 1].score()); -} - -TEST_F(HnswSearcherTest, TestRnnSearchInnerProduct) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - IndexMeta index_meta(IndexMeta::DataType::DT_FP32, dim); - index_meta.set_metric("InnerProduct", 0, ailego::Params()); - - ASSERT_EQ(0, builder->init(index_meta, ailego::Params())); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestRnnSearchInnerProduct"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 1.0; - } - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 50; - ctx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto &results = ctx->result(); - ASSERT_EQ(topk, results.size()); - - float radius = -results[topk / 2].score(); - ctx->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - // Test Reset Threshold - ctx->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(-radius, results[topk - 1].score()); -} - -TEST_F(HnswSearcherTest, TestRnnSearchCosine) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - - std::random_device rd; - std::mt19937 gen(rd()); - - std::uniform_real_distribution dist(-1.0, 1.0); - - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = dist(gen); - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp32Converter"); - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - - ASSERT_EQ(0, builder->init(index_meta, ailego::Params())); - ASSERT_EQ(0, builder->train(converted_holder)); - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestRnnSearchCosine"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 1.0; - } - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &new_query, &new_meta)); - - size_t topk = 50; - ctx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, ctx)); - auto &results = ctx->result(); - ASSERT_EQ(topk, results.size()); - - float radius = 0.5f; - ctx->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, ctx)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - // Test Reset Threshold - ctx->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, ctx)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(radius, results[topk - 1].score()); -} - -TEST_F(HnswSearcherTest, TestRnnSearchMipsSquaredEuclidean) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 10); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 10); - params.set(PARAM_HNSW_STREAMER_EF, 5); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - - IndexMeta index_meta(IndexMeta::DataType::DT_FP32, dim); - index_meta.set_metric("MipsSquaredEuclidean", 0, ailego::Params()); - - ailego::Params stg_params; - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestStreamerDump.index", true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - size_t doc_cnt = 1000UL; - auto streamer_ctx = streamer->create_context(); - ASSERT_TRUE(!!streamer_ctx); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - - streamer->add_impl(i, vec.data(), qmeta, streamer_ctx); - } - - { - // Test Reset Threshold - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 1.0; - } - - size_t topk = 50; - streamer_ctx->set_topk(topk); - ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); - auto &results = streamer_ctx->result(); - ASSERT_EQ(topk, results.size()); - - float radius = -results[topk / 2].score(); - streamer_ctx->set_threshold(radius); - ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - streamer_ctx->reset_threshold(); - ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(-radius, results[topk - 1].score()); - } - - auto path = _dir + "/TestStreamerDump"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - - auto read_storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - auto searcher_ctx = searcher->create_context(); - ASSERT_TRUE(!!searcher_ctx); - - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 1.0; - } - - { - size_t topk = 50; - searcher_ctx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, searcher_ctx)); - auto &results = searcher_ctx->result(); - ASSERT_EQ(topk, results.size()); - - float radius = -results[topk / 2].score(); - searcher_ctx->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, searcher_ctx)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - // Test Reset Threshold - searcher_ctx->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, searcher_ctx)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(-radius, results[topk - 1].score()); - } -} - -TEST_F(HnswSearcherTest, TestGeneral) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 5000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - ailego::Params params; - // params.set("proxima.hnsw.builder.max_neighbor_count", 16); - params.set("proxima.hnsw.builder.scaling_factor", 16); - params.set("proxima.hnsw.builder.ef_construction", 10); - params.set("proxima.hnsw.builder.thread_count", 2); - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 1); - ASSERT_EQ(0, searcher->init(searcherParams)); - - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto linearCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - ASSERT_TRUE(!!linearCtx); - ASSERT_TRUE(!!linearByPKeysCtx); - ASSERT_TRUE(!!knnCtx); - NumericalVector vec(dim); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 200; - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - int totalHits = 0; - int totalCnts = 0; - int topk1Hits = 0; - linearCtx->set_topk(topk); - linearByPKeysCtx->set_topk(topk); - knnCtx->set_topk(topk); - - // do linear search test - { - std::vector query(dim); - for (size_t i = 0; i < dim; ++i) { - query[i] = 3.1f; - } - ASSERT_EQ(0, searcher->search_bf_impl(query.data(), qmeta, linearCtx)); - auto &linearResult = linearCtx->result(); - ASSERT_EQ(3UL, linearResult[0].key()); - ASSERT_EQ(4UL, linearResult[1].key()); - ASSERT_EQ(2UL, linearResult[2].key()); - ASSERT_EQ(5UL, linearResult[3].key()); - ASSERT_EQ(1UL, linearResult[4].key()); - ASSERT_EQ(6UL, linearResult[5].key()); - ASSERT_EQ(0UL, linearResult[6].key()); - ASSERT_EQ(7UL, linearResult[7].key()); - for (size_t i = 8; i < topk; ++i) { - ASSERT_EQ(i, linearResult[i].key()); - } - } - - // do linear search by p_keys test - std::vector> p_keys; - p_keys.resize(1); - p_keys[0] = {8, 9, 10, 11, 3, 2, 1, 0}; - { - std::vector query(dim); - for (size_t i = 0; i < dim; ++i) { - query[i] = 3.1f; - } - ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(query.data(), p_keys, qmeta, - linearByPKeysCtx)); - auto &linearByPKeysResult = linearByPKeysCtx->result(); - ASSERT_EQ(8, linearByPKeysResult.size()); - ASSERT_EQ(3UL, linearByPKeysResult[0].key()); - ASSERT_EQ(2UL, linearByPKeysResult[1].key()); - ASSERT_EQ(1UL, linearByPKeysResult[2].key()); - ASSERT_EQ(0UL, linearByPKeysResult[3].key()); - ASSERT_EQ(8UL, linearByPKeysResult[4].key()); - ASSERT_EQ(9UL, linearByPKeysResult[5].key()); - ASSERT_EQ(10UL, linearByPKeysResult[6].key()); - ASSERT_EQ(11UL, linearByPKeysResult[7].key()); - } - - size_t step = 50; - for (size_t i = 0; i < doc_cnt; i += step) { - for (size_t j = 0; j < dim; ++j) { - vec[j] = i + 0.1f; - } - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - // TODO: check - // ASSERT_EQ(topk, knnResult.size()); - topk1Hits += i == knnResult[0].key(); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - for (size_t k = 0; k < topk; ++k) { - totalCnts++; - for (size_t j = 0; j < topk; ++j) { - if (linearResult[j].key() == knnResult[k].key()) { - totalHits++; - break; - } - } - } - } - float recall = totalHits * step * step * 1.0f / totalCnts; - float topk1Recall = topk1Hits * step * 1.0f / doc_cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; -#if 0 - printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " - "R@%zd=%f R@1=%f cost=%f\n", - knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall, - topk1Recall, cost); -#endif - EXPECT_GT(recall, 0.90f); - EXPECT_GT(topk1Recall, 0.90f); - // EXPECT_GT(cost, 2.0f); -} - -TEST_F(HnswSearcherTest, TestClearAndReload) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - ailego::Params params; - params.set("proxima.hnsw.builder.thread_count", 3); - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.check_crc_enable", true); - searcherParams.set("proxima.hnsw.searcher.max_scan_ratio", - 1.1f); // including upper layer - ASSERT_EQ(0, searcher->init(searcherParams)); - - - auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - ASSERT_TRUE(!!linearCtx); - ASSERT_TRUE(!!knnCtx); - NumericalVector vec(dim); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 100; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - auto &stats = searcher->stats(); - ASSERT_EQ(doc_cnt, stats.loaded_count()); - // ASSERT_GT(stats.loaded_costtime(), 0UL); - - //! cleanup - ASSERT_EQ(0, searcher->cleanup()); - ASSERT_EQ(nullptr, searcher->create_context()); - ASSERT_EQ(IndexError_Runtime, - searcher->load(storage, IndexMetric::Pointer())); - ASSERT_EQ(0UL, stats.loaded_count()); - - ASSERT_EQ(0, searcher->init(searcherParams)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - linearCtx = searcher->create_context(); - knnCtx = searcher->create_context(); - ASSERT_TRUE(!!linearCtx); - ASSERT_TRUE(!!knnCtx); - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto &knnResult1 = knnCtx->result(); - ASSERT_EQ(topk, knnResult1.size()); - auto &linearResult1 = linearCtx->result(); - ASSERT_EQ(topk, linearResult1.size()); - ASSERT_EQ(doc_cnt, stats.loaded_count()); - - //! unload - ASSERT_EQ(0, searcher->unload()); - ASSERT_EQ(nullptr, searcher->create_context()); - ASSERT_EQ(0UL, stats.loaded_count()); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - linearCtx = searcher->create_context(); - ASSERT_TRUE(!!linearCtx); - linearCtx->set_topk(topk); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto &linearResult2 = linearCtx->result(); - ASSERT_EQ(topk, linearResult2.size()); - ASSERT_EQ(doc_cnt, stats.loaded_count()); -} - -TEST_F(HnswSearcherTest, TestFilter) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 100UL; - std::vector> p_keys; - p_keys.resize(1); - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - p_keys[0].push_back(i); - } - ailego::Params params; - params.set("proxima.hnsw.builder.thread_count", 3); - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.check_crc_enable", true); - searcherParams.set("proxima.hnsw.searcher.max_scan_ratio", 1.0f); - ASSERT_EQ(0, searcher->init(searcherParams)); - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto linearCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - ASSERT_TRUE(!!linearCtx); - ASSERT_TRUE(!!linearByPKeysCtx); - ASSERT_TRUE(!!knnCtx); - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 10.1f; - } - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 10; - linearCtx->set_topk(topk); - linearByPKeysCtx->set_topk(topk); - knnCtx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(vec.data(), p_keys, qmeta, - linearByPKeysCtx)); - - auto filterFunc = [](uint64_t key) { - if (key == 10UL || key == 11UL) { - return true; - } - return false; - }; - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - ASSERT_EQ(10UL, knnResult[0].key()); - ASSERT_EQ(11UL, knnResult[1].key()); - ASSERT_EQ(9UL, knnResult[2].key()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(10UL, linearResult[0].key()); - ASSERT_EQ(11UL, linearResult[1].key()); - ASSERT_EQ(9UL, linearResult[2].key()); - - auto &linearByPKeysResult = linearByPKeysCtx->result(); - ASSERT_EQ(topk, linearByPKeysResult.size()); - ASSERT_EQ(10UL, linearByPKeysResult[0].key()); - ASSERT_EQ(11UL, linearByPKeysResult[1].key()); - ASSERT_EQ(9UL, linearByPKeysResult[2].key()); - - knnCtx->set_filter(filterFunc); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - auto &knnResult1 = knnCtx->result(); - ASSERT_EQ(topk, knnResult1.size()); - ASSERT_EQ(9UL, knnResult1[0].key()); - ASSERT_EQ(12UL, knnResult1[1].key()); - ASSERT_EQ(8UL, knnResult1[2].key()); - - linearCtx->set_filter(filterFunc); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto &linearResult1 = linearCtx->result(); - ASSERT_EQ(topk, linearResult1.size()); - ASSERT_EQ(9UL, linearResult1[0].key()); - ASSERT_EQ(12UL, linearResult1[1].key()); - ASSERT_EQ(8UL, linearResult1[2].key()); - - linearByPKeysCtx->set_filter(filterFunc); - ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(vec.data(), p_keys, qmeta, - linearByPKeysCtx)); - auto &linearByPKeysResult1 = linearByPKeysCtx->result(); - ASSERT_EQ(topk, linearByPKeysResult1.size()); - ASSERT_EQ(9UL, linearByPKeysResult1[0].key()); - ASSERT_EQ(12UL, linearByPKeysResult1[1].key()); - ASSERT_EQ(8UL, linearByPKeysResult1[2].key()); -} - -TEST_F(HnswSearcherTest, TestStreamerDump) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 10); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 10); - params.set(PARAM_HNSW_STREAMER_EF, 5); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - ailego::Params stg_params; - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestStreamerDump.index", true)); - ASSERT_EQ(0, streamer->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, streamer->open(storage)); - - NumericalVector vec(dim); - size_t cnt = 5000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - for (size_t i = 0; i < cnt; i++) { - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - streamer->add_impl(i, vec.data(), qmeta, ctx); - } - auto path = _dir + "/TestStreamerDump"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // do searcher knn - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - auto read_storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - int totalHits = 0; - int totalCnts = 0; - int topk1Hits = 0; - size_t step = 50; - for (size_t i = 0; i < cnt; i += step) { - for (size_t j = 0; j < dim; ++j) { - vec[j] = i + 0.1f; - } - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - // ASSERT_EQ(topk, knnResult.size()); - topk1Hits += i == knnResult[0].key(); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - for (size_t k = 0; k < topk; ++k) { - totalCnts++; - for (size_t j = 0; j < topk; ++j) { - if (linearResult[j].key() == knnResult[k].key()) { - totalHits++; - break; - } - } - } - } - float recall = totalHits * step * 1.0f / totalCnts; - float topk1Recall = topk1Hits * step * 1.0f / cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; -#if 0 - printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " - "R@%zd=%f R@1=%f cost=%f\n", - knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall, - topk1Recall, cost); -#endif - EXPECT_GT(recall, 0.90f); - EXPECT_GT(topk1Recall, 0.95f); - // EXPECT_GT(cost, 2.0f); -} - -TEST_F(HnswSearcherTest, TestSharedContext) { - auto gen_holder = [](int start, size_t doc_cnt) { - auto holder = - make_shared>(dim); - uint64_t key = start; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - key += 3; - holder->emplace(key, vec); - } - return holder; - }; - auto gen_index = [&gen_holder](int start, size_t docs, std::string path) { - auto holder = gen_holder(start, docs); - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ailego::Params params; - builder->init(*_index_meta_ptr, params); - builder->train(holder); - builder->build(holder); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - dumper->create(path); - builder->dump(dumper); - dumper->close(); - - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - auto name = rand() % 2 ? "FileReadStorage" : "MMapFileReadStorage"; - auto storage = IndexFactory::CreateStorage(name); - storage->open(path, false); - params.set("proxima.hnsw.searcher.visit_bloomfilter_enable", rand() % 2); - searcher->init(ailego::Params()); - searcher->load(storage, IndexMetric::Pointer()); - return searcher; - }; - - srand(ailego::Realtime::MilliSeconds()); - size_t docs1 = rand() % 500 + 100; - size_t docs2 = rand() % 5000 + 100; - size_t docs3 = rand() % 50000 + 100; - auto path1 = _dir + "/TestSharedContext.index1"; - auto path2 = _dir + "/TestSharedContext.index2"; - auto path3 = _dir + "/TestSharedContext.index3"; - auto searcher1 = gen_index(0, docs1, path1); - auto searcher2 = gen_index(1, docs2, path2); - auto searcher3 = gen_index(2, docs3, path3); - - srand(ailego::Realtime::MilliSeconds()); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - auto do_test = [&]() { - IndexSearcher::Context::Pointer ctx; - switch (rand() % 3) { - case 0: - ctx = searcher1->create_context(); - break; - case 1: - ctx = searcher2->create_context(); - break; - case 2: - ctx = searcher3->create_context(); - break; - } - ctx->set_topk(10); - - int ret = 0; - for (int i = 0; i < 100; ++i) { - NumericalVector query(dim); - for (size_t j = 0; j < dim; ++j) { - query[j] = i + 0.1f; - } - - auto code = rand() % 6; - switch (code) { - case 0: - ret = searcher1->search_impl(query.data(), qmeta, ctx); - break; - case 1: - ret = searcher2->search_impl(query.data(), qmeta, ctx); - break; - case 2: - ret = searcher3->search_impl(query.data(), qmeta, ctx); - break; - case 3: - ret = searcher1->search_bf_impl(query.data(), qmeta, ctx); - break; - case 4: - ret = searcher2->search_bf_impl(query.data(), qmeta, ctx); - break; - case 5: - ret = searcher3->search_bf_impl(query.data(), qmeta, ctx); - break; - } - - EXPECT_EQ(0, ret); - auto &results = ctx->result(); - EXPECT_EQ(10, results.size()); - for (int k = 0; k < 10; ++k) { - EXPECT_EQ(code % 3, results[k].key() % 3); - } - } - }; - auto t1 = std::async(std::launch::async, do_test); - auto t2 = std::async(std::launch::async, do_test); - t1.wait(); - t2.wait(); - - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - storage->init(ailego::Params()); - storage->open(_dir + "/TestSharedContext.index4", true); - streamer->init(*_index_meta_ptr, ailego::Params()); - streamer->open(storage); - NumericalVector query(dim); - auto ctx1 = streamer->create_context(); - EXPECT_EQ(IndexError_Unsupported, - searcher1->search_impl(query.data(), qmeta, ctx1)); - - auto ctx2 = searcher1->create_context(); - EXPECT_EQ(IndexError_Unsupported, - streamer->search_impl(query.data(), qmeta, ctx2)); -} - -TEST_F(HnswSearcherTest, TestProvider) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 5000UL; - std::vector keys(doc_cnt); - srand(ailego::Realtime::MilliSeconds()); - bool rand_key = rand() % 2; - bool rand_order = rand() % 2; - size_t step = rand() % 2 + 1; - LOG_DEBUG("randKey=%u randOrder=%u step=%zu", rand_key, rand_order, step); - if (rand_key) { - std::mt19937 mt; - std::uniform_int_distribution dt( - 0, std::numeric_limits::max()); - for (size_t i = 0; i < doc_cnt; ++i) { - keys[i] = dt(mt); - } - } else { - std::iota(keys.begin(), keys.end(), 0U); - std::transform(keys.begin(), keys.end(), keys.begin(), - [&](key_t k) { return step * k; }); - if (rand_order) { - uint32_t seed = ailego::Realtime::Seconds(); - std::shuffle(keys.begin(), keys.end(), std::default_random_engine(seed)); - } - } - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = keys[i]; - } - ASSERT_TRUE(holder->emplace(keys[i], vec)); - } - ailego::Params params; - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestProvider"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 1); - ASSERT_EQ(0, searcher->init(searcherParams)); - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - - auto provider = searcher->create_provider(); - for (size_t i = 0; i < keys.size(); ++i) { - const float *d1 = - reinterpret_cast(provider->get_vector(keys[i])); - ASSERT_TRUE(d1); - for (size_t j = 0; j < dim; ++j) { - ASSERT_FLOAT_EQ(d1[j], keys[i]); - } - } - - auto iter = provider->create_iterator(); - size_t cnt = 0; - while (iter->is_valid()) { - auto key = iter->key(); - const float *d = reinterpret_cast(iter->data()); - for (size_t j = 0; j < dim; ++j) { - ASSERT_FLOAT_EQ(d[j], key); - } - cnt++; - iter->next(); - } - ASSERT_EQ(cnt, doc_cnt); - - ASSERT_EQ(dim, provider->dimension()); - ASSERT_EQ(_index_meta_ptr->element_size(), provider->element_size()); - ASSERT_EQ(_index_meta_ptr->data_type(), provider->data_type()); -} - -TEST_F(HnswSearcherTest, TestMipsEuclideanMetric) { - constexpr size_t static dim = 32; - IndexMeta meta(IndexMeta::DataType::DT_FP32, dim); - meta.set_metric("MipsSquaredEuclidean", 0, ailego::Params()); - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - const size_t COUNT = 10000UL; - for (size_t i = 0; i < COUNT; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i / 100.0f; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(meta, ailego::Params())); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestMipsEuclideanMetric"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ailego::Params params; - params.set("proxima.hnsw.searcher.ef", 10); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(params)); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 1.0; - } - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 50; - ctx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto &results = ctx->result(); - EXPECT_EQ(results.size(), topk); - EXPECT_NEAR((uint64_t)(COUNT - 1), results[0].key(), 20); -} - -TEST_F(HnswSearcherTest, TestRandomPaddingTopk) { - std::mt19937 mt{}; - std::uniform_real_distribution gen(0.0f, 1.0f); - constexpr size_t static dim = 8; - IndexMeta meta(IndexMeta::DataType::DT_FP32, dim); - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - const size_t COUNT = 10000UL; - for (size_t i = 0; i < COUNT; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = gen(mt); - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - ASSERT_EQ(0, builder->init(meta, ailego::Params())); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestRandomPadding"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ailego::Params params; - params.set("proxima.hnsw.searcher.force_padding_result_enable", true); - params.set("proxima.hnsw.searcher.scan_ratio", 0.01f); - ASSERT_TRUE(searcher != nullptr); - ASSERT_EQ(0, searcher->init(params)); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = 1.0; - } - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - std::uniform_int_distribution gen_int(1, COUNT); - size_t topk = gen_int(mt); - ctx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto &results = ctx->result(); - EXPECT_EQ(results.size(), topk); - for (size_t i = 0; i < results.size(); ++i) { - for (size_t j = 0; j < i; ++j) { - EXPECT_NE(results[i].key(), results[j].key()); - } - } - - ctx->set_filter([](uint64_t key) { return true; }); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto &results1 = ctx->result(); - EXPECT_EQ(results1.size(), 0); -} - - -TEST_F(HnswSearcherTest, TestBruteForceSetupInContext) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 5000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - // params.set("proxima.hnsw.builder.max_neighbor_count", 16); - params.set("proxima.hnsw.builder.scaling_factor", 16); - params.set("proxima.hnsw.builder.ef_construction", 10); - params.set("proxima.hnsw.builder.thread_count", 2); - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 1); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - - NumericalVector vec(dim); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t topk = 200; - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - int totalHits = 0; - int totalCnts = 0; - int topk1Hits = 0; - - bool set_bf_threshold = false; - bool use_update = false; - - size_t step = 50; - for (size_t i = 0; i < doc_cnt; i += step) { - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - - ASSERT_TRUE(!!linearCtx); - ASSERT_TRUE(!!linearCtx); - - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - - for (size_t j = 0; j < dim; ++j) { - vec[j] = i + 0.1f; - } - auto t1 = ailego::Realtime::MicroSeconds(); - - if (set_bf_threshold) { - if (use_update) { - ailego::Params searcherParamsExtra; - - searcherParamsExtra.set("proxima.hnsw.searcher.brute_force_threshold", - doc_cnt); - knnCtx->update(searcherParamsExtra); - } else { - knnCtx->set_bruteforce_threshold(doc_cnt); - } - - use_update = !use_update; - } - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - - auto t2 = ailego::Realtime::MicroSeconds(); - - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - // auto t3 = ailego::Realtime::MicroSeconds(); - - if (set_bf_threshold) { - linearTotalTime += t2 - t1; - } else { - knnTotalTime += t2 - t1; - } - - set_bf_threshold = !set_bf_threshold; - - auto &knnResult = knnCtx->result(); - // TODO: check - // ASSERT_EQ(topk, knnResult.size()); - topk1Hits += i == knnResult[0].key(); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - for (size_t k = 0; k < topk; ++k) { - totalCnts++; - for (size_t j = 0; j < topk; ++j) { - if (linearResult[j].key() == knnResult[k].key()) { - totalHits++; - break; - } - } - } - } - float recall = totalHits * step * step * 1.0f / totalCnts; - float topk1Recall = topk1Hits * step * 1.0f / doc_cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; -#if 0 - printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " - "R@%zd=%f R@1=%f cost=%f\n", - knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall, - topk1Recall, cost); -#endif - EXPECT_GT(recall, 0.90f); - EXPECT_GT(topk1Recall, 0.90f); - // EXPECT_GT(cost, 2.0f); -} - -TEST_F(HnswSearcherTest, TestCosine) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp32Converter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestCosine.index", true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - NumericalVector vec(dim); - size_t cnt = 5000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - - float fixed_value = float(cnt) / 2; - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - vec[j] = fixed_value; - else - vec[j] = fixed_value + add_on; - } - - std::string new_vec; - IndexQueryMeta new_meta; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - } - - auto path = _dir + "/TestCosine"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - - ASSERT_TRUE(!!linearCtx); - ASSERT_TRUE(!!linearByPKeysCtx); - ASSERT_TRUE(!!knnCtx); - - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - int totalHits = 0; - int totalCnts = 0; - int topk1Hits = 0; - - NumericalVector qvec(dim); - for (size_t i = 0; i < query_cnt; i++) { - float add_on = i * 10; - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - qvec[j] = fixed_value; - else - qvec[j] = fixed_value + add_on; - } - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, - reformer->transform(qvec.data(), qmeta, &new_query, &new_meta)); - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, - searcher->search_bf_impl(new_query.data(), new_meta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - topk1Hits += i == knnResult[0].key(); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - for (size_t k = 0; k < topk; ++k) { - totalCnts++; - for (size_t j = 0; j < topk; ++j) { - if (linearResult[j].key() == knnResult[k].key()) { - totalHits++; - break; - } - } - } - } - - float recall = totalHits * 1.0f / totalCnts; - float topk1Recall = topk1Hits * 1.0f / query_cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; - - EXPECT_GT(recall, 0.90f); - EXPECT_GT(topk1Recall, 0.90f); - // EXPECT_GT(cost, 2.0f); -} - -TEST_F(HnswSearcherTest, TestFetchVector) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_TRUE(streamer != nullptr); - - IndexMeta index_meta(IndexMeta::DataType::DT_FP32, dim); - index_meta.set_metric("SquaredEuclidean", 0, ailego::Params()); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - ailego::Params stg_params; - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestFetchVector.index", true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - NumericalVector vec(dim); - size_t cnt = 2000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - - for (size_t i = 0; i < cnt; i++) { - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - - streamer->add_impl(i, vec.data(), qmeta, ctx); - } - - auto path = _dir + "/TestFetchVector"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - for (size_t i = 0; i < cnt; i++) { - const void *vector = searcher->get_vector(i); - ASSERT_NE(vector, nullptr); - - float vector_value = *(float *)(vector); - ASSERT_EQ(vector_value, i); - } - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - knnCtx->set_fetch_vector(true); - - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - - for (size_t i = 0; i < query_cnt; i++) { - for (size_t j = 0; j < dim; ++j) { - vec[j] = i; - } - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - ASSERT_NE(knnResult[0].vector(), nullptr); - float vector_value = *((float *)(knnResult[0].vector())); - ASSERT_EQ(vector_value, i); - } - - std::cout << "knnTotalTime: " << knnTotalTime << std::endl; - std::cout << "linearTotalTime: " << linearTotalTime << std::endl; -} - -TEST_F(HnswSearcherTest, TestFetchVectorCosine) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); - - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp32Converter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestFetchVectorCosine.index", true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - NumericalVector vec(dim); - size_t cnt = 2000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - IndexQueryMeta new_meta; - - const float epsilon = 1e-2; - float fixed_value = float(cnt) / 2; - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - vec[j] = fixed_value; - else - vec[j] = fixed_value + add_on; - } - - std::string new_vec; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - } - - auto path = _dir + "/TestFetchVectorCosine"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - - const void *vector = searcher->get_vector(i); - ASSERT_NE(vector, nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(vector, new_meta, &denormalized_vec); - - float vector_value = *((float *)(denormalized_vec.data()) + dim - 1); - EXPECT_NEAR(vector_value, fixed_value + add_on, epsilon); - } - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - knnCtx->set_fetch_vector(true); - - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - - NumericalVector qvec(dim); - for (size_t i = 0; i < query_cnt; i++) { - float add_on = i * 10; - - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - qvec[j] = fixed_value; - else - qvec[j] = fixed_value + add_on; - } - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, - reformer->transform(qvec.data(), qmeta, &new_query, &new_meta)); - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, - searcher->search_bf_impl(new_query.data(), new_meta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - ASSERT_NE(knnResult[0].vector(), nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(knnResult[0].vector(), new_meta, &denormalized_vec); - - float vector_value = *(((float *)(denormalized_vec.data()) + dim - 1)); - EXPECT_NEAR(vector_value, fixed_value + add_on, epsilon); - } - - std::cout << "knnTotalTime: " << knnTotalTime << std::endl; - std::cout << "linearTotalTime: " << linearTotalTime << std::endl; -} - - -TEST_F(HnswSearcherTest, TestFetchVectorCosineHalfFloatConverter) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); - - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP16, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineHalfFloatConverter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ( - 0, storage->open(_dir + "/TestFetchVectorCosineHalfFloatConverter.index", - true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - size_t cnt = 2000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP16, dim); - IndexQueryMeta new_meta; - - const float epsilon = 0.1; - - std::random_device rd; - std::mt19937 gen(rd()); - - std::uniform_real_distribution dist(-2.0, 2.0); - - std::vector> vecs; - for (size_t i = 0; i < cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - float value = dist(gen); - vec[j] = ailego::FloatHelper::ToFP16(value); - } - - std::string new_vec; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - - vecs.push_back(vec); - } - - auto path = _dir + "/TestFetchVectorCosineHalfFloatConverter"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - for (size_t i = 0; i < cnt; i++) { - uint16_t expected_vec_value = vecs[i][dim - 1]; - - const void *vector = searcher->get_vector(i); - ASSERT_NE(vector, nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(uint16_t)); - reformer->revert(vector, new_meta, &denormalized_vec); - - uint16_t vector_value = *((uint16_t *)(denormalized_vec.data()) + dim - 1); - float vector_value_float = ailego::FloatHelper::ToFP32(vector_value); - - float expected_vec_float = ailego::FloatHelper::ToFP32(expected_vec_value); - - EXPECT_NEAR(expected_vec_float, vector_value_float, epsilon); - } - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - knnCtx->set_fetch_vector(true); - - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - - NumericalVector qvec(dim); - - for (size_t i = 0; i < query_cnt; i++) { - auto &vec = vecs[i]; - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &new_query, &new_meta)); - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, - searcher->search_bf_impl(new_query.data(), new_meta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - ASSERT_NE(knnResult[0].vector(), nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(uint16_t)); - reformer->revert(knnResult[0].vector(), new_meta, &denormalized_vec); - - uint16_t expected_vec_value = vec[dim - 1]; - uint16_t vector_value = - *(((uint16_t *)(denormalized_vec.data()) + dim - 1)); - - float vector_value_float = ailego::FloatHelper::ToFP32(vector_value); - float expected_vec_float = ailego::FloatHelper::ToFP32(expected_vec_value); - - EXPECT_NEAR(expected_vec_float, vector_value_float, epsilon); - } - - std::cout << "knnTotalTime: " << knnTotalTime << std::endl; - std::cout << "linearTotalTime: " << linearTotalTime << std::endl; -} - -TEST_F(HnswSearcherTest, TestFetchVectorCosineFp16Converter) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); - - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp16Converter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestFetchVectorCosineFp16Converter.index", - true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - size_t cnt = 2000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - IndexQueryMeta new_meta; - - const float epsilon = 0.1; - - std::random_device rd; - std::mt19937 gen(rd()); - - std::uniform_real_distribution dist(-2.0, 2.0); - - std::vector> vecs; - for (size_t i = 0; i < cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = dist(gen); - } - - std::string new_vec; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - - vecs.push_back(vec); - } - - auto path = _dir + "/TestFetchVectorCosineFp16Converter"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - for (size_t i = 0; i < cnt; i++) { - float expected_vec_value = vecs[i][dim - 1]; - - const void *vector = searcher->get_vector(i); - ASSERT_NE(vector, nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(vector, new_meta, &denormalized_vec); - float vector_value = *((float *)(denormalized_vec.data()) + dim - 1); - - EXPECT_NEAR(expected_vec_value, vector_value, epsilon); - } - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - knnCtx->set_fetch_vector(true); - - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - - NumericalVector qvec(dim); - - for (size_t i = 0; i < query_cnt; i++) { - auto &vec = vecs[i]; - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &new_query, &new_meta)); - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, - searcher->search_bf_impl(new_query.data(), new_meta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - ASSERT_NE(knnResult[0].vector(), nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(knnResult[0].vector(), new_meta, &denormalized_vec); - - float expected_vec_value = vec[dim - 1]; - float vector_value = *(((float *)(denormalized_vec.data()) + dim - 1)); - - EXPECT_NEAR(expected_vec_value, vector_value, epsilon); - } - - std::cout << "knnTotalTime: " << knnTotalTime << std::endl; - std::cout << "linearTotalTime: " << linearTotalTime << std::endl; -} - -TEST_F(HnswSearcherTest, TestFetchVectorCosineInt8Converter) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); - - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineInt8Converter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestFetchVectorCosineInt8Converter.index", - true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - NumericalVector vec(dim); - size_t cnt = 2000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - IndexQueryMeta new_meta; - - const float epsilon = 1e-2; - float fixed_value = float(cnt) / 2; - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - vec[j] = fixed_value; - else - vec[j] = fixed_value + add_on; - } - - std::string new_vec; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - } - - auto path = _dir + "/TestFetchVectorCosineInt8Converter"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - - const void *vector = searcher->get_vector(i); - ASSERT_NE(vector, nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(vector, new_meta, &denormalized_vec); - - float vector_value = *((float *)(denormalized_vec.data()) + dim - 1); - EXPECT_NEAR(vector_value, fixed_value + add_on, epsilon); - } - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - knnCtx->set_fetch_vector(true); - - size_t topk = 200; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - - NumericalVector qvec(dim); - for (size_t i = 0; i < query_cnt; i++) { - float add_on = i * 10; - - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - qvec[j] = fixed_value; - else - qvec[j] = fixed_value + add_on; - } - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, - reformer->transform(qvec.data(), qmeta, &new_query, &new_meta)); - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, - searcher->search_bf_impl(new_query.data(), new_meta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - ASSERT_NE(knnResult[0].vector(), nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(knnResult[0].vector(), new_meta, &denormalized_vec); - - float vector_value = *(((float *)(denormalized_vec.data()) + dim - 1)); - EXPECT_NEAR(vector_value, fixed_value + add_on, epsilon); - } - - std::cout << "knnTotalTime: " << knnTotalTime << std::endl; - std::cout << "linearTotalTime: " << linearTotalTime << std::endl; -} - -TEST_F(HnswSearcherTest, TestFetchVectorCosineInt4Converter) { - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_NE(streamer, nullptr); - - ailego::Params params; - params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100); - params.set(PARAM_HNSW_STREAMER_EF, 100); - params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); - - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineInt4Converter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestFetchVectorCosineInt4Converter.index", - true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - NumericalVector vec(dim); - size_t cnt = 2000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - IndexQueryMeta new_meta; - - const float epsilon = 1e-2; - float fixed_value = float(cnt) / 2; - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - vec[j] = fixed_value; - else - vec[j] = fixed_value + add_on; - } - - std::string new_vec; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - } - - auto path = _dir + "/TestFetchVectorCosineInt4Converter"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 100); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - for (size_t i = 0; i < cnt; i++) { - float add_on = i * 10; - - const void *vector = searcher->get_vector(i); - ASSERT_NE(vector, nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(vector, new_meta, &denormalized_vec); - - float vector_value = *((float *)(denormalized_vec.data()) + dim - 1); - EXPECT_NEAR(vector_value, fixed_value + add_on, epsilon); - } - - size_t query_cnt = 200U; - auto linearCtx = searcher->create_context(); - auto knnCtx = searcher->create_context(); - auto linearByPKeysCtx = searcher->create_context(); - knnCtx->set_fetch_vector(true); - - size_t topk = 100; - linearCtx->set_topk(topk); - knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; - - NumericalVector qvec(dim); - for (size_t i = 0; i < query_cnt; i++) { - float add_on = i * 10; - - for (size_t j = 0; j < dim; ++j) { - if (j < dim / 4) - qvec[j] = fixed_value; - else - qvec[j] = fixed_value + add_on; - } - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, - reformer->transform(qvec.data(), qmeta, &new_query, &new_meta)); - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto t2 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, - searcher->search_bf_impl(new_query.data(), new_meta, linearCtx)); - auto t3 = ailego::Realtime::MicroSeconds(); - - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; - - auto &knnResult = knnCtx->result(); - ASSERT_EQ(topk, knnResult.size()); - - auto &linearResult = linearCtx->result(); - ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); - - ASSERT_NE(knnResult[0].vector(), nullptr); - - std::string denormalized_vec; - denormalized_vec.resize(dim * sizeof(float)); - reformer->revert(knnResult[0].vector(), new_meta, &denormalized_vec); - - float vector_value = *(((float *)(denormalized_vec.data()) + dim - 1)); - EXPECT_NEAR(vector_value, fixed_value + add_on, epsilon); - } - - std::cout << "knnTotalTime: " << knnTotalTime << std::endl; - std::cout << "linearTotalTime: " << linearTotalTime << std::endl; -} - -TEST_F(HnswSearcherTest, TestGroup) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 5000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i / 10.0; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGroup"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_NE(searcher, nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 50); - searcherParams.set("proxima.hnsw.searcher.max_scan_ratio", 0.8); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t group_topk = 20; - uint64_t total_time = 0; - - auto groupbyFunc = [](uint64_t key) { - uint32_t group_id = key / 10 % 10; - - // std::cout << "key: " << key << ", group id: " << group_id << std::endl; - - return std::string("g_") + std::to_string(group_id); - }; - - size_t group_num = 5; - - ctx->set_group_params(group_num, group_topk); - ctx->set_group_by(groupbyFunc); - - size_t query_value = doc_cnt / 2; - for (size_t j = 0; j < dim; ++j) { - vec[j] = float(query_value) / 10 + 0.1f; - } - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto t2 = ailego::Realtime::MicroSeconds(); - - total_time += t2 - t1; - - std::cout << "total time: " << total_time << std::endl; - - auto &group_result = ctx->group_result(); - - for (uint32_t i = 0; i < group_result.size(); ++i) { - // const std::string &group_id = group_result[i].group_id(); - auto &result = group_result[i].docs(); - - ASSERT_GT(result.size(), 0); - // std::cout << "Group ID: " << group_id << std::endl; - - // for (uint32_t j = 0; j < result.size(); ++j) { - // std::cout << "\tKey: " << result[j].key() << std::fixed - // << std::setprecision(3) << ", Score: " << result[j].score() - // << std::endl; - // } - } - - // do linear search by p_keys test - auto groupbyFuncLinear = [](uint64_t key) { - uint32_t group_id = key % 10; - - return std::string("g_") + std::to_string(group_id); - }; - - auto linear_pk_ctx = searcher->create_context(); - - linear_pk_ctx->set_group_params(group_num, group_topk); - linear_pk_ctx->set_group_by(groupbyFuncLinear); - - std::vector> p_keys; - p_keys.resize(1); - p_keys[0] = {4, 3, 2, 1, 5, 6, 7, 8, 9, 10}; - - ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(vec.data(), p_keys, qmeta, - linear_pk_ctx)); - auto &linear_by_pkeys_group_result = linear_pk_ctx->group_result(); - ASSERT_EQ(linear_by_pkeys_group_result.size(), group_num); - - for (uint32_t i = 0; i < linear_by_pkeys_group_result.size(); ++i) { - // const std::string &group_id = linear_by_pkeys_group_result[i].group_id(); - auto &result = linear_by_pkeys_group_result[i].docs(); - - ASSERT_GT(result.size(), 0); - // std::cout << "Group ID: " << group_id << std::endl; - - // for (uint32_t j = 0; j < result.size(); ++j) { - // std::cout << "\tKey: " << result[j].key() << std::fixed - // << std::setprecision(3) << ", Score: " << result[j].score() - // << std::endl; - // } - - ASSERT_EQ(10 - i, result[0].key()); - } -} - -TEST_F(HnswSearcherTest, TestGroupNotEnoughNum) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 5000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i / 10.0; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGroupNotEnoughNum"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_NE(searcher, nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 50); - searcherParams.set("proxima.hnsw.searcher.max_scan_ratio", 0.8); - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t group_topk = 20; - uint64_t total_time = 0; - - auto groupbyFunc = [](uint64_t key) { - uint32_t group_id = key / 10 % 10; - - // std::cout << "key: " << key << ", group id: " << group_id << std::endl; - - return std::string("g_") + std::to_string(group_id); - }; - - size_t group_num = 12; - ctx->set_group_params(group_num, group_topk); - ctx->set_group_by(groupbyFunc); - - size_t query_value = doc_cnt / 2; - for (size_t j = 0; j < dim; ++j) { - vec[j] = float(query_value) / 10 + 0.1f; - } - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto t2 = ailego::Realtime::MicroSeconds(); - total_time += t2 - t1; - - std::cout << "total time: " << total_time << std::endl; - - auto &group_result = ctx->group_result(); - ASSERT_EQ(group_result.size(), 10); - - for (uint32_t i = 0; i < group_result.size(); ++i) { - // const std::string &group_id = group_result[i].group_id(); - auto &result = group_result[i].docs(); - - ASSERT_GT(result.size(), 0); - // std::cout << "Group ID: " << group_id << std::endl; - - // for (uint32_t j = 0; j < result.size(); ++j) { - // std::cout << "\tKey: " << result[j].key() << std::fixed - // << std::setprecision(3) << ", Score: " << result[j].score() - // << std::endl; - // } - } -} - -TEST_F(HnswSearcherTest, TestGroupInBruteforceSearch) { - IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("HnswBuilder"); - ASSERT_NE(builder, nullptr); - auto holder = - make_shared>(dim); - size_t doc_cnt = 5000UL; - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = i / 10.0; - } - ASSERT_TRUE(holder->emplace(i, vec)); - } - - ailego::Params params; - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - string path = _dir + "/TestGroupInBruteforceSearch"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_NE(searcher, nullptr); - ailego::Params searcherParams; - searcherParams.set("proxima.hnsw.searcher.ef", 50); - searcherParams.set("proxima.hnsw.searcher.max_scan_ratio", 0.8); - searcherParams.set("proxima.hnsw.searcher.brute_force_threshold", - 2 * doc_cnt); - - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto storage = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, storage->open(path, false)); - ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); - - auto ctx = searcher->create_context(); - ASSERT_TRUE(!!ctx); - - NumericalVector vec(dim); - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); - size_t group_topk = 20; - uint64_t total_time = 0; - - auto groupbyFunc = [](uint64_t key) { - uint32_t group_id = key / 10 % 10; - - // std::cout << "key: " << key << ", group id: " << group_id << std::endl; - - return std::string("g_") + std::to_string(group_id); - }; - - size_t group_num = 5; - ctx->set_group_params(group_num, group_topk); - ctx->set_group_by(groupbyFunc); - - size_t query_value = doc_cnt / 2; - for (size_t j = 0; j < dim; ++j) { - vec[j] = float(query_value) / 10 + 0.1f; - } - - auto t1 = ailego::Realtime::MicroSeconds(); - ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto t2 = ailego::Realtime::MicroSeconds(); - total_time += t2 - t1; - - std::cout << "total time: " << total_time << std::endl; - - auto &group_result = ctx->group_result(); - ASSERT_EQ(group_result.size(), 5); - - for (uint32_t i = 0; i < group_result.size(); ++i) { - // const std::string &group_id = group_result[i].group_id(); - auto &result = group_result[i].docs(); - - ASSERT_GT(result.size(), 0); - // std::cout << "Group ID: " << group_id << std::endl; - - // for (uint32_t j = 0; j < result.size(); ++j) { - // std::cout << "\tKey: " << result[j].key() << std::fixed - // << std::setprecision(3) << ", Score: " << result[j].score() - // << std::endl; - // } - } -} - -TEST_F(HnswSearcherTest, TestBinaryConverter) { - uint32_t dimension = 256; - - IndexStreamer::Pointer streamer = - IndexFactory::CreateStreamer("HnswStreamer"); - ASSERT_TRUE(streamer != nullptr); - - ailego::Params params; - // params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 50); - // params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16); - // params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 10); - // params.set(PARAM_HNSW_STREAMER_EF, 5); - // params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); - - ailego::Params stg_params; - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dimension); - index_meta_raw.set_metric("InnerProduct", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("BinaryConverter"); - ASSERT_TRUE(converter != nullptr); - - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - auto storage = IndexFactory::CreateStorage("MMapFileStorage"); - ASSERT_EQ(0, storage->init(stg_params)); - ASSERT_EQ(0, storage->open(_dir + "/TestBinaryConverter.index", true)); - ASSERT_EQ(0, streamer->init(index_meta, params)); - ASSERT_EQ(0, streamer->open(storage)); - - size_t cnt = 5000U; - auto ctx = streamer->create_context(); - ASSERT_TRUE(!!ctx); - - IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dimension); - - std::random_device rd; - std::mt19937 gen(rd()); - - std::uniform_real_distribution dist(-2.0, 2.0); - std::vector> vecs; - - for (size_t i = 0; i < cnt; i++) { - NumericalVector vec(dimension); - for (size_t j = 0; j < dimension; ++j) { - vec[j] = dist(gen); - } - - std::string new_vec; - IndexQueryMeta new_meta; - - ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta)); - ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx)); - - vecs.push_back(vec); - } - - auto path = _dir + "/TestBinaryConverter"; - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, streamer->dump(dumper)); - ASSERT_EQ(0, streamer->close()); - ASSERT_EQ(0, dumper->close()); - - // test searcher - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - ASSERT_TRUE(searcher != nullptr); - - ailego::Params searcherParams; - ASSERT_EQ(0, searcher->init(searcherParams)); - - auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage"); - ASSERT_EQ(0, read_storage->open(path, false)); - ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer())); - - size_t query_cnt = 200U; - auto knnCtx = searcher->create_context(); - - float epison = 1e-6; - for (size_t i = 0; i < query_cnt; i++) { - auto &vec = vecs[i]; - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &new_query, &new_meta)); - - size_t topk = 50; - knnCtx->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, knnCtx)); - auto &results = knnCtx->result(); - ASSERT_EQ(topk, results.size()); - ASSERT_EQ(i, results[0].key()); - ASSERT_NEAR(0, results[0].score(), epison); - } -} - -} // namespace core -} // namespace zvec - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic pop -#endif \ No newline at end of file diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc index c04f67124..5fb5e2f4a 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc @@ -353,7 +353,7 @@ TEST_F(HnswStreamerTest, TestKnnSearch) { } float recall = totalHits * 1.0f / totalCnts; float topk1Recall = topk1Hits * 1.0f / cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; + // float cost = linearTotalTime * 1.0f / knnTotalTime; #if 0 printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " "R@%zd=%f R@1=%f cost=%f\n", @@ -439,7 +439,7 @@ TEST_F(HnswStreamerTest, TestAddAndSearch) { } float recall = totalHits * 1.0f / totalCnts; float topk1Recall = topk1Hits * 100.0f / cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; + // float cost = linearTotalTime * 1.0f / knnTotalTime; #if 0 printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " "R@%zd=%f R@1=%f cost=%f\n", @@ -1678,15 +1678,12 @@ TEST_F(HnswStreamerTest, TestDumpIndexAndAdd) { ASSERT_EQ(IndexError_Unsupported, code); // check dump index - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - auto container = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, container->init(ailego::Params())); - ASSERT_EQ(0, container->open(path1, false)); - ASSERT_NE(searcher, nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - ASSERT_EQ(0, searcher->load(container, IndexMetric::Pointer())); - auto iter = searcher->create_provider()->create_iterator(); + IndexStreamer::Pointer read_streamer = + IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(read_streamer, nullptr); + ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params)); + ASSERT_EQ(0, read_streamer->open(storage)); + auto iter = read_streamer->create_provider()->create_iterator(); size_t docs = 0; while (iter->is_valid()) { auto key = iter->key(); @@ -1777,15 +1774,12 @@ TEST_F(HnswStreamerTest, TestProvider) { streamer->close(); // check dump index - IndexSearcher::Pointer searcher = - IndexFactory::CreateSearcher("HnswSearcher"); - auto container = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_EQ(0, container->init(ailego::Params())); - ASSERT_EQ(0, container->open(path1, false)); - ASSERT_NE(searcher, nullptr); - ASSERT_EQ(0, searcher->init(ailego::Params())); - ASSERT_EQ(0, searcher->load(container, IndexMetric::Pointer())); - auto iter = searcher->create_provider()->create_iterator(); + IndexStreamer::Pointer read_streamer = + IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(read_streamer, nullptr); + ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params)); + ASSERT_EQ(0, read_streamer->open(storage)); + auto iter = read_streamer->create_provider()->create_iterator(); size_t cnt = 0; while (iter->is_valid()) { auto key = iter->key(); @@ -1812,29 +1806,6 @@ TEST_F(HnswStreamerTest, TestProvider) { iter->next(); } ASSERT_EQ(cnt, docs); - - - auto searcher_provider = searcher->create_provider(); - auto streamer_provider = streamer->create_provider(); - for (size_t i = 0; i < keys.size(); ++i) { - const float *d1 = - reinterpret_cast(searcher_provider->get_vector(keys[i])); - ASSERT_TRUE(d1); - for (size_t j = 0; j < dim; ++j) { - ASSERT_FLOAT_EQ(d1[j], keys[i]); - } - - const float *d2 = - reinterpret_cast(streamer_provider->get_vector(keys[i])); - ASSERT_TRUE(d2); - for (size_t j = 0; j < dim; ++j) { - ASSERT_FLOAT_EQ(d2[j], keys[i]); - } - } - - ASSERT_EQ(dim, streamer_provider->dimension()); - ASSERT_EQ(index_meta_ptr_->element_size(), streamer_provider->element_size()); - ASSERT_EQ(index_meta_ptr_->data_type(), streamer_provider->data_type()); } TEST_F(HnswStreamerTest, TestSharedContext) { @@ -2093,7 +2064,7 @@ TEST_F(HnswStreamerTest, TestBruteForceSetupInContext) { } float recall = totalHits * 1.0f / totalCnts; float topk1Recall = topk1Hits * 1.0f / cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; + // float cost = linearTotalTime * 1.0f / knnTotalTime; #if 0 printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " "R@%zd=%f R@1=%f cost=%f\n", @@ -2220,7 +2191,7 @@ TEST_F(HnswStreamerTest, TestKnnSearchCosine) { } float recall = totalHits * 1.0f / totalCnts; float topk1Recall = topk1Hits * 1.0f / query_cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; + // float cost = linearTotalTime * 1.0f / knnTotalTime; #if 0 printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " "R@%zd=%f R@1=%f cost=%f\n", @@ -3654,7 +3625,7 @@ TEST_F(HnswStreamerTest, TestAddAndSearchWithID) { } float recall = totalHits * 1.0f / totalCnts; float topk1Recall = topk1Hits * 100.0f / cnt; - float cost = linearTotalTime * 1.0f / knnTotalTime; + // float cost = linearTotalTime * 1.0f / knnTotalTime; #if 0 printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d " "R@%zd=%f R@1=%f cost=%f\n",