diff --git a/src/algorithm/hgraph/hgraph.cpp b/src/algorithm/hgraph/hgraph.cpp index f215c86588..062418a990 100644 --- a/src/algorithm/hgraph/hgraph.cpp +++ b/src/algorithm/hgraph/hgraph.cpp @@ -197,15 +197,18 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { FlattenInterfacePtr tune_source; if (is_tune_base_code or is_tune_precise_code or is_tune_raw_code) { - if (covers_active_ids(raw_vector_)) { - tune_source = raw_vector_; - } else if (covers_active_ids(high_precise_codes_) and - high_precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { + tune_source = raw_vector_; + if ((tune_source == nullptr or tune_source->TotalCount() == 0) and + high_precise_codes_ != nullptr and + high_precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { tune_source = high_precise_codes_; - } else if (covers_active_ids(basic_flatten_codes_) and - basic_flatten_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { + } + if ((tune_source == nullptr or tune_source->TotalCount() == 0) and + basic_flatten_codes_ != nullptr and + basic_flatten_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { tune_source = basic_flatten_codes_; - } else { + } + if (tune_source == nullptr) { return false; } } @@ -261,7 +264,7 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { new_code->Train(train_data.data(), train_count); Vector insert_buffer(dim_, 0, allocator_); - for (int64_t i = 0; i < total_count_; ++i) { + for (int64_t i = 0; i < current_count; ++i) { decode_tune_source(i, insert_buffer.data()); new_code->InsertVector(static_cast(insert_buffer.data()), i); } diff --git a/src/algorithm/hgraph/hgraph_add_test.cpp b/src/algorithm/hgraph/hgraph_add_test.cpp index a165a709e5..f324732fbe 100644 --- a/src/algorithm/hgraph/hgraph_add_test.cpp +++ b/src/algorithm/hgraph/hgraph_add_test.cpp @@ -15,8 +15,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -35,6 +37,38 @@ namespace { +class ArmableRejectingThreadPool final : public vsag::ThreadPool { +public: + void + WaitUntilEmpty() override { + } + + void + SetQueueSizeLimit(uint64_t) override { + } + + void + SetPoolSize(uint64_t) override { + } + + std::future + Enqueue(std::function task) override { + if (reject_submissions_.load(std::memory_order_acquire)) { + throw std::bad_alloc(); + } + task(); + return {}; + } + + void + SetRejectSubmissions(bool reject) { + reject_submissions_.store(reject, std::memory_order_release); + } + +private: + std::atomic reject_submissions_{false}; +}; + vsag::DatasetPtr MakeFloatDataset(std::vector& vectors, std::vector& ids, @@ -162,6 +196,62 @@ const std::string kBruteForceSearchParams = } // namespace +TEST_CASE("HGraph Tune accepts an incomplete source after Add failure", + "[ut][hgraph][add][hgraph_tune_incomplete_source]") { + constexpr int64_t dim = 4; + constexpr int64_t base_count = 2; + + auto rejecting_pool = std::make_shared(); + auto common_param = MakeCommonParam(dim); + common_param.thread_pool_ = std::make_shared(rejecting_pool); + auto hgraph_json = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "max_degree": 8, + "ef_construction": 32, + "build_thread_count": 1, + "store_raw_vector": true + })"); + auto index = MakeHGraphIndex(hgraph_json, common_param); + + std::vector base_vectors = { + 0.0F, + 0.0F, + 0.0F, + 0.0F, + 1.0F, + 1.0F, + 1.0F, + 1.0F, + }; + std::vector base_ids = {10, 20}; + auto base = MakeFloatDataset(base_vectors, base_ids, dim, base_count); + REQUIRE(index->Build(base).has_value()); + REQUIRE(index->GetNumElements() == base_count); + + std::vector add_vectors = {2.0F, 2.0F, 2.0F, 2.0F}; + std::vector add_ids = {30}; + auto add = MakeFloatDataset(add_vectors, add_ids, dim, 1); + + rejecting_pool->SetRejectSubmissions(true); + auto add_result = index->Add(add); + rejecting_pool->SetRejectSubmissions(false); + + REQUIRE_FALSE(add_result.has_value()); + REQUIRE(add_result.error().type == vsag::ErrorType::NO_ENOUGH_MEMORY); + REQUIRE(index->GetNumElements() == base_count + 1); + REQUIRE(index->CheckIdExist(add_ids[0])); + + auto tune_result = index->Tune(R"({ + "index_param": { + "base_quantization_type": "bf16", + "max_degree": 8, + "ef_construction": 32 + } + })"); + REQUIRE(tune_result.has_value()); + CHECK(tune_result.value()); +} + TEST_CASE("HGraph exact duplicate fallback supports every dense data type", "[ut][hgraph][duplicate][data_type]") { constexpr int64_t dim = 8;