diff --git a/assets/diverse-search/caselaw-jurisdiction-qps-vs-recall.png b/assets/diverse-search/caselaw-jurisdiction-qps-vs-recall.png new file mode 100644 index 000000000..d7fb4bae7 Binary files /dev/null and b/assets/diverse-search/caselaw-jurisdiction-qps-vs-recall.png differ diff --git a/assets/diverse-search/caselaw-qps-vs-recall.png b/assets/diverse-search/caselaw-qps-vs-recall.png new file mode 100644 index 000000000..39b9b271a Binary files /dev/null and b/assets/diverse-search/caselaw-qps-vs-recall.png differ diff --git a/assets/diverse-search/chart_data.json b/assets/diverse-search/chart_data.json new file mode 100644 index 000000000..32f984c51 --- /dev/null +++ b/assets/diverse-search/chart_data.json @@ -0,0 +1,52 @@ +{ + "_comment": "QPS-vs-recall data at L=100 for the diverse-search benchmark charts. Recall is deterministic and authoritative; QPS is the best (least-contended) of several single-threaded runs (see diverse-search-results.md). Update these values when rerunning benchmarks, then regenerate the PNGs with plot_qps_vs_recall.py.", + "methods": ["Standard", "Queue", "Design A", "Design B"], + "colors": { + "Standard": "#7f7f7f", + "Queue": "#1f77b4", + "Design A": "#2ca02c", + "Design B": "#d62728" + }, + "charts": [ + { + "output": "enron-qps-vs-recall.png", + "title": "Enron (concentrated attribute)", + "points": { + "Standard": {"recall": 70.72, "qps": 142.6}, + "Queue": {"recall": 77.92, "qps": 100.0}, + "Design A": {"recall": 86.05, "qps": 159.2}, + "Design B": {"recall": 93.01, "qps": 53.0} + } + }, + { + "output": "caselaw-qps-vs-recall.png", + "title": "Caselaw (diffuse attribute \u2014 doc_id)", + "points": { + "Standard": {"recall": 96.91, "qps": 160.3}, + "Queue": {"recall": 98.32, "qps": 84.3}, + "Design A": {"recall": 99.25, "qps": 128.8}, + "Design B": {"recall": 99.25, "qps": 105.9} + } + }, + { + "output": "caselaw-jurisdiction-qps-vs-recall.png", + "title": "Caselaw (concentrated + distance-correlated \u2014 court_jurisdiction)", + "points": { + "Standard": {"recall": 35.91, "qps": 165.6}, + "Queue": {"recall": 79.12, "qps": 152.1}, + "Design A": {"recall": 72.59, "qps": 174.5}, + "Design B": {"recall": 78.97, "qps": 96.0} + } + }, + { + "output": "yfcc-qps-vs-recall.png", + "title": "YFCC (concentrated attribute \u2014 camera)", + "points": { + "Standard": {"recall": 49.73, "qps": 227.1}, + "Queue": {"recall": 81.78, "qps": 186.1}, + "Design A": {"recall": 93.98, "qps": 192.6}, + "Design B": {"recall": 96.92, "qps": 111.9} + } + } + ] +} diff --git a/assets/diverse-search/enron-qps-vs-recall.png b/assets/diverse-search/enron-qps-vs-recall.png new file mode 100644 index 000000000..d72f15803 Binary files /dev/null and b/assets/diverse-search/enron-qps-vs-recall.png differ diff --git a/assets/diverse-search/plot_qps_vs_recall.py b/assets/diverse-search/plot_qps_vs_recall.py new file mode 100644 index 000000000..0d7214617 --- /dev/null +++ b/assets/diverse-search/plot_qps_vs_recall.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Generate the QPS-vs-recall scatter charts for the diverse-search benchmark doc. + +Reads chart_data.json (recall/QPS at L=100 for each method and dataset) and +writes one PNG per dataset next to this script. Kept in the repo so the charts +can be regenerated whenever the benchmark numbers change. + +Usage: + python plot_qps_vs_recall.py + +Requires matplotlib (pip install matplotlib). +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import matplotlib.pyplot as plt + +HERE = Path(__file__).resolve().parent +DATA_FILE = HERE / "chart_data.json" + + +def main() -> None: + data = json.loads(DATA_FILE.read_text(encoding="utf-8")) + methods = data["methods"] + colors = data["colors"] + + for chart in data["charts"]: + fig, ax = plt.subplots(figsize=(9, 6)) + + for method in methods: + point = chart["points"][method] + recall = point["recall"] + qps = point["qps"] + ax.scatter( + recall, + qps, + s=160, + color=colors[method], + label=method, + zorder=3, + ) + ax.annotate( + method, + (recall, qps), + textcoords="offset points", + xytext=(8, 8), + color=colors[method], + fontsize=10, + ) + + ax.set_title(f"{chart['title']}\nQPS vs recall at L=100") + ax.set_xlabel("Recall @10 (%) - higher is better") + ax.set_ylabel("QPS - higher is better") + ax.grid(True, color="0.9", zorder=0) + ax.legend(title="Approach", loc="lower left") + + out_path = HERE / chart["output"] + fig.tight_layout() + fig.savefig(out_path, dpi=100) + plt.close(fig) + print(f"wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/assets/diverse-search/yfcc-qps-vs-recall.png b/assets/diverse-search/yfcc-qps-vs-recall.png new file mode 100644 index 000000000..f1ca05dd3 Binary files /dev/null and b/assets/diverse-search/yfcc-qps-vs-recall.png differ diff --git a/diskann-benchmark/src/disk_index/search.rs b/diskann-benchmark/src/disk_index/search.rs index a1e84e79b..d45b199a4 100644 --- a/diskann-benchmark/src/disk_index/search.rs +++ b/diskann-benchmark/src/disk_index/search.rs @@ -37,7 +37,7 @@ use serde::{Deserialize, Serialize}; use crate::{ disk_index::json_spancollector::JsonSpanCollector, inputs::disk::{DiskIndexLoad, DiskSearchPhase}, - utils::{datafiles, SimilarityMeasure}, + utils::{attributes::FileAttributeProvider, datafiles, SimilarityMeasure}, }; #[derive(Serialize, Deserialize, Debug)] @@ -236,6 +236,35 @@ where logger.log_checkpoint("index_loaded"); + // Optional attribute-bucket diversity (Design A): loaded once and shared + // across all queries and search-list values. + let diverse_attribute_provider = match &search_params.attributes { + Some(attributes_file) => Some(std::sync::Arc::new(FileAttributeProvider::load( + attributes_file, + )?)), + None => None, + }; + let diverse_results_k = search_params.diverse_results_k; + let diverse_attribute_id = search_params.diverse_attribute_id; + // Effective output cap (per-bucket cap on the final result set). + let effective_diverse_results_k = + diverse_results_k.unwrap_or(search_params.recall_at as usize); + // Design B: when the diverse path is active and `adaptive_l` is configured, + // grow `L` per query from the observed bucket concentration. + let diverse_adaptive_l = search_params.search_mode.adaptive_l.as_ref().map(|a| { + diskann::graph::search::AdaptiveL::new(a.sample_count.into(), a.scale_factor) + .expect("validated adaptive L must construct") + }); + // Per-bucket cap for the yield sampler; defaults to the output cap when the + // config does not override it. + let diverse_yield_k = search_params + .search_mode + .adaptive_l + .as_ref() + .and_then(|a| a.yield_k) + .map(usize::from) + .unwrap_or(effective_diverse_results_k); + let pool = create_thread_pool(search_params.num_threads)?; let mut search_results_per_l = Vec::with_capacity(search_params.search_list.len()); let has_any_search_failed = AtomicBool::new(false); @@ -271,12 +300,23 @@ where // Construct the SearchMode from the JSON-driven // `adaptive_l` is now encapsulated in `DiskSearchMode`, so the // benchmark only supplies the per-query filter and post-processor. - let has_filter = search_params.vector_filters_file.is_some(); - let mode: SearchMode<'_> = search_params.search_mode.search_mode( - has_filter, - vf, - search_params.post_processor.as_ref(), - ); + let mode: SearchMode<'_> = match diverse_attribute_provider.as_ref() { + Some(provider) => SearchMode::diverse_attribute( + provider.clone(), + diverse_attribute_id, + effective_diverse_results_k, + diverse_yield_k, + diverse_adaptive_l.clone(), + ), + None => { + let has_filter = search_params.vector_filters_file.is_some(); + search_params.search_mode.search_mode( + has_filter, + vf, + search_params.post_processor.as_ref(), + ) + } + }; match searcher.search( q, diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 7ed521a87..c09fc45aa 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -170,6 +170,20 @@ pub(crate) struct DiskSearchPhase { pub(crate) num_nodes_to_cache: Option, pub(crate) search_io_limit: Option, pub(crate) post_processor: Option, + /// Plaintext file with one integer attribute per line; line `N` is vector `N`'s attribute. + /// + /// When present, the search phase runs attribute-bucket diversity via + /// post-processing (Design A): a plain greedy search over the top-`L` pool + /// followed by bucket selection keeping at most `diverse_results_k` results + /// per distinct attribute value, instead of the mode selected by + /// `search_mode`/`post_processor`. + pub(crate) attributes: Option, + /// The attribute dimension used for diversity (currently only a single attribute is used). + #[serde(default)] + pub(crate) diverse_attribute_id: usize, + /// The maximum number of results to keep per distinct attribute value. + #[serde(default)] + pub(crate) diverse_results_k: Option, } ///////// @@ -320,6 +334,18 @@ impl DiskSearchPhase { .context("invalid disk search post processor")?; } + if let Some(attributes) = self.attributes.as_mut() { + attributes + .resolve(checker) + .context("invalid attributes file")?; + match self.diverse_results_k { + Some(k) if k > 0 => {} + _ => anyhow::bail!( + "diverse_results_k must be a positive integer when attributes is set" + ), + } + } + Ok(()) } } @@ -367,6 +393,9 @@ impl Example for DiskIndexOperation { num_nodes_to_cache: None, search_io_limit: None, post_processor: None, + attributes: None, + diverse_attribute_id: 0, + diverse_results_k: None, }; Self { diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index 2137c6e57..873f74873 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -248,6 +248,12 @@ impl MultihopFilterSearchPhase { pub(crate) struct AdaptiveL { pub(crate) sample_count: NonZeroUsize, pub(crate) scale_factor: f64, + /// Per-bucket cap used by the Design-B diversity sampler when estimating + /// yield. Decoupled from `diverse_results_k` (the output cap); when unset it + /// defaults to the effective `diverse_results_k`. Only used by attribute + /// diversity search. + #[serde(default)] + pub(crate) yield_k: Option, } impl AdaptiveL { diff --git a/diskann-benchmark/src/utils/attributes.rs b/diskann-benchmark/src/utils/attributes.rs new file mode 100644 index 000000000..1d5751f9c --- /dev/null +++ b/diskann-benchmark/src/utils/attributes.rs @@ -0,0 +1,91 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! A file-backed [`AttributeValueProvider`] for diversity-aware benchmarks. + +use std::path::Path; + +use anyhow::Context; +use diskann::{neighbor::AttributeValueProvider, provider::HasId}; + +/// The reserved attribute bucket assigned to graph navigation nodes (frozen start points) +/// that fall outside the range of loaded per-vector attributes. +/// +/// Greedy graph search seeds traversal from the index's start points. Those points must +/// therefore carry an attribute, otherwise the diverse queue would drop them and search +/// would never expand beyond the entry node. Assigning them a dedicated bucket keeps them +/// traversable without merging them into any real attribute group. +const NAVIGATION_BUCKET: u32 = u32::MAX; + +/// An attribute value provider backed by a plaintext attribute file. +/// +/// The file is expected to contain one unsigned integer per line, where the value on the +/// `N`-th line (0-indexed) is the diversity attribute of the `N`-th vector. This matches the +/// on-disk layout produced by the labelling tools used elsewhere in the pipeline. +/// +/// Ids outside the range of loaded attributes (for example the graph's frozen start points) +/// are mapped to a reserved [`NAVIGATION_BUCKET`] so that greedy search can still traverse +/// the graph through them. +#[derive(Debug, Clone)] +pub(crate) struct FileAttributeProvider { + attributes: Vec, + num_buckets: usize, +} + +impl FileAttributeProvider { + /// Load attributes from the plaintext file at `path`. + /// + /// # Errors + /// + /// Returns an error if the file cannot be read or if any line fails to parse as a `u32`. + pub(crate) fn load(path: &Path) -> anyhow::Result { + let contents = std::fs::read_to_string(path) + .with_context(|| format!("while reading attribute file {}", path.display()))?; + + let attributes = contents + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .enumerate() + .map(|(line, value)| { + value.parse::().with_context(|| { + format!("invalid attribute value {value:?} on line {}", line + 1) + }) + }) + .collect::>>()?; + + let num_buckets = attributes + .iter() + .copied() + .collect::>() + .len(); + + Ok(Self { + attributes, + num_buckets, + }) + } +} + +impl HasId for FileAttributeProvider { + type Id = u32; +} + +impl AttributeValueProvider for FileAttributeProvider { + type Value = u32; + + fn get(&self, id: Self::Id) -> Option { + Some( + self.attributes + .get(id as usize) + .copied() + .unwrap_or(NAVIGATION_BUCKET), + ) + } + + fn num_buckets(&self) -> Option { + Some(self.num_buckets) + } +} diff --git a/diskann-benchmark/src/utils/mod.rs b/diskann-benchmark/src/utils/mod.rs index 9a25d2cc0..1a3be09cc 100644 --- a/diskann-benchmark/src/utils/mod.rs +++ b/diskann-benchmark/src/utils/mod.rs @@ -9,6 +9,8 @@ use diskann_benchmark_runner::{ }; use serde::{Deserialize, Serialize}; +#[cfg(feature = "disk-index")] +pub(crate) mod attributes; pub(crate) mod datafiles; pub(crate) mod filters; pub(crate) mod recall; diff --git a/diskann-bftree/Cargo.toml b/diskann-bftree/Cargo.toml index 79c8bb7ed..1dcaeb0b6 100644 --- a/diskann-bftree/Cargo.toml +++ b/diskann-bftree/Cargo.toml @@ -35,7 +35,6 @@ tokio = { workspace = true, features = ["full"] } [features] default = [] -experimental_diversity_search = ["diskann/experimental_diversity_search"] [lints.clippy] undocumented_unsafe_blocks = "warn" diff --git a/diskann-disk/Cargo.toml b/diskann-disk/Cargo.toml index d49fafa17..0782d5935 100644 --- a/diskann-disk/Cargo.toml +++ b/diskann-disk/Cargo.toml @@ -68,10 +68,6 @@ proptest.workspace = true default = [] perf_test = ["dep:opentelemetry"] virtual_storage = ["diskann-providers/virtual_storage"] -experimental_diversity_search = [ - "diskann/experimental_diversity_search", - "diskann-providers/experimental_diversity_search", -] [[bench]] name = "bench_main" diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 7b6cb6ab3..34f87d0bb 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -20,7 +20,7 @@ use diskann::{ self, ext::labeled::{self, QueryLabelProvider}, glue::{self, DefaultPostProcessor, SearchPostProcess, SearchStrategy}, - search::{AdaptiveL, InlineFilterSearch, Knn}, + search::{AdaptiveL, DiverseAdaptiveSearch, InlineFilterSearch, Knn}, search_output_buffer, DiskANNIndex, }, neighbor::{Neighbor, NeighborPriorityQueue}, @@ -50,7 +50,7 @@ use crate::{ data_model::{CachingStrategy, GraphHeader}, search::{ provider::disk_vertex_provider_factory::DiskVertexProviderFactory, - search_mode::SearchMode, + search_mode::{DynAttributeProvider, SearchMode}, traits::{VertexProvider, VertexProviderFactory}, }, storage::{api::AsyncDiskLoadContext, disk_index_reader::DiskIndexReader}, @@ -468,6 +468,120 @@ where } } +/// Post-processor implementing attribute-bucket diversity purely over the +/// candidate pool returned by a plain greedy search (Design A). +/// +/// Candidates are reranked with exact distances, sorted ascending, and then +/// greedily emitted while keeping at most `diverse_results_k` results per +/// distinct attribute value. The output buffer bounds the final count to `k`. +pub struct AttributeBucketDiversity<'a> { + filter: PostprocessStrategy<'a>, + provider: Arc, + diverse_results_k: usize, +} + +impl<'a> AttributeBucketDiversity<'a> { + pub fn new( + filter: PostprocessStrategy<'a>, + provider: Arc, + diverse_results_k: usize, + ) -> Self { + Self { + filter, + provider, + diverse_results_k, + } + } +} + +impl + SearchPostProcess< + DiskAccessor<'_, Data, VP>, + &[Data::VectorDataType], + ( + as DataProvider>::InternalId, + Data::AssociatedDataType, + ), + > for AttributeBucketDiversity<'_> +where + Data: GraphDataType, + VP: VertexProvider, +{ + type Error = ANNError; + async fn post_process( + &self, + accessor: &mut DiskAccessor<'_, Data, VP>, + query: &[Data::VectorDataType], + candidates: I, + output: &mut B, + ) -> Result + where + I: Iterator> + Send, + B: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + + Send + + ?Sized, + { + let provider = accessor.provider; + + // Reuse full-precision distances already computed during traversal + // (cached in `distance_cache`); only fetch vectors and recompute for + // candidates that missed the cache. + let mut uncached_ids = Vec::new(); + let mut reranked = { + let mut process = |n: u32| { + if let Some(entry) = accessor.scratch.distance_cache.get(&n) { + Some(Ok::<((u32, _), f32), ANNError>(((n, entry.1), entry.0))) + } else { + uncached_ids.push(n); + None + } + }; + match self.filter { + PostprocessStrategy::AcceptAll => candidates + .map(|n| n.id) + .filter_map(&mut process) + .collect::, _>>()?, + PostprocessStrategy::Apply(f) => candidates + .map(|n| n.id) + .filter(|id| f(id)) + .filter_map(&mut process) + .collect::, _>>()?, + } + }; + + if !uncached_ids.is_empty() { + ensure_vertex_loaded(&mut accessor.scratch.vertex_provider, &uncached_ids)?; + for n in &uncached_ids { + let v = accessor.scratch.vertex_provider.get_vector(n)?; + let d = provider.distance_comparer.evaluate_similarity(query, v); + let a = accessor.scratch.vertex_provider.get_associated_data(n)?; + reranked.push(((*n, *a), d)); + } + } + + // Sort by exact distance so the bucket selection walks candidates in + // true nearest-first order. + reranked + .sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); + + // Greedy bucket selection: keep at most `diverse_results_k` per bucket. + // `output.extend` truncates the stream to the buffer capacity (`k`). + let mut bucket_counts: HashMap = HashMap::new(); + let selected = reranked.into_iter().filter(|((id, _), _)| { + let bucket = self.provider.get(*id).unwrap_or(u32::MAX); + let count = bucket_counts.entry(bucket).or_insert(0); + if *count < self.diverse_results_k { + *count += 1; + true + } else { + false + } + }); + + Ok(output.extend(selected)) + } +} + impl SearchPostProcess< DiskAccessor<'_, Data, VP>, @@ -1190,6 +1304,57 @@ where &mut result_output_buffer, ))? } + SearchMode::DiverseAttribute { + provider, + diverse_attribute_id: _, + diverse_results_k, + yield_k, + adaptive_l, + } => { + // Attribute-bucket diversity: a greedy search collects the + // top-`L` pool, then `AttributeBucketDiversity` selects a + // bucket-diverse top-`k` from it in post-processing. + // + // Design A (`adaptive_l = None`): fixed `L`; the over-fetch is + // inherent in `L > k` and controlled by the search list size. + // + // Design B (`adaptive_l = Some(_)`): the traversal samples + // bucket concentration and grows `L` when few distinct buckets + // are seen, sizing the over-fetch per query. + let strategy = self.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let knn_search = Knn::new(k, l, beam_width)?; + let processor = AttributeBucketDiversity::new( + PostprocessStrategy::AcceptAll, + provider.clone(), + *diverse_results_k, + ); + match adaptive_l { + Some(adaptive_l) => { + let search = DiverseAdaptiveSearch::new( + knn_search, + provider.clone(), + *yield_k, + adaptive_l.clone(), + ); + self.runtime.block_on(self.index.search_with( + search, + &strategy, + processor, + &DefaultContext, + query, + &mut result_output_buffer, + ))? + } + None => self.runtime.block_on(self.index.search_with( + knn_search, + &strategy, + processor, + &DefaultContext, + query, + &mut result_output_buffer, + ))?, + } + } }; query_stats.total_comparisons = stats.cmps; query_stats.search_hops = stats.hops; @@ -1951,7 +2116,6 @@ mod disk_provider_tests { ); } - #[cfg(feature = "experimental_diversity_search")] #[test] fn test_disk_search_diversity_search() { use diskann::graph::DiverseSearchParams; diff --git a/diskann-disk/src/search/search_mode.rs b/diskann-disk/src/search/search_mode.rs index 49f7462f9..bdb9f1c41 100644 --- a/diskann-disk/src/search/search_mode.rs +++ b/diskann-disk/src/search/search_mode.rs @@ -11,7 +11,9 @@ use diskann::graph::ext::labeled::QueryLabelProvider; use diskann::graph::search::AdaptiveL; +use diskann::neighbor::AttributeValueProvider; use diskann_providers::model::graph::provider::DeterminantDiversityParams; +use std::sync::Arc; /// Owned closure used to filter vector IDs during disk search. /// @@ -19,6 +21,13 @@ use diskann_providers::model::graph::provider::DeterminantDiversityParams; /// on the disk path by construction). pub type SearchPredicate<'a> = Box bool + Send + Sync + 'a>; +/// Type-erased attribute provider used by [`SearchMode::DiverseAttribute`]. +/// +/// `SearchMode` is a non-generic enum, so the concrete provider is erased +/// behind this trait object. `Value` is pinned to `u32`, matching every +/// attribute-bucket provider in use. +pub type DynAttributeProvider = dyn AttributeValueProvider; + /// Top-level disk search mode. /// /// Three variants encode the algorithm + filter combination: @@ -51,6 +60,22 @@ pub enum SearchMode<'a> { filter: Option>, params: DeterminantDiversityParams, }, + + /// Attribute-bucket diversity done in post-processing: run a greedy graph + /// search over an enlarged candidate pool (`L`) and then greedily select, + /// in ascending distance order, at most `diverse_results_k` results per + /// distinct attribute value. + /// + /// `adaptive_l = Some(_)` enables Design B: the search samples bucket + /// concentration during traversal and grows `L` when few distinct buckets + /// are seen. `None` runs Design A with a fixed `L`. + DiverseAttribute { + provider: Arc, + diverse_attribute_id: usize, + diverse_results_k: usize, + yield_k: usize, + adaptive_l: Option, + }, } impl<'a> SearchMode<'a> { @@ -139,6 +164,33 @@ impl<'a> SearchMode<'a> { params, } } + + /// Attribute-bucket diversity via post-processing. The greedy search + /// collects the top-`L` pool, then a bucket-selection step keeps at most + /// `diverse_results_k` results per distinct attribute value. + /// + /// `adaptive_l = Some(_)` grows `L` mid-search from the observed bucket + /// concentration (Design B); `None` uses a fixed `L` (Design A). + /// + /// `yield_k` is the per-bucket cap used by the Design-B sampler when + /// estimating diversity yield; it is decoupled from `diverse_results_k` (the + /// output cap) so that a larger `yield_k` keeps `L` growing on + /// distance-correlated attributes. + pub fn diverse_attribute( + provider: Arc, + diverse_attribute_id: usize, + diverse_results_k: usize, + yield_k: usize, + adaptive_l: Option, + ) -> Self { + Self::DiverseAttribute { + provider, + diverse_attribute_id, + diverse_results_k, + yield_k, + adaptive_l, + } + } } #[cfg(test)] diff --git a/diskann-providers/Cargo.toml b/diskann-providers/Cargo.toml index 942c3919e..fc48e34aa 100644 --- a/diskann-providers/Cargo.toml +++ b/diskann-providers/Cargo.toml @@ -61,7 +61,6 @@ targets = [ default = [] perf_test = ["dep:opentelemetry"] testing = ["dep:tempfile"] -experimental_diversity_search = ["diskann/experimental_diversity_search"] virtual_storage = ["dep:vfs"] # Some 'cfg's in the source tree will be flagged by `cargo clippy -j 2 --workspace --no-deps --all-targets -- -D warnings` diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index 41389910e..a3526963c 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -2869,7 +2869,6 @@ pub(crate) mod tests { ); } - #[cfg(feature = "experimental_diversity_search")] #[tokio::test] async fn test_inmemory_search_diversity_search() { use diskann::neighbor::AttributeValueProvider; diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 14f6ba074..d5abf33fc 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -61,6 +61,3 @@ tracing = ["dep:tracing"] # Enable testing providers. testing = ["dep:dashmap"] - -# Enable experimental diversity search -experimental_diversity_search = [] diff --git a/diskann/src/graph/misc.rs b/diskann/src/graph/misc.rs index 067bd2d21..186f55406 100644 --- a/diskann/src/graph/misc.rs +++ b/diskann/src/graph/misc.rs @@ -32,7 +32,6 @@ pub enum InplaceDeleteMethod { } // Parameters for diverse search -#[cfg(feature = "experimental_diversity_search")] #[derive(Clone, Debug)] pub struct DiverseSearchParams

where @@ -43,7 +42,6 @@ where pub attribute_provider: std::sync::Arc

, } -#[cfg(feature = "experimental_diversity_search")] impl

DiverseSearchParams

where P: crate::neighbor::AttributeValueProvider, diff --git a/diskann/src/graph/mod.rs b/diskann/src/graph/mod.rs index 6bf4be7dd..e993b5258 100644 --- a/diskann/src/graph/mod.rs +++ b/diskann/src/graph/mod.rs @@ -23,7 +23,6 @@ pub use start_point::{SampleableForStart, StartPointStrategy}; mod misc; pub use misc::{ConsolidateKind, InplaceDeleteMethod}; -#[cfg(feature = "experimental_diversity_search")] pub use misc::DiverseSearchParams; pub mod glue; diff --git a/diskann/src/graph/search/diverse_adaptive_search.rs b/diskann/src/graph/search/diverse_adaptive_search.rs new file mode 100644 index 000000000..af0075a57 --- /dev/null +++ b/diskann/src/graph/search/diverse_adaptive_search.rs @@ -0,0 +1,496 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Attribute-diversity search with adaptive-L over-fetch (Design B). +//! +//! This is a plain greedy graph search whose only addition over [`Knn`] is a +//! traversal-time sampler that estimates *bucket concentration* — how many +//! distinct attribute buckets appear among the nodes visited so far — and grows +//! the search list `L` accordingly. The actual bucket-diverse top-`k` selection +//! is still performed by a downstream post-processor over the enlarged pool. +//! +//! # Motivation +//! +//! Post-processing-only diversity (Design A) relies on the caller picking an `L` +//! large enough that the top-`L` pool contains enough distinct buckets to fill a +//! diverse top-`k`. When a query lands in a region dominated by a few buckets, +//! a fixed `L` under-fetches and recall suffers. Design B sizes `L` per query +//! from the observed bucket concentration. +//! +//! # Yield metric +//! +//! The diversity *yield* of a sample measures how close the sampled nodes are +//! to a maximally diverse set, normalized to `[0, 1]`: +//! +//! ```text +//! yield = (Σ_b min(count_b, k)) / min(sample_visited, num_buckets · k) +//! ``` +//! +//! where `count_b` is the number of sampled nodes in bucket `b`, `k` is +//! `yield_k` (the per-bucket cap used for *this* sampling estimate), `num_buckets` +//! is the total number of distinct attribute buckets in the dataset (reported +//! by [`AttributeValueProvider::num_buckets`]), and `sample_visited` is the +//! number of nodes sampled. The denominator is the tight upper bound on the +//! numerator, coupling the yield to all three quantities so it spans the full +//! `[0, 1]` range regardless of how few buckets the dataset has or how small +//! the sample is: +//! +//! * few buckets, large sample → denominator is `num_buckets · k`, so covering +//! every bucket reaches `1.0`; +//! * many buckets, small sample → denominator is `sample_visited`, so a fully +//! distinct sample reaches `1.0` and diffuse data leaves `L` unchanged. +//! +//! `yield_k` is decoupled from the downstream `diverse_results_k` (the per-bucket +//! cap on the *final* result set). Setting `yield_k > diverse_results_k` requires +//! the sample to see several points per bucket before the yield saturates, which +//! keeps `L` growing on distance-correlated attributes — there the first point in +//! each bucket is often far from the query, so a low output `k` would otherwise +//! declare the buckets "covered" before the closer attribute-satisfying points +//! are reached. +//! +//! Unlike inline-filter specificity — where a value of `0.5` over a large + +//! sample is already good enough to stop enlarging `L` — a diversity yield of +//! `0.5` means only ~50% of the achievable diversity has been observed so far, +//! so `L` should still grow substantially. `L` is therefore scaled from the +//! *deficit* `1 − yield`: a yield near `1.0` leaves `L` unchanged, while a low +//! yield grows it toward the configured maximum multiplier. When the provider +//! cannot report `num_buckets`, the denominator falls back to `sample_visited` +//! alone: `L` still grows on concentrated samples (never worse than the +//! fixed-`L` Design A), though without the bucket total it may over-fetch on +//! few-bucket datasets. + +use std::sync::Arc; + +use diskann_utils::future::SendFuture; +use hashbrown::HashMap; + +use super::{AdaptiveL, Knn, Search, scratch::SearchScratch}; +use crate::{ + ANNResult, + error::IntoANNResult, + graph::{ + glue::{self, SearchAccessor, SearchPostProcess, SearchStrategy}, + index::{DiskANNIndex, InternalSearchStats, SearchStats}, + search_output_buffer::SearchOutputBuffer, + }, + neighbor::{AttributeValueProvider, Neighbor}, + provider::DataProvider, + utils::VectorId, +}; + +/// Attribute-diversity search with adaptive-L over-fetch. +/// +/// Runs a plain greedy beam search, sampling attribute buckets during traversal +/// to grow `L` when bucket concentration is low. The final bucket-diverse +/// selection is delegated to the post-processor supplied at search time. +pub struct DiverseAdaptiveSearch

+where + P: AttributeValueProvider + ?Sized, +{ + inner: Knn, + provider: Arc

, + yield_k: usize, + adaptive_l: AdaptiveL, +} + +impl

DiverseAdaptiveSearch

+where + P: AttributeValueProvider + ?Sized, +{ + /// Create new adaptive diverse search parameters. + /// + /// `yield_k` is the per-bucket cap used when estimating the diversity + /// *yield* during traversal. It is deliberately decoupled from the + /// downstream `diverse_results_k` (the per-bucket cap on the final result + /// set): requiring the sample to see several points per bucket before the + /// yield saturates keeps `L` growing on distance-correlated attributes, + /// where the first point in each bucket is often far from the query. + pub fn new( + inner: Knn, + provider: Arc

, + yield_k: usize, + adaptive_l: AdaptiveL, + ) -> Self { + Self { + inner, + provider, + yield_k, + adaptive_l, + } + } +} + +impl<'a, DP, S, T, P> Search<'a, DP, S, T> for DiverseAdaptiveSearch

+where + DP: DataProvider, + S: SearchStrategy<'a, DP, T, SearchAccessor: SearchAccessor>, + T: Copy + Send + Sync, + P: AttributeValueProvider + ?Sized, +{ + 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: 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_starting_points = accessor.num_starting_points().await?; + let mut scratch = index.search_scratch(self.inner.l_value().get(), num_starting_points); + + let stats = diverse_adaptive_search_internal( + index.max_degree_with_slack(), + &self.inner, + self.provider.as_ref(), + self.yield_k, + &self.adaptive_l, + &mut accessor, + &mut scratch, + ) + .await?; + + let result_count = processor + .post_process(&mut accessor, query, scratch.best.iter(), output) + .await + .into_ann_result()?; + + Ok(stats.finish(result_count as u32)) + } + } +} + +async fn diverse_adaptive_search_internal( + max_degree_with_slack: usize, + search_params: &Knn, + provider: &P, + yield_k: usize, + adaptive_l: &AdaptiveL, + accessor: &mut A, + scratch: &mut SearchScratch, +) -> ANNResult +where + I: VectorId, + A: SearchAccessor, + P: AttributeValueProvider + ?Sized, +{ + let beam_width = search_params.beam_width().get(); + let l_search = search_params.l_value().get(); + + // Total distinct buckets in the dataset, when the provider can report it. + // When unknown, the yield normalizes by the sample size alone (the + // pre-normalization heuristic): still over-fetches on concentrated samples, + // but may grow `L` more than necessary on few-bucket datasets. + let num_buckets = provider.num_buckets(); + + // Bucket-concentration sampling state. Sampling stops once L is (re)sized. + let mut bucket_counts: HashMap = HashMap::new(); + let mut sample_visited: usize = 0; + let mut l_adjusted = false; + + accessor + .start_point_distances(|id, distance| { + scratch.visited.insert(id); + scratch.best.insert(Neighbor::new(id, distance)); + scratch.cmps += 1; + }) + .await?; + + let mut neighbors = Vec::with_capacity(max_degree_with_slack); + + while scratch.best.has_notvisited_node() && !accessor.terminate_early() { + scratch.beam_nodes.clear(); + + while scratch.beam_nodes.len() < beam_width + && let Some(closest_node) = scratch.best.closest_notvisited() + { + scratch.beam_nodes.push(closest_node.id); + } + + if scratch.beam_nodes.is_empty() { + break; + } + + neighbors.clear(); + accessor + .expand_beam( + scratch.beam_nodes.iter().copied(), + glue::NotInMut::new(&mut scratch.visited), + |id, distance| neighbors.push(Neighbor::new(id, distance)), + ) + .await?; + + for neighbor in neighbors.iter() { + scratch.best.insert(*neighbor); + + // Sample the bucket of each newly visited node until L is resized. + if !l_adjusted { + if let Some(bucket) = provider.get(neighbor.id) { + *bucket_counts.entry(bucket).or_insert(0) += 1; + } + sample_visited += 1; + } + } + + scratch.cmps += neighbors.len() as u32; + scratch.hops += scratch.beam_nodes.len() as u32; + + // Once enough nodes are sampled, estimate the bucket-diversity yield + // and grow L from its deficit. `matched` is the number of "useful" + // diverse slots observed (Σ_b min(count_b, k)); dividing by the maximum + // achievable slots normalizes the yield to [0, 1]. A low yield => far + // from a full diverse set => larger L. + if !l_adjusted && sample_visited >= adaptive_l.sample_count() { + l_adjusted = true; + let matched = diverse_matched_slots(bucket_counts.values().copied(), yield_k); + let diversity_yield = diverse_yield(matched, num_buckets, yield_k, sample_visited); + let new_l = + compute_diverse_adaptive_l(l_search, diversity_yield, adaptive_l.scale_factor()); + if new_l > l_search { + scratch.resize(new_l); + } + } + } + + Ok(InternalSearchStats { + cmps: scratch.cmps, + hops: scratch.hops, + range_search_second_round: false, + }) +} + +/// Count the "useful" diverse slots observed in a bucket-count sample. +/// +/// This is `Σ_b min(count_b, k)`: each distinct bucket contributes at most `k` +/// (`yield_k`) slots. `yield_k` is the sampling cap, which may exceed the +/// downstream per-bucket result cap so that concentrated buckets keep counting +/// toward the deficit. +fn diverse_matched_slots(bucket_counts: I, yield_k: usize) -> usize +where + I: IntoIterator, +{ + bucket_counts + .into_iter() + .map(|count| count.min(yield_k)) + .sum() +} + +/// Normalize observed diverse slots into a `[0, 1]` yield. +/// +/// `matched` is `Σ_b min(count_b, k)` (see [`diverse_matched_slots`]). Its tight +/// upper bound is `min(sample_visited, num_buckets · k)`: `matched` can never +/// exceed the number of nodes sampled, nor the total capped slots across every +/// bucket. Dividing by that bound couples the yield to all three quantities that +/// bound the achievable diversity — `yield_k`, `num_buckets`, and the +/// sample size — so the yield spans the full `[0, 1]` range in both regimes: +/// +/// * few buckets, large sample (`num_buckets · k ≤ sample_visited`): the +/// denominator is `num_buckets · k`, so seeing every bucket up to its cap +/// reaches `1.0`. +/// * many buckets, small sample (`sample_visited < num_buckets · k`): the +/// denominator is `sample_visited`, so a fully distinct sample reaches `1.0` +/// and diffuse data keeps `L` unchanged (Design-A behavior). +/// +/// When `num_buckets` is `None` the denominator falls back to `sample_visited`: +/// the sample itself is the only available signal, so a concentrated sample +/// still grows `L` while a fully distinct one leaves it unchanged. This may +/// over-grow `L` on few-bucket datasets (the sample can look concentrated even +/// after every bucket has been seen), but never under-fetches relative to the +/// fixed-`L` Design-A baseline. +/// +/// The result is clamped to `[0, 1]`; a zero denominator yields `1.0` (nothing +/// to diversify over, so no over-fetch is warranted). +fn diverse_yield( + matched: usize, + num_buckets: Option, + yield_k: usize, + sample_visited: usize, +) -> f64 { + let achievable = match num_buckets { + Some(n) => n.saturating_mul(yield_k).min(sample_visited), + None => sample_visited, + }; + if achievable == 0 { + return 1.0; + } + (matched as f64 / achievable as f64).clamp(0.0, 1.0) +} + +/// Compute adaptive L from the observed diversity yield. +/// +/// `L` is scaled linearly from the yield *deficit* `1 − yield`, so higher yields +/// (closer to a full diverse set) grow `L` less: +/// yield = 1.0 → 1× L (no change; the sample already covers the buckets) +/// yield = 0.5 → midway between 1× and `max_multiplier`× L +/// yield = 0.0 → `max_multiplier`× L (maximum expansion) +/// +/// This differs from inline-filter [`compute_adaptive_l`](super::inline_filter_search) +/// on purpose: filter specificity of `0.5` is already good, whereas a diversity +/// yield of `0.5` means only half of the achievable diversity has been seen and +/// warrants substantial further search. +/// +/// The multiplier is clamped to `[1, max_multiplier]`. +pub(crate) fn compute_diverse_adaptive_l( + base_l: usize, + diversity_yield: f64, + max_multiplier: f64, +) -> usize { + let deficit = (1.0 - diversity_yield).clamp(0.0, 1.0); + let multiplier = (1.0 + (max_multiplier - 1.0) * deficit).clamp(1.0, max_multiplier); + (base_l as f64 * multiplier) as usize +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matched_slots_all_distinct_k1() { + // Every node in its own bucket, k=1: every node is a useful slot. + let counts = [1, 1, 1, 1]; + assert_eq!(diverse_matched_slots(counts, 1), 4); + } + + #[test] + fn matched_slots_single_bucket_k1() { + // All nodes in one bucket, k=1: only one useful slot. + let counts = [10]; + assert_eq!(diverse_matched_slots(counts, 1), 1); + } + + #[test] + fn matched_slots_caps_per_bucket() { + // k=2: each bucket contributes at most 2. + let counts = [5, 1, 3]; + assert_eq!(diverse_matched_slots(counts, 2), 2 + 1 + 2); + } + + #[test] + fn matched_slots_empty_sample() { + let counts: [usize; 0] = []; + assert_eq!(diverse_matched_slots(counts, 1), 0); + } + + #[test] + fn yield_is_normalized_to_unit_range() { + // 60 buckets, k=1, large sample (1000): denominator is num_buckets·k=60. + // Seeing all 60 distinct buckets is full diversity, independent of how + // many nodes were visited to get there. + let all_distinct = std::iter::repeat_n(1usize, 60); + let matched = diverse_matched_slots(all_distinct, 1); + assert_eq!(diverse_yield(matched, Some(60), 1, 1000), 1.0); + + // Half the buckets covered => yield 0.5. + let half = std::iter::repeat_n(1usize, 30); + let matched = diverse_matched_slots(half, 1); + assert_eq!(diverse_yield(matched, Some(60), 1, 1000), 0.5); + + // Everything in a single bucket => near-zero yield. + let one_bucket = std::iter::once(1000usize); + let matched = diverse_matched_slots(one_bucket, 1); + assert!((diverse_yield(matched, Some(60), 1, 1000) - 1.0 / 60.0).abs() < 1e-9); + } + + #[test] + fn yield_many_buckets_small_sample_reaches_one() { + // Diffuse dataset: far more buckets than the sample size. The + // denominator is capped at sample_visited, so a fully distinct sample + // yields 1.0 (no over-fetch) rather than being stuck near zero. + let all_distinct = std::iter::repeat_n(1usize, 100); + let matched = diverse_matched_slots(all_distinct, 1); + assert_eq!(diverse_yield(matched, Some(12_849), 1, 100), 1.0); + + // Same dataset, but the sample piles into a few buckets => low yield. + let concentrated = [40usize, 30, 30]; + let matched = diverse_matched_slots(concentrated, 1); + assert!((diverse_yield(matched, Some(12_849), 1, 100) - 3.0 / 100.0).abs() < 1e-9); + } + + #[test] + fn yield_caps_per_bucket_with_k() { + // k=2, 10 buckets, sample 20: denominator is num_buckets·k=20. One + // saturated bucket (>=2) plus one singleton contributes 2 + 1 = 3 + // useful slots => yield 3/20. + let counts = [5usize, 1]; + let matched = diverse_matched_slots(counts, 2); + assert_eq!(matched, 3); + assert!((diverse_yield(matched, Some(10), 2, 20) - 3.0 / 20.0).abs() < 1e-9); + } + + #[test] + fn yield_zero_buckets_is_full() { + // No buckets to diversify over => treat as fully satisfied, no growth. + assert_eq!(diverse_yield(0, Some(0), 1, 100), 1.0); + } + + #[test] + fn yield_unknown_buckets_uses_sample_size() { + // Provider cannot report num_buckets: denominator is sample_visited, so + // the yield reduces to the fraction of the sample that is useful + // (the pre-normalization heuristic). + // Fully distinct sample => yield 1.0, no growth. + let all_distinct = std::iter::repeat_n(1usize, 100); + let matched = diverse_matched_slots(all_distinct, 1); + assert_eq!(diverse_yield(matched, None, 1, 100), 1.0); + + // Concentrated sample => low yield => grow L, matching the old behavior + // even though num_buckets is unavailable. + let concentrated = [40usize, 30, 30]; + let matched = diverse_matched_slots(concentrated, 1); + assert!((diverse_yield(matched, None, 1, 100) - 3.0 / 100.0).abs() < 1e-9); + } + + #[test] + fn diverse_adaptive_l_scales_from_deficit() { + let base_l = 100; + let max_multiplier = 8.0; + + // Full diversity => no growth. + assert_eq!( + compute_diverse_adaptive_l(base_l, 1.0, max_multiplier), + base_l + ); + + // Zero yield => maximum expansion. + assert_eq!( + compute_diverse_adaptive_l(base_l, 0.0, max_multiplier), + base_l * 8 + ); + + // A yield of 0.5 is "bad" for diversity and must grow L a lot: halfway + // between 1x and 8x => 4.5x. + assert_eq!( + compute_diverse_adaptive_l(base_l, 0.5, max_multiplier), + (base_l as f64 * 4.5) as usize + ); + + // Out-of-range yields are clamped. + assert_eq!( + compute_diverse_adaptive_l(base_l, 1.5, max_multiplier), + base_l + ); + assert_eq!( + compute_diverse_adaptive_l(base_l, -0.5, max_multiplier), + base_l * 8 + ); + } +} diff --git a/diskann/src/graph/search/inline_filter_search.rs b/diskann/src/graph/search/inline_filter_search.rs index c535c87b6..b3cf05367 100644 --- a/diskann/src/graph/search/inline_filter_search.rs +++ b/diskann/src/graph/search/inline_filter_search.rs @@ -59,6 +59,18 @@ impl AdaptiveL { scale_factor, }) } + + /// Number of nodes to sample before estimating specificity and resizing L. + #[inline] + pub(crate) fn sample_count(&self) -> usize { + self.sample_count + } + + /// Maximum multiplier applied to the base L. + #[inline] + pub(crate) fn scale_factor(&self) -> f64 { + self.scale_factor + } } /// Inline filtered search: a standard k-NN search with an additional step @@ -300,7 +312,12 @@ where /// 0 matches in sample → `max_multiplier`× L (maximum expansion) /// /// Clamped to [1×, max_multiplier] range. -fn compute_adaptive_l(base_l: usize, visited: usize, matched: usize, max_multiplier: f64) -> usize { +pub(crate) fn compute_adaptive_l( + base_l: usize, + visited: usize, + matched: usize, + max_multiplier: f64, +) -> usize { if matched == 0 || visited == 0 { // No matches at all — use maximum multiplier return (base_l as f64 * max_multiplier) as usize; diff --git a/diskann/src/graph/search/mod.rs b/diskann/src/graph/search/mod.rs index 4e2aedc61..25aa3b142 100644 --- a/diskann/src/graph/search/mod.rs +++ b/diskann/src/graph/search/mod.rs @@ -117,9 +117,8 @@ pub use knn_search::{Knn, KnnSearchError, RecordedKnn}; pub use multihop_filter_search::MultihopFilterSearch; pub use range_search::{Range, RangeSearchError}; -// Feature-gated diverse search. -#[cfg(feature = "experimental_diversity_search")] mod diverse_search; - -#[cfg(feature = "experimental_diversity_search")] pub use diverse_search::Diverse; + +mod diverse_adaptive_search; +pub use diverse_adaptive_search::DiverseAdaptiveSearch; diff --git a/diskann/src/neighbor/diverse_priority_queue.rs b/diskann/src/neighbor/diverse_priority_queue.rs index 93f355500..a3febe9d1 100644 --- a/diskann/src/neighbor/diverse_priority_queue.rs +++ b/diskann/src/neighbor/diverse_priority_queue.rs @@ -276,6 +276,18 @@ pub trait AttributeValueProvider: crate::provider::HasId + Send + Sync + std::fm /// # Returns /// * `Option` - The attribute value if it exists, None otherwise fn get(&self, id: Self::Id) -> Option; + + /// Total number of distinct attribute buckets across the dataset. + /// + /// Adaptive-L diverse search uses this to normalize the observed diversity + /// yield into `[0, 1]`: without it, the yield of a sample is bounded by + /// `num_buckets / sample_count` and can never approach `1.0` when the + /// dataset has few buckets. Providers that cannot report this cheaply + /// return `None`, which disables adaptive-L growth (`L` stays at its base + /// value). + fn num_buckets(&self) -> Option { + None + } } #[cfg(test)] diff --git a/diskann/src/neighbor/mod.rs b/diskann/src/neighbor/mod.rs index 391a5435e..139559d6f 100644 --- a/diskann/src/neighbor/mod.rs +++ b/diskann/src/neighbor/mod.rs @@ -12,9 +12,7 @@ use crate::graph::{SearchOutputBuffer, search_output_buffer}; mod queue; pub use queue::{NeighborPriorityQueue, NeighborPriorityQueueIdType, NeighborQueue}; -#[cfg(feature = "experimental_diversity_search")] mod diverse_priority_queue; -#[cfg(feature = "experimental_diversity_search")] pub use diverse_priority_queue::{ Attribute, AttributeValueProvider, DiverseNeighborQueue, VectorIdWithAttribute, }; diff --git a/diskann/src/neighbor/queue.rs b/diskann/src/neighbor/queue.rs index 7c1047f4a..008cbdedf 100644 --- a/diskann/src/neighbor/queue.rs +++ b/diskann/src/neighbor/queue.rs @@ -421,7 +421,6 @@ impl NeighborPriorityQueue { /// /// # Arguments /// * `f` - A predicate that returns `true` for items to keep - #[cfg(feature = "experimental_diversity_search")] pub fn retain(&mut self, mut f: F) where F: FnMut(&Neighbor) -> bool, @@ -460,7 +459,6 @@ impl NeighborPriorityQueue { /// /// # Arguments /// * `len` - The new length of the queue - #[cfg(feature = "experimental_diversity_search")] pub fn truncate(&mut self, len: usize) { let new_size = len; if new_size < self.size { @@ -1291,7 +1289,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_retain() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1327,7 +1324,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_retain_empty() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1337,7 +1333,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_retain_remove_all() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1353,7 +1348,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_retain_remove_none() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1371,7 +1365,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_retain_resets_visited_state() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1411,7 +1404,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_truncate() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1433,7 +1425,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_truncate_larger_size() { let mut queue = NeighborPriorityQueue::::new(10); @@ -1447,7 +1438,6 @@ mod neighbor_priority_queue_test { } #[test] - #[cfg(feature = "experimental_diversity_search")] fn test_truncate_with_cursor() { let mut queue = NeighborPriorityQueue::::new(10); diff --git a/diverse-search-results.md b/diverse-search-results.md new file mode 100644 index 000000000..ab2d6b6c0 --- /dev/null +++ b/diverse-search-results.md @@ -0,0 +1,226 @@ +# Attribute-Diversity Search — Benchmark Results + +Comparison of attribute-bucket diverse search strategies across several dataset +and attribute regimes, all +evaluated against a **diverse ground truth** (GT built so that no more than +`diverse_results_k` results share the same attribute bucket). + +## Setup + +Common parameters for every run below: + +| Parameter | Value | +|---|---| +| Index | On-disk DiskANN, 200,000 base vectors | +| `recall_at` (k) | 10 | +| `diverse_results_k` | 1 (at most one result per attribute bucket) | +| `beam_width` | 4 | +| `num_threads` | 1 | +| Distance | squared L2 | +| Search lists (L) | 20, 40, 80, 100 | +| Ground truth | diverse GT (`groundtruth_diverse_k100.bin`) | + +> **QPS methodology.** Recall, IO counts, comparisons and hops are exact and +> fully reproducible. QPS, however, is single-threaded and disk-bound, and on the +> dev machine used here (shared with VS Code, browsers, Teams) it varied 10–30× +> run-to-run purely from disk-wait contention, with IO counts unchanged. The QPS +> figures below are the **best (least-contended) of several runs per cell**, +> which is closest to the true compute-/IO-bound cost. Treat QPS as indicative, +> not precise; the recall and IO columns are the authoritative comparison. + +### Datasets + +| Dataset | Dim | Diversity attribute | Attribute character | +|---|---|---|---| +| Enron | 1,369 | attribute id 0 | **Concentrated** — many vectors share a bucket | +| Caselaw | 1,536 | `doc_id` | **Diffuse, uncorrelated** — 200K vectors span 172,518 docs (18,383 multi-chunk) | +| Caselaw | 1,536 | `court_jurisdiction` | **Concentrated + distance-correlated** — 60 courts, largest bucket 47,358 (23.7%), median 1,668 | +| YFCC | 192 | `camera` model | **Concentrated, ~uncorrelated** — 67 distinct cameras, largest bucket 67,783 (33.9%), median bucket 110 | + +Two axes matter for diversity search: how **concentrated** the buckets are (does +the sampler need to over-fetch?) and how **correlated** the attribute is with +distance (are same-bucket vectors clustered together in the search pool?). The +caselaw `court_jurisdiction` attribute is the interesting case: cases from the +same court share citations and legal boilerplate, so their embeddings cluster — +the attribute is *both* concentrated and strongly correlated with distance. + +### Strategies compared + +- **Standard** — plain KNN search (no diversity enforcement), scored against the + diverse GT. Shows how far an ordinary search is from the diverse target. +- **Queue** — diversity enforced *during* graph traversal using a diversity-aware + neighbor queue with approximate (PQ) distances. Branch + `u/narendatha/diverse-search-benchmark`. +- **Design A (post-process)** — plain KNN over the top-L pool, then a greedy + bucket selection keeping ≤ `diverse_results_k` per bucket, using + full-precision distances. Fixed L (no over-fetch). +- **Design B (adaptive-L)** — like Design A, but a greedy walk first samples + bucket concentration; if buckets are concentrated it grows L (over-fetches) + before the same post-processing step. When buckets are diffuse it returns a 1× + multiplier and behaves exactly like Design A. + +--- + +## Enron (concentrated attribute) + +| L | Standard recall | Queue recall | Design A recall | Design B recall | Std QPS | Queue QPS | A QPS | B QPS | +|---|---|---|---|---|---|---|---|---| +| 20 | 55.71 | 61.11 | 59.98 | **79.92** | 375.5 | 261 | 542.9 | 385.6 | +| 40 | 64.28 | 71.08 | 73.93 | **87.38** | 245.3 | 221 | 268.0 | 183.2 | +| 80 | 69.67 | 76.46 | 84.00 | **92.12** | 173.2 | 116 | 165.6 | 107.2 | +| 100 | 70.72 | 77.92 | 86.05 | **93.01** | 142.6 | 100 | 159.2 | 53 | + +**QPS vs recall @ L=100** (up and to the right is better) + +![Enron QPS vs recall](assets/diverse-search/enron-qps-vs-recall.png) + +**Observations** +- Standard search is far below every diversity-aware method (it does not enforce + the one-per-bucket constraint), confirming the diverse GT is meaningfully + different from ordinary top-k. +- **Design B wins recall at every L** (+19 over queue at L=20), because the + attribute is concentrated: the adaptive sampler detects this and over-fetches a + larger pool before reranking. +- Design A beats the queue on both recall and QPS at every L. +- Design B carries a throughput cost (extra pool fetch + more IO), but delivers + the highest recall. At L=100 its IO count (320) is well over double Design A's + (112), so its lower QPS there is real, not measurement noise. +- Pareto note: Design B @ L=20 (79.9 recall) beats Design A @ L=40 (73.9 recall) + at comparable throughput. + +--- + +## Caselaw (diffuse, uncorrelated attribute — `doc_id`) + +| L | Standard recall | Queue recall | Design A recall | Design B recall | Std QPS | Queue QPS | A QPS | B QPS | +|---|---|---|---|---|---|---|---|---| +| 20 | 90.61 | 90.76 | 92.04 | 92.04 | 554.8 | 306.7 | 448.8 | 238.4 | +| 40 | 95.43 | 96.53 | 97.53 | 97.53 | 267.5 | 207.0 | 192.5 | 151.0 | +| 80 | 96.73 | 98.13 | 99.06 | 99.06 | 187.0 | 138.4 | 132.4 | 101.3 | +| 100 | 96.91 | 98.32 | 99.25 | 99.25 | 160.3 | 84.3 | 128.8 | 105.9 | + +**QPS vs recall @ L=100** (up and to the right is better) + +![Caselaw QPS vs recall](assets/diverse-search/caselaw-qps-vs-recall.png) + +**Observations** +- **Design A and Design B have identical recall AND identical IO/comparison/hop + counts.** The `doc_id` attribute is so diffuse (almost every vector is its own + bucket) that the concentration sampler returns a 1× multiplier — Design B + performs no over-fetch. Its slightly lower QPS is the pure cost of the + concentration-sampling walk, which here yields no benefit. +- Even standard search is already close to the diverse GT here, because with + near-unique buckets the diverse GT is almost the same as the ordinary top-k. + The diversity methods add only ~1–2 recall points. +- Design A/B edge out the queue method on recall (~1–1.3 points) via + exact-distance reranking, and Design A is also competitive on QPS + (e.g. 448.8 vs queue 306.7 at L=20). + +--- + +## Caselaw (concentrated + distance-correlated attribute — `court_jurisdiction`) + +Same 200K caselaw vectors and index as above, but the diversity attribute is now +the natural `court_jurisdiction` field (60 US courts) instead of `doc_id`. This +attribute is **both** concentrated (largest bucket 23.7%, median 1,668) **and +strongly correlated with distance** — cases from the same court cluster together +in embedding space. This is the hardest and most realistic regime: a query's +nearest neighbors are dominated by a few jurisdictions, so enforcing one-per-bucket +forces the search to reach far past the plain top-L pool. + +| L | Standard recall | Queue recall | Design A recall | Design B recall | Std QPS | Queue QPS | A QPS | B QPS | +|---|---|---|---|---|---|---|---|---| +| 20 | 34.74 | 71.35 | 46.30 | 57.30 | 387.9 | 434.9 | 413.5 | 317.1 | +| 40 | 35.37 | 76.19 | 59.68 | 68.23 | 394.3 | 278.7 | 303.8 | 203.3 | +| 80 | 35.79 | 78.72 | 69.74 | 76.62 | 212.2 | 176.4 | 211.1 | 114.2 | +| 100 | 35.91 | 79.12 | 72.59 | 78.97 | 165.6 | 152.1 | 174.5 | 96.0 | + +**QPS vs recall @ L=100** (up and to the right is better) + +![Caselaw jurisdiction QPS vs recall](assets/diverse-search/caselaw-jurisdiction-qps-vs-recall.png) + +**Observations** +- **This regime inverts the earlier ranking: the queue beats Design A at every L** + (e.g. 71.4 vs 46.3 at L=20). When the attribute is correlated with distance, the + plain top-L pool that Design A reranks is *saturated* with the dominant + jurisdictions, so no amount of exact-distance reranking can recover the missing + buckets. The queue enforces diversity *during* traversal, steering the walk out + of the dominant cluster early — exactly what post-processing cannot do. +- **Design B recovers most of the gap** via adaptive over-fetch (57.3 vs A's 46.3 + at L=20) but stays just below the queue at every L, including L=100 (79.0 vs + 79.1) — and only closes this much by fetching a larger pool (IO 193 vs queue's + 114), so its QPS is the lowest of all methods there. On the QPS-vs-recall plot + the queue **dominates Design B** (equal-or-better recall at ~1.6× the + throughput). +- Standard search collapses to ~35% and never improves with L — with buckets both + concentrated and distance-correlated, the ordinary top-k is almost entirely the + wrong jurisdictions. +- Takeaway: **for a distance-correlated attribute, in-traversal diversity (queue) + is the efficient choice; post-processing (Design A/B) must over-fetch heavily to + compete, and even then only approaches — never beats — the queue on the + recall/QPS frontier.** This is the opposite of the uncorrelated concentrated + case (YFCC) where Design B is the clear winner. + +--- + +## YFCC (concentrated attribute — `camera` model) + +YFCC-100M image embeddings (200K subset, 192-dim). The diversity attribute is the +`camera` model tag: 67 distinct cameras, but very skewed — the largest bucket holds +33.9% of all vectors and the median bucket only 110, so this is a strongly +concentrated regime (like Enron). + +| L | Standard recall | Queue recall | Design A recall | Design B recall | Std QPS | Queue QPS | A QPS | B QPS | +|---|---|---|---|---|---|---|---|---| +| 20 | 49.35 | 79.26 | 66.92 | **88.44** | 347.0 | 592.6 | 499.4 | 442.6 | +| 40 | 49.60 | 81.23 | 83.54 | **94.20** | 499.7 | 397.3 | 441.5 | 261.8 | +| 80 | 49.72 | 81.60 | 92.38 | **96.45** | 256.5 | 240.2 | 248.3 | 140.6 | +| 100 | 49.73 | 81.78 | 93.98 | **96.92** | 227.1 | 186.1 | 192.6 | 111.9 | + +**QPS vs recall @ L=100** (up and to the right is better) + +![YFCC QPS vs recall](assets/diverse-search/yfcc-qps-vs-recall.png) + +**Observations** +- Standard search is stuck at ~49.5 recall for every L — with a third of the + corpus in one camera bucket, ordinary top-k returns many same-bucket neighbors + that the diverse GT rejects. Growing L does not help it at all. +- **Design B wins recall at every L** (+9 to +15 over the queue), because the + adaptive sampler detects the concentration and over-fetches before reranking. +- **The queue method plateaus at ~81%** — its approximate (PQ) in-traversal + diversity saturates and extra L barely helps. Design A overtakes it from L=40 + onward and reaches 94% at L=100; Design B is ahead of the queue at every L. +- Design B's higher recall again costs throughput (its IO roughly triples by + L=100: 286 vs Design A's 112), so its lower QPS is real work, not noise. +- Pareto note: Design B @ L=20 (88.4 recall) already beats Design A @ L=40 + (83.5) and the queue at any L, at competitive throughput. + +--- + +## Summary + +The datasets exercise the two axes that govern diverse search — bucket +**concentration** and attribute/distance **correlation**: + +- **Concentrated, ~uncorrelated attribute (Enron, YFCC):** big win for Design B — + adaptive over-fetch recovers substantial recall (+19 over queue on Enron at + L=20; +13 on YFCC at L=40) at a throughput cost. On YFCC the queue method + plateaus near 81% while Design B climbs to 97%. +- **Diffuse attribute (Caselaw `doc_id`):** the sampler detects abundant diversity + and Design B performs no over-fetch — it matches Design A on recall and IO (its + only cost is the sampling walk), and both slightly beat the queue method. +- **Concentrated + distance-correlated attribute (Caselaw `court_jurisdiction`):** + the ranking *inverts*. Because same-bucket vectors cluster in distance, the plain + top-L pool is saturated with the dominant buckets, so post-processing (Design A) + falls far behind the queue (46 vs 71 recall at L=20). The queue enforces diversity + during traversal and is the efficient choice; Design B claws back most of the gap + via heavy over-fetch but stays just below the queue on recall (79.0 vs 79.1 at + L=100) at ~1.6× the IO, so the queue dominates it on the recall/QPS frontier. + +Recommendation: the best strategy depends on the attribute. For **concentrated, +distance-uncorrelated** attributes, **Design B** wins (adaptive over-fetch, no +per-dataset L tuning). For **distance-correlated** attributes, in-traversal +diversity (the **queue**) is more efficient than any post-processing approach — +diversity has to be enforced during the walk, not after it. Design A remains a +cheap, effective choice whenever the top-L pool already contains enough distinct +buckets (diffuse or weakly-correlated attributes).