Skip to content

refactor the filter code in the garnet provider #808

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hailangx merged 2 commits into from
Mar 5, 2026
Merged

refactor the filter code in the garnet provider #808

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d8ba5ab
refactor
Mar 2, 2026
6ba855e
remove the static
Mar 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 2 additions & 71 deletions diskann-garnet/src/dyn_index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,7 @@ use std::sync::Arc;

use diskann::{
ANNError, ANNResult,
graph::{
InplaceDeleteMethod, SearchOutputBuffer, glue::SearchStrategy, index::SearchStats, search,
},
graph::{InplaceDeleteMethod, glue::SearchStrategy, index::SearchStats, search},
provider::{Accessor, DataProvider},
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

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

DataProvider is no longer used in this file after removing FilteredSearchResults, but it’s still imported. CI runs with -Dwarnings, so this unused import will fail the build; please remove it (and keep Accessor which is still needed for get_element).

Suggested change
provider::{Accessor, DataProvider},
provider::Accessor,

Copilot uses AI. Check for mistakes.
utils::VectorRepr,
};
Expand All @@ -20,61 +18,6 @@ use crate::{
provider::{self, GarnetProvider},
};

/// Wraps `SearchResults` (FFI output buffer) with a provider and context so that
/// `SearchOutputBuffer<u32>` can be implemented: each internal ID (u32) is converted
/// to a `GarnetId` via `to_external_id` and written into the underlying buffer.
/// Start point (internal ID 0) is skipped since it has no external ID mapping.
///
/// TODO: Once diskann-providers >= 0.45.0 exports BetaAccessor/BetaComputer/Unwrap,
/// implement SearchStrategy<..., GarnetId> for BetaFilter directly and remove this wrapper.
struct FilteredSearchResults<'a, 'b, T: VectorRepr> {
inner: &'a mut SearchResults<'b>,
provider: &'a GarnetProvider<T>,
context: &'a Context,
}

impl<T: VectorRepr> SearchOutputBuffer<u32> for FilteredSearchResults<'_, '_, T> {
fn size_hint(&self) -> Option<usize> {
self.inner.size_hint()
}

fn push(&mut self, id: u32, distance: f32) -> diskann::graph::BufferState {
// Skip start point (internal ID 0) which has no external ID mapping
if id == 0 {
if let Some(sz) = self.inner.size_hint()
&& sz == 0
{
return diskann::graph::BufferState::Full;
} else {
return diskann::graph::BufferState::Available;
}
}

let eid = self.provider.to_external_id(self.context, id).unwrap();

self.inner.push(eid, distance)
}

fn current_len(&self) -> usize {
self.inner.current_len()
}

fn extend<Itr>(&mut self, itr: Itr) -> usize
where
Itr: IntoIterator<Item = (u32, f32)>,
{
let initial = self.inner.current_len();

for (id, dist) in itr {
if self.push(id, dist).is_full() {
break;
}
}

self.inner.current_len() - initial
}
}

/// Type-erased version of `DiskANNIndex<GarnetProvider>`.
/// All vector data is passed as untyped byte slices.
pub trait DynIndex: Send + Sync {
Expand Down Expand Up @@ -142,19 +85,7 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
let query = bytemuck::cast_slice::<u8, T>(data);
if let Some((labels, beta)) = filter {
let beta_filter = BetaFilter::new(FullPrecision, Arc::new(labels.clone()), beta);
let provider = self.inner.provider();
let mut filtered = FilteredSearchResults {
inner: output,
provider,
context,
};

let stats = self.search(&beta_filter, context, query, params, &mut filtered)?;

Ok(SearchStats {
result_count: filtered.current_len() as u32,
..stats
})
self.search(&beta_filter, context, query, params, output)
} else {
self.search(&FullPrecision, context, query, params, output)
}
Expand Down
1 change: 0 additions & 1 deletion diskann-garnet/src/labels.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
//!
//! ## Design
//!
//! Unlike CDB-DiskANN which calls back to C++ via `AreDocumentsIncluded()`,
//! Garnet sends a **pre-computed dense bitmap** where each bit position
//! corresponds to a DiskANN internal ID (`u32`). This makes `is_match()`
//! a simple bit lookup — no cross-FFI callback needed.
Expand Down