From 68d393241bed2b48640b7d2d1b3564cd16ac6604 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Wed, 12 Aug 2026 15:31:45 -0400 Subject: [PATCH 1/7] remove fields from scratch --- diskann/src/graph/search/diverse_search.rs | 2 -- .../src/graph/search/filtered_range_search.rs | 19 ++++++------ diskann/src/graph/search/range_search.rs | 29 ++++++++++--------- diskann/src/graph/search/scratch.rs | 20 ++----------- 4 files changed, 28 insertions(+), 42 deletions(-) diff --git a/diskann/src/graph/search/diverse_search.rs b/diskann/src/graph/search/diverse_search.rs index 44ffbb251a..5bbd54cd80 100644 --- a/diskann/src/graph/search/diverse_search.rs +++ b/diskann/src/graph/search/diverse_search.rs @@ -169,8 +169,6 @@ where ), id_scratch: Vec::with_capacity(index.max_degree_with_slack()), beam_nodes: Vec::with_capacity(self.inner.beam_width().get()), - range_frontier: std::collections::VecDeque::new(), - in_range: Vec::new(), hops: 0, cmps: 0, } diff --git a/diskann/src/graph/search/filtered_range_search.rs b/diskann/src/graph/search/filtered_range_search.rs index 9f1e80be6e..5f96808f97 100644 --- a/diskann/src/graph/search/filtered_range_search.rs +++ b/diskann/src/graph/search/filtered_range_search.rs @@ -183,15 +183,14 @@ where && matched_within_radius.len() < max_returned { // clear the visited set and repopulate it with all in-range points found so far, filtered and unfiltered - // also add these points to range_frontier for seeding the second-round search scratch.visited.clear(); - scratch.range_frontier.clear(); scratch .visited .extend(in_range.iter().map(|neighbor| *neighbor.id())); - scratch - .range_frontier - .extend(in_range.iter().map(|neighbor| *neighbor.id())); + + // Create a range frontier for seeding the second-round search + let range_frontier: std::collections::VecDeque<_> = + in_range.iter().map(|neighbor| *neighbor.id()).collect(); // Move to filtered range search let range_stats = filtered_range_search_internal( @@ -200,6 +199,7 @@ where &mut accessor, &mut scratch, &mut matched_within_radius, + range_frontier, ) .await?; @@ -262,6 +262,7 @@ pub(crate) async fn filtered_range_search_internal( accessor: &mut A, scratch: &mut SearchScratch, matched_in_range: &mut Vec>, + mut range_frontier: std::collections::VecDeque, ) -> ANNResult where A: FilteredAccessor, @@ -272,13 +273,13 @@ where let max_returned = search_params.max_returned().unwrap_or(usize::MAX); - while !scratch.range_frontier.is_empty() && matched_in_range.len() < max_returned { + while !range_frontier.is_empty() && matched_in_range.len() < max_returned { scratch.beam_nodes.clear(); // In this loop we are going to find the beam_width number of remaining nodes within the radius // Each of these nodes will be a frontier node. - while !scratch.range_frontier.is_empty() && scratch.beam_nodes.len() < beam_width { - let next = scratch.range_frontier.pop_front(); + while !range_frontier.is_empty() && scratch.beam_nodes.len() < beam_width { + let next = range_frontier.pop_front(); if let Some(next_node) = next { scratch.beam_nodes.push(next_node); } @@ -300,7 +301,7 @@ where for (decision, distance) in neighbors.iter().copied() { if distance <= navigation_radius { let id = decision.into_inner(); - scratch.range_frontier.push_back(id); + range_frontier.push_back(id); if distance <= search_params.radius() && decision.is_accept() diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 8c143e068a..e5008d6b87 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -5,6 +5,7 @@ //! Range-based search within a distance radius. +use std::collections::VecDeque; use std::num::NonZeroUsize; use diskann_utils::future::SendFuture; @@ -298,11 +299,9 @@ where for neighbor in in_range.iter() { scratch.visited.insert(*neighbor.id()); } - scratch.in_range = in_range; - let stats = if scratch.in_range.len() - >= ((starting_l as f32) * self.initial_slack()) as usize - && scratch.in_range.len() < max_returned + let stats = if in_range.len() >= ((starting_l as f32) * self.initial_slack()) as usize + && in_range.len() < max_returned { // Move to range search let range_stats = range_search_internal( @@ -310,6 +309,7 @@ where &self, &mut accessor, &mut scratch, + &mut in_range, ) .await?; @@ -338,7 +338,7 @@ where .post_process( &mut accessor, query, - scratch.in_range.iter().copied(), + in_range.iter().copied(), &mut filtered, ) .await @@ -413,27 +413,30 @@ pub(crate) async fn range_search_internal( search_params: &Range, accessor: &mut A, scratch: &mut SearchScratch, + in_range: &mut Vec>, ) -> ANNResult where A: SearchAccessor, { let beam_width = search_params.beam_width().get(); - for neighbor in &scratch.in_range { - scratch.range_frontier.push_back(*neighbor.id()); + let mut range_frontier: VecDeque = VecDeque::new(); + + for neighbor in in_range.iter() { + range_frontier.push_back(*neighbor.id()); } let mut neighbors = Vec::with_capacity(max_degree_with_slack); let max_returned = search_params.max_returned().unwrap_or(usize::MAX); - while !scratch.range_frontier.is_empty() && scratch.in_range.len() < max_returned { + while !range_frontier.is_empty() && in_range.len() < max_returned { scratch.beam_nodes.clear(); // In this loop we are going to find the beam_width number of remaining nodes within the radius // Each of these nodes will be a frontier node. - while !scratch.range_frontier.is_empty() && scratch.beam_nodes.len() < beam_width { - let next = scratch.range_frontier.pop_front(); + while !range_frontier.is_empty() && scratch.beam_nodes.len() < beam_width { + let next = range_frontier.pop_front(); if let Some(next_node) = next { scratch.beam_nodes.push(next_node); } @@ -451,10 +454,10 @@ where // The predicate ensures that the contents of `neighbors` are unique. for neighbor in neighbors.iter() { if *neighbor.distance() <= search_params.radius() * search_params.range_slack() - && scratch.in_range.len() < max_returned + && in_range.len() < max_returned { - scratch.in_range.push(*neighbor); - scratch.range_frontier.push_back(*neighbor.id()); + in_range.push(*neighbor); + range_frontier.push_back(*neighbor.id()); } } scratch.cmps += neighbors.len() as u32; diff --git a/diskann/src/graph/search/scratch.rs b/diskann/src/graph/search/scratch.rs index 2670f118c1..3f2627f7df 100644 --- a/diskann/src/graph/search/scratch.rs +++ b/diskann/src/graph/search/scratch.rs @@ -7,12 +7,7 @@ //! Scratch space for in-memory index based search -use std::collections::VecDeque; - -use crate::{ - neighbor::{Neighbor, NeighborPriorityQueue}, - utils::VectorId, -}; +use crate::{neighbor::NeighborPriorityQueue, utils::VectorId}; use diskann_utils::object_pool::AsPooled; use hashbrown::HashSet; @@ -60,14 +55,6 @@ where /// to temporarily hold beam of nodes in each hop. pub beam_nodes: Vec, - /// A queue of nodes to visit during range search - /// Does not need to be ordered by distance - pub range_frontier: VecDeque, - - /// A list of nodes that are in range of the query - /// Only used during range search - pub in_range: Vec>, - /// A tracker for how many hops we have taken during the current search pub hops: u32, @@ -126,8 +113,6 @@ where visited, id_scratch: Vec::new(), beam_nodes: Vec::new(), - in_range: Vec::new(), - range_frontier: VecDeque::new(), hops: 0, cmps: 0, } @@ -150,8 +135,6 @@ where self.visited.clear(); self.id_scratch.clear(); self.beam_nodes.clear(); - self.in_range.clear(); - self.range_frontier.clear(); self.hops = 0; self.cmps = 0; @@ -231,6 +214,7 @@ pub(crate) struct SearchScratchParams { #[cfg(test)] mod tests { use super::*; + use crate::neighbor::Neighbor; #[test] pub fn test_new() { From 9187f0f7c1ed340dffe868eaa9f7cbbd569a0053 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 14 Aug 2026 19:19:25 +0000 Subject: [PATCH 2/7] clean up range search a lot yay --- .../src/graph/search/filtered_range_search.rs | 142 ++++---- diskann/src/graph/search/range_search.rs | 339 +++++++++--------- .../graph/test/cases/filtered_range_search.rs | 195 +++++++--- diskann/src/graph/test/cases/range_search.rs | 216 +++++++++-- ...sults_respected_means_no_second_round.json | 35 -- ..._respected_and_second_round_triggered.json | 10 +- ...sults_respected_means_no_second_round.json | 35 -- .../cases/range_search/two_round_search.json | 4 +- 8 files changed, 606 insertions(+), 370 deletions(-) delete mode 100644 diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round.json delete mode 100644 diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json diff --git a/diskann/src/graph/search/filtered_range_search.rs b/diskann/src/graph/search/filtered_range_search.rs index 5f96808f97..ae9c1db4ca 100644 --- a/diskann/src/graph/search/filtered_range_search.rs +++ b/diskann/src/graph/search/filtered_range_search.rs @@ -17,12 +17,13 @@ use crate::{ index::{DiskANNIndex, InternalSearchStats, SearchStats}, search::inline_filter_search::{Ret, inline_filter_search_internal}, search::{ - Range, RangeSearchError, Search, range_search::DistanceFiltered, - range_search::RangeBuilder, record::NoopSearchRecord, scratch::SearchScratch, + Range, RangeSearchError, Search, + range_search::{InRange, InRangePushResult, RangeBuilder}, + record::NoopSearchRecord, + scratch::SearchScratch, }, search_output_buffer::SearchOutputBuffer, }, - neighbor::Neighbor, provider::DataProvider, }; @@ -55,6 +56,12 @@ impl FilteredRange { self.range_params.max_returned() } + /// Returns the maximum number of results to return. + #[inline] + pub fn effective_max_returned(&self, inc: usize) -> usize { + self.range_params.effective_max_returned(inc) + } + /// Returns the initial search list size. #[inline] pub fn starting_l(&self) -> NonZeroUsize { @@ -135,12 +142,12 @@ where .search_accessor(&index.data_provider, context, query) .into_ann_result()?; let num_start_ids = accessor.num_starting_points().await?; - let mut scratch = index.search_scratch(self.starting_l().get(), num_start_ids); - - // Perform an initial inline filtered search, store both filtered and unfiltered results + let starting_l = self.starting_l().get(); + let mut scratch = index.search_scratch(starting_l, num_start_ids); let search_knn = self.range_params.to_knn(); + // Perform an initial inline filtered search, store both filtered and unfiltered results let Ret { cmps, hops, @@ -155,42 +162,53 @@ where ) .await?; - let max_returned = self.max_returned().unwrap_or(usize::MAX); - - // merge matched_results with the best results from the first round, filtering by radius - - let mut in_range: Vec<_> = scratch - .best - .iter() - .take(self.starting_l().get()) - .chain(matched_results.iter().copied()) - .filter(|neighbor| *neighbor.distance() <= self.radius()) - .collect(); + let max_returned = self.effective_max_returned(num_start_ids); - in_range.sort_unstable_by(crate::neighbor::ord::fast_distance_total); - in_range.dedup_by(|left, right| left.id() == right.id()); + let mut in_outer_range = InRange::new_checked( + self.radius(), + None, + usize::MAX, + scratch.best.iter().take(starting_l).collect(), + ); // filter matched results by radius; this will be used to decide if `max_results` has been reached - let mut matched_within_radius = Vec::with_capacity(matched_results.len()); - for neighbor in matched_results.iter().copied() { - if *neighbor.distance() <= self.radius() { - matched_within_radius.push(neighbor); - } - } + let mut matched_in_outer_range = InRange::new_checked( + self.radius(), + self.inner_radius(), + max_returned, + matched_results.to_vec(), + ); - let stats = if in_range.len() - >= ((self.starting_l().get() as f32) * self.initial_slack()) as usize - && matched_within_radius.len() < max_returned + // merge matched_results with the best results from the first round, filtering by radius + for neighbor in matched_in_outer_range.neighbors.iter() { + in_outer_range.push_unchecked(*neighbor); + } + in_outer_range + .neighbors + .sort_unstable_by(crate::neighbor::ord::fast_distance_total); + in_outer_range + .neighbors + .dedup_by(|left, right| left.id() == right.id()); + + let stats = if in_outer_range.len() + >= ((starting_l as f32) * self.initial_slack()) as usize + && matched_in_outer_range.len() < max_returned { // clear the visited set and repopulate it with all in-range points found so far, filtered and unfiltered scratch.visited.clear(); - scratch - .visited - .extend(in_range.iter().map(|neighbor| *neighbor.id())); + scratch.visited.extend( + in_outer_range + .neighbors + .iter() + .map(|neighbor| *neighbor.id()), + ); // Create a range frontier for seeding the second-round search - let range_frontier: std::collections::VecDeque<_> = - in_range.iter().map(|neighbor| *neighbor.id()).collect(); + let mut range_frontier: std::collections::VecDeque<_> = in_outer_range + .neighbors + .iter() + .map(|neighbor| *neighbor.id()) + .collect(); // Move to filtered range search let range_stats = filtered_range_search_internal( @@ -198,8 +216,8 @@ where &self, &mut accessor, &mut scratch, - &mut matched_within_radius, - range_frontier, + &mut matched_in_outer_range, + &mut range_frontier, ) .await?; @@ -216,24 +234,14 @@ where } }; - // Post-process results directly into the output buffer, filtering by radius. - let inner_radius = self.inner_radius(); - - // Note matched_in_range is assumed to satisfy the radius filter, so we only apply the inner radius filter here - let mut filtered = DistanceFiltered::new(output, |dist| { - if let Some(ir) = inner_radius - && dist <= ir - { - false - } else { - true - } - }); - - let truncated_matched = matched_within_radius.iter().copied().take(max_returned); + let truncated_matched = matched_in_outer_range + .neighbors + .iter() + .copied() + .take(max_returned); let result_count = processor - .post_process(&mut accessor, query, truncated_matched, &mut filtered) + .post_process(&mut accessor, query, truncated_matched, output) .await .into_ann_result()?; @@ -261,8 +269,8 @@ pub(crate) async fn filtered_range_search_internal( search_params: &FilteredRange, accessor: &mut A, scratch: &mut SearchScratch, - matched_in_range: &mut Vec>, - mut range_frontier: std::collections::VecDeque, + matched_in_range: &mut InRange, + range_frontier: &mut std::collections::VecDeque, ) -> ANNResult where A: FilteredAccessor, @@ -273,7 +281,7 @@ where let max_returned = search_params.max_returned().unwrap_or(usize::MAX); - while !range_frontier.is_empty() && matched_in_range.len() < max_returned { + 'outer: while !range_frontier.is_empty() && matched_in_range.len() < max_returned { scratch.beam_nodes.clear(); // In this loop we are going to find the beam_width number of remaining nodes within the radius @@ -299,18 +307,26 @@ where // but only accepted IDs are added to in-range results. let navigation_radius = search_params.radius() * search_params.range_slack(); for (decision, distance) in neighbors.iter().copied() { - if distance <= navigation_radius { - let id = decision.into_inner(); - range_frontier.push_back(id); - - if distance <= search_params.radius() - && decision.is_accept() - && matched_in_range.len() < max_returned - { - matched_in_range.push(Neighbor::new(id, distance)); + match matched_in_range.push_with_filter(distance, decision) { + InRangePushResult::Full => { + break 'outer; + } + InRangePushResult::Accepted + | InRangePushResult::RejectedbyInner + | InRangePushResult::RejectedbyFilter => { + let id = decision.into_inner(); + range_frontier.push_back(id); + } + InRangePushResult::RejectedbyOuter => { + // add to the range frontier if it's within the tolerance + if distance <= navigation_radius { + let id = decision.into_inner(); + range_frontier.push_back(id); + } } } } + scratch.cmps += neighbors.len() as u32; scratch.hops += scratch.beam_nodes.len() as u32; } diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index e5008d6b87..8f68c0eba1 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -21,7 +21,7 @@ use crate::{ Knn, Search, filtered_range_search::FilteredRange, record::NoopSearchRecord, scratch::SearchScratch, }, - search_output_buffer::{self, SearchOutputBuffer}, + search_output_buffer::SearchOutputBuffer, }, neighbor::Neighbor, provider::DataProvider, @@ -140,6 +140,19 @@ impl Range { self.max_returned } + /// Returns either usize::MAX or the user-specified maximum number + /// results to return incremented by a user-inputted value. + /// Useful to enforce max results exactly when start points + /// are filtered out during post-processing. + #[inline] + pub fn effective_max_returned(&self, inc: usize) -> usize { + if let Some(max) = self.max_returned { + max.saturating_add(inc) + } else { + usize::MAX + } + } + /// Returns the initial search list size. #[inline] pub fn starting_l(&self) -> NonZeroUsize { @@ -183,6 +196,25 @@ impl Range { } /// Builder for [`Range`] search parameters. +/// +/// `max_returned`: If specified, the search will stop and return results once this number of points has been found within both the inner and outer radii. +/// Note that due to adding some extra slack for start points, occasionally slightly more than this number of points may be returned. Since the initial +/// search phase does not respect `max_returned`, this parameter may not be set lower than `starting_l`. +/// +/// `starting_l`: the L_search parameter for the initial search phase. Must be greater than zero. +/// +/// `beam_width`: the beam width for parallel graph exploration. If not specified, defaults to 1. Must be greater than zero if specified. +/// +/// `radius`: the outer radius for the range search. Points within this distance from the query are candidates for inclusion in the results. +/// +/// `inner_radius`: the inner radius for the range search. Points closer than this distance from the query are excluded from the results. Must be less than or equal to `radius` if specified. +/// +/// `initial_slack`: after the initial knn search phase, a decision is made on whether to continue to the second round of search. This decision is based on whether +/// the number of points found within the outer radius is greater than `starting_l * initial_slack`, so lower values of `initial_slack` will make it more likely to continue to +/// the second round of search. Must be between 0.0 and 1.0. +/// +/// `range_slack`: during the second round of search, points that are within `radius * range_slack` are expanded to search for candidates within the range, so greater +/// values of `range_slack` will mean more expansions. Must be greater than or equal to 1.0. #[derive(Debug, Clone, Copy)] pub struct RangeBuilder { max_returned: Option, @@ -272,7 +304,9 @@ where .search_accessor(&index.data_provider, context, query) .into_ann_result()?; let num_start_ids = accessor.num_starting_points().await?; - let mut scratch = index.search_scratch(self.starting_l().get(), num_start_ids); + + let starting_l = self.starting_l().get(); + let mut scratch = index.search_scratch(starting_l, num_start_ids); let initial_stats = index .search_internal( @@ -283,63 +317,68 @@ where ) .await?; - let mut in_range = Vec::with_capacity(self.starting_l().get()); - - let starting_l = self.starting_l().get(); - let max_returned = self.max_returned().unwrap_or(usize::MAX); - - for neighbor in scratch.best.iter().take(starting_l) { - if *neighbor.distance() <= self.radius() { - in_range.push(neighbor); - } - } + let in_outer_range = InRange::new_checked( + self.radius(), + None, + starting_l, + scratch.best.iter().collect(), + ); + + // Increment the max results by the number of starting points, in case + // they are filtered out later and leave us with fewer than the requested + // number of results. + let max_returned = self.effective_max_returned(num_start_ids); + + let outer_range_len = in_outer_range.len(); + let outer_range_ids: Vec<_> = in_outer_range + .neighbors + .iter() + .map(|neighbor| *neighbor.id()) + .collect(); + + let mut in_range = InRange::new_checked( + self.radius(), + self.inner_radius(), + max_returned, + in_outer_range.neighbors, + ); + + let stats = if outer_range_len >= ((starting_l as f32) * self.initial_slack()) as usize + && outer_range_len <= max_returned + { + // clear the visited set and repopulate it with just the in-range points + scratch.visited.clear(); + scratch.visited.extend(outer_range_ids.iter().copied()); - // clear the visited set and repopulate it with just the in-range points - scratch.visited.clear(); - for neighbor in in_range.iter() { - scratch.visited.insert(*neighbor.id()); - } + // Create a range frontier for seeding the second-round search + let mut range_frontier: VecDeque<_> = outer_range_ids.into_iter().collect(); - let stats = if in_range.len() >= ((starting_l as f32) * self.initial_slack()) as usize - && in_range.len() < max_returned - { // Move to range search let range_stats = range_search_internal( index.max_degree_with_slack(), &self, &mut accessor, &mut scratch, + &mut range_frontier, &mut in_range, ) .await?; InternalSearchStats { - cmps: initial_stats.cmps, - hops: initial_stats.hops + range_stats.hops, + cmps: range_stats.cmps, + hops: range_stats.hops, range_search_second_round: true, } } else { initial_stats }; - // Post-process results directly into the output buffer, filtering by radius. - let radius = self.radius(); - let inner_radius = self.inner_radius(); - let mut filtered = DistanceFiltered::new(output, |dist| { - if let Some(ir) = inner_radius - && dist <= ir - { - return false; - } - dist <= radius - }); - let result_count = processor .post_process( &mut accessor, query, - in_range.iter().copied(), - &mut filtered, + in_range.neighbors.iter().copied(), + output, ) .await .into_ann_result()?; @@ -354,49 +393,101 @@ where } } -/// A [`SearchOutputBuffer`] wrapper that filters results by distance before -/// forwarding them to an inner buffer. -pub(super) struct DistanceFiltered<'a, F, B: ?Sized> { - predicate: F, - inner: &'a mut B, +pub(super) enum InRangePushResult { + Accepted, + RejectedbyOuter, + RejectedbyInner, // this state is only returned when the outer radius is respected + RejectedbyFilter, // this state is only returned when outer AND inner radius are respected + Full, // this state is returned *before* checking inner/outer radius } -impl<'a, F, B: ?Sized> DistanceFiltered<'a, F, B> { - pub(super) fn new(inner: &'a mut B, predicate: F) -> Self { - Self { predicate, inner } - } +pub(super) struct InRange { + pub neighbors: Vec>, + radius: f32, + inner_radius: Option, + max_returned: usize, } -impl SearchOutputBuffer for DistanceFiltered<'_, F, B> -where - F: FnMut(f32) -> bool, - B: SearchOutputBuffer + ?Sized, -{ - fn size_hint(&self) -> Option { - self.inner.size_hint() +impl InRange { + pub(super) fn push(&mut self, neighbor: Neighbor) -> InRangePushResult { + if self.neighbors.len() >= self.max_returned { + return InRangePushResult::Full; + } + let dist = *neighbor.distance(); + if dist > self.radius { + return InRangePushResult::RejectedbyOuter; + } + if let Some(inner) = self.inner_radius + && dist <= inner + { + return InRangePushResult::RejectedbyInner; + } + self.neighbors.push(neighbor); + InRangePushResult::Accepted } - fn push(&mut self, neighbor: Neighbor) -> search_output_buffer::BufferState { - if (self.predicate)(*neighbor.distance()) { - self.inner.push(neighbor) - } else { - match self.inner.size_hint() { - Some(0) => search_output_buffer::BufferState::Full, - _ => search_output_buffer::BufferState::Available, - } + pub(super) fn push_with_filter( + &mut self, + distance: f32, + decision: glue::Decision, + ) -> InRangePushResult { + if self.neighbors.len() >= self.max_returned { + return InRangePushResult::Full; + } + if distance > self.radius { + return InRangePushResult::RejectedbyOuter; + } + if let Some(inner) = self.inner_radius + && distance <= inner + { + return InRangePushResult::RejectedbyInner; + } + if decision.is_reject() { + return InRangePushResult::RejectedbyFilter; } + self.neighbors + .push(Neighbor::new(decision.into_inner(), distance)); + InRangePushResult::Accepted } - fn current_len(&self) -> usize { - self.inner.current_len() + pub(super) fn push_unchecked(&mut self, neighbor: Neighbor) { + self.neighbors.push(neighbor); } - fn extend(&mut self, itr: Itr) -> usize - where - Itr: IntoIterator>, - { - self.inner - .extend(itr.into_iter().filter(|n| (self.predicate)(*n.distance()))) + pub(super) fn len(&self) -> usize { + self.neighbors.len() + } + + /// Create a new InRange with the given parameters, + /// filtering the candidate neighbors and truncating + /// as needed to respect the maximum number of results. + pub(super) fn new_checked( + radius: f32, + inner_radius: Option, + max_returned: usize, + candidate_neighbors: Vec>, + ) -> Self { + Self { + neighbors: candidate_neighbors + .into_iter() + .filter(|n| { + let dist = *n.distance(); + if dist > radius { + return false; + } + if let Some(inner) = inner_radius + && dist <= inner + { + return false; + } + true + }) + .take(max_returned) + .collect(), + radius, + inner_radius, + max_returned, + } } } @@ -413,24 +504,17 @@ pub(crate) async fn range_search_internal( search_params: &Range, accessor: &mut A, scratch: &mut SearchScratch, - in_range: &mut Vec>, + range_frontier: &mut VecDeque, + in_range: &mut InRange, ) -> ANNResult where A: SearchAccessor, { let beam_width = search_params.beam_width().get(); - let mut range_frontier: VecDeque = VecDeque::new(); - - for neighbor in in_range.iter() { - range_frontier.push_back(*neighbor.id()); - } - let mut neighbors = Vec::with_capacity(max_degree_with_slack); - let max_returned = search_params.max_returned().unwrap_or(usize::MAX); - - while !range_frontier.is_empty() && in_range.len() < max_returned { + 'outer: while !range_frontier.is_empty() { scratch.beam_nodes.clear(); // In this loop we are going to find the beam_width number of remaining nodes within the radius @@ -451,15 +535,26 @@ where ) .await?; - // The predicate ensures that the contents of `neighbors` are unique. + let navigation_radius = search_params.radius() * search_params.range_slack(); for neighbor in neighbors.iter() { - if *neighbor.distance() <= search_params.radius() * search_params.range_slack() - && in_range.len() < max_returned - { - in_range.push(*neighbor); - range_frontier.push_back(*neighbor.id()); + match in_range.push(*neighbor) { + InRangePushResult::Full => { + break 'outer; // we've reached the maximum number of results, stop processing + } + InRangePushResult::Accepted | InRangePushResult::RejectedbyInner => { + // both these results guarantee within outer radius + range_frontier.push_back(*neighbor.id()); + } + InRangePushResult::RejectedbyOuter + if *neighbor.distance() <= navigation_radius => + { + // add to the range frontier if it's within the tolerance + range_frontier.push_back(*neighbor.id()); + } + _ => {} } } + scratch.cmps += neighbors.len() as u32; scratch.hops += scratch.beam_nodes.len() as u32; } @@ -478,8 +573,6 @@ where #[cfg(test)] mod tests { use super::*; - use crate::graph::search_output_buffer::BufferState; - use crate::neighbor::Neighbor; #[test] fn range_builder_defaults_match_new() { @@ -551,78 +644,4 @@ mod tests { .is_err() ); } - - #[test] - fn distance_filtered_push_accepts_passing_items() { - let mut inner: Vec> = Vec::new(); - let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); - - assert_eq!(filtered.push(Neighbor::new(1, 0.5)), BufferState::Available); - assert_eq!(filtered.current_len(), 1); - assert_eq!(*inner[0].id(), 1); - assert_eq!(*inner[0].distance(), 0.5); - } - - #[test] - fn distance_filtered_push_rejects_failing_items() { - let mut inner: Vec> = Vec::new(); - let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); - - assert_eq!(filtered.push(Neighbor::new(1, 1.5)), BufferState::Available); - assert_eq!(filtered.current_len(), 0); - } - - #[test] - fn distance_filtered_extend_filters_correctly() { - let mut inner: Vec> = Vec::new(); - let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); - assert!(filtered.size_hint().is_none()); - - let items = [(1u32, 0.3), (2, 1.5), (3, 0.7), (4, 2.0), (5, 0.9)].map(Neighbor::from_tuple); - let count = filtered.extend(items); - - assert_eq!(count, 3); - assert_eq!(inner.len(), 3); - assert_eq!(*inner[0].id(), 1); - assert_eq!(*inner[1].id(), 3); - assert_eq!(*inner[2].id(), 5); - } - - #[test] - fn distance_filtered_respects_inner_capacity() { - let mut ids = [0u32; 2]; - let mut dists = [0.0f32; 2]; - let mut inner = search_output_buffer::IdDistance::new(&mut ids, &mut dists); - let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); - assert_eq!(filtered.size_hint(), Some(2)); - - let items = [(1u32, 0.1), (2, 0.2), (3, 0.3)].map(Neighbor::from_tuple); - let count = filtered.extend(items); - - assert_eq!(count, 2); - assert_eq!(ids, [1, 2]); - } - - #[test] - fn distance_filtered_inner_radius_pattern() { - let mut inner: Vec> = Vec::new(); - let radius = 1.0f32; - let inner_radius = Some(0.3f32); - let mut filtered = DistanceFiltered::new(&mut inner, |dist| { - if let Some(ir) = inner_radius - && dist <= ir - { - return false; - } - dist < radius - }); - - let items = [(1u32, 0.1), (2, 0.5), (3, 0.3), (4, 1.0), (5, 0.8)].map(Neighbor::from_tuple); - let count = filtered.extend(items); - - // 0.1 and 0.3 are <= inner_radius, 1.0 is not < radius - assert_eq!(count, 2); - assert_eq!(*inner[0].id(), 2); - assert_eq!(*inner[1].id(), 5); - } } diff --git a/diskann/src/graph/test/cases/filtered_range_search.rs b/diskann/src/graph/test/cases/filtered_range_search.rs index 8b9e9a02f5..389dfbeb76 100644 --- a/diskann/src/graph/test/cases/filtered_range_search.rs +++ b/diskann/src/graph/test/cases/filtered_range_search.rs @@ -247,57 +247,6 @@ fn empty_results() { ); } -#[test] -fn max_results_respected_means_no_second_round() { - let description = "Test of `max_results` that sets `initial_l_search` equal \ - to `max_results`, with a permissive radius so that `max_results` is met \ - without needing a second round."; - let mut test_root = root(); - let mut path = test_root.path(); - let name = path.push("max_results_respected_means_no_second_round"); - - let grid_size = 5; - let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); - let radius = 50.0; - let starting_l = 4; - let max_results = 4; - let filter = AlwaysTrueFilter; - - let filtered_range = FilteredRange::builder(starting_l, radius) - .max_returned(Some(max_results)) - .build_filtered() - .unwrap(); - - let (filtered_stats, filtered_results) = - run_filtered_range_search(&index, query.as_slice(), filtered_range, &filter); - - let baseline = RangeSearchBaseline::new( - &filtered_range.range(), - &filtered_results, - filtered_stats, - Grid::Three, - grid_size, - description, - query.clone(), - ); - - let expected = get_or_save_test_results(&name, &baseline); - assert_eq_verbose!(expected, baseline); - - assert!( - filtered_results.len() <= max_results, - "result count {} exceeds max_results {}", - filtered_results.len(), - max_results - ); - assert!( - !filtered_stats.range_search_second_round, - "If max_results is respected, a second round should not be triggered" - ); - assert_range_invariants(&filtered_results, radius, None); - assert_no_duplicates(&filtered_results); -} - #[test] fn max_results_respected_and_second_round_triggered() { let description = "Test of `max_results` that sets `initial_l_search` \ @@ -444,3 +393,147 @@ fn divisible_by_four_filter_no_second_round_from_l_search() { assert_no_duplicates(&filtered_results); assert_divisible_by_four(&filtered_results); } + +///////////////////////////////////////////////////////////// +// Tests for the `initial_slack` and `range_slack` knobs. // +// `initial_slack` gates whether the second round runs at // +// all; `range_slack` widens the navigation radius used to // +// seed the frontier during the second round. // +///////////////////////////////////////////////////////////// + +#[test] +fn initial_slack_low_triggers_second_round() { + let description = "Low `initial_slack` lowers the bar for entering the second \ + round: with starting_l = 4 and slack = 0.5 the threshold is 2, so the second \ + round is triggered."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("initial_slack_low_triggers_second_round"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 50.0; + let starting_l = 4; + let filter = AlwaysTrueFilter; + + let filtered_range = FilteredRange::builder(starting_l, radius) + .initial_slack(0.5) + .build_filtered() + .unwrap(); + + let (filtered_stats, filtered_results) = + run_filtered_range_search(&index, query.as_slice(), filtered_range, &filter); + + let baseline = RangeSearchBaseline::new( + &filtered_range.range(), + &filtered_results, + filtered_stats, + Grid::Three, + grid_size, + description, + query.clone(), + ); + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert!( + filtered_stats.range_search_second_round, + "low initial_slack should trigger a second round" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); +} + +#[test] +fn range_slack_low_constrains_frontier() { + let description = "Low `range_slack` (1.0) keeps the navigation radius equal to \ + the search radius, so only points inside the radius are used to expand the \ + frontier during the second round."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("range_slack_low_constrains_frontier"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 30.0; + let starting_l = 4; + let filter = AlwaysTrueFilter; + + let filtered_range = FilteredRange::builder(starting_l, radius) + .initial_slack(0.5) + .range_slack(1.0) + .build_filtered() + .unwrap(); + + let (filtered_stats, filtered_results) = + run_filtered_range_search(&index, query.as_slice(), filtered_range, &filter); + + let baseline = RangeSearchBaseline::new( + &filtered_range.range(), + &filtered_results, + filtered_stats, + Grid::Three, + grid_size, + description, + query.clone(), + ); + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert!( + filtered_stats.range_search_second_round, + "low initial_slack should trigger a second round" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); +} + +#[test] +fn range_slack_high_expands_frontier() { + let description = "High `range_slack` (1.3) widens the navigation radius so that \ + points beyond the search radius still seed the frontier. Results remain within \ + the search radius, but coverage should be at least as good as with low slack."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("range_slack_high_expands_frontier"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 30.0; + let starting_l = 4; + let filter = AlwaysTrueFilter; + + let filtered_range = FilteredRange::builder(starting_l, radius) + .initial_slack(0.5) + .range_slack(1.3) + .build_filtered() + .unwrap(); + + let (filtered_stats, filtered_results) = + run_filtered_range_search(&index, query.as_slice(), filtered_range, &filter); + + let baseline = RangeSearchBaseline::new( + &filtered_range.range(), + &filtered_results, + filtered_stats, + Grid::Three, + grid_size, + description, + query.clone(), + ); + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert!( + filtered_stats.range_search_second_round, + "low initial_slack should trigger a second round" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); +} diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index 4ee3410a41..111dd92b7d 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -347,20 +347,20 @@ fn empty_results() { } #[test] -fn max_results_respected_means_no_second_round() { - let description = "Two round search test to validate that max_results = \ - starting_l means no second round is triggered."; +fn max_results_respected_and_second_round_triggered() { + let description = "Two round search test to validate that max_results > \ + starting_l means a second round is triggered."; let rt = current_thread_runtime(); let mut test_root = root(); let mut path = test_root.path(); - let name = path.push("max_results_respected_means_no_second_round"); + let name = path.push("max_results_respected_and_second_round_triggered"); let grid_size = 5; let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); let radius = 1.0e9; // every point will be in range with this radius let starting_l = 4; // small set to trigger multiple rounds - let max_results = 4; // max_returned = starting_l, so second round should not be triggered + let max_results = 5; // max_returned greater than starting_l, so second round should be triggered let range_search = Range::builder(starting_l, radius) .max_returned(Some(max_results)) @@ -405,34 +405,96 @@ fn max_results_respected_means_no_second_round() { ); assert!( - !stats.range_search_second_round, - "If max_results is respected, a second round should not be triggered" + stats.range_search_second_round, + "If max_results is respected, a second round should be triggered" ); + assert_range_invariants(&results, radius, None); assert_no_duplicates(&results); } #[test] -fn max_results_respected_and_second_round_triggered() { - let description = "Two round search test to validate that max_results > \ - starting_l means a second round is triggered."; +fn initial_slack_low_triggers_second_round() { + let _description = "Test that low initial_slack triggers second round. \ + With initial_slack=0.5 and starting_l=4, the threshold is 2, so any \ + outer_range_len >= 2 will trigger the second round."; let rt = current_thread_runtime(); let mut test_root = root(); let mut path = test_root.path(); - let name = path.push("max_results_respected_and_second_round_triggered"); + let name = path.push("initial_slack_low_triggers_second_round"); let grid_size = 5; let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); - let radius = 1.0e9; // every point will be in range with this radius - let starting_l = 4; // small set to trigger multiple rounds - let max_results = 5; // max_returned greater than starting_l, so second round should be triggered + let radius = 50.0; + let starting_l = 4; + let low_slack = 0.5; let range_search = Range::builder(starting_l, radius) - .max_returned(Some(max_results)) + .initial_slack(low_slack) .build() .unwrap(); + let mut results: Vec> = Vec::new(); + let stats = rt + .block_on(index.search( + range_search, + &test_provider::Strategy::new(), + &test_provider::Context::new(), + query.as_slice(), + &mut results, + )) + .unwrap(); + + let baseline = RangeSearchBaseline { + description: "Low initial_slack triggers second round search.".to_string(), + grid_dims: Grid::Three.dim(), + grid_size, + query: query.clone(), + radius, + inner_radius: None, + starting_l, + results: results.iter().map(|n| n.as_tuple()).collect(), + comparisons: stats.cmps as usize, + hops: stats.hops as usize, + result_count: results.len(), + range_search_second_round: stats.range_search_second_round, + }; + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert!( + stats.range_search_second_round, + "low initial_slack ({}) should trigger second round", + low_slack + ); + + assert_range_invariants(&results, radius, None); + assert_no_duplicates(&results); +} + +#[test] +fn initial_slack_high_avoids_second_round() { + let _description = "Test that high initial_slack avoids second round. \ + With initial_slack=1.0 and starting_l=4, the threshold is 4, making it \ + harder to trigger the second round."; + + let rt = current_thread_runtime(); + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("initial_slack_high_avoids_second_round"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 50.0; + let starting_l = 4; + let high_slack = 1.0; + + let range_search = Range::builder(starting_l, radius) + .initial_slack(high_slack) + .build() + .unwrap(); let mut results: Vec> = Vec::new(); let stats = rt @@ -446,7 +508,64 @@ fn max_results_respected_and_second_round_triggered() { .unwrap(); let baseline = RangeSearchBaseline { - description: description.to_string(), + description: "High initial_slack avoids second round search.".to_string(), + grid_dims: Grid::Three.dim(), + grid_size, + query: query.clone(), + radius, + inner_radius: None, + starting_l, + results: results.iter().map(|n| n.as_tuple()).collect(), + comparisons: stats.cmps as usize, + hops: stats.hops as usize, + result_count: results.len(), + range_search_second_round: stats.range_search_second_round, + }; + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert_range_invariants(&results, radius, None); + assert_no_duplicates(&results); +} + +#[test] +fn range_slack_low_constrains_frontier() { + let _description = "Test that low range_slack constrains frontier expansion. \ + With range_slack=1.0, the frontier only expands to nodes within radius, \ + resulting in fewer total results found."; + + let rt = current_thread_runtime(); + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("range_slack_low_constrains_frontier"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 30.0; + let starting_l = 4; + let initial_slack = 0.5; + let low_range_slack = 1.0; + + let range_search = Range::builder(starting_l, radius) + .initial_slack(initial_slack) + .range_slack(low_range_slack) + .build() + .unwrap(); + let mut results: Vec> = Vec::new(); + + let stats = rt + .block_on(index.search( + range_search, + &test_provider::Strategy::new(), + &test_provider::Context::new(), + query.as_slice(), + &mut results, + )) + .unwrap(); + + let baseline = RangeSearchBaseline { + description: "Low range_slack constrains frontier expansion.".to_string(), grid_dims: Grid::Three.dim(), grid_size, query: query.clone(), @@ -464,15 +583,70 @@ fn max_results_respected_and_second_round_triggered() { assert_eq_verbose!(expected, baseline); assert!( - results.len() <= max_results, - "result count {} exceeds max_results {}", - results.len(), - max_results + stats.range_search_second_round, + "low initial_slack should trigger second round" ); + assert_range_invariants(&results, radius, None); + assert_no_duplicates(&results); +} + +#[test] +fn range_slack_high_expands_frontier() { + let _description = "Test that high range_slack expands frontier exploration. \ + With range_slack=1.3, the frontier expands to nodes up to radius * 1.3, \ + potentially finding more results than lower range_slack."; + + let rt = current_thread_runtime(); + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("range_slack_high_expands_frontier"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 30.0; + let starting_l = 4; + let initial_slack = 0.5; + let high_range_slack = 1.3; + + let range_search = Range::builder(starting_l, radius) + .initial_slack(initial_slack) + .range_slack(high_range_slack) + .build() + .unwrap(); + let mut results: Vec> = Vec::new(); + + let stats = rt + .block_on(index.search( + range_search, + &test_provider::Strategy::new(), + &test_provider::Context::new(), + query.as_slice(), + &mut results, + )) + .unwrap(); + + let baseline = RangeSearchBaseline { + description: "High range_slack expands frontier exploration.".to_string(), + grid_dims: Grid::Three.dim(), + grid_size, + query: query.clone(), + radius, + inner_radius: None, + starting_l, + results: results.iter().map(|n| n.as_tuple()).collect(), + comparisons: stats.cmps as usize, + hops: stats.hops as usize, + result_count: results.len(), + range_search_second_round: stats.range_search_second_round, + }; + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + assert!( stats.range_search_second_round, - "If max_results is respected, a second round should be triggered" + "low initial_slack should trigger second round" ); assert_range_invariants(&results, radius, None); diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round.json b/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round.json deleted file mode 100644 index dc78b37acc..0000000000 --- a/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round.json +++ /dev/null @@ -1,35 +0,0 @@ -{ - "file": "diskann/src/graph/test/cases/filtered_range_search.rs", - "test": "graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round", - "payload": { - "comparisons": 10, - "description": "Test of `max_results` that sets `initial_l_search` equal to `max_results`, with a permissive radius so that `max_results` is met without needing a second round.", - "grid_dims": 3, - "grid_size": 5, - "hops": 5, - "inner_radius": null, - "query": [ - 5.0, - 5.0, - 5.0 - ], - "radius": 50.0, - "range_search_second_round": false, - "result_count": 3, - "results": [ - [ - 124, - 3.0 - ], - [ - 99, - 6.0 - ], - [ - 119, - 6.0 - ] - ], - "starting_l": 4 - } -} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json index a0214085bc..357e576228 100644 --- a/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json +++ b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json @@ -2,11 +2,11 @@ "file": "diskann/src/graph/test/cases/range_search.rs", "test": "graph/test/cases/range_search/max_results_respected_and_second_round_triggered", "payload": { - "comparisons": 11, + "comparisons": 12, "description": "Two round search test to validate that max_results > starting_l means a second round is triggered.", "grid_dims": 3, "grid_size": 5, - "hops": 12, + "hops": 7, "inner_radius": null, "query": [ 5.0, @@ -15,7 +15,7 @@ ], "radius": 1000000000.0, "range_search_second_round": true, - "result_count": 4, + "result_count": 5, "results": [ [ 124, @@ -32,6 +32,10 @@ [ 99, 6.0 + ], + [ + 98, + 9.0 ] ], "starting_l": 4 diff --git a/diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json deleted file mode 100644 index 4d63c852a6..0000000000 --- a/diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json +++ /dev/null @@ -1,35 +0,0 @@ -{ - "file": "diskann/src/graph/test/cases/range_search.rs", - "test": "graph/test/cases/range_search/max_results_respected_means_no_second_round", - "payload": { - "comparisons": 11, - "description": "Two round search test to validate that max_results = starting_l means no second round is triggered.", - "grid_dims": 3, - "grid_size": 5, - "hops": 5, - "inner_radius": null, - "query": [ - 5.0, - 5.0, - 5.0 - ], - "radius": 1000000000.0, - "range_search_second_round": false, - "result_count": 3, - "results": [ - [ - 124, - 3.0 - ], - [ - 123, - 6.0 - ], - [ - 119, - 6.0 - ] - ], - "starting_l": 4 - } -} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/range_search/two_round_search.json b/diskann/test/generated/graph/test/cases/range_search/two_round_search.json index 7a2bb189c2..f9757a8130 100644 --- a/diskann/test/generated/graph/test/cases/range_search/two_round_search.json +++ b/diskann/test/generated/graph/test/cases/range_search/two_round_search.json @@ -2,11 +2,11 @@ "file": "diskann/src/graph/test/cases/range_search.rs", "test": "graph/test/cases/range_search/two_round_search", "payload": { - "comparisons": 11, + "comparisons": 129, "description": "Two round search test to validate that a low starting L with a large radius triggers a second round of range search.", "grid_dims": 3, "grid_size": 5, - "hops": 120, + "hops": 115, "inner_radius": null, "query": [ 5.0, From cdcbb849e9631913c36fd88ab54f9a6fefbbc94b Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 14 Aug 2026 19:40:14 +0000 Subject: [PATCH 3/7] add missing test jsons --- ...itial_slack_low_triggers_second_round.json | 459 ++++++++++++++++++ .../range_slack_high_expands_frontier.json | 263 ++++++++++ .../range_slack_low_constrains_frontier.json | 263 ++++++++++ ...nitial_slack_high_avoids_second_round.json | 459 ++++++++++++++++++ ...itial_slack_low_triggers_second_round.json | 459 ++++++++++++++++++ .../range_slack_high_expands_frontier.json | 263 ++++++++++ .../range_slack_low_constrains_frontier.json | 263 ++++++++++ 7 files changed, 2429 insertions(+) create mode 100644 diskann/test/generated/graph/test/cases/filtered_range_search/initial_slack_low_triggers_second_round.json create mode 100644 diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_high_expands_frontier.json create mode 100644 diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_low_constrains_frontier.json create mode 100644 diskann/test/generated/graph/test/cases/range_search/initial_slack_high_avoids_second_round.json create mode 100644 diskann/test/generated/graph/test/cases/range_search/initial_slack_low_triggers_second_round.json create mode 100644 diskann/test/generated/graph/test/cases/range_search/range_slack_high_expands_frontier.json create mode 100644 diskann/test/generated/graph/test/cases/range_search/range_slack_low_constrains_frontier.json diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/initial_slack_low_triggers_second_round.json b/diskann/test/generated/graph/test/cases/filtered_range_search/initial_slack_low_triggers_second_round.json new file mode 100644 index 0000000000..ac686bf9fd --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/initial_slack_low_triggers_second_round.json @@ -0,0 +1,459 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/initial_slack_low_triggers_second_round", + "payload": { + "comparisons": 121, + "description": "Low `initial_slack` lowers the bar for entering the second round: with starting_l = 4 and slack = 0.5 the threshold is 2, so the second round is triggered.", + "grid_dims": 3, + "grid_size": 5, + "hops": 115, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 50.0, + "range_search_second_round": true, + "result_count": 109, + "results": [ + [ + 124, + 3.0 + ], + [ + 99, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 123, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 94, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 73, + 14.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 49, + 18.0 + ], + [ + 109, + 18.0 + ], + [ + 121, + 18.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 68, + 17.0 + ], + [ + 84, + 21.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 48, + 21.0 + ], + [ + 72, + 19.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 24, + 27.0 + ], + [ + 104, + 27.0 + ], + [ + 120, + 27.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 43, + 24.0 + ], + [ + 59, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 79, + 30.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 23, + 30.0 + ], + [ + 47, + 26.0 + ], + [ + 71, + 26.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 14, + 35.0 + ], + [ + 18, + 33.0 + ], + [ + 34, + 33.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 54, + 35.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 78, + 33.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ], + [ + 90, + 33.0 + ], + [ + 22, + 35.0 + ], + [ + 46, + 33.0 + ], + [ + 70, + 35.0 + ], + [ + 102, + 35.0 + ], + [ + 106, + 33.0 + ], + [ + 110, + 35.0 + ], + [ + 9, + 42.0 + ], + [ + 13, + 38.0 + ], + [ + 17, + 38.0 + ], + [ + 29, + 42.0 + ], + [ + 33, + 36.0 + ], + [ + 37, + 34.0 + ], + [ + 41, + 36.0 + ], + [ + 53, + 38.0 + ], + [ + 57, + 34.0 + ], + [ + 61, + 34.0 + ], + [ + 65, + 38.0 + ], + [ + 77, + 38.0 + ], + [ + 81, + 36.0 + ], + [ + 85, + 38.0 + ], + [ + 21, + 42.0 + ], + [ + 45, + 42.0 + ], + [ + 101, + 42.0 + ], + [ + 105, + 42.0 + ], + [ + 8, + 45.0 + ], + [ + 12, + 43.0 + ], + [ + 16, + 45.0 + ], + [ + 28, + 45.0 + ], + [ + 32, + 41.0 + ], + [ + 36, + 41.0 + ], + [ + 40, + 45.0 + ], + [ + 52, + 43.0 + ], + [ + 56, + 41.0 + ], + [ + 60, + 43.0 + ], + [ + 76, + 45.0 + ], + [ + 80, + 45.0 + ], + [ + 7, + 50.0 + ], + [ + 11, + 50.0 + ], + [ + 27, + 50.0 + ], + [ + 31, + 48.0 + ], + [ + 35, + 50.0 + ], + [ + 51, + 50.0 + ], + [ + 55, + 50.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_high_expands_frontier.json b/diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_high_expands_frontier.json new file mode 100644 index 0000000000..7f47b13047 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_high_expands_frontier.json @@ -0,0 +1,263 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/range_slack_high_expands_frontier", + "payload": { + "comparisons": 102, + "description": "High `range_slack` (1.3) widens the navigation radius so that points beyond the search radius still seed the frontier. Results remain within the search radius, but coverage should be at least as good as with low slack.", + "grid_dims": 3, + "grid_size": 5, + "hops": 90, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 30.0, + "range_search_second_round": true, + "result_count": 60, + "results": [ + [ + 124, + 3.0 + ], + [ + 99, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 123, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 94, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 73, + 14.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 49, + 18.0 + ], + [ + 109, + 18.0 + ], + [ + 121, + 18.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 68, + 17.0 + ], + [ + 84, + 21.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 48, + 21.0 + ], + [ + 72, + 19.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 24, + 27.0 + ], + [ + 104, + 27.0 + ], + [ + 120, + 27.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 43, + 24.0 + ], + [ + 59, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 79, + 30.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 23, + 30.0 + ], + [ + 47, + 26.0 + ], + [ + 71, + 26.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_low_constrains_frontier.json b/diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_low_constrains_frontier.json new file mode 100644 index 0000000000..cca26485e1 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/range_slack_low_constrains_frontier.json @@ -0,0 +1,263 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/range_slack_low_constrains_frontier", + "payload": { + "comparisons": 84, + "description": "Low `range_slack` (1.0) keeps the navigation radius equal to the search radius, so only points inside the radius are used to expand the frontier during the second round.", + "grid_dims": 3, + "grid_size": 5, + "hops": 66, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 30.0, + "range_search_second_round": true, + "result_count": 60, + "results": [ + [ + 124, + 3.0 + ], + [ + 99, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 123, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 94, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 73, + 14.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 49, + 18.0 + ], + [ + 109, + 18.0 + ], + [ + 121, + 18.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 68, + 17.0 + ], + [ + 84, + 21.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 48, + 21.0 + ], + [ + 72, + 19.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 24, + 27.0 + ], + [ + 104, + 27.0 + ], + [ + 120, + 27.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 43, + 24.0 + ], + [ + 59, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 79, + 30.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 23, + 30.0 + ], + [ + 47, + 26.0 + ], + [ + 71, + 26.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/range_search/initial_slack_high_avoids_second_round.json b/diskann/test/generated/graph/test/cases/range_search/initial_slack_high_avoids_second_round.json new file mode 100644 index 0000000000..7754fe7237 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/range_search/initial_slack_high_avoids_second_round.json @@ -0,0 +1,459 @@ +{ + "file": "diskann/src/graph/test/cases/range_search.rs", + "test": "graph/test/cases/range_search/initial_slack_high_avoids_second_round", + "payload": { + "comparisons": 129, + "description": "High initial_slack avoids second round search.", + "grid_dims": 3, + "grid_size": 5, + "hops": 115, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 50.0, + "range_search_second_round": true, + "result_count": 109, + "results": [ + [ + 124, + 3.0 + ], + [ + 123, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 99, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 94, + 9.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 73, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 121, + 18.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 109, + 18.0 + ], + [ + 49, + 18.0 + ], + [ + 48, + 21.0 + ], + [ + 68, + 17.0 + ], + [ + 72, + 19.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 120, + 27.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 84, + 21.0 + ], + [ + 104, + 27.0 + ], + [ + 24, + 27.0 + ], + [ + 23, + 30.0 + ], + [ + 43, + 24.0 + ], + [ + 47, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 71, + 26.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 59, + 26.0 + ], + [ + 79, + 30.0 + ], + [ + 18, + 33.0 + ], + [ + 22, + 35.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 46, + 33.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 70, + 35.0 + ], + [ + 78, + 33.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ], + [ + 90, + 33.0 + ], + [ + 102, + 35.0 + ], + [ + 106, + 33.0 + ], + [ + 110, + 35.0 + ], + [ + 14, + 35.0 + ], + [ + 34, + 33.0 + ], + [ + 54, + 35.0 + ], + [ + 13, + 38.0 + ], + [ + 17, + 38.0 + ], + [ + 21, + 42.0 + ], + [ + 33, + 36.0 + ], + [ + 37, + 34.0 + ], + [ + 41, + 36.0 + ], + [ + 45, + 42.0 + ], + [ + 53, + 38.0 + ], + [ + 57, + 34.0 + ], + [ + 61, + 34.0 + ], + [ + 65, + 38.0 + ], + [ + 77, + 38.0 + ], + [ + 81, + 36.0 + ], + [ + 85, + 38.0 + ], + [ + 101, + 42.0 + ], + [ + 105, + 42.0 + ], + [ + 9, + 42.0 + ], + [ + 29, + 42.0 + ], + [ + 8, + 45.0 + ], + [ + 12, + 43.0 + ], + [ + 16, + 45.0 + ], + [ + 28, + 45.0 + ], + [ + 32, + 41.0 + ], + [ + 36, + 41.0 + ], + [ + 40, + 45.0 + ], + [ + 52, + 43.0 + ], + [ + 56, + 41.0 + ], + [ + 60, + 43.0 + ], + [ + 76, + 45.0 + ], + [ + 80, + 45.0 + ], + [ + 7, + 50.0 + ], + [ + 11, + 50.0 + ], + [ + 27, + 50.0 + ], + [ + 31, + 48.0 + ], + [ + 35, + 50.0 + ], + [ + 51, + 50.0 + ], + [ + 55, + 50.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/range_search/initial_slack_low_triggers_second_round.json b/diskann/test/generated/graph/test/cases/range_search/initial_slack_low_triggers_second_round.json new file mode 100644 index 0000000000..9b353f1119 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/range_search/initial_slack_low_triggers_second_round.json @@ -0,0 +1,459 @@ +{ + "file": "diskann/src/graph/test/cases/range_search.rs", + "test": "graph/test/cases/range_search/initial_slack_low_triggers_second_round", + "payload": { + "comparisons": 129, + "description": "Low initial_slack triggers second round search.", + "grid_dims": 3, + "grid_size": 5, + "hops": 115, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 50.0, + "range_search_second_round": true, + "result_count": 109, + "results": [ + [ + 124, + 3.0 + ], + [ + 123, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 99, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 94, + 9.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 73, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 121, + 18.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 109, + 18.0 + ], + [ + 49, + 18.0 + ], + [ + 48, + 21.0 + ], + [ + 68, + 17.0 + ], + [ + 72, + 19.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 120, + 27.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 84, + 21.0 + ], + [ + 104, + 27.0 + ], + [ + 24, + 27.0 + ], + [ + 23, + 30.0 + ], + [ + 43, + 24.0 + ], + [ + 47, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 71, + 26.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 59, + 26.0 + ], + [ + 79, + 30.0 + ], + [ + 18, + 33.0 + ], + [ + 22, + 35.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 46, + 33.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 70, + 35.0 + ], + [ + 78, + 33.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ], + [ + 90, + 33.0 + ], + [ + 102, + 35.0 + ], + [ + 106, + 33.0 + ], + [ + 110, + 35.0 + ], + [ + 14, + 35.0 + ], + [ + 34, + 33.0 + ], + [ + 54, + 35.0 + ], + [ + 13, + 38.0 + ], + [ + 17, + 38.0 + ], + [ + 21, + 42.0 + ], + [ + 33, + 36.0 + ], + [ + 37, + 34.0 + ], + [ + 41, + 36.0 + ], + [ + 45, + 42.0 + ], + [ + 53, + 38.0 + ], + [ + 57, + 34.0 + ], + [ + 61, + 34.0 + ], + [ + 65, + 38.0 + ], + [ + 77, + 38.0 + ], + [ + 81, + 36.0 + ], + [ + 85, + 38.0 + ], + [ + 101, + 42.0 + ], + [ + 105, + 42.0 + ], + [ + 9, + 42.0 + ], + [ + 29, + 42.0 + ], + [ + 8, + 45.0 + ], + [ + 12, + 43.0 + ], + [ + 16, + 45.0 + ], + [ + 28, + 45.0 + ], + [ + 32, + 41.0 + ], + [ + 36, + 41.0 + ], + [ + 40, + 45.0 + ], + [ + 52, + 43.0 + ], + [ + 56, + 41.0 + ], + [ + 60, + 43.0 + ], + [ + 76, + 45.0 + ], + [ + 80, + 45.0 + ], + [ + 7, + 50.0 + ], + [ + 11, + 50.0 + ], + [ + 27, + 50.0 + ], + [ + 31, + 48.0 + ], + [ + 35, + 50.0 + ], + [ + 51, + 50.0 + ], + [ + 55, + 50.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/range_search/range_slack_high_expands_frontier.json b/diskann/test/generated/graph/test/cases/range_search/range_slack_high_expands_frontier.json new file mode 100644 index 0000000000..8a6e61e830 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/range_search/range_slack_high_expands_frontier.json @@ -0,0 +1,263 @@ +{ + "file": "diskann/src/graph/test/cases/range_search.rs", + "test": "graph/test/cases/range_search/range_slack_high_expands_frontier", + "payload": { + "comparisons": 110, + "description": "High range_slack expands frontier exploration.", + "grid_dims": 3, + "grid_size": 5, + "hops": 90, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 30.0, + "range_search_second_round": true, + "result_count": 60, + "results": [ + [ + 124, + 3.0 + ], + [ + 123, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 99, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 94, + 9.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 73, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 121, + 18.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 109, + 18.0 + ], + [ + 49, + 18.0 + ], + [ + 48, + 21.0 + ], + [ + 68, + 17.0 + ], + [ + 72, + 19.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 120, + 27.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 84, + 21.0 + ], + [ + 104, + 27.0 + ], + [ + 24, + 27.0 + ], + [ + 23, + 30.0 + ], + [ + 43, + 24.0 + ], + [ + 47, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 71, + 26.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 59, + 26.0 + ], + [ + 79, + 30.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/range_search/range_slack_low_constrains_frontier.json b/diskann/test/generated/graph/test/cases/range_search/range_slack_low_constrains_frontier.json new file mode 100644 index 0000000000..2819c4c2c8 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/range_search/range_slack_low_constrains_frontier.json @@ -0,0 +1,263 @@ +{ + "file": "diskann/src/graph/test/cases/range_search.rs", + "test": "graph/test/cases/range_search/range_slack_low_constrains_frontier", + "payload": { + "comparisons": 92, + "description": "Low range_slack constrains frontier expansion.", + "grid_dims": 3, + "grid_size": 5, + "hops": 66, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 30.0, + "range_search_second_round": true, + "result_count": 60, + "results": [ + [ + 124, + 3.0 + ], + [ + 123, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 99, + 6.0 + ], + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 94, + 9.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 73, + 14.0 + ], + [ + 93, + 12.0 + ], + [ + 97, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 121, + 18.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 109, + 18.0 + ], + [ + 49, + 18.0 + ], + [ + 48, + 21.0 + ], + [ + 68, + 17.0 + ], + [ + 72, + 19.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 96, + 21.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 120, + 27.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 84, + 21.0 + ], + [ + 104, + 27.0 + ], + [ + 24, + 27.0 + ], + [ + 23, + 30.0 + ], + [ + 43, + 24.0 + ], + [ + 47, + 26.0 + ], + [ + 63, + 22.0 + ], + [ + 67, + 22.0 + ], + [ + 71, + 26.0 + ], + [ + 83, + 24.0 + ], + [ + 87, + 22.0 + ], + [ + 91, + 24.0 + ], + [ + 95, + 30.0 + ], + [ + 103, + 30.0 + ], + [ + 107, + 26.0 + ], + [ + 111, + 26.0 + ], + [ + 115, + 30.0 + ], + [ + 19, + 30.0 + ], + [ + 39, + 26.0 + ], + [ + 59, + 26.0 + ], + [ + 79, + 30.0 + ], + [ + 38, + 29.0 + ], + [ + 42, + 29.0 + ], + [ + 58, + 29.0 + ], + [ + 62, + 27.0 + ], + [ + 66, + 29.0 + ], + [ + 82, + 29.0 + ], + [ + 86, + 29.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file From 79ee33904a60d7180f74e21d9a4399585348f4f5 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 14 Aug 2026 19:41:50 +0000 Subject: [PATCH 4/7] somehow another format check --- diskann/src/graph/search/range_search.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 8f68c0eba1..340c03508d 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -545,9 +545,7 @@ where // both these results guarantee within outer radius range_frontier.push_back(*neighbor.id()); } - InRangePushResult::RejectedbyOuter - if *neighbor.distance() <= navigation_radius => - { + InRangePushResult::RejectedbyOuter if *neighbor.distance() <= navigation_radius => { // add to the range frontier if it's within the tolerance range_frontier.push_back(*neighbor.id()); } From 6ae9b8839058ba3c3114c26c8617c5f8098282fa Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Mon, 17 Aug 2026 16:08:20 -0700 Subject: [PATCH 5/7] Proposed change sketch. --- .../src/graph/search/filtered_range_search.rs | 88 +++++----- diskann/src/graph/search/range_search.rs | 153 +++++++----------- ..._respected_and_second_round_triggered.json | 4 +- 3 files changed, 100 insertions(+), 145 deletions(-) diff --git a/diskann/src/graph/search/filtered_range_search.rs b/diskann/src/graph/search/filtered_range_search.rs index ae9c1db4ca..eee403cbda 100644 --- a/diskann/src/graph/search/filtered_range_search.rs +++ b/diskann/src/graph/search/filtered_range_search.rs @@ -18,12 +18,13 @@ use crate::{ search::inline_filter_search::{Ret, inline_filter_search_internal}, search::{ Range, RangeSearchError, Search, - range_search::{InRange, InRangePushResult, RangeBuilder}, + range_search::{InRange, RangeBuilder}, record::NoopSearchRecord, scratch::SearchScratch, }, search_output_buffer::SearchOutputBuffer, }, + neighbor::Neighbor, provider::DataProvider, }; @@ -164,31 +165,30 @@ where let max_returned = self.effective_max_returned(num_start_ids); - let mut in_outer_range = InRange::new_checked( + // Filter matched results by radius. + // + // This will be used to decide if `max_results` has been reached. + let mut matched_in_outer_range = InRange::new( self.radius(), - None, - usize::MAX, - scratch.best.iter().take(starting_l).collect(), + self.inner_radius(), + max_returned, + matched_results.into_iter(), ); - // filter matched results by radius; this will be used to decide if `max_results` has been reached - let mut matched_in_outer_range = InRange::new_checked( + // Merge `matched_results` with the best results from the first round, + // filtering by radius + let mut in_outer_range = InRange::new( self.radius(), - self.inner_radius(), - max_returned, - matched_results.to_vec(), + None, + usize::MAX, + scratch + .best + .iter() + .take(starting_l) + .chain(matched_in_outer_range.iter()), ); - // merge matched_results with the best results from the first round, filtering by radius - for neighbor in matched_in_outer_range.neighbors.iter() { - in_outer_range.push_unchecked(*neighbor); - } - in_outer_range - .neighbors - .sort_unstable_by(crate::neighbor::ord::fast_distance_total); - in_outer_range - .neighbors - .dedup_by(|left, right| left.id() == right.id()); + in_outer_range.dedup(); let stats = if in_outer_range.len() >= ((starting_l as f32) * self.initial_slack()) as usize @@ -196,16 +196,13 @@ where { // clear the visited set and repopulate it with all in-range points found so far, filtered and unfiltered scratch.visited.clear(); - scratch.visited.extend( - in_outer_range - .neighbors - .iter() - .map(|neighbor| *neighbor.id()), - ); + scratch + .visited + .extend(in_outer_range.iter().map(|neighbor| *neighbor.id())); // Create a range frontier for seeding the second-round search let mut range_frontier: std::collections::VecDeque<_> = in_outer_range - .neighbors + .take() .iter() .map(|neighbor| *neighbor.id()) .collect(); @@ -234,11 +231,7 @@ where } }; - let truncated_matched = matched_in_outer_range - .neighbors - .iter() - .copied() - .take(max_returned); + let truncated_matched = matched_in_outer_range.iter().take(max_returned); let result_count = processor .post_process(&mut accessor, query, truncated_matched, output) @@ -279,9 +272,7 @@ where let mut neighbors = Vec::with_capacity(max_degree_with_slack); - let max_returned = search_params.max_returned().unwrap_or(usize::MAX); - - 'outer: while !range_frontier.is_empty() && matched_in_range.len() < max_returned { + while !range_frontier.is_empty() && !matched_in_range.is_full() { scratch.beam_nodes.clear(); // In this loop we are going to find the beam_width number of remaining nodes within the radius @@ -307,23 +298,16 @@ where // but only accepted IDs are added to in-range results. let navigation_radius = search_params.radius() * search_params.range_slack(); for (decision, distance) in neighbors.iter().copied() { - match matched_in_range.push_with_filter(distance, decision) { - InRangePushResult::Full => { - break 'outer; - } - InRangePushResult::Accepted - | InRangePushResult::RejectedbyInner - | InRangePushResult::RejectedbyFilter => { - let id = decision.into_inner(); - range_frontier.push_back(id); - } - InRangePushResult::RejectedbyOuter => { - // add to the range frontier if it's within the tolerance - if distance <= navigation_radius { - let id = decision.into_inner(); - range_frontier.push_back(id); - } - } + if matched_in_range.is_full() { + break; + } + + if let glue::Decision::Accept(id) = decision + && matched_in_range.push(Neighbor::new(id.into_inner(), distance)) + { + range_frontier.push_back(id.into_inner()); + } else if distance <= navigation_radius { + range_frontier.push_back(decision.into_inner()); } } diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 340c03508d..99bd08ae2c 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -317,41 +317,33 @@ where ) .await?; - let in_outer_range = InRange::new_checked( - self.radius(), - None, - starting_l, - scratch.best.iter().collect(), - ); + let in_outer_range = InRange::new(self.radius(), None, starting_l, scratch.best.iter()); // Increment the max results by the number of starting points, in case // they are filtered out later and leave us with fewer than the requested // number of results. let max_returned = self.effective_max_returned(num_start_ids); - let outer_range_len = in_outer_range.len(); - let outer_range_ids: Vec<_> = in_outer_range - .neighbors - .iter() - .map(|neighbor| *neighbor.id()) - .collect(); - - let mut in_range = InRange::new_checked( + let mut in_range = InRange::new( self.radius(), self.inner_radius(), max_returned, - in_outer_range.neighbors, + in_outer_range.iter(), ); - let stats = if outer_range_len >= ((starting_l as f32) * self.initial_slack()) as usize - && outer_range_len <= max_returned + let stats = if in_outer_range.len() + >= ((starting_l as f32) * self.initial_slack()) as usize + && in_outer_range.len() <= max_returned { // clear the visited set and repopulate it with just the in-range points scratch.visited.clear(); - scratch.visited.extend(outer_range_ids.iter().copied()); + scratch + .visited + .extend(in_outer_range.iter().map(|n| *n.id())); // Create a range frontier for seeding the second-round search - let mut range_frontier: VecDeque<_> = outer_range_ids.into_iter().collect(); + let mut range_frontier: VecDeque<_> = + in_outer_range.take().into_iter().map(|n| *n.id()).collect(); // Move to range search let range_stats = range_search_internal( @@ -374,12 +366,7 @@ where }; let result_count = processor - .post_process( - &mut accessor, - query, - in_range.neighbors.iter().copied(), - output, - ) + .post_process(&mut accessor, query, in_range.iter(), output) .await .into_ann_result()?; @@ -393,82 +380,73 @@ where } } -pub(super) enum InRangePushResult { - Accepted, - RejectedbyOuter, - RejectedbyInner, // this state is only returned when the outer radius is respected - RejectedbyFilter, // this state is only returned when outer AND inner radius are respected - Full, // this state is returned *before* checking inner/outer radius -} - pub(super) struct InRange { - pub neighbors: Vec>, + neighbors: Vec>, radius: f32, inner_radius: Option, max_returned: usize, } impl InRange { - pub(super) fn push(&mut self, neighbor: Neighbor) -> InRangePushResult { - if self.neighbors.len() >= self.max_returned { - return InRangePushResult::Full; - } - let dist = *neighbor.distance(); - if dist > self.radius { - return InRangePushResult::RejectedbyOuter; - } - if let Some(inner) = self.inner_radius - && dist <= inner - { - return InRangePushResult::RejectedbyInner; + #[must_use] + pub(super) fn push(&mut self, neighbor: Neighbor) -> bool { + let d = *neighbor.distance(); + if self.neighbors.len() < self.max_returned && self.check(d) { + self.neighbors.push(neighbor); + true + } else { + false } - self.neighbors.push(neighbor); - InRangePushResult::Accepted } - pub(super) fn push_with_filter( - &mut self, - distance: f32, - decision: glue::Decision, - ) -> InRangePushResult { - if self.neighbors.len() >= self.max_returned { - return InRangePushResult::Full; - } - if distance > self.radius { - return InRangePushResult::RejectedbyOuter; - } - if let Some(inner) = self.inner_radius - && distance <= inner - { - return InRangePushResult::RejectedbyInner; - } - if decision.is_reject() { - return InRangePushResult::RejectedbyFilter; - } + pub(super) fn dedup(&mut self) + where + I: Ord, + { self.neighbors - .push(Neighbor::new(decision.into_inner(), distance)); - InRangePushResult::Accepted - } - - pub(super) fn push_unchecked(&mut self, neighbor: Neighbor) { - self.neighbors.push(neighbor); + .sort_unstable_by(crate::neighbor::ord::fast_distance_total); + self.neighbors + .dedup_by(|left, right| left.id() == right.id()); } pub(super) fn len(&self) -> usize { self.neighbors.len() } + pub(super) fn is_full(&self) -> bool { + self.len() == self.max_returned + } + + pub(super) fn take(self) -> Vec> { + self.neighbors + } + + pub(super) fn iter(&self) -> impl ExactSizeIterator> + where + I: Copy, + { + self.neighbors.iter().copied() + } + + #[must_use] + pub(super) fn check(&self, distance: f32) -> bool { + distance <= self.radius && self.inner_radius.map_or(true, |inner| distance > inner) + } + /// Create a new InRange with the given parameters, /// filtering the candidate neighbors and truncating /// as needed to respect the maximum number of results. - pub(super) fn new_checked( + pub(super) fn new( radius: f32, inner_radius: Option, max_returned: usize, - candidate_neighbors: Vec>, - ) -> Self { + candidates: Itr, + ) -> Self + where + Itr: IntoIterator>, + { Self { - neighbors: candidate_neighbors + neighbors: candidates .into_iter() .filter(|n| { let dist = *n.distance(); @@ -514,7 +492,7 @@ where let mut neighbors = Vec::with_capacity(max_degree_with_slack); - 'outer: while !range_frontier.is_empty() { + while !range_frontier.is_empty() && !in_range.is_full() { scratch.beam_nodes.clear(); // In this loop we are going to find the beam_width number of remaining nodes within the radius @@ -537,19 +515,12 @@ where let navigation_radius = search_params.radius() * search_params.range_slack(); for neighbor in neighbors.iter() { - match in_range.push(*neighbor) { - InRangePushResult::Full => { - break 'outer; // we've reached the maximum number of results, stop processing - } - InRangePushResult::Accepted | InRangePushResult::RejectedbyInner => { - // both these results guarantee within outer radius - range_frontier.push_back(*neighbor.id()); - } - InRangePushResult::RejectedbyOuter if *neighbor.distance() <= navigation_radius => { - // add to the range frontier if it's within the tolerance - range_frontier.push_back(*neighbor.id()); - } - _ => {} + if in_range.is_full() { + break; + } + + if in_range.push(*neighbor) || *neighbor.distance() <= navigation_radius { + range_frontier.push_back(*neighbor.id()); } } diff --git a/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json index 357e576228..a43d226f47 100644 --- a/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json +++ b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json @@ -2,11 +2,11 @@ "file": "diskann/src/graph/test/cases/range_search.rs", "test": "graph/test/cases/range_search/max_results_respected_and_second_round_triggered", "payload": { - "comparisons": 12, + "comparisons": 15, "description": "Two round search test to validate that max_results > starting_l means a second round is triggered.", "grid_dims": 3, "grid_size": 5, - "hops": 7, + "hops": 8, "inner_radius": null, "query": [ 5.0, From 027cb1faaaa1fe01f6c7092c7e458ad53cde7f35 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Wed, 19 Aug 2026 13:16:31 +0000 Subject: [PATCH 6/7] merge with main and mhildebr/review, add tests for InRange --- diskann/src/graph/search/range_search.rs | 79 ++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 99bd08ae2c..f59e402c95 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -543,6 +543,85 @@ where mod tests { use super::*; + fn neighbor(id: u32, distance: f32) -> Neighbor { + Neighbor::new(id, distance) + } + + #[test] + fn in_range_applies_inner_and_outer_radius_boundaries() { + let candidates = [ + neighbor(0, 0.1), + neighbor(1, 0.2), + neighbor(2, 0.3), + neighbor(3, 0.5), + neighbor(4, 0.6), + ]; + + let in_range = InRange::new(0.5, Some(0.2), usize::MAX, candidates); + + let ids: Vec<_> = in_range.iter().map(|candidate| *candidate.id()).collect(); + assert_eq!(ids, [2, 3]); + assert!(!in_range.check(0.2)); + assert!(in_range.check(0.5)); + } + + #[test] + fn in_range_filters_before_truncating_and_preserves_order() { + let candidates = [ + neighbor(0, 0.8), + neighbor(1, 0.3), + neighbor(2, 0.7), + neighbor(3, 0.2), + neighbor(4, 0.1), + ]; + + let in_range = InRange::new(0.5, None, 2, candidates); + + let ids: Vec<_> = in_range.iter().map(|candidate| *candidate.id()).collect(); + assert_eq!(ids, [1, 3]); + assert_eq!(in_range.len(), 2); + assert!(in_range.is_full()); + } + + #[test] + fn in_range_push_rejects_out_of_range_and_over_capacity() { + let mut in_range = InRange::new(0.5, Some(0.1), 2, []); + + assert!(!in_range.push(neighbor(0, 0.1))); + assert!(!in_range.push(neighbor(1, 0.6))); + assert!(in_range.push(neighbor(2, 0.2))); + assert!(!in_range.is_full()); + assert!(in_range.push(neighbor(3, 0.5))); + assert!(in_range.is_full()); + assert!(!in_range.push(neighbor(4, 0.3))); + } + + #[test] + fn in_range_dedup_sorts_by_distance_and_removes_repeated_ids() { + let mut in_range = InRange::new( + 1.0, + None, + usize::MAX, + [ + neighbor(2, 0.4), + neighbor(1, 0.2), + neighbor(1, 0.2), + neighbor(3, 0.3), + ], + ); + + in_range.dedup(); + + let neighbors = in_range.take(); + let ids: Vec<_> = neighbors.iter().map(|candidate| *candidate.id()).collect(); + let distances: Vec<_> = neighbors + .iter() + .map(|candidate| *candidate.distance()) + .collect(); + assert_eq!(ids, [1, 3, 2]); + assert_eq!(distances, [0.2, 0.3, 0.4]); + } + #[test] fn range_builder_defaults_match_new() { let from_new = Range::new(100, 0.5).unwrap(); From 474ad46b47a2051c14e9ee95282c28aec6124cb1 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Wed, 19 Aug 2026 14:19:58 +0000 Subject: [PATCH 7/7] added result limiting post-processor and corresponding tests --- .../src/graph/search/filtered_range_search.rs | 17 +- diskann/src/graph/search/range_search.rs | 194 ++++++++++++++++-- diskann/src/graph/test/cases/range_search.rs | 91 +++++++- 3 files changed, 277 insertions(+), 25 deletions(-) diff --git a/diskann/src/graph/search/filtered_range_search.rs b/diskann/src/graph/search/filtered_range_search.rs index eee403cbda..f681092578 100644 --- a/diskann/src/graph/search/filtered_range_search.rs +++ b/diskann/src/graph/search/filtered_range_search.rs @@ -18,7 +18,7 @@ use crate::{ search::inline_filter_search::{Ret, inline_filter_search_internal}, search::{ Range, RangeSearchError, Search, - range_search::{InRange, RangeBuilder}, + range_search::{InRange, LimitedOutputBuffer, RangeBuilder}, record::NoopSearchRecord, scratch::SearchScratch, }, @@ -172,7 +172,7 @@ where self.radius(), self.inner_radius(), max_returned, - matched_results.into_iter(), + matched_results, ); // Merge `matched_results` with the best results from the first round, @@ -188,7 +188,7 @@ where .chain(matched_in_outer_range.iter()), ); - in_outer_range.dedup(); + in_outer_range.sort_and_dedup(); let stats = if in_outer_range.len() >= ((starting_l as f32) * self.initial_slack()) as usize @@ -231,10 +231,15 @@ where } }; - let truncated_matched = matched_in_outer_range.iter().take(max_returned); - + let mut limited_output = + LimitedOutputBuffer::new(output, self.max_returned().unwrap_or(usize::MAX)); let result_count = processor - .post_process(&mut accessor, query, truncated_matched, output) + .post_process( + &mut accessor, + query, + matched_in_outer_range.iter(), + &mut limited_output, + ) .await .into_ann_result()?; diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index f59e402c95..f80b0da26e 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -21,12 +21,72 @@ use crate::{ Knn, Search, filtered_range_search::FilteredRange, record::NoopSearchRecord, scratch::SearchScratch, }, - search_output_buffer::SearchOutputBuffer, + search_output_buffer::{BufferState, SearchOutputBuffer}, }, neighbor::Neighbor, provider::DataProvider, }; +/// Limits the number of additional results written to an output buffer. +pub(super) struct LimitedOutputBuffer<'a, B: ?Sized> { + inner: &'a mut B, + remaining: usize, +} + +impl<'a, B: ?Sized> LimitedOutputBuffer<'a, B> { + pub(super) fn new(inner: &'a mut B, limit: usize) -> Self { + Self { + inner, + remaining: limit, + } + } +} + +impl SearchOutputBuffer for LimitedOutputBuffer<'_, B> +where + B: SearchOutputBuffer + ?Sized, +{ + fn size_hint(&self) -> Option { + Some( + self.inner + .size_hint() + .map_or(self.remaining, |inner| inner.min(self.remaining)), + ) + } + + fn push(&mut self, neighbor: Neighbor) -> BufferState { + if self.remaining == 0 { + return BufferState::Full; + } + + let previous_len = self.inner.current_len(); + let state = self.inner.push(neighbor); + if self.inner.current_len() > previous_len { + self.remaining -= 1; + } + + if self.remaining == 0 { + BufferState::Full + } else { + state + } + } + + fn current_len(&self) -> usize { + self.inner.current_len() + } + + fn extend(&mut self, itr: Itr) -> usize + where + Itr: IntoIterator>, + { + let count = self.inner.extend(itr.into_iter().take(self.remaining)); + debug_assert!(count <= self.remaining); + self.remaining = self.remaining.saturating_sub(count); + count + } +} + /// Error type for [`Range`] parameter validation. #[derive(Debug, Error)] pub enum RangeSearchError { @@ -197,24 +257,32 @@ impl Range { /// Builder for [`Range`] search parameters. /// -/// `max_returned`: If specified, the search will stop and return results once this number of points has been found within both the inner and outer radii. -/// Note that due to adding some extra slack for start points, occasionally slightly more than this number of points may be returned. Since the initial -/// search phase does not respect `max_returned`, this parameter may not be set lower than `starting_l`. +/// `max_returned`: If specified, the search will stop and return results once this number +/// of points has been found within both the inner and outer radii. Extra candidate slack is +/// reserved for start points, then the post-processed output is capped at this value after +/// start points are removed. Since the initial search phase does not respect `max_returned`, +/// this parameter may not be set lower than `starting_l`. /// /// `starting_l`: the L_search parameter for the initial search phase. Must be greater than zero. /// -/// `beam_width`: the beam width for parallel graph exploration. If not specified, defaults to 1. Must be greater than zero if specified. +/// `beam_width`: the beam width for parallel graph exploration. If not specified, defaults +/// to 1. Must be greater than zero if specified. /// -/// `radius`: the outer radius for the range search. Points within this distance from the query are candidates for inclusion in the results. +/// `radius`: the outer radius for the range search. Points within this distance from the +/// query are candidates for inclusion in the results. /// -/// `inner_radius`: the inner radius for the range search. Points closer than this distance from the query are excluded from the results. Must be less than or equal to `radius` if specified. +/// `inner_radius`: the inner radius for the range search. Points closer than this distance from the +/// query are excluded from the results. Must be less than or equal to `radius` if specified. /// -/// `initial_slack`: after the initial knn search phase, a decision is made on whether to continue to the second round of search. This decision is based on whether -/// the number of points found within the outer radius is greater than `starting_l * initial_slack`, so lower values of `initial_slack` will make it more likely to continue to -/// the second round of search. Must be between 0.0 and 1.0. +/// `initial_slack`: after the initial knn search phase, a decision is made on whether to continue +/// to the second round of search. This decision is based on whether the number of points found +/// within the outer radius is greater than `starting_l * initial_slack`, so lower values of +/// `initial_slack` will make it more likely to continue to the second round of search. Must be +/// between 0.0 and 1.0. /// -/// `range_slack`: during the second round of search, points that are within `radius * range_slack` are expanded to search for candidates within the range, so greater -/// values of `range_slack` will mean more expansions. Must be greater than or equal to 1.0. +/// `range_slack`: during the second round of search, points that are within `radius * range_slack` +/// are expanded to search for candidates within the range, so greater values of `range_slack` will +/// mean more expansions. Must be greater than or equal to 1.0. #[derive(Debug, Clone, Copy)] pub struct RangeBuilder { max_returned: Option, @@ -365,8 +433,10 @@ where initial_stats }; + let mut limited_output = + LimitedOutputBuffer::new(output, self.max_returned().unwrap_or(usize::MAX)); let result_count = processor - .post_process(&mut accessor, query, in_range.iter(), output) + .post_process(&mut accessor, query, in_range.iter(), &mut limited_output) .await .into_ann_result()?; @@ -399,7 +469,7 @@ impl InRange { } } - pub(super) fn dedup(&mut self) + pub(super) fn sort_and_dedup(&mut self) where I: Ord, { @@ -430,7 +500,7 @@ impl InRange { #[must_use] pub(super) fn check(&self, distance: f32) -> bool { - distance <= self.radius && self.inner_radius.map_or(true, |inner| distance > inner) + distance <= self.radius && self.inner_radius.is_none_or(|inner| distance > inner) } /// Create a new InRange with the given parameters, @@ -542,11 +612,101 @@ where #[cfg(test)] mod tests { use super::*; + use crate::graph::IdDistance; fn neighbor(id: u32, distance: f32) -> Neighbor { Neighbor::new(id, distance) } + #[test] + fn limited_output_buffer_applies_limit_after_filtering() { + let candidates = [ + neighbor(0, 0.0), + neighbor(1, 0.1), + neighbor(2, 0.2), + neighbor(3, 0.3), + ]; + let mut output = Vec::new(); + let mut limited = LimitedOutputBuffer::new(&mut output, 2); + + let written = limited.extend( + candidates + .into_iter() + .filter(|candidate| *candidate.id() != 0), + ); + + assert_eq!(written, 2); + assert_eq!(limited.current_len(), 2); + assert_eq!(limited.size_hint(), Some(0)); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + [(1, 0.1), (2, 0.2)] + ); + } + + #[test] + fn limited_output_buffer_push_reports_full_at_limit() { + let mut output = Vec::new(); + let mut limited = LimitedOutputBuffer::new(&mut output, 2); + + assert!(limited.push(neighbor(1, 0.1)).is_available()); + assert!(limited.push(neighbor(2, 0.2)).is_full()); + assert!(limited.push(neighbor(3, 0.3)).is_full()); + assert_eq!(limited.size_hint(), Some(0)); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + [(1, 0.1), (2, 0.2)] + ); + } + + #[test] + fn limited_output_buffer_zero_limit_writes_nothing() { + let mut output = Vec::new(); + let mut limited = LimitedOutputBuffer::new(&mut output, 0); + + assert!(limited.push(neighbor(1, 0.1)).is_full()); + assert_eq!(limited.extend([neighbor(2, 0.2)]), 0); + assert_eq!(limited.size_hint(), Some(0)); + assert!(output.is_empty()); + } + + #[test] + fn limited_output_buffer_respects_smaller_inner_capacity() { + let mut ids = [0; 1]; + let mut distances = [0.0; 1]; + let mut output = IdDistance::new(&mut ids, &mut distances); + let mut limited = LimitedOutputBuffer::new(&mut output, 3); + + assert!(limited.push(neighbor(1, 0.1)).is_full()); + assert!(limited.push(neighbor(2, 0.2)).is_full()); + assert_eq!(limited.current_len(), 1); + assert_eq!(limited.size_hint(), Some(0)); + assert_eq!(ids, [1]); + assert_eq!(distances, [0.1]); + } + + #[test] + fn limited_output_buffer_limit_counts_only_additional_results() { + let mut output = vec![neighbor(0, 0.0)]; + let mut limited = LimitedOutputBuffer::new(&mut output, 2); + + assert_eq!(limited.extend([neighbor(1, 0.1), neighbor(2, 0.2)]), 2); + assert_eq!(limited.current_len(), 3); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + [(0, 0.0), (1, 0.1), (2, 0.2)] + ); + } + #[test] fn in_range_applies_inner_and_outer_radius_boundaries() { let candidates = [ @@ -597,7 +757,7 @@ mod tests { } #[test] - fn in_range_dedup_sorts_by_distance_and_removes_repeated_ids() { + fn in_range_sort_and_dedup_sorts_by_distance_and_removes_repeated_ids() { let mut in_range = InRange::new( 1.0, None, @@ -610,7 +770,7 @@ mod tests { ], ); - in_range.dedup(); + in_range.sort_and_dedup(); let neighbors = in_range.take(); let ids: Vec<_> = neighbors.iter().map(|candidate| *candidate.id()).collect(); diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index 111dd92b7d..4638e7d0c2 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -9,18 +9,26 @@ //! and empty result handling. Integration tests use baselines for regression //! protection. -use std::sync::Arc; +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, +}; use diskann_vector::distance::Metric; use crate::{ graph::{ - self, DiskANNIndex, + self, DiskANNIndex, SearchOutputBuffer, + glue::SearchPostProcess, index::SearchStats, search::Range, test::{provider as test_provider, synthetic::Grid}, }, neighbor::Neighbor, + provider::HasId, test::{ TestRoot, cmp::{assert_eq_verbose, verbose_eq}, @@ -29,6 +37,43 @@ use crate::{ }, }; +#[derive(Clone)] +struct RecordingCopyIds { + candidate_count: Arc, + saw_start_point: Arc, +} + +impl SearchPostProcess for RecordingCopyIds +where + A: HasId, +{ + type Error = Infallible; + + fn post_process( + &self, + _accessor: &mut A, + _query: T, + candidates: I, + output: &mut B, + ) -> impl std::future::Future> + Send + where + I: Iterator> + Send, + B: SearchOutputBuffer + Send + ?Sized, + { + let candidates: Vec<_> = candidates.collect(); + self.candidate_count + .store(candidates.len(), Ordering::Relaxed); + self.saw_start_point.store( + candidates + .iter() + .any(|candidate| *candidate.id() == u32::MAX), + Ordering::Relaxed, + ); + let count = output.extend(candidates); + std::future::ready(Ok(count)) + } +} + #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub(super) struct RangeSearchBaseline { /// A description of what to expect, what trends to observe, and anything else @@ -413,6 +458,48 @@ fn max_results_respected_and_second_round_triggered() { assert_no_duplicates(&results); } +#[test] +fn max_results_caps_non_start_candidates_after_range_collection() { + let rt = current_thread_runtime(); + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let starting_l = 4; + let max_results = 5; + let range_search = Range::builder(starting_l, f32::MAX) + .inner_radius(Some(0.0)) + .max_returned(Some(max_results)) + .build() + .unwrap(); + + let candidate_count = Arc::new(AtomicUsize::new(0)); + let saw_start_point = Arc::new(AtomicBool::new(false)); + let processor = RecordingCopyIds { + candidate_count: Arc::clone(&candidate_count), + saw_start_point: Arc::clone(&saw_start_point), + }; + let mut results = Vec::>::new(); + + let stats = rt + .block_on(index.search_with( + range_search, + &test_provider::Strategy::new(), + processor, + &test_provider::Context::new(), + query.as_slice(), + &mut results, + )) + .unwrap(); + + assert_eq!(candidate_count.load(Ordering::Relaxed), max_results + 1); + assert!(!saw_start_point.load(Ordering::Relaxed)); + assert_eq!(results.len(), max_results); + assert_eq!(stats.result_count as usize, max_results); + assert!(stats.range_search_second_round); + assert!(results.iter().all(|result| *result.id() != u32::MAX)); + assert_range_invariants(&results, f32::MAX, Some(0.0)); + assert_no_duplicates(&results); +} + #[test] fn initial_slack_low_triggers_second_round() { let _description = "Test that low initial_slack triggers second round. \