diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 66380ad21..01c859ca5 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -36,6 +36,8 @@ 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 +93,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 +204,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 = 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); } @@ -211,7 +221,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 +347,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 @@ -401,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 b375567df..6e2f4dce2 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 = 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 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 = 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 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); +} 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..add732b7c --- /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": 12, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 1000000000.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..4c93a1589 --- /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": 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