diff --git a/diskann-benchmark-core/src/search/graph/mod.rs b/diskann-benchmark-core/src/search/graph/mod.rs index cfcecb0db..3e280fae2 100644 --- a/diskann-benchmark-core/src/search/graph/mod.rs +++ b/diskann-benchmark-core/src/search/graph/mod.rs @@ -5,6 +5,7 @@ pub mod determinant_diversity; pub mod knn; +pub mod multi_attribute_diversity; pub mod multihop; pub mod range; diff --git a/diskann-benchmark-core/src/search/graph/multi_attribute_diversity.rs b/diskann-benchmark-core/src/search/graph/multi_attribute_diversity.rs new file mode 100644 index 000000000..5dd000032 --- /dev/null +++ b/diskann-benchmark-core/src/search/graph/multi_attribute_diversity.rs @@ -0,0 +1,203 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::sync::Arc; + +use diskann::{ + ANNResult, + graph::{self, glue}, + provider, +}; +use diskann_benchmark_runner::utils::{MicroSeconds, percentiles}; +use diskann_providers::model::graph::provider::async_::MultiAttributeDiversitySearchParams; +use diskann_utils::{future::AsyncFriendly, views::Matrix}; + +use crate::{ + recall, + search::{self, Search, graph::Strategy}, + utils, +}; + +#[derive(Debug, Clone, Copy)] +pub struct Parameters { + pub inner: graph::search::Knn, + pub processor: MultiAttributeDiversitySearchParams, +} + +#[derive(Debug)] +pub struct MultiAttributeDiversity +where + DP: provider::DataProvider, +{ + index: Arc>, + queries: Arc>, + strategy: Strategy, +} + +impl MultiAttributeDiversity +where + DP: provider::DataProvider, +{ + pub fn new( + index: Arc>, + queries: Arc>, + strategy: Strategy, + ) -> anyhow::Result> { + strategy.length_compatible(queries.nrows())?; + + Ok(Arc::new(Self { + index, + queries, + strategy, + })) + } +} + +impl Search for MultiAttributeDiversity +where + DP: provider::DataProvider, + S: glue::DefaultSearchStrategy + Clone + AsyncFriendly, + MultiAttributeDiversitySearchParams: + for<'a> glue::SearchPostProcess, [T], DP::ExternalId> + Send + Sync, + T: AsyncFriendly + Clone, +{ + type Id = DP::ExternalId; + type Parameters = Parameters; + type Output = super::knn::Metrics; + + fn num_queries(&self) -> usize { + self.queries.nrows() + } + + fn id_count(&self, parameters: &Self::Parameters) -> search::IdCount { + search::IdCount::Fixed(parameters.inner.k_value()) + } + + async fn search( + &self, + parameters: &Self::Parameters, + buffer: &mut O, + index: usize, + ) -> ANNResult + where + O: graph::SearchOutputBuffer + Send, + { + let context = DP::Context::default(); + let stats = self + .index + .search_with( + parameters.inner, + self.strategy.get(index)?, + parameters.processor, + &context, + self.queries.row(index), + buffer, + ) + .await?; + + Ok(super::knn::Metrics { + comparisons: stats.cmps, + hops: stats.hops, + }) + } +} + +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct Summary { + pub setup: search::Setup, + pub parameters: Parameters, + pub end_to_end_latencies: Vec, + pub mean_latencies: Vec, + pub p90_latencies: Vec, + pub p99_latencies: Vec, + pub recall: recall::RecallMetrics, + pub mean_cmps: f64, + pub mean_hops: f64, +} + +pub struct Aggregator<'a, I> { + groundtruth: &'a dyn crate::recall::Rows, + recall_k: usize, + recall_n: usize, +} + +impl<'a, I> Aggregator<'a, I> { + pub fn new( + groundtruth: &'a dyn crate::recall::Rows, + recall_k: usize, + recall_n: usize, + ) -> Self { + Self { + groundtruth, + recall_k, + recall_n, + } + } +} + +impl search::Aggregate for Aggregator<'_, I> +where + I: crate::recall::RecallCompatible, +{ + type Output = Summary; + + fn aggregate( + &mut self, + run: search::Run, + mut results: Vec>, + ) -> anyhow::Result { + let recall = match results.first() { + Some(first) => crate::recall::knn( + self.groundtruth, + None, + first.ids().as_rows(), + self.recall_k, + self.recall_n, + true, + )?, + 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(), + recall, + mean_latencies, + p90_latencies, + p99_latencies, + mean_cmps: utils::average_all( + results + .iter() + .flat_map(|r| r.output().iter().map(|o| o.comparisons)), + ), + mean_hops: utils::average_all( + results.iter().flat_map(|r| r.output().iter().map(|o| o.hops)), + ), + }) + } +} diff --git a/diskann-benchmark/example/async-multi-attribute-diversity.json b/diskann-benchmark/example/async-multi-attribute-diversity.json new file mode 100644 index 000000000..c41034e40 --- /dev/null +++ b/diskann-benchmark/example/async-multi-attribute-diversity.json @@ -0,0 +1,88 @@ +{ + "search_directories": [ + "test_data/disk_index_search" + ], + "jobs": [ + { + "type": "async-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": "topk", + "queries": "disk_index_sample_query_10pts.fbin", + "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin", + "reps": 5, + "num_threads": [ + 1 + ], + "runs": [ + { + "search_n": 20, + "search_l": [ + 20, + 30, + 40 + ], + "recall_k": 10 + } + ] + } + } + }, + { + "type": "async-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": "topk", + "queries": "disk_index_sample_query_10pts.fbin", + "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin", + "reps": 5, + "num_threads": [ + 1 + ], + "runs": [ + { + "search_n": 20, + "search_l": [ + 20, + 30, + 40 + ], + "recall_k": 10 + } + ], + "multi_attribute_diversity_eta": 0.01, + "multi_attribute_diversity_power": 2.0, + "multi_attribute_diversity_results_k": 10 + } + } + } + ] +} diff --git a/diskann-benchmark/example/openai-disk-multi-attribute-diversity-compare.json b/diskann-benchmark/example/openai-disk-multi-attribute-diversity-compare.json new file mode 100644 index 000000000..84249ed56 --- /dev/null +++ b/diskann-benchmark/example/openai-disk-multi-attribute-diversity-compare.json @@ -0,0 +1,52 @@ +{ + "search_directories": [ + "C:/data/openai" + ], + "jobs": [ + { + "type": "disk-index", + "content": { + "source": { + "disk-index-source": "Load", + "data_type": "float32", + "load_path": "C:/data/openai/openai_index_normal" + }, + "search_phase": { + "queries": "openai_query.bin", + "groundtruth": "openai_gt_50.bin", + "search_list": [100, 200, 400], + "beam_width": 4, + "recall_at": 10, + "num_threads": 8, + "is_flat_search": false, + "distance": "squared_l2", + "vector_filters_file": null + } + } + }, + { + "type": "disk-index", + "content": { + "source": { + "disk-index-source": "Load", + "data_type": "float32", + "load_path": "C:/data/openai/openai_index_normal" + }, + "search_phase": { + "queries": "openai_query.bin", + "groundtruth": "openai_gt_50.bin", + "search_list": [100, 200, 400], + "beam_width": 4, + "recall_at": 10, + "num_threads": 8, + "is_flat_search": false, + "distance": "squared_l2", + "vector_filters_file": null, + "is_multi_attribute_diversity_search": true, + "multi_attribute_diversity_eta": 0.01, + "multi_attribute_diversity_power": 2.0 + } + } + } + ] +} diff --git a/diskann-benchmark/src/backend/disk_index/search.rs b/diskann-benchmark/src/backend/disk_index/search.rs index ea1f0f130..6b70a4daa 100644 --- a/diskann-benchmark/src/backend/disk_index/search.rs +++ b/diskann-benchmark/src/backend/disk_index/search.rs @@ -45,6 +45,9 @@ pub(super) struct DiskSearchStats { pub(crate) is_determinant_diversity_search: bool, pub(crate) determinant_diversity_eta: Option, pub(crate) determinant_diversity_power: Option, + pub(crate) is_multi_attribute_diversity_search: bool, + pub(crate) multi_attribute_diversity_eta: Option, + pub(crate) multi_attribute_diversity_power: Option, pub(crate) distance: SimilarityMeasure, pub(crate) uses_vector_filters: bool, pub(super) num_nodes_to_cache: Option, @@ -241,6 +244,15 @@ where // Execute search iterations for &l in search_params.search_list.iter() { + let use_diversity_post_process = search_params.is_determinant_diversity_search + || search_params.is_multi_attribute_diversity_search; + let diversity_eta = search_params + .multi_attribute_diversity_eta + .or(search_params.determinant_diversity_eta); + let diversity_power = search_params + .multi_attribute_diversity_power + .or(search_params.determinant_diversity_power); + let mut statistics_vec: Vec = vec![QueryStatistics::default(); num_queries]; let mut result_counts: Vec = vec![0; num_queries]; @@ -279,9 +291,9 @@ where Some(search_params.beam_width), vector_filter, search_params.is_flat_search, - search_params.is_determinant_diversity_search, - search_params.determinant_diversity_eta, - search_params.determinant_diversity_power, + use_diversity_post_process, + diversity_eta, + diversity_power, ) { Ok(search_result) => { *stats = search_result.stats.query_statistics; @@ -350,6 +362,9 @@ where is_determinant_diversity_search: search_params.is_determinant_diversity_search, determinant_diversity_eta: search_params.determinant_diversity_eta, determinant_diversity_power: search_params.determinant_diversity_power, + is_multi_attribute_diversity_search: search_params.is_multi_attribute_diversity_search, + multi_attribute_diversity_eta: search_params.multi_attribute_diversity_eta, + multi_attribute_diversity_power: search_params.multi_attribute_diversity_power, distance: search_params.distance, uses_vector_filters: search_params.vector_filters_file.is_some(), num_nodes_to_cache: search_params.num_nodes_to_cache, @@ -450,6 +465,22 @@ impl fmt::Display for DiskSearchStats { _ => "None".to_string(), } )?; + writeln!( + f, + "Multi-div search, : {}", + self.is_multi_attribute_diversity_search + )?; + writeln!( + f, + "Multi-div params, : {}", + match ( + self.multi_attribute_diversity_eta, + self.multi_attribute_diversity_power, + ) { + (Some(eta), Some(power)) => format!("eta={eta}, power={power}"), + _ => "None".to_string(), + } + )?; writeln!(f, "Distance, : {}", self.distance)?; writeln!(f, "Vector filters, : {}", self.uses_vector_filters)?; writeln!( diff --git a/diskann-benchmark/src/backend/index/benchmarks.rs b/diskann-benchmark/src/backend/index/benchmarks.rs index a733ada8e..4f5d607f0 100644 --- a/diskann-benchmark/src/backend/index/benchmarks.rs +++ b/diskann-benchmark/src/backend/index/benchmarks.rs @@ -26,7 +26,9 @@ use diskann_providers::{ index::diskann_async, model::{ configuration::IndexConfiguration, - graph::provider::async_::{common, DeterminantDiversitySearchParams}, + graph::provider::async_::{ + common, DeterminantDiversitySearchParams, MultiAttributeDiversitySearchParams, + }, }, }; use diskann_utils::{ @@ -355,6 +357,8 @@ where S: glue::DefaultSearchStrategy + Clone + AsyncFriendly, DeterminantDiversitySearchParams: for<'a> glue::SearchPostProcess, [T], DP::ExternalId> + Send + Sync, + MultiAttributeDiversitySearchParams: + for<'a> glue::SearchPostProcess, [T], DP::ExternalId> + Send + Sync, { match &input { SearchPhase::Topk(search_phase) => { @@ -378,6 +382,24 @@ where ); let search_results = if let (Some(eta), Some(power)) = ( + search_phase.multi_attribute_diversity_eta, + search_phase.multi_attribute_diversity_power, + ) { + let knn = benchmark_core::search::graph::multi_attribute_diversity::MultiAttributeDiversity::new( + index, + queries, + benchmark_core::search::graph::Strategy::broadcast(search_strategy), + )?; + + search::knn::run_multi_attribute_diversity( + &knn, + &groundtruth, + steps, + eta, + power, + search_phase.multi_attribute_diversity_results_k, + )? + } else if let (Some(eta), Some(power)) = ( search_phase.determinant_diversity_eta, search_phase.determinant_diversity_power, ) { diff --git a/diskann-benchmark/src/backend/index/result.rs b/diskann-benchmark/src/backend/index/result.rs index 1f9c2e50a..25a3dbb7b 100644 --- a/diskann-benchmark/src/backend/index/result.rs +++ b/diskann-benchmark/src/backend/index/result.rs @@ -122,6 +122,38 @@ pub(super) struct SearchResults { } impl SearchResults { + fn from_topk_parts( + setup: benchmark_core::search::Setup, + search_n: usize, + search_l: usize, + end_to_end_latencies: Vec, + mean_latencies: Vec, + p90_latencies: Vec, + p99_latencies: Vec, + recall: benchmark_core::recall::RecallMetrics, + mean_cmps: f64, + mean_hops: f64, + ) -> Self { + let qps = end_to_end_latencies + .iter() + .map(|latency| recall.num_queries as f64 / latency.as_seconds()) + .collect(); + + Self { + num_tasks: setup.tasks.into(), + search_n, + search_l, + qps, + search_latencies: end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + recall: (&recall).into(), + mean_cmps: mean_cmps as f32, + mean_hops: mean_hops as f32, + } + } + pub fn new(summary: benchmark_core::search::graph::knn::Summary) -> Self { let benchmark_core::search::graph::knn::Summary { setup, @@ -136,24 +168,18 @@ impl SearchResults { .. } = summary; - let qps = end_to_end_latencies - .iter() - .map(|latency| recall.num_queries as f64 / latency.as_seconds()) - .collect(); - - Self { - num_tasks: setup.tasks.into(), - search_n: parameters.k_value().get(), - search_l: parameters.l_value().get(), - qps, - search_latencies: end_to_end_latencies, + Self::from_topk_parts( + setup, + parameters.k_value().get(), + parameters.l_value().get(), + end_to_end_latencies, mean_latencies, p90_latencies, p99_latencies, - recall: (&recall).into(), - mean_cmps: mean_cmps as f32, - mean_hops: mean_hops as f32, - } + recall, + mean_cmps, + mean_hops, + ) } pub fn new_determinant_diversity( @@ -172,24 +198,48 @@ impl SearchResults { .. } = summary; - let qps = end_to_end_latencies - .iter() - .map(|latency| recall.num_queries as f64 / latency.as_seconds()) - .collect(); + Self::from_topk_parts( + setup, + parameters.inner.k_value().get(), + parameters.inner.l_value().get(), + end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + recall, + mean_cmps, + mean_hops, + ) + } - Self { - num_tasks: setup.tasks.into(), - search_n: parameters.inner.k_value().get(), - search_l: parameters.inner.l_value().get(), - qps, - search_latencies: end_to_end_latencies, + pub fn new_multi_attribute_diversity( + summary: benchmark_core::search::graph::multi_attribute_diversity::Summary, + ) -> Self { + let benchmark_core::search::graph::multi_attribute_diversity::Summary { + setup, + parameters, + end_to_end_latencies, mean_latencies, p90_latencies, p99_latencies, - recall: (&recall).into(), - mean_cmps: mean_cmps as f32, - mean_hops: mean_hops as f32, - } + recall, + mean_cmps, + mean_hops, + .. + } = summary; + + Self::from_topk_parts( + setup, + parameters.inner.k_value().get(), + parameters.inner.l_value().get(), + end_to_end_latencies, + mean_latencies, + p90_latencies, + p99_latencies, + recall, + mean_cmps, + mean_hops, + ) } } diff --git a/diskann-benchmark/src/backend/index/search/knn.rs b/diskann-benchmark/src/backend/index/search/knn.rs index 30560a6cd..e4540c67b 100644 --- a/diskann-benchmark/src/backend/index/search/knn.rs +++ b/diskann-benchmark/src/backend/index/search/knn.rs @@ -6,7 +6,9 @@ use std::{num::NonZeroUsize, sync::Arc}; use diskann_benchmark_core::{self as benchmark_core, search as core_search}; -use diskann_providers::model::graph::provider::async_::DeterminantDiversitySearchParams; +use diskann_providers::model::graph::provider::async_::{ + DeterminantDiversitySearchParams, MultiAttributeDiversitySearchParams, +}; use crate::{backend::index::result::SearchResults, inputs::async_::GraphSearch}; @@ -36,9 +38,11 @@ pub(crate) fn run( groundtruth: &dyn benchmark_core::recall::Rows, steps: SearchSteps<'_>, ) -> anyhow::Result> { - run_search(runner, groundtruth, steps, |setup, search_l, search_n| { + run_search_with_builder(steps, |setup, search_l, search_n| { let search_params = diskann::graph::search::Knn::new(search_n, search_l, None).unwrap(); - core_search::Run::new(search_params, setup) + Ok(core_search::Run::new(search_params, setup)) + }, |parameters, recall_k, search_n| { + runner.search_all(parameters, groundtruth, recall_k, search_n) }) } @@ -64,7 +68,7 @@ pub(crate) fn run_determinant_diversity( power: f64, results_k: Option, ) -> anyhow::Result> { - run_search_determinant_diversity(runner, groundtruth, steps, |setup, search_l, search_n| { + run_search_with_builder(steps, |setup, search_l, search_n| { let base = diskann::graph::search::Knn::new(search_n, search_l, None).unwrap(); let processor = DeterminantDiversitySearchParams::new(results_k.unwrap_or(search_n), eta, power) @@ -78,49 +82,50 @@ pub(crate) fn run_determinant_diversity( processor, }; Ok(core_search::Run::new(search_params, setup)) + }, |parameters, recall_k, search_n| { + runner.search_all(parameters, groundtruth, recall_k, search_n) }) } -fn run_search( - runner: &dyn Knn, +type MultiAttributeRun = + core_search::Run; + +pub(crate) fn run_multi_attribute_diversity( + runner: &dyn MultiAttributeDiversityKnn, groundtruth: &dyn benchmark_core::recall::Rows, steps: SearchSteps<'_>, - builder: F, + eta: f64, + power: f64, + results_k: Option, ) -> anyhow::Result> -where - F: Fn(core_search::Setup, usize, usize) -> Run, { - let mut all = Vec::new(); + run_search_with_builder(steps, |setup, search_l, search_n| { + let base = diskann::graph::search::Knn::new(search_n, search_l, None).unwrap(); + let processor = + MultiAttributeDiversitySearchParams::new(results_k.unwrap_or(search_n), eta, power) + .map_err(|err| { + anyhow::anyhow!("Invalid multi-attribute-diversity parameters: {err}") + })?; - for threads in steps.num_tasks.iter() { - for run in steps.runs.iter() { - let setup = core_search::Setup { - threads: *threads, - tasks: *threads, - reps: steps.reps, + let search_params = + diskann_benchmark_core::search::graph::multi_attribute_diversity::Parameters { + inner: base, + processor, }; - - let parameters: Vec<_> = run - .search_l - .iter() - .map(|&search_l| builder(setup.clone(), search_l, run.search_n)) - .collect(); - - all.extend(runner.search_all(parameters, groundtruth, run.recall_k, run.search_n)?); - } - } - - Ok(all) + Ok(core_search::Run::new(search_params, setup)) + }, |parameters, recall_k, search_n| { + runner.search_all(parameters, groundtruth, recall_k, search_n) + }) } -fn run_search_determinant_diversity( - runner: &dyn DeterminantDiversityKnn, - groundtruth: &dyn benchmark_core::recall::Rows, +fn run_search_with_builder( steps: SearchSteps<'_>, builder: F, + mut execute: E, ) -> anyhow::Result> where - F: Fn(core_search::Setup, usize, usize) -> anyhow::Result, + F: Fn(core_search::Setup, usize, usize) -> anyhow::Result>, + E: FnMut(Vec>, usize, usize) -> anyhow::Result>, { let mut all = Vec::new(); @@ -138,7 +143,7 @@ where .map(|&search_l| builder(setup.clone(), search_l, run.search_n)) .collect::>>()?; - all.extend(runner.search_all(parameters, groundtruth, run.recall_k, run.search_n)?); + all.extend(execute(parameters, run.recall_k, run.search_n)?); } } @@ -155,6 +160,16 @@ pub(crate) trait DeterminantDiversityKnn { ) -> anyhow::Result>; } +pub(crate) trait MultiAttributeDiversityKnn { + fn search_all( + &self, + parameters: Vec, + groundtruth: &dyn benchmark_core::recall::Rows, + recall_k: usize, + recall_n: usize, + ) -> anyhow::Result>; +} + /////////// // Impls // /////////// @@ -244,3 +259,38 @@ where .collect()) } } + +impl MultiAttributeDiversityKnn + for Arc> +where + DP: diskann::provider::DataProvider, + core_search::graph::multi_attribute_diversity::MultiAttributeDiversity: + core_search::Search< + Id = DP::InternalId, + Parameters = diskann_benchmark_core::search::graph::multi_attribute_diversity::Parameters, + Output = core_search::graph::knn::Metrics, + >, +{ + fn search_all( + &self, + parameters: Vec, + groundtruth: &dyn benchmark_core::recall::Rows, + recall_k: usize, + recall_n: usize, + ) -> anyhow::Result> { + let results = core_search::search_all( + self.clone(), + parameters.into_iter(), + core_search::graph::multi_attribute_diversity::Aggregator::new( + groundtruth, + recall_k, + recall_n, + ), + )?; + + Ok(results + .into_iter() + .map(SearchResults::new_multi_attribute_diversity) + .collect()) + } +} diff --git a/diskann-benchmark/src/backend/index/spherical.rs b/diskann-benchmark/src/backend/index/spherical.rs index 1918d6676..dcfcf48f5 100644 --- a/diskann-benchmark/src/backend/index/spherical.rs +++ b/diskann-benchmark/src/backend/index/spherical.rs @@ -68,6 +68,7 @@ mod imp { index::diskann_async::{self}, model::graph::provider::async_::{ common::NoDeletes, inmem, DeterminantDiversitySearchParams, + MultiAttributeDiversitySearchParams, }, }; use diskann_quantization::alloc::GlobalAllocator; @@ -335,6 +336,24 @@ mod imp { for &layout in self.input.query_layouts.iter() { let strategy = inmem::spherical::Quantized::search(layout.into()); let search_results = if let (Some(eta), Some(power)) = ( + search_phase.multi_attribute_diversity_eta, + search_phase.multi_attribute_diversity_power, + ) { + let knn = benchmark_core::search::graph::multi_attribute_diversity::MultiAttributeDiversity::new( + index.clone(), + queries.clone(), + benchmark_core::search::graph::Strategy::broadcast(strategy), + )?; + + search::knn::run_multi_attribute_diversity( + &knn, + &groundtruth, + steps, + eta, + power, + search_phase.multi_attribute_diversity_results_k, + )? + } else if let (Some(eta), Some(power)) = ( search_phase.determinant_diversity_eta, search_phase.determinant_diversity_power, ) { diff --git a/diskann-benchmark/src/inputs/async_.rs b/diskann-benchmark/src/inputs/async_.rs index 54b795635..5b6dd7df1 100644 --- a/diskann-benchmark/src/inputs/async_.rs +++ b/diskann-benchmark/src/inputs/async_.rs @@ -126,6 +126,9 @@ pub(crate) struct TopkSearchPhase { pub(crate) determinant_diversity_eta: Option, pub(crate) determinant_diversity_power: Option, pub(crate) determinant_diversity_results_k: Option, + pub(crate) multi_attribute_diversity_eta: Option, + pub(crate) multi_attribute_diversity_power: Option, + pub(crate) multi_attribute_diversity_results_k: Option, // Enable sweeping threads pub(crate) num_threads: Vec, pub(crate) runs: Vec, @@ -148,6 +151,20 @@ impl CheckDeserialization for TopkSearchPhase { )); } + if self.multi_attribute_diversity_eta.is_some() + != self.multi_attribute_diversity_power.is_some() + { + return Err(anyhow!( + "multi_attribute_diversity_eta and multi_attribute_diversity_power must either both be set or both be omitted" + )); + } + + if self.determinant_diversity_eta.is_some() && self.multi_attribute_diversity_eta.is_some() { + return Err(anyhow!( + "determinant_diversity and multi_attribute_diversity are mutually exclusive" + )); + } + if let Some(eta) = self.determinant_diversity_eta { if eta < 0.0 { return Err(anyhow!( @@ -172,6 +189,30 @@ impl CheckDeserialization for TopkSearchPhase { } } + if let Some(eta) = self.multi_attribute_diversity_eta { + if eta < 0.0 { + return Err(anyhow!( + "multi_attribute_diversity_eta must be >= 0.0, got {}", + eta + )); + } + } + + if let Some(power) = self.multi_attribute_diversity_power { + if power <= 0.0 { + return Err(anyhow!( + "multi_attribute_diversity_power must be > 0.0, got {}", + power + )); + } + } + + if let Some(k) = self.multi_attribute_diversity_results_k { + if k == 0 { + return Err(anyhow!("multi_attribute_diversity_results_k must be > 0")); + } + } + Ok(()) } } @@ -200,6 +241,9 @@ impl Example for TopkSearchPhase { determinant_diversity_eta: None, determinant_diversity_power: None, determinant_diversity_results_k: None, + multi_attribute_diversity_eta: None, + multi_attribute_diversity_power: None, + multi_attribute_diversity_results_k: None, num_threads: THREAD_COUNTS.to_vec(), runs, } diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index f5058f321..9a697bce1 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -87,6 +87,12 @@ pub(crate) struct DiskSearchPhase { pub(crate) determinant_diversity_eta: Option, #[serde(default)] pub(crate) determinant_diversity_power: Option, + #[serde(default)] + pub(crate) is_multi_attribute_diversity_search: bool, + #[serde(default)] + pub(crate) multi_attribute_diversity_eta: Option, + #[serde(default)] + pub(crate) multi_attribute_diversity_power: Option, pub(crate) distance: SimilarityMeasure, pub(crate) vector_filters_file: Option, pub(crate) num_nodes_to_cache: Option, @@ -259,6 +265,42 @@ impl CheckDeserialization for DiskSearchPhase { ); } + if self.is_multi_attribute_diversity_search { + if self.is_flat_search { + anyhow::bail!( + "is_multi_attribute_diversity_search is not supported when is_flat_search is true" + ); + } + + let eta = self.multi_attribute_diversity_eta.unwrap_or(0.01); + let power = self.multi_attribute_diversity_power.unwrap_or(2.0); + + if eta < 0.0 || !eta.is_finite() { + anyhow::bail!("multi_attribute_diversity_eta must be >= 0.0 and finite, got {eta}"); + } + + if power <= 0.0 || !power.is_finite() { + anyhow::bail!( + "multi_attribute_diversity_power must be > 0.0 and finite, got {power}" + ); + } + + self.multi_attribute_diversity_eta = Some(eta); + self.multi_attribute_diversity_power = Some(power); + } else if self.multi_attribute_diversity_eta.is_some() + || self.multi_attribute_diversity_power.is_some() + { + anyhow::bail!( + "multi_attribute_diversity_eta/multi_attribute_diversity_power may only be set when is_multi_attribute_diversity_search is true" + ); + } + + if self.is_determinant_diversity_search && self.is_multi_attribute_diversity_search { + anyhow::bail!( + "is_determinant_diversity_search and is_multi_attribute_diversity_search are mutually exclusive" + ); + } + if let Some(n) = self.num_nodes_to_cache { if n == 0 { anyhow::bail!("num_nodes_to_cache must be positive if specified"); @@ -306,6 +348,9 @@ impl Example for DiskIndexOperation { is_determinant_diversity_search: false, determinant_diversity_eta: None, determinant_diversity_power: None, + is_multi_attribute_diversity_search: false, + multi_attribute_diversity_eta: None, + multi_attribute_diversity_power: None, distance: SimilarityMeasure::SquaredL2, vector_filters_file: None, num_nodes_to_cache: None, @@ -438,6 +483,22 @@ impl DiskSearchPhase { )?, _ => write_field!(f, "Determinant Diversity Params", "none")?, } + write_field!( + f, + "Multi-Attr Diversity Search", + self.is_multi_attribute_diversity_search + )?; + match ( + self.multi_attribute_diversity_eta, + self.multi_attribute_diversity_power, + ) { + (Some(eta), Some(power)) => write_field!( + f, + "Multi-Attr Diversity Params", + format!("eta={eta}, power={power}") + )?, + _ => write_field!(f, "Multi-Attr Diversity Params", "none")?, + } write_field!(f, "Distance", self.distance)?; match &self.vector_filters_file { Some(vf) => write_field!(f, "Vector Filters File", vf.display())?, diff --git a/diskann-providers/src/model/graph/provider/async_/mod.rs b/diskann-providers/src/model/graph/provider/async_/mod.rs index 774c5530d..f8a81c956 100644 --- a/diskann-providers/src/model/graph/provider/async_/mod.rs +++ b/diskann-providers/src/model/graph/provider/async_/mod.rs @@ -14,6 +14,11 @@ pub use determinant_diversity_post_process::{ DeterminantDiversityError, DeterminantDiversitySearchParams, determinant_diversity_post_process, }; +mod multi_attribute_diversity_post_process; +pub use multi_attribute_diversity_post_process::{ + MultiAttributeDiversityError, MultiAttributeDiversitySearchParams, +}; + pub mod distances; pub mod memory_vector_provider; diff --git a/diskann-providers/src/model/graph/provider/async_/multi_attribute_diversity_post_process.rs b/diskann-providers/src/model/graph/provider/async_/multi_attribute_diversity_post_process.rs new file mode 100644 index 000000000..887241758 --- /dev/null +++ b/diskann-providers/src/model/graph/provider/async_/multi_attribute_diversity_post_process.rs @@ -0,0 +1,113 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Multi-attribute-diversity search post-processing. + +use std::future::Future; + +use diskann::{ + ANNError, + graph::{SearchOutputBuffer, glue}, + neighbor::Neighbor, + provider::BuildQueryComputer, +}; + +use super::postprocess::{AsDeletionCheck, DeletionCheck}; + +#[derive(Debug)] +pub enum MultiAttributeDiversityError { + InvalidTopK { top_k: usize }, + InvalidEta { eta: f64 }, + InvalidPower { power: f64 }, +} + +impl std::fmt::Display for MultiAttributeDiversityError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidTopK { top_k } => write!(f, "top_k must be > 0, got {top_k}"), + Self::InvalidEta { eta } => write!(f, "eta must be >= 0.0, got {eta}"), + Self::InvalidPower { power } => write!(f, "power must be > 0.0, got {power}"), + } + } +} + +impl std::error::Error for MultiAttributeDiversityError {} + +#[derive(Debug, Clone, Copy)] +pub struct MultiAttributeDiversitySearchParams { + pub top_k: usize, + pub multi_attribute_diversity_eta: f64, + pub multi_attribute_diversity_power: f64, +} + +impl MultiAttributeDiversitySearchParams { + pub fn new( + top_k: usize, + multi_attribute_diversity_eta: f64, + multi_attribute_diversity_power: f64, + ) -> Result { + if top_k == 0 { + return Err(MultiAttributeDiversityError::InvalidTopK { top_k }); + } + + if multi_attribute_diversity_eta < 0.0 || !multi_attribute_diversity_eta.is_finite() { + return Err(MultiAttributeDiversityError::InvalidEta { + eta: multi_attribute_diversity_eta, + }); + } + + if multi_attribute_diversity_power <= 0.0 || !multi_attribute_diversity_power.is_finite() { + return Err(MultiAttributeDiversityError::InvalidPower { + power: multi_attribute_diversity_power, + }); + } + + Ok(Self { + top_k, + multi_attribute_diversity_eta, + multi_attribute_diversity_power, + }) + } +} + +impl glue::SearchPostProcess for MultiAttributeDiversitySearchParams +where + A: BuildQueryComputer<[T], Id = u32> + AsDeletionCheck, +{ + type Error = ANNError; + + fn post_process( + &self, + accessor: &mut A, + _query: &[T], + _computer: &>::QueryComputer, + candidates: I, + output: &mut B, + ) -> impl Future> + Send + where + I: Iterator> + Send, + B: SearchOutputBuffer + Send + ?Sized, + { + let checker = accessor.as_deletion_check(); + let mut values: Vec<_> = candidates + .filter_map(|candidate| { + if checker.deletion_check(candidate.id) { + None + } else { + Some((candidate.id, candidate.distance)) + } + }) + .collect(); + + values.sort_by(|left, right| { + left.1 + .partial_cmp(&right.1) + .unwrap_or(std::cmp::Ordering::Equal) + }); + values.truncate(self.top_k); + + std::future::ready(Ok(output.extend(values))) + } +}