diff --git a/diskann-benchmark-core/src/search/graph/filtered_range.rs b/diskann-benchmark-core/src/search/graph/filtered_range.rs new file mode 100644 index 000000000..94a6853c5 --- /dev/null +++ b/diskann-benchmark-core/src/search/graph/filtered_range.rs @@ -0,0 +1,379 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{num::NonZeroUsize, sync::Arc}; + +use diskann::{ + ANNResult, + graph::{self, ext::labeled, glue}, + provider, +}; +use diskann_benchmark_runner::utils::{MicroSeconds, percentiles}; +use diskann_utils::{future::AsyncFriendly, views::Matrix}; + +use crate::{ + recall, + search::{self, Search, graph::Strategy}, +}; + +/// A built-in helper for benchmarking the filtered range search method +/// [`graph::DiskANNIndex::search`]. +/// +/// This is intended to be used in conjunction with [`search::search`] or +/// [`search::search_all`] and provides some basic additional metrics for the +/// latter. Result aggregation for [`search::search_all`] is provided by the +/// [`Aggregator`] type. +/// +/// The provided implementation of [`Search`] accepts +/// [`graph::search::Range`] and returns [`Metrics`] as additional output. +#[derive(Debug)] +pub struct FilteredRange +where + DP: provider::DataProvider, +{ + index: Arc>, + queries: Arc>, + strategy: Strategy, + labels: Arc<[Arc>]>, +} + +impl FilteredRange +where + DP: provider::DataProvider, +{ + /// Construct a new [`FilteredRange`] searcher. + /// + /// If `strategy` is one of the container variants of [`Strategy`], its length + /// must match the number of rows in `queries`. If this is the case, then the + /// strategies will have a querywise correspondence (see [`search::SearchResults`]) + /// with the query matrix. + /// + /// # Errors + /// + /// Returns an error if the number of elements in `strategy` is not compatible with + /// the number of rows in `queries`. + pub fn new( + index: Arc>, + queries: Arc>, + strategy: Strategy, + labels: Arc<[Arc>]>, + ) -> anyhow::Result> { + strategy.length_compatible(queries.nrows())?; + + if labels.len() != queries.nrows() { + Err(anyhow::anyhow!( + "Number of label providers ({}) must be equal to the number of queries ({})", + labels.len(), + queries.nrows() + )) + } else { + Ok(Arc::new(Self { + index, + queries, + strategy, + labels, + })) + } + } +} + +/// Placeholder for filtered range search metrics. +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct Metrics {} + +impl Search for FilteredRange +where + DP: provider::DataProvider, + S: for<'a> glue::DefaultSearchStrategy< + 'a, + DP, + &'a [T], + DP::ExternalId, + SearchAccessor: glue::SearchAccessor, + > + Clone + + AsyncFriendly, + T: AsyncFriendly + Clone, +{ + type Id = DP::ExternalId; + type Parameters = graph::search::Range; + type Output = Metrics; + + fn num_queries(&self) -> usize { + self.queries.nrows() + } + + fn id_count(&self, parameters: &Self::Parameters) -> search::IdCount { + search::IdCount::Dynamic(NonZeroUsize::new(parameters.starting_l())) + } + + async fn search( + &self, + parameters: &Self::Parameters, + buffer: &mut O, + index: usize, + ) -> ANNResult + where + O: graph::SearchOutputBuffer + Send, + { + let context = DP::Context::default(); + let filtered_range_search = + graph::search::FilteredRange::builder(parameters.starting_l(), parameters.radius()) + .max_returned(parameters.max_returned()) + .beam_width(parameters.beam_width()) + .inner_radius(parameters.inner_radius()) + .initial_slack(parameters.initial_slack()) + .range_slack(parameters.range_slack()) + .build_filtered()?; + let strategy = + labeled::Filtered::new(self.strategy.get(index)?.clone(), &*self.labels[index]); + let _ = self + .index + .search( + filtered_range_search, + &strategy, + &context, + self.queries.row(index), + buffer, + ) + .await?; + + Ok(Metrics {}) + } +} + +/// An [`search::Aggregate`]d summary of multiple [`FilteredRange`] search runs returned by +/// the provided [`Aggregator`]. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct Summary { + pub setup: search::Setup, + pub parameters: graph::search::Range, + pub end_to_end_latencies: Vec, + pub mean_latencies: Vec, + pub p90_latencies: Vec, + pub p99_latencies: Vec, + pub average_precision: recall::AveragePrecisionMetrics, +} + +/// A [`search::Aggregate`] for collecting the results of multiple [`FilteredRange`] search +/// runs. +pub struct Aggregator<'a, I> { + groundtruth: &'a dyn crate::recall::Rows, +} + +impl<'a, I> Aggregator<'a, I> { + pub fn new(groundtruth: &'a dyn crate::recall::Rows) -> Self { + Self { groundtruth } + } +} + +impl search::Aggregate for Aggregator<'_, I> +where + I: crate::recall::RecallCompatible, +{ + type Output = Summary; + + #[inline(never)] + fn aggregate( + &mut self, + run: search::Run, + mut results: Vec>, + ) -> anyhow::Result { + let average_precision = match results.first() { + Some(first) => { + crate::recall::average_precision(first.ids().as_rows(), self.groundtruth)? + } + None => anyhow::bail!("Results must be non-empty"), + }; + + let mut mean_latencies = Vec::with_capacity(results.len()); + let mut p90_latencies = Vec::with_capacity(results.len()); + let mut p99_latencies = Vec::with_capacity(results.len()); + + results.iter_mut().for_each(|r| { + match percentiles::compute_percentiles(r.latencies_mut()) { + Ok(values) => { + let percentiles::Percentiles { mean, p90, p99, .. } = values; + mean_latencies.push(mean); + p90_latencies.push(p90); + p99_latencies.push(p99); + } + Err(_) => { + let zero = MicroSeconds::new(0); + mean_latencies.push(0.0); + p90_latencies.push(zero); + p99_latencies.push(zero); + } + } + }); + + Ok(Summary { + setup: run.setup().clone(), + parameters: *run.parameters(), + end_to_end_latencies: results.iter().map(|r| r.end_to_end_latency()).collect(), + mean_latencies, + p90_latencies, + p99_latencies, + average_precision, + }) + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use diskann::graph::{ext::labeled::QueryLabelProvider, test::provider}; + + #[derive(Debug)] + struct NoOdds; + + impl labeled::QueryLabelProvider for NoOdds { + fn is_match(&self, id: u32) -> bool { + id.is_multiple_of(2) + } + } + + #[test] + fn test_filtered_range() { + let index = search::graph::test_grid_provider(); + + let mut queries = Matrix::new(0.0f32, 5, index.provider().dim()); + queries.row_mut(0).copy_from_slice(&[0.0, 0.0, 0.0, 0.0]); + queries.row_mut(1).copy_from_slice(&[4.0, 0.0, 0.0, 0.0]); + queries.row_mut(2).copy_from_slice(&[0.0, 4.0, 0.0, 0.0]); + queries.row_mut(3).copy_from_slice(&[0.0, 0.0, 4.0, 0.0]); + queries.row_mut(4).copy_from_slice(&[0.0, 0.0, 0.0, 4.0]); + + let queries = Arc::new(queries); + let labels: Arc<[_]> = (0..queries.nrows()) + .map(|_| -> Arc> { Arc::new(NoOdds {}) }) + .collect(); + + let filtered_range = FilteredRange::new( + index, + queries.clone(), + Strategy::broadcast(provider::Strategy::new()), + labels, + ) + .unwrap(); + + // Test the standard search interface. + let rt = crate::tokio::runtime(2).unwrap(); + let results = search::search( + filtered_range.clone(), + graph::search::Range::builder(10, 2.0) + .initial_slack(0.8) + .range_slack(1.2) + .build() + .unwrap(), + NonZeroUsize::new(2).unwrap(), + &rt, + ) + .unwrap(); + + assert_eq!(results.len(), queries.nrows()); + let rows = results.ids().as_rows(); + + // Check that only even IDs are returned. + for r in 0..rows.nrows() { + for &id in rows.row(r) { + assert_eq!(id % 2, 0, "Found odd ID {} in row {}", id, r); + } + } + + const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap(); + let setup = search::Setup { + threads: TWO, + tasks: TWO, + reps: TWO, + }; + + // Try the aggregated strategy. + let parameters = [ + search::Run::new( + graph::search::Range::builder(10, 2.0) + .initial_slack(0.8) + .range_slack(1.2) + .build() + .unwrap(), + setup.clone(), + ), + search::Run::new( + graph::search::Range::builder(15, 2.0) + .initial_slack(0.8) + .range_slack(1.2) + .build() + .unwrap(), + setup.clone(), + ), + ]; + + let all = search::search_all(filtered_range, parameters, Aggregator::new(rows)).unwrap(); + + assert_eq!(all.len(), 2); + for summary in all { + assert_eq!(summary.setup, setup); + assert_eq!(summary.end_to_end_latencies.len(), TWO.get()); + assert_eq!(summary.mean_latencies.len(), TWO.get()); + assert_eq!(summary.p90_latencies.len(), TWO.get()); + assert_eq!(summary.p99_latencies.len(), TWO.get()); + + let ap = summary.average_precision; + assert_eq!(ap.num_queries, queries.nrows()); + assert_eq!( + ap.average_precision, 1.0, + "we used a search as the groundtruth" + ); + } + } + + #[test] + fn test_filtered_range_error() { + let index = search::graph::test_grid_provider(); + let queries = Arc::new(Matrix::new(0.0f32, 2, index.provider().dim())); + + let labels: Arc<[_]> = (0..queries.nrows() + 1) + .map(|_| -> Arc> { Arc::new(NoOdds {}) }) + .collect(); + + let strategy = provider::Strategy::new(); + + // Error for a mismatch between strategies and queries. + let err = FilteredRange::new( + index.clone(), + queries.clone(), + Strategy::collection([strategy.clone()]), + labels.clone(), + ) + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("1 strategy was provided when 2 were expected"), + "failed with {msg}" + ); + + // Error for a mismatch between label providers and queries. + let err = FilteredRange::new( + index, + queries.clone(), + Strategy::broadcast(strategy), + labels, + ) + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains( + "Number of label providers (3) must be equal to the number of queries (2)" + ), + "failed with {msg}" + ); + } +} diff --git a/diskann-benchmark-core/src/search/graph/mod.rs b/diskann-benchmark-core/src/search/graph/mod.rs index 8063f0875..5653fa6bf 100644 --- a/diskann-benchmark-core/src/search/graph/mod.rs +++ b/diskann-benchmark-core/src/search/graph/mod.rs @@ -3,6 +3,7 @@ * Licensed under the MIT license. */ +pub mod filtered_range; pub mod inline; pub mod knn; pub mod multihop; diff --git a/diskann-benchmark-core/src/search/graph/range.rs b/diskann-benchmark-core/src/search/graph/range.rs index edf95e29a..a9d94d511 100644 --- a/diskann-benchmark-core/src/search/graph/range.rs +++ b/diskann-benchmark-core/src/search/graph/range.rs @@ -263,7 +263,11 @@ mod tests { let rt = crate::tokio::runtime(2).unwrap(); let results = search::search( range.clone(), - graph::search::Range::with_options(None, 10, None, 2.0, None, 0.8, 1.2).unwrap(), + graph::search::Range::builder(10, 2.0) + .initial_slack(0.8) + .range_slack(1.2) + .build() + .unwrap(), NonZeroUsize::new(2).unwrap(), &rt, ) @@ -282,11 +286,19 @@ mod tests { // Try the aggregated strategy. let parameters = [ search::Run::new( - graph::search::Range::with_options(None, 10, None, 2.0, None, 0.8, 1.2).unwrap(), + graph::search::Range::builder(10, 2.0) + .initial_slack(0.8) + .range_slack(1.2) + .build() + .unwrap(), setup.clone(), ), search::Run::new( - graph::search::Range::with_options(None, 15, None, 2.0, None, 0.8, 1.2).unwrap(), + graph::search::Range::builder(15, 2.0) + .initial_slack(0.8) + .range_slack(1.2) + .build() + .unwrap(), setup.clone(), ), ]; diff --git a/diskann-benchmark-core/src/search/mod.rs b/diskann-benchmark-core/src/search/mod.rs index 36b4be3ad..b4fdcd2ac 100644 --- a/diskann-benchmark-core/src/search/mod.rs +++ b/diskann-benchmark-core/src/search/mod.rs @@ -123,6 +123,7 @@ //! //! * [`graph::search::Knn`]: K-nearest neighbors search for [`diskann::graph::DiskANNIndex`]. //! * [`graph::search::Range`]: Range search for [`diskann::graph::DiskANNIndex`]. +//! * [`graph::search::FilteredRange`]: Filtered range search for [`diskann::graph::DiskANNIndex`]. //! * [`graph::MultiHop`]: Multi-hop filtered search for [`diskann::graph::DiskANNIndex`]. pub(crate) mod ids; diff --git a/diskann-benchmark/example/graph-index-filter-range.json b/diskann-benchmark/example/graph-index-filter-range.json new file mode 100644 index 000000000..e1277be08 --- /dev/null +++ b/diskann-benchmark/example/graph-index-filter-range.json @@ -0,0 +1,54 @@ +{ + "search_directories": [ + "test_data/disk_index_search" + ], + "jobs": [ + { + "type": "graph-index-build", + "content": { + "source": { + "index-source": "Build", + "data_type": "float32", + "data": "disk_index_siftsmall_learn_256pts_data.fbin", + "distance": "squared_l2", + "max_degree": 32, + "l_build": 50, + "alpha": 1.2, + "backedge_ratio": 1.0, + "num_threads": 1, + "start_point_strategy": "medoid", + "num_insert_attempts": 1, + "saturate_inserts": false + }, + "search_phase": { + "search-type": "filtered-range", + "queries": "disk_index_sample_query_10pts.fbin", + "query_predicates": "query.10.label.jsonl", + "data_labels": "data.256.label.jsonl", + "groundtruth": "groundtruth_rad_125000_filtered.rangeres", + "reps": 5, + "num_threads": [ + 1 + ], + "runs": [ + { + "radius": 125000, + "inner_radius": null, + "max_returned": null, + "beam_width": null, + "initial_search_l": [ + 1, + 5, + 10, + 20, + 30 + ], + "initial_search_slack": 1.0, + "range_search_slack": 1.0 + } + ] + } + } + } + ] +} \ No newline at end of file diff --git a/diskann-benchmark/src/index/benchmarks.rs b/diskann-benchmark/src/index/benchmarks.rs index 4307f2e5b..21aefb8cf 100644 --- a/diskann-benchmark/src/index/benchmarks.rs +++ b/diskann-benchmark/src/index/benchmarks.rs @@ -76,6 +76,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> FullPrecision::::new() .search(plugins::Topk) .search(plugins::Range) + .search(plugins::FilteredRange) .search(plugins::TopkBetaFilter) .search(plugins::TopkMultihopFilter) .search(plugins::TopkInlineFilter) @@ -88,7 +89,10 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> )?; registry.register( "graph-index-full-precision-u8", - FullPrecision::::new().search(plugins::Topk), + FullPrecision::::new() + .search(plugins::Topk) + .search(plugins::TopkInlineFilter) + .search(plugins::FilteredRange), )?; registry.register( "graph-index-full-precision-i8", @@ -563,6 +567,72 @@ where } } +//-------------------// +// Filtered Range // +//-------------------// + +impl search::Plugin> for plugins::FilteredRange +where + DP: DataProvider + QueryType, + S: for<'a> glue::DefaultSearchStrategy< + 'a, + DP, + &'a [DP::Element], + SearchAccessor: glue::SearchAccessor, + > + Clone + + AsyncFriendly, +{ + fn is_match(&self, phase: &SearchPhase) -> bool { + plugins::FilteredRange::is_match(phase) + } + + fn kind(&self) -> &'static str { + plugins::FilteredRange::as_str() + } + + fn run( + &self, + index: Arc>, + phase: &SearchPhase, + strategy: &Strategy, + ) -> anyhow::Result { + let filtered_range = phase.as_filtered_range()?; + + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&filtered_range.queries), + )?); + + let groundtruth = + datafiles::load_range_groundtruth(datafiles::BinFile(&filtered_range.groundtruth))?; + + let steps = search::range::RangeSearchSteps::new( + filtered_range.reps, + &filtered_range.num_threads, + &filtered_range.runs, + ); + + let bit_maps = generate_bitmaps( + &filtered_range.query_predicates, + &filtered_range.data_labels, + )?; + + let labels: Arc<[_]> = bit_maps + .into_iter() + .map(utils::filters::as_query_label_provider) + .collect(); + + let filtered_range = benchmark_core::search::graph::filtered_range::FilteredRange::new( + index, + queries, + benchmark_core::search::graph::Strategy::broadcast(strategy.inner()), + labels, + )?; + + let result = search::range::run(&filtered_range, &groundtruth, steps)?; + Ok(AggregatedSearchResults::Range(result)) + } +} + //------------// // BetaFilter // //------------// diff --git a/diskann-benchmark/src/index/inmem/spherical.rs b/diskann-benchmark/src/index/inmem/spherical.rs index 5eadb03cc..438e205e9 100644 --- a/diskann-benchmark/src/index/inmem/spherical.rs +++ b/diskann-benchmark/src/index/inmem/spherical.rs @@ -26,6 +26,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> imp::SphericalQ::<1>::new() .search(plugins::Topk) .search(plugins::Range) + .search(plugins::FilteredRange) .search(plugins::TopkBetaFilter) .search(plugins::TopkMultihopFilter) .search(plugins::TopkInlineFilter), @@ -36,6 +37,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> imp::SphericalQ::<2>::new() .search(plugins::Topk) .search(plugins::Range) + .search(plugins::FilteredRange) .search(plugins::TopkBetaFilter) .search(plugins::TopkMultihopFilter) .search(plugins::TopkInlineFilter), @@ -46,6 +48,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> imp::SphericalQ::<4>::new() .search(plugins::Topk) .search(plugins::Range) + .search(plugins::FilteredRange) .search(plugins::TopkBetaFilter) .search(plugins::TopkMultihopFilter) .search(plugins::TopkInlineFilter), @@ -428,6 +431,63 @@ mod imp { } } + impl search::plugins::Plugin + for search::plugins::FilteredRange + { + fn is_match(&self, phase: &SearchPhase) -> bool { + search::plugins::FilteredRange::is_match(phase) + } + + fn kind(&self) -> &'static str { + SearchPhaseKind::FilteredRange.as_str() + } + + fn run( + &self, + index: Arc>, + phase: &SearchPhase, + query_layout: &exhaustive::SphericalQuery, + ) -> anyhow::Result { + let filtered_range = phase.as_filtered_range()?; + + let queries: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile( + &filtered_range.queries, + ))?); + + let groundtruth = + datafiles::load_range_groundtruth(datafiles::BinFile(&filtered_range.groundtruth))?; + + let steps = search::range::RangeSearchSteps::new( + filtered_range.reps, + &filtered_range.num_threads, + &filtered_range.runs, + ); + + let bit_maps = generate_bitmaps( + &filtered_range.query_predicates, + &filtered_range.data_labels, + )?; + + let labels: Arc<[_]> = bit_maps + .into_iter() + .map(utils::filters::as_query_label_provider) + .collect(); + + let filtered_range = benchmark_core::search::graph::filtered_range::FilteredRange::new( + index.clone(), + queries.clone(), + benchmark_core::search::graph::Strategy::broadcast( + inmem::spherical::Quantized::search((*query_layout).into()), + ), + labels, + )?; + + let result = search::range::run(&filtered_range, &groundtruth, steps)?; + + Ok(AggregatedSearchResults::Range(result)) + } + } + impl search::plugins::Plugin for search::plugins::TopkBetaFilter { diff --git a/diskann-benchmark/src/index/result.rs b/diskann-benchmark/src/index/result.rs index 8bd90b33e..070c3df98 100644 --- a/diskann-benchmark/src/index/result.rs +++ b/diskann-benchmark/src/index/result.rs @@ -269,6 +269,35 @@ impl RangeSearchResults { average_precision: (&average_precision).into(), } } + + pub fn new_filtered(summary: benchmark_core::search::graph::filtered_range::Summary) -> Self { + let benchmark_core::search::graph::filtered_range::Summary { + setup, + parameters, + end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + average_precision, + .. + } = summary; + + let qps = end_to_end_latencies + .iter() + .map(|latency| average_precision.num_queries as f64 / latency.as_seconds()) + .collect(); + + Self { + num_tasks: setup.tasks.into(), + initial_l: parameters.starting_l(), + qps, + search_latencies: end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + average_precision: (&average_precision).into(), + } + } } fn format_range_search_results_table( diff --git a/diskann-benchmark/src/index/search/plugins.rs b/diskann-benchmark/src/index/search/plugins.rs index de050004f..2caf73406 100644 --- a/diskann-benchmark/src/index/search/plugins.rs +++ b/diskann-benchmark/src/index/search/plugins.rs @@ -26,10 +26,10 @@ //! * [`Plugins::is_match`]: check whether any registered plugin accepts a requested `Kind`. //! * [`Plugins::run`]: dispatch to the first registered plugin matching `Kind`. //! -//! The built-in ZST plugins in this module (`Topk`, `Range`, `BetaFilter`, and -//! `MultihopFilter`) target the async benchmark inputs and fold their outputs into the closed -//! [`AggregatedSearchResults`] families. That closed result boundary is deliberate: plugins are -//! open for new search flavors, while result aggregation remains a curated +//! The built-in ZST plugins in this module (`Topk`, `Range`, `FilteredRange`, `BetaFilter`, +//! and `MultihopFilter`) target the async benchmark inputs and fold their outputs into the +//! closed [`AggregatedSearchResults`] families. That closed result boundary is deliberate: +//! plugins are open for new search flavors, while result aggregation remains a curated //! reporting/evaluation boundary. use std::sync::Arc; @@ -192,6 +192,20 @@ impl Range { } } +/// A search plugin for filtered range search. +#[derive(Debug, Clone, Copy)] +pub(crate) struct FilteredRange; + +impl FilteredRange { + pub(crate) fn is_match(phase: &SearchPhase) -> bool { + phase.as_filtered_range().is_ok() + } + + pub(crate) const fn as_str() -> &'static str { + "filtered-range" + } +} + /// A search plugin for beta-filtered search. #[derive(Debug, Clone, Copy)] pub(crate) struct TopkBetaFilter; diff --git a/diskann-benchmark/src/index/search/range.rs b/diskann-benchmark/src/index/search/range.rs index 9224f69e9..17e7d68ac 100644 --- a/diskann-benchmark/src/index/search/range.rs +++ b/diskann-benchmark/src/index/search/range.rs @@ -5,6 +5,7 @@ use std::{num::NonZeroUsize, sync::Arc}; +use diskann::graph::search::Range as RangeParameters; use diskann_benchmark_core::{self as benchmark_core, search as core_search}; use crate::{index::result::RangeSearchResults, inputs::graph_index::GraphRangeSearch}; @@ -30,18 +31,18 @@ impl<'a> RangeSearchSteps<'a> { } } -type Run = core_search::Run; - pub(crate) trait Range { + type Parameters: Clone; + fn search_all( &self, - parameters: Vec, + parameters: Vec>, groundtruth: &dyn benchmark_core::recall::Rows, ) -> anyhow::Result>; } pub(crate) fn run( - runner: &dyn Range, + runner: &dyn Range, groundtruth: &dyn benchmark_core::recall::Rows, steps: RangeSearchSteps<'_>, ) -> anyhow::Result> { @@ -69,7 +70,6 @@ pub(crate) fn run( Ok(all) } - /////////// // Impls // /////////// @@ -79,13 +79,15 @@ where DP: diskann::provider::DataProvider, core_search::graph::Range: core_search::Search< Id = DP::InternalId, - Parameters = diskann::graph::search::Range, + Parameters = RangeParameters, Output = core_search::graph::range::Metrics, >, { + type Parameters = RangeParameters; + fn search_all( &self, - parameters: Vec>, + parameters: Vec>, groundtruth: &dyn benchmark_core::recall::Rows, ) -> anyhow::Result> { let results = core_search::search_all( @@ -97,3 +99,33 @@ where Ok(results.into_iter().map(RangeSearchResults::new).collect()) } } + +impl Range + for Arc> +where + DP: diskann::provider::DataProvider, + core_search::graph::filtered_range::FilteredRange: core_search::Search< + Id = DP::InternalId, + Parameters = RangeParameters, + Output = core_search::graph::filtered_range::Metrics, + >, +{ + type Parameters = RangeParameters; + + fn search_all( + &self, + parameters: Vec>, + groundtruth: &dyn benchmark_core::recall::Rows, + ) -> anyhow::Result> { + let results = core_search::search_all( + self.clone(), + parameters.into_iter(), + core_search::graph::filtered_range::Aggregator::new(groundtruth), + )?; + + Ok(results + .into_iter() + .map(RangeSearchResults::new_filtered) + .collect()) + } +} diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index 2137c6e57..09a7a28ec 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -81,15 +81,13 @@ impl GraphRangeSearch { self.initial_search_l .iter() .map(|&l| { - Range::with_options( - self.max_returned, - l, - self.beam_width, - self.radius, - self.inner_radius, - self.initial_search_slack, - self.range_search_slack, - ) + Range::builder(l, self.radius) + .max_returned(self.max_returned) + .beam_width(self.beam_width) + .inner_radius(self.inner_radius) + .initial_slack(self.initial_search_slack) + .range_slack(self.range_search_slack) + .build() }) .collect() } @@ -168,6 +166,33 @@ pub(crate) struct RangeSearchPhase { pub(crate) runs: Vec, } +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct FilteredRangeSearchPhase { + pub(crate) queries: InputFile, + pub(crate) query_predicates: InputFile, + pub(crate) data_labels: InputFile, + pub(crate) groundtruth: InputFile, + pub(crate) reps: NonZeroUsize, + // Enable sweeping threads + pub(crate) num_threads: Vec, + pub(crate) runs: Vec, +} + +impl FilteredRangeSearchPhase { + pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + self.queries.resolve(checker)?; + self.query_predicates.resolve(checker)?; + self.data_labels.resolve(checker)?; + self.groundtruth.resolve(checker)?; + for (i, run) in self.runs.iter_mut().enumerate() { + run.validate(checker) + .with_context(|| format!("search run {}", i))?; + } + + Ok(()) + } +} + impl RangeSearchPhase { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; @@ -406,6 +431,7 @@ impl Example for TopkDeterminantDiversityPhase { pub(crate) enum SearchPhase { Topk(TopkSearchPhase), Range(RangeSearchPhase), + FilteredRange(FilteredRangeSearchPhase), TopkBetaFilter(BetaSearchPhase), TopkMultihopFilter(MultihopFilterSearchPhase), TopkInlineFilter(InlineFilterSearchPhase), @@ -434,6 +460,7 @@ impl SearchPhase { match self { Self::Topk(_) => SearchPhaseKind::Topk, Self::Range(_) => SearchPhaseKind::Range, + Self::FilteredRange(_) => SearchPhaseKind::FilteredRange, Self::TopkBetaFilter(_) => SearchPhaseKind::TopkBetaFilter, Self::TopkMultihopFilter(_) => SearchPhaseKind::TopkMultihopFilter, Self::TopkInlineFilter(_) => SearchPhaseKind::TopkInlineFilter, @@ -461,6 +488,18 @@ impl SearchPhase { } } + pub(crate) fn as_filtered_range( + &self, + ) -> Result<&FilteredRangeSearchPhase, WrongSearchPhaseKind> { + match self { + Self::FilteredRange(phase) => Ok(phase), + _ => Err(WrongSearchPhaseKind::new( + SearchPhaseKind::FilteredRange, + self.kind(), + )), + } + } + pub(crate) fn as_topk_beta_filter(&self) -> Result<&BetaSearchPhase, WrongSearchPhaseKind> { match self { Self::TopkBetaFilter(phase) => Ok(phase), @@ -513,6 +552,7 @@ impl SearchPhase { match self { SearchPhase::Topk(phase) => phase.validate(checker), SearchPhase::Range(phase) => phase.validate(checker), + SearchPhase::FilteredRange(phase) => phase.validate(checker), SearchPhase::TopkBetaFilter(phase) => phase.validate(checker), SearchPhase::TopkMultihopFilter(phase) => phase.validate(checker), SearchPhase::TopkInlineFilter(phase) => phase.validate(checker), @@ -525,6 +565,7 @@ impl SearchPhase { pub(crate) enum SearchPhaseKind { Topk, Range, + FilteredRange, TopkBetaFilter, TopkMultihopFilter, TopkInlineFilter, @@ -536,6 +577,7 @@ impl SearchPhaseKind { match self { Self::Topk => "topk", Self::Range => "range", + Self::FilteredRange => "filtered-range", Self::TopkBetaFilter => "topk-beta-filter", Self::TopkMultihopFilter => "topk-multihop-filter", Self::TopkInlineFilter => "topk-inline-filter", diff --git a/diskann-benchmark/src/main.rs b/diskann-benchmark/src/main.rs index 55f6bd019..d4f108e67 100644 --- a/diskann-benchmark/src/main.rs +++ b/diskann-benchmark/src/main.rs @@ -700,6 +700,22 @@ mod tests { run_integration_test(raw); } + #[test] + fn graph_index_range_integration() { + // First, parse and modify the input file to establish paths relative to the + // directory building the dispatcher. + let raw = value_from_file(&example_directory().join("graph-index-range.json")); + run_integration_test(raw); + } + + #[test] + fn graph_index_filter_range_integration() { + // First, parse and modify the input file to establish paths relative to the + // directory building the dispatcher. + let raw = value_from_file(&example_directory().join("graph-index-filter-range.json")); + run_integration_test(raw); + } + /// Filtered disk search end-to-end: drives the disk-index backend through /// `disk-index-filter.json` #[test] diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index 41389910e..5c1d246c0 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -1549,16 +1549,10 @@ pub(crate) mod tests { // Test with an inner radius assert!(inner_radius <= radius); - let range_search = Range::with_options( - None, - starting_l_value, - None, - radius, - Some(inner_radius), - 1.0, - 1.0, - ) - .unwrap(); + let range_search = Range::builder(starting_l_value, radius) + .inner_radius(Some(inner_radius)) + .build() + .unwrap(); let mut results: Vec> = Vec::new(); let _ = index .search(range_search, &FullPrecision, ctx, query, &mut results) diff --git a/diskann/src/graph/search/filtered_range_search.rs b/diskann/src/graph/search/filtered_range_search.rs new file mode 100644 index 000000000..ed5a929d4 --- /dev/null +++ b/diskann/src/graph/search/filtered_range_search.rs @@ -0,0 +1,358 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Filtered range-based search with a distance radius + +use diskann_utils::future::SendFuture; + +use super::{Knn, Search, scratch::SearchScratch}; +use crate::{ + ANNResult, + error::IntoANNResult, + graph::{ + glue::{self, FilteredAccessor, SearchStrategy}, + index::{DiskANNIndex, InternalSearchStats, SearchStats}, + search::inline_filter_search::{Ret, inline_filter_search_internal}, + search::{Range, RangeSearchError, range_search::RangeBuilder, record::NoopSearchRecord}, + search_output_buffer::SearchOutputBuffer, + }, + neighbor::Neighbor, + provider::DataProvider, +}; + +use super::range_search::DistanceFiltered; + +/// Parameters for range-based search. +/// +/// Finds all points within a specified distance radius from the query. +#[derive(Debug, Clone, Copy)] +pub struct FilteredRange { + range_params: Range, +} + +impl FilteredRange { + /// Create range search with default slack values. + pub fn new(starting_l: usize, radius: f32) -> Result { + Self::builder(starting_l, radius).build_filtered() + } + + /// Create a builder for filtered range search parameters. + /// + /// The builder starts with the same defaults as [`Self::new`]. + pub fn builder(starting_l: usize, radius: f32) -> RangeBuilder { + Range::builder(starting_l, radius) + } + + /// Returns the maximum number of results to return. + #[inline] + pub fn max_returned(&self) -> Option { + self.range_params.max_returned() + } + + /// Returns the initial search list size. + #[inline] + pub fn starting_l(&self) -> usize { + self.range_params.starting_l() + } + + /// Returns the optional beam width. + #[inline] + pub fn beam_width(&self) -> Option { + self.range_params.beam_width() + } + + /// Returns the outer radius. + #[inline] + pub fn radius(&self) -> f32 { + self.range_params.radius() + } + + /// Returns the inner radius (points closer are excluded). + #[inline] + pub fn inner_radius(&self) -> Option { + self.range_params.inner_radius() + } + + /// Returns the initial search slack factor. + #[inline] + pub fn initial_slack(&self) -> f32 { + self.range_params.initial_slack() + } + + /// Returns the range search slack factor. + #[inline] + pub fn range_slack(&self) -> f32 { + self.range_params.range_slack() + } + + /// Returns the underlying range search parameters. + #[inline] + pub fn range(&self) -> Range { + self.range_params + } +} + +impl RangeBuilder { + /// Build validated [`FilteredRange`] parameters. + pub fn build_filtered(self) -> Result { + let range_params = self.build()?; + Ok(FilteredRange { range_params }) + } +} + +impl<'a, DP, S, T> Search<'a, DP, S, T> for FilteredRange +where + DP: DataProvider, + S: SearchStrategy<'a, DP, T, SearchAccessor: FilteredAccessor>, + T: Copy + Send + Sync, +{ + type Output = SearchStats; + + fn search( + self, + index: &'a DiskANNIndex, + strategy: &'a S, + processor: PP, + context: &'a DP::Context, + query: T, + output: &mut OB, + ) -> impl SendFuture> + where + O: Send, + PP: glue::SearchPostProcess + Send + Sync, + OB: SearchOutputBuffer + Send + ?Sized, + { + async move { + let mut accessor = strategy + .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(), num_start_ids); + + // Perform an initial inline filtered search, store both filtered and unfiltered results + + let search_knn = Knn::new(self.starting_l(), self.starting_l(), self.beam_width())?; + + let Ret { + cmps, + hops, + matched_results, + } = inline_filter_search_internal( + index.max_degree_with_slack(), + &search_knn, + &mut accessor, + &mut scratch, + &mut NoopSearchRecord::new(), + None, + ) + .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()) + .chain(matched_results.iter().copied()) + .filter(|neighbor| neighbor.distance <= self.radius()) + .collect(); + + in_range.sort_unstable_by(|left, right| { + left.id + .cmp(&right.id) + .then_with(|| left.distance.total_cmp(&right.distance)) + }); + in_range.dedup_by_key(|neighbor| neighbor.id); + + in_range.sort_unstable_by(|left, right| { + left.distance + .total_cmp(&right.distance) + .then_with(|| left.id.cmp(&right.id)) + }); + + 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); + } + } + + // clear the visited set and repopulate it with all in-range points found so far, filtered and unfiltered + scratch.visited.clear(); + scratch.range_frontier.clear(); + for neighbor in in_range.iter() { + scratch.visited.insert(neighbor.id); + scratch.range_frontier.push_back(neighbor.id); + } + + let stats = if in_range.len() + >= ((self.starting_l() as f32) * self.initial_slack()) as usize + && matched_within_radius.len() < max_returned + { + // Move to filtered range search + let range_stats = filtered_range_search_internal( + index.max_degree_with_slack(), + &self, + &mut accessor, + &mut scratch, + &mut matched_within_radius, + ) + .await?; + + InternalSearchStats { + cmps, + hops: hops + range_stats.hops, + range_search_second_round: true, + } + } else { + InternalSearchStats { + cmps, + hops, + range_search_second_round: false, + } + }; + + // 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 truncated_matched = matched_within_radius.iter().copied().take(max_returned); + + let result_count = processor + .post_process(&mut accessor, query, truncated_matched, &mut filtered) + .await + .into_ann_result()?; + + Ok(SearchStats { + cmps: stats.cmps, + hops: stats.hops, + result_count: result_count as u32, + range_search_second_round: stats.range_search_second_round, + }) + } + } +} + +///////////////////////////// +// Internal Implementation // +///////////////////////////// + +/// Internal range search implementation. +/// +/// Expands the search frontier to find all points within the specified radius. +/// Called after the initial graph search has identified starting candidates. +pub(crate) async fn filtered_range_search_internal( + max_degree_with_slack: usize, + search_params: &FilteredRange, + accessor: &mut A, + scratch: &mut SearchScratch, + matched_in_range: &mut Vec>, +) -> ANNResult +where + A: FilteredAccessor, +{ + let beam_width = search_params.beam_width().unwrap_or(1); + + 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() && 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(); + if let Some(next_node) = next { + scratch.beam_nodes.push(next_node); + } + } + + neighbors.clear(); + accessor + .expand_beam_filtered( + scratch.beam_nodes.iter().copied(), + glue::NotInMut::new(&mut scratch.visited), + |id, distance| neighbors.push((id, distance)), + ) + .await?; + + // The predicate ensures that the contents of `neighbors` are unique. + // We still traverse both accepted and rejected IDs via frontier expansion, + // but only accepted IDs are added to in-range results. + for (decision, distance) in neighbors.iter().copied() { + if distance <= search_params.radius() * search_params.range_slack() { + let is_accept = decision.is_accept(); + let id = decision.into_inner(); + + scratch.range_frontier.push_back(id); + + if is_accept && matched_in_range.len() < max_returned { + matched_in_range.push(Neighbor::new(id, distance)); + } + } + } + scratch.cmps += neighbors.len() as u32; + scratch.hops += scratch.beam_nodes.len() as u32; + } + + Ok(InternalSearchStats { + cmps: scratch.cmps, + hops: scratch.hops, + range_search_second_round: true, + }) +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_range_search_validation() { + // Valid + assert!(FilteredRange::new(100, 0.5).is_ok()); + + // Invalid: zero l + assert!(FilteredRange::new(0, 0.5).is_err()); + + // Invalid slack values + assert!( + FilteredRange::builder(100, 0.5) + .initial_slack(1.5) + .build_filtered() + .is_err() + ); + assert!( + FilteredRange::builder(100, 0.5) + .range_slack(0.5) + .build_filtered() + .is_err() + ); + + // Invalid inner radius > radius + assert!( + FilteredRange::builder(100, 0.5) + .inner_radius(Some(1.0)) + .build_filtered() + .is_err() + ); + } +} diff --git a/diskann/src/graph/search/inline_filter_search.rs b/diskann/src/graph/search/inline_filter_search.rs index c535c87b6..4f7747c9d 100644 --- a/diskann/src/graph/search/inline_filter_search.rs +++ b/diskann/src/graph/search/inline_filter_search.rs @@ -159,16 +159,16 @@ where } #[derive(Debug)] -struct Ret +pub(crate) struct Ret where I: Eq, { - cmps: u32, - hops: u32, - matched_results: Vec>, + pub(crate) cmps: u32, + pub(crate) hops: u32, + pub(crate) matched_results: Vec>, } -async fn inline_filter_search_internal( +pub(crate) async fn inline_filter_search_internal( max_degree_with_slack: usize, search_params: &Knn, accessor: &mut A, diff --git a/diskann/src/graph/search/mod.rs b/diskann/src/graph/search/mod.rs index 4e2aedc61..55867486e 100644 --- a/diskann/src/graph/search/mod.rs +++ b/diskann/src/graph/search/mod.rs @@ -42,6 +42,7 @@ use crate::{ provider::DataProvider, }; +mod filtered_range_search; mod inline_filter_search; mod knn_search; pub(crate) mod multihop_filter_search; @@ -63,6 +64,7 @@ pub(crate) mod scratch; /// See the specific search types for detailed documentation: /// - [`Knn`] - Standard k-nearest neighbor search /// - [`Range`] - Range-based search within a distance radius +/// - [`FilteredRange`] - Filtered range search /// - [`Diverse`] - Diversity-aware search (feature-gated) /// - [`MultihopFilterSearch`] - Label-filtered search with multi-hop expansion /// - [`InlineFilterSearch`] - Inline filtered search with optional adaptive L sizing @@ -112,6 +114,7 @@ where OB: graph::search_output_buffer::SearchOutputBuffer + Send + ?Sized; } +pub use filtered_range_search::FilteredRange; pub use inline_filter_search::{AdaptiveL, InlineFilterSearch}; pub use knn_search::{Knn, KnnSearchError, RecordedKnn}; pub use multihop_filter_search::MultihopFilterSearch; diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 01c859ca5..431d28e5b 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -68,15 +68,41 @@ pub struct Range { range_slack: f32, } +/// Builder for [`Range`] search parameters. +#[derive(Debug, Clone, Copy)] +pub struct RangeBuilder { + max_returned: Option, + starting_l: usize, + beam_width: Option, + radius: f32, + inner_radius: Option, + initial_slack: f32, + range_slack: f32, +} + impl Range { /// Create range search with default slack values. pub fn new(starting_l: usize, radius: f32) -> Result { - Self::with_options(None, starting_l, None, radius, None, 1.0, 1.0) + Self::builder(starting_l, radius).build() + } + + /// Create a builder for range search parameters. + /// + /// The builder starts with the same defaults as [`Self::new`]. + pub fn builder(starting_l: usize, radius: f32) -> RangeBuilder { + RangeBuilder { + max_returned: None, + starting_l, + beam_width: None, + radius, + inner_radius: None, + initial_slack: 1.0, + range_slack: 1.0, + } } - /// Create range search with full options. #[allow(clippy::too_many_arguments)] - pub fn with_options( + fn validate_and_create( max_returned: Option, starting_l: usize, beam_width: Option, @@ -164,6 +190,51 @@ impl Range { } } +impl RangeBuilder { + /// Set maximum results to return (`None` means unlimited). + pub fn max_returned(mut self, value: Option) -> Self { + self.max_returned = value; + self + } + + /// Set the beam width. + pub fn beam_width(mut self, value: Option) -> Self { + self.beam_width = value; + self + } + + /// Set the inner radius. + pub fn inner_radius(mut self, value: Option) -> Self { + self.inner_radius = value; + self + } + + /// Set the initial-search slack factor. + pub fn initial_slack(mut self, value: f32) -> Self { + self.initial_slack = value; + self + } + + /// Set the range-search slack factor. + pub fn range_slack(mut self, value: f32) -> Self { + self.range_slack = value; + self + } + + /// Build validated [`Range`] parameters. + pub fn build(self) -> Result { + Range::validate_and_create( + self.max_returned, + self.starting_l, + self.beam_width, + self.radius, + self.inner_radius, + self.initial_slack, + self.range_slack, + ) + } +} + impl<'a, DP, S, T> Search<'a, DP, S, T> for Range where DP: DataProvider, @@ -276,13 +347,13 @@ where /// A [`SearchOutputBuffer`] wrapper that filters results by distance before /// forwarding them to an inner buffer. -struct DistanceFiltered<'a, F, B: ?Sized> { +pub(crate) struct DistanceFiltered<'a, F, B: ?Sized> { predicate: F, inner: &'a mut B, } impl<'a, F, B: ?Sized> DistanceFiltered<'a, F, B> { - fn new(inner: &'a mut B, predicate: F) -> Self { + pub(crate) fn new(inner: &'a mut B, predicate: F) -> Self { Self { predicate, inner } } } @@ -398,6 +469,49 @@ mod tests { use crate::graph::search_output_buffer::BufferState; use crate::neighbor::Neighbor; + #[test] + fn range_builder_defaults_match_new() { + let from_new = Range::new(100, 0.5).unwrap(); + let from_builder = Range::builder(100, 0.5).build().unwrap(); + + assert_eq!(from_builder.max_returned(), from_new.max_returned()); + assert_eq!(from_builder.starting_l(), from_new.starting_l()); + assert_eq!(from_builder.beam_width(), from_new.beam_width()); + assert_eq!(from_builder.radius(), from_new.radius()); + assert_eq!(from_builder.inner_radius(), from_new.inner_radius()); + assert_eq!(from_builder.initial_slack(), from_new.initial_slack()); + assert_eq!(from_builder.range_slack(), from_new.range_slack()); + } + + #[test] + fn range_builder_custom_options_match_expected_values() { + let built = Range::builder(100, 0.8) + .max_returned(Some(10)) + .beam_width(Some(8)) + .inner_radius(Some(0.3)) + .initial_slack(0.9) + .range_slack(1.2) + .build() + .unwrap(); + + assert_eq!(built.max_returned(), Some(10)); + assert_eq!(built.starting_l(), 100); + assert_eq!(built.beam_width(), Some(8)); + assert_eq!(built.radius(), 0.8); + assert_eq!(built.inner_radius(), Some(0.3)); + assert_eq!(built.initial_slack(), 0.9); + assert_eq!(built.range_slack(), 1.2); + } + + #[test] + fn range_builder_validation_error() { + let err = Range::builder(100, 0.5) + .beam_width(Some(0)) + .build() + .unwrap_err(); + assert!(matches!(err, RangeSearchError::BeamWidthZero)); + } + #[test] fn test_range_search_validation() { // Valid @@ -407,14 +521,23 @@ mod tests { assert!(Range::new(0, 0.5).is_err()); // Invalid slack values - assert!(Range::with_options(None, 100, None, 0.5, None, 1.5, 1.0).is_err()); - assert!(Range::with_options(None, 100, None, 0.5, None, 1.0, 0.5).is_err()); + assert!(Range::builder(100, 0.5).initial_slack(1.5).build().is_err()); + assert!(Range::builder(100, 0.5).range_slack(0.5).build().is_err()); // 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()); + assert!( + Range::builder(100, 0.5) + .inner_radius(Some(1.0)) + .build() + .is_err() + ); + + assert!( + Range::builder(100, 0.5) + .max_returned(Some(1.0)) + .build() + .is_err() + ); } #[test] diff --git a/diskann/src/graph/test/cases/filtered_range_search.rs b/diskann/src/graph/test/cases/filtered_range_search.rs new file mode 100644 index 000000000..b9c459f08 --- /dev/null +++ b/diskann/src/graph/test/cases/filtered_range_search.rs @@ -0,0 +1,484 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Tests for filtered range search using an always-true filter. +//! +//! These cover the filtered-range cases directly and validate the filtered-range +//! behavior against its own baselines. + +use std::sync::Arc; + +use super::range_search::{ + RangeSearchBaseline, assert_no_duplicates, assert_range_invariants, + setup_grid_index_and_default_query, +}; +use crate::{ + graph::{ + self, DiskANNIndex, + ext::labeled, + search::FilteredRange, + test::{provider as test_provider, synthetic::Grid}, + }, + neighbor::Neighbor, + test::{ + TestRoot, cmp::assert_eq_verbose, get_or_save_test_results, tokio::current_thread_runtime, + }, +}; + +fn root() -> TestRoot { + TestRoot::new("graph/test/cases/filtered_range_search") +} + +#[derive(Debug)] +struct AlwaysTrueFilter; + +impl labeled::QueryLabelProvider for AlwaysTrueFilter { + fn is_match(&self, _: u32) -> bool { + true + } +} + +fn run_filtered_range_search( + index: &Arc>, + query: &[f32], + filtered_range: FilteredRange, + filter: &impl labeled::QueryLabelProvider, +) -> (graph::index::SearchStats, Vec>) { + let rt = current_thread_runtime(); + let mut results: Vec> = Vec::new(); + + let stats = rt + .block_on(index.search( + filtered_range, + &labeled::Filtered::new(test_provider::Strategy::new(), filter), + &test_provider::Context::new(), + query, + &mut results, + )) + .unwrap(); + + (stats, results) +} + +#[derive(Debug)] +struct DivisibleByFourFilter; + +impl labeled::QueryLabelProvider for DivisibleByFourFilter { + fn is_match(&self, id: u32) -> bool { + id.is_multiple_of(4) + } +} + +fn assert_divisible_by_four(results: &[Neighbor]) { + for n in results { + assert_eq!( + n.id % 4, + 0, + "result id {} does not satisfy divisible-by-4 filter", + n.id + ); + } +} + +/////////////////////////////////////////////////////// +// Tests with an always-true filter that validate // +// that filtered range search gives correct answers // +// and respects parameters such as inner_radius and // +// max_results. Note that identical behavior to the // +// unfiltered range search is NOT expected even when // +// using an always-true filter, because the filtered // +// initial search returns every predicate-satisfying // +// point it finds as opposed to only those in // +// `scratch.best`. // +/////////////////////////////////////////////////////// + +#[test] +fn basic_range_search() { + let description = "Basic range search test with an always-true filter. Validates / + that the filtered range search returns correct results without any duplicates."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("basic_range_search"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 12.0; + let starting_l = 32; + let filter = AlwaysTrueFilter; + let filtered_range = FilteredRange::new(starting_l, radius).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_size, + description, + query.clone(), + ); + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); +} + +#[test] +fn inner_radius_filtering() { + let description = "Inner radius filtering test to validate that the \ + range search correctly excludes neighbors within the inner radius."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("inner_radius_filtering"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 20.0; + let inner_radius = 6.0; + let starting_l = 32; + let filter = AlwaysTrueFilter; + + let filtered_range = FilteredRange::builder(starting_l, radius) + .inner_radius(Some(inner_radius)) + .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_size, + description, + query.clone(), + ); + + let expected = get_or_save_test_results(&name, &baseline); + assert_eq_verbose!(expected, baseline); + + assert_range_invariants(&filtered_results, radius, Some(inner_radius)); + assert_no_duplicates(&filtered_results); +} + +#[test] +fn two_round_search() { + let description = "Two round search test to validate that a / + low starting L with a large radius triggers a second round / + of range search."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("two_round_search"); + + 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::new(starting_l, radius).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_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 starting_l with large radius should trigger a second round" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); +} + +#[test] +fn empty_results() { + let rt = current_thread_runtime(); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 0.01; + let starting_l = 32; + let filter = AlwaysTrueFilter; + + let filtered_range = FilteredRange::new(starting_l, radius).unwrap(); + let mut filtered_results: Vec> = Vec::new(); + + let filtered_stats = rt + .block_on(index.search( + filtered_range, + &labeled::Filtered::new(test_provider::Strategy::new(), &filter), + &test_provider::Context::new(), + query.as_slice(), + &mut filtered_results, + )) + .unwrap(); + + assert!( + filtered_results.is_empty(), + "no points should be within the radius {}", + radius + ); + assert!( + !filtered_stats.range_search_second_round, + "empty results shouldn't trigger a second round" + ); +} + +#[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_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` / + below `max_results` so that a second round of search is triggered. Note that / + the result set size is expected to be below `max_results` due to filtering / + out start points during `post_process`."; + + 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; + let starting_l = 4; + let max_results = 200; + 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_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 be triggered" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); +} + +////////////////////////////////////////////////////////////// +// Next, some tests with a somewhat more complex filter. // +// We check that second round behavior is as expected, and // +// that the filter is applied correctly. We also check that // +// putting a cap on the returned results correctly uses the // +// filter-satisfying results, not all results. // +////////////////////////////////////////////////////////////// + +#[test] +fn divisible_by_four_filter_second_round_triggered() { + let description = "Test that a small starting L triggers a second / + round of search when using the divisible-by-4 filter. Tests for / + no duplicates, respect for the radius, and respect for the filter."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("divisible_by_four_filter_second_round_triggered"); + + 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 = DivisibleByFourFilter; + + let filtered_range = FilteredRange::new(starting_l, radius).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_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, + "small starting_l should trigger a second round with divisible-by-4 filter" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); + assert_divisible_by_four(&filtered_results); +} + +#[test] +fn divisible_by_four_filter_no_second_round_from_l_search() { + let description = "Test of divisible-by-4 filter with a larger starting L. \ + Since only 1/4 of discovered points will satisfy the filter, the second \ + round of search should not be triggered."; + + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("divisible_by_four_filter_no_second_round_from_l_search"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 12.0; + let starting_l = 32; + let filter = DivisibleByFourFilter; + + let filtered_range = FilteredRange::new(starting_l, radius).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_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, + "larger starting_l should avoid a second round with divisible-by-4 filter" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); + assert_divisible_by_four(&filtered_results); +} + +#[test] +fn divisible_by_four_filter_no_second_round_from_max_results() { + let description = "A test of the divisible-by-4 filter with `max_results` / + set to starting_l / 4. The second round of search should not be triggered / + since the correct number of results should be discovered."; + let mut test_root = root(); + let mut path = test_root.path(); + let name = path.push("divisible_by_four_filter_no_second_round_from_max_results"); + + let grid_size = 5; + let (index, query) = setup_grid_index_and_default_query(grid_size, Grid::Three); + let radius = 50.0; + let starting_l = 16; + let max_results = starting_l / 4; + let filter = DivisibleByFourFilter; + + 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_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, + "max_results = starting_l / 4 should prevent a second round" + ); + assert_range_invariants(&filtered_results, radius, None); + assert_no_duplicates(&filtered_results); + assert_divisible_by_four(&filtered_results); +} diff --git a/diskann/src/graph/test/cases/mod.rs b/diskann/src/graph/test/cases/mod.rs index 69f502d1f..37857e5c8 100644 --- a/diskann/src/graph/test/cases/mod.rs +++ b/diskann/src/graph/test/cases/mod.rs @@ -4,6 +4,7 @@ */ mod consolidate; +mod filtered_range_search; mod grid_insert; mod grid_search; mod helpers; diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index 6e2f4dce2..b75de22ef 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -16,6 +16,7 @@ use diskann_vector::distance::Metric; use crate::{ graph::{ self, DiskANNIndex, + index::SearchStats, search::Range, test::{provider as test_provider, synthetic::Grid}, }, @@ -28,11 +29,71 @@ use crate::{ }, }; +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub(super) struct RangeSearchBaseline { + /// A description of what to expect, what trends to observe, and anything else + /// a reviewer may need to either understand why this test is checked in or to validate + /// any changes that occur in the checked-in file. + pub(super) description: String, + pub(super) grid_size: usize, + pub(super) query: Vec, + pub(super) radius: f32, + pub(super) inner_radius: Option, + pub(super) starting_l: usize, + pub(super) results: Vec<(u32, f32)>, + pub(super) comparisons: usize, + pub(super) hops: usize, + pub(super) result_count: usize, + pub(super) range_search_second_round: bool, +} + +impl RangeSearchBaseline { + pub(super) fn new( + range: &Range, + results: &[Neighbor], + stats: SearchStats, + grid_size: usize, + description: impl Into, + query: Vec, + ) -> Self { + Self { + description: description.into(), + grid_size, + query, + radius: range.radius(), + inner_radius: range.inner_radius(), + starting_l: range.starting_l(), + results: results.iter().map(|n| (n.id, n.distance)).collect(), + comparisons: stats.cmps as usize, + hops: stats.hops as usize, + result_count: stats.result_count as usize, + range_search_second_round: stats.range_search_second_round, + } + } +} + +verbose_eq!(RangeSearchBaseline { + description, + grid_size, + query, + radius, + inner_radius, + starting_l, + results, + comparisons, + hops, + result_count, + range_search_second_round, +}); + fn root() -> TestRoot { TestRoot::new("graph/test/cases/range_search") } -fn setup_grid_index(grid_size: usize, dims: Grid) -> Arc> { +pub(super) fn setup_grid_index( + grid_size: usize, + dims: Grid, +) -> Arc> { let provider = test_provider::Provider::grid(dims, grid_size).unwrap(); let index_config = graph::config::Builder::new( @@ -47,7 +108,7 @@ fn setup_grid_index(grid_size: usize, dims: Grid) -> Arc (Arc>, Vec) { @@ -56,41 +117,18 @@ fn setup_grid_index_and_default_query( (index, query) } -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -struct RangeSearchBaseline { - grid_size: usize, - query: Vec, - radius: f32, - inner_radius: Option, - starting_l: usize, - results: Vec<(u32, f32)>, - comparisons: usize, - hops: usize, - result_count: usize, - range_search_second_round: bool, -} - -verbose_eq!(RangeSearchBaseline { - grid_size, - query, - radius, - inner_radius, - starting_l, - results, - comparisons, - hops, - result_count, - range_search_second_round, -}); - -fn assert_no_duplicates(results: &[Neighbor]) { +pub(super) fn assert_no_duplicates(results: &[Neighbor]) { let mut seen = std::collections::HashSet::new(); for n in results { assert!(seen.insert(n.id), "duplicate result id {}", n.id); } } -fn assert_range_invariants(results: &[Neighbor], radius: f32, inner_radius: Option) { +pub(super) fn assert_range_invariants( + results: &[Neighbor], + radius: f32, + inner_radius: Option, +) { for n in results { assert!( n.distance <= radius, @@ -113,6 +151,10 @@ fn assert_range_invariants(results: &[Neighbor], radius: f32, inner_radius: #[test] fn basic_range_search() { + let description = "Basic range search test to validate that the range / + search returns results within the specified radius and that there are / + no duplicate results."; + let rt = current_thread_runtime(); let mut test_root = root(); let mut path = test_root.path(); @@ -137,6 +179,7 @@ fn basic_range_search() { .unwrap(); let baseline = RangeSearchBaseline { + description: description.to_string(), grid_size, query: query.clone(), radius, @@ -158,6 +201,9 @@ fn basic_range_search() { #[test] fn inner_radius_filtering() { + let description = "Inner radius filtering test to validate that the \ + range search correctly excludes neighbors within the inner radius."; + let rt = current_thread_runtime(); let mut test_root = root(); let mut path = test_root.path(); @@ -169,8 +215,10 @@ fn inner_radius_filtering() { let inner_radius = 6.0; // exclude closest neighbors let starting_l = 32; - let range_search = - Range::with_options(None, starting_l, None, radius, Some(inner_radius), 1.0, 1.0).unwrap(); + let range_search = Range::builder(starting_l, radius) + .inner_radius(Some(inner_radius)) + .build() + .unwrap(); let mut results: Vec> = Vec::new(); let stats = rt @@ -184,6 +232,7 @@ fn inner_radius_filtering() { .unwrap(); let baseline = RangeSearchBaseline { + description: description.to_string(), grid_size, query: query.clone(), radius, @@ -205,6 +254,10 @@ fn inner_radius_filtering() { #[test] fn two_round_search() { + let description = "Two round search test to validate that a / + low starting L with a large radius triggers a second round / + of range search."; + let rt = current_thread_runtime(); let mut test_root = root(); let mut path = test_root.path(); @@ -229,6 +282,7 @@ fn two_round_search() { .unwrap(); let baseline = RangeSearchBaseline { + description: description.to_string(), grid_size, query: query.clone(), radius, diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/basic_range_search.json b/diskann/test/generated/graph/test/cases/filtered_range_search/basic_range_search.json new file mode 100644 index 000000000..b9361871c --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/basic_range_search.json @@ -0,0 +1,66 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/basic_range_search", + "payload": { + "comparisons": 53, + "description": "Basic range search test with an always-true filter. Validates /\n that the filtered range search returns correct results without any duplicates.", + "grid_size": 5, + "hops": 33, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 12.0, + "range_search_second_round": false, + "result_count": 11, + "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 + ], + [ + 93, + 12.0 + ] + ], + "starting_l": 32 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_l_search.json b/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_l_search.json new file mode 100644 index 000000000..709d57bb2 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_l_search.json @@ -0,0 +1,26 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_l_search", + "payload": { + "comparisons": 53, + "description": "Test of divisible-by-4 filter with a larger starting L. Since only 1/4 of discovered points will satisfy the filter, the second round of search should not be triggered.", + "grid_size": 5, + "hops": 33, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 12.0, + "range_search_second_round": false, + "result_count": 1, + "results": [ + [ + 124, + 3.0 + ] + ], + "starting_l": 32 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_max_results.json b/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_max_results.json new file mode 100644 index 000000000..d1d3f3b15 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_max_results.json @@ -0,0 +1,38 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/divisible_by_four_filter_no_second_round_from_max_results", + "payload": { + "comparisons": 31, + "description": "A test of the divisible-by-4 filter with `max_results` /\n set to starting_l / 4. The second round of search should not be triggered /\n since the correct number of results should be discovered.", + "grid_size": 5, + "hops": 17, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 50.0, + "range_search_second_round": false, + "result_count": 4, + "results": [ + [ + 124, + 3.0 + ], + [ + 68, + 17.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ] + ], + "starting_l": 16 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_second_round_triggered.json b/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_second_round_triggered.json new file mode 100644 index 000000000..36fef33db --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/divisible_by_four_filter_second_round_triggered.json @@ -0,0 +1,134 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/divisible_by_four_filter_second_round_triggered", + "payload": { + "comparisons": 10, + "description": "Test that a small starting L triggers a second /\n round of search when using the divisible-by-4 filter. Tests for /\n no duplicates, respect for the radius, and respect for the filter.", + "grid_size": 5, + "hops": 120, + "inner_radius": null, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 50.0, + "range_search_second_round": true, + "result_count": 28, + "results": [ + [ + 124, + 3.0 + ], + [ + 44, + 21.0 + ], + [ + 64, + 19.0 + ], + [ + 68, + 17.0 + ], + [ + 84, + 21.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 104, + 27.0 + ], + [ + 108, + 21.0 + ], + [ + 112, + 19.0 + ], + [ + 116, + 21.0 + ], + [ + 48, + 21.0 + ], + [ + 72, + 19.0 + ], + [ + 96, + 21.0 + ], + [ + 120, + 27.0 + ], + [ + 24, + 27.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 + ] + ], + "starting_l": 4 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/inner_radius_filtering.json b/diskann/test/generated/graph/test/cases/filtered_range_search/inner_radius_filtering.json new file mode 100644 index 000000000..77d511d12 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/inner_radius_filtering.json @@ -0,0 +1,110 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/inner_radius_filtering", + "payload": { + "comparisons": 53, + "description": "Inner radius filtering test to validate that the range search correctly excludes neighbors within the inner radius.", + "grid_size": 5, + "hops": 33, + "inner_radius": 6.0, + "query": [ + 5.0, + 5.0, + 5.0 + ], + "radius": 20.0, + "range_search_second_round": false, + "result_count": 22, + "results": [ + [ + 98, + 9.0 + ], + [ + 118, + 9.0 + ], + [ + 94, + 9.0 + ], + [ + 122, + 11.0 + ], + [ + 114, + 11.0 + ], + [ + 74, + 11.0 + ], + [ + 93, + 12.0 + ], + [ + 69, + 14.0 + ], + [ + 89, + 14.0 + ], + [ + 113, + 14.0 + ], + [ + 117, + 14.0 + ], + [ + 73, + 14.0 + ], + [ + 97, + 14.0 + ], + [ + 68, + 17.0 + ], + [ + 88, + 17.0 + ], + [ + 92, + 17.0 + ], + [ + 109, + 18.0 + ], + [ + 49, + 18.0 + ], + [ + 121, + 18.0 + ], + [ + 72, + 19.0 + ], + [ + 64, + 19.0 + ], + [ + 112, + 19.0 + ] + ], + "starting_l": 32 + } +} \ No newline at end of file diff --git a/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_and_second_round_triggered.json b/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_and_second_round_triggered.json new file mode 100644 index 000000000..8d5d1ab25 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_and_second_round_triggered.json @@ -0,0 +1,458 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/max_results_respected_and_second_round_triggered", + "payload": { + "comparisons": 10, + "description": "Test of `max_results` that sets `initial_l_search` /\n below `max_results` so that a second round of search is triggered. Note that /\n the result set size is expected to be below `max_results` due to filtering /\n out start points during `post_process`.", + "grid_size": 5, + "hops": 120, + "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/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 new file mode 100644 index 000000000..303db39c5 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/max_results_respected_means_no_second_round.json @@ -0,0 +1,34 @@ +{ + "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 /\n to `max_results`, with a permissive radius so that `max_results` is met / \n without needing a second round.", + "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/filtered_range_search/two_round_search.json b/diskann/test/generated/graph/test/cases/filtered_range_search/two_round_search.json new file mode 100644 index 000000000..98391a3d3 --- /dev/null +++ b/diskann/test/generated/graph/test/cases/filtered_range_search/two_round_search.json @@ -0,0 +1,458 @@ +{ + "file": "diskann/src/graph/test/cases/filtered_range_search.rs", + "test": "graph/test/cases/filtered_range_search/two_round_search", + "payload": { + "comparisons": 10, + "description": "Two round search test to validate that a /\n low starting L with a large radius triggers a second round /\n of range search.", + "grid_size": 5, + "hops": 120, + "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/range_search/basic_range_search.json b/diskann/test/generated/graph/test/cases/range_search/basic_range_search.json index d24b9c789..328e0e70e 100644 --- a/diskann/test/generated/graph/test/cases/range_search/basic_range_search.json +++ b/diskann/test/generated/graph/test/cases/range_search/basic_range_search.json @@ -3,6 +3,7 @@ "test": "graph/test/cases/range_search/basic_range_search", "payload": { "comparisons": 54, + "description": "Basic range search test to validate that the range /\n search returns results within the specified radius and that there are /\n no duplicate results.", "grid_size": 5, "hops": 33, "inner_radius": null, diff --git a/diskann/test/generated/graph/test/cases/range_search/inner_radius_filtering.json b/diskann/test/generated/graph/test/cases/range_search/inner_radius_filtering.json index 2d95adda5..b9fb9ef2a 100644 --- a/diskann/test/generated/graph/test/cases/range_search/inner_radius_filtering.json +++ b/diskann/test/generated/graph/test/cases/range_search/inner_radius_filtering.json @@ -3,6 +3,7 @@ "test": "graph/test/cases/range_search/inner_radius_filtering", "payload": { "comparisons": 54, + "description": "Inner radius filtering test to validate that the range search correctly excludes neighbors within the inner radius.", "grid_size": 5, "hops": 33, "inner_radius": 6.0, 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 992d3cce2..24a8a97d5 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 @@ -3,6 +3,7 @@ "test": "graph/test/cases/range_search/two_round_search", "payload": { "comparisons": 11, + "description": "Two round search test to validate that a /\n low starting L with a large radius triggers a second round /\n of range search.", "grid_size": 5, "hops": 120, "inner_radius": null, diff --git a/test_data/disk_index_search/groundtruth_rad_125000_filtered.rangeres b/test_data/disk_index_search/groundtruth_rad_125000_filtered.rangeres new file mode 100644 index 000000000..6004efabe --- /dev/null +++ b/test_data/disk_index_search/groundtruth_rad_125000_filtered.rangeres @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:de8d1e9b55afb10baa6f0ab84b6cd4bffff3f67826026ea0d241fae02dc698e7 +size 380 diff --git a/test_data/yfcc/groundtruth_rad_110000_filtered.rangeres b/test_data/yfcc/groundtruth_rad_110000_filtered.rangeres new file mode 100644 index 000000000..5395dd7ea --- /dev/null +++ b/test_data/yfcc/groundtruth_rad_110000_filtered.rangeres @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:185d5505233a233b9603f16fe12d330dbf394204007349823df7f7966fbea198 +size 6832 diff --git a/test_data/yfcc/groundtruth_rad_75000.rangeres b/test_data/yfcc/groundtruth_rad_75000.rangeres new file mode 100644 index 000000000..29457faa8 --- /dev/null +++ b/test_data/yfcc/groundtruth_rad_75000.rangeres @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:50a497a2769288285344db2e6bc98634174ea356aa6dc67119fcbe0603c66774 +size 8068