[diskann-garnet] Implement continue_search() - #1357
[diskann-garnet] Implement continue_search()#1357Jack Moffitt (metajack) wants to merge 1 commit into
continue_search()#1357Conversation
There was a problem hiding this comment.
Pull request overview
This PR implements continue_search() for the Garnet FFI by introducing an overflow/continuation mechanism to safely return variable-length external IDs when the caller-provided ID buffer is too small. It extends the existing search FFI surface to return a continuation pointer for subsequent calls, updates neighbor queries to use max_degree(), and bumps the diskann-garnet package version.
Changes:
- Add
Continuationand implementcontinue_search()to drain overflowed search results across multiple FFI calls. - Extend
SearchResultsto trackkand store overflow IDs/distances when output buffers can’t fit all results. - Wire continuation out-parameters through
search_vector,search_element, andsearch_neighbors, and update/extend tests accordingly; bump version to5.0.1.
Reviewed changes
Copilot reviewed 7 out of 8 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| diskann-garnet/src/provider.rs | Exposes max_degree() from the provider for neighbor retrieval sizing. |
| diskann-garnet/src/lib.rs | Implements Continuation + continue_search(), adds overflow handling to SearchResults, and updates FFI search functions to return continuation pointers. |
| diskann-garnet/src/ffi_tests.rs | Updates FFI tests to use the new continuation out-parameter and adds coverage for continuation being written. |
| diskann-garnet/src/ffi_recall_tests.rs | Updates recall tests to pass and clean up continuation pointers. |
| diskann-garnet/src/dyn_index.rs | Extends DynIndex with max_degree() and plumbs it to the provider. |
| diskann-garnet/diskann-garnet.nuspec | Bumps NuGet package version to 5.0.1. |
| diskann-garnet/Cargo.toml | Bumps crate version to 5.0.1. |
| Cargo.lock | Updates locked version for diskann-garnet to 5.0.1. |
Suppressed comments (4)
diskann-garnet/src/lib.rs:941
search_elementdoes not initialize/validate thecontinuationout-parameter. On success without overflow it currently never writes to it (caller must pre-initialize), and on error returns it can leave the out-parameter stale/uninitialized.
let index = unsafe { &*index_ptr.cast::<Index>() };
let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) };
let id = GarnetId::from(id_bytes);
let ctx = Context::new(ctx);
diskann-garnet/src/lib.rs:1021
continue_searchonly writesnew_continuationwhen more results remain. If the continuation is exhausted, the out-parameter is left untouched, so callers that don't pre-initialize it may treat garbage as a live pointer (double-free/UB). Also,slice::from_raw_parts_mutis invoked unconditionally, which is UB if a C caller passes a null pointer with a 0 length (a common FFI pattern).
if continuation.is_null() || new_continuation.is_null() {
return -1;
}
let output_ids = unsafe { slice::from_raw_parts_mut(output_ids, output_ids_len) };
diskann-garnet/src/lib.rs:850
search_vectorwrites through thecontinuationout-parameter on success, but it is never validated and it is not initialized on early-return error paths (e.g., failedinterpret_vector). This can segfault if Garnet passes a null out-parameter, and it can leave the caller with an uninitialized/stale continuation pointer on errors.
This issue also appears in the following locations of the same file:
- line 937
- line 1017
- line 1176
continuation: *mut *mut c_void,
) -> i32 {
let index = unsafe { &*index_ptr.cast::<Index>() };
let v = if let Some(v) = interpret_vector(index.quant_type, &vector_data, vector_len) {
diskann-garnet/src/lib.rs:1179
search_neighborscan return-1before it ever writes to thecontinuationout-parameter (e.g., ifneighbors()fails), and it doesn't validate the out-parameter before writing on success. This can leave callers with stale/uninitialized continuation pointers or cause a segfault if a null out-parameter is passed.
let index = unsafe { &*index_ptr.cast::<Index>() };
let ctx = Context::new(ctx);
let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) };
let id = GarnetId::from(id_bytes);
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1357 +/- ##
========================================
Coverage 91.55% 91.56%
========================================
Files 521 521
Lines 100371 100539 +168
========================================
+ Hits 91898 92059 +161
- Misses 8473 8480 +7
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks Jack - just a few small API questions!
| _output_distances: *mut f32, | ||
| _output_distances_len: usize, | ||
| _new_continuation: *mut c_void, | ||
| continuation: *mut c_void, |
There was a problem hiding this comment.
Why not modify the existing continuation in place? This effectively moves out of the old continuation and into new_continuation, but doesn't take steps to make continuation null. The caller then has to know to not attempt to free continuation.
7711a7a to
8ad6351
Compare
| new_continuation: *mut *mut c_void, | ||
| ) -> i32 { | ||
| -1 | ||
| let index = unsafe { &*index_ptr.cast::<Index>() }; |
Implements
continue_search()in the Garnet FFI.The buffers for ids that Garnet passes may be insufficient since external IDs are user-provided byte strings of arbitrary length. The distances buffer will always be correctly sized. In the case the id buffer is too small, a
Continuationis boxed and returned, which can be used by potentially repeated calls tocontinue_search()to retrieve the rest of the results.This set up is used all the
search_XFFI methods. It is not used inrandom_members()since due to how other things are handled with that it can just use multiple calls to get more random members if needed.