Skip to content

[diskann-garnet] Implement continue_search() - #1357

Open
Jack Moffitt (metajack) wants to merge 1 commit into
mainfrom
push-pssntsrowrpl
Open

[diskann-garnet] Implement continue_search()#1357
Jack Moffitt (metajack) wants to merge 1 commit into
mainfrom
push-pssntsrowrpl

Conversation

@metajack

Copy link
Copy Markdown
Contributor

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 Continuation is boxed and returned, which can be used by potentially repeated calls to continue_search() to retrieve the rest of the results.

This set up is used all the search_X FFI methods. It is not used in random_members() since due to how other things are handled with that it can just use multiple calls to get more random members if needed.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Continuation and implement continue_search() to drain overflowed search results across multiple FFI calls.
  • Extend SearchResults to track k and store overflow IDs/distances when output buffers can’t fit all results.
  • Wire continuation out-parameters through search_vector, search_element, and search_neighbors, and update/extend tests accordingly; bump version to 5.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_element does not initialize/validate the continuation out-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_search only writes new_continuation when 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_mut is 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_vector writes through the continuation out-parameter on success, but it is never validated and it is not initialized on early-return error paths (e.g., failed interpret_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_neighbors can return -1 before it ever writes to the continuation out-parameter (e.g., if neighbors() 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.

Comment thread diskann-garnet/src/lib.rs
Comment thread diskann-garnet/src/lib.rs Fixed
Comment thread diskann-garnet/src/lib.rs Fixed
Comment thread diskann-garnet/src/lib.rs Fixed
Comment thread diskann-garnet/src/lib.rs Fixed
Comment thread diskann-garnet/src/lib.rs Fixed
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.35450% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.56%. Comparing base (158126e) to head (7711a7a).

Files with missing lines Patch % Lines
diskann-garnet/src/lib.rs 97.26% 5 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@           Coverage Diff            @@
##             main    #1357    +/-   ##
========================================
  Coverage   91.55%   91.56%            
========================================
  Files         521      521            
  Lines      100371   100539   +168     
========================================
+ Hits        91898    92059   +161     
- Misses       8473     8480     +7     
Flag Coverage Δ
miri 91.56% <97.35%> (+<0.01%) ⬆️
unittests 91.24% <97.35%> (+<0.01%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
diskann-garnet/src/dyn_index.rs 80.35% <100.00%> (+0.54%) ⬆️
diskann-garnet/src/provider.rs 79.17% <100.00%> (+0.04%) ⬆️
diskann-garnet/src/lib.rs 94.97% <97.26%> (+0.36%) ⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Jack - just a few small API questions!

Comment thread diskann-garnet/src/lib.rs
_output_distances: *mut f32,
_output_distances_len: usize,
_new_continuation: *mut c_void,
continuation: *mut c_void,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread diskann-garnet/src/lib.rs Outdated
Comment thread diskann-garnet/src/lib.rs
Comment thread diskann-garnet/src/lib.rs
Comment thread diskann-garnet/src/lib.rs
new_continuation: *mut *mut c_void,
) -> i32 {
-1
let index = unsafe { &*index_ptr.cast::<Index>() };
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants