diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index c0d561db..5fd7c75e 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -61,7 +61,6 @@ class QGBuilder { static constexpr size_t kMaxCandidatePoolSize = 750; static constexpr size_t kMaxPrunedSize = 300; std::vector new_neighbors_; // new neighbors for current iteration - std::vector pruned_neighbors_; std::vector degrees_; uint32_t seed_; // base seed for graph initialization std::optional cache_; @@ -69,6 +68,9 @@ class QGBuilder { struct BuildWorker { VisitedSet visited; CandidateList candidates; + CandidateList pruned; + CandidateList completion_result; + CandidateList successful_result; std::vector rotated_neighbors; std::vector neighbor_ids; @@ -79,13 +81,13 @@ class QGBuilder { const BuildCache* build_cache() const; void find_candidates(PID, CandidateList&, VisitedSet&) const; void random_init(); - void search_new_neighbors(bool record_pruned); - void heuristic_prune(PID, CandidateList&, CandidateList&, bool record_pruned); - void add_reverse_edges(bool record_pruned); + void search_new_neighbors(); + void heuristic_prune(CandidateList&, CandidateList&, CandidateList* recorded); + void add_reverse_edges(bool final_pass); void add_pruned_edges( const CandidateList&, const CandidateList&, CandidateList&, float ); - void complete_degrees(); + void complete_degree(PID, BuildWorker&); void run_pass(bool final_pass); void run_passes(size_t); void write_row(PID, BuildWorker&); diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index 26f974b3..bb6c13b1 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -267,13 +267,11 @@ void QGBuilder::initialize_storage(const float* data) { // Allocate refinement scratch after PiPNN releases its dense workspace. // Clear old graph-dependent state on reset while retaining vector capacities. new_neighbors_.resize(num_nodes_); - pruned_neighbors_.resize(num_nodes_); reverse_buffer_.resize(num_nodes_); for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const size_t i = static_cast(index); new_neighbors_[i].clear(); - pruned_neighbors_[i].clear(); reverse_buffer_[i].clear(); } if (workers_.empty()) { @@ -402,7 +400,7 @@ void QGBuilder::add_pruned_edges( } void QGBuilder::heuristic_prune( - PID cur_id, CandidateList& pool, CandidateList& pruned_results, bool record_pruned + CandidateList& pool, CandidateList& pruned_results, CandidateList* recorded ) { if (pool.empty()) { return; @@ -439,8 +437,8 @@ void QGBuilder::heuristic_prune( auto djk = query.distance_to(pool[k].id); if (djk < dik) { - if (record_pruned && pruned_neighbors_[cur_id].size() < kMaxPrunedSize) { - pruned_neighbors_[cur_id].emplace_back(pool[k]); + if (recorded != nullptr && recorded->size() < kMaxPrunedSize) { + recorded->emplace_back(pool[k]); } pruned[k] = true; } @@ -453,9 +451,9 @@ void QGBuilder::heuristic_prune( /** * @brief search for new neighbor in qg * - * @param record_pruned Record pruned candidates for final degree completion. + * Search pruning does not record candidates for final degree completion. */ -void QGBuilder::search_new_neighbors(bool record_pruned) { +void QGBuilder::search_new_neighbors() { { index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_threads_) schedule(dynamic) @@ -506,14 +504,14 @@ void QGBuilder::search_new_neighbors(bool record_pruned) { // prune and update qg new_neighbors_[cur_id].reserve(degree_bound_); - heuristic_prune(cur_id, candidates, new_neighbors_[cur_id], record_pruned); + heuristic_prune(candidates, new_neighbors_[cur_id], /*recorded=*/nullptr); } catch (...) { failure.capture(); } } failure.rethrow(); } } -void QGBuilder::add_reverse_edges(bool record_pruned) { +void QGBuilder::add_reverse_edges(bool final_pass) { std::vector locks(num_nodes_); for (auto& row : reverse_buffer_) { row.clear(); @@ -586,7 +584,19 @@ void QGBuilder::add_reverse_edges(bool record_pruned) { new_neighbors_[data_id].end() ); std::sort(tmp_pool.begin(), tmp_pool.end()); - heuristic_prune(data_id, tmp_pool, new_neighbors_[data_id], record_pruned); + auto& worker = workers_[static_cast(omp_get_thread_num())]; + worker.pruned.clear(); + if (final_pass) { + worker.pruned.reserve(kMaxPrunedSize); + } + heuristic_prune( + tmp_pool, new_neighbors_[data_id], final_pass ? &worker.pruned : nullptr + ); + // Reverse collection has finished. Complete this row against the old + // graph before reusing its worker scratch; graph writes stay deferred. + if (final_pass) { + complete_degree(data_id, worker); + } } catch (...) { failure.capture(); } } failure.rethrow(); @@ -607,7 +617,6 @@ void QGBuilder::random_init() { try { const size_t i = static_cast(index); new_neighbors_[i].clear(); - pruned_neighbors_[i].clear(); std::mt19937 rng(seed_ ^ static_cast(i)); std::uniform_int_distribution distribution(min_id, max_id); std::unordered_set neighbor_set; @@ -637,110 +646,76 @@ void QGBuilder::random_init() { } } -/** - * @brief Complete each vertex's degree to the degree bound (multiple of 32), using - * pruned candidates and random fallback edges. - * - */ -void QGBuilder::complete_degrees() { - { - index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - if (failure.failed()) { - continue; - } - try { - const size_t i = static_cast(index); - CandidateList& cur_neighbors = new_neighbors_[i]; - size_t cur_degree = cur_neighbors.size(); - - // skip vertices with enough neighbors - if (cur_degree >= degree_bound_) { - continue; - } - - CandidateList& pruned_list = pruned_neighbors_[i]; - CandidateList new_result; - new_result.reserve(degree_bound_); - - std::sort(pruned_list.begin(), pruned_list.end()); - - // use binary search to get refined results - CandidateList successful_result; - successful_result.reserve(degree_bound_); - float left = 0.5; - float right = 1.0; - size_t iter = 0; - while (iter++ < kMaxBsIter) { - float mid = (left + right) / 2; - add_pruned_edges(cur_neighbors, pruned_list, new_result, mid); - if (new_result.size() < degree_bound_) { - left = mid; - } else { - right = mid; - successful_result = new_result; - } - } - - // Every successful right endpoint was already evaluated at that exact - // threshold. - if (!successful_result.empty()) { - new_result.swap(successful_result); - } else { - add_pruned_edges(cur_neighbors, pruned_list, new_result, right); - } +// Complete one final-pass row using only this worker's reverse-prune candidates. +void QGBuilder::complete_degree(PID i, BuildWorker& worker) { + CandidateList& cur_neighbors = new_neighbors_[i]; + size_t cur_degree = cur_neighbors.size(); - // if the vertex still doesn't have enough neighbors, use random vertices - if (new_result.size() < degree_bound_) { - std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); - std::uniform_int_distribution distribution( - 0, static_cast(num_nodes_) - 1 - ); - std::unordered_set ids; - ids.reserve(degree_bound_); - for (auto& neighbor : new_result) { - ids.emplace(neighbor.id); - } - BuildQuery query(qg_, build_cache()); - query.reset(i); - while (new_result.size() < degree_bound_) { - const PID rand_id = distribution(rng); - if (rand_id != static_cast(i) && - ids.find(rand_id) == ids.end()) { - new_result.emplace_back(rand_id, query.distance_to(rand_id)); - ids.emplace(rand_id); - } - } - } + // skip vertices with enough neighbors + if (cur_degree >= degree_bound_) { + return; + } - cur_neighbors = new_result; - } catch (...) { failure.capture(); } + CandidateList& pruned_list = worker.pruned; + CandidateList& new_result = worker.completion_result; + new_result.reserve(degree_bound_); + + std::sort(pruned_list.begin(), pruned_list.end()); + + // use binary search to get refined results + CandidateList& successful_result = worker.successful_result; + successful_result.clear(); + successful_result.reserve(degree_bound_); + float left = 0.5; + float right = 1.0; + size_t iter = 0; + while (iter++ < kMaxBsIter) { + float mid = (left + right) / 2; + add_pruned_edges(cur_neighbors, pruned_list, new_result, mid); + if (new_result.size() < degree_bound_) { + left = mid; + } else { + right = mid; + successful_result = new_result; } - failure.rethrow(); } -} -void QGBuilder::run_pass(bool final_pass) { - if (final_pass) { - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - pruned_neighbors_[i].clear(); - pruned_neighbors_[i].reserve(kMaxPrunedSize); - } + // Every successful right endpoint was already evaluated at that exact + // threshold. + if (!successful_result.empty()) { + new_result.swap(successful_result); + } else { + add_pruned_edges(cur_neighbors, pruned_list, new_result, right); } - const bool record_pruned = final_pass; - search_new_neighbors(record_pruned); + // if the vertex still doesn't have enough neighbors, use random vertices + if (new_result.size() < degree_bound_) { + std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); + std::uniform_int_distribution distribution( + 0, static_cast(num_nodes_) - 1 + ); + std::unordered_set ids; + ids.reserve(degree_bound_); + for (auto& neighbor : new_result) { + ids.emplace(neighbor.id); + } + BuildQuery query(qg_, build_cache()); + query.reset(i); + while (new_result.size() < degree_bound_) { + const PID rand_id = distribution(rng); + if (rand_id != static_cast(i) && ids.find(rand_id) == ids.end()) { + new_result.emplace_back(rand_id, query.distance_to(rand_id)); + ids.emplace(rand_id); + } + } + } - add_reverse_edges(record_pruned); + cur_neighbors = new_result; +} - // Only the final pass completes degrees; every pass writes updated rows below. - if (final_pass) { - complete_degrees(); - } +void QGBuilder::run_pass(bool final_pass) { + search_new_neighbors(); + add_reverse_edges(final_pass); // update qg { diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index a2ddeb5d..4e1e356e 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -45,6 +45,41 @@ namespace rabitqlib::symqg { struct QGConstructionTestAccess { + static void check_search_scratch(QGBuilder& builder) { + for (auto& worker : builder.workers_) { + worker.pruned.clear(); + } + builder.search_new_neighbors(); + for (const auto& worker : builder.workers_) { + EXPECT_TRUE(worker.pruned.empty()); + } + } + + static CandidateList record_collinear_prune(QGBuilder& builder) { + CandidateList pool; + for (PID id = 1; id < builder.num_nodes_; ++id) { + pool.emplace_back(id, build_distance(builder, 0, id)); + } + CandidateList result; + CandidateList recorded; + builder.heuristic_prune(pool, result, &recorded); + EXPECT_EQ(result.size(), 1U); + return recorded; + } + + static void prepare_reverse_fan_in(QGBuilder& builder) { + for (auto& row : builder.new_neighbors_) { + row.clear(); + } + for (PID id = 1; id < builder.num_nodes_; ++id) { + builder.new_neighbors_[id].emplace_back(0, build_distance(builder, id, 0)); + } + } + + static void reverse_prune(QGBuilder& builder, bool record) { + builder.add_reverse_edges(record); + } + static void use_legacy_padding(QuantizedGraph& graph) { graph.padded_dim_ = round_up_to_multiple(graph.dim_, 64); graph.initialize(); @@ -107,12 +142,9 @@ struct QGConstructionTestAccess { } while (selected < 32); expected.emplace_back(selected, build_distance(builder, 0, selected)); } - for (auto& row : builder.new_neighbors_) { - row.assign(32, AnnCandidate(1, 0.0F)); - } builder.new_neighbors_[0] = current; - builder.pruned_neighbors_[0] = pruned; - builder.complete_degrees(); + builder.workers_[0].pruned = pruned; + builder.complete_degree(0, builder.workers_[0]); const auto& actual = builder.new_neighbors_[0]; EXPECT_EQ(actual.size(), expected.size()); for (size_t i = 0; i < std::min(actual.size(), expected.size()); ++i) { @@ -220,9 +252,7 @@ struct QGConstructionTestAccess { } } } - static void search(QGBuilder& builder) { - builder.search_new_neighbors(/*record_pruned=*/false); - } + static void search(QGBuilder& builder) { builder.search_new_neighbors(); } static auto encoded_neighbors(const QuantizedGraph& graph, PID id) { return graph.get_neighbors(id); } @@ -299,6 +329,58 @@ struct QGConstructionTestAccess { namespace { +TEST(QGConstructionTest, OnlyFinalReversePruningCompletesRowsWithoutWritingGraph) { + constexpr size_t kCount = 400; + constexpr size_t kDim = 64; + std::vector data(kCount * kDim, 0.0F); + for (size_t i = 0; i < kCount; ++i) { + data[i * kDim] = static_cast(i + 1); + } + for (size_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + QuantizedGraph graph( + kCount, kDim, 32, METRIC_L2, RotatorType::FhtKacRotator, 0, 42 + ); + QGBuilder builder(graph, 64, data.data(), threads, QGInitialization::Random, 42); + QGConstructionTestAccess::check_search_scratch(builder); + + // Collinear candidates must retain the first 300 occluded edges, in order. + const auto recorded = QGConstructionTestAccess::record_collinear_prune(builder); + ASSERT_EQ(recorded.size(), 300U); + for (size_t i = 0; i < recorded.size(); ++i) { + EXPECT_EQ(recorded[i].id, i + 2); + EXPECT_FLOAT_EQ(recorded[i].distance, static_cast((i + 2) * (i + 2))); + } + + const auto original_graph = QGConstructionTestAccess::rows(graph); + QGConstructionTestAccess::prepare_reverse_fan_in(builder); + QGConstructionTestAccess::reverse_prune(builder, false); + ASSERT_EQ(QGConstructionTestAccess::neighbors(builder)[0].size(), 1U); + EXPECT_EQ(QGConstructionTestAccess::rows(graph), original_graph); + + QGConstructionTestAccess::prepare_reverse_fan_in(builder); + QGConstructionTestAccess::reverse_prune(builder, true); + const auto& rows = QGConstructionTestAccess::neighbors(builder); + for (size_t id = 0; id < kCount; ++id) { + ASSERT_EQ(rows[id].size(), 32U); + std::vector ids; + for (const auto& neighbor : rows[id]) { + EXPECT_NE(neighbor.id, id); + EXPECT_LT(neighbor.id, kCount); + ids.push_back(neighbor.id); + } + std::sort(ids.begin(), ids.end()); + EXPECT_EQ(std::adjacent_find(ids.begin(), ids.end()), ids.end()); + } + // At threshold1 collinear candidates fill row0 in distance order. + for (size_t i = 0; i < rows[0].size(); ++i) { + EXPECT_EQ(rows[0][i].id, i + 1); + EXPECT_FLOAT_EQ(rows[0][i].distance, static_cast((i + 1) * (i + 1))); + } + EXPECT_EQ(QGConstructionTestAccess::rows(graph), original_graph); + } +} + TEST(QGConstructionTest, RefinementMatchesOriginalFinalEndpointAndSeededFallback) { constexpr size_t kDim = 65; constexpr size_t kCount = 65; @@ -344,7 +426,9 @@ TEST(QGConstructionTest, RefinementMatchesOriginalFinalEndpointAndSeededFallback ), (std::array{2, 3, 0}) ); - // No probe succeeds; the previously untested endpoint 1.0 does. + // Reuse the same worker after successful completion: no stale + // successful result may survive when every new probe fails. + // The previously untested endpoint 1.0 succeeds. EXPECT_EQ( QGConstructionTestAccess::check_refinement_endpoint( builder, 0.9921875F, false