feat(ivf): add graph bucket searcher for large bucket approximate search#2494
feat(ivf): add graph bucket searcher for large bucket approximate search#2494LHT129 wants to merge 1 commit into
Conversation
|
/label S-waiting-on-review |
LHT129
added
kind/feature
Merge Protections🟢 All 3 merge protections satisfied — ready to merge. Show 3 satisfied protections🟢 Require kind label
🟢 Require version label
🟢 Require linked issue for feature/bug PRs
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces GraphBucketSearcher to enable graph-based bucket searching in IVF indexes when bucket sizes exceed a specified threshold. The review feedback highlights critical issues in the serialization and deserialization of bucket graphs, which currently bypass offset registration and cause cursor mismatches. Additionally, the reviewer suggests optimizing performance in the hot query path by allocating heaps on the stack instead of using std::make_shared, ensuring proper memory tracking by replacing std::vector with the custom Vector container, and removing a redundant loop in the neighbor selection logic.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
f12780e to
9986050
Compare
|
Automated pull request review failed. Review effort: git exited with 128: Cloning into '/tmp/vsag-pull-request-reviews/01a9c713-f061-4741-9fd5-b498ecd10567/source'... No GitHub review was submitted. |
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds an IVF bucket-level approximate search path using per-bucket graphs and beam search for large buckets, while preserving exact flat scans for small buckets / default configs.
Changes:
- Introduces
GraphBucketSearcherimplementingIVFBucketSearcherwith graph-based beam search and flat fallback. - Extends IVF parameters and search parameters to support
graph_build_threshold,graph_params, andef_search. - Adds build + (de)serialization logic for per-bucket graphs.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| src/inner_string_params.h | Adds string keys for graph params, build threshold, and ef_search. |
| src/algorithm/ivf/ivf_parameter.h | Adds graph-related IVF build params and ef_search search param. |
| src/algorithm/ivf/ivf_parameter.cpp | Parses/serializes new IVF params and parses ef_search. |
| src/algorithm/ivf/ivf.h | Adds members for bucket graphs + build helper. |
| src/algorithm/ivf/ivf.cpp | Wires graph searcher, builds graphs post-Build(), and serializes/deserializes graphs. |
| src/algorithm/ivf/graph_bucket_searcher.h | Declares new bucket searcher for graph/flat modes. |
| src/algorithm/ivf/graph_bucket_searcher.cpp | Implements graph beam search and flat scan fallback. |
| src/algorithm/ivf/CMakeLists.txt | Adds graph_bucket_searcher.cpp to build. |
1169ad1 to
0f2c515
Compare
0f2c515 to
f80c0c4
Compare
459dd61 to
cc26e14
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (3)
src/algorithm/ivf/ivf.cpp:683
- Bucket graphs are built using GetInnerIds() offsets, but BucketInterface explicitly allows holes where GetInnerIds()[i] == std::numeric_limits::max(). This code divides that sentinel by buckets_per_data_ and uses it to index into the base dataset, which can lead to out-of-bounds access/crashes. The neighbor selection loop also currently allows selecting hole positions as neighbors.
auto* inner_ids = bucket_->GetInnerIds(b);
for (InnerIdType i = 0; i < static_cast<InnerIdType>(bucket_size); ++i) {
uint64_t orig_idx = static_cast<uint64_t>(inner_ids[i]) / buckets_per_data_;
const void* query_vec =
get_ivf_graph_build_data(base, common_param_.data_type_, orig_idx, this->dim_);
src/algorithm/ivf/ivf.cpp:971
- Streaming serialization adds IVF_BUCKET_GRAPH to the manifest whenever graph_build_threshold_ > 0 (collect_streaming_header), but the body only writes the block when graph_param_ != nullptr and bucket_graphs_ is non-empty. This can produce an invalid stream where the body has fewer blocks than the manifest (Factory::GetStreamingMetadata will throw).
if (graph_build_threshold_ > 0 && graph_param_ != nullptr && !bucket_graphs_.empty()) {
auto tag = static_cast<uint32_t>(StreamSerializationTag::IVF_BUCKET_GRAPH);
WriteStreamingBlock(
writer, tag, StreamSerializationTagCritical(tag), [this](StreamWriter& w) {
StreamWriter::WriteObj(w, graph_build_threshold_);
src/algorithm/ivf/graph_bucket_searcher.cpp:161
- Neighbor expansion should skip hole entries (ids[neighbor_id] == std::numeric_limits::max()) before computing distances and recording visits; otherwise QueryOneById / origin_id computation can be done on an invalid slot.
for (uint32_t i = 0; i < neighbor_count; ++i) {
auto neighbor_id = neighbors[i];
if (vl.Get(neighbor_id)) {
continue;
}
vl.Set(neighbor_id);
float d = bucket->QueryOneById(computer, bucket_id, neighbor_id);
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (890 changed lines across 11 files).
Submitted 1 inline comment.
Reviewed commit cc26e14.
|
Fixed hole sentinel check: SearchGraph now skips hole entries (ids[i] == numeric_limits::max()) when finding entry point. Also changed default graph IO from block_memory_io to memory_io to avoid 128 MiB allocation per bucket. |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1005 changed lines across 13 files).
Submitted 6 inline comments.
Reviewed commit 5de1759.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1005 changed lines across 13 files).
Submitted 7 inline comments.
Reviewed commit 3805cad.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1005 changed lines across 13 files).
Submitted 3 inline comments.
Reviewed commit 6ae932a.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1009 changed lines across 13 files).
Submitted 8 inline comments.
Reviewed commit 389b0f1.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1025 changed lines across 13 files).
Submitted 7 inline comments.
Reviewed commit 8541e91.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1027 changed lines across 13 files).
Submitted 8 inline comments.
Reviewed commit 353cb12.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1064 changed lines across 13 files).
Submitted 8 inline comments.
Reviewed commit 297ca19.
| }; | ||
|
|
||
| // Find a valid (non-hole) entry point | ||
| InnerIdType entry = 0; |
There was a problem hiding this comment.
[P1] Guarantee reachability from the single graph entry
Every search starts from this one offset, but construction stores only directed nearest-neighbor edges and never connects components. Two separated groups larger than max_degree therefore form disconnected subgraphs, so a query in the group without this entry cannot return its true neighbors even when ef_search equals the bucket size. Build a navigable connected graph, retain entries for every component, or use flat fallback.
| graph->Deserialize(block); | ||
| if (bid >= 0 && bid < static_cast<BucketIdType>(bucket_count)) { | ||
| auto total = graph->TotalCount(); | ||
| for (InnerIdType nid = 0; nid < total; ++nid) { |
There was a problem hiding this comment.
[P1] Validate graph capacity before iterating TotalCount
total_count and the compressed graph's serialized vertex_num are independent fields. A malformed payload with total_count > vertex_num makes GetNeighborSize(nid) index past neighbor_sets_ in this validation loop before an error can be raised; the regular deserializer repeats the same pattern. Verify count/capacity against the target bucket before calling any neighbor accessor.
| } | ||
| } | ||
| } | ||
| this->graph_param = |
There was a problem hiding this comment.
[P1] Reject flat graph parameters with a null IO parameter
With graph_params: {"io_params": {}}, GetGraphParameterByJson returns a non-null flat graph parameter whose nested io_parameter_ is null, so the check below passes and GraphInterface::MakeInstance later dereferences it during Build. The persisted-parameter parser has the same exposure. Validate the concrete flat parameter and return INVALID_ARGUMENT/INVALID_BINARY instead of crashing.
| if (graph_json.Contains(GRAPH_STORAGE_TYPE_KEY)) { | ||
| const auto graph_storage_type_str = graph_json[GRAPH_STORAGE_TYPE_KEY].GetString(); | ||
| if (graph_storage_type_str == GRAPH_STORAGE_TYPE_VALUE_COMPRESSED) { | ||
| graph_storage_type = GraphStorageTypes::GRAPH_STORAGE_TYPE_VALUE_COMPRESSED; |
There was a problem hiding this comment.
[P2] Enforce the compressed graph's 255-neighbor limit
Compressed adjacency stores its element count in uint8_t, but this path accepts any max_degree. With graph_storage_type: "compressed", max_degree: 256, and a bucket of at least 257 vectors, graph construction selects 256 neighbors and EliasFanoEncoder::Encode throws after Add has already populated the index. Reject degrees above 255 during parameter parsing.
| get_ivf_graph_build_data(base, common_param_.data_type_, orig_idx, this->dim_); | ||
| CHECK_ARGUMENT(query_vec != nullptr, "base dataset has no graph build vector data"); | ||
| auto computer = bucket_->FactoryComputer(query_vec); | ||
| bucket_->ScanBucketById(dists.data(), computer, b); |
There was a problem hiding this comment.
[P2] Avoid quadratic graph construction for large buckets
This full bucket scan runs once for every node and is followed by another complete candidate pass, producing O(bucket_size²) distance evaluations plus heap work. A qualifying 100,000-vector bucket requires roughly 10 billion comparisons, making the feature impractical specifically for the large buckets it targets. Use bounded approximate or otherwise subquadratic construction.
| const auto* ids = bucket->GetInnerIds(bucket_id); | ||
|
|
||
| uint64_t ef = param.ef; | ||
| if (ef < static_cast<uint64_t>(topk)) { |
There was a problem hiding this comment.
[P2] Keep unlimited range search off a full graph walk
Unlimited RangeSearch supplies the total index size as topk, so this raises ef and the later cap sets it to the complete bucket size. The traversal consequently processes every reachable node through scalar QueryOneById calls and heap operations, replacing the existing batched flat scan with a more expensive full graph walk. Decouple traversal width from the output limit or use flat range scanning.
| ef = static_cast<uint64_t>(bucket_size); | ||
| } | ||
|
|
||
| VisitedList vl(bucket_size, allocator_); |
There was a problem hiding this comment.
[P2] Reuse visited state across bucket searches
Every probed bucket allocates and initializes visited storage proportional to the complete bucket, even when traversal visits only ef_search nodes. Large buckets therefore retain linear setup cost and allocator pressure for ordinary bounded-KNN queries. Reuse a per-thread versioned visited list or use sparse traversal state.
| "{ATTR_PARAMS_KEY}": { | ||
| "{ATTR_HAS_BUCKETS_KEY}": true | ||
| }, | ||
| "{GRAPH_BUILD_THRESHOLD_KEY}": 0, |
There was a problem hiding this comment.
[P3] Document the public graph-bucket behavior
This exposes graph_build_threshold, nested graph configuration, and ef_search, and changes qualifying buckets from direct scoring to approximate graph traversal, but neither canonical English nor Chinese IVF page documents them. Update both parameter tables and explain recall/connectivity, build cost, memory, flat fallback, and the Build-only graph lifecycle.
Signed-off-by: LHT129 <tianlan.lht@antgroup.com>
Summary
Add GraphBucketSearcher that uses beam search on per-bucket graphs for large buckets while falling back to flat brute-force scan for small ones.
Closes #2493
Changes
Files Changed