diff --git a/Cargo.lock b/Cargo.lock index 20d6ad95d..a921caa50 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -731,6 +731,7 @@ dependencies = [ "byteorder", "clap", "diskann", + "diskann-benchmark-core", "diskann-label-filter", "diskann-providers", "diskann-quantization", diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs index fbaee73df..072b3b1f6 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs @@ -176,8 +176,7 @@ impl RunBook { } /// Returns a reference to the stages in this runbook. - #[cfg(test)] - pub(super) fn stages(&self) -> &[Stage] { + pub fn stages(&self) -> &[Stage] { &self.stages } diff --git a/diskann-tools/Cargo.toml b/diskann-tools/Cargo.toml index 644ee1173..92bcbc949 100644 --- a/diskann-tools/Cargo.toml +++ b/diskann-tools/Cargo.toml @@ -15,6 +15,7 @@ clap = { workspace = true, features = ["derive"] } diskann-providers = { workspace = true, default-features = false } # see `linalg/Cargo.toml` diskann-vector = { workspace = true } diskann-utils = { workspace = true } +diskann-benchmark-core = { workspace = true, features = ["bigann"] } bytemuck.workspace = true num_cpus.workspace = true rayon.workspace = true diff --git a/diskann-tools/src/bin/compute_streaming_groundtruth.rs b/diskann-tools/src/bin/compute_streaming_groundtruth.rs new file mode 100644 index 000000000..3cf8661ba --- /dev/null +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -0,0 +1,342 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Computes KNN ground truth for every search stage in a BigANN-style runbook. +//! +//! The tool simulates the insert / replace / delete operations in the runbook, +//! tracking the set of active base-vector IDs at each search stage. For each +//! search stage it computes the exact top-k nearest neighbours for each query +//! against the currently active set, and writes the result to +//! `/step.gt` -- the naming convention expected by +//! [`diskann_benchmark_core::streaming::executors::bigann::ScanDirectory`]. + +use std::{ + collections::HashMap, + path::{Path, PathBuf}, + time::Instant, +}; + +use anyhow::Context; +use clap::Parser; +use diskann::neighbor::{Neighbor, NeighborPriorityQueue}; +use diskann::utils::VectorRepr; +use diskann_benchmark_core::streaming::executors::bigann::{FindGroundtruth, RunBook}; +use diskann_benchmark_core::streaming::{self, Executor}; +use diskann_providers::storage::{FileStorageProvider, StorageReadProvider}; +use diskann_tools::utils::{ + init_subscriber, write_ground_truth, CMDResult, CMDToolError, DataType, +}; +use diskann_utils::io::read_bin; +use diskann_utils::views::Matrix; +use diskann_vector::{distance::Metric, DistanceFunction}; +use rayon::prelude::*; + +fn main() -> CMDResult<()> { + init_subscriber(); + let args = Args::parse(); + match args.data_type { + DataType::Float => run::(&args), + DataType::Fp16 => run::(&args), + DataType::Uint8 => run::(&args), + DataType::Int8 => run::(&args), + } +} + +fn run(args: &Args) -> CMDResult<()> { + let storage = FileStorageProvider; + + tracing::info!("Loading dataset from {}", args.base_file); + let dataset = + read_bin::(&mut storage.open_reader(&args.base_file)?).map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + tracing::info!("Loading queries from {}", args.query_file); + let queries = + read_bin::(&mut storage.open_reader(&args.query_file)?).map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + let n_base = dataset.nrows(); + let n_queries = queries.nrows(); + let recall_at = args.recall_at as usize; + + tracing::info!( + "Dataset: {} vectors, Queries: {} vectors, dim: {}, recall@{}", + n_base, + n_queries, + dataset.ncols(), + recall_at, + ); + + let output_dir = Path::new(&args.output_dir); + std::fs::create_dir_all(output_dir) + .with_context(|| format!("creating output directory {}", output_dir.display())) + .map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + let gt_suffix = format!("gt{}", recall_at); + + // FindGroundtruth impl that always returns the expected output path whether + // or not it exists yet -- we are about to generate the files. + struct AllowMissing { + dir: PathBuf, + suffix: String, + } + impl FindGroundtruth for AllowMissing { + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { + Ok(self.dir.join(format!("step{}.{}", stage, self.suffix))) + } + } + + let mut finder = AllowMissing { + dir: output_dir.to_path_buf(), + suffix: gt_suffix, + }; + + tracing::info!( + "Parsing runbook {} for dataset \"{}\"", + args.runbook_file, + args.dataset_name + ); + let mut runbook = RunBook::load( + Path::new(&args.runbook_file), + &args.dataset_name, + &mut finder, + ) + .map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + tracing::info!("Runbook has {} stages", runbook.len()); + let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); + + let mut stream = GroundtruthStream { + storage: &storage, + dataset: &dataset, + queries: &queries, + distance_fn, + n_base, + n_queries, + recall_at, + active: vec![false; n_base], + ext_id: vec![0u32; n_base], + ext_to_offset: HashMap::new(), + }; + + runbook + .run_with(&mut stream, |_| Ok(())) + .map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + tracing::info!("Done."); + Ok(()) +} + +struct GroundtruthStream<'a, V: VectorRepr + Send + Sync> { + storage: &'a FileStorageProvider, + dataset: &'a Matrix, + queries: &'a Matrix, + distance_fn: V::Distance, + n_base: usize, + n_queries: usize, + recall_at: usize, + active: Vec, + ext_id: Vec, + ext_to_offset: HashMap, +} + +impl<'a, V> streaming::Stream + for GroundtruthStream<'a, V> +where + V: VectorRepr + Send + Sync, +{ + type Output = (); + + fn search( + &mut self, + args: diskann_benchmark_core::streaming::executors::bigann::Search<'_>, + ) -> anyhow::Result { + let timer = Instant::now(); + let n_active = self.active.iter().filter(|&&b| b).count(); + + tracing::info!( + "Computing top-{} groundtruth for {} active vectors against {} queries", + self.recall_at, + n_active, + self.n_queries, + ); + + let active_ids: Vec = self + .active + .iter() + .enumerate() + .filter_map(|(i, &on)| if on { Some(i) } else { None }) + .collect(); + + // it's generally find to use the global threadpool in diskann-tools + #[allow(clippy::disallowed_methods)] + let results: Vec> = (0..self.n_queries) + .into_par_iter() + .map(|qi| { + let query = self.queries.row(qi); + let mut pq = NeighborPriorityQueue::new(self.recall_at); + for &offset in &active_ids { + let dist = self + .distance_fn + .evaluate_similarity(self.dataset.row(offset), query); + pq.insert(Neighbor { + id: self.ext_id[offset], + distance: dist, + }); + } + pq + }) + .collect(); + + let under_k: Vec = results + .iter() + .enumerate() + .filter_map(|(qi, pq)| (pq.size() < self.recall_at).then_some(qi)) + .collect(); + if !under_k.is_empty() { + let preview: Vec = under_k.iter().copied().take(20).collect(); + let suffix = if under_k.len() > preview.len() { + format!(" (showing first {} query indices)", preview.len()) + } else { + String::new() + }; + return Err(anyhow::anyhow!( + "{}: {} / {} queries have fewer than {} results (active set = {}). Query indices: {:?}{}", + args.groundtruth.display(), + under_k.len(), + self.n_queries, + self.recall_at, + n_active, + preview, + suffix, + )); + } + + write_ground_truth::<()>( + self.storage, + args.groundtruth.to_str().ok_or_else(|| { + anyhow::anyhow!("Non-UTF8 groundtruth path: {}", args.groundtruth.display()) + })?, + self.n_queries, + self.recall_at, + results, + None, + ) + .map_err(|e| anyhow::anyhow!(e.to_string()))?; + + tracing::info!( + "Groundtruth written to {} in {:?}", + args.groundtruth.display(), + timer.elapsed(), + ); + + Ok(()) + } + + fn insert( + &mut self, + args: diskann_benchmark_core::streaming::executors::bigann::Insert, + ) -> anyhow::Result { + for id in args.offsets.clone() { + if id < self.n_base { + self.active[id] = true; + self.ext_id[id] = id as u32; + self.ext_to_offset.insert(id as u32, id); + } + } + Ok(()) + } + + fn replace( + &mut self, + args: diskann_benchmark_core::streaming::executors::bigann::Replace, + ) -> anyhow::Result { + for eid in args.ids.clone() { + if let Some(&offset) = self.ext_to_offset.get(&(eid as u32)) { + self.active[offset] = false; + self.ext_to_offset.remove(&(eid as u32)); + } + } + + for (offset, eid) in args.offsets.clone().zip(args.ids.clone()) { + if offset < self.n_base { + self.active[offset] = true; + self.ext_id[offset] = eid as u32; + self.ext_to_offset.insert(eid as u32, offset); + } + } + + Ok(()) + } + + fn delete( + &mut self, + args: diskann_benchmark_core::streaming::executors::bigann::Delete, + ) -> anyhow::Result { + for eid in args.ids.clone() { + if let Some(&offset) = self.ext_to_offset.get(&(eid as u32)) { + self.active[offset] = false; + self.ext_to_offset.remove(&(eid as u32)); + } + } + Ok(()) + } + + fn maintain(&mut self, _args: ()) -> anyhow::Result { + Ok(()) + } + + fn needs_maintenance(&mut self) -> bool { + false + } +} + +#[derive(Debug, Parser)] +struct Args { + /// Data type of the base and query vectors. + #[arg(long = "data-type", default_value = "float")] + pub data_type: DataType, + + /// Distance function to use. + #[arg(long = "dist-fn", default_value = "l2")] + pub distance_function: Metric, + + /// File containing the full base dataset in binary format. + #[arg(long = "base-file", short, required = true)] + pub base_file: String, + + /// File containing the query vectors in binary format. + #[arg(long = "query-file", short, required = true)] + pub query_file: String, + + /// Path to the BigANN runbook YAML file. + #[arg(long = "runbook-file", required = true)] + pub runbook_file: String, + + /// Dataset name within the runbook YAML file. + #[arg(long = "dataset-name", required = true)] + pub dataset_name: String, + + /// Number of nearest neighbours to compute per query (k). + /// + /// Output files are named step.gt. + #[arg(long = "recall-at", short = 'K', required = true)] + pub recall_at: u32, + + /// Directory to write the groundtruth files into. + /// + /// Files are written as `step.gt`, matching the + /// naming convention expected by `ScanDirectory`. + #[arg(long = "gt-dir", short, required = true)] + pub output_dir: String, +} diff --git a/diskann-tools/src/utils/ground_truth.rs b/diskann-tools/src/utils/ground_truth.rs index 77c2f663c..f42535ca8 100644 --- a/diskann-tools/src/utils/ground_truth.rs +++ b/diskann-tools/src/utils/ground_truth.rs @@ -520,7 +520,7 @@ fn write_range_search_ground_truth( +pub fn write_ground_truth( storage_provider: &impl StorageWriteProvider, ground_truth_file: &str, number_of_queries: usize,