From 6bc1497c46c0b1cb5b26fa7d0bf8c3563733a877 Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Mon, 10 Aug 2026 10:41:48 -0500 Subject: [PATCH 1/8] [diskann-garnet] Support VRANDMEMBER --- Cargo.lock | 2 +- diskann-garnet/Cargo.toml | 2 +- diskann-garnet/diskann-garnet.nuspec | 2 +- diskann-garnet/src/dyn_index.rs | 12 ++ diskann-garnet/src/lib.rs | 228 +++++++++++++++++++++++++++ diskann-garnet/src/provider.rs | 52 +++++- 6 files changed, 294 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b28f10fe42..c0914502ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -603,7 +603,7 @@ dependencies = [ [[package]] name = "diskann-garnet" -version = "4.0.4" +version = "5.0.0" dependencies = [ "bytemuck", "crossbeam", diff --git a/diskann-garnet/Cargo.toml b/diskann-garnet/Cargo.toml index 4d79cbd1f5..810eeaf3e2 100644 --- a/diskann-garnet/Cargo.toml +++ b/diskann-garnet/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "diskann-garnet" -version = "4.0.4" +version = "5.0.0" edition = "2024" authors.workspace = true license.workspace = true diff --git a/diskann-garnet/diskann-garnet.nuspec b/diskann-garnet/diskann-garnet.nuspec index cf87ccfb73..e385f9f182 100644 --- a/diskann-garnet/diskann-garnet.nuspec +++ b/diskann-garnet/diskann-garnet.nuspec @@ -2,7 +2,7 @@ diskann-garnet - 4.0.4 + 5.0.0 docs/README.md Microsoft https://github.com/microsoft/DiskANN diff --git a/diskann-garnet/src/dyn_index.rs b/diskann-garnet/src/dyn_index.rs index fe9858bac8..44589a9766 100644 --- a/diskann-garnet/src/dyn_index.rs +++ b/diskann-garnet/src/dyn_index.rs @@ -70,6 +70,9 @@ pub(crate) trait DynIndex: Send + Sync { fn train_quantizer(&self, context: &Context) -> bool; fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize); + + fn random_members(&self, context: &Context, count: u32, output: &mut SearchResults<'_>) + -> bool; } impl DynIndex for DiskANNIndex> { @@ -187,4 +190,13 @@ impl DynIndex for DiskANNIndex> { .provider() .backfill_quant_vectors(context, task_idx, task_count); } + + fn random_members( + &self, + context: &Context, + count: u32, + output: &mut SearchResults<'_>, + ) -> bool { + self.inner.provider().random_members(context, count, output) + } } diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 2aa509f3aa..7c53ee47fc 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -120,6 +120,10 @@ impl SearchResults<'_> { id_index, } } + + fn push_id(&mut self, id: GarnetId) -> diskann::graph::BufferState { + self.push(Neighbor::new(id, 0.0)) + } } impl SearchOutputBuffer for SearchResults<'_> { @@ -835,6 +839,38 @@ pub unsafe extern "C" fn check_external_id_valid( index.inner.external_id_exists(&ctx, &id) } +/// # Safety +/// +/// FFI +#[unsafe(no_mangle)] +pub unsafe extern "C" fn random_members( + ctx: u64, + index_ptr: *const c_void, + count: u32, + output_ids: *mut u8, + output_ids_len: usize, +) -> bool { + let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); + + // Dummy buffer for distances + let mut output_distances = vec![0f32; output_ids_len / 5]; + let mut output = SearchResults::new( + output_ids, + output_ids_len, + output_distances.as_mut_ptr(), + output_distances.len(), + ); + + index.inner.random_members(&ctx, count, &mut output) +} + +/// # Safety +/// +/// FFI +#[unsafe(no_mangle)] +pub unsafe extern "C" fn search_neighbors() {} + #[cfg(test)] mod tests { use std::{mem, ptr}; @@ -844,6 +880,7 @@ mod tests { neighbor::Neighbor, }; use diskann_vector::distance::Metric; + use rand::Rng; use crate::{ Index, IndexState, PolyCow, SearchResults, VectorQuantType, drop_index, @@ -1069,4 +1106,195 @@ mod tests { drop_index(0, index_ptr); } } + + #[test] + fn random_members() { + let store = Store::new(); + let mut quant_needed = false; + + let index_ptr = unsafe { + super::create_index( + 0, + 2, + 0, + VectorQuantType::NoQuant, + Metric::L2.into(), + 10, + 8, + store.callbacks().read_callback(), + store.callbacks().write_callback(), + store.callbacks().delete_callback(), + store.callbacks().rmw_callback(), + store.callbacks().filter_callback(), + &mut quant_needed, + ) + }; + + assert!(!index_ptr.is_null()); + + let ctx = Context::new(0); + let mut rng = rand::rng(); + + for id in 0..100 { + let mut v = vec![0u8; 2]; + rng.fill(v.as_mut_slice()); + let v = v.into_iter().map(|i| i as f32).collect::>(); + + let eid = GarnetId::from(bytemuck::bytes_of(&id)); + assert_eq!( + unsafe { + super::insert( + ctx.get(), + index_ptr, + eid.as_ptr(), + eid.len(), + bytemuck::cast_slice::(&v).as_ptr(), + v.len(), + ptr::null(), + 0, + ) + }, + 1 + ); + } + + // Check basic correctness + let mut output_ids = vec![u32::MAX; 20]; + assert!(unsafe { + super::random_members( + ctx.get(), + index_ptr, + 10, + bytemuck::cast_slice_mut::(output_ids.as_mut_slice()).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + ) + }); + assert!( + output_ids + .iter() + .enumerate() + .all(|(i, e)| if i.is_multiple_of(2) { + *e == 4 + } else { + *e < 100 + }) + ); + + // Check undersized buffer + output_ids.fill(u32::MAX); + assert!(unsafe { + super::random_members( + ctx.get(), + index_ptr, + 20, + bytemuck::cast_slice_mut::(output_ids.as_mut_slice()).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + ) + }); + assert!( + output_ids + .iter() + .enumerate() + .all(|(i, e)| if i.is_multiple_of(2) { + *e == 4 + } else { + *e < 100 + }) + ); + + // Check oversized buffer + output_ids.fill(u32::MAX); + assert!(unsafe { + super::random_members( + ctx.get(), + index_ptr, + 5, + bytemuck::cast_slice_mut::(output_ids.as_mut_slice()).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + ) + }); + assert!(output_ids.iter().enumerate().all(|(i, e)| if i < 10 { + if i.is_multiple_of(2) { + *e == 4 + } else { + *e < 100 + } + } else { + *e == u32::MAX + })); + + // Delete 50 vectors at random + let ids = rand::seq::index::sample(&mut rng, 100, 50); + for id in ids { + let id = id as u32; + let eid = GarnetId::from(bytemuck::bytes_of(&id)); + + assert!(unsafe { super::remove(ctx.get(), index_ptr, eid.as_ptr(), eid.len()) }); + } + + // Check basic correctness + output_ids.fill(u32::MAX); + assert!(unsafe { + super::random_members( + ctx.get(), + index_ptr, + 10, + bytemuck::cast_slice_mut::(output_ids.as_mut_slice()).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + ) + }); + assert!( + output_ids + .iter() + .enumerate() + .all(|(i, e)| if i.is_multiple_of(2) { + *e == 4 + } else { + *e < 100 + }) + ); + + // Check undersized buffer + output_ids.fill(u32::MAX); + assert!(unsafe { + super::random_members( + ctx.get(), + index_ptr, + 20, + bytemuck::cast_slice_mut::(output_ids.as_mut_slice()).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + ) + }); + assert!( + output_ids + .iter() + .enumerate() + .all(|(i, e)| if i.is_multiple_of(2) { + *e == 4 + } else { + *e < 100 + }) + ); + + // Check oversized buffer + output_ids.fill(u32::MAX); + assert!(unsafe { + super::random_members( + ctx.get(), + index_ptr, + 5, + bytemuck::cast_slice_mut::(output_ids.as_mut_slice()).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + ) + }); + assert!(output_ids.iter().enumerate().all(|(i, e)| if i < 10 { + if i.is_multiple_of(2) { + *e == 4 + } else { + *e < 100 + } + } else { + *e == u32::MAX + })); + } } diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index e305ec626a..5585edd8cd 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -34,6 +34,7 @@ use diskann_vector::{ }; use std::{ any::TypeId, + collections::HashSet, future, marker::PhantomData, mem, @@ -46,7 +47,7 @@ use std::{ use thiserror::Error; use crate::{ - VectorQuantType, + SearchResults, VectorQuantType, alloc::AlignToEight, fsm::{FreeSpaceMap, FsmError}, garnet::{Callbacks, Context, GarnetError, GarnetId, Term}, @@ -672,6 +673,55 @@ impl GarnetProvider { } } + pub(crate) fn random_members( + &self, + context: &Context, + count: u32, + output: &mut SearchResults<'_>, + ) -> bool { + let mut rng = rand::rng(); + + let id_space = self.max_internal_id() as usize + 1; + let total_vectors = self.fsm.total_used(); + let mut remaining = (count as usize).min(total_vectors); + let mut chosen = HashSet::new(); + + // Deletions leave holes in the ID space, so scale the first batch by the density of + // live IDs, then grow it until the request is satisfied or the whole space is covered. + let mut batch = remaining + .saturating_mul(id_space) + .div_ceil(total_vectors.max(1)) + .clamp(1, id_space); + + while remaining > 0 { + for samp in rand::seq::index::sample(&mut rng, id_space, batch) { + let samp = samp as u32; + if !chosen.insert(samp) { + // Already considered on an earlier round. + continue; + } + let Ok(eid) = self.to_external_id(context, samp) else { + // Deleted or otherwise unreadable. + continue; + }; + + let state = output.push_id(eid); + remaining -= 1; + if remaining == 0 || state == diskann::graph::BufferState::Full { + return true; + } + } + + if batch == id_space { + // The whole ID space was scanned; fewer live IDs exist than were requested. + break; + } + batch = batch.saturating_mul(2).min(id_space); + } + + true + } + /// Returns the quantizer associated with the index. fn quantizer(&self) -> Option<&dyn GarnetQuantizer> { if let Some(quantizer) = &self.quantizer { From 30aaa3eca98621856b5f154cd77cacadfd3eff74 Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Mon, 10 Aug 2026 12:12:34 -0500 Subject: [PATCH 2/8] [diskann-garnet] Support VLINKS --- diskann-garnet/src/dyn_index.rs | 7 ++ diskann-garnet/src/lib.rs | 128 +++++++++++++++++++++++++++++++- diskann-garnet/src/provider.rs | 29 ++++++++ 3 files changed, 163 insertions(+), 1 deletion(-) diff --git a/diskann-garnet/src/dyn_index.rs b/diskann-garnet/src/dyn_index.rs index 44589a9766..5d805502cd 100644 --- a/diskann-garnet/src/dyn_index.rs +++ b/diskann-garnet/src/dyn_index.rs @@ -11,6 +11,7 @@ use crate::{ use diskann::{ ANNResult, graph::{InplaceDeleteMethod, index::SearchStats, search}, + neighbor::Neighbor, provider::DataProvider, utils::VectorRepr, }; @@ -73,6 +74,8 @@ pub(crate) trait DynIndex: Send + Sync { fn random_members(&self, context: &Context, count: u32, output: &mut SearchResults<'_>) -> bool; + + fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult>>; } impl DynIndex for DiskANNIndex> { @@ -199,4 +202,8 @@ impl DynIndex for DiskANNIndex> { ) -> bool { self.inner.provider().random_members(context, count, output) } + + fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult>> { + self.inner.provider().neighbors(context, id) + } } diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 7c53ee47fc..2256463faa 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -869,7 +869,37 @@ pub unsafe extern "C" fn random_members( /// /// FFI #[unsafe(no_mangle)] -pub unsafe extern "C" fn search_neighbors() {} +pub unsafe extern "C" fn search_neighbors( + ctx: u64, + index_ptr: *const c_void, + id_data: *const u8, + id_len: usize, + output_ids: *mut u8, + output_ids_len: usize, + output_distances: *mut f32, + output_distances_len: usize, + _continuation: *mut c_void, +) -> i32 { + let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); + let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + let id = GarnetId::from(id_bytes); + + let mut output = SearchResults::new( + output_ids, + output_ids_len, + output_distances, + output_distances_len, + ); + + let Ok(neighbors) = index.inner.neighbors(&ctx, &id) else { + return -1; + }; + + output.extend(neighbors); + + output.current_len() as i32 +} #[cfg(test)] mod tests { @@ -1297,4 +1327,100 @@ mod tests { *e == u32::MAX })); } + + #[test] + fn search_neighbors() { + let store = Store::new(); + let mut quant_needed = false; + + let index_ptr = unsafe { + super::create_index( + 0, + 2, + 0, + VectorQuantType::NoQuant, + Metric::L2.into(), + 10, + 8, + store.callbacks().read_callback(), + store.callbacks().write_callback(), + store.callbacks().delete_callback(), + store.callbacks().rmw_callback(), + store.callbacks().filter_callback(), + &mut quant_needed, + ) + }; + + assert!(!index_ptr.is_null()); + + let ctx = Context::new(0); + let mut rng = rand::rng(); + + for id in 0..100 { + let mut v = vec![0u8; 2]; + rng.fill(v.as_mut_slice()); + let v = v.into_iter().map(|i| i as f32).collect::>(); + + let eid = GarnetId::from(bytemuck::bytes_of(&id)); + assert_eq!( + unsafe { + super::insert( + ctx.get(), + index_ptr, + eid.as_ptr(), + eid.len(), + bytemuck::cast_slice::(&v).as_ptr(), + v.len(), + ptr::null(), + 0, + ) + }, + 1 + ); + } + + let mut output_ids = vec![u32::MAX; 20]; + let mut output_dists = vec![f32::MAX; 10]; + let good_id = GarnetId::from(bytemuck::bytes_of(&25u32)); + let bad_id = GarnetId::from(bytemuck::bytes_of(&250u32)); + + // check the good case + let count = unsafe { + super::search_neighbors( + ctx.get(), + index_ptr, + good_id.as_ptr(), + good_id.len(), + bytemuck::cast_slice_mut(&mut output_ids).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + output_dists.as_mut_ptr(), + output_dists.len(), + ptr::null_mut(), + ) + }; + + assert!(count > 0 && count <= 8, "count = {count}"); + + for i in 0..count as usize { + assert_eq!(output_ids[i * 2], 4); + assert!(output_ids[i * 2 + 1] < 100); + assert!(output_dists[i] < f32::MAX); + } + + let count = unsafe { + super::search_neighbors( + ctx.get(), + index_ptr, + bad_id.as_ptr(), + bad_id.len(), + bytemuck::cast_slice_mut(&mut output_ids).as_mut_ptr(), + output_ids.len() * mem::size_of::(), + output_dists.as_mut_ptr(), + output_dists.len(), + ptr::null_mut(), + ) + }; + + assert!(count < 0); + } } diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 5585edd8cd..4ddae764a9 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -722,6 +722,35 @@ impl GarnetProvider { true } + pub(crate) fn neighbors( + &self, + context: &Context, + id: &GarnetId, + ) -> ANNResult>> { + let iid = self.to_internal_id(context, id)?; + let v = self.get_full_vector(context, iid)?; + let mut neighbors = AdjacencyList::with_capacity(self.max_degree + 1); + + if !self.get_neighbors(context, iid, &mut neighbors) { + return Err(GarnetProviderError::Garnet(GarnetError::Read).into()); + } + + let d = ::distance(self.metric_type, Some(self.dim)); + let mut result = Vec::with_capacity(self.max_degree); + for &nbr_id in neighbors.iter() { + if nbr_id == 0 { + // Skip the start point + continue; + } + let nbr_v = self.get_full_vector(context, nbr_id)?; + let nbr_eid = self.to_external_id(context, nbr_id)?; + let nbr_d = d.evaluate_similarity(&v, &nbr_v); + result.push(Neighbor::new(nbr_eid, nbr_d)); + } + + Ok(result) + } + /// Returns the quantizer associated with the index. fn quantizer(&self) -> Option<&dyn GarnetQuantizer> { if let Some(quantizer) = &self.quantizer { From 0ebac0fb2697329e087841dbfff828ff228112ba Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Thu, 13 Aug 2026 16:10:53 -0500 Subject: [PATCH 3/8] [diskann-garnet] Add logCallback --- diskann-garnet/src/ffi_recall_tests.rs | 1 + diskann-garnet/src/ffi_tests.rs | 1 + diskann-garnet/src/garnet.rs | 18 ++++++++++++++ diskann-garnet/src/lib.rs | 9 ++++++- diskann-garnet/src/provider.rs | 8 +++++++ diskann-garnet/src/test_utils.rs | 33 ++++++++++++++++++++++++-- 6 files changed, 67 insertions(+), 3 deletions(-) diff --git a/diskann-garnet/src/ffi_recall_tests.rs b/diskann-garnet/src/ffi_recall_tests.rs index 637363a433..43d5ef75fe 100644 --- a/diskann-garnet/src/ffi_recall_tests.rs +++ b/diskann-garnet/src/ffi_recall_tests.rs @@ -175,6 +175,7 @@ mod tests { callbacks.delete_callback(), callbacks.rmw_callback(), callbacks.filter_callback(), + callbacks.log_callback(), &mut quant_needed, ) }; diff --git a/diskann-garnet/src/ffi_tests.rs b/diskann-garnet/src/ffi_tests.rs index 1582413859..8e2356d5ad 100644 --- a/diskann-garnet/src/ffi_tests.rs +++ b/diskann-garnet/src/ffi_tests.rs @@ -60,6 +60,7 @@ mod tests { callbacks.delete_callback(), callbacks.rmw_callback(), callbacks.filter_callback(), + callbacks.log_callback(), &mut quant_needed, ) }; diff --git a/diskann-garnet/src/garnet.rs b/diskann-garnet/src/garnet.rs index b43d798710..8f7bbed272 100644 --- a/diskann-garnet/src/garnet.rs +++ b/diskann-garnet/src/garnet.rs @@ -86,6 +86,7 @@ pub(crate) type ReadModifyWriteCallback = pub(crate) type ReadDataCallback = unsafe extern "C" fn(u32, *mut c_void, *const u8, usize); pub(crate) type RmwDataCallback = unsafe extern "C" fn(*mut c_void, *mut u8, usize); pub(crate) type FilterCallback = unsafe extern "C" fn(u64, u32) -> bool; +pub(crate) type LogCallback = unsafe extern "C" fn(u64, *const u8, usize); #[derive(Copy, Clone)] pub(crate) struct Callbacks { @@ -94,6 +95,7 @@ pub(crate) struct Callbacks { delete_callback: DeleteCallback, rmw_callback: ReadModifyWriteCallback, filter_callback: FilterCallback, + log_callback: LogCallback, } impl Callbacks { @@ -103,6 +105,7 @@ impl Callbacks { delete_callback: DeleteCallback, rmw_callback: ReadModifyWriteCallback, filter_callback: FilterCallback, + log_callback: LogCallback, ) -> Self { Self { read_callback, @@ -110,6 +113,7 @@ impl Callbacks { delete_callback, rmw_callback, filter_callback, + log_callback, } } @@ -138,6 +142,11 @@ impl Callbacks { self.filter_callback } + #[cfg(test)] + pub(crate) fn log_callback(&self) -> LogCallback { + self.log_callback + } + #[cfg(test)] pub(crate) fn exists_iid(&self, ctx: &Context, id: u32) -> bool { let key = [4, id]; @@ -492,6 +501,15 @@ impl Callbacks { pub(crate) fn matches_filter(&self, ctx: &Context, id: u32) -> bool { unsafe { (self.filter_callback)(ctx.inner, id) } } + + /// Log a message to Garnet. + /// + /// The context bits can be set with appropriate `Term` to flag which area the log message concerns. + pub(crate) fn log(&self, ctx: &Context, msg: &str) { + unsafe { + (self.log_callback)(ctx.inner, msg.as_ptr(), msg.len()); + } + } } unsafe extern "C" fn read_call<'a, F, T>(index: u32, ptr: *mut c_void, data: *const u8, len: usize) diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 2256463faa..ca70901f6f 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -29,7 +29,7 @@ use diskann_vector::distance::Metric; use crate::{ alloc::AlignToEight, - garnet::FilterCallback, + garnet::{FilterCallback, LogCallback}, provider::{GarnetProvider, GarnetProviderError}, }; use crate::{ @@ -236,6 +236,7 @@ pub unsafe extern "C" fn create_index( delete_callback: DeleteCallback, rmw_callback: ReadModifyWriteCallback, filter_callback: FilterCallback, + log_callback: LogCallback, quantization_needed: *mut bool, ) -> *const c_void { unsafe { *quantization_needed = false }; @@ -267,6 +268,7 @@ pub unsafe extern "C" fn create_index( delete_callback, rmw_callback, filter_callback, + log_callback, ); match quant_type { @@ -995,6 +997,7 @@ mod tests { store.callbacks().delete_callback(), store.callbacks().rmw_callback(), store.callbacks().filter_callback(), + store.callbacks().log_callback(), &mut quant_needed, ) }; @@ -1027,6 +1030,7 @@ mod tests { store.callbacks().delete_callback(), store.callbacks().rmw_callback(), store.callbacks().filter_callback(), + store.callbacks().log_callback(), &mut quant_needed, ) }; @@ -1086,6 +1090,7 @@ mod tests { store.callbacks().delete_callback(), store.callbacks().rmw_callback(), store.callbacks().filter_callback(), + store.callbacks().log_callback(), &mut quant_needed, ) }; @@ -1156,6 +1161,7 @@ mod tests { store.callbacks().delete_callback(), store.callbacks().rmw_callback(), store.callbacks().filter_callback(), + store.callbacks().log_callback(), &mut quant_needed, ) }; @@ -1347,6 +1353,7 @@ mod tests { store.callbacks().delete_callback(), store.callbacks().rmw_callback(), store.callbacks().filter_callback(), + store.callbacks().log_callback(), &mut quant_needed, ) }; diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 4ddae764a9..cb6acb5be4 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -647,6 +647,10 @@ impl GarnetProvider { // NOTE: This return is unrecoverable in the current design, as there is no way // to signal that backfill has failed. The index will operate on full precision // vectors only from now on. + self.callbacks.log( + &context.term(Term::Quantized), + "Error quantizing start point; failed to finish backfill. Index will operate full precision only mode.", + ); return; }; self.start_point_quant_cache.insert(0, point); @@ -665,6 +669,10 @@ impl GarnetProvider { ) { // NOTE: This return is unrecoverable in the current design, as there is no way to // signal that backfill failed. + self.callbacks.log( + &context.term(Term::Quantized), + "Error saving quantizer state; failed to finish backfill. Index will operate full precision only mode.", + ); return; } diff --git a/diskann-garnet/src/test_utils.rs b/diskann-garnet/src/test_utils.rs index cb8babeb59..7a3e55173c 100644 --- a/diskann-garnet/src/test_utils.rs +++ b/diskann-garnet/src/test_utils.rs @@ -8,11 +8,15 @@ use core::slice; use dashmap::DashMap; use std::{ ffi::c_void, - sync::atomic::{AtomicUsize, Ordering}, + sync::{ + Mutex, + atomic::{AtomicUsize, Ordering}, + }, }; thread_local! { pub static STORE: DashMap, Vec> = DashMap::new(); + pub static LOGS: Mutex> = const { Mutex::new(Vec::new()) }; pub static FULL_READS: AtomicUsize = const { AtomicUsize::new(0) }; pub static QUANT_READS: AtomicUsize = const { AtomicUsize::new(0) }; } @@ -36,11 +40,22 @@ impl Store { } pub fn callbacks(&self) -> Callbacks { - Callbacks::new(test_read, test_write, test_delete, test_rmw, test_filter) + Callbacks::new( + test_read, + test_write, + test_delete, + test_rmw, + test_filter, + test_log, + ) } pub fn clear(&self) { STORE.with(|s| s.clear()); + LOGS.with(|l| { + let mut guard = l.lock().unwrap(); + guard.clear(); + }); FULL_READS.with(|fr| fr.store(0, Ordering::Release)); QUANT_READS.with(|qr| qr.store(0, Ordering::Release)); } @@ -96,6 +111,13 @@ impl Store { pub fn quant_reads(&self) -> usize { QUANT_READS.with(|qr| qr.load(Ordering::Acquire)) } + + pub fn log(&self, context: u64, msg: &str) { + LOGS.with(|l| { + let mut guard = l.lock().unwrap(); + guard.push((context, msg.to_owned())); + }); + } } unsafe extern "C" fn test_read( @@ -180,6 +202,13 @@ unsafe extern "C" fn test_filter(_context: u64, _internal_id: u32) -> bool { true } +unsafe extern "C" fn test_log(context: u64, msg: *const u8, msg_len: usize) { + let store = Store::attach(); + let msg_slice = unsafe { slice::from_raw_parts(msg, msg_len) }; + let msg = str::from_utf8(msg_slice).unwrap(); + store.log(context, msg) +} + mod tests { use std::collections::HashMap; From 173850920d90bdafbed31eba91de7c85e71f0cd4 Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Thu, 13 Aug 2026 17:46:48 -0500 Subject: [PATCH 4/8] [diskann-garnet] Add valueLengthHint to readCallback --- diskann-garnet/src/fsm.rs | 11 ++- diskann-garnet/src/garnet.rs | 25 +++-- diskann-garnet/src/provider.rs | 156 ++++++++++++++++++++++--------- diskann-garnet/src/test_utils.rs | 9 +- 4 files changed, 139 insertions(+), 62 deletions(-) diff --git a/diskann-garnet/src/fsm.rs b/diskann-garnet/src/fsm.rs index 8f8079317b..b5fb57006d 100644 --- a/diskann-garnet/src/fsm.rs +++ b/diskann-garnet/src/fsm.rs @@ -143,7 +143,7 @@ impl FreeSpaceMap { let block_key = Self::block_key(0); if this .callbacks - .exists_wid(&ctx.term(Term::Metadata), block_key) + .exists_wid(&ctx.term(Term::Metadata), block_key, BLOCK_SIZE_BYTES) { this.load_state(ctx)?; } else { @@ -159,10 +159,11 @@ impl FreeSpaceMap { /// Load all state from Garnet by scanning the FSM blocks. fn load_state(&mut self, ctx: &Context) -> Result<(), FsmError> { let mut max_block_id = 0; - while self - .callbacks - .exists_wid(&ctx.term(Term::Metadata), Self::block_key(max_block_id)) - { + while self.callbacks.exists_wid( + &ctx.term(Term::Metadata), + Self::block_key(max_block_id), + BLOCK_SIZE_BYTES, + ) { max_block_id += 1; } diff --git a/diskann-garnet/src/garnet.rs b/diskann-garnet/src/garnet.rs index 8f7bbed272..e3f64388f1 100644 --- a/diskann-garnet/src/garnet.rs +++ b/diskann-garnet/src/garnet.rs @@ -77,7 +77,7 @@ impl Context { impl ExecutionContext for Context {} pub(crate) type ReadCallback = - unsafe extern "C" fn(u64, u32, *const u8, usize, ReadDataCallback, *mut c_void); + unsafe extern "C" fn(u64, u32, u32, *const u8, usize, ReadDataCallback, *mut c_void); pub(crate) type WriteCallback = unsafe extern "C" fn(u64, *const u8, usize, *const u8, usize) -> bool; pub(crate) type DeleteCallback = unsafe extern "C" fn(u64, *const u8, usize) -> bool; @@ -148,34 +148,34 @@ impl Callbacks { } #[cfg(test)] - pub(crate) fn exists_iid(&self, ctx: &Context, id: u32) -> bool { + pub(crate) fn exists_iid(&self, ctx: &Context, id: u32, length_hint: usize) -> bool { let key = [4, id]; // SAFETY: Key bytes are preceded by 4 bytes of space. - unsafe { self.exists_raw(ctx, bytemuck::bytes_of(&key)) } + unsafe { self.exists_raw(ctx, bytemuck::bytes_of(&key), length_hint) } } - pub(crate) fn exists_wid(&self, ctx: &Context, key: u64) -> bool { + pub(crate) fn exists_wid(&self, ctx: &Context, key: u64, length_hint: usize) -> bool { // NOTE: the length is bit-shifted so that we have a u32 in the lower half of the u64. let mut key = [8 << 32, key]; let key_bytes = bytemuck::bytes_of_mut(&mut key); // SAFETY: Key bytes are preceded by 8 bytes of extra space. - unsafe { self.exists_raw(ctx, &key_bytes[4..]) } + unsafe { self.exists_raw(ctx, &key_bytes[4..], length_hint) } } #[expect( dead_code, reason = "currently unused, but may be needed in the future" )] - pub(crate) fn exists_eid(&self, ctx: &Context, id: &GarnetId) -> bool { + pub(crate) fn exists_eid(&self, ctx: &Context, id: &GarnetId, length_hint: usize) -> bool { // SAFETY: GarnetId ensures there are 4 bytes preceding the key bytes. - unsafe { self.exists_raw(ctx, id) } + unsafe { self.exists_raw(ctx, id, length_hint) } } /// Check for a key's existance in Garnet. /// /// NOTE: The key bytes must be preceded by 4 valid bytes that Garnet can write into. /// This invariant must be checked by the caller. - unsafe fn exists_raw(&self, ctx: &Context, key: &[u8]) -> bool { + unsafe fn exists_raw(&self, ctx: &Context, key: &[u8], length_hint: usize) -> bool { let mut called = false; let mut cb = |_, _: &[u8]| { called = true; @@ -185,6 +185,7 @@ impl Callbacks { (self.read_callback)( ctx.inner, 1, + length_hint as u32, key.as_ptr(), key.len(), make_read_call(&cb), @@ -256,6 +257,7 @@ impl Callbacks { /// This invariant must be checked by the caller. #[must_use] unsafe fn read_single_raw(&self, ctx: &Context, key: &[u8], value: &mut [u8]) -> bool { + let length_hint = value.len() as u32; let mut found = false; let mut cb = |_, data: &[u8]| { found = true; @@ -266,6 +268,7 @@ impl Callbacks { (self.read_callback)( ctx.inner, 1, + length_hint, key.as_ptr(), key.len(), make_read_call(&cb), @@ -281,6 +284,7 @@ impl Callbacks { &self, ctx: &Context, ids: &[u32], + length_hint: usize, mut f: F, ) where F: FnMut(u32, &'a [T]), @@ -293,6 +297,7 @@ impl Callbacks { (self.read_callback)( ctx.inner, ids.len() as u32 / 2, + length_hint as u32, bytemuck::must_cast_slice::<_, u8>(ids).as_ptr(), mem::size_of_val(ids), make_read_call(&f), @@ -325,11 +330,15 @@ impl Callbacks { result = Some(bytemuck::cast_slice::(data).to_owned()); }; + // NOTE: We hint the length as 8192 bytes, which will often overestimate. The only varsize + // things to read are the quant state and the external ID map. Quant state is + // maximum `117 + 6 * dim` bytes, which is several kilobytes in practice. // SAFETY: Key bytes are preceded by 4 bytes of extra space. unsafe { (self.read_callback)( ctx.inner, 1, + 8192, bytemuck::bytes_of(&key).as_ptr(), mem::size_of_val(&key), make_read_call(&cb), diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index cb6acb5be4..c83f897cdc 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -912,6 +912,27 @@ impl GarnetProvider { Ok(()) } + + /// The size of a stored full vector. + fn full_vector_size(&self) -> usize { + self.dim * mem::size_of::() + } + + /// The size of a stored quant vector. + fn quant_vector_size(&self) -> usize { + if let Some(quantizer) = &self.quantizer { + quantizer.bytes() + } else { + 0 + } + } + + /// Provides an estimate of quantizer state size. + /// This is allowed to be wrong, but should ideally be an overestimate. + #[cfg(test)] + fn quant_state_size(&self) -> usize { + self.dim * 6 + 128 + } } impl DataProvider for GarnetProvider { @@ -1248,19 +1269,28 @@ impl SearchAccessor for DynamicAccessor<'_, T> { } } - let ctx = if self.quantized { - self.context.term(Term::Quantized) + let (ctx, length_hint) = if self.quantized { + ( + self.context.term(Term::Quantized), + self.provider.quant_vector_size(), + ) } else { - self.context.term(Term::Vector) + ( + self.context.term(Term::Vector), + self.provider.full_vector_size(), + ) }; if !self.filtered_ids.is_empty() { - self.provider - .callbacks - .read_multi_lpiid(&ctx, &self.filtered_ids, |i, v| { + self.provider.callbacks.read_multi_lpiid( + &ctx, + &self.filtered_ids, + length_hint, + |i, v| { let dist = self.computer.evaluate_similarity(v); on_neighbors(self.filtered_ids[i as usize * 2 + 1], dist); - }); + }, + ); } } @@ -1420,6 +1450,7 @@ impl<'a, 'b, T: VectorRepr> SearchPostProcessStep, &'b [T provider.callbacks.read_multi_lpiid( &accessor.context.term(Term::Vector), &accessor.filtered_ids, + provider.full_vector_size(), |i, v| { let dist = f.evaluate_similarity(query, bytemuck::cast_slice::(v)); reranked.push(Neighbor::new( @@ -1498,16 +1529,24 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { } } - let ctx = if self.quantized { - self.context.term(Term::Quantized) + let (ctx, length_hint) = if self.quantized { + ( + self.context.term(Term::Quantized), + self.provider.quant_vector_size(), + ) } else { - self.context.term(Term::Vector) + ( + self.context.term(Term::Vector), + self.provider.full_vector_size(), + ) }; if !self.filtered_ids.is_empty() { - self.provider - .callbacks - .read_multi_lpiid(&ctx, &self.filtered_ids, |i, v| { + self.provider.callbacks.read_multi_lpiid( + &ctx, + &self.filtered_ids, + length_hint, + |i, v| { let dist = self.computer.evaluate_similarity(v); let decision = if self.filtered_decisions[i as usize] { Decision::accept(self.filtered_ids[i as usize * 2 + 1]) @@ -1515,7 +1554,8 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { Decision::reject(self.filtered_ids[i as usize * 2 + 1]) }; on_neighbors(decision, dist); - }); + }, + ); } } @@ -1554,19 +1594,28 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { } } - let ctx = if self.quantized { - self.context.term(Term::Quantized) + let (ctx, length_hint) = if self.quantized { + ( + self.context.term(Term::Quantized), + self.provider.quant_vector_size(), + ) } else { - self.context.term(Term::Vector) + ( + self.context.term(Term::Vector), + self.provider.full_vector_size(), + ) }; if !self.filtered_ids.is_empty() { - self.provider - .callbacks - .read_multi_lpiid(&ctx, &self.filtered_ids, |i, v| { + self.provider.callbacks.read_multi_lpiid( + &ctx, + &self.filtered_ids, + length_hint, + |i, v| { let dist = self.computer.evaluate_similarity(v); on_neighbors(Accept::new(self.filtered_ids[i as usize * 2 + 1]), dist); - }); + }, + ); } } @@ -1706,19 +1755,28 @@ where } } - let ctx = if self.quantized { - self.context.term(Term::Quantized) + let (ctx, length_hint) = if self.quantized { + ( + self.context.term(Term::Quantized), + self.provider.quant_vector_size(), + ) } else { - self.context.term(Term::Vector) + ( + self.context.term(Term::Vector), + self.provider.full_vector_size(), + ) }; if !self.filtered_ids.is_empty() { - self.provider - .callbacks - .read_multi_lpiid(&ctx, &self.filtered_ids, |id, v| { + self.provider.callbacks.read_multi_lpiid( + &ctx, + &self.filtered_ids, + length_hint, + |id, v| { self.set .insert(self.filtered_ids[id as usize * 2 + 1], v.into()); - }); + }, + ); } Ok((self.set.view(), &self.distance)) @@ -1995,9 +2053,11 @@ mod tests { // There should be no saved quant state. assert!( - !provider - .callbacks - .exists_iid(&ctx.term(Term::Metadata), QUANT_STATE_KEY), + !provider.callbacks.exists_iid( + &ctx.term(Term::Metadata), + QUANT_STATE_KEY, + provider.quant_state_size() + ), "quant state should not be stored yet" ); @@ -2069,9 +2129,11 @@ mod tests { // There should be saved quant state. assert!( - provider - .callbacks - .exists_iid(&ctx.term(Term::Metadata), QUANT_STATE_KEY), + provider.callbacks.exists_iid( + &ctx.term(Term::Metadata), + QUANT_STATE_KEY, + provider.quant_state_size() + ), "quant state missing" ); @@ -2167,9 +2229,11 @@ mod tests { // There should be saved quant state. assert!( - provider - .callbacks - .exists_iid(&ctx.term(Term::Metadata), QUANT_STATE_KEY), + provider.callbacks.exists_iid( + &ctx.term(Term::Metadata), + QUANT_STATE_KEY, + provider.quant_state_size() + ), "quant state missing" ); @@ -2183,11 +2247,11 @@ mod tests { // Every quant vector should be present in the store for id in 0..last_inserted_id { - assert!( - provider - .callbacks - .exists_iid(&ctx.term(Term::Quantized), id) - ); + assert!(provider.callbacks.exists_iid( + &ctx.term(Term::Quantized), + id, + provider.quant_vector_size() + )); } // Searches should still work and use quantized vectors @@ -2256,9 +2320,11 @@ mod tests { // There should be saved quant state. assert!( - provider - .callbacks - .exists_iid(&ctx.term(Term::Metadata), QUANT_STATE_KEY), + provider.callbacks.exists_iid( + &ctx.term(Term::Metadata), + QUANT_STATE_KEY, + provider.quant_state_size() + ), "quant state missing" ); diff --git a/diskann-garnet/src/test_utils.rs b/diskann-garnet/src/test_utils.rs index 7a3e55173c..1aaa6e30c4 100644 --- a/diskann-garnet/src/test_utils.rs +++ b/diskann-garnet/src/test_utils.rs @@ -123,6 +123,7 @@ impl Store { unsafe extern "C" fn test_read( ctx: u64, count: u32, + _length_hint: u32, id_bytes: *const u8, id_len: usize, cb: ReadDataCallback, @@ -225,7 +226,7 @@ mod tests { let ctx = Context::new(0); // Reading a non-existant key should fail. - assert!(!callbacks.exists_iid(&ctx, 0)); + assert!(!callbacks.exists_iid(&ctx, 0, 10)); // Round tripping a write should work. assert!(callbacks.write_iid(&ctx, 0, b"test")); @@ -240,9 +241,9 @@ mod tests { assert_eq!(val, b"again"); // Exists and delete should work. - assert!(callbacks.exists_iid(&ctx, 0)); + assert!(callbacks.exists_iid(&ctx, 0, 10)); assert!(callbacks.delete_iid(&ctx, 0)); - assert!(!callbacks.exists_iid(&ctx, 0)); + assert!(!callbacks.exists_iid(&ctx, 0, 10)); // Different contexts should stay separate. assert!(callbacks.write_iid(&ctx.term(Term::Vector), 0, b"0000")); @@ -257,7 +258,7 @@ mod tests { assert!(callbacks.write_iid(&ctx.term(Term::Vector), 1, b"2222")); let ids = [4u32, 0, 4, 1, 4, 2]; let mut results = HashMap::new(); - callbacks.read_multi_lpiid(&ctx.term(Term::Vector), &ids, |i, v| { + callbacks.read_multi_lpiid(&ctx.term(Term::Vector), &ids, 10, |i, v| { results.insert(i, v.to_owned()); }); assert_eq!(results.get(&0), Some(b"0000".to_vec()).as_ref()); From bc7a159b7bbec59cf9239bdc7cf0c615eadc2dd4 Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Thu, 13 Aug 2026 18:26:39 -0500 Subject: [PATCH 5/8] [diskann-garnet] Add result to backfill_quant_vectors --- diskann-garnet/src/dyn_index.rs | 12 ++++++++--- diskann-garnet/src/lib.rs | 4 ++-- diskann-garnet/src/provider.rs | 35 ++++++++++++++++++++++----------- 3 files changed, 34 insertions(+), 17 deletions(-) diff --git a/diskann-garnet/src/dyn_index.rs b/diskann-garnet/src/dyn_index.rs index 5d805502cd..3d4da83c32 100644 --- a/diskann-garnet/src/dyn_index.rs +++ b/diskann-garnet/src/dyn_index.rs @@ -70,7 +70,8 @@ pub(crate) trait DynIndex: Send + Sync { fn train_quantizer(&self, context: &Context) -> bool; - fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize); + fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize) + -> bool; fn random_members(&self, context: &Context, count: u32, output: &mut SearchResults<'_>) -> bool; @@ -188,10 +189,15 @@ impl DynIndex for DiskANNIndex> { self.inner.provider().train_quantizer(context) } - fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize) { + fn backfill_quant_vectors( + &self, + context: &Context, + task_idx: usize, + task_count: usize, + ) -> bool { self.inner .provider() - .backfill_quant_vectors(context, task_idx, task_count); + .backfill_quant_vectors(context, task_idx, task_count) } fn random_members( diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index ca70901f6f..5928ac7fc6 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -566,12 +566,12 @@ pub unsafe extern "C" fn backfill_quant_vectors( index_ptr: *const c_void, task_index: usize, task_count: usize, -) { +) -> bool { let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(context); index .inner - .backfill_quant_vectors(&ctx, task_index, task_count); + .backfill_quant_vectors(&ctx, task_index, task_count) } /// # Safety diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index c83f897cdc..900acf6ab3 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -575,22 +575,36 @@ impl GarnetProvider { context: &Context, task_idx: usize, task_count: usize, - ) { + ) -> bool { let quantizer = match &self.quantizer { Some(q) => q, - None => return, + None => { + self.callbacks.log( + &context.term(Term::Quantized), + "Error: backfill_quant_vectors: Quantizer not found. Index will operate full precision only mode.", + ); + return false; + } }; let max_id = self.fsm.max_id_for_backfill() as usize; if max_id >= u32::MAX as usize { // The max_id was somehow not sampled, so bail. - return; + self.callbacks.log( + &context.term(Term::Quantized), + "Error: backfill_quant_vectors: Couldn't calculate max id to backfill. Index will operate full precision only mode.", + ); + return false; } // If we have more tasks than vectors to backfill, we exit the extra tasks early. let task_count = task_count.min(max_id + 1); if task_idx >= task_count { - return; + self.callbacks.log( + &context.term(Term::Quantized), + "Error: backfill_quant_vectors: Bad task index. Index will operate full precision only mode.", + ); + return false; } // Evenly divide the ID range from 0..max_id and determine this thread's backfill @@ -644,14 +658,11 @@ impl GarnetProvider { let point = if let Ok(p) = Poly::from_iter(q.iter().copied(), AlignToEight) { p } else { - // NOTE: This return is unrecoverable in the current design, as there is no way - // to signal that backfill has failed. The index will operate on full precision - // vectors only from now on. self.callbacks.log( &context.term(Term::Quantized), "Error quantizing start point; failed to finish backfill. Index will operate full precision only mode.", ); - return; + return false; }; self.start_point_quant_cache.insert(0, point); } @@ -667,18 +678,18 @@ impl GarnetProvider { data[0] = 1; }, ) { - // NOTE: This return is unrecoverable in the current design, as there is no way to - // signal that backfill failed. self.callbacks.log( &context.term(Term::Quantized), "Error saving quantizer state; failed to finish backfill. Index will operate full precision only mode.", ); - return; + return false; } // Signal to the index that it is now safe to operate in quantized mode. self.all_quantized.store(true, Ordering::Release); } + + true } pub(crate) fn random_members( @@ -2214,7 +2225,7 @@ mod tests { // Run backfill for job_id in 0..4 { - provider.backfill_quant_vectors(&ctx, job_id, 4); + assert!(provider.backfill_quant_vectors(&ctx, job_id, 4)); } // Drop and re-create the index, keeping the same backing store From 4ccec761f5a2003d66679caf02e4d2ee8d9539ea Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Fri, 14 Aug 2026 10:51:30 -0500 Subject: [PATCH 6/8] [diskann-garnet] Change filter callback to take slices --- diskann-garnet/src/garnet.rs | 16 +++- diskann-garnet/src/lib.rs | 28 +++--- diskann-garnet/src/provider.rs | 153 +++++++++++++++++++++++-------- diskann-garnet/src/test_utils.rs | 2 +- 4 files changed, 141 insertions(+), 58 deletions(-) diff --git a/diskann-garnet/src/garnet.rs b/diskann-garnet/src/garnet.rs index e3f64388f1..d1547b1f4b 100644 --- a/diskann-garnet/src/garnet.rs +++ b/diskann-garnet/src/garnet.rs @@ -85,7 +85,7 @@ pub(crate) type ReadModifyWriteCallback = unsafe extern "C" fn(u64, *const u8, usize, usize, RmwDataCallback, *mut c_void) -> bool; pub(crate) type ReadDataCallback = unsafe extern "C" fn(u32, *mut c_void, *const u8, usize); pub(crate) type RmwDataCallback = unsafe extern "C" fn(*mut c_void, *mut u8, usize); -pub(crate) type FilterCallback = unsafe extern "C" fn(u64, u32) -> bool; +pub(crate) type FilterCallback = unsafe extern "C" fn(u64, *const u8, usize) -> bool; pub(crate) type LogCallback = unsafe extern "C" fn(u64, *const u8, usize); #[derive(Copy, Clone)] @@ -507,8 +507,18 @@ impl Callbacks { /// Evaluate the filter callback on an ID. #[must_use] - pub(crate) fn matches_filter(&self, ctx: &Context, id: u32) -> bool { - unsafe { (self.filter_callback)(ctx.inner, id) } + pub(crate) fn matches_filter(&self, ctx: &Context, data: &[u8]) -> bool { + unsafe { + (self.filter_callback)( + ctx.inner, + if data.is_empty() { + std::ptr::null() + } else { + data.as_ptr() + }, + data.len(), + ) + } } /// Log a message to Garnet. diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 5928ac7fc6..8a3385cc35 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -471,20 +471,21 @@ pub unsafe extern "C" fn insert( return InsertResult::Fail.into(); }; - // Write attributes to garnet - let attr_data = if attribute_len > 0 && !attribute_data.is_null() { - unsafe { slice::from_raw_parts(attribute_data, attribute_len) } - } else { - &[] - }; - if index.inner.set_attributes(&ctx, &id, attr_data).is_err() { - return InsertResult::Fail.into(); - } - let old_ready = ctx.quantizer_ready(); // Insert the vector if index.inner.insert(&ctx, &id, &v).is_ok() { + // Write attributes to garnet. These are written after insert since + // they are keyed on internal id. + let attr_data = if attribute_len > 0 && !attribute_data.is_null() { + unsafe { slice::from_raw_parts(attribute_data, attribute_len) } + } else { + &[] + }; + if index.inner.set_attributes(&ctx, &id, attr_data).is_err() { + return InsertResult::Fail.into(); + } + let ready = ctx.quantizer_ready(); if !old_ready && ready { InsertResult::SuccessStartTraining.into() @@ -1103,8 +1104,6 @@ mod tests { let ctx = Context::new(0); let v = [0.0f32, 0.0f32]; - assert!(store.get(ctx.term(Term::Attributes).get(), &eid).is_none()); - assert_eq!( unsafe { super::insert( @@ -1120,8 +1119,9 @@ mod tests { }, 1 ); + let iid = store.get(ctx.term(Term::IntMap).get(), &eid).unwrap(); assert_eq!( - store.get(ctx.term(Term::Attributes).get(), &eid), + store.get(ctx.term(Term::Attributes).get(), &iid), Some(metadata.as_slice().to_owned()) ); @@ -1135,7 +1135,7 @@ mod tests { 0, ) }); - assert!(store.get(ctx.term(Term::Attributes).get(), &eid).is_none()); + assert!(store.get(ctx.term(Term::Attributes).get(), &iid).is_none()); unsafe { drop_index(0, index_ptr); diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 900acf6ab3..879b5f0831 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -66,6 +66,10 @@ const QUANT_STATE_KEY: u32 = u32::from_be_bytes(*b"_qnt"); /// Starting capacity of the pre-allocated rerank buffers. const RERANK_BUFFER_LENGTH: usize = 1024; +/// Size hint passed to Garnet when batch reading attributes. Attributes are variable +/// length, so this is only an estimate used to size Garnet's read buffer. +const ATTRIBUTE_LENGTH_HINT: usize = 1024; + #[derive(Clone)] struct AdjList(AdjacencyList); @@ -424,9 +428,18 @@ impl GarnetProvider { id: &GarnetId, data: &[u8], ) -> Result<(), GarnetProviderError> { + let mut iid = u32::MAX; + if !self.callbacks.read_single_eid( + &context.term(Term::IntMap), + id, + bytemuck::bytes_of_mut(&mut iid), + ) { + return Err(GarnetError::Read.into()); + } + if self .callbacks - .write_eid(&context.term(Term::Attributes), id, data) + .write_iid(&context.term(Term::Attributes), iid, data) { Ok(()) } else { @@ -439,9 +452,18 @@ impl GarnetProvider { context: &Context, id: &GarnetId, ) -> Result<(), GarnetProviderError> { + let mut iid = u32::MAX; + if !self.callbacks.read_single_eid( + &context.term(Term::IntMap), + id, + bytemuck::bytes_of_mut(&mut iid), + ) { + return Err(GarnetError::Read.into()); + } + if self .callbacks - .delete_eid(&context.term(Term::Attributes), id) + .delete_iid(&context.term(Term::Attributes), iid) { Ok(()) } else { @@ -1066,7 +1088,7 @@ impl Delete for GarnetProvider { // It is not an error to fail deleting attributes; they may not exist. let _: bool = self .callbacks - .delete_eid(&context.term(Term::Attributes), gid); + .delete_iid(&context.term(Term::Attributes), id); // TODO: inplace_delete needs access to neighbors. Delete these once that bug is fixed. // See https://github.com/microsoft/DiskANN/issues/1153. @@ -1203,6 +1225,32 @@ impl<'a, T: VectorRepr> DynamicAccessor<'a, T> { } } } + + /// Batch read the attributes for `filtered_ids` and record the filter result for each + /// into `filtered_decisions`. + /// + /// Garnet skips ids with no stored attributes, so those keep `default_decision`. + fn compute_filter_decisions(&mut self, default_decision: bool) { + let Self { + provider, + context, + filtered_ids, + filtered_decisions, + .. + } = self; + + filtered_decisions.clear(); + filtered_decisions.resize(filtered_ids.len() / 2, default_decision); + + provider.callbacks.read_multi_lpiid::<_, u8>( + &context.term(Term::Attributes), + filtered_ids, + ATTRIBUTE_LENGTH_HINT, + |i, attrs| { + filtered_decisions[i as usize] = provider.callbacks.matches_filter(context, attrs); + }, + ); + } } impl HasId for DynamicAccessor<'_, T> { @@ -1516,12 +1564,13 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { // borrow. We put it back at the end to save the allocation. let mut id_buffer = mem::take(&mut **self.id_buffer); + let default_decision = self.provider.callbacks.matches_filter(self.context, &[]); + for nl_id in ids { self.provider .get_neighbors(self.context, nl_id, &mut id_buffer); self.filtered_ids.clear(); - self.filtered_decisions.clear(); for id in id_buffer.iter().copied().filter(|id| pred.eval_mut(id)) { if id == Self::START_ID { @@ -1531,15 +1580,15 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { }; on_neighbors(Decision::reject(id), dist); } else { - let matches = self.provider.callbacks.matches_filter(self.context, id); - self.filtered_ids.push(4); self.filtered_ids.push(id); - - self.filtered_decisions.push(matches); } } + if self.filtered_ids.is_empty() { + continue; + } + let (ctx, length_hint) = if self.quantized { ( self.context.term(Term::Quantized), @@ -1552,22 +1601,23 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { ) }; - if !self.filtered_ids.is_empty() { - self.provider.callbacks.read_multi_lpiid( - &ctx, - &self.filtered_ids, - length_hint, - |i, v| { - let dist = self.computer.evaluate_similarity(v); - let decision = if self.filtered_decisions[i as usize] { - Decision::accept(self.filtered_ids[i as usize * 2 + 1]) - } else { - Decision::reject(self.filtered_ids[i as usize * 2 + 1]) - }; - on_neighbors(decision, dist); - }, - ); - } + self.compute_filter_decisions(default_decision); + + // Read vectors and calculate distances + self.provider.callbacks.read_multi_lpiid( + &ctx, + &self.filtered_ids, + length_hint, + |i, v| { + let dist = self.computer.evaluate_similarity(v); + let decision = if self.filtered_decisions[i as usize] { + Decision::accept(self.filtered_ids[i as usize * 2 + 1]) + } else { + Decision::reject(self.filtered_ids[i as usize * 2 + 1]) + }; + on_neighbors(decision, dist); + }, + ); } **self.id_buffer = id_buffer; @@ -1589,6 +1639,8 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { // borrow. We put it back at the end to save the allocation. let mut id_buffer = mem::take(&mut **self.id_buffer); + let default_decision = self.provider.callbacks.matches_filter(self.context, &[]); + for nl_id in ids { self.provider .get_neighbors(self.context, nl_id, &mut id_buffer); @@ -1596,13 +1648,36 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { for id in id_buffer.iter().copied() { if id != Self::START_ID && pred.eval(&id) { - let matches = self.provider.callbacks.matches_filter(self.context, id); + self.filtered_ids.push(4); + self.filtered_ids.push(id); + } + } - if matches && pred.eval_mut(&Accept::new(id)) { - self.filtered_ids.push(4); - self.filtered_ids.push(id); - } + if self.filtered_ids.is_empty() { + continue; + } + + self.compute_filter_decisions(default_decision); + + // Remove non-matching ids + let mut index = 0; + for (i, &matches) in self.filtered_decisions.iter().enumerate() { + if !matches { + continue; } + + let id = self.filtered_ids[i * 2 + 1]; + + if pred.eval_mut(&Accept::new(id)) { + self.filtered_ids[index * 2] = 4; + self.filtered_ids[index * 2 + 1] = id; + index += 1; + } + } + self.filtered_ids.truncate(index * 2); + + if self.filtered_ids.is_empty() { + continue; } let (ctx, length_hint) = if self.quantized { @@ -1617,17 +1692,15 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { ) }; - if !self.filtered_ids.is_empty() { - self.provider.callbacks.read_multi_lpiid( - &ctx, - &self.filtered_ids, - length_hint, - |i, v| { - let dist = self.computer.evaluate_similarity(v); - on_neighbors(Accept::new(self.filtered_ids[i as usize * 2 + 1]), dist); - }, - ); - } + self.provider.callbacks.read_multi_lpiid( + &ctx, + &self.filtered_ids, + length_hint, + |i, v| { + let dist = self.computer.evaluate_similarity(v); + on_neighbors(Accept::new(self.filtered_ids[i as usize * 2 + 1]), dist); + }, + ); } **self.id_buffer = id_buffer; diff --git a/diskann-garnet/src/test_utils.rs b/diskann-garnet/src/test_utils.rs index 1aaa6e30c4..827d637549 100644 --- a/diskann-garnet/src/test_utils.rs +++ b/diskann-garnet/src/test_utils.rs @@ -199,7 +199,7 @@ unsafe extern "C" fn test_rmw( true } -unsafe extern "C" fn test_filter(_context: u64, _internal_id: u32) -> bool { +unsafe extern "C" fn test_filter(_context: u64, _data: *const u8, _len: usize) -> bool { true } From 3cfb5de96e981b73ef636ef251a4db1e49b7a5b8 Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Thu, 20 Aug 2026 10:09:06 -0500 Subject: [PATCH 7/8] [diskann-garnet] Add beam_width param to search_vector/element --- diskann-garnet/src/ffi_recall_tests.rs | 1 + diskann-garnet/src/ffi_tests.rs | 3 +++ diskann-garnet/src/lib.rs | 12 ++++++++++-- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/diskann-garnet/src/ffi_recall_tests.rs b/diskann-garnet/src/ffi_recall_tests.rs index 43d5ef75fe..b855d2bf14 100644 --- a/diskann-garnet/src/ffi_recall_tests.rs +++ b/diskann-garnet/src/ffi_recall_tests.rs @@ -220,6 +220,7 @@ mod tests { output_id_buffer.len(), output_dists.as_mut_ptr(), output_dists.len(), + 1, continuation, ) }; diff --git a/diskann-garnet/src/ffi_tests.rs b/diskann-garnet/src/ffi_tests.rs index 8e2356d5ad..4d37dde00d 100644 --- a/diskann-garnet/src/ffi_tests.rs +++ b/diskann-garnet/src/ffi_tests.rs @@ -481,6 +481,7 @@ mod tests { output_id_buffer.len(), output_dists.as_mut_ptr(), output_dists.len(), + 1, ptr::null_mut(), ) }; @@ -588,6 +589,7 @@ mod tests { output_id_buffer.len(), output_dists.as_mut_ptr(), output_dists.len(), + 1, ptr::null_mut(), ) }; @@ -708,6 +710,7 @@ mod tests { output_id_buffer.len(), output_dists.as_mut_ptr(), output_dists.len(), + 1, ptr::null_mut(), ) }; diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 8a3385cc35..df88c1d531 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -632,6 +632,7 @@ pub unsafe extern "C" fn search_vector( output_ids_len: usize, output_distances: *mut f32, output_distances_len: usize, + beam_width: u32, _continuation: *mut c_void, ) -> i32 { let index = unsafe { &*index_ptr.cast::() }; @@ -651,7 +652,10 @@ pub unsafe extern "C" fn search_vector( output_distances_len, ); - let knn_params = match search::Knn::new(search_exploration_factor as usize, None) { + let knn_params = match search::Knn::new( + search_exploration_factor as usize, + Some(beam_width as usize), + ) { Ok(params) => params, Err(_) => return -1, }; @@ -702,6 +706,7 @@ pub unsafe extern "C" fn search_element( output_ids_len: usize, output_distances: *mut f32, output_distances_len: usize, + beam_width: u32, _continuation: *mut c_void, ) -> i32 { let index = unsafe { &*index_ptr.cast::() }; @@ -716,7 +721,10 @@ pub unsafe extern "C" fn search_element( output_distances_len, ); - let knn_params = match search::Knn::new(search_exploration_factor as usize, None) { + let knn_params = match search::Knn::new( + search_exploration_factor as usize, + Some(beam_width as usize), + ) { Ok(knn) => knn, Err(_) => return -1, }; From ded76cc26151c85ee6c10d6d1dba0e62da7f06fc Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Thu, 20 Aug 2026 10:35:35 -0500 Subject: [PATCH 8/8] [diskann-garnet] Update documentation and comments --- diskann-garnet/docs/ffi-design.rs | 315 +++++++++++++++++++++++------- diskann-garnet/src/dyn_index.rs | 24 +++ diskann-garnet/src/lib.rs | 97 ++++++++- 3 files changed, 365 insertions(+), 71 deletions(-) diff --git a/diskann-garnet/docs/ffi-design.rs b/diskann-garnet/docs/ffi-design.rs index b1a1c4acb9..13c8dba05c 100644 --- a/diskann-garnet/docs/ffi-design.rs +++ b/diskann-garnet/docs/ffi-design.rs @@ -3,6 +3,108 @@ * Licensed under the MIT license. */ +//! Reference listing of the FFI surface exported to Garnet. +//! +//! This file is not compiled. It mirrors the `#[unsafe(no_mangle)] extern "C"` symbols in +//! `src/lib.rs` and the callback types in `src/garnet.rs`, and must be updated alongside them. + +/// Element type of the vector data passed across the FFI. Must match the definition on the +/// C# side. +#[repr(C)] +enum VectorValueType { + Invalid = 0, + FP32, + XB8, +} + +/// Quantizer selection for an index. Must match the definition on the C# side. +/// +/// `NoQuant`, `Bin`, and `Q8` map to the quantizations Redis exposes and take `f32` vector +/// data. The `X`-prefixed variants are DiskANN extensions taking `u8`/`i8` vector data. +#[repr(C)] +enum VectorQuantType { + Invalid = 0, + NoQuant, + Bin, + Q8, + XNoQuantU8, + XNoQuantI8, + XBinI8, + XBinU8, +} + +/// Status returned by `insert`, encoded as a `u8`. +/// +/// `SuccessStartTraining` signals that the insert crossed the threshold at which the quantizer +/// can be trained, and that Garnet should call `build_quant_table`. +enum InsertResult { + Fail = 0, + Success = 1, + SuccessStartTraining = 2, +} + +/// Read one or more keys from Garnet. +/// +/// `keys` holds `key_count` keys, each 4-byte length prefixed. `value_length_hint` is the +/// expected size in bytes of a single value, and may be an overestimate. +/// +/// For every key that is present, Garnet invokes `read_data` with the index of that key within +/// the batch, the opaque `read_data_state`, and the value bytes. Missing keys are skipped, so +/// a key's existence is determined by whether `read_data` fires for it. Values must be aligned +/// to at least 8 bytes. +type ReadCallback = unsafe extern "C" fn( + context: u64, + key_count: u32, + value_length_hint: u32, + keys: *const u8, + keys_len: usize, + read_data: ReadDataCallback, + read_data_state: *mut c_void, +); + +/// Delivers a single value to the caller of `ReadCallback`. +type ReadDataCallback = + unsafe extern "C" fn(index: u32, state: *mut c_void, value: *const u8, value_len: usize); + +/// Write a value for a key. Returns true on success. +type WriteCallback = unsafe extern "C" fn( + context: u64, + key: *const u8, + key_len: usize, + value: *const u8, + value_len: usize, +) -> bool; + +/// Delete a key. Returns true on success. +type DeleteCallback = unsafe extern "C" fn(context: u64, key: *const u8, key_len: usize) -> bool; + +/// Atomically read, modify, and write the value for a key. +/// +/// Garnet invokes `modify` with the opaque `modify_state` and a mutable view of the current +/// value. If the key does not exist, a zero-initialized value of `write_len` bytes is created +/// and passed instead. Returns true on success. +type ReadModifyWriteCallback = unsafe extern "C" fn( + context: u64, + key: *const u8, + key_len: usize, + write_len: usize, + modify: RmwDataCallback, + modify_state: *mut c_void, +) -> bool; + +/// Mutates the value in place on behalf of `ReadModifyWriteCallback`. +type RmwDataCallback = unsafe extern "C" fn(state: *mut c_void, value: *mut u8, value_len: usize); + +/// Evaluate the filter of the in-flight search against a vector's attribute data. +/// +/// `attributes` is null when the vector has no attributes. Returns true if the vector passes. +type FilterCallback = + unsafe extern "C" fn(context: u64, attributes: *const u8, attributes_len: usize) -> bool; + +/// Emit a UTF-8 log message. The `Term` bits of the context indicate which area of the index +/// the message concerns. +type LogCallback = unsafe extern "C" fn(context: u64, message: *const u8, message_len: usize); + /// Create a new empty index /// Takes the params of VADD (see: https://redis.io/docs/latest/commands/vadd/), maps to a reasonable interpretation /// @@ -10,163 +112,193 @@ /// /// Expectation is any state necessary to recover an index is stored via read/write callbacks - including quantizers. /// -/// reduce_dims == 0 to indicate no reduction requested (and can be ignored even if provided, if that is reasonable). -/// -/// quant_type needs option that map from NoQuant, Bin, and Q8 (as that's what provided in Redis) in addition to any custom index. They don't need to be exact, just reasonable. +/// reduce_dim == 0 to indicate no reduction requested. Dimensionality reduction is not +/// implemented, so this parameter is currently ignored. /// /// metric_type is passed as a raw i32. Valid values are: /// - 0: Cosine /// - 1: InnerProduct /// - 2: L2 (Euclidean distance) /// - 3: CosineNormalized -/// Invalid values will cause the function to return null. /// -/// Returning a single pointer conceal all the generics behind an opaque handle +/// Returns an opaque handle that conceals all the generics, or null on error. The +/// handle must be released with `drop_index`. +/// +/// Sets the `quantization_needed` out-param if the index requires Garnet to drive the quantizer +/// lifecycle via `build_quant_table` and `backfill_quant_vectors`. This can be false even when a +/// quantizer is in use, since not every quantizer requires training and backfill. #[unsafe(no_mangle)] extern "C" fn create_index( context: u64, - dimensions: u32, - reduce_dims: u32, - quant_type: SomeCStyleEnumeration, + dim: u32, + reduce_dim: u32, + quant_type: VectorQuantType, metric_type: i32, - build_exploration_factor: u32, - num_links: u32, - read_callback: unsafe extern "C" fn(u64, *const u8, usize, *mut u8, usize) -> i32, - write_callback: unsafe extern "C" fn(u64, *const u8, usize, *const u8, usize) -> bool, - delete_callback: unsafe extern "C" fn(u64, *const u8, usize) -> bool, -) -> *mut c_void; + l_build: u32, + max_degree: u32, + read_callback: ReadCallback, + write_callback: WriteCallback, + delete_callback: DeleteCallback, + rmw_callback: ReadModifyWriteCallback, + filter_callback: FilterCallback, + log_callback: LogCallback, + quantization_needed: *mut bool, +) -> *const c_void; /// Drop a previously created index /// +/// This is the only valid way to release a handle returned by `create_index`. +/// /// Not called if any other operation against the index may be in flight or started. #[unsafe(no_mangle)] -extern "C" fn drop_index( - context: u64, - index: *const c_void -); +extern "C" fn drop_index(context: u64, index_ptr: *const c_void); /// Insert a vector into an index. /// -/// Returns true if the vector is added, false if it is not. +/// Returns an `InsertResult` discriminant. `Fail` may result from the vector already being in +/// the index, or from writes failing. /// -/// False may result from the vector already being in the index, or writes failing. +/// vector_len is a count of elements, not bytes; the element type follows from the index's +/// `quant_type`. The pointer need not be aligned for that element type. /// /// Note that insert has to be aware of quantizer weirdness, if buffering has to happen it happens here. If we transition from not-quantizing to quantizing, it also has to happen here. /// -/// For now, attribute_data/attribute_len can be ignored - just want space for them. +/// Attributes are optional; pass a null pointer or a zero length to insert without them. #[unsafe(no_mangle)] extern "C" fn insert( context: u64, - index: *const c_void, + index_ptr: *const c_void, id_data: *const u8, id_len: usize, - vector_data: *const f32, + vector_data: *const u8, vector_len: usize, attribute_data: *const u8, - attribute_len: usize + attribute_len: usize, +) -> u8; + +/// Train the quantizer. +/// +/// Garnet calls this once per `insert` that returned `InsertResult::SuccessStartTraining`. +/// Because inserts are concurrent it may be invoked more than once, and the implementation +/// ensures the tables are only built once. +/// +/// Returns true once the tables are built, after which Garnet issues `backfill_quant_vectors` +/// calls from a thread pool. Returns false on failure, in which case it may be retried. +#[unsafe(no_mangle)] +extern "C" fn build_quant_table(context: u64, index_ptr: *const c_void) -> bool; + +/// Quantize vectors that were inserted before the quantizer was trained. +/// +/// Once `build_quant_table` succeeds, Garnet invokes this an arbitrary number of times from a +/// thread pool. Each invocation receives its own `task_index` and the total `task_count` so +/// that it can select and size its share of the work. +/// +/// Returns true on success and false otherwise. +#[unsafe(no_mangle)] +extern "C" fn backfill_quant_vectors( + context: u64, + index_ptr: *const c_void, + task_index: usize, + task_count: usize, ) -> bool; /// Update attribute data on a vector already in the index. /// /// To implement VSETATTR (https://redis.io/docs/latest/commands/vsetattr/). /// -/// We can skip implementing this for now since we don't need filters yet, just needs to be spec'd. +/// An empty attribute value deletes the attributes. /// /// Return true if vector was in index and attribute was updated (even if attribute did not change), false otherwise. #[unsafe(no_mangle)] extern "C" fn set_attribute( context: u64, - index: *const c_void, + index_ptr: *const c_void, id_data: *const u8, id_len: usize, attribute_data: *const u8, - attribute_len: usize + attribute_len: usize, ) -> bool; /// Find similar vectors, takes parameters of VSIM (https://redis.io/docs/latest/commands/vsim/) and maps to a reasonable interpretation. /// /// Works with vector values. /// -/// vector_data is unquantized, vector_len will always match dimensions from create_index. +/// vector_data is unquantized, vector_len will always match dimensions from create_index. As +/// with `insert`, it is a count of elements whose type follows from the index's `quant_type`. /// -/// delta will be [0, 1]. +/// delta is not implemented and is currently ignored. /// /// Maximum number of results is indicated by output_distances_len, elements are i32 length prefixed in byte blobs in output_ids. /// /// distances are [0, 1]. /// -/// Returns number of results, sets continuation to non-zero if there are more to fetch. +/// search_exploration_factor is the search list size and must be non-zero. beam_width is the +/// number of nodes explored per hop and must also be non-zero. /// -/// Filtering can be ignored for now, just reserving space in FFI. +/// Passing a non-null `bitmap_data` with a non-zero `bitmap_len` selects filtered search, with +/// `max_filtering_effort` bounding the extra work spent satisfying the filter. Attribute data +/// is evaluated through the `filter_callback` supplied to `create_index`. /// -/// Various search & effort values can be ignored for now, but will eventually be mapped to something sensible. Exist for compat with Redis. +/// Returns the number of results, or -1 on error. Continuations are not implemented, so the +/// `continuation` parameter is currently ignored. #[unsafe(no_mangle)] extern "C" fn search_vector( context: u64, index_ptr: *const c_void, - vector_data: *const f32, + vector_data: *const u8, vector_len: usize, - delta: float, - search_exploration_factor: i32, - filter_data: *const u8, - filter_len: usize, + delta: f32, + search_exploration_factor: u32, + bitmap_data: *const u8, + bitmap_len: usize, max_filtering_effort: usize, output_ids: *mut u8, output_ids_len: usize, output_distances: *mut f32, output_distances_len: usize, + beam_width: u32, continuation: *mut c_void, ) -> i32; - /// Find similar vectors, takes parameters of VSIM (https://redis.io/docs/latest/commands/vsim/) and maps to a reasonable interpretation. /// -/// Works with item id -/// -/// delta will be [0, 1]. -/// -/// Maximum number of results is indicated by output_distances_len, elements are i32 length prefixed in byte blobs in output_ids. -/// -/// distances are [0, 1]. -/// -/// Returns number of results. -/// -/// Filtering can be ignored for now, just reserving space in FFI. -/// -/// Various search & effort values can be ignored for now, but will eventually be mapped to something sensible. Exist for compat with Redis. +/// Works with item id. Parameters and return value are otherwise as documented on +/// `search_vector`. #[unsafe(no_mangle)] extern "C" fn search_element( context: u64, index_ptr: *const c_void, id_data: *const u8, - id_length: usize, - delta: float, - search_exploration_factor: i32, - filter_data: *const u8, - filter_len: usize, - max_filtering_effort: i32, + id_len: usize, + delta: f32, + search_exploration_factor: u32, + bitmap_data: *const u8, + bitmap_len: usize, + max_filtering_effort: usize, output_ids: *mut u8, output_ids_len: usize, output_distances: *mut f32, output_distances_len: usize, - continuation: *mut c_void + beam_width: u32, + continuation: *mut c_void, ) -> i32; /// Continues fetching results if not all were available after a call to search_xxx /// -/// Returns number of results placed in output_xxx +/// Returns the number of results placed in output_xxx, or -1 on error, and sets +/// new_continuation to non-zero if even more results are available. /// -/// Sets new_continuation to non-zero if even more results are available. +/// NOTE: This is not implemented and always returns -1. #[unsafe(no_mangle)] extern "C" fn continue_search( context: u64, index_ptr: *const c_void, - continuation: usize, + continuation: *mut c_void, output_ids: *mut u8, output_ids_len: usize, output_distances: *mut f32, output_distances_len: usize, - new_continuation: *mut c_void + new_continuation: *mut c_void, ) -> i32; /// Remove vector from index. @@ -175,22 +307,32 @@ extern "C" fn continue_search( /// /// Returns true if element was removed from index. #[unsafe(no_mangle)] -extern "C" fn delete( +extern "C" fn remove( context: u64, index_ptr: *const c_void, - vector_data: *const u8, - vector_len: usize + id_data: *const u8, + id_len: usize, ) -> bool; /// Return number of vectors stored in index. /// /// Equivalent to VCARD (https://redis.io/docs/latest/commands/vcard/) can be approximate, must be fast. #[unsafe(no_mangle)] -extern "C" fn card( +extern "C" fn card(context: u64, index_ptr: *const c_void) -> u64; + +/// Check whether an internal ID refers to a live vector. +/// +/// `internal_id_data` must be exactly 4 bytes holding a native-endian u32; any other length +/// returns false. +/// +/// Returns true if the vector exists in the index, false otherwise. +#[unsafe(no_mangle)] +extern "C" fn check_internal_id_valid( context: u64, index_ptr: *const c_void, -) -> u64; - + internal_id_data: *const u8, + internal_id_len: usize, +) -> bool; /// Check if a vector exists in the index. /// @@ -198,11 +340,44 @@ extern "C" fn card( /// /// Returns true if the vector exists in the index, false otherwise. #[unsafe(no_mangle)] -extern "C" fn has_vector( +extern "C" fn check_external_id_valid( context: u64, index_ptr: *const c_void, id_data: *const u8, id_len: usize, ) -> bool; -// To inspect neighbor lists and vector data, Garnet just has to be aware of the format - not a big deal, no need for FFI. \ No newline at end of file +/// Return up to `count` random members of the index. +/// +/// For implementing VRANDMEMBER (https://redis.io/docs/latest/commands/vrandmember/). No +/// distances are produced; ids are written to output_ids as i32 length prefixed byte blobs. +/// +/// Returns true on success and false otherwise. +#[unsafe(no_mangle)] +extern "C" fn random_members( + context: u64, + index_ptr: *const c_void, + count: u32, + output_ids: *mut u8, + output_ids_len: usize, +) -> bool; + +/// Return the neighbor list of a vector, with the distance to each neighbor. +/// +/// For implementing VLINKS (https://redis.io/docs/latest/commands/vlinks/). Output buffers are +/// filled as they are for search_xxx. +/// +/// Returns the number of neighbors written, or -1 on error. Continuations are not implemented, +/// so the `continuation` parameter is currently ignored. +#[unsafe(no_mangle)] +extern "C" fn search_neighbors( + context: u64, + index_ptr: *const c_void, + id_data: *const u8, + id_len: usize, + output_ids: *mut u8, + output_ids_len: usize, + output_distances: *mut f32, + output_distances_len: usize, + continuation: *mut c_void, +) -> i32; diff --git a/diskann-garnet/src/dyn_index.rs b/diskann-garnet/src/dyn_index.rs index 3d4da83c32..ecfd780642 100644 --- a/diskann-garnet/src/dyn_index.rs +++ b/diskann-garnet/src/dyn_index.rs @@ -20,12 +20,16 @@ use diskann_providers::index::wrapped_async::DiskANNIndex; /// Type-erased version of `DiskANNIndex`. /// All vector data is passed as untyped byte slices. pub(crate) trait DynIndex: Send + Sync { + /// Inserts a vector with id into the index fn insert(&self, context: &Context, id: &GarnetId, data: &[u8]) -> ANNResult<()>; + /// Sets the attributes for a vector fn set_attributes(&self, context: &Context, id: &GarnetId, data: &[u8]) -> ANNResult<()>; + /// Deletes the attributes for a vector fn delete_attributes(&self, context: &Context, id: &GarnetId) -> ANNResult<()>; + /// Searches for the nearest neighbors of a vector fn search_vector( &self, context: &Context, @@ -34,6 +38,7 @@ pub(crate) trait DynIndex: Send + Sync { output: &mut SearchResults<'_>, ) -> ANNResult; + /// Searches for the nearest neighbors of an existing vector in the index fn search_element( &self, context: &Context, @@ -42,6 +47,7 @@ pub(crate) trait DynIndex: Send + Sync { output: &mut SearchResults<'_>, ) -> ANNResult; + /// Filtered search for a vector fn filtered_search_vector( &self, context: &Context, @@ -50,6 +56,7 @@ pub(crate) trait DynIndex: Send + Sync { output: &mut SearchResults<'_>, ) -> ANNResult; + /// Filtered search for an existing vector in the index fn filtered_search_element( &self, context: &Context, @@ -58,24 +65,41 @@ pub(crate) trait DynIndex: Send + Sync { output: &mut SearchResults<'_>, ) -> ANNResult; + /// Delete a vector from the index fn remove(&self, context: &Context, id: &GarnetId) -> ANNResult<()>; + /// Return an approximate count of vectors in the index fn approximate_count(&self) -> u64; + /// Set a start point if one doesn't already exist. + /// If there is already a start point, this is a no-op. fn maybe_set_start_point(&self, context: &Context, data: &[u8]) -> ANNResult<()>; + /// Check if a vector exists by its internal id. + /// Returns true if the vector exists and false otherwise. fn internal_id_exists(&self, context: &Context, id: u32) -> bool; + /// Check if a vector exists by its external id. + /// Returns true if the vector exists false otherwise. fn external_id_exists(&self, context: &Context, id: &GarnetId) -> bool; + /// Train the quantizer. + /// Returns true if training was successful and false otherwise. fn train_quantizer(&self, context: &Context) -> bool; + /// Quantize a group of previously inserted vectors. + /// This function will be called `task_count` times with `task_idx` as a zero-based + /// identifier of the group. This will attempt to quantize `total_vectors / task_count` + /// vectors and returns true if it was successful and false otherwise. fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize) -> bool; + /// Return `count` random vectors from the index. + /// Returns true on success and false otherwise. fn random_members(&self, context: &Context, count: u32, output: &mut SearchResults<'_>) -> bool; + /// Returns the neighbors of and distances from the target vector fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult>>; } diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index df88c1d531..a49f59d824 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -55,10 +55,14 @@ mod test_utils; const ADAPTIVE_L_SAMPLES: usize = 1000; +/// State of index readiness #[derive(Debug, PartialEq)] enum IndexState { + /// No starting points are present in the graph NoStartPoints, + /// Some thread is currently in the process of setting start points SettingStartPoints, + /// Start points set; index ready for normal operation Ready, } impl From for IndexState { @@ -73,12 +77,19 @@ impl From for IndexState { } } +/// Index wrapper type. +/// An `&Arc` is what will be given out over the FFI. pub(crate) struct Index { + /// The type-erased index inner: Box, + /// The quantizer type of the index quant_type: VectorQuantType, + /// A marker for index readiness; uses `IndexState` as the value state: AtomicUsize, } +/// Element type of vectors in the index +/// NOTE: This must match the definition on the C# side. #[derive(Copy, Clone, Debug, PartialEq, Eq)] #[repr(C)] pub enum VectorValueType { @@ -87,6 +98,8 @@ pub enum VectorValueType { XB8, } +/// Quantizer type of the index +/// NOTE: This must match the definition on the C# side. #[derive(Copy, Clone, Debug, PartialEq, Eq)] #[repr(C)] pub enum VectorQuantType { @@ -100,6 +113,10 @@ pub enum VectorQuantType { XBinU8, } +/// Helper struct to manage the FFI buffers for handling search results +/// +/// NOTE: The ids will be 4-byte length prefixed, and external IDs are arbitrary length +/// byte strings. struct SearchResults<'a> { ids: &'a mut [u8], dists: &'a mut [f32], @@ -108,6 +125,7 @@ struct SearchResults<'a> { } impl SearchResults<'_> { + /// Construct from the raw pointers fn new(ids: *mut u8, ids_len: usize, dists: *mut f32, dists_len: usize) -> Self { let ids = unsafe { slice::from_raw_parts_mut(ids, ids_len) }; let dists = unsafe { slice::from_raw_parts_mut(dists, dists_len) }; @@ -121,6 +139,8 @@ impl SearchResults<'_> { } } + /// Push an ID only into the results. + /// This is primarily used by `random_members` which does not use distances. fn push_id(&mut self, id: GarnetId) -> diskann::graph::BufferState { self.push(Neighbor::new(id, 0.0)) } @@ -177,6 +197,8 @@ impl SearchOutputBuffer for SearchResults<'_> { } } +/// Helper generic function to create the correct type-erased `Arc`. +/// This also returns a bool indicating whether quantization is needed. fn create_index_impl( quant_type: VectorQuantType, config: config::Config, @@ -219,6 +241,18 @@ fn create_index_impl( )) } +/// Create an index. +/// +/// Constructs a type-erased DiskANN index object as a `Arc` and return a pointer +/// to the leaked Arc. This pointer must be freed with `drop_index()`. +/// +/// Returns `ptr::null()` if there is an error. Sets the `quantization_needed` outvar if +/// the index requires quantization callbacks during its lifecycle. +/// +/// Note that `quantization_needed` can be set to false even when a quantizer is used. The +/// flag controls whether supplemental control is needed from Garnet to manage quantizers +/// which require training and backfill. +/// /// # Safety /// /// FFI @@ -324,6 +358,10 @@ pub unsafe extern "C" fn create_index( } } +/// Drop an index. +/// +/// This is the only valid way to free an index pointer created with `create_index()`. +/// /// # Safety /// /// FFI @@ -333,6 +371,7 @@ pub unsafe extern "C" fn drop_index(_ctx: u64, index_ptr: *const c_void) { let _ = unsafe { Arc::from_raw(index_ptr.cast::()) }; } +/// `Cow` type for `Poly<[u8], AlignToEight>` types. enum PolyCow<'a> { Owned(Poly<[u8], AlignToEight>), Borrowed(&'a [u8]), @@ -361,6 +400,11 @@ impl<'a> From> for PolyCow<'a> { } } +/// Helper function to interpret the vector pointer and size into a usable Rust +/// type. This will return either a borrowed or owned vector depending on how +/// the pointer is aligned. Since Garnet doesn't guarantee the alignment, if it +/// is not 4-byte aligned, we must allocate an appropriately aligned buffer to +/// access it as its correct element type. fn interpret_vector<'a>( quant_type: VectorQuantType, vector_data: &'a *const u8, @@ -405,6 +449,11 @@ fn interpret_vector<'a>( Some(v) } +/// Return type for `insert()`. +/// +/// `Fail` and `Success` are obvious. `SuccessStartTraining` is used when enough vectors have +/// been inserted to start training the quantizer. That return value signals to Garnet that +/// `build_quant_table` should be called. #[derive(Debug, Clone, Copy, PartialEq)] enum InsertResult { Fail, @@ -497,6 +546,7 @@ pub unsafe extern "C" fn insert( } } +/// Ensures the index is ready to be used, and if not, runs the `init` function. fn ensure_index_ready_or_init(index: &Index, init: F) -> Option where F: FnOnce() -> Option, @@ -543,7 +593,7 @@ where /// Once this function returns `true`, Garnet will invoke several `backfill_quant_vectors()` /// calls from a thread pool. If it returns false, it may be re-invoked to try again. /// -/// # Safety +/// # Safety /// /// FFI #[unsafe(no_mangle)] @@ -558,6 +608,8 @@ pub unsafe extern "C" fn build_quant_table(context: u64, index_ptr: *const c_voi /// times from a thread pool. Each invocation is told its index and the total number of /// invocations so that each invocation can correctly pick and size its work. /// +/// Returns true for success and false otherwise. +/// /// # Safety /// /// FFI @@ -575,6 +627,12 @@ pub unsafe extern "C" fn backfill_quant_vectors( .backfill_quant_vectors(&ctx, task_index, task_count) } +/// Set the attributes for a vector. +/// +/// Setting attributes with `attribute_len == 0` is equivalent to deleting them. +/// +/// Returns true for success and false otherwise. +/// /// # Safety /// /// FFI @@ -614,6 +672,8 @@ pub unsafe extern "C" fn set_attribute( true } +/// Search the closest vectors to the given query vector. +/// /// # Safety /// /// FFI @@ -688,6 +748,8 @@ pub unsafe extern "C" fn search_vector( } } +/// Search the closest vectors to the given existing vector in the index. +/// /// # Safety /// /// FFI @@ -760,6 +822,12 @@ pub unsafe extern "C" fn search_element( } } +/// Continue getting results for a previously executed search. +/// +/// NOTE: This is currently unimplemented. +/// +/// Positive return values are the count of vectors returned. `-1` will be returned on errors. +/// /// # Safety /// /// FFI @@ -777,6 +845,10 @@ pub unsafe extern "C" fn continue_search( -1 } +/// Remove a vector from the index. +/// +/// Returns true on success and false otherwise. +/// /// # Safety /// /// FFI @@ -799,6 +871,8 @@ pub unsafe extern "C" fn remove( index.inner.remove(&ctx, &id).is_ok() } +/// Return the approximate count of vectors in the index. +/// /// # Safety /// /// FFI @@ -809,6 +883,10 @@ pub unsafe extern "C" fn card(_ctx: u64, index_ptr: *const c_void) -> u64 { index.inner.approximate_count() } +/// Check if a given internal ID is a valid vector. +/// +/// Returns true if the vector exists, and false otherwise. +/// /// # Safety /// /// FFI @@ -832,6 +910,10 @@ pub unsafe extern "C" fn check_internal_id_valid( index.inner.internal_id_exists(&ctx, id) } +/// Check if a given external ID is a valid vector. +/// +/// Returns true if the vector exists, and false otherwise. +/// /// # Safety /// /// FFI @@ -850,6 +932,12 @@ pub unsafe extern "C" fn check_external_id_valid( index.inner.external_id_exists(&ctx, &id) } +/// Returns random vectors from the index. +/// +/// This is primarily a debugging aid. +/// +/// Returns true on success and false otherwise. +/// /// # Safety /// /// FFI @@ -876,6 +964,13 @@ pub unsafe extern "C" fn random_members( index.inner.random_members(&ctx, count, &mut output) } +/// Return the neighbors for an index vector. +/// +/// This is primarily a debugging aid. It returns both the neighbors' IDs and their +/// distance from the given vector. +/// +/// Returns the number of results or `-1` on error. +/// /// # Safety /// /// FFI