From 04c3126f0f647745896183dba0c9c0d6ae6d6937 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Mon, 13 Jul 2026 15:31:46 +0000 Subject: [PATCH 1/5] add new test, clean up range search a bit --- diskann/src/graph/search/range_search.rs | 18 ++- diskann/src/graph/test/cases/range_search.rs | 117 +++++++++++++++++++ 2 files changed, 132 insertions(+), 3 deletions(-) diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 66380ad21..9682e0561 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -36,6 +36,9 @@ pub enum RangeSearchError { RangeSearchSlackValueError, #[error("inner_radius must be less than or equal to radius")] InnerRadiusValueError, + #[error("max_returned must be greater than or equal to starting_l")] + MaxReturnedLessThanInitialL, + } impl From for ANNError { @@ -91,6 +94,11 @@ impl Range { if starting_l == 0 { return Err(RangeSearchError::LZero); } + if let Some(max) = max_returned + && max < starting_l + { + return Err(RangeSearchError::MaxReturnedLessThanInitialL); + } if !(0.0..=1.0).contains(&initial_slack) { return Err(RangeSearchError::StartingListSlackValueError); } @@ -197,7 +205,10 @@ where let mut in_range = Vec::with_capacity(self.starting_l().into_usize()); - for neighbor in scratch.best.iter().take(self.starting_l().into_usize()) { + let starting_l = self.starting_l().into_usize(); + let max_returned = search_params.max_returned().unwrap_or(usize::MAX); + + for neighbor in scratch.best.iter().take(starting_l) { if neighbor.distance <= self.radius() { in_range.push(neighbor); } @@ -211,7 +222,8 @@ where scratch.in_range = in_range; let stats = if scratch.in_range.len() - >= ((self.starting_l() as f32) * self.initial_slack()) as usize + >= ((starting_l as f32) * self.initial_slack()) as usize + && scratch.in_range.len() < max_returned { // Move to range search let range_stats = range_search_internal( @@ -336,7 +348,7 @@ where let max_returned = search_params.max_returned().unwrap_or(usize::MAX); - while !scratch.range_frontier.is_empty() { + while !scratch.range_frontier.is_empty() && scratch.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 diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index b375567df..dd0cbad18 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -284,3 +284,120 @@ fn empty_results() { "empty results shouldn't trigger a second round" ); } + +#[test] +fn max_results_respected_means_no_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_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; // 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 limited to starting_l, so no second round should be triggered + + let range_search = + Range::with_options(Some(max_results), starting_l, None, radius, None, 1.0, 1.0).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 { + grid_size, + query: query.clone(), + radius, + inner_radius: None, + starting_l, + results: results.iter().map(|n| (n.id, n.distance)).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!( + results.len() <= max_results, + "result count {} exceeds max_results {}", + results.len(), + max_results + ); + + assert!( + !stats.range_search_second_round, + "If max_results is respected, a second round should not be triggered" + ); + assert_range_invariants(&results, radius, None); + assert_no_duplicates(&results); +} + +#[test] +fn max_results_respected_and_second_round_triggered() { + 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 grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 50.0; // 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 range_search = + Range::with_options(Some(max_results), starting_l, None, radius, None, 1.0, 1.0).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 { + grid_size, + query: query.clone(), + radius, + inner_radius: None, + starting_l, + results: results.iter().map(|n| (n.id, n.distance)).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!( + results.len() <= max_results, + "result count {} exceeds max_results {}", + results.len(), + max_results + ); + + assert!( + 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); +} From fb3fbc98225621a717f31d37af2dc5a16541c725 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Mon, 13 Jul 2026 16:03:39 +0000 Subject: [PATCH 2/5] clippy, fmt --- diskann/src/graph/search/range_search.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 9682e0561..de19fd4e5 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -38,7 +38,6 @@ pub enum RangeSearchError { InnerRadiusValueError, #[error("max_returned must be greater than or equal to starting_l")] MaxReturnedLessThanInitialL, - } impl From for ANNError { @@ -205,8 +204,8 @@ where let mut in_range = Vec::with_capacity(self.starting_l().into_usize()); - let starting_l = self.starting_l().into_usize(); - let max_returned = search_params.max_returned().unwrap_or(usize::MAX); + let starting_l = self.starting_l().into_usize(); + let max_returned = self.max_returned().unwrap_or(usize::MAX); for neighbor in scratch.best.iter().take(starting_l) { if neighbor.distance <= self.radius() { From cfebc4db3095a0003f4eb2e56677f9d1efcbe1d3 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Mon, 13 Jul 2026 16:04:06 +0000 Subject: [PATCH 3/5] add new baselines --- ..._respected_and_second_round_triggered.json | 37 +++++++++++++++++++ ...sults_respected_means_no_second_round.json | 33 +++++++++++++++++ 2 files changed, 70 insertions(+) create mode 100644 diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json create mode 100644 diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json 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 new file mode 100644 index 000000000..f171a0167 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_and_second_round_triggered.json @@ -0,0 +1,37 @@ +{ + "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, + "grid_size": 5, + "hops": 15, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 50.0, + "range_search_second_round": true, + "result_count": 4, + "results": [ + [ + 124, + 3.0 + ], + [ + 123, + 6.0 + ], + [ + 119, + 6.0 + ], + [ + 99, + 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_means_no_second_round.json b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json new file mode 100644 index 000000000..124883e82 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/range_search/max_results_respected_means_no_second_round.json @@ -0,0 +1,33 @@ +{ + "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, + "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 + ], + [ + 123, + 6.0 + ], + [ + 119, + 6.0 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file From 9957c19ad5a245784895b147e45b0188381a4882 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Mon, 13 Jul 2026 16:38:57 +0000 Subject: [PATCH 4/5] regenerate baselines --- .../max_results_respected_and_second_round_triggered.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f171a0167..9a8b2f816 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 @@ -4,7 +4,7 @@ "payload": { "comparisons": 11, "grid_size": 5, - "hops": 15, + "hops": 12, "inner_radius": null, "query": [ 5.0, From 112abe8789a9ba68441b1ad0252fd9a1c59f060e Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Mon, 13 Jul 2026 19:41:40 +0000 Subject: [PATCH 5/5] regenerate baselines --- diskann/src/graph/search/range_search.rs | 3 +++ diskann/src/graph/test/cases/range_search.rs | 6 +++--- .../max_results_respected_and_second_round_triggered.json | 2 +- .../max_results_respected_means_no_second_round.json | 2 +- 4 files changed, 8 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index de19fd4e5..01c859ca5 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -412,6 +412,9 @@ mod tests { // Invalid inner radius > radius assert!(Range::with_options(None, 100, None, 0.5, Some(1.0), 1.0, 1.0).is_err()); + + // Invalid max_results < initial_l_search + assert!(Range::with_options(Some(50), 100, None, 0.5, None, 1.0, 1.0).is_err()); } #[test] diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index dd0cbad18..6e2f4dce2 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -294,9 +294,9 @@ fn 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; // every point will be in range with this radius + 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 limited to starting_l, so no second round should be triggered + let max_results = 4; // max_returned = starting_l, so second round should not be triggered let range_search = Range::with_options(Some(max_results), starting_l, None, radius, None, 1.0, 1.0).unwrap(); @@ -352,7 +352,7 @@ fn 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 = 50.0; // every point will be in range with this radius + 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 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 9a8b2f816..add732b7c 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 @@ -11,7 +11,7 @@ 5.0, 5.0 ], - "radius": 50.0, + "radius": 1000000000.0, "range_search_second_round": true, "result_count": 4, "results": [ 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 index 124883e82..4c93a1589 100644 --- 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 @@ -11,7 +11,7 @@ 5.0, 5.0 ], - "radius": 50.0, + "radius": 1000000000.0, "range_search_second_round": false, "result_count": 3, "results": [