From a56aecd56e91b622842970e9391cf974b747caf8 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Thu, 21 May 2026 10:51:59 -0700 Subject: [PATCH 1/3] Refactor `flat_search`. --- .../src/search/provider/disk_provider.rs | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 1dd518781..b837b1baf 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -920,25 +920,35 @@ where where OB: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + Send, { + // We process IDs up to `BATCH_SIZE` elements at a time. + const BATCH_SIZE: usize = 32; + let mut id_buffer = Vec::with_capacity(BATCH_SIZE); + let provider = self.index.provider(); let mut accessor = strategy .search_accessor(provider, &DefaultContext) .into_ann_result()?; - let computer = accessor.build_query_computer(query).into_ann_result()?; let mut best = NeighborPriorityQueue::new(neighbors_before_reranking); let mut cmps = 0u32; - let num_points = provider.num_points as u32; - for id in 0..num_points { - if vector_filter(&id) { - let element = accessor.get_element(id).await.into_ann_result()?; - let dist = computer.evaluate_similarity(element); - best.insert(Neighbor::new(id, dist)); - cmps += 1; + let mut iter = (0..provider.num_points as u32).filter(vector_filter); + + loop { + id_buffer.clear(); + id_buffer.extend(iter.by_ref().take(BATCH_SIZE)); + + if id_buffer.is_empty() { + break; } + + accessor.pq_distances(&id_buffer, |dist, id| best.insert(Neighbor::new(id, dist)))?; + cmps += id_buffer.len() as u32; } + // FIXME: This is a temporary bridge. We don't really need the query computer, but + // we do need to satisfy the trait definition until PR 1067 lands. + let computer = accessor.build_query_computer(query).into_ann_result()?; let result_count = strategy .default_post_processor() .post_process(&mut accessor, query, &computer, best.iter(), output) From 0d57050296da5f2d6b8d30cd97f4359e55a026d1 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Thu, 21 May 2026 11:02:44 -0700 Subject: [PATCH 2/3] Use a dynamic batch size to avoid panics. --- diskann-disk/src/search/pq/pq_scratch.rs | 8 ++++++++ diskann-disk/src/search/provider/disk_provider.rs | 12 ++++++------ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/diskann-disk/src/search/pq/pq_scratch.rs b/diskann-disk/src/search/pq/pq_scratch.rs index e90af5212..1df666e18 100644 --- a/diskann-disk/src/search/pq/pq_scratch.rs +++ b/diskann-disk/src/search/pq/pq_scratch.rs @@ -78,6 +78,12 @@ impl PQScratch { self.query_scratch.copy_from_slice(&query[..dim]); Ok(()) } + + /// Return the largest number of PQ vectors whose distances can be computed using this + /// scratch data structure. + pub(crate) fn max_vectors(&self) -> usize { + self.aligned_dist_scratch.len() + } } #[cfg(test)] @@ -112,6 +118,8 @@ mod tests { 0 ); + assert_eq!(pq_scratch.max_vectors(), graph_degree); + // Test set() method let query: Vec = (1..=dim).map(|i| i as f32).collect(); pq_scratch.set(&query).unwrap(); diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index b837b1baf..d95252374 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -920,23 +920,23 @@ where where OB: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + Send, { - // We process IDs up to `BATCH_SIZE` elements at a time. - const BATCH_SIZE: usize = 32; - let mut id_buffer = Vec::with_capacity(BATCH_SIZE); - let provider = self.index.provider(); let mut accessor = strategy .search_accessor(provider, &DefaultContext) .into_ann_result()?; + // Derive the batch size from the scratch data structure. Providing too many vectors + // will panic. + let batch_size = accessor.scratch.pq_scratch.max_vectors(); + let mut id_buffer = Vec::with_capacity(batch_size); + let mut best = NeighborPriorityQueue::new(neighbors_before_reranking); let mut cmps = 0u32; let mut iter = (0..provider.num_points as u32).filter(vector_filter); - loop { id_buffer.clear(); - id_buffer.extend(iter.by_ref().take(BATCH_SIZE)); + id_buffer.extend(iter.by_ref().take(batch_size)); if id_buffer.is_empty() { break; From 19a4a3ca1547b0822ab076654bd2a2084e5e7605 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Thu, 21 May 2026 11:12:58 -0700 Subject: [PATCH 3/3] Protect against infinite loops. --- diskann-disk/src/search/provider/disk_provider.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index d95252374..da54a58db 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -928,6 +928,17 @@ where // Derive the batch size from the scratch data structure. Providing too many vectors // will panic. let batch_size = accessor.scratch.pq_scratch.max_vectors(); + + // This check should always hold since `graph_degree` comes from + // `diskann::graph::Config` and is forced to be non-zero. But this is defensive + // against misconfiguration. + if batch_size == 0 { + return Err(ANNError::message( + diskann::ANNErrorKind::IndexError, + "pq scratch must support at least one vector", + )); + } + let mut id_buffer = Vec::with_capacity(batch_size); let mut best = NeighborPriorityQueue::new(neighbors_before_reranking);