From 59b2e0a09189b423e578d8dce64e37c7e4ba121e Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 17 Jul 2026 18:38:49 +0000 Subject: [PATCH 01/10] add utility for streaming groundtruth, reorganize runbook structs --- Cargo.lock | 5 +- diskann-benchmark-core/Cargo.toml | 4 +- diskann-benchmark-core/src/lib.rs | 42 - diskann-benchmark-core/src/streaming/api.rs | 334 +------ .../src/streaming/executors/bigann/mod.rs | 9 +- .../src/streaming/executors/bigann/parsing.rs | 805 ----------------- .../src/streaming/executors/bigann/runbook.rs | 830 ------------------ .../streaming/executors/bigann/validate.rs | 678 -------------- .../dataset.txt | 1 - .../expected.txt | 4 - .../runbook.yaml | 10 - .../bigann-ux/dataset-not-found/dataset.txt | 1 - .../bigann-ux/dataset-not-found/expected.txt | 1 - .../bigann-ux/dataset-not-found/runbook.yaml | 13 - .../dataset-value-not-mapping/dataset.txt | 1 - .../dataset-value-not-mapping/expected.txt | 5 - .../dataset-value-not-mapping/runbook.yaml | 10 - .../delete-invalid-range/dataset.txt | 1 - .../delete-invalid-range/expected.txt | 6 - .../delete-invalid-range/runbook.yaml | 10 - .../insert-invalid-range/dataset.txt | 1 - .../insert-invalid-range/expected.txt | 6 - .../insert-invalid-range/runbook.yaml | 6 - .../insert-start-equals-end/dataset.txt | 1 - .../insert-start-equals-end/expected.txt | 6 - .../insert-start-equals-end/runbook.yaml | 6 - .../bigann-ux/missing-max-pts/dataset.txt | 1 - .../bigann-ux/missing-max-pts/expected.txt | 4 - .../bigann-ux/missing-max-pts/runbook.yaml | 5 - .../bigann-ux/non-integer-max-pts/dataset.txt | 1 - .../non-integer-max-pts/expected.txt | 4 - .../non-integer-max-pts/runbook.yaml | 6 - .../replace-ids-start-equals-end/dataset.txt | 1 - .../replace-ids-start-equals-end/expected.txt | 6 - .../replace-ids-start-equals-end/runbook.yaml | 12 - .../replace-invalid-ids-range/dataset.txt | 1 - .../replace-invalid-ids-range/expected.txt | 6 - .../replace-invalid-ids-range/runbook.yaml | 12 - .../replace-invalid-tags-range/dataset.txt | 1 - .../replace-invalid-tags-range/expected.txt | 6 - .../replace-invalid-tags-range/runbook.yaml | 12 - .../replace-tags-start-equals-end/dataset.txt | 1 - .../expected.txt | 6 - .../runbook.yaml | 12 - .../runbook-does-not-exist/dataset.txt | 1 - .../runbook-does-not-exist/expected.txt | 4 - .../runbook-not-yaml-map/dataset.txt | 0 .../runbook-not-yaml-map/expected.txt | 1 - .../runbook-not-yaml-map/runbook.yaml | 2 - .../stage-key-not-integer/dataset.txt | 1 - .../stage-key-not-integer/expected.txt | 4 - .../stage-key-not-integer/runbook.yaml | 10 - .../stage-key-not-number/dataset.txt | 1 - .../stage-key-not-number/expected.txt | 4 - .../stage-key-not-number/runbook.yaml | 10 - .../unrecognized-dataset-key/dataset.txt | 1 - .../unrecognized-dataset-key/expected.txt | 4 - .../unrecognized-dataset-key/runbook.yaml | 7 - .../unrecognized-operation/dataset.txt | 1 - .../unrecognized-operation/expected.txt | 5 - .../unrecognized-operation/runbook.yaml | 6 - .../index/bftree/full_precision_streaming.rs | 10 +- .../src/index/bftree/spherical_streaming.rs | 7 +- diskann-tools/Cargo.toml | 2 +- .../src/bin/compute_streaming_groundtruth.rs | 320 +++++++ diskann-tools/src/utils/ground_truth.rs | 2 +- diskann-utils/Cargo.toml | 6 + diskann-utils/src/lib.rs | 34 + 68 files changed, 377 insertions(+), 2959 deletions(-) delete mode 100644 diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs delete mode 100644 diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs delete mode 100644 diskann-benchmark-core/src/streaming/executors/bigann/validate.rs delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml delete mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt delete mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml create mode 100644 diskann-tools/src/bin/compute_streaming_groundtruth.rs diff --git a/Cargo.lock b/Cargo.lock index b516798e8..b812bfe79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -691,7 +691,6 @@ dependencies = [ "diskann-vector", "futures-util", "itertools 0.13.0", - "serde_yaml", "tempfile", "thiserror 2.0.17", "tokio", @@ -943,8 +942,10 @@ dependencies = [ name = "diskann-utils" version = "0.55.0" dependencies = [ + "anyhow", "bytemuck", "cfg-if", + "diskann-benchmark-runner", "diskann-vector", "diskann-wide", "half", @@ -952,6 +953,8 @@ dependencies = [ "rand_distr", "rayon", "rstest", + "serde_yaml", + "tempfile", "thiserror 2.0.17", ] diff --git a/diskann-benchmark-core/Cargo.toml b/diskann-benchmark-core/Cargo.toml index 5b6740ef8..7cd914500 100644 --- a/diskann-benchmark-core/Cargo.toml +++ b/diskann-benchmark-core/Cargo.toml @@ -17,13 +17,11 @@ futures-util = { workspace = true, default-features = false } thiserror.workspace = true tokio.workspace = true -serde_yaml = { version = "0.9.34", optional = true } - [features] default = [] # BigANN Runbook support requires parsing a YAML file. -bigann = ["dep:serde_yaml"] +bigann = ["diskann-utils/bigann"] [lints] workspace = true diff --git a/diskann-benchmark-core/src/lib.rs b/diskann-benchmark-core/src/lib.rs index 337616878..9c7f82d1a 100644 --- a/diskann-benchmark-core/src/lib.rs +++ b/diskann-benchmark-core/src/lib.rs @@ -50,45 +50,3 @@ pub mod tokio; pub mod build; pub mod search; pub mod streaming; - -/// # Notes on Testing -/// -/// Some components in this framework (e.g. runbook parsing), have UX tests to report errors -/// encountered during parsing. The necessary input files and expected outputs are checked -/// in to the repo in the `tests` directory. -/// -/// If error message change, the expected results can generally be regenerated by running -/// the test suite with the environment variable -/// ```text -/// DISKANN_TEST=overwrite -/// ``` -/// set. -#[cfg(all(test, feature = "bigann"))] -mod ux { - const ENV_VAR: &str = "DISKANN_TEST"; - - /// Check if we should overwrite expected files. - pub(crate) fn should_overwrite() -> bool { - match std::env::var(ENV_VAR) { - Ok(v) if v == "overwrite" => true, - Ok(v) => panic!( - "Unknown value for {}: \"{}\". Expected \"overwrite\"", - ENV_VAR, v - ), - Err(std::env::VarError::NotPresent) => false, - Err(std::env::VarError::NotUnicode(_)) => { - panic!("Value for {} is not unicode", ENV_VAR) - } - } - } - - pub(crate) fn help() -> String { - format!("{}=overwrite", ENV_VAR) - } - - pub(crate) fn test_dir() -> std::path::PathBuf { - let manifest: &std::path::Path = env!("CARGO_MANIFEST_DIR").as_ref(); - manifest.join("tests") - // let test_data_path: PathBuf = format!("{}/tests/{}", manifest_dir, TEST_DATA_DIR).into(); - } -} diff --git a/diskann-benchmark-core/src/streaming/api.rs b/diskann-benchmark-core/src/streaming/api.rs index 28f24fe2e..c373709a9 100644 --- a/diskann-benchmark-core/src/streaming/api.rs +++ b/diskann-benchmark-core/src/streaming/api.rs @@ -3,336 +3,4 @@ * Licensed under the MIT license. */ -use std::any::Any; - -/// A streaming interface for performing dynamic (streaming) operations on an index. -/// -/// Streams are characterized by five operations: -/// -/// * `search`: This is, after all, the whole reason for building an index. -/// * `insert`: Insert new points into the index that do not already exist. -/// * `replace`: Replace existing points in the index with new data. -/// * `delete`: Remove points from the index. -/// * `maintain`: Perform maintenance operations on the index. Examples may -/// include fully removing deleted points from internal references. -/// -/// This trait is parameterized by an [`Arguments`] proxy trait, which defines the -/// argument types for each of the operations. The motivation here is to allow nesting -/// of [`Stream`] implementations that progressively modify or adapt the arguments -/// for better code reuse. An example of this is -/// [`crate::streaming::executors::bigann::WithData`], which is a stream layer adapting -/// the raw ranges used by [`crate::streaming::executors::bigann::RunBook`] into -/// actual data slices. -/// -/// Runners for [`Stream`]s use the [`Executor`] trait to invoke the stream operations -/// in a structured way. -pub trait Stream -where - A: Arguments, -{ - /// Output type for all operations. The `'static` is to allow results to be - /// aggregated in [`Any`] for type erasure in higher level [`Executor`]s. - type Output: 'static; - - /// Perform a search operation. - fn search(&mut self, args: A::Search<'_>) -> anyhow::Result; - - /// Perform an insert operation. - fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result; - - /// Perform a replace operation. - fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result; - - /// Perform a delete operation. - fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result; - - /// Perform a maintain operation. - fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result; - - /// Indicate whether or not maintenance is needed. [`Executor`] implementations - /// are responsible periodically checking this. - fn needs_maintenance(&mut self) -> bool; -} - -/// Operation arguments to [`Stream`]. -pub trait Arguments: 'static { - /// Argument to [`Stream::search`]. - type Search<'a>; - /// Argument to [`Stream::insert`]. - type Insert<'a>; - /// Argument to [`Stream::replace`]. - type Replace<'a>; - /// Argument to [`Stream::delete`]. - type Delete<'a>; - /// Argument to [`Stream::maintain`]. - type Maintain<'a>; -} - -/// A sequential executor for [`Stream`]s. -/// -/// Implementations invoke the operations of a [`Stream`] in a structured way (which should -/// be reflected in the associated documentation) and aggregate the results. -pub trait Executor { - /// The argument collection type for the underlying [`Stream`]. - type Args: Arguments; - - /// Execute a series of operations on `stream`. As outputs are produced, they will be - /// passed to `collect` for aggregation. - /// - /// Since dynamic execution may be long-running, this allows implementations of `collect` - /// to perform operations like status updates or partial saving of results as they are - /// generated. - /// - /// See also: [`Executor::run`]. - fn run_with(&mut self, stream: &mut S, collect: F) -> anyhow::Result<()> - where - S: Stream, - O: 'static, - F: FnMut(O) -> anyhow::Result<()>; - - /// Execute a series of operations on `stream`. The outputs of each operation will be - /// collected in the retuned `Vec` in-order. - /// - /// See also: [`Executor::run_with`]. - fn run(&mut self, stream: &mut S) -> anyhow::Result> - where - S: Stream, - { - let mut outputs = Vec::new(); - self.run_with(stream, |output| { - outputs.push(output); - Ok(()) - })?; - Ok(outputs) - } -} - -/// A type-erased [`Stream`] implementation that wraps stream (outputs)[`Stream::Output`] in -/// [`Box`]. -/// -/// This is useful as the final layer in a stack of nested [`Stream`]s that allows the top -/// level [`Executor`] to operate without knowledge of the concrete output types. -/// -/// From a performance perspective, this is usually fine since -/// -/// 1. Individual [`Stream`] operations are typically expensive relative to the cost of boxing. -/// 2. Many concrete [`Stream`] implementations will use the same [`Executor`] implementation. -/// Thus, boxing can help reduce code bloat without significant performance impact. -#[derive(Debug)] -pub struct AnyStream<'a, T>(&'a mut T); - -impl<'a, T> AnyStream<'a, T> { - /// Wrap `stream` in an [`AnyStream`]. - pub fn new(stream: &'a mut T) -> Self { - Self(stream) - } -} - -fn boxed(x: T) -> Box -where - T: Any, -{ - Box::new(x) -} - -impl Stream for AnyStream<'_, T> -where - A: Arguments, - T: Stream, -{ - type Output = Box; - - fn search(&mut self, args: A::Search<'_>) -> anyhow::Result { - self.0.search(args).map(boxed) - } - - fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result { - self.0.insert(args).map(boxed) - } - - fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result { - self.0.replace(args).map(boxed) - } - - fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result { - self.0.delete(args).map(boxed) - } - - fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result { - self.0.maintain(args).map(boxed) - } - - fn needs_maintenance(&mut self) -> bool { - self.0.needs_maintenance() - } -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - // The main thing we're testing here is the implementation of [`Executor::run`] and - // to a lesser extent `AnyStream`. - - struct Search; - struct Insert; - struct Replace; - struct Delete; - struct Maintain; - - #[derive(Debug, PartialEq)] - enum Op { - Search, - Insert, - Replace, - Delete, - Maintain, - } - - struct TestArgs; - - impl Arguments for TestArgs { - type Search<'a> = Search; - type Insert<'a> = Insert; - type Replace<'a> = Replace; - type Delete<'a> = Delete; - type Maintain<'a> = Maintain; - } - - struct TestStream { - needs_maintenance: bool, - } - - impl TestStream { - fn new(needs_maintenance: bool) -> Self { - Self { needs_maintenance } - } - } - - impl Stream for TestStream { - type Output = Op; - - fn search(&mut self, _args: Search) -> anyhow::Result { - Ok(Op::Search) - } - - fn insert(&mut self, _args: Insert) -> anyhow::Result { - Ok(Op::Insert) - } - - fn replace(&mut self, _args: Replace) -> anyhow::Result { - Ok(Op::Replace) - } - - fn delete(&mut self, _args: Delete) -> anyhow::Result { - Ok(Op::Delete) - } - - fn maintain(&mut self, _args: Maintain) -> anyhow::Result { - Ok(Op::Maintain) - } - - fn needs_maintenance(&mut self) -> bool { - self.needs_maintenance - } - } - - struct TestExecutor; - - impl Executor for TestExecutor { - type Args = TestArgs; - - fn run_with(&mut self, stream: &mut S, mut collect: F) -> anyhow::Result<()> - where - S: Stream, - O: 'static, - F: FnMut(O) -> anyhow::Result<()>, - { - collect(stream.search(Search)?)?; - collect(stream.insert(Insert)?)?; - collect(stream.replace(Replace)?)?; - collect(stream.delete(Delete)?)?; - collect(stream.maintain(Maintain)?)?; - Ok(()) - } - } - - #[test] - fn test_executor_run() -> anyhow::Result<()> { - let mut stream = TestStream::new(false); - let mut executor = TestExecutor; - - let outputs = executor.run(&mut stream)?; - - assert_eq!(outputs.len(), 5); - assert!(matches!(outputs[0], Op::Search)); - assert!(matches!(outputs[1], Op::Insert)); - assert!(matches!(outputs[2], Op::Replace)); - assert!(matches!(outputs[3], Op::Delete)); - assert!(matches!(outputs[4], Op::Maintain)); - - Ok(()) - } - - #[test] - fn test_any_stream() { - let mut stream = TestStream::new(false); - let mut any_stream = AnyStream::new(&mut stream); - - assert!( - !any_stream.needs_maintenance(), - "AnyStream should forward `needs_maintenance`" - ); - assert_eq!( - any_stream - .search(Search) - .unwrap() - .downcast_ref::() - .unwrap(), - &Op::Search - ); - assert_eq!( - any_stream - .insert(Insert) - .unwrap() - .downcast_ref::() - .unwrap(), - &Op::Insert - ); - assert_eq!( - any_stream - .replace(Replace) - .unwrap() - .downcast_ref::() - .unwrap(), - &Op::Replace - ); - assert_eq!( - any_stream - .delete(Delete) - .unwrap() - .downcast_ref::() - .unwrap(), - &Op::Delete - ); - assert_eq!( - any_stream - .maintain(Maintain) - .unwrap() - .downcast_ref::() - .unwrap(), - &Op::Maintain - ); - - let mut stream = TestStream::new(true); - let mut any_stream = AnyStream::new(&mut stream); - assert!( - any_stream.needs_maintenance(), - "AnyStream should forward `needs_maintenance`" - ); - } -} +pub use diskann_utils::streaming::{AnyStream, Arguments, Executor, Stream}; diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs b/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs index 5a7f5c72a..5ccc9bf41 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs @@ -5,18 +5,15 @@ //! # BigANN Runbook Support //! -//! This module provides an executor for processing BigANN-style runbooks. -//! Please refer to the [`RunBook`] documentation for more details. +//! Runbook parsing and execution are provided by `diskann-utils`. +//! This module provides benchmark-specific adaptors over those shared types. //! //! Users can also leverage the [`WithData`] adaptor to convert raw index ranges //! into actual data points for a dataset. -mod parsing; -mod runbook; -mod validate; mod withdata; -pub use runbook::{ +pub use diskann_utils::streaming::executors::bigann::{ Args, Delete, FindGroundtruth, Insert, Replace, RunBook, ScanDirectory, Search, Stage, }; pub use withdata::{DataArgs, WithData}; diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs b/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs deleted file mode 100644 index 3070bd47e..000000000 --- a/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs +++ /dev/null @@ -1,805 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::path::{Path, PathBuf}; - -use anyhow::Context; -use serde_yaml::{Mapping, Value}; - -/// See [`super::RunBook::load`] for documentation of the file format parsed -/// by this function. -pub(super) fn load( - path: &Path, - dataset: &str, - groundtruth: &mut dyn super::FindGroundtruth, -) -> anyhow::Result { - load_mapping(parse_yaml(path)?, path, dataset, groundtruth) -} - -fn load_mapping( - mapping: Mapping, - path: &Path, - dataset: &str, - groundtruth: &mut dyn super::FindGroundtruth, -) -> anyhow::Result { - // Multiple datasets can exist in the same YAML file. - // - // Fortunately, all datasets are keyed by their dataset name which allows us to - // quickly find whether or not it exists. - let dataset_value = match mapping.get(dataset) { - Some(value) => value, - None => return Err(DumpKeys::new(mapping, dataset, path).into()), - }; - - // Try to coerce the dataset value into a `Mapping`. - let dataset_mapping: &Mapping = match dataset_value.try_as() { - Ok(mapping) => mapping, - Err(_) => anyhow::bail!( - "dataset \"{}\" exists in file \"{}\", but its associated payload is not a YAML map", - dataset, - path.display(), - ), - }; - - let mut raw = parse_stages(dataset_mapping).with_context(|| { - format!( - "parsing dataset \"{}\" in file \"{}\"", - dataset, - path.display() - ) - })?; - raw.stages.sort_by_key(|s| s.index); - - // Translate from raw ranges into higher level steps. - let context = |index: usize| { - format!( - "precessing stage {} of dataset \"{}\" in file \"{}\"", - index, - dataset, - path.display() - ) - }; - - let stages: anyhow::Result> = raw - .stages - .iter() - .map(|stage| { - let stage = match &stage.operation { - Operation::Search => super::Stage::Search { - groundtruth: groundtruth - .find_groundtruth(stage.index) - .with_context(|| context(stage.index))?, - }, - Operation::Insert(insert) => super::Stage::Insert { - dataset_offsets_and_ids: insert.start..insert.end, - }, - Operation::Replace(replace) => super::Stage::Replace { - dataset_offsets: replace.ids_start..replace.ids_end, - ids: replace.tags_start..replace.tags_end, - }, - Operation::Delete(delete) => super::Stage::Delete { - ids: delete.start..delete.end, - }, - }; - Ok(stage) - }) - .collect(); - - super::RunBook::new(stages?, raw.max_points) -} - -fn parse_yaml(path: &Path) -> anyhow::Result { - let f = std::fs::File::open(path) - .with_context(|| format!("while opening file \"{}\"", path.display()))?; - - Ok(serde_yaml::from_reader(std::io::BufReader::new(f))?) -} - -fn parse_stages(mapping: &Mapping) -> anyhow::Result { - let mut stages = Vec::::new(); - let mut max_points = None; - mapping.iter().try_for_each(|(key, value)| match key { - Value::String(s) => match s.as_str() { - "max_pts" => { - let points: usize = value - .try_as() - .map_err(|kind| anyhow::anyhow!("failed to parse \"max_pts\" as a {}", kind))?; - - max_points = Some(points); - Ok(()) - } - "gt_url" => Ok(()), - _ => anyhow::bail!("Unrecognized runbook key: \"{}\"", s), - }, - Value::Number(stage) => match stage.as_i64() { - Some(stage) => { - stages.push( - handle_stage(stage as usize, value) - .with_context(|| format!("processing stage {}", stage))?, - ); - Ok(()) - } - None => anyhow::bail!("Stage \"{}\" must be an integer", stage), - }, - _ => anyhow::bail!("Unrecognized key of type {}", classify(key),), - })?; - - let max_points = match max_points { - Some(points) => points, - None => anyhow::bail!("key \"max_pts\" not found"), - }; - - Ok(Raw { max_points, stages }) -} - -fn handle_stage(index: usize, value: &Value) -> anyhow::Result { - let mapping: &Mapping = value - .try_as() - .map_err(|_| anyhow::anyhow!("YAML type is not a map"))?; - - let kind: &str = mapping.get_as("operation")?; - let operation = match Kind::try_parse(kind)? { - Kind::Search => Operation::Search, - Kind::Insert => Operation::Insert(mapping.try_into()?), - Kind::Replace => Operation::Replace(mapping.try_into()?), - Kind::Delete => Operation::Delete(mapping.try_into()?), - }; - - Ok(Stage { index, operation }) -} - -struct Raw { - max_points: usize, - stages: Vec, -} - -struct Stage { - index: usize, - operation: Operation, -} - -enum Operation { - Search, - Insert(Insert), - Replace(Replace), - Delete(Delete), -} - -#[derive(Debug)] -struct DumpKeys { - mapping: Mapping, - dataset: String, - file: PathBuf, -} - -impl DumpKeys { - #[inline(never)] - fn new(mapping: Mapping, dataset: &str, file: &Path) -> Self { - Self { - mapping, - dataset: dataset.to_string(), - file: file.into(), - } - } -} - -impl std::fmt::Display for DumpKeys { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "dataset \"{}\" not found in file \"{}\" - possible alternatives: [", - self.dataset, - self.file.display(), - )?; - - let len = self.mapping.len(); - self.mapping.keys().enumerate().try_for_each(|(i, key)| { - let mut write = |key: &dyn std::fmt::Display| { - if i + 1 == len { - write!(f, "{}", key) - } else { - write!(f, "{}, ", key) - } - }; - - match key { - Value::Null => write(&"null"), - Value::Bool(b) => write(b), - Value::Number(number) => write(number), - Value::String(s) => write(s), - Value::Sequence(_) => write(&""), - Value::Mapping(_) => write(&""), - Value::Tagged(_) => write(&""), - } - })?; - write!(f, "]") - } -} - -impl std::error::Error for DumpKeys {} - -trait TryAs<'a, T> { - fn try_as(&'a self) -> Result; -} - -impl<'a> TryAs<'a, usize> for Value { - fn try_as(&'a self) -> Result { - self.as_i64().map(|i| i as usize).ok_or("usize") - } -} - -impl<'a> TryAs<'a, &'a str> for Value { - fn try_as(&'a self) -> Result<&'a str, &'static str> { - self.as_str().ok_or("string") - } -} - -impl<'a> TryAs<'a, &'a Mapping> for Value { - fn try_as(&'a self) -> Result<&'a Mapping, &'static str> { - self.as_mapping().ok_or("map") - } -} - -trait MappingExt { - fn get_as<'a, T>(&'a self, index: &str) -> anyhow::Result - where - Value: TryAs<'a, T>; -} - -impl MappingExt for Mapping { - fn get_as<'a, T>(&'a self, key: &str) -> anyhow::Result - where - Value: TryAs<'a, T>, - { - match self.get(key) { - Some(value) => match value.try_as() { - Ok(v) => Ok(v), - Err(expected) => Err(anyhow::anyhow!( - "key \"{}\" exists but it is not a {}", - key, - expected, - )), - }, - None => Err(anyhow::anyhow!("key \"{}\" not found", key)), - } - } -} - -#[derive(Debug, Clone, Copy)] -enum Kind { - Search, - Insert, - Replace, - Delete, -} - -impl Kind { - fn try_parse(string: &str) -> anyhow::Result { - match string { - "search" => Ok(Kind::Search), - "insert" => Ok(Kind::Insert), - "replace" => Ok(Kind::Replace), - "delete" => Ok(Kind::Delete), - _ => Err(anyhow::anyhow!("unrecognized operation: {}", string)), - } - } -} - -#[derive(Debug)] -struct Replace { - ids_start: usize, - ids_end: usize, - tags_start: usize, - tags_end: usize, -} - -impl TryFrom<&Mapping> for Replace { - type Error = anyhow::Error; - fn try_from(mapping: &Mapping) -> anyhow::Result { - let inner = || -> anyhow::Result { - let this = Self { - ids_start: mapping.get_as("ids_start")?, - ids_end: mapping.get_as("ids_end")?, - tags_start: mapping.get_as("tags_start")?, - tags_end: mapping.get_as("tags_end")?, - }; - if this.ids_start >= this.ids_end { - anyhow::bail!( - "ids_start ({}) must be less than ids_end ({})", - this.ids_start, - this.ids_end - ); - } - if this.tags_start >= this.tags_end { - anyhow::bail!( - "tags_start ({}) must be less than tags_end ({})", - this.tags_start, - this.tags_end - ); - } - Ok(this) - }; - - inner().context("trying to parse an \"replace\"") - } -} - -#[derive(Debug)] -struct Insert { - start: usize, - end: usize, -} - -impl TryFrom<&Mapping> for Insert { - type Error = anyhow::Error; - fn try_from(mapping: &Mapping) -> anyhow::Result { - let inner = || -> anyhow::Result { - let this = Self { - start: mapping.get_as("start")?, - end: mapping.get_as("end")?, - }; - if this.start >= this.end { - anyhow::bail!( - "start ({}) must be less than end ({})", - this.start, - this.end - ); - } - Ok(this) - }; - - inner().context("trying to parse an \"insert\"") - } -} - -#[derive(Debug)] -struct Delete { - start: usize, - end: usize, -} - -impl TryFrom<&Mapping> for Delete { - type Error = anyhow::Error; - fn try_from(mapping: &Mapping) -> anyhow::Result { - let inner = || -> anyhow::Result { - let this = Self { - start: mapping.get_as("start")?, - end: mapping.get_as("end")?, - }; - if this.start >= this.end { - anyhow::bail!( - "start ({}) must be less than end ({})", - this.start, - this.end - ); - } - Ok(this) - }; - - inner().context("trying to parse \"delete\"") - } -} - -fn classify(value: &Value) -> &'static str { - match value { - Value::Null => "null", - Value::Bool(_) => "bool", - Value::Number(_) => "number", - Value::String(_) => "string", - Value::Sequence(_) => "sequence", - Value::Mapping(_) => "mapping", - Value::Tagged(_) => "tagged", - } -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - use std::{collections::HashMap, io::Write}; - - use tempfile::NamedTempFile; - - use crate::streaming::executors::bigann::Stage; - - /// A test implementation of [`super::super::FindGroundtruth`] that returns - /// pre-configured paths for each stage index. - struct MockGroundtruth { - paths: HashMap, - } - - impl MockGroundtruth { - fn new(stages: impl IntoIterator) -> Self { - Self { - paths: stages.into_iter().collect(), - } - } - } - - impl super::super::FindGroundtruth for MockGroundtruth { - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { - self.paths - .get(&stage) - .cloned() - .ok_or_else(|| anyhow::anyhow!("no groundtruth configured for stage {}", stage)) - } - } - - /// Helper to create a temporary YAML file with the given content. - fn create_yaml_file(content: &str) -> anyhow::Result { - let mut file = NamedTempFile::new()?; - file.write_all(content.as_bytes())?; - file.flush()?; - Ok(file) - } - - #[test] - fn test_load_simple_insert_only_runbook() { - let yaml = r#" -test_dataset: - max_pts: 1000 - 0: - operation: insert - start: 0 - end: 500 -"#; - - let file = create_yaml_file(yaml).unwrap(); - let mut groundtruth = MockGroundtruth::new([]); - - let runbook = load(file.path(), "test_dataset", &mut groundtruth).unwrap(); - - assert_eq!(runbook.max_points(), 1000); - assert_eq!(runbook.len(), 1); - - assert_eq!( - runbook.stages()[0], - Stage::Insert { - dataset_offsets_and_ids: 0..500 - } - ); - } - - #[test] - fn test_load_runbook_with_search_stage() { - let yaml = r#" -my_dataset: - max_pts: 2000 - 0: - operation: insert - start: 0 - end: 1000 - 1: - operation: search -"#; - - let file = create_yaml_file(yaml).unwrap(); - let mut groundtruth = - MockGroundtruth::new([(1, PathBuf::from("/path/to/groundtruth.bin"))]); - - let runbook = load(file.path(), "my_dataset", &mut groundtruth).unwrap(); - - assert_eq!(runbook.max_points(), 2000); - assert_eq!(runbook.len(), 2); - - assert_eq!( - runbook.stages()[1], - Stage::Search { - groundtruth: PathBuf::from("/path/to/groundtruth.bin") - } - ); - } - - #[test] - fn test_load_runbook_with_all_operation_types() { - let yaml = r#" -full_dataset: - max_pts: 5000 - 0: - operation: insert - start: 0 - end: 1000 - 1: - operation: search - 2: - operation: replace - ids_start: 1000 - ids_end: 1500 - tags_start: 0 - tags_end: 500 - 3: - operation: delete - start: 500 - end: 1000 -"#; - - let file = create_yaml_file(yaml).unwrap(); - let mut groundtruth = MockGroundtruth::new([(1, PathBuf::from("/gt/step1.bin"))]); - - let runbook = load(file.path(), "full_dataset", &mut groundtruth).unwrap(); - - assert_eq!(runbook.max_points(), 5000); - assert_eq!(runbook.len(), 4); - - // Check insert - assert_eq!( - runbook.stages()[0], - Stage::Insert { - dataset_offsets_and_ids: 0..1000 - } - ); - - // Check search - assert_eq!( - runbook.stages()[1], - Stage::Search { - groundtruth: PathBuf::from("/gt/step1.bin") - } - ); - - // Check replace - assert_eq!( - runbook.stages()[2], - Stage::Replace { - dataset_offsets: 1000..1500, - ids: 0..500 - } - ); - - // Check delete - assert_eq!(runbook.stages()[3], Stage::Delete { ids: 500..1000 }); - } - - #[test] - fn test_load_stages_out_of_order_are_sorted() { - let yaml = r#" -unordered: - max_pts: 1000 - 2: - operation: delete - start: 500 - end: 600 - 0: - operation: insert - start: 0 - end: 500 - 1: - operation: insert - start: 500 - end: 1000 -"#; - - let file = create_yaml_file(yaml).unwrap(); - let mut groundtruth = MockGroundtruth::new([]); - - let runbook = load(file.path(), "unordered", &mut groundtruth).unwrap(); - - assert_eq!(runbook.len(), 3); - - // Stages should be in order 0, 1, 2 regardless of YAML order - assert_eq!( - runbook.stages()[0], - Stage::Insert { - dataset_offsets_and_ids: 0..500 - } - ); - - assert_eq!( - runbook.stages()[1], - Stage::Insert { - dataset_offsets_and_ids: 500..1000 - } - ); - - assert_eq!(runbook.stages()[2], Stage::Delete { ids: 500..600 }); - } - - #[test] - fn test_load_multiple_datasets_in_file() { - let yaml = r#" -dataset_a: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - -dataset_b: - max_pts: 200 - 0: - operation: insert - start: 0 - end: 200 -"#; - - let file = create_yaml_file(yaml).unwrap(); - - // Load dataset_a - let mut groundtruth_a = MockGroundtruth::new([]); - let runbook_a = load(file.path(), "dataset_a", &mut groundtruth_a).unwrap(); - assert_eq!(runbook_a.max_points(), 100); - - // Load dataset_b - let mut groundtruth_b = MockGroundtruth::new([]); - let runbook_b = load(file.path(), "dataset_b", &mut groundtruth_b).unwrap(); - assert_eq!(runbook_b.max_points(), 200); - } - - #[test] - fn test_load_gt_url_is_ignored() { - let yaml = r#" -with_gt_url: - max_pts: 100 - gt_url: "https://example.com/groundtruth.bin" - 0: - operation: insert - start: 0 - end: 100 -"#; - - let file = create_yaml_file(yaml).unwrap(); - let mut groundtruth = MockGroundtruth::new([]); - - // Should succeed - gt_url is parsed but ignored - let runbook = load(file.path(), "with_gt_url", &mut groundtruth).unwrap(); - assert_eq!(runbook.max_points(), 100); - } -} - -#[cfg(test)] -mod ux_tests { - use super::*; - - // Exposed by the `ux-tools` feature of `diskann_benchmark_runner` - use diskann_benchmark_runner::ux as runner_ux; - - //---------------------------// - // File-Based UX Error Tests // - //---------------------------// - // - // These tests use checked-in YAML files and expected output files to verify error messages. - // This approach makes it easy to: - // 1. Add new test cases (just add a new directory with runbook.yaml, dataset.txt, expected.txt) - // 2. See the full error output for review - // 3. Regenerate expected output when error messages change - // - // ## Directory Structure - // - // Each test case is a directory under `tests/bigann-ux` containing: - // - `runbook.yaml` - The YAML runbook file to parse - // - `dataset.txt` - Contains the dataset name to load (single line) - // - `expected.txt` - The expected error message output - // - // ## Regenerating Expected Results - // - // Run tests with the environment variable: - // ``` - // DISKANN_TEST=overwrite - // ``` - // to regenerate the `expected.txt` files. Use `git diff` to review changes. - - const TEST_DATA_DIR: &str = "bigann-ux"; - const RUNBOOK_FILE: &str = "runbook.yaml"; - const DATASET_FILE: &str = "dataset.txt"; - const EXPECTED_FILE: &str = "expected.txt"; - const PATH_PLACEHOLDER: &str = ""; - - /// Replace the test directory path with a placeholder to make tests portable. - /// This handles both forward and backslash path separators. - /// - /// Additionally: - /// * All backslashes are replaced with forward slashes. - /// * Common OS-specific "file not found" error messages are normalized. - fn fixup_paths_and_os_errors(s: &str, test_dir: &Path) -> String { - // Try both the native path and with normalized separators - let native_path = test_dir.display().to_string(); - let forward_slash_path = native_path.replace('\\', "/"); - - const NOT_FOUND_WINDOWS: &str = "The system cannot find the file specified."; - const NOT_FOUND_LINUX: &str = "No such file or directory"; - - s.replace(&native_path, PATH_PLACEHOLDER) - .replace(&forward_slash_path, PATH_PLACEHOLDER) - .replace('\\', "/") // Normalize any remaining backslashes - .replace(NOT_FOUND_WINDOWS, NOT_FOUND_LINUX) // Normalize error messages - } - - /// A groundtruth finder that always fails - used for error path testing. - struct FailingGroundtruth; - - impl super::super::FindGroundtruth for FailingGroundtruth { - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { - Err(anyhow::anyhow!( - "groundtruth not available for stage {}", - stage - )) - } - } - - /// Run a single file-based test case. - fn run_file_test(test_dir: &Path) { - let runbook_path = test_dir.join(RUNBOOK_FILE); - let dataset_path = test_dir.join(DATASET_FILE); - let expected_path = test_dir.join(EXPECTED_FILE); - - // Read the dataset name - let dataset = std::fs::read_to_string(&dataset_path) - .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", dataset_path, e)); - let dataset = dataset.trim(); - - // Try to load the runbook - we expect an error - let mut groundtruth = FailingGroundtruth; - let result = load(&runbook_path, dataset, &mut groundtruth); - - let actual_output = match result { - Ok(_) => panic!( - "Test {:?} expected an error but parsing succeeded", - test_dir.file_name().unwrap() - ), - Err(err) => format!("{:?}", err), - }; - - // Replace test directory path with placeholder for portability - let actual_portable = fixup_paths_and_os_errors(&actual_output, test_dir); - let actual_normalized = runner_ux::strip_backtrace(runner_ux::normalize(actual_portable)); - - if crate::ux::should_overwrite() { - std::fs::write(&expected_path, &actual_normalized) - .unwrap_or_else(|e| panic!("Failed to write {:?}: {}", expected_path, e)); - println!("Overwrote {:?}", expected_path); - } else { - let expected = std::fs::read_to_string(&expected_path) - .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", expected_path, e)); - let expected_normalized = runner_ux::normalize(expected); - - if actual_normalized != expected_normalized { - panic!( - "Test {:?} failed.\n\nExpected:\n---\n{}\n---\n\nActual:\n---\n{}\n---\nIf this is expected, run with {} to update the expected output.", - test_dir.file_name().unwrap(), - expected_normalized, - actual_normalized, - crate::ux::help(), - ); - } - } - } - - /// Run all file-based tests in the test_data directory. - fn run_all_file_tests() { - let test_data_path = crate::ux::test_dir().join(TEST_DATA_DIR); - if !test_data_path.exists() { - println!( - "No test_data directory found at {:?}, skipping file-based tests", - test_data_path - ); - return; - } - - let mut found_tests = false; - for entry in std::fs::read_dir(&test_data_path) - .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", test_data_path, e)) - { - let entry = entry.unwrap(); - if entry.file_type().unwrap().is_dir() { - found_tests = true; - println!("Running file-based test: {:?}", entry.file_name()); - run_file_test(&entry.path()); - } - } - - if !found_tests { - panic!("No test directories found in {:?}", test_data_path); - } - } - - #[test] - fn file_based_error_tests() { - run_all_file_tests(); - } -} diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs deleted file mode 100644 index fbaee73df..000000000 --- a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs +++ /dev/null @@ -1,830 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::{ - any::Any, - ops::Range, - path::{Path, PathBuf}, -}; - -use anyhow::Context; - -use crate::streaming::{self, Executor, Stream}; - -use super::{parsing, validate}; - -/// An executor for [BigANN-style runbooks](https://github.com/harsha-simhadri/big-ann-benchmarks/tree/main/neurips23/streaming). -/// -/// If using this struct as a [`streaming::Executor`], consider using the -/// [`super::WithData`] adaptor to provide dataset and query matrices. -#[derive(Debug, Clone)] -pub struct RunBook { - // The individual runbook stages. - stages: Vec, - // The maximum number of active points at any point in the runbook. - max_points: usize, - // The maximum tag referenced in the runbook. - max_tag: Option, -} - -impl RunBook { - /// Loads a [`RunBook`] from a YAML file at `path` for the specified `dataset`. - /// - /// # Groundtruth Resolution - /// - /// When loading a runbook, search stages may require groundtruth files to be present. - /// The index of these stages will be provided to the `groundtruth` argument for resolution. - /// If working with the standard BigANN benchmark format, consider using [`ScanDirectory`] - /// and providing the directory where the BigANN framework downloads the groundtruth files. - /// - /// Note that the implementation here currently does not attempt to download any groundtruth - /// files using the `gt_url` field; it is merely parsed and ignored. - /// - /// # YAML File Format - /// - /// The top-level structure is a mapping from dataset names (strings) to their - /// corresponding runbook definitions: - /// ```yaml - /// dataset_name_1: - /// # runbook definition... - /// dataset_name_2: - /// # runbook definition... - /// ``` - /// The `dataset` parameter specifies which dataset's runbook to load from the file. - /// - /// Each runbook definition has the following format: - /// - /// ```yaml - /// max_pts: # required - /// [gt_url]: # ignored - /// 0: - /// # stage definition ... - /// 1: - /// # stage definition ... - /// ... - /// ``` - /// Entries need not be in order, but stages must be sequentially numbered starting - /// from `0`. Each stage takes one of four forms (described below). - /// - /// ## Search Stage - /// - /// Merely specifies that a search should be performed. The queries must be provided - /// externally. The groundtruth file for this stage is located via the provided - /// [`FindGroundtruth::find_groundtruth`] method, which will be provided with the stage index. - /// - /// ```yaml - /// : - /// operation: "search" - /// ``` - /// - /// ## Insert Stage - /// - /// Insert vectors from the underlying dataset into the index. The vectors to insert - /// are specified by the range `start..end`, which serves as both the offsets of the - /// vectors in the dataset and their external ids. - /// - /// ```yaml - /// : - /// operation: "insert" - /// start: - /// end: - /// ``` - /// - /// ## Replace Stage - /// - /// Replace vectors in the index with vectors from the underlying dataset. Unlike insertions, - /// replace operations distinguish between the dataset offsets (`ids_start..ids_end`) - /// and the external ids (tags) of the vectors to replace (`tags_start..tags_end`). - /// - /// These operations indicate that the vectors in the index tagged by `tags_start..tags_end` should - /// be replaced with the vectors from the dataset at offsets `ids_start..ids_end`. - /// - /// ```yaml - /// : - /// operation: "replace" - /// ids_start: - /// ids_end: - /// tags_start: - /// tags_end: - /// ``` - /// - /// ## Delete Stage - /// - /// Delete vectors from the index with external ids in the range `start..end`. - /// - /// ```yaml - /// : - /// operation: "delete" - /// start: - /// end: - /// ``` - /// - /// # File Validation - /// - /// The loading framework does a limited amount of validation on the YAML file: - /// - /// 1. The specified `dataset` must exist in the top-level mapping. - /// 2. The `max_pts` key must be present and be a valid integer. Its value is verified and if found to - /// be inaccurate, it is updated to the correct value internally. - /// 3. Each stage must be sequentially numbered starting from `0` with no gaps. - /// 4. Each stage must have a valid operation with all required fields present and of the correct type. - /// 5. Groundtruth files must be resolvable via the provided [`FindGroundtruth`] implementation. - pub fn load( - path: &Path, - dataset: &str, - groundtruth: &mut dyn FindGroundtruth, - ) -> anyhow::Result { - parsing::load(path, dataset, groundtruth) - } - - pub(super) fn new(stages: Vec, max_points: usize) -> anyhow::Result { - let mut this = Self { - stages, - max_points, - max_tag: None, - }; - - let mut validator = validate::Validate::new(); - this.run_with(&mut validator, |_| Ok(()))?; - - this.max_points = this.max_points.max(validator.max_active()); - this.max_tag = validator.max_tag(); - - Ok(this) - } - - /// Returns the maximum number of points specified in the runbook. - pub fn max_points(&self) -> usize { - self.max_points - } - - /// Returns the maximum tag specified in the runbook. - pub fn max_tag(&self) -> Option { - self.max_tag - } - - /// Returns the number of stages in the runbook. - pub fn len(&self) -> usize { - self.stages.len() - } - - /// Returns `true` if the runbook contains no stages. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Returns a reference to the stages in this runbook. - #[cfg(test)] - pub(super) fn stages(&self) -> &[Stage] { - &self.stages - } - - /// Executes the runbook by iterating through each stage. - /// - /// When calling this method, the dynamic type of `stream`'s output must be - /// compatible with the dynamic type expected by `collect`. - fn run_with_internal( - &self, - stream: &mut dyn streaming::Stream>, - collect: &mut dyn FnMut(Box) -> anyhow::Result<()>, - ) -> anyhow::Result<()> { - for (i, stage) in self.stages.iter().enumerate() { - #[derive(Clone, Copy)] - struct OnStage(usize, usize); - - impl std::fmt::Display for OnStage { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "on stage {} of {}", self.0, self.1) - } - } - - let context = OnStage(i, self.len()); - - if stream.needs_maintenance() { - collect(stream.maintain(()).context(context)?).context(context)?; - } - - let output = match stage { - Stage::Search { groundtruth } => { - let args = Search { groundtruth }; - stream.search(args).context(context)? - } - Stage::Insert { - dataset_offsets_and_ids, - } => { - let args = Insert { - offsets: dataset_offsets_and_ids.clone(), - ids: dataset_offsets_and_ids.clone(), - }; - stream.insert(args).context(context)? - } - Stage::Replace { - dataset_offsets, - ids, - } => { - let args = Replace { - offsets: dataset_offsets.clone(), - ids: ids.clone(), - }; - stream.replace(args).context(context)? - } - Stage::Delete { ids } => { - let args = Delete { ids: ids.clone() }; - stream.delete(args).context(context)? - } - }; - - collect(output).context(context)?; - } - - Ok(()) - } -} - -/// An operation in a BigANN runbook. -#[derive(Debug, Clone, PartialEq)] -pub enum Stage { - /// Perform a search operation using the specified groundtruth file. - Search { - /// The resolved path to the groundtruth file for this search stage. - groundtruth: PathBuf, - }, - /// Insert vectors from the dataset into the index. - Insert { - /// The offsets in the dataset, which also serve as the external ids for - /// the inserted vectors. - dataset_offsets_and_ids: Range, - }, - /// Replace vectors in the index with new vectors from the dataset. - Replace { - /// The offsets in the dataset for the replacement vectors. - dataset_offsets: Range, - /// The external ids of the vectors to be replaced. - ids: Range, - }, - /// Delete vectors from the index. - Delete { - /// The external ids of the vectors to delete. - ids: Range, - }, -} - -/// Arguments for a BigANN runbook "search" stage. -#[derive(Debug)] -pub struct Search<'a> { - /// The resolved path to the file containing the groundtruth for this stage. - pub groundtruth: &'a Path, -} - -/// Arguments for a BigANN runbook "insert" stage. -pub struct Insert { - /// The range of offsets in the dataset for vectors to insert. - pub offsets: Range, - /// The external ids to assign to the inserted vectors. - pub ids: Range, -} - -/// Arguments for a BigANN runbook "replace" stage. -pub struct Replace { - /// The range of offsets in the dataset for the replacement vectors. - pub offsets: Range, - /// The external ids of the vectors to replace. - pub ids: Range, -} - -/// Arguments for a BigANN runbook "delete" stage. -pub struct Delete { - /// The range of external ids to delete from the index. - pub ids: Range, -} - -/// The argument type for a [`RunBook`]. -/// -/// See also: [`super::WithData`], [`super::DataArgs`]. -#[derive(Debug, Clone, Copy)] -pub struct Args; - -impl streaming::Arguments for Args { - type Search<'a> = Search<'a>; - type Insert<'a> = Insert; - type Replace<'a> = Replace; - type Delete<'a> = Delete; - type Maintain<'a> = (); -} - -impl streaming::Executor for RunBook { - type Args = Args; - - fn run_with(&mut self, stream: &mut S, mut collect: F) -> anyhow::Result<()> - where - S: Stream, - O: 'static, - F: FnMut(O) -> anyhow::Result<()>, - { - self.run_with_internal(&mut streaming::AnyStream::new(stream), &mut |any| { - let typed = *any - .downcast::() - .expect("the dynamic type should be configured correctly"); - collect(typed) - }) - } -} - -/// A trait for resolving groundtruth files for search stages in a [`RunBook`]. -/// -/// Implementors of this trait provide the logic to locate groundtruth files -/// given a stage index. See [`ScanDirectory`] for a common implementation. -pub trait FindGroundtruth { - /// Resolves the groundtruth file path for the specified `stage` index. - /// - /// This method is only called for "search" stages in a runbook. - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result; -} - -/// A [`FindGroundtruth`] implementation that scans a directory for groundtruth files -/// matching the expected BigANN naming convention: `step{stage}.gt[0-9]*` where -/// `{stage}` substitutes the stage index formatted as a string. -#[derive(Debug)] -pub struct ScanDirectory { - directory: PathBuf, - - // Cached contents of the files in `directory`. - files: Vec, -} - -impl ScanDirectory { - /// Creates a new [`ScanDirectory`] instance for the specified `directory`. - /// - /// # Notes - /// - /// This constructor scans the directory and caches its contents. - /// If the directory contents change after creation, the instance will not - /// reflect those changes. - /// - /// This is meant for the common benchmarking scenario where the benchmarking - /// machine is generally static while benchmarks are executed. - pub fn new(directory: impl Into) -> anyhow::Result { - Self::new_(directory.into()) - } - - fn new_(directory: PathBuf) -> anyhow::Result { - // Read all files in the directory. - let read_dir = std::fs::read_dir(&directory).with_context(|| { - format!( - "while trying to read the contents of {}", - directory.display() - ) - })?; - - let files = read_dir - .filter_map(|entry| { - if let Ok(entry) = entry - && let Ok(file_type) = entry.file_type() - && file_type.is_file() - { - Some(entry.file_name().to_string_lossy().into()) - } else { - None - } - }) - .collect(); - - Ok(Self { directory, files }) - } -} - -impl FindGroundtruth for ScanDirectory { - /// Finds the groundtruth file for the specified `stage` by scanning the directory. - /// - /// # Errors - /// - /// Returns an error if no matching file is found or if multiple matches exist. - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { - let prefix = format!("step{}.gt", stage); - - enum Matches<'a> { - None, - One(&'a str), - Many(Vec<&'a str>), - } - - impl<'a> Matches<'a> { - fn push(&mut self, file: &'a str) { - *self = match std::mem::replace(self, Self::None) { - Self::None => Self::One(file), - Self::One(first) => Self::Many(vec![first, file]), - Self::Many(mut all) => { - all.push(file); - Self::Many(all) - } - }; - } - } - - let mut matches = Matches::None; - - for file in self.files.iter() { - if file.starts_with(&prefix) { - let suffix = &file[prefix.len()..]; - if suffix.chars().all(|c| c.is_ascii_digit()) { - matches.push(file); - } - } - } - - match matches { - Matches::One(m) => Ok(self.directory.join(m)), - Matches::None => Err(anyhow::anyhow!( - "No groundtruth found for step {} in \"{}\", expected pattern: \"step{}.gt[0-9]*\"", - stage, - self.directory.display(), - stage, - )), - Matches::Many(matches) => Err(anyhow::anyhow!( - "Multiple groundtruth files found for step {} in \"{}\": {:?}", - stage, - self.directory.display(), - matches, - )), - } - } -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - use std::fs::File; - - use tempfile::TempDir; - - use crate::streaming::Executor; - - //---------// - // Runbook // - //---------// - - fn _assert_is_clone(_x: &T) {} - - struct MockStream { - stages: Vec, - current_stage: usize, - asked_for_maintenance: bool, - } - - impl MockStream { - fn new(stages: Vec) -> Self { - Self { - stages, - current_stage: 0, - asked_for_maintenance: false, - } - } - - fn increment(&mut self) -> usize { - let output = self.current_stage; - self.current_stage += 1; - output - } - - fn current(&self) -> &Stage { - &self.stages[self.current_stage] - } - } - - impl streaming::Stream for MockStream { - type Output = Option; - - fn search(&mut self, args: Search<'_>) -> anyhow::Result> { - if let Stage::Search { groundtruth } = self.current() { - assert_eq!(args.groundtruth, groundtruth.as_path()); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Search stage, instead got {:?}", - self.current() - )) - } - } - - fn insert(&mut self, args: Insert) -> anyhow::Result> { - if let Stage::Insert { - dataset_offsets_and_ids, - } = self.current() - { - assert_eq!(&args.offsets, dataset_offsets_and_ids); - assert_eq!(&args.ids, dataset_offsets_and_ids); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Insert stage, instead got {:?}", - self.current() - )) - } - } - - fn replace(&mut self, args: Replace) -> anyhow::Result> { - if let Stage::Replace { - dataset_offsets, - ids, - } = self.current() - { - assert_eq!(&args.offsets, dataset_offsets); - assert_eq!(&args.ids, ids); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Replace stage, instead got {:?}", - self.current() - )) - } - } - - fn delete(&mut self, args: Delete) -> anyhow::Result> { - if let Stage::Delete { ids } = self.current() { - assert_eq!(&args.ids, ids); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Delete stage, instead got {:?}", - self.current() - )) - } - } - - fn maintain(&mut self, _args: ()) -> anyhow::Result> { - assert!( - self.asked_for_maintenance, - "Stream was not expected to need maintenance" - ); - self.asked_for_maintenance = false; - Ok(None) - } - - fn needs_maintenance(&mut self) -> bool { - let needs = self.asked_for_maintenance; - self.asked_for_maintenance = true; - needs - } - } - - #[test] - fn test_runbook() { - let stages = vec![ - Stage::Insert { - dataset_offsets_and_ids: 0..100, - }, - Stage::Search { - groundtruth: PathBuf::from("gt0"), - }, - Stage::Replace { - dataset_offsets: 100..200, - ids: 0..100, - }, - Stage::Delete { ids: 50..75 }, - Stage::Search { - groundtruth: PathBuf::from("gt1"), - }, - ]; - - let mut runbook = RunBook::new(stages.clone(), 1000).unwrap(); - assert_eq!(runbook.len(), stages.len()); - assert!(!runbook.is_empty()); - assert_eq!(runbook.max_points(), 1000); - - let mut stream = MockStream::new(stages.clone()); - let outputs = runbook.run(&mut stream).unwrap(); - - // Verify that the outputs match the expected stage indices. - let expected_outputs: Vec = (0..stages.len()).collect(); - let non_maintenance: Vec<_> = outputs.iter().filter_map(|o| *o).collect(); - assert_eq!(non_maintenance, expected_outputs); - } - - #[test] - fn test_load_runbook_from_yaml() { - use std::io::Write; - - let temp_dir = TempDir::new().unwrap(); - - // Create groundtruth files for search stages (stages 1 and 7) - File::create(temp_dir.path().join("step1.gt100")).unwrap(); - File::create(temp_dir.path().join("step7.gt100")).unwrap(); - - // Create a YAML runbook with multiple insert, replace, and delete stages - let yaml_content = r#" -test_dataset: - max_pts: 100 - gt_url: "http://example.com/groundtruth" - 0: - operation: "insert" - start: 0 - end: 1000 - 1: - operation: "search" - 2: - operation: "insert" - start: 1000 - end: 2000 - 3: - operation: "delete" - start: 200 - end: 400 - 4: - operation: "replace" - ids_start: 2000 - ids_end: 2500 - tags_start: 400 - tags_end: 900 - 5: - operation: "insert" - start: 2500 - end: 3000 - 6: - operation: "delete" - start: 500 - end: 700 - 7: - operation: "search" -"#; - - let yaml_path = temp_dir.path().join("runbook.yaml"); - { - let mut file = File::create(&yaml_path).unwrap(); - file.write_all(yaml_content.as_bytes()).unwrap(); - } - - // Load the runbook - let mut groundtruth_finder = ScanDirectory::new(temp_dir.path()).unwrap(); - let runbook = RunBook::load(&yaml_path, "test_dataset", &mut groundtruth_finder).unwrap(); - _assert_is_clone(&runbook); - - // Verify the runbook was loaded correctly - assert_eq!(runbook.len(), 8); - assert_eq!(runbook.max_points(), 2300); - assert_eq!(runbook.max_tag(), Some(2999)); - - let stages = runbook.stages(); - - // Stage 0: Insert 0..1000 - assert_eq!( - stages[0], - Stage::Insert { - dataset_offsets_and_ids: 0..1000 - } - ); - - // Stage 1: Search - assert!( - matches!(&stages[1], Stage::Search { groundtruth } if groundtruth.file_name().unwrap() == "step1.gt100") - ); - - // Stage 2: Insert 1000..2000 - assert_eq!( - stages[2], - Stage::Insert { - dataset_offsets_and_ids: 1000..2000 - } - ); - - // Stage 3: Delete 200..400 - assert_eq!(stages[3], Stage::Delete { ids: 200..400 }); - - // Stage 4: Replace offsets 2000..2500, ids 400..900 - assert_eq!( - stages[4], - Stage::Replace { - dataset_offsets: 2000..2500, - ids: 400..900 - } - ); - - // Stage 5: Insert 2500..3000 - assert_eq!( - stages[5], - Stage::Insert { - dataset_offsets_and_ids: 2500..3000 - } - ); - - // Stage 6: Delete 500..700 - assert_eq!(stages[6], Stage::Delete { ids: 500..700 }); - - // Stage 7: Search - assert!( - matches!(&stages[7], Stage::Search { groundtruth } if groundtruth.file_name().unwrap() == "step7.gt100") - ); - } - - //---------------------// - // ScanDirectory Tests // - //---------------------// - - #[test] - fn scan_directory_finds_groundtruth_file() { - let temp_dir = TempDir::new().unwrap(); - - // Create a groundtruth file matching the expected pattern - File::create(temp_dir.path().join("step0.gt100")).unwrap(); - - let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); - let result = scanner.find_groundtruth(0).unwrap(); - assert_eq!(result, temp_dir.path().join("step0.gt100")); - } - - #[test] - fn scan_directory_finds_groundtruth_without_suffix_digits() { - let temp_dir = TempDir::new().unwrap(); - - // Create a groundtruth file with no digits after ".gt" - File::create(temp_dir.path().join("step5.gt")).unwrap(); - - let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); - let result = scanner.find_groundtruth(5).unwrap(); - assert_eq!(result, temp_dir.path().join("step5.gt")); - } - - #[test] - fn scan_directory_errors_when_no_groundtruth_found() { - let temp_dir = TempDir::new().unwrap(); - - // Create some files that don't match the pattern - File::create(temp_dir.path().join("other_file.bin")).unwrap(); - File::create(temp_dir.path().join("step0.other")).unwrap(); - - let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); - let err = scanner.find_groundtruth(0).unwrap_err(); - - let msg = err.to_string(); - assert!(msg.contains("No groundtruth found"), "Got: {}", msg); - } - - #[test] - fn scan_directory_errors_when_multiple_groundtruth_files() { - let temp_dir = TempDir::new().unwrap(); - - // Create multiple groundtruth files for the same stage - File::create(temp_dir.path().join("step0.gt100")).unwrap(); - File::create(temp_dir.path().join("step0.gt200")).unwrap(); - File::create(temp_dir.path().join("step0.gt300")).unwrap(); - - let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); - let err = scanner.find_groundtruth(0).unwrap_err(); - - let msg = err.to_string(); - assert!(msg.contains("Multiple groundtruth files"), "Got: {}", msg); - } - - #[test] - fn scan_directory_ignores_non_digit_suffix() { - let temp_dir = TempDir::new().unwrap(); - - // Create files with non-digit suffixes (should be ignored) - File::create(temp_dir.path().join("step0.gtabc")).unwrap(); - File::create(temp_dir.path().join("step0.gt100")).unwrap(); - - let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); - let result = scanner.find_groundtruth(0).unwrap(); - - // Should find only the valid one - assert_eq!(result, temp_dir.path().join("step0.gt100")); - } - - #[test] - fn scan_directory_errors_on_nonexistent_directory() { - let _ = ScanDirectory::new("/nonexistent/path/that/does/not/exist").unwrap_err(); - } - - #[test] - fn scan_directory_handles_different_stage_indices() { - let temp_dir = TempDir::new().unwrap(); - - File::create(temp_dir.path().join("step0.gt")).unwrap(); - File::create(temp_dir.path().join("step5.gt")).unwrap(); - File::create(temp_dir.path().join("step10.gt")).unwrap(); - - let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); - - assert_eq!( - scanner.find_groundtruth(0).unwrap(), - temp_dir.path().join("step0.gt") - ); - assert_eq!( - scanner.find_groundtruth(5).unwrap(), - temp_dir.path().join("step5.gt") - ); - assert_eq!( - scanner.find_groundtruth(10).unwrap(), - temp_dir.path().join("step10.gt") - ); - - // Stage 1 doesn't exist - assert!(scanner.find_groundtruth(1).is_err()); - } -} diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs b/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs deleted file mode 100644 index 80504111a..000000000 --- a/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs +++ /dev/null @@ -1,678 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::ops::Range; - -use super::Args; -use crate::streaming; - -#[derive(Debug, Clone, Copy, PartialEq)] -struct SimpleRange { - start: usize, - end: usize, -} - -impl SimpleRange { - fn new(start: usize, end: usize) -> Self { - Self { - start, - end: end.max(start), - } - } - - fn len(&self) -> usize { - self.end - self.start - } - - fn is_empty(&self) -> bool { - self.len() == 0 - } - - fn overlaps(&self, other: SimpleRange) -> bool { - self.contains(other.start) || other.contains(self.start) - } - - fn remove(self, other: SimpleRange) -> Remaining { - // If the other range is empty - there's nothing to remove. - if other.is_empty() { - if self.is_empty() { - return Remaining::None; - } else { - return Remaining::One(self); - } - } - - // Five cases to consider: - // - // 1. The two ranges are disjoint - there's nothing to remove. - // 2. `other` strictly contains `self - we remove all of `self`. - // 3. `other` overlaps with the just the start of `self` - return the last part of `self`. - // 4. `other` overlaps with the just the end of `self` - return the last part of `self`. - // 5. `other` is inside of `self` - we must return two pieces. - - if !self.overlaps(other) { - Remaining::One(self) - } else if other.strictly_contains(self) { - Remaining::None - } else if other.start <= self.start { - Remaining::One(Self::new(other.end, self.end)) - } else if other.end >= self.end { - Remaining::One(Self::new(self.start, other.start)) - } else { - assert!(self.strictly_contains(other)); - - Remaining::Two( - Self::new(self.start, other.start), - Self::new(other.end, self.end), - ) - } - } - - fn contains(&self, i: usize) -> bool { - self.as_range().contains(&i) - } - - fn strictly_contains(&self, other: SimpleRange) -> bool { - other.start >= self.start && other.end <= self.end - } - - fn as_range(&self) -> Range { - self.start..self.end - } -} - -impl std::fmt::Display for SimpleRange { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}..{}", self.start, self.end) - } -} - -#[derive(Debug, Clone, Copy, PartialEq)] -enum Remaining { - None, - One(SimpleRange), - Two(SimpleRange, SimpleRange), -} - -impl From> for SimpleRange { - fn from(range: Range) -> Self { - Self { - start: range.start, - end: range.end, - } - } -} - -/// A test `stream` that ensures a Runbook is well formed. -pub(super) struct Validate { - active: Vec, - max_active: usize, - currently_active: usize, - max_tag: Option, -} - -impl Validate { - pub(super) fn new() -> Self { - Self { - active: Vec::new(), - max_active: 0, - currently_active: 0, - max_tag: None, - } - } - - #[cfg(test)] - fn new_test(active: Vec) -> Self { - let currently_active = active.iter().map(|r| r.len()).sum(); - let max_tag = active.iter().filter_map(|r| r.end.checked_sub(1)).max(); - - Self { - active, - max_active: currently_active, - currently_active, - max_tag, - } - } - - /// Return the maximum number of active points seen so far. - pub(super) fn max_active(&self) -> usize { - self.max_active - } - - /// Return the maximum tag seen so far. - pub(super) fn max_tag(&self) -> Option { - self.max_tag - } - - fn update_max(&mut self, tag: usize) { - match self.max_tag.as_mut() { - Some(max) => *max = (*max).max(tag), - None => self.max_tag = Some(tag), - } - } - - /// Merge adjacent ranges where one range's end equals another's start. - /// - /// Assumes the ranges in `active` are already disjoint and ordered by start. - /// After compaction, no two consecutive ranges will have `r1.end == r2.start`. - fn compact(&mut self) { - if self.active.len() <= 1 { - return; - } - - let mut write = 0; - for read in 1..self.active.len() { - let write_end = self.active[write].end; - let read_start = self.active[read].start; - - if write_end == read_start { - // Merge: extend the current range to include the next one - self.active[write].end = self.active[read].end; - } else { - assert!( - read_start > write_end, - "internal invariant has been violated {} > {}", - read_start, - write_end, - ); - - // No merge: move to the next write position - write += 1; - if write != read { - self.active[write] = self.active[read]; - } - } - } - - self.active.truncate(write + 1); - } - - fn find(&self, tags: SimpleRange) -> Slot { - for (i, range) in self.active.iter().enumerate() { - if tags.overlaps(*range) { - return Slot::In(i); - } - if range.start >= tags.end { - return Slot::Before(i); - } - } - Slot::Back - } - - fn insert(&mut self, tags: SimpleRange) -> anyhow::Result<()> { - match self.find(tags) { - Slot::Back => self.active.push(tags), - Slot::Before(i) => self.active.insert(i, tags), - Slot::In(_) => anyhow::bail!("tag range {} is already present for insertion", tags), - } - - // We've successfully added tags, so this update is accurate. - self.currently_active += tags.len(); - self.max_active = self.max_active.max(self.currently_active); - if let Some(tag) = tags.end.checked_sub(1) { - self.update_max(tag); - } - - self.compact(); - Ok(()) - } - - fn replace(&mut self, tags: SimpleRange) -> anyhow::Result<()> { - match self.find(tags) { - Slot::Back | Slot::Before(_) => Err(anyhow::anyhow!( - "tag range {} not valid for replacement", - tags - )), - Slot::In(i) => { - // Due to compaction, if `tags` is truly present in the list of tags, it - // cannot be split across two ranges. - let overlaps = self.active[i]; - if !overlaps.strictly_contains(tags) { - Err(anyhow::anyhow!("could not match the entire range {}", tags)) - } else { - Ok(()) - } - } - } - } - - fn delete(&mut self, tags: SimpleRange) -> anyhow::Result<()> { - match self.find(tags) { - Slot::Back | Slot::Before(_) => { - Err(anyhow::anyhow!("tag range {} not valid for deletion", tags)) - } - Slot::In(i) => { - let current = self.active[i]; - if !current.strictly_contains(tags) { - return Err(anyhow::anyhow!( - "could not match the entire range {} for deletion", - tags - )); - } - - match current.remove(tags) { - Remaining::None => { - self.active.remove(i); - } - Remaining::One(remaining) => self.active[i] = remaining, - Remaining::Two(first, last) => { - self.active.splice(i..(i + 1), [first, last]); - } - } - - // We've successfully remove tags, so this update is accurate. - self.currently_active -= tags.len(); - Ok(()) - } - } - } -} - -/// Location where a range would be located in the sorted ranges vector. -#[derive(Debug, Clone, Copy, PartialEq)] -enum Slot { - /// Range would be placed in the back. It is disjoint from existing ranges. - Back, - - /// Range would be before this index. This implies that the range is disjoin from - /// existing ranges. - Before(usize), - - /// Range partial overlaps with this index. This should be the first such range that - /// overlaps. - In(usize), -} - -impl streaming::Stream for Validate { - type Output = (); - - fn search(&mut self, _args: super::Search<'_>) -> anyhow::Result<()> { - Ok(()) - } - - fn insert(&mut self, args: super::Insert) -> anyhow::Result<()> { - self.insert(args.ids.into()) - } - - fn replace(&mut self, args: super::Replace) -> anyhow::Result<()> { - self.replace(args.ids.into()) - } - - fn delete(&mut self, args: super::Delete) -> anyhow::Result<()> { - self.delete(args.ids.into()) - } - - fn maintain(&mut self, _args: ()) -> anyhow::Result<()> { - Ok(()) - } - - fn needs_maintenance(&mut self) -> bool { - false - } -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - fn range(start: usize, end: usize) -> SimpleRange { - SimpleRange::new(start, end) - } - - #[test] - fn test_range() { - let r = range(0, 1); - assert_eq!(r.start, 0); - assert_eq!(r.end, 1); - assert_eq!(r.len(), 1); - - let r = range(10, 5); - assert_eq!(r.start, 10); - assert_eq!(r.end, 10); - assert_eq!(r.len(), 0); - - let r = range(5, 100); - assert_eq!(r.start, 5); - assert_eq!(r.end, 100); - assert_eq!(r.len(), 95); - } - - #[test] - fn simple_range() { - // Disjoint ranges. - { - let a = range(0, 10); - let b = range(20, 30); - - assert!(!a.overlaps(b), "a = {a}, b = {b}"); - assert!(!b.overlaps(a), "a = {a}, b = {b}"); - - assert!(!a.strictly_contains(b)); - assert!(!b.strictly_contains(a)); - - assert_eq!(a.remove(b), Remaining::One(a)); - assert_eq!(b.remove(a), Remaining::One(b)); - } - - // Disjoint ranges - b - { - let a = range(0, 10); - let b = range(10, 30); - - assert!(!a.overlaps(b), "a = {a}, b = {b}"); - assert!(!b.overlaps(a), "a = {a}, b = {b}"); - - assert!(!a.strictly_contains(b)); - assert!(!b.strictly_contains(a)); - - assert_eq!(a.remove(b), Remaining::One(a)); - assert_eq!(b.remove(a), Remaining::One(b)); - } - - // Partial overlap. - { - let a = range(10, 20); - let b = range(5, 15); - - assert!(a.overlaps(b), "a = {a}, b = {b}"); - assert!(b.overlaps(a), "a = {a}, b = {b}"); - - assert!(!a.strictly_contains(b)); - assert!(!b.strictly_contains(a)); - - assert_eq!(a.remove(b), Remaining::One(range(15, 20))); - assert_eq!(b.remove(a), Remaining::One(range(5, 10))); - } - - // Total overlap. - { - let a = range(10, 20); - let b = range(10, 20); - - assert!(a.overlaps(b), "a = {a}, b = {b}"); - assert!(a.strictly_contains(b)); - assert_eq!(a.remove(b), Remaining::None); - } - - // Totally inside. - { - let a = range(10, 20); - let b = range(12, 18); - - assert!(a.overlaps(b), "a = {a}, b = {b}"); - - assert!(a.strictly_contains(b)); - assert!(!b.strictly_contains(a)); - - assert_eq!(a.remove(b), Remaining::Two(range(10, 12), range(18, 20))); - assert_eq!(b.remove(a), Remaining::None); - } - - // Totally inside - touching left. - { - let a = range(10, 20); - let b = range(10, 15); - - assert!(a.overlaps(b), "a = {a}, b = {b}"); - - assert!(a.strictly_contains(b)); - assert!(!b.strictly_contains(a)); - - assert_eq!(a.remove(b), Remaining::One(range(15, 20))); - assert_eq!(b.remove(a), Remaining::None); - } - - // Totally inside - touching right. - { - let a = range(10, 20); - let b = range(15, 20); - - assert!(a.overlaps(b), "a = {a}, b = {b}"); - - assert!(a.strictly_contains(b)); - assert!(!b.strictly_contains(a)); - - assert_eq!(a.remove(b), Remaining::One(range(10, 15))); - assert_eq!(b.remove(a), Remaining::None); - } - - // Empty ranges. - { - let a = range(10, 10); - assert_eq!(a.remove(a), Remaining::None); - - let a = range(10, 20); - for j in [5, 10, 15, 20, 25] { - let b = range(j, j); - assert_eq!(a.remove(b), Remaining::One(a)); - } - } - } - - #[test] - fn test_compact() { - // Empty - { - let mut v = Validate::new_test(vec![]); - v.compact(); - assert_eq!(&v.active, &[]); - } - - // One - { - let start = vec![range(10, 20)]; - let mut v = Validate::new_test(start.clone()); - v.compact(); - assert_eq!(v.active, start); - } - - // No change. - { - let start = vec![range(10, 20), range(30, 40), range(50, 60)]; - - let mut v = Validate::new_test(start.clone()); - v.compact(); - - assert_eq!(v.active, start); - } - - // Merge in multiple locations with empty elements. - { - let active = vec![ - range(0, 5), // + Merged - range(5, 5), // | - range(5, 5), // | - range(5, 10), // | - range(20, 30), - range(35, 40), // + Merged - range(40, 41), // | - range(41, 41), // | - range(41, 50), // | - range(55, 60), - range(70, 70), // + Merged - range(70, 75), // | - range(75, 100), // | - ]; - - let mut v = Validate::new_test(active); - v.compact(); - - let expected = [ - range(0, 10), - range(20, 30), - range(35, 50), - range(55, 60), - range(70, 100), - ]; - - assert_eq!(&*v.active, &expected); - } - } - - #[test] - fn test_find() { - // Empty - { - let v = Validate::new_test(vec![]); - assert_eq!(v.find(range(0, 10)), Slot::Back); - } - - // One - { - let v = Validate::new_test(vec![range(10, 20)]); - - assert_eq!(v.find(range(0, 5)), Slot::Before(0)); - assert_eq!(v.find(range(5, 10)), Slot::Before(0)); - assert_eq!(v.find(range(8, 12)), Slot::In(0)); - assert_eq!(v.find(range(10, 15)), Slot::In(0)); - assert_eq!(v.find(range(15, 20)), Slot::In(0)); - assert_eq!(v.find(range(18, 22)), Slot::In(0)); - assert_eq!(v.find(range(20, 25)), Slot::Back); - } - - // Multiple - { - let v = Validate::new_test(vec![range(10, 20), range(30, 40), range(50, 60)]); - - assert_eq!(v.find(range(0, 5)), Slot::Before(0)); - assert_eq!(v.find(range(5, 10)), Slot::Before(0)); - assert_eq!(v.find(range(15, 25)), Slot::In(0)); - - assert_eq!(v.find(range(20, 25)), Slot::Before(1)); - assert_eq!(v.find(range(28, 35)), Slot::In(1)); - assert_eq!(v.find(range(35, 45)), Slot::In(1)); - - assert_eq!(v.find(range(45, 50)), Slot::Before(2)); - assert_eq!(v.find(range(55, 65)), Slot::In(2)); - assert_eq!(v.find(range(65, 70)), Slot::Back); - } - } - - #[test] - fn test_insert() { - let mut v = Validate::new_test(vec![]); - - v.insert(range(10, 20)).unwrap(); - assert_eq!(&*v.active, &[range(10, 20)]); - assert_eq!(v.max_active(), 10); - assert_eq!(v.max_tag(), Some(19)); - - v.insert(range(30, 40)).unwrap(); - assert_eq!(&*v.active, &[range(10, 20), range(30, 40)]); - assert_eq!(v.max_active(), 20); - assert_eq!(v.max_tag(), Some(39)); - - v.insert(range(20, 30)).unwrap(); - assert_eq!(&*v.active, &[range(10, 40)]); - assert!(v.max_active() == 30); - assert_eq!(v.max_tag(), Some(39)); - - let result = v.insert(range(15, 25)); - assert!(result.is_err()); - - v.insert(range(0, 5)).unwrap(); - assert_eq!(&*v.active, &[range(0, 5), range(10, 40)]); - assert_eq!(v.max_active(), 35); - assert_eq!(v.max_tag(), Some(39)); - - v.insert(range(7, 8)).unwrap(); - assert_eq!(&*v.active, &[range(0, 5), range(7, 8), range(10, 40)]); - assert_eq!(v.max_active(), 36); - assert_eq!(v.max_tag(), Some(39)); - - v.insert(range(4, 8)).unwrap_err(); - v.insert(range(50, 60)).unwrap(); - assert_eq!( - &*v.active, - &[range(0, 5), range(7, 8), range(10, 40), range(50, 60)] - ); - assert_eq!(v.max_active(), 46); - assert_eq!(v.max_tag(), Some(59)); - - v.insert(range(5, 7)).unwrap(); - assert_eq!(&*v.active, &[range(0, 8), range(10, 40), range(50, 60)]); - assert_eq!(v.max_active(), 48); - assert_eq!(v.max_tag(), Some(59)); - - v.insert(range(40, 50)).unwrap(); - assert_eq!(&*v.active, &[range(0, 8), range(10, 60)]); - assert_eq!(v.max_active(), 58); - assert_eq!(v.max_tag(), Some(59)); - - v.insert(range(8, 10)).unwrap(); - assert_eq!(&*v.active, &[range(0, 60)]); - assert_eq!(v.max_active(), 60); - assert_eq!(v.max_tag(), Some(59)); - } - - #[test] - fn test_replace() { - let mut v = Validate::new_test(vec![range(10, 20), range(30, 40)]); - - // Success - v.replace(range(12, 18)).unwrap(); - v.replace(range(10, 20)).unwrap(); - v.replace(range(15, 15)).unwrap(); - v.replace(range(30, 35)).unwrap(); - - // Failure - v.replace(range(15, 25)).unwrap_err(); - v.replace(range(5, 15)).unwrap_err(); - v.replace(range(25, 35)).unwrap_err(); - v.replace(range(40, 50)).unwrap_err(); - } - - #[test] - fn test_delete() { - let mut v = Validate::new_test(vec![range(10, 20), range(30, 40)]); - - assert_eq!(v.max_active(), 20); - - // Partial deletions - v.delete(range(15, 18)).unwrap(); - assert_eq!(&*v.active, &[range(10, 15), range(18, 20), range(30, 40)]); - - v.delete(range(35, 36)).unwrap(); - assert_eq!( - &*v.active, - &[range(10, 15), range(18, 20), range(30, 35), range(36, 40)] - ); - - // Partial at beginning - v.delete(range(10, 12)).unwrap(); - assert_eq!( - &*v.active, - &[range(12, 15), range(18, 20), range(30, 35), range(36, 40)] - ); - - // Partial at end - v.delete(range(32, 35)).unwrap(); - assert_eq!( - &*v.active, - &[range(12, 15), range(18, 20), range(30, 32), range(36, 40)] - ); - - // Full deletions - v.delete(range(18, 20)).unwrap(); - assert_eq!(&*v.active, &[range(12, 15), range(30, 32), range(36, 40)]); - - v.delete(range(36, 40)).unwrap(); - assert_eq!(&*v.active, &[range(12, 15), range(30, 32)]); - - assert_eq!(v.max_active(), 20); - assert_eq!(v.max_tag(), Some(39)); - - // Failure cases. - v.delete(range(20, 25)).unwrap_err(); - v.delete(range(0, 10)).unwrap_err(); - v.delete(range(10, 14)).unwrap_err(); - v.delete(range(31, 33)).unwrap_err(); - v.delete(range(40, 50)).unwrap_err(); - } -} diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt deleted file mode 100644 index ebc0da280..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Unrecognized key of type sequence \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml deleted file mode 100644 index 1c17b82fb..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - [1, 2, 3]: - operation: delete - start: 50 - end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt deleted file mode 100644 index 55341d4f6..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -my_dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt deleted file mode 100644 index 4718a6dec..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt +++ /dev/null @@ -1 +0,0 @@ -dataset "my_dataset" not found in file "/runbook.yaml" - possible alternatives: [other_dataset, another_dataset] \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml deleted file mode 100644 index 30553b966..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml +++ /dev/null @@ -1,13 +0,0 @@ -other_dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - -another_dataset: - max_pts: 200 - 0: - operation: insert - start: 0 - end: 200 diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt deleted file mode 100644 index 3eca2590b..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt +++ /dev/null @@ -1,5 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: YAML type is not a map \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml deleted file mode 100644 index a10034f9a..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - 1: - - operation: delete - start: 50 - end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt deleted file mode 100644 index 1bfa6bbee..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse "delete" - 2: start (80) must be less than end (40) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml deleted file mode 100644 index 79b218eb4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: delete - start: 80 - end: 40 diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt deleted file mode 100644 index 2794d487e..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 0 - 1: trying to parse an "insert" - 2: start (50) must be less than end (25) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml deleted file mode 100644 index e5901a5d4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 50 - end: 25 diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt deleted file mode 100644 index 7bd551584..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 0 - 1: trying to parse an "insert" - 2: start (50) must be less than end (50) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml deleted file mode 100644 index 93dfb43f1..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 50 - end: 50 diff --git a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt deleted file mode 100644 index 08b6dac7a..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - key "max_pts" not found \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml deleted file mode 100644 index c95ee86f6..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml +++ /dev/null @@ -1,5 +0,0 @@ -dataset: - 0: - operation: insert - start: 0 - end: 100 diff --git a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt deleted file mode 100644 index 44c1c4328..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - failed to parse "max_pts" as a usize \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml deleted file mode 100644 index 0dbab4cdd..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: "not_an_integer" - 0: - operation: insert - start: 0 - end: 100 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt deleted file mode 100644 index b7397a054..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: ids_start (50) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml deleted file mode 100644 index 4b07f23ff..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 50 - ids_end: 50 - tags_start: 0 - tags_end: 10 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt deleted file mode 100644 index cb5e83c8c..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: ids_start (100) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml deleted file mode 100644 index a1f70e2b6..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 100 - ids_end: 50 - tags_start: 0 - tags_end: 10 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt deleted file mode 100644 index 014298aa0..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: tags_start (80) must be less than tags_end (40) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml deleted file mode 100644 index de5f1426b..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 0 - ids_end: 50 - tags_start: 80 - tags_end: 40 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt deleted file mode 100644 index ffbbf0d4f..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: tags_start (25) must be less than tags_end (25) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml deleted file mode 100644 index 498be2272..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 0 - ids_end: 50 - tags_start: 25 - tags_end: 25 diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt deleted file mode 100644 index 8781fc588..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -while opening file "/runbook.yaml" - -Caused by: - No such file or directory (os error 2) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/dataset.txt deleted file mode 100644 index e69de29bb..000000000 diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt deleted file mode 100644 index 837a2194a..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt +++ /dev/null @@ -1 +0,0 @@ -invalid type: string "item1 item2", expected a YAML mapping \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml deleted file mode 100644 index de1ea531d..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml +++ /dev/null @@ -1,2 +0,0 @@ -item1 -item2 \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt deleted file mode 100644 index 9143f3e00..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Stage "1.1" must be an integer \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml deleted file mode 100644 index a37c6f0c1..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - 1.1: - operation: delete - start: 50 - end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt deleted file mode 100644 index 8bafcb6c3..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Unrecognized runbook key: "not_an_integer" \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml deleted file mode 100644 index 581e83e11..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - not_an_integer: - operation: delete - start: 50 - end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt deleted file mode 100644 index 527ba9779..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Unrecognized runbook key: "unknown_key" \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml deleted file mode 100644 index 4418781d4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml +++ /dev/null @@ -1,7 +0,0 @@ -dataset: - max_pts: 10 - unknown_key: 42 - 0: - operation: insert - start: 0 - end: 100 diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt deleted file mode 100644 index 5c20d5c1c..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt +++ /dev/null @@ -1,5 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 0 - 1: unrecognized operation: upsert \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml deleted file mode 100644 index f1b77d852..000000000 --- a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: upsert - start: 0 - end: 50 diff --git a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs index 1d39e59cb..6b475b1f9 100644 --- a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs +++ b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs @@ -48,6 +48,8 @@ use crate::{ //////////////////////// type BfTreeFPIndex = Arc>>; +type BfTreeFPStreamer = bigann::WithData>; +type BfTreeFPStreamingResult = anyhow::Result<(BfTreeFPStreamer, BfTreeFPIndex)>; /// The bf_tree streaming index implementation. /// @@ -224,13 +226,7 @@ where } } -fn bftree_streaming( - input: &BfTreeDynamicRun, - max_points: usize, -) -> anyhow::Result<( - bigann::WithData>, - BfTreeFPIndex, -)> +fn bftree_streaming(input: &BfTreeDynamicRun, max_points: usize) -> BfTreeFPStreamingResult where T: bytemuck::Pod + VectorRepr + SampleableForStart, { diff --git a/diskann-benchmark/src/index/bftree/spherical_streaming.rs b/diskann-benchmark/src/index/bftree/spherical_streaming.rs index d676b9994..bc33a0431 100644 --- a/diskann-benchmark/src/index/bftree/spherical_streaming.rs +++ b/diskann-benchmark/src/index/bftree/spherical_streaming.rs @@ -52,6 +52,8 @@ use crate::{ type BfTreeSQProvider = BfTreeProvider; type BfTreeSQIndex = Arc>; +type BfTreeSQStreamer = bigann::WithData>; +type BfTreeSQStreamingResult = anyhow::Result<(BfTreeSQStreamer, BfTreeSQIndex)>; fn new_quantizer( quantizer: SphericalQuantizer, @@ -229,10 +231,7 @@ impl Benchmark for StreamingSpherical { fn bftree_sq_streaming_impl( input: &BfTreeSphericalDynamicRun, max_points: usize, -) -> anyhow::Result<( - bigann::WithData>, - BfTreeSQIndex, -)> { +) -> BfTreeSQStreamingResult { let topk = match input.search_phase() { SearchPhase::Topk(topk) => topk, _ => anyhow::bail!("Only TopK is currently supported by the streaming index"), diff --git a/diskann-tools/Cargo.toml b/diskann-tools/Cargo.toml index 644ee1173..bcac24631 100644 --- a/diskann-tools/Cargo.toml +++ b/diskann-tools/Cargo.toml @@ -14,7 +14,7 @@ byteorder.workspace = true 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-utils = { 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..561f310a2 --- /dev/null +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -0,0 +1,320 @@ +/* + * 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_utils::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_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::streaming::executors::bigann::{FindGroundtruth, RunBook, Stage}; +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 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()); + + // Boolean active-vector mask indexed by base offset. + let mut active: Vec = vec![false; n_base]; + // For each active dataset offset, the external ID that should appear in the groundtruth. + // For plain inserts external_id == dataset_offset; for Replace they diverge. + let mut ext_id: Vec = vec![0u32; n_base]; + // Reverse map: external_id -> dataset_offset, needed to resolve Delete/Replace removals. + let mut ext_to_offset: HashMap = HashMap::new(); + + println!("Using distance function: {:?}", args.distance_function); + + let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); + + for (stage_idx, stage) in runbook.stages().iter().enumerate() { + match stage { + Stage::Insert { + dataset_offsets_and_ids, + } => { + for id in dataset_offsets_and_ids.clone() { + if id < n_base { + active[id] = true; + ext_id[id] = id as u32; + ext_to_offset.insert(id as u32, id); + } + } + tracing::info!( + "Stage {}: insert {}..{} ({} active)", + stage_idx, + dataset_offsets_and_ids.start, + dataset_offsets_and_ids.end, + active.iter().filter(|&&b| b).count(), + ); + } + Stage::Delete { ids } => { + for eid in ids.clone() { + if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { + active[offset] = false; + ext_to_offset.remove(&(eid as u32)); + } + } + tracing::info!( + "Stage {}: delete {}..{} ({} active)", + stage_idx, + ids.start, + ids.end, + active.iter().filter(|&&b| b).count(), + ); + } + Stage::Replace { + dataset_offsets, + ids, + } => { + // Remove old vectors by external ID. + for eid in ids.clone() { + if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { + active[offset] = false; + ext_to_offset.remove(&(eid as u32)); + } + } + // Insert new vectors: they inherit the external IDs from `ids`. + for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) { + if offset < n_base { + active[offset] = true; + ext_id[offset] = eid as u32; + ext_to_offset.insert(eid as u32, offset); + } + } + tracing::info!( + "Stage {}: replace ids {}..{} with offsets {}..{} ({} active)", + stage_idx, + ids.start, + ids.end, + dataset_offsets.start, + dataset_offsets.end, + active.iter().filter(|&&b| b).count(), + ); + } + Stage::Search { + groundtruth: output_path, + } => { + let timer = Instant::now(); + let n_active = active.iter().filter(|&&b| b).count(); + + tracing::info!( + "Stage {}: computing top-{} groundtruth for {} active vectors against {} queries", + stage_idx, recall_at, n_active, n_queries, + ); + + let active_ids: Vec = active + .iter() + .enumerate() + .filter_map(|(i, &on)| if on { Some(i) } else { None }) + .collect(); + + // Compute KNN in parallel over queries. + // Use ext_id[offset] as the neighbor ID so that groundtruth IDs + // match external IDs (which differ from dataset offsets after Replace). + + // Using the global threadpool is fine here + #[allow(clippy::disallowed_methods)] + let results: Vec> = (0..n_queries) + .into_par_iter() + .map(|qi| { + let query = queries.row(qi); + let mut pq = NeighborPriorityQueue::new(recall_at); + for &offset in &active_ids { + let dist = distance_fn.evaluate_similarity(dataset.row(offset), query); + pq.insert(Neighbor { + id: ext_id[offset], + distance: dist, + }); + } + pq + }) + .collect(); + + // Warn about queries that got fewer than K results (active set smaller than K). + let under_k: Vec = results + .iter() + .enumerate() + .filter_map(|(qi, pq)| { + if pq.size() < recall_at { + Some(qi) + } else { + None + } + }) + .collect(); + if !under_k.is_empty() { + tracing::warn!( + "Stage {}: {} / {} queries have fewer than {} results (active set = {}). \ + Query indices: {:?}", + stage_idx, + under_k.len(), + n_queries, + recall_at, + n_active, + under_k, + ); + } + + write_ground_truth::<()>( + &storage, + output_path.to_str().ok_or_else(|| CMDToolError { + details: format!("Non-UTF8 output path: {}", output_path.display()), + })?, + n_queries, + recall_at, + results, + None, + ) + .map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + tracing::info!( + "Stage {}: groundtruth written to {} in {:?}", + stage_idx, + output_path.display(), + timer.elapsed(), + ); + } + } + } + + tracing::info!("Done."); + Ok(()) +} + +#[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, diff --git a/diskann-utils/Cargo.toml b/diskann-utils/Cargo.toml index ea358601e..6d789ae32 100644 --- a/diskann-utils/Cargo.toml +++ b/diskann-utils/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true edition.workspace = true [dependencies] +anyhow.workspace = true half = { workspace = true, features = ["rand_distr", "num-traits"] } cfg-if.workspace = true rayon = { workspace = true, optional = true } @@ -17,6 +18,7 @@ rand = { workspace = true } rand_distr = { workspace = true } diskann-wide = { workspace = true } bytemuck = { workspace = true, features = ["must_cast"] } +serde_yaml = { version = "0.9.34", optional = true } [lints] workspace = true @@ -32,6 +34,8 @@ workspace = true cfg-if.workspace = true rand.workspace = true rstest.workspace = true +tempfile.workspace = true +diskann-benchmark-runner = { workspace = true, features = ["ux-tools"] } [features] @@ -41,3 +45,5 @@ default = ["rayon"] rayon = ["dep:rayon"] # Enable testing utilities like test_data_root() testing = [] +# Enable BigANN runbook parsing and execution support. +bigann = ["dep:serde_yaml"] diff --git a/diskann-utils/src/lib.rs b/diskann-utils/src/lib.rs index 089e0dece..54d6ceff6 100644 --- a/diskann-utils/src/lib.rs +++ b/diskann-utils/src/lib.rs @@ -17,6 +17,7 @@ pub mod future; pub mod io; pub mod object_pool; pub mod sampling; +pub mod streaming; // Views pub mod strided; @@ -45,3 +46,36 @@ pub fn test_data_directory() -> &'static str { pub fn test_data_root() -> std::path::PathBuf { workspace_root().join(test_data_directory()) } + +/// # Notes on Testing +/// +/// BigANN parsing UX tests use checked-in fixtures under the crate's `tests` directory. +/// Expected outputs can be regenerated by running tests with `DISKANN_TEST=overwrite`. +#[cfg(all(test, feature = "bigann"))] +mod ux { + const ENV_VAR: &str = "DISKANN_TEST"; + + /// Check if we should overwrite expected files. + pub(crate) fn should_overwrite() -> bool { + match std::env::var(ENV_VAR) { + Ok(v) if v == "overwrite" => true, + Ok(v) => panic!( + "Unknown value for {}: \"{}\". Expected \"overwrite\"", + ENV_VAR, v + ), + Err(std::env::VarError::NotPresent) => false, + Err(std::env::VarError::NotUnicode(_)) => { + panic!("Value for {} is not unicode", ENV_VAR) + } + } + } + + pub(crate) fn help() -> String { + format!("{}=overwrite", ENV_VAR) + } + + pub(crate) fn test_dir() -> std::path::PathBuf { + let manifest: &std::path::Path = env!("CARGO_MANIFEST_DIR").as_ref(); + manifest.join("tests") + } +} From 87a3ad54fd83da23562adb26456a1fec1c95d483 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 17 Jul 2026 18:40:01 +0000 Subject: [PATCH 02/10] add moved tests folder --- .../dataset-key-not-integer-or-string/dataset.txt | 1 + .../dataset-key-not-integer-or-string/expected.txt | 4 ++++ .../dataset-key-not-integer-or-string/runbook.yaml | 10 ++++++++++ .../tests/bigann-ux/dataset-not-found/dataset.txt | 1 + .../tests/bigann-ux/dataset-not-found/expected.txt | 1 + .../tests/bigann-ux/dataset-not-found/runbook.yaml | 13 +++++++++++++ .../bigann-ux/dataset-value-not-mapping/dataset.txt | 1 + .../dataset-value-not-mapping/expected.txt | 5 +++++ .../dataset-value-not-mapping/runbook.yaml | 10 ++++++++++ .../bigann-ux/delete-invalid-range/dataset.txt | 1 + .../bigann-ux/delete-invalid-range/expected.txt | 6 ++++++ .../bigann-ux/delete-invalid-range/runbook.yaml | 10 ++++++++++ .../bigann-ux/insert-invalid-range/dataset.txt | 1 + .../bigann-ux/insert-invalid-range/expected.txt | 6 ++++++ .../bigann-ux/insert-invalid-range/runbook.yaml | 6 ++++++ .../bigann-ux/insert-start-equals-end/dataset.txt | 1 + .../bigann-ux/insert-start-equals-end/expected.txt | 6 ++++++ .../bigann-ux/insert-start-equals-end/runbook.yaml | 6 ++++++ .../tests/bigann-ux/missing-max-pts/dataset.txt | 1 + .../tests/bigann-ux/missing-max-pts/expected.txt | 4 ++++ .../tests/bigann-ux/missing-max-pts/runbook.yaml | 5 +++++ .../tests/bigann-ux/non-integer-max-pts/dataset.txt | 1 + .../bigann-ux/non-integer-max-pts/expected.txt | 4 ++++ .../bigann-ux/non-integer-max-pts/runbook.yaml | 6 ++++++ .../replace-ids-start-equals-end/dataset.txt | 1 + .../replace-ids-start-equals-end/expected.txt | 6 ++++++ .../replace-ids-start-equals-end/runbook.yaml | 12 ++++++++++++ .../bigann-ux/replace-invalid-ids-range/dataset.txt | 1 + .../replace-invalid-ids-range/expected.txt | 6 ++++++ .../replace-invalid-ids-range/runbook.yaml | 12 ++++++++++++ .../replace-invalid-tags-range/dataset.txt | 1 + .../replace-invalid-tags-range/expected.txt | 6 ++++++ .../replace-invalid-tags-range/runbook.yaml | 12 ++++++++++++ .../replace-tags-start-equals-end/dataset.txt | 1 + .../replace-tags-start-equals-end/expected.txt | 6 ++++++ .../replace-tags-start-equals-end/runbook.yaml | 12 ++++++++++++ .../bigann-ux/runbook-does-not-exist/dataset.txt | 1 + .../bigann-ux/runbook-does-not-exist/expected.txt | 4 ++++ .../bigann-ux/runbook-not-yaml-map/dataset.txt | 0 .../bigann-ux/runbook-not-yaml-map/expected.txt | 1 + .../bigann-ux/runbook-not-yaml-map/runbook.yaml | 2 ++ .../bigann-ux/stage-key-not-integer/dataset.txt | 1 + .../bigann-ux/stage-key-not-integer/expected.txt | 4 ++++ .../bigann-ux/stage-key-not-integer/runbook.yaml | 10 ++++++++++ .../bigann-ux/stage-key-not-number/dataset.txt | 1 + .../bigann-ux/stage-key-not-number/expected.txt | 4 ++++ .../bigann-ux/stage-key-not-number/runbook.yaml | 10 ++++++++++ .../bigann-ux/unrecognized-dataset-key/dataset.txt | 1 + .../bigann-ux/unrecognized-dataset-key/expected.txt | 4 ++++ .../bigann-ux/unrecognized-dataset-key/runbook.yaml | 7 +++++++ .../bigann-ux/unrecognized-operation/dataset.txt | 1 + .../bigann-ux/unrecognized-operation/expected.txt | 5 +++++ .../bigann-ux/unrecognized-operation/runbook.yaml | 6 ++++++ 53 files changed, 248 insertions(+) create mode 100644 diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml create mode 100644 diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt create mode 100644 diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt create mode 100644 diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml diff --git a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt new file mode 100644 index 000000000..ebc0da280 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Unrecognized key of type sequence \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml new file mode 100644 index 000000000..1c17b82fb --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + [1, 2, 3]: + operation: delete + start: 50 + end: 60 diff --git a/diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt b/diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt new file mode 100644 index 000000000..55341d4f6 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt @@ -0,0 +1 @@ +my_dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt b/diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt new file mode 100644 index 000000000..4718a6dec --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt @@ -0,0 +1 @@ +dataset "my_dataset" not found in file "/runbook.yaml" - possible alternatives: [other_dataset, another_dataset] \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml b/diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml new file mode 100644 index 000000000..30553b966 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml @@ -0,0 +1,13 @@ +other_dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + +another_dataset: + max_pts: 200 + 0: + operation: insert + start: 0 + end: 200 diff --git a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt new file mode 100644 index 000000000..3eca2590b --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt @@ -0,0 +1,5 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: YAML type is not a map \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml new file mode 100644 index 000000000..a10034f9a --- /dev/null +++ b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + 1: + - operation: delete + start: 50 + end: 60 diff --git a/diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt b/diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt b/diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt new file mode 100644 index 000000000..1bfa6bbee --- /dev/null +++ b/diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse "delete" + 2: start (80) must be less than end (40) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml b/diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml new file mode 100644 index 000000000..79b218eb4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: delete + start: 80 + end: 40 diff --git a/diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt b/diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt b/diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt new file mode 100644 index 000000000..2794d487e --- /dev/null +++ b/diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 0 + 1: trying to parse an "insert" + 2: start (50) must be less than end (25) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml b/diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml new file mode 100644 index 000000000..e5901a5d4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 50 + end: 25 diff --git a/diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt b/diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt b/diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt new file mode 100644 index 000000000..7bd551584 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 0 + 1: trying to parse an "insert" + 2: start (50) must be less than end (50) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml b/diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml new file mode 100644 index 000000000..93dfb43f1 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 50 + end: 50 diff --git a/diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt b/diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt b/diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt new file mode 100644 index 000000000..08b6dac7a --- /dev/null +++ b/diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + key "max_pts" not found \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml b/diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml new file mode 100644 index 000000000..c95ee86f6 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml @@ -0,0 +1,5 @@ +dataset: + 0: + operation: insert + start: 0 + end: 100 diff --git a/diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt b/diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt b/diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt new file mode 100644 index 000000000..44c1c4328 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + failed to parse "max_pts" as a usize \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml b/diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml new file mode 100644 index 000000000..0dbab4cdd --- /dev/null +++ b/diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: "not_an_integer" + 0: + operation: insert + start: 0 + end: 100 diff --git a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt new file mode 100644 index 000000000..b7397a054 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: ids_start (50) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml new file mode 100644 index 000000000..4b07f23ff --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 50 + ids_end: 50 + tags_start: 0 + tags_end: 10 diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt new file mode 100644 index 000000000..cb5e83c8c --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: ids_start (100) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml new file mode 100644 index 000000000..a1f70e2b6 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 100 + ids_end: 50 + tags_start: 0 + tags_end: 10 diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt new file mode 100644 index 000000000..014298aa0 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: tags_start (80) must be less than tags_end (40) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml new file mode 100644 index 000000000..de5f1426b --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 0 + ids_end: 50 + tags_start: 80 + tags_end: 40 diff --git a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt new file mode 100644 index 000000000..ffbbf0d4f --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: tags_start (25) must be less than tags_end (25) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml new file mode 100644 index 000000000..498be2272 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 0 + ids_end: 50 + tags_start: 25 + tags_end: 25 diff --git a/diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt b/diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt b/diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt new file mode 100644 index 000000000..8781fc588 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt @@ -0,0 +1,4 @@ +while opening file "/runbook.yaml" + +Caused by: + No such file or directory (os error 2) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt new file mode 100644 index 000000000..e69de29bb diff --git a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt new file mode 100644 index 000000000..837a2194a --- /dev/null +++ b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt @@ -0,0 +1 @@ +invalid type: string "item1 item2", expected a YAML mapping \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml new file mode 100644 index 000000000..de1ea531d --- /dev/null +++ b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml @@ -0,0 +1,2 @@ +item1 +item2 \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt b/diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt b/diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt new file mode 100644 index 000000000..9143f3e00 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Stage "1.1" must be an integer \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml b/diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml new file mode 100644 index 000000000..a37c6f0c1 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + 1.1: + operation: delete + start: 50 + end: 60 diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt b/diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt b/diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt new file mode 100644 index 000000000..8bafcb6c3 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Unrecognized runbook key: "not_an_integer" \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml b/diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml new file mode 100644 index 000000000..581e83e11 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + not_an_integer: + operation: delete + start: 50 + end: 60 diff --git a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt new file mode 100644 index 000000000..527ba9779 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Unrecognized runbook key: "unknown_key" \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml new file mode 100644 index 000000000..4418781d4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml @@ -0,0 +1,7 @@ +dataset: + max_pts: 10 + unknown_key: 42 + 0: + operation: insert + start: 0 + end: 100 diff --git a/diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt b/diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt b/diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt new file mode 100644 index 000000000..5c20d5c1c --- /dev/null +++ b/diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt @@ -0,0 +1,5 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 0 + 1: unrecognized operation: upsert \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml b/diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml new file mode 100644 index 000000000..f1b77d852 --- /dev/null +++ b/diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: 100 + 0: + operation: upsert + start: 0 + end: 50 From bea772bfefe91c322ba25c045209413b160fc7ec Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 17 Jul 2026 19:06:11 +0000 Subject: [PATCH 03/10] add forgotten files --- diskann-utils/src/streaming/api.rs | 127 ++++ .../src/streaming/executors/bigann/mod.rs | 17 + .../src/streaming/executors/bigann/parsing.rs | 502 +++++++++++++++ .../src/streaming/executors/bigann/runbook.rs | 597 ++++++++++++++++++ .../streaming/executors/bigann/validate.rs | 272 ++++++++ diskann-utils/src/streaming/executors/mod.rs | 12 + diskann-utils/src/streaming/mod.rs | 14 + 7 files changed, 1541 insertions(+) create mode 100644 diskann-utils/src/streaming/api.rs create mode 100644 diskann-utils/src/streaming/executors/bigann/mod.rs create mode 100644 diskann-utils/src/streaming/executors/bigann/parsing.rs create mode 100644 diskann-utils/src/streaming/executors/bigann/runbook.rs create mode 100644 diskann-utils/src/streaming/executors/bigann/validate.rs create mode 100644 diskann-utils/src/streaming/executors/mod.rs create mode 100644 diskann-utils/src/streaming/mod.rs diff --git a/diskann-utils/src/streaming/api.rs b/diskann-utils/src/streaming/api.rs new file mode 100644 index 000000000..70251fd92 --- /dev/null +++ b/diskann-utils/src/streaming/api.rs @@ -0,0 +1,127 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::any::Any; + +/// A streaming interface for performing dynamic (streaming) operations on an index. +pub trait Stream +where + A: Arguments, +{ + /// Output type for all operations. The `'static` is to allow results to be + /// aggregated in [`Any`] for type erasure in higher level [`Executor`]s. + type Output: 'static; + + /// Perform a search operation. + fn search(&mut self, args: A::Search<'_>) -> anyhow::Result; + + /// Perform an insert operation. + fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result; + + /// Perform a replace operation. + fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result; + + /// Perform a delete operation. + fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result; + + /// Perform a maintain operation. + fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result; + + /// Indicate whether or not maintenance is needed. [`Executor`] implementations + /// are responsible periodically checking this. + fn needs_maintenance(&mut self) -> bool; +} + +/// Operation arguments to [`Stream`]. +pub trait Arguments: 'static { + /// Argument to [`Stream::search`]. + type Search<'a>; + /// Argument to [`Stream::insert`]. + type Insert<'a>; + /// Argument to [`Stream::replace`]. + type Replace<'a>; + /// Argument to [`Stream::delete`]. + type Delete<'a>; + /// Argument to [`Stream::maintain`]. + type Maintain<'a>; +} + +/// A sequential executor for [`Stream`]s. +pub trait Executor { + /// The argument collection type for the underlying [`Stream`]. + type Args: Arguments; + + /// Execute a series of operations on `stream`. As outputs are produced, they will be + /// passed to `collect` for aggregation. + fn run_with(&mut self, stream: &mut S, collect: F) -> anyhow::Result<()> + where + S: Stream, + O: 'static, + F: FnMut(O) -> anyhow::Result<()>; + + /// Execute a series of operations on `stream`. The outputs of each operation will be + /// collected in the returned `Vec` in-order. + fn run(&mut self, stream: &mut S) -> anyhow::Result> + where + S: Stream, + { + let mut outputs = Vec::new(); + self.run_with(stream, |output| { + outputs.push(output); + Ok(()) + })?; + Ok(outputs) + } +} + +/// A type-erased [`Stream`] implementation that wraps stream outputs in [`Box`]. +#[derive(Debug)] +pub struct AnyStream<'a, T>(&'a mut T); + +impl<'a, T> AnyStream<'a, T> { + /// Wrap `stream` in an [`AnyStream`]. + pub fn new(stream: &'a mut T) -> Self { + Self(stream) + } +} + +fn boxed(x: T) -> Box +where + T: Any, +{ + Box::new(x) +} + +impl Stream for AnyStream<'_, T> +where + A: Arguments, + T: Stream, +{ + type Output = Box; + + fn search(&mut self, args: A::Search<'_>) -> anyhow::Result { + self.0.search(args).map(boxed) + } + + fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result { + self.0.insert(args).map(boxed) + } + + fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result { + self.0.replace(args).map(boxed) + } + + fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result { + self.0.delete(args).map(boxed) + } + + fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result { + self.0.maintain(args).map(boxed) + } + + fn needs_maintenance(&mut self) -> bool { + self.0.needs_maintenance() + } +} diff --git a/diskann-utils/src/streaming/executors/bigann/mod.rs b/diskann-utils/src/streaming/executors/bigann/mod.rs new file mode 100644 index 000000000..191467cb9 --- /dev/null +++ b/diskann-utils/src/streaming/executors/bigann/mod.rs @@ -0,0 +1,17 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! # BigANN Runbook Support +//! +//! This module provides an executor for processing BigANN-style runbooks. +//! Please refer to the [`RunBook`] documentation for more details. + +mod parsing; +mod runbook; +mod validate; + +pub use runbook::{ + Args, Delete, FindGroundtruth, Insert, Replace, RunBook, ScanDirectory, Search, Stage, +}; diff --git a/diskann-utils/src/streaming/executors/bigann/parsing.rs b/diskann-utils/src/streaming/executors/bigann/parsing.rs new file mode 100644 index 000000000..b98689343 --- /dev/null +++ b/diskann-utils/src/streaming/executors/bigann/parsing.rs @@ -0,0 +1,502 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::path::{Path, PathBuf}; + +use anyhow::Context; +use serde_yaml::{Mapping, Value}; + +/// See [`super::RunBook::load`] for documentation of the file format parsed +/// by this function. +pub(super) fn load( + path: &Path, + dataset: &str, + groundtruth: &mut dyn super::FindGroundtruth, +) -> anyhow::Result { + load_mapping(parse_yaml(path)?, path, dataset, groundtruth) +} + +fn load_mapping( + mapping: Mapping, + path: &Path, + dataset: &str, + groundtruth: &mut dyn super::FindGroundtruth, +) -> anyhow::Result { + let dataset_value = match mapping.get(dataset) { + Some(value) => value, + None => return Err(DumpKeys::new(mapping, dataset, path).into()), + }; + + let dataset_mapping: &Mapping = match dataset_value.try_as() { + Ok(mapping) => mapping, + Err(_) => anyhow::bail!( + "dataset \"{}\" exists in file \"{}\", but its associated payload is not a YAML map", + dataset, + path.display(), + ), + }; + + let mut raw = parse_stages(dataset_mapping).with_context(|| { + format!( + "parsing dataset \"{}\" in file \"{}\"", + dataset, + path.display() + ) + })?; + raw.stages.sort_by_key(|s| s.index); + + let context = |index: usize| { + format!( + "precessing stage {} of dataset \"{}\" in file \"{}\"", + index, + dataset, + path.display() + ) + }; + + let stages: anyhow::Result> = raw + .stages + .iter() + .map(|stage| { + let stage = match &stage.operation { + Operation::Search => super::Stage::Search { + groundtruth: groundtruth + .find_groundtruth(stage.index) + .with_context(|| context(stage.index))?, + }, + Operation::Insert(insert) => super::Stage::Insert { + dataset_offsets_and_ids: insert.start..insert.end, + }, + Operation::Replace(replace) => super::Stage::Replace { + dataset_offsets: replace.ids_start..replace.ids_end, + ids: replace.tags_start..replace.tags_end, + }, + Operation::Delete(delete) => super::Stage::Delete { + ids: delete.start..delete.end, + }, + }; + Ok(stage) + }) + .collect(); + + super::RunBook::new(stages?, raw.max_points) +} + +fn parse_yaml(path: &Path) -> anyhow::Result { + let f = std::fs::File::open(path) + .with_context(|| format!("while opening file \"{}\"", path.display()))?; + + Ok(serde_yaml::from_reader(std::io::BufReader::new(f))?) +} + +fn parse_stages(mapping: &Mapping) -> anyhow::Result { + let mut stages = Vec::::new(); + let mut max_points = None; + mapping.iter().try_for_each(|(key, value)| match key { + Value::String(s) => match s.as_str() { + "max_pts" => { + let points: usize = value + .try_as() + .map_err(|kind| anyhow::anyhow!("failed to parse \"max_pts\" as a {}", kind))?; + + max_points = Some(points); + Ok(()) + } + "gt_url" => Ok(()), + _ => anyhow::bail!("Unrecognized runbook key: \"{}\"", s), + }, + Value::Number(stage) => match stage.as_i64() { + Some(stage) => { + stages.push( + handle_stage(stage as usize, value) + .with_context(|| format!("processing stage {}", stage))?, + ); + Ok(()) + } + None => anyhow::bail!("Stage \"{}\" must be an integer", stage), + }, + _ => anyhow::bail!("Unrecognized key of type {}", classify(key),), + })?; + + let max_points = match max_points { + Some(points) => points, + None => anyhow::bail!("key \"max_pts\" not found"), + }; + + Ok(Raw { max_points, stages }) +} + +fn handle_stage(index: usize, value: &Value) -> anyhow::Result { + let mapping: &Mapping = value + .try_as() + .map_err(|_| anyhow::anyhow!("YAML type is not a map"))?; + + let kind: &str = mapping.get_as("operation")?; + let operation = match Kind::try_parse(kind)? { + Kind::Search => Operation::Search, + Kind::Insert => Operation::Insert(mapping.try_into()?), + Kind::Replace => Operation::Replace(mapping.try_into()?), + Kind::Delete => Operation::Delete(mapping.try_into()?), + }; + + Ok(Stage { index, operation }) +} + +struct Raw { + max_points: usize, + stages: Vec, +} + +struct Stage { + index: usize, + operation: Operation, +} + +enum Operation { + Search, + Insert(Insert), + Replace(Replace), + Delete(Delete), +} + +#[derive(Debug)] +struct DumpKeys { + mapping: Mapping, + dataset: String, + file: PathBuf, +} + +impl DumpKeys { + #[inline(never)] + fn new(mapping: Mapping, dataset: &str, file: &Path) -> Self { + Self { + mapping, + dataset: dataset.to_string(), + file: file.into(), + } + } +} + +impl std::fmt::Display for DumpKeys { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "dataset \"{}\" not found in file \"{}\" - possible alternatives: [", + self.dataset, + self.file.display(), + )?; + + let len = self.mapping.len(); + self.mapping.keys().enumerate().try_for_each(|(i, key)| { + let mut write = |key: &dyn std::fmt::Display| { + if i + 1 == len { + write!(f, "{}", key) + } else { + write!(f, "{}, ", key) + } + }; + + match key { + Value::Null => write(&"null"), + Value::Bool(b) => write(b), + Value::Number(number) => write(number), + Value::String(s) => write(s), + Value::Sequence(_) => write(&""), + Value::Mapping(_) => write(&""), + Value::Tagged(_) => write(&""), + } + })?; + write!(f, "]") + } +} + +impl std::error::Error for DumpKeys {} + +trait TryAs<'a, T> { + fn try_as(&'a self) -> Result; +} + +impl<'a> TryAs<'a, usize> for Value { + fn try_as(&'a self) -> Result { + self.as_i64().map(|i| i as usize).ok_or("usize") + } +} + +impl<'a> TryAs<'a, &'a str> for Value { + fn try_as(&'a self) -> Result<&'a str, &'static str> { + self.as_str().ok_or("string") + } +} + +impl<'a> TryAs<'a, &'a Mapping> for Value { + fn try_as(&'a self) -> Result<&'a Mapping, &'static str> { + self.as_mapping().ok_or("map") + } +} + +trait MappingExt { + fn get_as<'a, T>(&'a self, index: &str) -> anyhow::Result + where + Value: TryAs<'a, T>; +} + +impl MappingExt for Mapping { + fn get_as<'a, T>(&'a self, key: &str) -> anyhow::Result + where + Value: TryAs<'a, T>, + { + match self.get(key) { + Some(value) => match value.try_as() { + Ok(v) => Ok(v), + Err(expected) => Err(anyhow::anyhow!( + "key \"{}\" exists but it is not a {}", + key, + expected, + )), + }, + None => Err(anyhow::anyhow!("key \"{}\" not found", key)), + } + } +} + +#[derive(Debug, Clone, Copy)] +enum Kind { + Search, + Insert, + Replace, + Delete, +} + +impl Kind { + fn try_parse(string: &str) -> anyhow::Result { + match string { + "search" => Ok(Kind::Search), + "insert" => Ok(Kind::Insert), + "replace" => Ok(Kind::Replace), + "delete" => Ok(Kind::Delete), + _ => Err(anyhow::anyhow!("unrecognized operation: {}", string)), + } + } +} + +#[derive(Debug)] +struct Replace { + ids_start: usize, + ids_end: usize, + tags_start: usize, + tags_end: usize, +} + +impl TryFrom<&Mapping> for Replace { + type Error = anyhow::Error; + fn try_from(mapping: &Mapping) -> anyhow::Result { + let inner = || -> anyhow::Result { + let this = Self { + ids_start: mapping.get_as("ids_start")?, + ids_end: mapping.get_as("ids_end")?, + tags_start: mapping.get_as("tags_start")?, + tags_end: mapping.get_as("tags_end")?, + }; + if this.ids_start >= this.ids_end { + anyhow::bail!( + "ids_start ({}) must be less than ids_end ({})", + this.ids_start, + this.ids_end + ); + } + if this.tags_start >= this.tags_end { + anyhow::bail!( + "tags_start ({}) must be less than tags_end ({})", + this.tags_start, + this.tags_end + ); + } + Ok(this) + }; + + inner().context("trying to parse an \"replace\"") + } +} + +#[derive(Debug)] +struct Insert { + start: usize, + end: usize, +} + +impl TryFrom<&Mapping> for Insert { + type Error = anyhow::Error; + fn try_from(mapping: &Mapping) -> anyhow::Result { + let inner = || -> anyhow::Result { + let this = Self { + start: mapping.get_as("start")?, + end: mapping.get_as("end")?, + }; + if this.start >= this.end { + anyhow::bail!( + "start ({}) must be less than end ({})", + this.start, + this.end + ); + } + Ok(this) + }; + + inner().context("trying to parse an \"insert\"") + } +} + +#[derive(Debug)] +struct Delete { + start: usize, + end: usize, +} + +impl TryFrom<&Mapping> for Delete { + type Error = anyhow::Error; + fn try_from(mapping: &Mapping) -> anyhow::Result { + let inner = || -> anyhow::Result { + let this = Self { + start: mapping.get_as("start")?, + end: mapping.get_as("end")?, + }; + if this.start >= this.end { + anyhow::bail!( + "start ({}) must be less than end ({})", + this.start, + this.end + ); + } + Ok(this) + }; + + inner().context("trying to parse \"delete\"") + } +} + +fn classify(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "bool", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Sequence(_) => "sequence", + Value::Mapping(_) => "mapping", + Value::Tagged(_) => "tagged", + } +} + +#[cfg(test)] +mod ux_tests { + use super::*; + + use diskann_benchmark_runner::ux as runner_ux; + + const TEST_DATA_DIR: &str = "bigann-ux"; + const RUNBOOK_FILE: &str = "runbook.yaml"; + const DATASET_FILE: &str = "dataset.txt"; + const EXPECTED_FILE: &str = "expected.txt"; + const PATH_PLACEHOLDER: &str = ""; + + fn fixup_paths_and_os_errors(s: &str, test_dir: &Path) -> String { + let native_path = test_dir.display().to_string(); + let forward_slash_path = native_path.replace('\\', "/"); + + const NOT_FOUND_WINDOWS: &str = "The system cannot find the file specified."; + const NOT_FOUND_LINUX: &str = "No such file or directory"; + + s.replace(&native_path, PATH_PLACEHOLDER) + .replace(&forward_slash_path, PATH_PLACEHOLDER) + .replace('\\', "/") + .replace(NOT_FOUND_WINDOWS, NOT_FOUND_LINUX) + } + + struct FailingGroundtruth; + + impl super::super::FindGroundtruth for FailingGroundtruth { + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { + Err(anyhow::anyhow!( + "groundtruth not available for stage {}", + stage + )) + } + } + + fn run_file_test(test_dir: &Path) { + let runbook_path = test_dir.join(RUNBOOK_FILE); + let dataset_path = test_dir.join(DATASET_FILE); + let expected_path = test_dir.join(EXPECTED_FILE); + + let dataset = std::fs::read_to_string(&dataset_path) + .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", dataset_path, e)); + let dataset = dataset.trim(); + + let mut groundtruth = FailingGroundtruth; + let result = load(&runbook_path, dataset, &mut groundtruth); + + let actual_output = match result { + Ok(_) => panic!( + "Test {:?} expected an error but parsing succeeded", + test_dir.file_name().unwrap() + ), + Err(err) => format!("{:?}", err), + }; + + let actual_portable = fixup_paths_and_os_errors(&actual_output, test_dir); + let actual_normalized = runner_ux::strip_backtrace(runner_ux::normalize(actual_portable)); + + if crate::ux::should_overwrite() { + std::fs::write(&expected_path, &actual_normalized) + .unwrap_or_else(|e| panic!("Failed to write {:?}: {}", expected_path, e)); + println!("Overwrote {:?}", expected_path); + } else { + let expected = std::fs::read_to_string(&expected_path) + .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", expected_path, e)); + let expected_normalized = runner_ux::normalize(expected); + + if actual_normalized != expected_normalized { + panic!( + "Test {:?} failed.\n\nExpected:\n---\n{}\n---\n\nActual:\n---\n{}\n---\nIf this is expected, run with {} to update the expected output.", + test_dir.file_name().unwrap(), + expected_normalized, + actual_normalized, + crate::ux::help(), + ); + } + } + } + + fn run_all_file_tests() { + let test_data_path = crate::ux::test_dir().join(TEST_DATA_DIR); + if !test_data_path.exists() { + println!( + "No test_data directory found at {:?}, skipping file-based tests", + test_data_path + ); + return; + } + + let mut found_tests = false; + for entry in std::fs::read_dir(&test_data_path) + .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", test_data_path, e)) + { + let entry = entry.unwrap(); + if entry.file_type().unwrap().is_dir() { + found_tests = true; + println!("Running file-based test: {:?}", entry.file_name()); + run_file_test(&entry.path()); + } + } + + if !found_tests { + panic!("No test directories found in {:?}", test_data_path); + } + } + + #[test] + fn file_based_error_tests() { + run_all_file_tests(); + } +} diff --git a/diskann-utils/src/streaming/executors/bigann/runbook.rs b/diskann-utils/src/streaming/executors/bigann/runbook.rs new file mode 100644 index 000000000..e9d4bada6 --- /dev/null +++ b/diskann-utils/src/streaming/executors/bigann/runbook.rs @@ -0,0 +1,597 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + any::Any, + ops::Range, + path::{Path, PathBuf}, +}; + +use anyhow::Context; + +use crate::streaming::{self, Executor, Stream}; + +use super::{parsing, validate}; + +/// An executor for [BigANN-style runbooks](https://github.com/harsha-simhadri/big-ann-benchmarks/tree/main/neurips23/streaming). +#[derive(Debug, Clone)] +pub struct RunBook { + // The individual runbook stages. + stages: Vec, + // The maximum number of active points at any point in the runbook. + max_points: usize, + // The maximum tag referenced in the runbook. + max_tag: Option, +} + +impl RunBook { + /// Loads a [`RunBook`] from a YAML file at `path` for the specified `dataset`. + /// + /// # Groundtruth Resolution + /// + /// When loading a runbook, search stages may require groundtruth files to be present. + /// The index of these stages will be provided to the `groundtruth` argument for resolution. + /// If working with the standard BigANN benchmark format, consider using [`ScanDirectory`] + /// and providing the directory where the BigANN framework downloads the groundtruth files. + /// + /// Note that the implementation here currently does not attempt to download any groundtruth + /// files using the `gt_url` field; it is merely parsed and ignored. + /// + /// # YAML File Format + /// + /// The top-level structure is a mapping from dataset names (strings) to their + /// corresponding runbook definitions: + /// ```yaml + /// dataset_name_1: + /// # runbook definition... + /// dataset_name_2: + /// # runbook definition... + /// ``` + /// The `dataset` parameter specifies which dataset's runbook to load from the file. + /// + /// Each runbook definition has the following format: + /// + /// ```yaml + /// max_pts: # required + /// [gt_url]: # ignored + /// 0: + /// # stage definition ... + /// 1: + /// # stage definition ... + /// ... + /// ``` + /// Entries need not be in order, but stages must be sequentially numbered starting + /// from `0`. Each stage takes one of four forms (described below). + /// + /// ## Search Stage + /// + /// Merely specifies that a search should be performed. The queries must be provided + /// externally. The groundtruth file for this stage is located via the provided + /// [`FindGroundtruth::find_groundtruth`] method, which will be provided with the stage index. + /// + /// ```yaml + /// : + /// operation: "search" + /// ``` + /// + /// ## Insert Stage + /// + /// Insert vectors from the underlying dataset into the index. The vectors to insert + /// are specified by the range `start..end`, which serves as both the offsets of the + /// vectors in the dataset and their external ids. + /// + /// ```yaml + /// : + /// operation: "insert" + /// start: + /// end: + /// ``` + /// + /// ## Replace Stage + /// + /// Replace vectors in the index with vectors from the underlying dataset. Unlike insertions, + /// replace operations distinguish between the dataset offsets (`ids_start..ids_end`) + /// and the external ids (tags) of the vectors to replace (`tags_start..tags_end`). + /// + /// These operations indicate that the vectors in the index tagged by `tags_start..tags_end` should + /// be replaced with the vectors from the dataset at offsets `ids_start..ids_end`. + /// + /// ```yaml + /// : + /// operation: "replace" + /// ids_start: + /// ids_end: + /// tags_start: + /// tags_end: + /// ``` + /// + /// ## Delete Stage + /// + /// Delete vectors from the index with external ids in the range `start..end`. + /// + /// ```yaml + /// : + /// operation: "delete" + /// start: + /// end: + /// ``` + /// + /// # File Validation + /// + /// The loading framework does a limited amount of validation on the YAML file: + /// + /// 1. The specified `dataset` must exist in the top-level mapping. + /// 2. The `max_pts` key must be present and be a valid integer. Its value is verified and if found to + /// be inaccurate, it is updated to the correct value internally. + /// 3. Each stage must be sequentially numbered starting from `0` with no gaps. + /// 4. Each stage must have a valid operation with all required fields present and of the correct type. + /// 5. Groundtruth files must be resolvable via the provided [`FindGroundtruth`] implementation. + pub fn load( + path: &Path, + dataset: &str, + groundtruth: &mut dyn FindGroundtruth, + ) -> anyhow::Result { + parsing::load(path, dataset, groundtruth) + } + + pub(super) fn new(stages: Vec, max_points: usize) -> anyhow::Result { + let mut this = Self { + stages, + max_points, + max_tag: None, + }; + + let mut validator = validate::Validate::new(); + this.run_with(&mut validator, |_| Ok(()))?; + + this.max_points = this.max_points.max(validator.max_active()); + this.max_tag = validator.max_tag(); + + Ok(this) + } + + /// Returns the maximum number of points specified in the runbook. + pub fn max_points(&self) -> usize { + self.max_points + } + + /// Returns the maximum tag specified in the runbook. + pub fn max_tag(&self) -> Option { + self.max_tag + } + + /// Returns the number of stages in the runbook. + pub fn len(&self) -> usize { + self.stages.len() + } + + /// Returns `true` if the runbook contains no stages. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns a reference to the stages in this runbook. + pub fn stages(&self) -> &[Stage] { + &self.stages + } + + /// Executes the runbook by iterating through each stage. + /// + /// When calling this method, the dynamic type of `stream`'s output must be + /// compatible with the dynamic type expected by `collect`. + fn run_with_internal( + &self, + stream: &mut dyn streaming::Stream>, + collect: &mut dyn FnMut(Box) -> anyhow::Result<()>, + ) -> anyhow::Result<()> { + for (i, stage) in self.stages.iter().enumerate() { + #[derive(Clone, Copy)] + struct OnStage(usize, usize); + + impl std::fmt::Display for OnStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "on stage {} of {}", self.0, self.1) + } + } + + let context = OnStage(i, self.len()); + + if stream.needs_maintenance() { + collect(stream.maintain(()).context(context)?).context(context)?; + } + + let output = match stage { + Stage::Search { groundtruth } => { + let args = Search { groundtruth }; + stream.search(args).context(context)? + } + Stage::Insert { + dataset_offsets_and_ids, + } => { + let args = Insert { + offsets: dataset_offsets_and_ids.clone(), + ids: dataset_offsets_and_ids.clone(), + }; + stream.insert(args).context(context)? + } + Stage::Replace { + dataset_offsets, + ids, + } => { + let args = Replace { + offsets: dataset_offsets.clone(), + ids: ids.clone(), + }; + stream.replace(args).context(context)? + } + Stage::Delete { ids } => { + let args = Delete { ids: ids.clone() }; + stream.delete(args).context(context)? + } + }; + + collect(output).context(context)?; + } + + Ok(()) + } +} + +/// An operation in a BigANN runbook. +#[derive(Debug, Clone, PartialEq)] +pub enum Stage { + /// Perform a search operation using the specified groundtruth file. + Search { + /// The resolved path to the groundtruth file for this search stage. + groundtruth: PathBuf, + }, + /// Insert vectors from the dataset into the index. + Insert { + /// The offsets in the dataset, which also serve as the external ids for + /// the inserted vectors. + dataset_offsets_and_ids: Range, + }, + /// Replace vectors in the index with new vectors from the dataset. + Replace { + /// The offsets in the dataset for the replacement vectors. + dataset_offsets: Range, + /// The external ids of the vectors to be replaced. + ids: Range, + }, + /// Delete vectors from the index. + Delete { + /// The external ids of the vectors to delete. + ids: Range, + }, +} + +/// Arguments for a BigANN runbook "search" stage. +#[derive(Debug)] +pub struct Search<'a> { + /// The resolved path to the file containing the groundtruth for this stage. + pub groundtruth: &'a Path, +} + +/// Arguments for a BigANN runbook "insert" stage. +pub struct Insert { + /// The range of offsets in the dataset for vectors to insert. + pub offsets: Range, + /// The external ids to assign to the inserted vectors. + pub ids: Range, +} + +/// Arguments for a BigANN runbook "replace" stage. +pub struct Replace { + /// The range of offsets in the dataset for the replacement vectors. + pub offsets: Range, + /// The external ids of the vectors to replace. + pub ids: Range, +} + +/// Arguments for a BigANN runbook "delete" stage. +pub struct Delete { + /// The range of external ids to delete from the index. + pub ids: Range, +} + +/// The argument type for a [`RunBook`]. +#[derive(Debug, Clone, Copy)] +pub struct Args; + +impl streaming::Arguments for Args { + type Search<'a> = Search<'a>; + type Insert<'a> = Insert; + type Replace<'a> = Replace; + type Delete<'a> = Delete; + type Maintain<'a> = (); +} + +impl streaming::Executor for RunBook { + type Args = Args; + + fn run_with(&mut self, stream: &mut S, mut collect: F) -> anyhow::Result<()> + where + S: Stream, + O: 'static, + F: FnMut(O) -> anyhow::Result<()>, + { + self.run_with_internal(&mut streaming::AnyStream::new(stream), &mut |any| { + let typed = *any + .downcast::() + .expect("the dynamic type should be configured correctly"); + collect(typed) + }) + } +} + +/// A trait for resolving groundtruth files for search stages in a [`RunBook`]. +/// +/// Implementors of this trait provide the logic to locate groundtruth files +/// given a stage index. See [`ScanDirectory`] for a common implementation. +pub trait FindGroundtruth { + /// Resolves the groundtruth file path for the specified `stage` index. + /// + /// This method is only called for "search" stages in a runbook. + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result; +} + +/// A [`FindGroundtruth`] implementation that scans a directory for groundtruth files +/// matching the expected BigANN naming convention: `step{stage}.gt[0-9]*` where +/// `{stage}` substitutes the stage index formatted as a string. +#[derive(Debug)] +pub struct ScanDirectory { + directory: PathBuf, + + // Cached contents of the files in `directory`. + files: Vec, +} + +impl ScanDirectory { + /// Creates a new [`ScanDirectory`] instance for the specified `directory`. + /// + /// # Notes + /// + /// This constructor scans the directory and caches its contents. + /// If the directory contents change after creation, the instance will not + /// reflect those changes. + /// + /// This is meant for the common benchmarking scenario where the benchmarking + /// machine is generally static while benchmarks are executed. + pub fn new(directory: impl Into) -> anyhow::Result { + Self::new_(directory.into()) + } + + fn new_(directory: PathBuf) -> anyhow::Result { + // Read all files in the directory. + let read_dir = std::fs::read_dir(&directory).with_context(|| { + format!( + "while trying to read the contents of {}", + directory.display() + ) + })?; + + let files = read_dir + .filter_map(|entry| { + let entry = entry.ok()?; + let file_type = entry.file_type().ok()?; + if file_type.is_file() { + Some(entry.file_name().to_string_lossy().into()) + } else { + None + } + }) + .collect(); + + Ok(Self { directory, files }) + } +} + +impl FindGroundtruth for ScanDirectory { + /// Finds the groundtruth file for the specified `stage` by scanning the directory. + /// + /// # Errors + /// + /// Returns an error if no matching file is found or if multiple matches exist. + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { + let prefix = format!("step{}.gt", stage); + + enum Matches<'a> { + None, + One(&'a str), + Many(Vec<&'a str>), + } + + impl<'a> Matches<'a> { + fn push(&mut self, file: &'a str) { + *self = match std::mem::replace(self, Self::None) { + Self::None => Self::One(file), + Self::One(first) => Self::Many(vec![first, file]), + Self::Many(mut all) => { + all.push(file); + Self::Many(all) + } + }; + } + } + + let mut matches = Matches::None; + + for file in &self.files { + if file.starts_with(&prefix) { + let suffix = &file[prefix.len()..]; + if suffix.chars().all(|c| c.is_ascii_digit()) { + matches.push(file); + } + } + } + + match matches { + Matches::One(m) => Ok(self.directory.join(m)), + Matches::None => Err(anyhow::anyhow!( + "No groundtruth found for step {} in \"{}\", expected pattern: \"step{}.gt[0-9]*\"", + stage, + self.directory.display(), + stage, + )), + Matches::Many(matches) => Err(anyhow::anyhow!( + "Multiple groundtruth files found for step {} in \"{}\": {:?}", + stage, + self.directory.display(), + matches, + )), + } + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use crate::streaming::Executor; + + //---------// + // Runbook // + //---------// + + fn _assert_is_clone(_x: &T) {} + + struct MockStream { + stages: Vec, + current_stage: usize, + asked_for_maintenance: bool, + } + + impl MockStream { + fn new(stages: Vec) -> Self { + Self { + stages, + current_stage: 0, + asked_for_maintenance: false, + } + } + + fn increment(&mut self) -> usize { + let output = self.current_stage; + self.current_stage += 1; + output + } + + fn current(&self) -> &Stage { + &self.stages[self.current_stage] + } + } + + impl streaming::Stream for MockStream { + type Output = Option; + + fn search(&mut self, args: Search<'_>) -> anyhow::Result> { + if let Stage::Search { groundtruth } = self.current() { + assert_eq!(args.groundtruth, groundtruth.as_path()); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Search stage, instead got {:?}", + self.current() + )) + } + } + + fn insert(&mut self, args: Insert) -> anyhow::Result> { + if let Stage::Insert { + dataset_offsets_and_ids, + } = self.current() + { + assert_eq!(&args.offsets, dataset_offsets_and_ids); + assert_eq!(&args.ids, dataset_offsets_and_ids); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Insert stage, instead got {:?}", + self.current() + )) + } + } + + fn replace(&mut self, args: Replace) -> anyhow::Result> { + if let Stage::Replace { + dataset_offsets, + ids, + } = self.current() + { + assert_eq!(&args.offsets, dataset_offsets); + assert_eq!(&args.ids, ids); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Replace stage, instead got {:?}", + self.current() + )) + } + } + + fn delete(&mut self, args: Delete) -> anyhow::Result> { + if let Stage::Delete { ids } = self.current() { + assert_eq!(&args.ids, ids); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Delete stage, instead got {:?}", + self.current() + )) + } + } + + fn maintain(&mut self, _args: ()) -> anyhow::Result> { + assert!( + self.asked_for_maintenance, + "Stream was not expected to need maintenance" + ); + self.asked_for_maintenance = false; + Ok(None) + } + + fn needs_maintenance(&mut self) -> bool { + let needs = self.asked_for_maintenance; + self.asked_for_maintenance = true; + needs + } + } + + #[test] + fn test_runbook() { + let stages = vec![ + Stage::Insert { + dataset_offsets_and_ids: 0..100, + }, + Stage::Search { + groundtruth: PathBuf::from("gt0"), + }, + Stage::Replace { + dataset_offsets: 100..200, + ids: 0..100, + }, + Stage::Delete { ids: 50..75 }, + Stage::Search { + groundtruth: PathBuf::from("gt1"), + }, + ]; + + let mut runbook = RunBook::new(stages.clone(), 1000).unwrap(); + assert_eq!(runbook.len(), stages.len()); + assert!(!runbook.is_empty()); + assert_eq!(runbook.max_points(), 1000); + + let mut stream = MockStream::new(stages.clone()); + let outputs = runbook.run(&mut stream).unwrap(); + + let expected_outputs: Vec = (0..stages.len()).collect(); + let non_maintenance: Vec<_> = outputs.iter().filter_map(|o| *o).collect(); + assert_eq!(non_maintenance, expected_outputs); + } +} diff --git a/diskann-utils/src/streaming/executors/bigann/validate.rs b/diskann-utils/src/streaming/executors/bigann/validate.rs new file mode 100644 index 000000000..05ca5933a --- /dev/null +++ b/diskann-utils/src/streaming/executors/bigann/validate.rs @@ -0,0 +1,272 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::ops::Range; + +use super::Args; +use crate::streaming; + +#[derive(Debug, Clone, Copy, PartialEq)] +struct SimpleRange { + start: usize, + end: usize, +} + +impl SimpleRange { + fn new(start: usize, end: usize) -> Self { + Self { + start, + end: end.max(start), + } + } + + fn len(&self) -> usize { + self.end - self.start + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn overlaps(&self, other: SimpleRange) -> bool { + self.contains(other.start) || other.contains(self.start) + } + + fn remove(self, other: SimpleRange) -> Remaining { + if other.is_empty() { + if self.is_empty() { + return Remaining::None; + } + return Remaining::One(self); + } + + if !self.overlaps(other) { + Remaining::One(self) + } else if other.strictly_contains(self) { + Remaining::None + } else if other.start <= self.start { + Remaining::One(Self::new(other.end, self.end)) + } else if other.end >= self.end { + Remaining::One(Self::new(self.start, other.start)) + } else { + assert!(self.strictly_contains(other)); + + Remaining::Two( + Self::new(self.start, other.start), + Self::new(other.end, self.end), + ) + } + } + + fn contains(&self, i: usize) -> bool { + self.as_range().contains(&i) + } + + fn strictly_contains(&self, other: SimpleRange) -> bool { + other.start >= self.start && other.end <= self.end + } + + fn as_range(&self) -> Range { + self.start..self.end + } +} + +impl std::fmt::Display for SimpleRange { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}..{}", self.start, self.end) + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum Remaining { + None, + One(SimpleRange), + Two(SimpleRange, SimpleRange), +} + +impl From> for SimpleRange { + fn from(range: Range) -> Self { + Self { + start: range.start, + end: range.end, + } + } +} + +pub(super) struct Validate { + active: Vec, + max_active: usize, + currently_active: usize, + max_tag: Option, +} + +impl Validate { + pub(super) fn new() -> Self { + Self { + active: Vec::new(), + max_active: 0, + currently_active: 0, + max_tag: None, + } + } + + /// Return the maximum number of active points seen so far. + pub(super) fn max_active(&self) -> usize { + self.max_active + } + + /// Return the maximum tag seen so far. + pub(super) fn max_tag(&self) -> Option { + self.max_tag + } + + fn update_max(&mut self, tag: usize) { + match self.max_tag.as_mut() { + Some(max) => *max = (*max).max(tag), + None => self.max_tag = Some(tag), + } + } + + fn compact(&mut self) { + if self.active.len() <= 1 { + return; + } + + let mut write = 0; + for read in 1..self.active.len() { + let write_end = self.active[write].end; + let read_start = self.active[read].start; + + if write_end == read_start { + self.active[write].end = self.active[read].end; + } else { + assert!( + read_start > write_end, + "internal invariant has been violated {} > {}", + read_start, + write_end, + ); + + write += 1; + if write != read { + self.active[write] = self.active[read]; + } + } + } + + self.active.truncate(write + 1); + } + + fn find(&self, tags: SimpleRange) -> Slot { + for (i, range) in self.active.iter().enumerate() { + if tags.overlaps(*range) { + return Slot::In(i); + } + if range.start >= tags.end { + return Slot::Before(i); + } + } + Slot::Back + } + + fn insert(&mut self, tags: SimpleRange) -> anyhow::Result<()> { + match self.find(tags) { + Slot::Back => self.active.push(tags), + Slot::Before(i) => self.active.insert(i, tags), + Slot::In(_) => anyhow::bail!("tag range {} is already present for insertion", tags), + } + + self.currently_active += tags.len(); + self.max_active = self.max_active.max(self.currently_active); + if let Some(tag) = tags.end.checked_sub(1) { + self.update_max(tag); + } + + self.compact(); + Ok(()) + } + + fn replace(&mut self, tags: SimpleRange) -> anyhow::Result<()> { + match self.find(tags) { + Slot::Back | Slot::Before(_) => Err(anyhow::anyhow!( + "tag range {} not valid for replacement", + tags + )), + Slot::In(i) => { + let overlaps = self.active[i]; + if !overlaps.strictly_contains(tags) { + Err(anyhow::anyhow!("could not match the entire range {}", tags)) + } else { + Ok(()) + } + } + } + } + + fn delete(&mut self, tags: SimpleRange) -> anyhow::Result<()> { + match self.find(tags) { + Slot::Back | Slot::Before(_) => { + Err(anyhow::anyhow!("tag range {} not valid for deletion", tags)) + } + Slot::In(i) => { + let current = self.active[i]; + if !current.strictly_contains(tags) { + return Err(anyhow::anyhow!( + "could not match the entire range {} for deletion", + tags + )); + } + + match current.remove(tags) { + Remaining::None => { + self.active.remove(i); + } + Remaining::One(remaining) => self.active[i] = remaining, + Remaining::Two(first, last) => { + self.active.splice(i..(i + 1), [first, last]); + } + } + + self.currently_active -= tags.len(); + Ok(()) + } + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum Slot { + Back, + Before(usize), + In(usize), +} + +impl streaming::Stream for Validate { + type Output = (); + + fn search(&mut self, _args: super::Search<'_>) -> anyhow::Result<()> { + Ok(()) + } + + fn insert(&mut self, args: super::Insert) -> anyhow::Result<()> { + self.insert(args.ids.into()) + } + + fn replace(&mut self, args: super::Replace) -> anyhow::Result<()> { + self.replace(args.ids.into()) + } + + fn delete(&mut self, args: super::Delete) -> anyhow::Result<()> { + self.delete(args.ids.into()) + } + + fn maintain(&mut self, _args: ()) -> anyhow::Result<()> { + Ok(()) + } + + fn needs_maintenance(&mut self) -> bool { + false + } +} diff --git a/diskann-utils/src/streaming/executors/mod.rs b/diskann-utils/src/streaming/executors/mod.rs new file mode 100644 index 000000000..18876a5e5 --- /dev/null +++ b/diskann-utils/src/streaming/executors/mod.rs @@ -0,0 +1,12 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! # Provided Executors. +//! +//! Built-in implementations of [`Executor`] for common streaming inputs. + +#[cfg(feature = "bigann")] +#[cfg_attr(docsrs, doc(cfg(feature = "bigann")))] +pub mod bigann; diff --git a/diskann-utils/src/streaming/mod.rs b/diskann-utils/src/streaming/mod.rs new file mode 100644 index 000000000..9bd3b311c --- /dev/null +++ b/diskann-utils/src/streaming/mod.rs @@ -0,0 +1,14 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! # Support for Streaming Operations. +//! +//! Streaming operations are defined by sequences of insertions, deletions, and replacements +//! with the goal to help study how an algorithm performs under dynamic workloads. + +mod api; +pub use api::{AnyStream, Arguments, Executor, Stream}; + +pub mod executors; From 3550ea475d0eda363fa6153afdb3f60e557d2ea1 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Tue, 21 Jul 2026 16:25:47 +0000 Subject: [PATCH 04/10] move runbook infra back to diskann-benchmark-core, restore files --- Cargo.lock | 1 + diskann-benchmark-core/Cargo.toml | 4 +- diskann-benchmark-core/src/lib.rs | 42 ++ .../src/streaming/executors/bigann/mod.rs | 9 +- .../src/streaming/executors/bigann/parsing.rs | 502 +++++++++++++++ .../src/streaming/executors/bigann/runbook.rs | 597 ++++++++++++++++++ .../streaming/executors/bigann/validate.rs | 272 ++++++++ .../dataset.txt | 1 + .../expected.txt | 4 + .../runbook.yaml | 10 + .../bigann-ux/dataset-not-found/dataset.txt | 1 + .../bigann-ux/dataset-not-found/expected.txt | 1 + .../bigann-ux/dataset-not-found/runbook.yaml | 13 + .../dataset-value-not-mapping/dataset.txt | 1 + .../dataset-value-not-mapping/expected.txt | 5 + .../dataset-value-not-mapping/runbook.yaml | 10 + .../delete-invalid-range/dataset.txt | 1 + .../delete-invalid-range/expected.txt | 6 + .../delete-invalid-range/runbook.yaml | 10 + .../insert-invalid-range/dataset.txt | 1 + .../insert-invalid-range/expected.txt | 6 + .../insert-invalid-range/runbook.yaml | 6 + .../insert-start-equals-end/dataset.txt | 1 + .../insert-start-equals-end/expected.txt | 6 + .../insert-start-equals-end/runbook.yaml | 6 + .../bigann-ux/missing-max-pts/dataset.txt | 1 + .../bigann-ux/missing-max-pts/expected.txt | 4 + .../bigann-ux/missing-max-pts/runbook.yaml | 5 + .../bigann-ux/non-integer-max-pts/dataset.txt | 1 + .../non-integer-max-pts/expected.txt | 4 + .../non-integer-max-pts/runbook.yaml | 6 + .../replace-ids-start-equals-end/dataset.txt | 1 + .../replace-ids-start-equals-end/expected.txt | 6 + .../replace-ids-start-equals-end/runbook.yaml | 12 + .../replace-invalid-ids-range/dataset.txt | 1 + .../replace-invalid-ids-range/expected.txt | 6 + .../replace-invalid-ids-range/runbook.yaml | 12 + .../replace-invalid-tags-range/dataset.txt | 1 + .../replace-invalid-tags-range/expected.txt | 6 + .../replace-invalid-tags-range/runbook.yaml | 12 + .../replace-tags-start-equals-end/dataset.txt | 1 + .../expected.txt | 6 + .../runbook.yaml | 12 + .../runbook-does-not-exist/dataset.txt | 1 + .../runbook-does-not-exist/expected.txt | 4 + .../runbook-not-yaml-map/dataset.txt | 0 .../runbook-not-yaml-map/expected.txt | 1 + .../runbook-not-yaml-map/runbook.yaml | 2 + .../stage-key-not-integer/dataset.txt | 1 + .../stage-key-not-integer/expected.txt | 4 + .../stage-key-not-integer/runbook.yaml | 10 + .../stage-key-not-number/dataset.txt | 1 + .../stage-key-not-number/expected.txt | 4 + .../stage-key-not-number/runbook.yaml | 10 + .../unrecognized-dataset-key/dataset.txt | 1 + .../unrecognized-dataset-key/expected.txt | 4 + .../unrecognized-dataset-key/runbook.yaml | 7 + .../unrecognized-operation/dataset.txt | 1 + .../unrecognized-operation/expected.txt | 5 + .../unrecognized-operation/runbook.yaml | 6 + diskann-tools/Cargo.toml | 2 +- .../src/bin/compute_streaming_groundtruth.rs | 320 ---------- diskann-tools/src/utils/ground_truth.rs | 2 +- 63 files changed, 1673 insertions(+), 326 deletions(-) create mode 100644 diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs create mode 100644 diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs create mode 100644 diskann-benchmark-core/src/streaming/executors/bigann/validate.rs create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml create mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt create mode 100644 diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml delete mode 100644 diskann-tools/src/bin/compute_streaming_groundtruth.rs diff --git a/Cargo.lock b/Cargo.lock index b812bfe79..89e7ddf0e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -691,6 +691,7 @@ dependencies = [ "diskann-vector", "futures-util", "itertools 0.13.0", + "serde_yaml", "tempfile", "thiserror 2.0.17", "tokio", diff --git a/diskann-benchmark-core/Cargo.toml b/diskann-benchmark-core/Cargo.toml index 7cd914500..5b6740ef8 100644 --- a/diskann-benchmark-core/Cargo.toml +++ b/diskann-benchmark-core/Cargo.toml @@ -17,11 +17,13 @@ futures-util = { workspace = true, default-features = false } thiserror.workspace = true tokio.workspace = true +serde_yaml = { version = "0.9.34", optional = true } + [features] default = [] # BigANN Runbook support requires parsing a YAML file. -bigann = ["diskann-utils/bigann"] +bigann = ["dep:serde_yaml"] [lints] workspace = true diff --git a/diskann-benchmark-core/src/lib.rs b/diskann-benchmark-core/src/lib.rs index 9c7f82d1a..337616878 100644 --- a/diskann-benchmark-core/src/lib.rs +++ b/diskann-benchmark-core/src/lib.rs @@ -50,3 +50,45 @@ pub mod tokio; pub mod build; pub mod search; pub mod streaming; + +/// # Notes on Testing +/// +/// Some components in this framework (e.g. runbook parsing), have UX tests to report errors +/// encountered during parsing. The necessary input files and expected outputs are checked +/// in to the repo in the `tests` directory. +/// +/// If error message change, the expected results can generally be regenerated by running +/// the test suite with the environment variable +/// ```text +/// DISKANN_TEST=overwrite +/// ``` +/// set. +#[cfg(all(test, feature = "bigann"))] +mod ux { + const ENV_VAR: &str = "DISKANN_TEST"; + + /// Check if we should overwrite expected files. + pub(crate) fn should_overwrite() -> bool { + match std::env::var(ENV_VAR) { + Ok(v) if v == "overwrite" => true, + Ok(v) => panic!( + "Unknown value for {}: \"{}\". Expected \"overwrite\"", + ENV_VAR, v + ), + Err(std::env::VarError::NotPresent) => false, + Err(std::env::VarError::NotUnicode(_)) => { + panic!("Value for {} is not unicode", ENV_VAR) + } + } + } + + pub(crate) fn help() -> String { + format!("{}=overwrite", ENV_VAR) + } + + pub(crate) fn test_dir() -> std::path::PathBuf { + let manifest: &std::path::Path = env!("CARGO_MANIFEST_DIR").as_ref(); + manifest.join("tests") + // let test_data_path: PathBuf = format!("{}/tests/{}", manifest_dir, TEST_DATA_DIR).into(); + } +} diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs b/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs index 5ccc9bf41..5a7f5c72a 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/mod.rs @@ -5,15 +5,18 @@ //! # BigANN Runbook Support //! -//! Runbook parsing and execution are provided by `diskann-utils`. -//! This module provides benchmark-specific adaptors over those shared types. +//! This module provides an executor for processing BigANN-style runbooks. +//! Please refer to the [`RunBook`] documentation for more details. //! //! Users can also leverage the [`WithData`] adaptor to convert raw index ranges //! into actual data points for a dataset. +mod parsing; +mod runbook; +mod validate; mod withdata; -pub use diskann_utils::streaming::executors::bigann::{ +pub use runbook::{ Args, Delete, FindGroundtruth, Insert, Replace, RunBook, ScanDirectory, Search, Stage, }; pub use withdata::{DataArgs, WithData}; diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs b/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs new file mode 100644 index 000000000..b98689343 --- /dev/null +++ b/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs @@ -0,0 +1,502 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::path::{Path, PathBuf}; + +use anyhow::Context; +use serde_yaml::{Mapping, Value}; + +/// See [`super::RunBook::load`] for documentation of the file format parsed +/// by this function. +pub(super) fn load( + path: &Path, + dataset: &str, + groundtruth: &mut dyn super::FindGroundtruth, +) -> anyhow::Result { + load_mapping(parse_yaml(path)?, path, dataset, groundtruth) +} + +fn load_mapping( + mapping: Mapping, + path: &Path, + dataset: &str, + groundtruth: &mut dyn super::FindGroundtruth, +) -> anyhow::Result { + let dataset_value = match mapping.get(dataset) { + Some(value) => value, + None => return Err(DumpKeys::new(mapping, dataset, path).into()), + }; + + let dataset_mapping: &Mapping = match dataset_value.try_as() { + Ok(mapping) => mapping, + Err(_) => anyhow::bail!( + "dataset \"{}\" exists in file \"{}\", but its associated payload is not a YAML map", + dataset, + path.display(), + ), + }; + + let mut raw = parse_stages(dataset_mapping).with_context(|| { + format!( + "parsing dataset \"{}\" in file \"{}\"", + dataset, + path.display() + ) + })?; + raw.stages.sort_by_key(|s| s.index); + + let context = |index: usize| { + format!( + "precessing stage {} of dataset \"{}\" in file \"{}\"", + index, + dataset, + path.display() + ) + }; + + let stages: anyhow::Result> = raw + .stages + .iter() + .map(|stage| { + let stage = match &stage.operation { + Operation::Search => super::Stage::Search { + groundtruth: groundtruth + .find_groundtruth(stage.index) + .with_context(|| context(stage.index))?, + }, + Operation::Insert(insert) => super::Stage::Insert { + dataset_offsets_and_ids: insert.start..insert.end, + }, + Operation::Replace(replace) => super::Stage::Replace { + dataset_offsets: replace.ids_start..replace.ids_end, + ids: replace.tags_start..replace.tags_end, + }, + Operation::Delete(delete) => super::Stage::Delete { + ids: delete.start..delete.end, + }, + }; + Ok(stage) + }) + .collect(); + + super::RunBook::new(stages?, raw.max_points) +} + +fn parse_yaml(path: &Path) -> anyhow::Result { + let f = std::fs::File::open(path) + .with_context(|| format!("while opening file \"{}\"", path.display()))?; + + Ok(serde_yaml::from_reader(std::io::BufReader::new(f))?) +} + +fn parse_stages(mapping: &Mapping) -> anyhow::Result { + let mut stages = Vec::::new(); + let mut max_points = None; + mapping.iter().try_for_each(|(key, value)| match key { + Value::String(s) => match s.as_str() { + "max_pts" => { + let points: usize = value + .try_as() + .map_err(|kind| anyhow::anyhow!("failed to parse \"max_pts\" as a {}", kind))?; + + max_points = Some(points); + Ok(()) + } + "gt_url" => Ok(()), + _ => anyhow::bail!("Unrecognized runbook key: \"{}\"", s), + }, + Value::Number(stage) => match stage.as_i64() { + Some(stage) => { + stages.push( + handle_stage(stage as usize, value) + .with_context(|| format!("processing stage {}", stage))?, + ); + Ok(()) + } + None => anyhow::bail!("Stage \"{}\" must be an integer", stage), + }, + _ => anyhow::bail!("Unrecognized key of type {}", classify(key),), + })?; + + let max_points = match max_points { + Some(points) => points, + None => anyhow::bail!("key \"max_pts\" not found"), + }; + + Ok(Raw { max_points, stages }) +} + +fn handle_stage(index: usize, value: &Value) -> anyhow::Result { + let mapping: &Mapping = value + .try_as() + .map_err(|_| anyhow::anyhow!("YAML type is not a map"))?; + + let kind: &str = mapping.get_as("operation")?; + let operation = match Kind::try_parse(kind)? { + Kind::Search => Operation::Search, + Kind::Insert => Operation::Insert(mapping.try_into()?), + Kind::Replace => Operation::Replace(mapping.try_into()?), + Kind::Delete => Operation::Delete(mapping.try_into()?), + }; + + Ok(Stage { index, operation }) +} + +struct Raw { + max_points: usize, + stages: Vec, +} + +struct Stage { + index: usize, + operation: Operation, +} + +enum Operation { + Search, + Insert(Insert), + Replace(Replace), + Delete(Delete), +} + +#[derive(Debug)] +struct DumpKeys { + mapping: Mapping, + dataset: String, + file: PathBuf, +} + +impl DumpKeys { + #[inline(never)] + fn new(mapping: Mapping, dataset: &str, file: &Path) -> Self { + Self { + mapping, + dataset: dataset.to_string(), + file: file.into(), + } + } +} + +impl std::fmt::Display for DumpKeys { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "dataset \"{}\" not found in file \"{}\" - possible alternatives: [", + self.dataset, + self.file.display(), + )?; + + let len = self.mapping.len(); + self.mapping.keys().enumerate().try_for_each(|(i, key)| { + let mut write = |key: &dyn std::fmt::Display| { + if i + 1 == len { + write!(f, "{}", key) + } else { + write!(f, "{}, ", key) + } + }; + + match key { + Value::Null => write(&"null"), + Value::Bool(b) => write(b), + Value::Number(number) => write(number), + Value::String(s) => write(s), + Value::Sequence(_) => write(&""), + Value::Mapping(_) => write(&""), + Value::Tagged(_) => write(&""), + } + })?; + write!(f, "]") + } +} + +impl std::error::Error for DumpKeys {} + +trait TryAs<'a, T> { + fn try_as(&'a self) -> Result; +} + +impl<'a> TryAs<'a, usize> for Value { + fn try_as(&'a self) -> Result { + self.as_i64().map(|i| i as usize).ok_or("usize") + } +} + +impl<'a> TryAs<'a, &'a str> for Value { + fn try_as(&'a self) -> Result<&'a str, &'static str> { + self.as_str().ok_or("string") + } +} + +impl<'a> TryAs<'a, &'a Mapping> for Value { + fn try_as(&'a self) -> Result<&'a Mapping, &'static str> { + self.as_mapping().ok_or("map") + } +} + +trait MappingExt { + fn get_as<'a, T>(&'a self, index: &str) -> anyhow::Result + where + Value: TryAs<'a, T>; +} + +impl MappingExt for Mapping { + fn get_as<'a, T>(&'a self, key: &str) -> anyhow::Result + where + Value: TryAs<'a, T>, + { + match self.get(key) { + Some(value) => match value.try_as() { + Ok(v) => Ok(v), + Err(expected) => Err(anyhow::anyhow!( + "key \"{}\" exists but it is not a {}", + key, + expected, + )), + }, + None => Err(anyhow::anyhow!("key \"{}\" not found", key)), + } + } +} + +#[derive(Debug, Clone, Copy)] +enum Kind { + Search, + Insert, + Replace, + Delete, +} + +impl Kind { + fn try_parse(string: &str) -> anyhow::Result { + match string { + "search" => Ok(Kind::Search), + "insert" => Ok(Kind::Insert), + "replace" => Ok(Kind::Replace), + "delete" => Ok(Kind::Delete), + _ => Err(anyhow::anyhow!("unrecognized operation: {}", string)), + } + } +} + +#[derive(Debug)] +struct Replace { + ids_start: usize, + ids_end: usize, + tags_start: usize, + tags_end: usize, +} + +impl TryFrom<&Mapping> for Replace { + type Error = anyhow::Error; + fn try_from(mapping: &Mapping) -> anyhow::Result { + let inner = || -> anyhow::Result { + let this = Self { + ids_start: mapping.get_as("ids_start")?, + ids_end: mapping.get_as("ids_end")?, + tags_start: mapping.get_as("tags_start")?, + tags_end: mapping.get_as("tags_end")?, + }; + if this.ids_start >= this.ids_end { + anyhow::bail!( + "ids_start ({}) must be less than ids_end ({})", + this.ids_start, + this.ids_end + ); + } + if this.tags_start >= this.tags_end { + anyhow::bail!( + "tags_start ({}) must be less than tags_end ({})", + this.tags_start, + this.tags_end + ); + } + Ok(this) + }; + + inner().context("trying to parse an \"replace\"") + } +} + +#[derive(Debug)] +struct Insert { + start: usize, + end: usize, +} + +impl TryFrom<&Mapping> for Insert { + type Error = anyhow::Error; + fn try_from(mapping: &Mapping) -> anyhow::Result { + let inner = || -> anyhow::Result { + let this = Self { + start: mapping.get_as("start")?, + end: mapping.get_as("end")?, + }; + if this.start >= this.end { + anyhow::bail!( + "start ({}) must be less than end ({})", + this.start, + this.end + ); + } + Ok(this) + }; + + inner().context("trying to parse an \"insert\"") + } +} + +#[derive(Debug)] +struct Delete { + start: usize, + end: usize, +} + +impl TryFrom<&Mapping> for Delete { + type Error = anyhow::Error; + fn try_from(mapping: &Mapping) -> anyhow::Result { + let inner = || -> anyhow::Result { + let this = Self { + start: mapping.get_as("start")?, + end: mapping.get_as("end")?, + }; + if this.start >= this.end { + anyhow::bail!( + "start ({}) must be less than end ({})", + this.start, + this.end + ); + } + Ok(this) + }; + + inner().context("trying to parse \"delete\"") + } +} + +fn classify(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "bool", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Sequence(_) => "sequence", + Value::Mapping(_) => "mapping", + Value::Tagged(_) => "tagged", + } +} + +#[cfg(test)] +mod ux_tests { + use super::*; + + use diskann_benchmark_runner::ux as runner_ux; + + const TEST_DATA_DIR: &str = "bigann-ux"; + const RUNBOOK_FILE: &str = "runbook.yaml"; + const DATASET_FILE: &str = "dataset.txt"; + const EXPECTED_FILE: &str = "expected.txt"; + const PATH_PLACEHOLDER: &str = ""; + + fn fixup_paths_and_os_errors(s: &str, test_dir: &Path) -> String { + let native_path = test_dir.display().to_string(); + let forward_slash_path = native_path.replace('\\', "/"); + + const NOT_FOUND_WINDOWS: &str = "The system cannot find the file specified."; + const NOT_FOUND_LINUX: &str = "No such file or directory"; + + s.replace(&native_path, PATH_PLACEHOLDER) + .replace(&forward_slash_path, PATH_PLACEHOLDER) + .replace('\\', "/") + .replace(NOT_FOUND_WINDOWS, NOT_FOUND_LINUX) + } + + struct FailingGroundtruth; + + impl super::super::FindGroundtruth for FailingGroundtruth { + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { + Err(anyhow::anyhow!( + "groundtruth not available for stage {}", + stage + )) + } + } + + fn run_file_test(test_dir: &Path) { + let runbook_path = test_dir.join(RUNBOOK_FILE); + let dataset_path = test_dir.join(DATASET_FILE); + let expected_path = test_dir.join(EXPECTED_FILE); + + let dataset = std::fs::read_to_string(&dataset_path) + .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", dataset_path, e)); + let dataset = dataset.trim(); + + let mut groundtruth = FailingGroundtruth; + let result = load(&runbook_path, dataset, &mut groundtruth); + + let actual_output = match result { + Ok(_) => panic!( + "Test {:?} expected an error but parsing succeeded", + test_dir.file_name().unwrap() + ), + Err(err) => format!("{:?}", err), + }; + + let actual_portable = fixup_paths_and_os_errors(&actual_output, test_dir); + let actual_normalized = runner_ux::strip_backtrace(runner_ux::normalize(actual_portable)); + + if crate::ux::should_overwrite() { + std::fs::write(&expected_path, &actual_normalized) + .unwrap_or_else(|e| panic!("Failed to write {:?}: {}", expected_path, e)); + println!("Overwrote {:?}", expected_path); + } else { + let expected = std::fs::read_to_string(&expected_path) + .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", expected_path, e)); + let expected_normalized = runner_ux::normalize(expected); + + if actual_normalized != expected_normalized { + panic!( + "Test {:?} failed.\n\nExpected:\n---\n{}\n---\n\nActual:\n---\n{}\n---\nIf this is expected, run with {} to update the expected output.", + test_dir.file_name().unwrap(), + expected_normalized, + actual_normalized, + crate::ux::help(), + ); + } + } + } + + fn run_all_file_tests() { + let test_data_path = crate::ux::test_dir().join(TEST_DATA_DIR); + if !test_data_path.exists() { + println!( + "No test_data directory found at {:?}, skipping file-based tests", + test_data_path + ); + return; + } + + let mut found_tests = false; + for entry in std::fs::read_dir(&test_data_path) + .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", test_data_path, e)) + { + let entry = entry.unwrap(); + if entry.file_type().unwrap().is_dir() { + found_tests = true; + println!("Running file-based test: {:?}", entry.file_name()); + run_file_test(&entry.path()); + } + } + + if !found_tests { + panic!("No test directories found in {:?}", test_data_path); + } + } + + #[test] + fn file_based_error_tests() { + run_all_file_tests(); + } +} diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs new file mode 100644 index 000000000..e9d4bada6 --- /dev/null +++ b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs @@ -0,0 +1,597 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + any::Any, + ops::Range, + path::{Path, PathBuf}, +}; + +use anyhow::Context; + +use crate::streaming::{self, Executor, Stream}; + +use super::{parsing, validate}; + +/// An executor for [BigANN-style runbooks](https://github.com/harsha-simhadri/big-ann-benchmarks/tree/main/neurips23/streaming). +#[derive(Debug, Clone)] +pub struct RunBook { + // The individual runbook stages. + stages: Vec, + // The maximum number of active points at any point in the runbook. + max_points: usize, + // The maximum tag referenced in the runbook. + max_tag: Option, +} + +impl RunBook { + /// Loads a [`RunBook`] from a YAML file at `path` for the specified `dataset`. + /// + /// # Groundtruth Resolution + /// + /// When loading a runbook, search stages may require groundtruth files to be present. + /// The index of these stages will be provided to the `groundtruth` argument for resolution. + /// If working with the standard BigANN benchmark format, consider using [`ScanDirectory`] + /// and providing the directory where the BigANN framework downloads the groundtruth files. + /// + /// Note that the implementation here currently does not attempt to download any groundtruth + /// files using the `gt_url` field; it is merely parsed and ignored. + /// + /// # YAML File Format + /// + /// The top-level structure is a mapping from dataset names (strings) to their + /// corresponding runbook definitions: + /// ```yaml + /// dataset_name_1: + /// # runbook definition... + /// dataset_name_2: + /// # runbook definition... + /// ``` + /// The `dataset` parameter specifies which dataset's runbook to load from the file. + /// + /// Each runbook definition has the following format: + /// + /// ```yaml + /// max_pts: # required + /// [gt_url]: # ignored + /// 0: + /// # stage definition ... + /// 1: + /// # stage definition ... + /// ... + /// ``` + /// Entries need not be in order, but stages must be sequentially numbered starting + /// from `0`. Each stage takes one of four forms (described below). + /// + /// ## Search Stage + /// + /// Merely specifies that a search should be performed. The queries must be provided + /// externally. The groundtruth file for this stage is located via the provided + /// [`FindGroundtruth::find_groundtruth`] method, which will be provided with the stage index. + /// + /// ```yaml + /// : + /// operation: "search" + /// ``` + /// + /// ## Insert Stage + /// + /// Insert vectors from the underlying dataset into the index. The vectors to insert + /// are specified by the range `start..end`, which serves as both the offsets of the + /// vectors in the dataset and their external ids. + /// + /// ```yaml + /// : + /// operation: "insert" + /// start: + /// end: + /// ``` + /// + /// ## Replace Stage + /// + /// Replace vectors in the index with vectors from the underlying dataset. Unlike insertions, + /// replace operations distinguish between the dataset offsets (`ids_start..ids_end`) + /// and the external ids (tags) of the vectors to replace (`tags_start..tags_end`). + /// + /// These operations indicate that the vectors in the index tagged by `tags_start..tags_end` should + /// be replaced with the vectors from the dataset at offsets `ids_start..ids_end`. + /// + /// ```yaml + /// : + /// operation: "replace" + /// ids_start: + /// ids_end: + /// tags_start: + /// tags_end: + /// ``` + /// + /// ## Delete Stage + /// + /// Delete vectors from the index with external ids in the range `start..end`. + /// + /// ```yaml + /// : + /// operation: "delete" + /// start: + /// end: + /// ``` + /// + /// # File Validation + /// + /// The loading framework does a limited amount of validation on the YAML file: + /// + /// 1. The specified `dataset` must exist in the top-level mapping. + /// 2. The `max_pts` key must be present and be a valid integer. Its value is verified and if found to + /// be inaccurate, it is updated to the correct value internally. + /// 3. Each stage must be sequentially numbered starting from `0` with no gaps. + /// 4. Each stage must have a valid operation with all required fields present and of the correct type. + /// 5. Groundtruth files must be resolvable via the provided [`FindGroundtruth`] implementation. + pub fn load( + path: &Path, + dataset: &str, + groundtruth: &mut dyn FindGroundtruth, + ) -> anyhow::Result { + parsing::load(path, dataset, groundtruth) + } + + pub(super) fn new(stages: Vec, max_points: usize) -> anyhow::Result { + let mut this = Self { + stages, + max_points, + max_tag: None, + }; + + let mut validator = validate::Validate::new(); + this.run_with(&mut validator, |_| Ok(()))?; + + this.max_points = this.max_points.max(validator.max_active()); + this.max_tag = validator.max_tag(); + + Ok(this) + } + + /// Returns the maximum number of points specified in the runbook. + pub fn max_points(&self) -> usize { + self.max_points + } + + /// Returns the maximum tag specified in the runbook. + pub fn max_tag(&self) -> Option { + self.max_tag + } + + /// Returns the number of stages in the runbook. + pub fn len(&self) -> usize { + self.stages.len() + } + + /// Returns `true` if the runbook contains no stages. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns a reference to the stages in this runbook. + pub fn stages(&self) -> &[Stage] { + &self.stages + } + + /// Executes the runbook by iterating through each stage. + /// + /// When calling this method, the dynamic type of `stream`'s output must be + /// compatible with the dynamic type expected by `collect`. + fn run_with_internal( + &self, + stream: &mut dyn streaming::Stream>, + collect: &mut dyn FnMut(Box) -> anyhow::Result<()>, + ) -> anyhow::Result<()> { + for (i, stage) in self.stages.iter().enumerate() { + #[derive(Clone, Copy)] + struct OnStage(usize, usize); + + impl std::fmt::Display for OnStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "on stage {} of {}", self.0, self.1) + } + } + + let context = OnStage(i, self.len()); + + if stream.needs_maintenance() { + collect(stream.maintain(()).context(context)?).context(context)?; + } + + let output = match stage { + Stage::Search { groundtruth } => { + let args = Search { groundtruth }; + stream.search(args).context(context)? + } + Stage::Insert { + dataset_offsets_and_ids, + } => { + let args = Insert { + offsets: dataset_offsets_and_ids.clone(), + ids: dataset_offsets_and_ids.clone(), + }; + stream.insert(args).context(context)? + } + Stage::Replace { + dataset_offsets, + ids, + } => { + let args = Replace { + offsets: dataset_offsets.clone(), + ids: ids.clone(), + }; + stream.replace(args).context(context)? + } + Stage::Delete { ids } => { + let args = Delete { ids: ids.clone() }; + stream.delete(args).context(context)? + } + }; + + collect(output).context(context)?; + } + + Ok(()) + } +} + +/// An operation in a BigANN runbook. +#[derive(Debug, Clone, PartialEq)] +pub enum Stage { + /// Perform a search operation using the specified groundtruth file. + Search { + /// The resolved path to the groundtruth file for this search stage. + groundtruth: PathBuf, + }, + /// Insert vectors from the dataset into the index. + Insert { + /// The offsets in the dataset, which also serve as the external ids for + /// the inserted vectors. + dataset_offsets_and_ids: Range, + }, + /// Replace vectors in the index with new vectors from the dataset. + Replace { + /// The offsets in the dataset for the replacement vectors. + dataset_offsets: Range, + /// The external ids of the vectors to be replaced. + ids: Range, + }, + /// Delete vectors from the index. + Delete { + /// The external ids of the vectors to delete. + ids: Range, + }, +} + +/// Arguments for a BigANN runbook "search" stage. +#[derive(Debug)] +pub struct Search<'a> { + /// The resolved path to the file containing the groundtruth for this stage. + pub groundtruth: &'a Path, +} + +/// Arguments for a BigANN runbook "insert" stage. +pub struct Insert { + /// The range of offsets in the dataset for vectors to insert. + pub offsets: Range, + /// The external ids to assign to the inserted vectors. + pub ids: Range, +} + +/// Arguments for a BigANN runbook "replace" stage. +pub struct Replace { + /// The range of offsets in the dataset for the replacement vectors. + pub offsets: Range, + /// The external ids of the vectors to replace. + pub ids: Range, +} + +/// Arguments for a BigANN runbook "delete" stage. +pub struct Delete { + /// The range of external ids to delete from the index. + pub ids: Range, +} + +/// The argument type for a [`RunBook`]. +#[derive(Debug, Clone, Copy)] +pub struct Args; + +impl streaming::Arguments for Args { + type Search<'a> = Search<'a>; + type Insert<'a> = Insert; + type Replace<'a> = Replace; + type Delete<'a> = Delete; + type Maintain<'a> = (); +} + +impl streaming::Executor for RunBook { + type Args = Args; + + fn run_with(&mut self, stream: &mut S, mut collect: F) -> anyhow::Result<()> + where + S: Stream, + O: 'static, + F: FnMut(O) -> anyhow::Result<()>, + { + self.run_with_internal(&mut streaming::AnyStream::new(stream), &mut |any| { + let typed = *any + .downcast::() + .expect("the dynamic type should be configured correctly"); + collect(typed) + }) + } +} + +/// A trait for resolving groundtruth files for search stages in a [`RunBook`]. +/// +/// Implementors of this trait provide the logic to locate groundtruth files +/// given a stage index. See [`ScanDirectory`] for a common implementation. +pub trait FindGroundtruth { + /// Resolves the groundtruth file path for the specified `stage` index. + /// + /// This method is only called for "search" stages in a runbook. + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result; +} + +/// A [`FindGroundtruth`] implementation that scans a directory for groundtruth files +/// matching the expected BigANN naming convention: `step{stage}.gt[0-9]*` where +/// `{stage}` substitutes the stage index formatted as a string. +#[derive(Debug)] +pub struct ScanDirectory { + directory: PathBuf, + + // Cached contents of the files in `directory`. + files: Vec, +} + +impl ScanDirectory { + /// Creates a new [`ScanDirectory`] instance for the specified `directory`. + /// + /// # Notes + /// + /// This constructor scans the directory and caches its contents. + /// If the directory contents change after creation, the instance will not + /// reflect those changes. + /// + /// This is meant for the common benchmarking scenario where the benchmarking + /// machine is generally static while benchmarks are executed. + pub fn new(directory: impl Into) -> anyhow::Result { + Self::new_(directory.into()) + } + + fn new_(directory: PathBuf) -> anyhow::Result { + // Read all files in the directory. + let read_dir = std::fs::read_dir(&directory).with_context(|| { + format!( + "while trying to read the contents of {}", + directory.display() + ) + })?; + + let files = read_dir + .filter_map(|entry| { + let entry = entry.ok()?; + let file_type = entry.file_type().ok()?; + if file_type.is_file() { + Some(entry.file_name().to_string_lossy().into()) + } else { + None + } + }) + .collect(); + + Ok(Self { directory, files }) + } +} + +impl FindGroundtruth for ScanDirectory { + /// Finds the groundtruth file for the specified `stage` by scanning the directory. + /// + /// # Errors + /// + /// Returns an error if no matching file is found or if multiple matches exist. + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { + let prefix = format!("step{}.gt", stage); + + enum Matches<'a> { + None, + One(&'a str), + Many(Vec<&'a str>), + } + + impl<'a> Matches<'a> { + fn push(&mut self, file: &'a str) { + *self = match std::mem::replace(self, Self::None) { + Self::None => Self::One(file), + Self::One(first) => Self::Many(vec![first, file]), + Self::Many(mut all) => { + all.push(file); + Self::Many(all) + } + }; + } + } + + let mut matches = Matches::None; + + for file in &self.files { + if file.starts_with(&prefix) { + let suffix = &file[prefix.len()..]; + if suffix.chars().all(|c| c.is_ascii_digit()) { + matches.push(file); + } + } + } + + match matches { + Matches::One(m) => Ok(self.directory.join(m)), + Matches::None => Err(anyhow::anyhow!( + "No groundtruth found for step {} in \"{}\", expected pattern: \"step{}.gt[0-9]*\"", + stage, + self.directory.display(), + stage, + )), + Matches::Many(matches) => Err(anyhow::anyhow!( + "Multiple groundtruth files found for step {} in \"{}\": {:?}", + stage, + self.directory.display(), + matches, + )), + } + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use crate::streaming::Executor; + + //---------// + // Runbook // + //---------// + + fn _assert_is_clone(_x: &T) {} + + struct MockStream { + stages: Vec, + current_stage: usize, + asked_for_maintenance: bool, + } + + impl MockStream { + fn new(stages: Vec) -> Self { + Self { + stages, + current_stage: 0, + asked_for_maintenance: false, + } + } + + fn increment(&mut self) -> usize { + let output = self.current_stage; + self.current_stage += 1; + output + } + + fn current(&self) -> &Stage { + &self.stages[self.current_stage] + } + } + + impl streaming::Stream for MockStream { + type Output = Option; + + fn search(&mut self, args: Search<'_>) -> anyhow::Result> { + if let Stage::Search { groundtruth } = self.current() { + assert_eq!(args.groundtruth, groundtruth.as_path()); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Search stage, instead got {:?}", + self.current() + )) + } + } + + fn insert(&mut self, args: Insert) -> anyhow::Result> { + if let Stage::Insert { + dataset_offsets_and_ids, + } = self.current() + { + assert_eq!(&args.offsets, dataset_offsets_and_ids); + assert_eq!(&args.ids, dataset_offsets_and_ids); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Insert stage, instead got {:?}", + self.current() + )) + } + } + + fn replace(&mut self, args: Replace) -> anyhow::Result> { + if let Stage::Replace { + dataset_offsets, + ids, + } = self.current() + { + assert_eq!(&args.offsets, dataset_offsets); + assert_eq!(&args.ids, ids); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Replace stage, instead got {:?}", + self.current() + )) + } + } + + fn delete(&mut self, args: Delete) -> anyhow::Result> { + if let Stage::Delete { ids } = self.current() { + assert_eq!(&args.ids, ids); + Ok(Some(self.increment())) + } else { + Err(anyhow::anyhow!( + "Expected Delete stage, instead got {:?}", + self.current() + )) + } + } + + fn maintain(&mut self, _args: ()) -> anyhow::Result> { + assert!( + self.asked_for_maintenance, + "Stream was not expected to need maintenance" + ); + self.asked_for_maintenance = false; + Ok(None) + } + + fn needs_maintenance(&mut self) -> bool { + let needs = self.asked_for_maintenance; + self.asked_for_maintenance = true; + needs + } + } + + #[test] + fn test_runbook() { + let stages = vec![ + Stage::Insert { + dataset_offsets_and_ids: 0..100, + }, + Stage::Search { + groundtruth: PathBuf::from("gt0"), + }, + Stage::Replace { + dataset_offsets: 100..200, + ids: 0..100, + }, + Stage::Delete { ids: 50..75 }, + Stage::Search { + groundtruth: PathBuf::from("gt1"), + }, + ]; + + let mut runbook = RunBook::new(stages.clone(), 1000).unwrap(); + assert_eq!(runbook.len(), stages.len()); + assert!(!runbook.is_empty()); + assert_eq!(runbook.max_points(), 1000); + + let mut stream = MockStream::new(stages.clone()); + let outputs = runbook.run(&mut stream).unwrap(); + + let expected_outputs: Vec = (0..stages.len()).collect(); + let non_maintenance: Vec<_> = outputs.iter().filter_map(|o| *o).collect(); + assert_eq!(non_maintenance, expected_outputs); + } +} diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs b/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs new file mode 100644 index 000000000..05ca5933a --- /dev/null +++ b/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs @@ -0,0 +1,272 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::ops::Range; + +use super::Args; +use crate::streaming; + +#[derive(Debug, Clone, Copy, PartialEq)] +struct SimpleRange { + start: usize, + end: usize, +} + +impl SimpleRange { + fn new(start: usize, end: usize) -> Self { + Self { + start, + end: end.max(start), + } + } + + fn len(&self) -> usize { + self.end - self.start + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn overlaps(&self, other: SimpleRange) -> bool { + self.contains(other.start) || other.contains(self.start) + } + + fn remove(self, other: SimpleRange) -> Remaining { + if other.is_empty() { + if self.is_empty() { + return Remaining::None; + } + return Remaining::One(self); + } + + if !self.overlaps(other) { + Remaining::One(self) + } else if other.strictly_contains(self) { + Remaining::None + } else if other.start <= self.start { + Remaining::One(Self::new(other.end, self.end)) + } else if other.end >= self.end { + Remaining::One(Self::new(self.start, other.start)) + } else { + assert!(self.strictly_contains(other)); + + Remaining::Two( + Self::new(self.start, other.start), + Self::new(other.end, self.end), + ) + } + } + + fn contains(&self, i: usize) -> bool { + self.as_range().contains(&i) + } + + fn strictly_contains(&self, other: SimpleRange) -> bool { + other.start >= self.start && other.end <= self.end + } + + fn as_range(&self) -> Range { + self.start..self.end + } +} + +impl std::fmt::Display for SimpleRange { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}..{}", self.start, self.end) + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum Remaining { + None, + One(SimpleRange), + Two(SimpleRange, SimpleRange), +} + +impl From> for SimpleRange { + fn from(range: Range) -> Self { + Self { + start: range.start, + end: range.end, + } + } +} + +pub(super) struct Validate { + active: Vec, + max_active: usize, + currently_active: usize, + max_tag: Option, +} + +impl Validate { + pub(super) fn new() -> Self { + Self { + active: Vec::new(), + max_active: 0, + currently_active: 0, + max_tag: None, + } + } + + /// Return the maximum number of active points seen so far. + pub(super) fn max_active(&self) -> usize { + self.max_active + } + + /// Return the maximum tag seen so far. + pub(super) fn max_tag(&self) -> Option { + self.max_tag + } + + fn update_max(&mut self, tag: usize) { + match self.max_tag.as_mut() { + Some(max) => *max = (*max).max(tag), + None => self.max_tag = Some(tag), + } + } + + fn compact(&mut self) { + if self.active.len() <= 1 { + return; + } + + let mut write = 0; + for read in 1..self.active.len() { + let write_end = self.active[write].end; + let read_start = self.active[read].start; + + if write_end == read_start { + self.active[write].end = self.active[read].end; + } else { + assert!( + read_start > write_end, + "internal invariant has been violated {} > {}", + read_start, + write_end, + ); + + write += 1; + if write != read { + self.active[write] = self.active[read]; + } + } + } + + self.active.truncate(write + 1); + } + + fn find(&self, tags: SimpleRange) -> Slot { + for (i, range) in self.active.iter().enumerate() { + if tags.overlaps(*range) { + return Slot::In(i); + } + if range.start >= tags.end { + return Slot::Before(i); + } + } + Slot::Back + } + + fn insert(&mut self, tags: SimpleRange) -> anyhow::Result<()> { + match self.find(tags) { + Slot::Back => self.active.push(tags), + Slot::Before(i) => self.active.insert(i, tags), + Slot::In(_) => anyhow::bail!("tag range {} is already present for insertion", tags), + } + + self.currently_active += tags.len(); + self.max_active = self.max_active.max(self.currently_active); + if let Some(tag) = tags.end.checked_sub(1) { + self.update_max(tag); + } + + self.compact(); + Ok(()) + } + + fn replace(&mut self, tags: SimpleRange) -> anyhow::Result<()> { + match self.find(tags) { + Slot::Back | Slot::Before(_) => Err(anyhow::anyhow!( + "tag range {} not valid for replacement", + tags + )), + Slot::In(i) => { + let overlaps = self.active[i]; + if !overlaps.strictly_contains(tags) { + Err(anyhow::anyhow!("could not match the entire range {}", tags)) + } else { + Ok(()) + } + } + } + } + + fn delete(&mut self, tags: SimpleRange) -> anyhow::Result<()> { + match self.find(tags) { + Slot::Back | Slot::Before(_) => { + Err(anyhow::anyhow!("tag range {} not valid for deletion", tags)) + } + Slot::In(i) => { + let current = self.active[i]; + if !current.strictly_contains(tags) { + return Err(anyhow::anyhow!( + "could not match the entire range {} for deletion", + tags + )); + } + + match current.remove(tags) { + Remaining::None => { + self.active.remove(i); + } + Remaining::One(remaining) => self.active[i] = remaining, + Remaining::Two(first, last) => { + self.active.splice(i..(i + 1), [first, last]); + } + } + + self.currently_active -= tags.len(); + Ok(()) + } + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum Slot { + Back, + Before(usize), + In(usize), +} + +impl streaming::Stream for Validate { + type Output = (); + + fn search(&mut self, _args: super::Search<'_>) -> anyhow::Result<()> { + Ok(()) + } + + fn insert(&mut self, args: super::Insert) -> anyhow::Result<()> { + self.insert(args.ids.into()) + } + + fn replace(&mut self, args: super::Replace) -> anyhow::Result<()> { + self.replace(args.ids.into()) + } + + fn delete(&mut self, args: super::Delete) -> anyhow::Result<()> { + self.delete(args.ids.into()) + } + + fn maintain(&mut self, _args: ()) -> anyhow::Result<()> { + Ok(()) + } + + fn needs_maintenance(&mut self) -> bool { + false + } +} diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt new file mode 100644 index 000000000..ebc0da280 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Unrecognized key of type sequence \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml new file mode 100644 index 000000000..1c17b82fb --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + [1, 2, 3]: + operation: delete + start: 50 + end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt new file mode 100644 index 000000000..55341d4f6 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/dataset.txt @@ -0,0 +1 @@ +my_dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt new file mode 100644 index 000000000..4718a6dec --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/expected.txt @@ -0,0 +1 @@ +dataset "my_dataset" not found in file "/runbook.yaml" - possible alternatives: [other_dataset, another_dataset] \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml new file mode 100644 index 000000000..30553b966 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-not-found/runbook.yaml @@ -0,0 +1,13 @@ +other_dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + +another_dataset: + max_pts: 200 + 0: + operation: insert + start: 0 + end: 200 diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt new file mode 100644 index 000000000..3eca2590b --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/expected.txt @@ -0,0 +1,5 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: YAML type is not a map \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml new file mode 100644 index 000000000..a10034f9a --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + 1: + - operation: delete + start: 50 + end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt new file mode 100644 index 000000000..1bfa6bbee --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse "delete" + 2: start (80) must be less than end (40) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml new file mode 100644 index 000000000..79b218eb4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/delete-invalid-range/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: delete + start: 80 + end: 40 diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt new file mode 100644 index 000000000..2794d487e --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 0 + 1: trying to parse an "insert" + 2: start (50) must be less than end (25) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml new file mode 100644 index 000000000..e5901a5d4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/insert-invalid-range/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 50 + end: 25 diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt new file mode 100644 index 000000000..7bd551584 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 0 + 1: trying to parse an "insert" + 2: start (50) must be less than end (50) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml new file mode 100644 index 000000000..93dfb43f1 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/insert-start-equals-end/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 50 + end: 50 diff --git a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt new file mode 100644 index 000000000..08b6dac7a --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + key "max_pts" not found \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml new file mode 100644 index 000000000..c95ee86f6 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/missing-max-pts/runbook.yaml @@ -0,0 +1,5 @@ +dataset: + 0: + operation: insert + start: 0 + end: 100 diff --git a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt new file mode 100644 index 000000000..44c1c4328 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + failed to parse "max_pts" as a usize \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml new file mode 100644 index 000000000..0dbab4cdd --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/non-integer-max-pts/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: "not_an_integer" + 0: + operation: insert + start: 0 + end: 100 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt new file mode 100644 index 000000000..b7397a054 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: ids_start (50) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml new file mode 100644 index 000000000..4b07f23ff --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 50 + ids_end: 50 + tags_start: 0 + tags_end: 10 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt new file mode 100644 index 000000000..cb5e83c8c --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: ids_start (100) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml new file mode 100644 index 000000000..a1f70e2b6 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 100 + ids_end: 50 + tags_start: 0 + tags_end: 10 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt new file mode 100644 index 000000000..014298aa0 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: tags_start (80) must be less than tags_end (40) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml new file mode 100644 index 000000000..de5f1426b --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 0 + ids_end: 50 + tags_start: 80 + tags_end: 40 diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt new file mode 100644 index 000000000..ffbbf0d4f --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/expected.txt @@ -0,0 +1,6 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 1 + 1: trying to parse an "replace" + 2: tags_start (25) must be less than tags_end (25) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml new file mode 100644 index 000000000..498be2272 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml @@ -0,0 +1,12 @@ +dataset: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + 1: + operation: replace + ids_start: 0 + ids_end: 50 + tags_start: 25 + tags_end: 25 diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt new file mode 100644 index 000000000..8781fc588 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/runbook-does-not-exist/expected.txt @@ -0,0 +1,4 @@ +while opening file "/runbook.yaml" + +Caused by: + No such file or directory (os error 2) \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/dataset.txt new file mode 100644 index 000000000..e69de29bb diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt new file mode 100644 index 000000000..837a2194a --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/expected.txt @@ -0,0 +1 @@ +invalid type: string "item1 item2", expected a YAML mapping \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml new file mode 100644 index 000000000..de1ea531d --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml @@ -0,0 +1,2 @@ +item1 +item2 \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt new file mode 100644 index 000000000..9143f3e00 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Stage "1.1" must be an integer \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml new file mode 100644 index 000000000..a37c6f0c1 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-integer/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + 1.1: + operation: delete + start: 50 + end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt new file mode 100644 index 000000000..8bafcb6c3 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Unrecognized runbook key: "not_an_integer" \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml new file mode 100644 index 000000000..581e83e11 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/stage-key-not-number/runbook.yaml @@ -0,0 +1,10 @@ +dataset: + max_pts: 10 + 0: + operation: insert + start: 0 + end: 100 + not_an_integer: + operation: delete + start: 50 + end: 60 diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt new file mode 100644 index 000000000..527ba9779 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/expected.txt @@ -0,0 +1,4 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + Unrecognized runbook key: "unknown_key" \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml new file mode 100644 index 000000000..4418781d4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml @@ -0,0 +1,7 @@ +dataset: + max_pts: 10 + unknown_key: 42 + 0: + operation: insert + start: 0 + end: 100 diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt new file mode 100644 index 000000000..122af2cf4 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/dataset.txt @@ -0,0 +1 @@ +dataset \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt new file mode 100644 index 000000000..5c20d5c1c --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/expected.txt @@ -0,0 +1,5 @@ +parsing dataset "dataset" in file "/runbook.yaml" + +Caused by: + 0: processing stage 0 + 1: unrecognized operation: upsert \ No newline at end of file diff --git a/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml new file mode 100644 index 000000000..f1b77d852 --- /dev/null +++ b/diskann-benchmark-core/tests/bigann-ux/unrecognized-operation/runbook.yaml @@ -0,0 +1,6 @@ +dataset: + max_pts: 100 + 0: + operation: upsert + start: 0 + end: 50 diff --git a/diskann-tools/Cargo.toml b/diskann-tools/Cargo.toml index bcac24631..644ee1173 100644 --- a/diskann-tools/Cargo.toml +++ b/diskann-tools/Cargo.toml @@ -14,7 +14,7 @@ byteorder.workspace = true clap = { workspace = true, features = ["derive"] } diskann-providers = { workspace = true, default-features = false } # see `linalg/Cargo.toml` diskann-vector = { workspace = true } -diskann-utils = { workspace = true, features = ["bigann"] } +diskann-utils = { workspace = true } 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 deleted file mode 100644 index 561f310a2..000000000 --- a/diskann-tools/src/bin/compute_streaming_groundtruth.rs +++ /dev/null @@ -1,320 +0,0 @@ -/* - * 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_utils::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_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::streaming::executors::bigann::{FindGroundtruth, RunBook, Stage}; -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 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()); - - // Boolean active-vector mask indexed by base offset. - let mut active: Vec = vec![false; n_base]; - // For each active dataset offset, the external ID that should appear in the groundtruth. - // For plain inserts external_id == dataset_offset; for Replace they diverge. - let mut ext_id: Vec = vec![0u32; n_base]; - // Reverse map: external_id -> dataset_offset, needed to resolve Delete/Replace removals. - let mut ext_to_offset: HashMap = HashMap::new(); - - println!("Using distance function: {:?}", args.distance_function); - - let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); - - for (stage_idx, stage) in runbook.stages().iter().enumerate() { - match stage { - Stage::Insert { - dataset_offsets_and_ids, - } => { - for id in dataset_offsets_and_ids.clone() { - if id < n_base { - active[id] = true; - ext_id[id] = id as u32; - ext_to_offset.insert(id as u32, id); - } - } - tracing::info!( - "Stage {}: insert {}..{} ({} active)", - stage_idx, - dataset_offsets_and_ids.start, - dataset_offsets_and_ids.end, - active.iter().filter(|&&b| b).count(), - ); - } - Stage::Delete { ids } => { - for eid in ids.clone() { - if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { - active[offset] = false; - ext_to_offset.remove(&(eid as u32)); - } - } - tracing::info!( - "Stage {}: delete {}..{} ({} active)", - stage_idx, - ids.start, - ids.end, - active.iter().filter(|&&b| b).count(), - ); - } - Stage::Replace { - dataset_offsets, - ids, - } => { - // Remove old vectors by external ID. - for eid in ids.clone() { - if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { - active[offset] = false; - ext_to_offset.remove(&(eid as u32)); - } - } - // Insert new vectors: they inherit the external IDs from `ids`. - for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) { - if offset < n_base { - active[offset] = true; - ext_id[offset] = eid as u32; - ext_to_offset.insert(eid as u32, offset); - } - } - tracing::info!( - "Stage {}: replace ids {}..{} with offsets {}..{} ({} active)", - stage_idx, - ids.start, - ids.end, - dataset_offsets.start, - dataset_offsets.end, - active.iter().filter(|&&b| b).count(), - ); - } - Stage::Search { - groundtruth: output_path, - } => { - let timer = Instant::now(); - let n_active = active.iter().filter(|&&b| b).count(); - - tracing::info!( - "Stage {}: computing top-{} groundtruth for {} active vectors against {} queries", - stage_idx, recall_at, n_active, n_queries, - ); - - let active_ids: Vec = active - .iter() - .enumerate() - .filter_map(|(i, &on)| if on { Some(i) } else { None }) - .collect(); - - // Compute KNN in parallel over queries. - // Use ext_id[offset] as the neighbor ID so that groundtruth IDs - // match external IDs (which differ from dataset offsets after Replace). - - // Using the global threadpool is fine here - #[allow(clippy::disallowed_methods)] - let results: Vec> = (0..n_queries) - .into_par_iter() - .map(|qi| { - let query = queries.row(qi); - let mut pq = NeighborPriorityQueue::new(recall_at); - for &offset in &active_ids { - let dist = distance_fn.evaluate_similarity(dataset.row(offset), query); - pq.insert(Neighbor { - id: ext_id[offset], - distance: dist, - }); - } - pq - }) - .collect(); - - // Warn about queries that got fewer than K results (active set smaller than K). - let under_k: Vec = results - .iter() - .enumerate() - .filter_map(|(qi, pq)| { - if pq.size() < recall_at { - Some(qi) - } else { - None - } - }) - .collect(); - if !under_k.is_empty() { - tracing::warn!( - "Stage {}: {} / {} queries have fewer than {} results (active set = {}). \ - Query indices: {:?}", - stage_idx, - under_k.len(), - n_queries, - recall_at, - n_active, - under_k, - ); - } - - write_ground_truth::<()>( - &storage, - output_path.to_str().ok_or_else(|| CMDToolError { - details: format!("Non-UTF8 output path: {}", output_path.display()), - })?, - n_queries, - recall_at, - results, - None, - ) - .map_err(|e| CMDToolError { - details: e.to_string(), - })?; - - tracing::info!( - "Stage {}: groundtruth written to {} in {:?}", - stage_idx, - output_path.display(), - timer.elapsed(), - ); - } - } - } - - tracing::info!("Done."); - Ok(()) -} - -#[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 f42535ca8..77c2f663c 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( +fn write_ground_truth( storage_provider: &impl StorageWriteProvider, ground_truth_file: &str, number_of_queries: usize, From 18990ff552ec8879ed443037e81e22bf6593da8f Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Tue, 21 Jul 2026 16:33:38 +0000 Subject: [PATCH 05/10] move changes back into diskann-tools --- Cargo.lock | 5 +- diskann-benchmark-core/src/streaming/api.rs | 334 +++++++++- diskann-tools/Cargo.toml | 1 + .../src/bin/compute_streaming_groundtruth.rs | 320 ++++++++++ diskann-tools/src/utils/ground_truth.rs | 2 +- diskann-utils/Cargo.toml | 6 - diskann-utils/src/lib.rs | 34 - diskann-utils/src/streaming/api.rs | 127 ---- .../src/streaming/executors/bigann/mod.rs | 17 - .../src/streaming/executors/bigann/parsing.rs | 502 --------------- .../src/streaming/executors/bigann/runbook.rs | 597 ------------------ .../streaming/executors/bigann/validate.rs | 272 -------- diskann-utils/src/streaming/executors/mod.rs | 12 - diskann-utils/src/streaming/mod.rs | 14 - .../dataset.txt | 1 - .../expected.txt | 4 - .../runbook.yaml | 10 - .../bigann-ux/dataset-not-found/dataset.txt | 1 - .../bigann-ux/dataset-not-found/expected.txt | 1 - .../bigann-ux/dataset-not-found/runbook.yaml | 13 - .../dataset-value-not-mapping/dataset.txt | 1 - .../dataset-value-not-mapping/expected.txt | 5 - .../dataset-value-not-mapping/runbook.yaml | 10 - .../delete-invalid-range/dataset.txt | 1 - .../delete-invalid-range/expected.txt | 6 - .../delete-invalid-range/runbook.yaml | 10 - .../insert-invalid-range/dataset.txt | 1 - .../insert-invalid-range/expected.txt | 6 - .../insert-invalid-range/runbook.yaml | 6 - .../insert-start-equals-end/dataset.txt | 1 - .../insert-start-equals-end/expected.txt | 6 - .../insert-start-equals-end/runbook.yaml | 6 - .../bigann-ux/missing-max-pts/dataset.txt | 1 - .../bigann-ux/missing-max-pts/expected.txt | 4 - .../bigann-ux/missing-max-pts/runbook.yaml | 5 - .../bigann-ux/non-integer-max-pts/dataset.txt | 1 - .../non-integer-max-pts/expected.txt | 4 - .../non-integer-max-pts/runbook.yaml | 6 - .../replace-ids-start-equals-end/dataset.txt | 1 - .../replace-ids-start-equals-end/expected.txt | 6 - .../replace-ids-start-equals-end/runbook.yaml | 12 - .../replace-invalid-ids-range/dataset.txt | 1 - .../replace-invalid-ids-range/expected.txt | 6 - .../replace-invalid-ids-range/runbook.yaml | 12 - .../replace-invalid-tags-range/dataset.txt | 1 - .../replace-invalid-tags-range/expected.txt | 6 - .../replace-invalid-tags-range/runbook.yaml | 12 - .../replace-tags-start-equals-end/dataset.txt | 1 - .../expected.txt | 6 - .../runbook.yaml | 12 - .../runbook-does-not-exist/dataset.txt | 1 - .../runbook-does-not-exist/expected.txt | 4 - .../runbook-not-yaml-map/dataset.txt | 0 .../runbook-not-yaml-map/expected.txt | 1 - .../runbook-not-yaml-map/runbook.yaml | 2 - .../stage-key-not-integer/dataset.txt | 1 - .../stage-key-not-integer/expected.txt | 4 - .../stage-key-not-integer/runbook.yaml | 10 - .../stage-key-not-number/dataset.txt | 1 - .../stage-key-not-number/expected.txt | 4 - .../stage-key-not-number/runbook.yaml | 10 - .../unrecognized-dataset-key/dataset.txt | 1 - .../unrecognized-dataset-key/expected.txt | 4 - .../unrecognized-dataset-key/runbook.yaml | 7 - .../unrecognized-operation/dataset.txt | 1 - .../unrecognized-operation/expected.txt | 5 - .../unrecognized-operation/runbook.yaml | 6 - 67 files changed, 656 insertions(+), 1835 deletions(-) create mode 100644 diskann-tools/src/bin/compute_streaming_groundtruth.rs delete mode 100644 diskann-utils/src/streaming/api.rs delete mode 100644 diskann-utils/src/streaming/executors/bigann/mod.rs delete mode 100644 diskann-utils/src/streaming/executors/bigann/parsing.rs delete mode 100644 diskann-utils/src/streaming/executors/bigann/runbook.rs delete mode 100644 diskann-utils/src/streaming/executors/bigann/validate.rs delete mode 100644 diskann-utils/src/streaming/executors/mod.rs delete mode 100644 diskann-utils/src/streaming/mod.rs delete mode 100644 diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml delete mode 100644 diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt delete mode 100644 diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt delete mode 100644 diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml diff --git a/Cargo.lock b/Cargo.lock index 89e7ddf0e..7db2dd33c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -918,6 +918,7 @@ dependencies = [ "byteorder", "clap", "diskann", + "diskann-benchmark-core", "diskann-label-filter", "diskann-providers", "diskann-quantization", @@ -943,10 +944,8 @@ dependencies = [ name = "diskann-utils" version = "0.55.0" dependencies = [ - "anyhow", "bytemuck", "cfg-if", - "diskann-benchmark-runner", "diskann-vector", "diskann-wide", "half", @@ -954,8 +953,6 @@ dependencies = [ "rand_distr", "rayon", "rstest", - "serde_yaml", - "tempfile", "thiserror 2.0.17", ] diff --git a/diskann-benchmark-core/src/streaming/api.rs b/diskann-benchmark-core/src/streaming/api.rs index c373709a9..28f24fe2e 100644 --- a/diskann-benchmark-core/src/streaming/api.rs +++ b/diskann-benchmark-core/src/streaming/api.rs @@ -3,4 +3,336 @@ * Licensed under the MIT license. */ -pub use diskann_utils::streaming::{AnyStream, Arguments, Executor, Stream}; +use std::any::Any; + +/// A streaming interface for performing dynamic (streaming) operations on an index. +/// +/// Streams are characterized by five operations: +/// +/// * `search`: This is, after all, the whole reason for building an index. +/// * `insert`: Insert new points into the index that do not already exist. +/// * `replace`: Replace existing points in the index with new data. +/// * `delete`: Remove points from the index. +/// * `maintain`: Perform maintenance operations on the index. Examples may +/// include fully removing deleted points from internal references. +/// +/// This trait is parameterized by an [`Arguments`] proxy trait, which defines the +/// argument types for each of the operations. The motivation here is to allow nesting +/// of [`Stream`] implementations that progressively modify or adapt the arguments +/// for better code reuse. An example of this is +/// [`crate::streaming::executors::bigann::WithData`], which is a stream layer adapting +/// the raw ranges used by [`crate::streaming::executors::bigann::RunBook`] into +/// actual data slices. +/// +/// Runners for [`Stream`]s use the [`Executor`] trait to invoke the stream operations +/// in a structured way. +pub trait Stream +where + A: Arguments, +{ + /// Output type for all operations. The `'static` is to allow results to be + /// aggregated in [`Any`] for type erasure in higher level [`Executor`]s. + type Output: 'static; + + /// Perform a search operation. + fn search(&mut self, args: A::Search<'_>) -> anyhow::Result; + + /// Perform an insert operation. + fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result; + + /// Perform a replace operation. + fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result; + + /// Perform a delete operation. + fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result; + + /// Perform a maintain operation. + fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result; + + /// Indicate whether or not maintenance is needed. [`Executor`] implementations + /// are responsible periodically checking this. + fn needs_maintenance(&mut self) -> bool; +} + +/// Operation arguments to [`Stream`]. +pub trait Arguments: 'static { + /// Argument to [`Stream::search`]. + type Search<'a>; + /// Argument to [`Stream::insert`]. + type Insert<'a>; + /// Argument to [`Stream::replace`]. + type Replace<'a>; + /// Argument to [`Stream::delete`]. + type Delete<'a>; + /// Argument to [`Stream::maintain`]. + type Maintain<'a>; +} + +/// A sequential executor for [`Stream`]s. +/// +/// Implementations invoke the operations of a [`Stream`] in a structured way (which should +/// be reflected in the associated documentation) and aggregate the results. +pub trait Executor { + /// The argument collection type for the underlying [`Stream`]. + type Args: Arguments; + + /// Execute a series of operations on `stream`. As outputs are produced, they will be + /// passed to `collect` for aggregation. + /// + /// Since dynamic execution may be long-running, this allows implementations of `collect` + /// to perform operations like status updates or partial saving of results as they are + /// generated. + /// + /// See also: [`Executor::run`]. + fn run_with(&mut self, stream: &mut S, collect: F) -> anyhow::Result<()> + where + S: Stream, + O: 'static, + F: FnMut(O) -> anyhow::Result<()>; + + /// Execute a series of operations on `stream`. The outputs of each operation will be + /// collected in the retuned `Vec` in-order. + /// + /// See also: [`Executor::run_with`]. + fn run(&mut self, stream: &mut S) -> anyhow::Result> + where + S: Stream, + { + let mut outputs = Vec::new(); + self.run_with(stream, |output| { + outputs.push(output); + Ok(()) + })?; + Ok(outputs) + } +} + +/// A type-erased [`Stream`] implementation that wraps stream (outputs)[`Stream::Output`] in +/// [`Box`]. +/// +/// This is useful as the final layer in a stack of nested [`Stream`]s that allows the top +/// level [`Executor`] to operate without knowledge of the concrete output types. +/// +/// From a performance perspective, this is usually fine since +/// +/// 1. Individual [`Stream`] operations are typically expensive relative to the cost of boxing. +/// 2. Many concrete [`Stream`] implementations will use the same [`Executor`] implementation. +/// Thus, boxing can help reduce code bloat without significant performance impact. +#[derive(Debug)] +pub struct AnyStream<'a, T>(&'a mut T); + +impl<'a, T> AnyStream<'a, T> { + /// Wrap `stream` in an [`AnyStream`]. + pub fn new(stream: &'a mut T) -> Self { + Self(stream) + } +} + +fn boxed(x: T) -> Box +where + T: Any, +{ + Box::new(x) +} + +impl Stream for AnyStream<'_, T> +where + A: Arguments, + T: Stream, +{ + type Output = Box; + + fn search(&mut self, args: A::Search<'_>) -> anyhow::Result { + self.0.search(args).map(boxed) + } + + fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result { + self.0.insert(args).map(boxed) + } + + fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result { + self.0.replace(args).map(boxed) + } + + fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result { + self.0.delete(args).map(boxed) + } + + fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result { + self.0.maintain(args).map(boxed) + } + + fn needs_maintenance(&mut self) -> bool { + self.0.needs_maintenance() + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + // The main thing we're testing here is the implementation of [`Executor::run`] and + // to a lesser extent `AnyStream`. + + struct Search; + struct Insert; + struct Replace; + struct Delete; + struct Maintain; + + #[derive(Debug, PartialEq)] + enum Op { + Search, + Insert, + Replace, + Delete, + Maintain, + } + + struct TestArgs; + + impl Arguments for TestArgs { + type Search<'a> = Search; + type Insert<'a> = Insert; + type Replace<'a> = Replace; + type Delete<'a> = Delete; + type Maintain<'a> = Maintain; + } + + struct TestStream { + needs_maintenance: bool, + } + + impl TestStream { + fn new(needs_maintenance: bool) -> Self { + Self { needs_maintenance } + } + } + + impl Stream for TestStream { + type Output = Op; + + fn search(&mut self, _args: Search) -> anyhow::Result { + Ok(Op::Search) + } + + fn insert(&mut self, _args: Insert) -> anyhow::Result { + Ok(Op::Insert) + } + + fn replace(&mut self, _args: Replace) -> anyhow::Result { + Ok(Op::Replace) + } + + fn delete(&mut self, _args: Delete) -> anyhow::Result { + Ok(Op::Delete) + } + + fn maintain(&mut self, _args: Maintain) -> anyhow::Result { + Ok(Op::Maintain) + } + + fn needs_maintenance(&mut self) -> bool { + self.needs_maintenance + } + } + + struct TestExecutor; + + impl Executor for TestExecutor { + type Args = TestArgs; + + fn run_with(&mut self, stream: &mut S, mut collect: F) -> anyhow::Result<()> + where + S: Stream, + O: 'static, + F: FnMut(O) -> anyhow::Result<()>, + { + collect(stream.search(Search)?)?; + collect(stream.insert(Insert)?)?; + collect(stream.replace(Replace)?)?; + collect(stream.delete(Delete)?)?; + collect(stream.maintain(Maintain)?)?; + Ok(()) + } + } + + #[test] + fn test_executor_run() -> anyhow::Result<()> { + let mut stream = TestStream::new(false); + let mut executor = TestExecutor; + + let outputs = executor.run(&mut stream)?; + + assert_eq!(outputs.len(), 5); + assert!(matches!(outputs[0], Op::Search)); + assert!(matches!(outputs[1], Op::Insert)); + assert!(matches!(outputs[2], Op::Replace)); + assert!(matches!(outputs[3], Op::Delete)); + assert!(matches!(outputs[4], Op::Maintain)); + + Ok(()) + } + + #[test] + fn test_any_stream() { + let mut stream = TestStream::new(false); + let mut any_stream = AnyStream::new(&mut stream); + + assert!( + !any_stream.needs_maintenance(), + "AnyStream should forward `needs_maintenance`" + ); + assert_eq!( + any_stream + .search(Search) + .unwrap() + .downcast_ref::() + .unwrap(), + &Op::Search + ); + assert_eq!( + any_stream + .insert(Insert) + .unwrap() + .downcast_ref::() + .unwrap(), + &Op::Insert + ); + assert_eq!( + any_stream + .replace(Replace) + .unwrap() + .downcast_ref::() + .unwrap(), + &Op::Replace + ); + assert_eq!( + any_stream + .delete(Delete) + .unwrap() + .downcast_ref::() + .unwrap(), + &Op::Delete + ); + assert_eq!( + any_stream + .maintain(Maintain) + .unwrap() + .downcast_ref::() + .unwrap(), + &Op::Maintain + ); + + let mut stream = TestStream::new(true); + let mut any_stream = AnyStream::new(&mut stream); + assert!( + any_stream.needs_maintenance(), + "AnyStream should forward `needs_maintenance`" + ); + } +} 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..0732b6337 --- /dev/null +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -0,0 +1,320 @@ +/* + * 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_benchmark_core::streaming::executors::bigann::{FindGroundtruth, RunBook, Stage}; +use diskann::neighbor::{Neighbor, NeighborPriorityQueue}; +use diskann::utils::VectorRepr; +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_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 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()); + + // Boolean active-vector mask indexed by base offset. + let mut active: Vec = vec![false; n_base]; + // For each active dataset offset, the external ID that should appear in the groundtruth. + // For plain inserts external_id == dataset_offset; for Replace they diverge. + let mut ext_id: Vec = vec![0u32; n_base]; + // Reverse map: external_id -> dataset_offset, needed to resolve Delete/Replace removals. + let mut ext_to_offset: HashMap = HashMap::new(); + + println!("Using distance function: {:?}", args.distance_function); + + let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); + + for (stage_idx, stage) in runbook.stages().iter().enumerate() { + match stage { + Stage::Insert { + dataset_offsets_and_ids, + } => { + for id in dataset_offsets_and_ids.clone() { + if id < n_base { + active[id] = true; + ext_id[id] = id as u32; + ext_to_offset.insert(id as u32, id); + } + } + tracing::info!( + "Stage {}: insert {}..{} ({} active)", + stage_idx, + dataset_offsets_and_ids.start, + dataset_offsets_and_ids.end, + active.iter().filter(|&&b| b).count(), + ); + } + Stage::Delete { ids } => { + for eid in ids.clone() { + if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { + active[offset] = false; + ext_to_offset.remove(&(eid as u32)); + } + } + tracing::info!( + "Stage {}: delete {}..{} ({} active)", + stage_idx, + ids.start, + ids.end, + active.iter().filter(|&&b| b).count(), + ); + } + Stage::Replace { + dataset_offsets, + ids, + } => { + // Remove old vectors by external ID. + for eid in ids.clone() { + if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { + active[offset] = false; + ext_to_offset.remove(&(eid as u32)); + } + } + // Insert new vectors: they inherit the external IDs from `ids`. + for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) { + if offset < n_base { + active[offset] = true; + ext_id[offset] = eid as u32; + ext_to_offset.insert(eid as u32, offset); + } + } + tracing::info!( + "Stage {}: replace ids {}..{} with offsets {}..{} ({} active)", + stage_idx, + ids.start, + ids.end, + dataset_offsets.start, + dataset_offsets.end, + active.iter().filter(|&&b| b).count(), + ); + } + Stage::Search { + groundtruth: output_path, + } => { + let timer = Instant::now(); + let n_active = active.iter().filter(|&&b| b).count(); + + tracing::info!( + "Stage {}: computing top-{} groundtruth for {} active vectors against {} queries", + stage_idx, recall_at, n_active, n_queries, + ); + + let active_ids: Vec = active + .iter() + .enumerate() + .filter_map(|(i, &on)| if on { Some(i) } else { None }) + .collect(); + + // Compute KNN in parallel over queries. + // Use ext_id[offset] as the neighbor ID so that groundtruth IDs + // match external IDs (which differ from dataset offsets after Replace). + + // Using the global threadpool is fine here + #[allow(clippy::disallowed_methods)] + let results: Vec> = (0..n_queries) + .into_par_iter() + .map(|qi| { + let query = queries.row(qi); + let mut pq = NeighborPriorityQueue::new(recall_at); + for &offset in &active_ids { + let dist = distance_fn.evaluate_similarity(dataset.row(offset), query); + pq.insert(Neighbor { + id: ext_id[offset], + distance: dist, + }); + } + pq + }) + .collect(); + + // Warn about queries that got fewer than K results (active set smaller than K). + let under_k: Vec = results + .iter() + .enumerate() + .filter_map(|(qi, pq)| { + if pq.size() < recall_at { + Some(qi) + } else { + None + } + }) + .collect(); + if !under_k.is_empty() { + tracing::warn!( + "Stage {}: {} / {} queries have fewer than {} results (active set = {}). \ + Query indices: {:?}", + stage_idx, + under_k.len(), + n_queries, + recall_at, + n_active, + under_k, + ); + } + + write_ground_truth::<()>( + &storage, + output_path.to_str().ok_or_else(|| CMDToolError { + details: format!("Non-UTF8 output path: {}", output_path.display()), + })?, + n_queries, + recall_at, + results, + None, + ) + .map_err(|e| CMDToolError { + details: e.to_string(), + })?; + + tracing::info!( + "Stage {}: groundtruth written to {} in {:?}", + stage_idx, + output_path.display(), + timer.elapsed(), + ); + } + } + } + + tracing::info!("Done."); + Ok(()) +} + +#[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, diff --git a/diskann-utils/Cargo.toml b/diskann-utils/Cargo.toml index 6d789ae32..ea358601e 100644 --- a/diskann-utils/Cargo.toml +++ b/diskann-utils/Cargo.toml @@ -8,7 +8,6 @@ license.workspace = true edition.workspace = true [dependencies] -anyhow.workspace = true half = { workspace = true, features = ["rand_distr", "num-traits"] } cfg-if.workspace = true rayon = { workspace = true, optional = true } @@ -18,7 +17,6 @@ rand = { workspace = true } rand_distr = { workspace = true } diskann-wide = { workspace = true } bytemuck = { workspace = true, features = ["must_cast"] } -serde_yaml = { version = "0.9.34", optional = true } [lints] workspace = true @@ -34,8 +32,6 @@ workspace = true cfg-if.workspace = true rand.workspace = true rstest.workspace = true -tempfile.workspace = true -diskann-benchmark-runner = { workspace = true, features = ["ux-tools"] } [features] @@ -45,5 +41,3 @@ default = ["rayon"] rayon = ["dep:rayon"] # Enable testing utilities like test_data_root() testing = [] -# Enable BigANN runbook parsing and execution support. -bigann = ["dep:serde_yaml"] diff --git a/diskann-utils/src/lib.rs b/diskann-utils/src/lib.rs index 54d6ceff6..089e0dece 100644 --- a/diskann-utils/src/lib.rs +++ b/diskann-utils/src/lib.rs @@ -17,7 +17,6 @@ pub mod future; pub mod io; pub mod object_pool; pub mod sampling; -pub mod streaming; // Views pub mod strided; @@ -46,36 +45,3 @@ pub fn test_data_directory() -> &'static str { pub fn test_data_root() -> std::path::PathBuf { workspace_root().join(test_data_directory()) } - -/// # Notes on Testing -/// -/// BigANN parsing UX tests use checked-in fixtures under the crate's `tests` directory. -/// Expected outputs can be regenerated by running tests with `DISKANN_TEST=overwrite`. -#[cfg(all(test, feature = "bigann"))] -mod ux { - const ENV_VAR: &str = "DISKANN_TEST"; - - /// Check if we should overwrite expected files. - pub(crate) fn should_overwrite() -> bool { - match std::env::var(ENV_VAR) { - Ok(v) if v == "overwrite" => true, - Ok(v) => panic!( - "Unknown value for {}: \"{}\". Expected \"overwrite\"", - ENV_VAR, v - ), - Err(std::env::VarError::NotPresent) => false, - Err(std::env::VarError::NotUnicode(_)) => { - panic!("Value for {} is not unicode", ENV_VAR) - } - } - } - - pub(crate) fn help() -> String { - format!("{}=overwrite", ENV_VAR) - } - - pub(crate) fn test_dir() -> std::path::PathBuf { - let manifest: &std::path::Path = env!("CARGO_MANIFEST_DIR").as_ref(); - manifest.join("tests") - } -} diff --git a/diskann-utils/src/streaming/api.rs b/diskann-utils/src/streaming/api.rs deleted file mode 100644 index 70251fd92..000000000 --- a/diskann-utils/src/streaming/api.rs +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::any::Any; - -/// A streaming interface for performing dynamic (streaming) operations on an index. -pub trait Stream -where - A: Arguments, -{ - /// Output type for all operations. The `'static` is to allow results to be - /// aggregated in [`Any`] for type erasure in higher level [`Executor`]s. - type Output: 'static; - - /// Perform a search operation. - fn search(&mut self, args: A::Search<'_>) -> anyhow::Result; - - /// Perform an insert operation. - fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result; - - /// Perform a replace operation. - fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result; - - /// Perform a delete operation. - fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result; - - /// Perform a maintain operation. - fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result; - - /// Indicate whether or not maintenance is needed. [`Executor`] implementations - /// are responsible periodically checking this. - fn needs_maintenance(&mut self) -> bool; -} - -/// Operation arguments to [`Stream`]. -pub trait Arguments: 'static { - /// Argument to [`Stream::search`]. - type Search<'a>; - /// Argument to [`Stream::insert`]. - type Insert<'a>; - /// Argument to [`Stream::replace`]. - type Replace<'a>; - /// Argument to [`Stream::delete`]. - type Delete<'a>; - /// Argument to [`Stream::maintain`]. - type Maintain<'a>; -} - -/// A sequential executor for [`Stream`]s. -pub trait Executor { - /// The argument collection type for the underlying [`Stream`]. - type Args: Arguments; - - /// Execute a series of operations on `stream`. As outputs are produced, they will be - /// passed to `collect` for aggregation. - fn run_with(&mut self, stream: &mut S, collect: F) -> anyhow::Result<()> - where - S: Stream, - O: 'static, - F: FnMut(O) -> anyhow::Result<()>; - - /// Execute a series of operations on `stream`. The outputs of each operation will be - /// collected in the returned `Vec` in-order. - fn run(&mut self, stream: &mut S) -> anyhow::Result> - where - S: Stream, - { - let mut outputs = Vec::new(); - self.run_with(stream, |output| { - outputs.push(output); - Ok(()) - })?; - Ok(outputs) - } -} - -/// A type-erased [`Stream`] implementation that wraps stream outputs in [`Box`]. -#[derive(Debug)] -pub struct AnyStream<'a, T>(&'a mut T); - -impl<'a, T> AnyStream<'a, T> { - /// Wrap `stream` in an [`AnyStream`]. - pub fn new(stream: &'a mut T) -> Self { - Self(stream) - } -} - -fn boxed(x: T) -> Box -where - T: Any, -{ - Box::new(x) -} - -impl Stream for AnyStream<'_, T> -where - A: Arguments, - T: Stream, -{ - type Output = Box; - - fn search(&mut self, args: A::Search<'_>) -> anyhow::Result { - self.0.search(args).map(boxed) - } - - fn insert(&mut self, args: A::Insert<'_>) -> anyhow::Result { - self.0.insert(args).map(boxed) - } - - fn replace(&mut self, args: A::Replace<'_>) -> anyhow::Result { - self.0.replace(args).map(boxed) - } - - fn delete(&mut self, args: A::Delete<'_>) -> anyhow::Result { - self.0.delete(args).map(boxed) - } - - fn maintain(&mut self, args: A::Maintain<'_>) -> anyhow::Result { - self.0.maintain(args).map(boxed) - } - - fn needs_maintenance(&mut self) -> bool { - self.0.needs_maintenance() - } -} diff --git a/diskann-utils/src/streaming/executors/bigann/mod.rs b/diskann-utils/src/streaming/executors/bigann/mod.rs deleted file mode 100644 index 191467cb9..000000000 --- a/diskann-utils/src/streaming/executors/bigann/mod.rs +++ /dev/null @@ -1,17 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! # BigANN Runbook Support -//! -//! This module provides an executor for processing BigANN-style runbooks. -//! Please refer to the [`RunBook`] documentation for more details. - -mod parsing; -mod runbook; -mod validate; - -pub use runbook::{ - Args, Delete, FindGroundtruth, Insert, Replace, RunBook, ScanDirectory, Search, Stage, -}; diff --git a/diskann-utils/src/streaming/executors/bigann/parsing.rs b/diskann-utils/src/streaming/executors/bigann/parsing.rs deleted file mode 100644 index b98689343..000000000 --- a/diskann-utils/src/streaming/executors/bigann/parsing.rs +++ /dev/null @@ -1,502 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::path::{Path, PathBuf}; - -use anyhow::Context; -use serde_yaml::{Mapping, Value}; - -/// See [`super::RunBook::load`] for documentation of the file format parsed -/// by this function. -pub(super) fn load( - path: &Path, - dataset: &str, - groundtruth: &mut dyn super::FindGroundtruth, -) -> anyhow::Result { - load_mapping(parse_yaml(path)?, path, dataset, groundtruth) -} - -fn load_mapping( - mapping: Mapping, - path: &Path, - dataset: &str, - groundtruth: &mut dyn super::FindGroundtruth, -) -> anyhow::Result { - let dataset_value = match mapping.get(dataset) { - Some(value) => value, - None => return Err(DumpKeys::new(mapping, dataset, path).into()), - }; - - let dataset_mapping: &Mapping = match dataset_value.try_as() { - Ok(mapping) => mapping, - Err(_) => anyhow::bail!( - "dataset \"{}\" exists in file \"{}\", but its associated payload is not a YAML map", - dataset, - path.display(), - ), - }; - - let mut raw = parse_stages(dataset_mapping).with_context(|| { - format!( - "parsing dataset \"{}\" in file \"{}\"", - dataset, - path.display() - ) - })?; - raw.stages.sort_by_key(|s| s.index); - - let context = |index: usize| { - format!( - "precessing stage {} of dataset \"{}\" in file \"{}\"", - index, - dataset, - path.display() - ) - }; - - let stages: anyhow::Result> = raw - .stages - .iter() - .map(|stage| { - let stage = match &stage.operation { - Operation::Search => super::Stage::Search { - groundtruth: groundtruth - .find_groundtruth(stage.index) - .with_context(|| context(stage.index))?, - }, - Operation::Insert(insert) => super::Stage::Insert { - dataset_offsets_and_ids: insert.start..insert.end, - }, - Operation::Replace(replace) => super::Stage::Replace { - dataset_offsets: replace.ids_start..replace.ids_end, - ids: replace.tags_start..replace.tags_end, - }, - Operation::Delete(delete) => super::Stage::Delete { - ids: delete.start..delete.end, - }, - }; - Ok(stage) - }) - .collect(); - - super::RunBook::new(stages?, raw.max_points) -} - -fn parse_yaml(path: &Path) -> anyhow::Result { - let f = std::fs::File::open(path) - .with_context(|| format!("while opening file \"{}\"", path.display()))?; - - Ok(serde_yaml::from_reader(std::io::BufReader::new(f))?) -} - -fn parse_stages(mapping: &Mapping) -> anyhow::Result { - let mut stages = Vec::::new(); - let mut max_points = None; - mapping.iter().try_for_each(|(key, value)| match key { - Value::String(s) => match s.as_str() { - "max_pts" => { - let points: usize = value - .try_as() - .map_err(|kind| anyhow::anyhow!("failed to parse \"max_pts\" as a {}", kind))?; - - max_points = Some(points); - Ok(()) - } - "gt_url" => Ok(()), - _ => anyhow::bail!("Unrecognized runbook key: \"{}\"", s), - }, - Value::Number(stage) => match stage.as_i64() { - Some(stage) => { - stages.push( - handle_stage(stage as usize, value) - .with_context(|| format!("processing stage {}", stage))?, - ); - Ok(()) - } - None => anyhow::bail!("Stage \"{}\" must be an integer", stage), - }, - _ => anyhow::bail!("Unrecognized key of type {}", classify(key),), - })?; - - let max_points = match max_points { - Some(points) => points, - None => anyhow::bail!("key \"max_pts\" not found"), - }; - - Ok(Raw { max_points, stages }) -} - -fn handle_stage(index: usize, value: &Value) -> anyhow::Result { - let mapping: &Mapping = value - .try_as() - .map_err(|_| anyhow::anyhow!("YAML type is not a map"))?; - - let kind: &str = mapping.get_as("operation")?; - let operation = match Kind::try_parse(kind)? { - Kind::Search => Operation::Search, - Kind::Insert => Operation::Insert(mapping.try_into()?), - Kind::Replace => Operation::Replace(mapping.try_into()?), - Kind::Delete => Operation::Delete(mapping.try_into()?), - }; - - Ok(Stage { index, operation }) -} - -struct Raw { - max_points: usize, - stages: Vec, -} - -struct Stage { - index: usize, - operation: Operation, -} - -enum Operation { - Search, - Insert(Insert), - Replace(Replace), - Delete(Delete), -} - -#[derive(Debug)] -struct DumpKeys { - mapping: Mapping, - dataset: String, - file: PathBuf, -} - -impl DumpKeys { - #[inline(never)] - fn new(mapping: Mapping, dataset: &str, file: &Path) -> Self { - Self { - mapping, - dataset: dataset.to_string(), - file: file.into(), - } - } -} - -impl std::fmt::Display for DumpKeys { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "dataset \"{}\" not found in file \"{}\" - possible alternatives: [", - self.dataset, - self.file.display(), - )?; - - let len = self.mapping.len(); - self.mapping.keys().enumerate().try_for_each(|(i, key)| { - let mut write = |key: &dyn std::fmt::Display| { - if i + 1 == len { - write!(f, "{}", key) - } else { - write!(f, "{}, ", key) - } - }; - - match key { - Value::Null => write(&"null"), - Value::Bool(b) => write(b), - Value::Number(number) => write(number), - Value::String(s) => write(s), - Value::Sequence(_) => write(&""), - Value::Mapping(_) => write(&""), - Value::Tagged(_) => write(&""), - } - })?; - write!(f, "]") - } -} - -impl std::error::Error for DumpKeys {} - -trait TryAs<'a, T> { - fn try_as(&'a self) -> Result; -} - -impl<'a> TryAs<'a, usize> for Value { - fn try_as(&'a self) -> Result { - self.as_i64().map(|i| i as usize).ok_or("usize") - } -} - -impl<'a> TryAs<'a, &'a str> for Value { - fn try_as(&'a self) -> Result<&'a str, &'static str> { - self.as_str().ok_or("string") - } -} - -impl<'a> TryAs<'a, &'a Mapping> for Value { - fn try_as(&'a self) -> Result<&'a Mapping, &'static str> { - self.as_mapping().ok_or("map") - } -} - -trait MappingExt { - fn get_as<'a, T>(&'a self, index: &str) -> anyhow::Result - where - Value: TryAs<'a, T>; -} - -impl MappingExt for Mapping { - fn get_as<'a, T>(&'a self, key: &str) -> anyhow::Result - where - Value: TryAs<'a, T>, - { - match self.get(key) { - Some(value) => match value.try_as() { - Ok(v) => Ok(v), - Err(expected) => Err(anyhow::anyhow!( - "key \"{}\" exists but it is not a {}", - key, - expected, - )), - }, - None => Err(anyhow::anyhow!("key \"{}\" not found", key)), - } - } -} - -#[derive(Debug, Clone, Copy)] -enum Kind { - Search, - Insert, - Replace, - Delete, -} - -impl Kind { - fn try_parse(string: &str) -> anyhow::Result { - match string { - "search" => Ok(Kind::Search), - "insert" => Ok(Kind::Insert), - "replace" => Ok(Kind::Replace), - "delete" => Ok(Kind::Delete), - _ => Err(anyhow::anyhow!("unrecognized operation: {}", string)), - } - } -} - -#[derive(Debug)] -struct Replace { - ids_start: usize, - ids_end: usize, - tags_start: usize, - tags_end: usize, -} - -impl TryFrom<&Mapping> for Replace { - type Error = anyhow::Error; - fn try_from(mapping: &Mapping) -> anyhow::Result { - let inner = || -> anyhow::Result { - let this = Self { - ids_start: mapping.get_as("ids_start")?, - ids_end: mapping.get_as("ids_end")?, - tags_start: mapping.get_as("tags_start")?, - tags_end: mapping.get_as("tags_end")?, - }; - if this.ids_start >= this.ids_end { - anyhow::bail!( - "ids_start ({}) must be less than ids_end ({})", - this.ids_start, - this.ids_end - ); - } - if this.tags_start >= this.tags_end { - anyhow::bail!( - "tags_start ({}) must be less than tags_end ({})", - this.tags_start, - this.tags_end - ); - } - Ok(this) - }; - - inner().context("trying to parse an \"replace\"") - } -} - -#[derive(Debug)] -struct Insert { - start: usize, - end: usize, -} - -impl TryFrom<&Mapping> for Insert { - type Error = anyhow::Error; - fn try_from(mapping: &Mapping) -> anyhow::Result { - let inner = || -> anyhow::Result { - let this = Self { - start: mapping.get_as("start")?, - end: mapping.get_as("end")?, - }; - if this.start >= this.end { - anyhow::bail!( - "start ({}) must be less than end ({})", - this.start, - this.end - ); - } - Ok(this) - }; - - inner().context("trying to parse an \"insert\"") - } -} - -#[derive(Debug)] -struct Delete { - start: usize, - end: usize, -} - -impl TryFrom<&Mapping> for Delete { - type Error = anyhow::Error; - fn try_from(mapping: &Mapping) -> anyhow::Result { - let inner = || -> anyhow::Result { - let this = Self { - start: mapping.get_as("start")?, - end: mapping.get_as("end")?, - }; - if this.start >= this.end { - anyhow::bail!( - "start ({}) must be less than end ({})", - this.start, - this.end - ); - } - Ok(this) - }; - - inner().context("trying to parse \"delete\"") - } -} - -fn classify(value: &Value) -> &'static str { - match value { - Value::Null => "null", - Value::Bool(_) => "bool", - Value::Number(_) => "number", - Value::String(_) => "string", - Value::Sequence(_) => "sequence", - Value::Mapping(_) => "mapping", - Value::Tagged(_) => "tagged", - } -} - -#[cfg(test)] -mod ux_tests { - use super::*; - - use diskann_benchmark_runner::ux as runner_ux; - - const TEST_DATA_DIR: &str = "bigann-ux"; - const RUNBOOK_FILE: &str = "runbook.yaml"; - const DATASET_FILE: &str = "dataset.txt"; - const EXPECTED_FILE: &str = "expected.txt"; - const PATH_PLACEHOLDER: &str = ""; - - fn fixup_paths_and_os_errors(s: &str, test_dir: &Path) -> String { - let native_path = test_dir.display().to_string(); - let forward_slash_path = native_path.replace('\\', "/"); - - const NOT_FOUND_WINDOWS: &str = "The system cannot find the file specified."; - const NOT_FOUND_LINUX: &str = "No such file or directory"; - - s.replace(&native_path, PATH_PLACEHOLDER) - .replace(&forward_slash_path, PATH_PLACEHOLDER) - .replace('\\', "/") - .replace(NOT_FOUND_WINDOWS, NOT_FOUND_LINUX) - } - - struct FailingGroundtruth; - - impl super::super::FindGroundtruth for FailingGroundtruth { - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { - Err(anyhow::anyhow!( - "groundtruth not available for stage {}", - stage - )) - } - } - - fn run_file_test(test_dir: &Path) { - let runbook_path = test_dir.join(RUNBOOK_FILE); - let dataset_path = test_dir.join(DATASET_FILE); - let expected_path = test_dir.join(EXPECTED_FILE); - - let dataset = std::fs::read_to_string(&dataset_path) - .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", dataset_path, e)); - let dataset = dataset.trim(); - - let mut groundtruth = FailingGroundtruth; - let result = load(&runbook_path, dataset, &mut groundtruth); - - let actual_output = match result { - Ok(_) => panic!( - "Test {:?} expected an error but parsing succeeded", - test_dir.file_name().unwrap() - ), - Err(err) => format!("{:?}", err), - }; - - let actual_portable = fixup_paths_and_os_errors(&actual_output, test_dir); - let actual_normalized = runner_ux::strip_backtrace(runner_ux::normalize(actual_portable)); - - if crate::ux::should_overwrite() { - std::fs::write(&expected_path, &actual_normalized) - .unwrap_or_else(|e| panic!("Failed to write {:?}: {}", expected_path, e)); - println!("Overwrote {:?}", expected_path); - } else { - let expected = std::fs::read_to_string(&expected_path) - .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", expected_path, e)); - let expected_normalized = runner_ux::normalize(expected); - - if actual_normalized != expected_normalized { - panic!( - "Test {:?} failed.\n\nExpected:\n---\n{}\n---\n\nActual:\n---\n{}\n---\nIf this is expected, run with {} to update the expected output.", - test_dir.file_name().unwrap(), - expected_normalized, - actual_normalized, - crate::ux::help(), - ); - } - } - } - - fn run_all_file_tests() { - let test_data_path = crate::ux::test_dir().join(TEST_DATA_DIR); - if !test_data_path.exists() { - println!( - "No test_data directory found at {:?}, skipping file-based tests", - test_data_path - ); - return; - } - - let mut found_tests = false; - for entry in std::fs::read_dir(&test_data_path) - .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", test_data_path, e)) - { - let entry = entry.unwrap(); - if entry.file_type().unwrap().is_dir() { - found_tests = true; - println!("Running file-based test: {:?}", entry.file_name()); - run_file_test(&entry.path()); - } - } - - if !found_tests { - panic!("No test directories found in {:?}", test_data_path); - } - } - - #[test] - fn file_based_error_tests() { - run_all_file_tests(); - } -} diff --git a/diskann-utils/src/streaming/executors/bigann/runbook.rs b/diskann-utils/src/streaming/executors/bigann/runbook.rs deleted file mode 100644 index e9d4bada6..000000000 --- a/diskann-utils/src/streaming/executors/bigann/runbook.rs +++ /dev/null @@ -1,597 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::{ - any::Any, - ops::Range, - path::{Path, PathBuf}, -}; - -use anyhow::Context; - -use crate::streaming::{self, Executor, Stream}; - -use super::{parsing, validate}; - -/// An executor for [BigANN-style runbooks](https://github.com/harsha-simhadri/big-ann-benchmarks/tree/main/neurips23/streaming). -#[derive(Debug, Clone)] -pub struct RunBook { - // The individual runbook stages. - stages: Vec, - // The maximum number of active points at any point in the runbook. - max_points: usize, - // The maximum tag referenced in the runbook. - max_tag: Option, -} - -impl RunBook { - /// Loads a [`RunBook`] from a YAML file at `path` for the specified `dataset`. - /// - /// # Groundtruth Resolution - /// - /// When loading a runbook, search stages may require groundtruth files to be present. - /// The index of these stages will be provided to the `groundtruth` argument for resolution. - /// If working with the standard BigANN benchmark format, consider using [`ScanDirectory`] - /// and providing the directory where the BigANN framework downloads the groundtruth files. - /// - /// Note that the implementation here currently does not attempt to download any groundtruth - /// files using the `gt_url` field; it is merely parsed and ignored. - /// - /// # YAML File Format - /// - /// The top-level structure is a mapping from dataset names (strings) to their - /// corresponding runbook definitions: - /// ```yaml - /// dataset_name_1: - /// # runbook definition... - /// dataset_name_2: - /// # runbook definition... - /// ``` - /// The `dataset` parameter specifies which dataset's runbook to load from the file. - /// - /// Each runbook definition has the following format: - /// - /// ```yaml - /// max_pts: # required - /// [gt_url]: # ignored - /// 0: - /// # stage definition ... - /// 1: - /// # stage definition ... - /// ... - /// ``` - /// Entries need not be in order, but stages must be sequentially numbered starting - /// from `0`. Each stage takes one of four forms (described below). - /// - /// ## Search Stage - /// - /// Merely specifies that a search should be performed. The queries must be provided - /// externally. The groundtruth file for this stage is located via the provided - /// [`FindGroundtruth::find_groundtruth`] method, which will be provided with the stage index. - /// - /// ```yaml - /// : - /// operation: "search" - /// ``` - /// - /// ## Insert Stage - /// - /// Insert vectors from the underlying dataset into the index. The vectors to insert - /// are specified by the range `start..end`, which serves as both the offsets of the - /// vectors in the dataset and their external ids. - /// - /// ```yaml - /// : - /// operation: "insert" - /// start: - /// end: - /// ``` - /// - /// ## Replace Stage - /// - /// Replace vectors in the index with vectors from the underlying dataset. Unlike insertions, - /// replace operations distinguish between the dataset offsets (`ids_start..ids_end`) - /// and the external ids (tags) of the vectors to replace (`tags_start..tags_end`). - /// - /// These operations indicate that the vectors in the index tagged by `tags_start..tags_end` should - /// be replaced with the vectors from the dataset at offsets `ids_start..ids_end`. - /// - /// ```yaml - /// : - /// operation: "replace" - /// ids_start: - /// ids_end: - /// tags_start: - /// tags_end: - /// ``` - /// - /// ## Delete Stage - /// - /// Delete vectors from the index with external ids in the range `start..end`. - /// - /// ```yaml - /// : - /// operation: "delete" - /// start: - /// end: - /// ``` - /// - /// # File Validation - /// - /// The loading framework does a limited amount of validation on the YAML file: - /// - /// 1. The specified `dataset` must exist in the top-level mapping. - /// 2. The `max_pts` key must be present and be a valid integer. Its value is verified and if found to - /// be inaccurate, it is updated to the correct value internally. - /// 3. Each stage must be sequentially numbered starting from `0` with no gaps. - /// 4. Each stage must have a valid operation with all required fields present and of the correct type. - /// 5. Groundtruth files must be resolvable via the provided [`FindGroundtruth`] implementation. - pub fn load( - path: &Path, - dataset: &str, - groundtruth: &mut dyn FindGroundtruth, - ) -> anyhow::Result { - parsing::load(path, dataset, groundtruth) - } - - pub(super) fn new(stages: Vec, max_points: usize) -> anyhow::Result { - let mut this = Self { - stages, - max_points, - max_tag: None, - }; - - let mut validator = validate::Validate::new(); - this.run_with(&mut validator, |_| Ok(()))?; - - this.max_points = this.max_points.max(validator.max_active()); - this.max_tag = validator.max_tag(); - - Ok(this) - } - - /// Returns the maximum number of points specified in the runbook. - pub fn max_points(&self) -> usize { - self.max_points - } - - /// Returns the maximum tag specified in the runbook. - pub fn max_tag(&self) -> Option { - self.max_tag - } - - /// Returns the number of stages in the runbook. - pub fn len(&self) -> usize { - self.stages.len() - } - - /// Returns `true` if the runbook contains no stages. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Returns a reference to the stages in this runbook. - pub fn stages(&self) -> &[Stage] { - &self.stages - } - - /// Executes the runbook by iterating through each stage. - /// - /// When calling this method, the dynamic type of `stream`'s output must be - /// compatible with the dynamic type expected by `collect`. - fn run_with_internal( - &self, - stream: &mut dyn streaming::Stream>, - collect: &mut dyn FnMut(Box) -> anyhow::Result<()>, - ) -> anyhow::Result<()> { - for (i, stage) in self.stages.iter().enumerate() { - #[derive(Clone, Copy)] - struct OnStage(usize, usize); - - impl std::fmt::Display for OnStage { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "on stage {} of {}", self.0, self.1) - } - } - - let context = OnStage(i, self.len()); - - if stream.needs_maintenance() { - collect(stream.maintain(()).context(context)?).context(context)?; - } - - let output = match stage { - Stage::Search { groundtruth } => { - let args = Search { groundtruth }; - stream.search(args).context(context)? - } - Stage::Insert { - dataset_offsets_and_ids, - } => { - let args = Insert { - offsets: dataset_offsets_and_ids.clone(), - ids: dataset_offsets_and_ids.clone(), - }; - stream.insert(args).context(context)? - } - Stage::Replace { - dataset_offsets, - ids, - } => { - let args = Replace { - offsets: dataset_offsets.clone(), - ids: ids.clone(), - }; - stream.replace(args).context(context)? - } - Stage::Delete { ids } => { - let args = Delete { ids: ids.clone() }; - stream.delete(args).context(context)? - } - }; - - collect(output).context(context)?; - } - - Ok(()) - } -} - -/// An operation in a BigANN runbook. -#[derive(Debug, Clone, PartialEq)] -pub enum Stage { - /// Perform a search operation using the specified groundtruth file. - Search { - /// The resolved path to the groundtruth file for this search stage. - groundtruth: PathBuf, - }, - /// Insert vectors from the dataset into the index. - Insert { - /// The offsets in the dataset, which also serve as the external ids for - /// the inserted vectors. - dataset_offsets_and_ids: Range, - }, - /// Replace vectors in the index with new vectors from the dataset. - Replace { - /// The offsets in the dataset for the replacement vectors. - dataset_offsets: Range, - /// The external ids of the vectors to be replaced. - ids: Range, - }, - /// Delete vectors from the index. - Delete { - /// The external ids of the vectors to delete. - ids: Range, - }, -} - -/// Arguments for a BigANN runbook "search" stage. -#[derive(Debug)] -pub struct Search<'a> { - /// The resolved path to the file containing the groundtruth for this stage. - pub groundtruth: &'a Path, -} - -/// Arguments for a BigANN runbook "insert" stage. -pub struct Insert { - /// The range of offsets in the dataset for vectors to insert. - pub offsets: Range, - /// The external ids to assign to the inserted vectors. - pub ids: Range, -} - -/// Arguments for a BigANN runbook "replace" stage. -pub struct Replace { - /// The range of offsets in the dataset for the replacement vectors. - pub offsets: Range, - /// The external ids of the vectors to replace. - pub ids: Range, -} - -/// Arguments for a BigANN runbook "delete" stage. -pub struct Delete { - /// The range of external ids to delete from the index. - pub ids: Range, -} - -/// The argument type for a [`RunBook`]. -#[derive(Debug, Clone, Copy)] -pub struct Args; - -impl streaming::Arguments for Args { - type Search<'a> = Search<'a>; - type Insert<'a> = Insert; - type Replace<'a> = Replace; - type Delete<'a> = Delete; - type Maintain<'a> = (); -} - -impl streaming::Executor for RunBook { - type Args = Args; - - fn run_with(&mut self, stream: &mut S, mut collect: F) -> anyhow::Result<()> - where - S: Stream, - O: 'static, - F: FnMut(O) -> anyhow::Result<()>, - { - self.run_with_internal(&mut streaming::AnyStream::new(stream), &mut |any| { - let typed = *any - .downcast::() - .expect("the dynamic type should be configured correctly"); - collect(typed) - }) - } -} - -/// A trait for resolving groundtruth files for search stages in a [`RunBook`]. -/// -/// Implementors of this trait provide the logic to locate groundtruth files -/// given a stage index. See [`ScanDirectory`] for a common implementation. -pub trait FindGroundtruth { - /// Resolves the groundtruth file path for the specified `stage` index. - /// - /// This method is only called for "search" stages in a runbook. - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result; -} - -/// A [`FindGroundtruth`] implementation that scans a directory for groundtruth files -/// matching the expected BigANN naming convention: `step{stage}.gt[0-9]*` where -/// `{stage}` substitutes the stage index formatted as a string. -#[derive(Debug)] -pub struct ScanDirectory { - directory: PathBuf, - - // Cached contents of the files in `directory`. - files: Vec, -} - -impl ScanDirectory { - /// Creates a new [`ScanDirectory`] instance for the specified `directory`. - /// - /// # Notes - /// - /// This constructor scans the directory and caches its contents. - /// If the directory contents change after creation, the instance will not - /// reflect those changes. - /// - /// This is meant for the common benchmarking scenario where the benchmarking - /// machine is generally static while benchmarks are executed. - pub fn new(directory: impl Into) -> anyhow::Result { - Self::new_(directory.into()) - } - - fn new_(directory: PathBuf) -> anyhow::Result { - // Read all files in the directory. - let read_dir = std::fs::read_dir(&directory).with_context(|| { - format!( - "while trying to read the contents of {}", - directory.display() - ) - })?; - - let files = read_dir - .filter_map(|entry| { - let entry = entry.ok()?; - let file_type = entry.file_type().ok()?; - if file_type.is_file() { - Some(entry.file_name().to_string_lossy().into()) - } else { - None - } - }) - .collect(); - - Ok(Self { directory, files }) - } -} - -impl FindGroundtruth for ScanDirectory { - /// Finds the groundtruth file for the specified `stage` by scanning the directory. - /// - /// # Errors - /// - /// Returns an error if no matching file is found or if multiple matches exist. - fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { - let prefix = format!("step{}.gt", stage); - - enum Matches<'a> { - None, - One(&'a str), - Many(Vec<&'a str>), - } - - impl<'a> Matches<'a> { - fn push(&mut self, file: &'a str) { - *self = match std::mem::replace(self, Self::None) { - Self::None => Self::One(file), - Self::One(first) => Self::Many(vec![first, file]), - Self::Many(mut all) => { - all.push(file); - Self::Many(all) - } - }; - } - } - - let mut matches = Matches::None; - - for file in &self.files { - if file.starts_with(&prefix) { - let suffix = &file[prefix.len()..]; - if suffix.chars().all(|c| c.is_ascii_digit()) { - matches.push(file); - } - } - } - - match matches { - Matches::One(m) => Ok(self.directory.join(m)), - Matches::None => Err(anyhow::anyhow!( - "No groundtruth found for step {} in \"{}\", expected pattern: \"step{}.gt[0-9]*\"", - stage, - self.directory.display(), - stage, - )), - Matches::Many(matches) => Err(anyhow::anyhow!( - "Multiple groundtruth files found for step {} in \"{}\": {:?}", - stage, - self.directory.display(), - matches, - )), - } - } -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - use crate::streaming::Executor; - - //---------// - // Runbook // - //---------// - - fn _assert_is_clone(_x: &T) {} - - struct MockStream { - stages: Vec, - current_stage: usize, - asked_for_maintenance: bool, - } - - impl MockStream { - fn new(stages: Vec) -> Self { - Self { - stages, - current_stage: 0, - asked_for_maintenance: false, - } - } - - fn increment(&mut self) -> usize { - let output = self.current_stage; - self.current_stage += 1; - output - } - - fn current(&self) -> &Stage { - &self.stages[self.current_stage] - } - } - - impl streaming::Stream for MockStream { - type Output = Option; - - fn search(&mut self, args: Search<'_>) -> anyhow::Result> { - if let Stage::Search { groundtruth } = self.current() { - assert_eq!(args.groundtruth, groundtruth.as_path()); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Search stage, instead got {:?}", - self.current() - )) - } - } - - fn insert(&mut self, args: Insert) -> anyhow::Result> { - if let Stage::Insert { - dataset_offsets_and_ids, - } = self.current() - { - assert_eq!(&args.offsets, dataset_offsets_and_ids); - assert_eq!(&args.ids, dataset_offsets_and_ids); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Insert stage, instead got {:?}", - self.current() - )) - } - } - - fn replace(&mut self, args: Replace) -> anyhow::Result> { - if let Stage::Replace { - dataset_offsets, - ids, - } = self.current() - { - assert_eq!(&args.offsets, dataset_offsets); - assert_eq!(&args.ids, ids); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Replace stage, instead got {:?}", - self.current() - )) - } - } - - fn delete(&mut self, args: Delete) -> anyhow::Result> { - if let Stage::Delete { ids } = self.current() { - assert_eq!(&args.ids, ids); - Ok(Some(self.increment())) - } else { - Err(anyhow::anyhow!( - "Expected Delete stage, instead got {:?}", - self.current() - )) - } - } - - fn maintain(&mut self, _args: ()) -> anyhow::Result> { - assert!( - self.asked_for_maintenance, - "Stream was not expected to need maintenance" - ); - self.asked_for_maintenance = false; - Ok(None) - } - - fn needs_maintenance(&mut self) -> bool { - let needs = self.asked_for_maintenance; - self.asked_for_maintenance = true; - needs - } - } - - #[test] - fn test_runbook() { - let stages = vec![ - Stage::Insert { - dataset_offsets_and_ids: 0..100, - }, - Stage::Search { - groundtruth: PathBuf::from("gt0"), - }, - Stage::Replace { - dataset_offsets: 100..200, - ids: 0..100, - }, - Stage::Delete { ids: 50..75 }, - Stage::Search { - groundtruth: PathBuf::from("gt1"), - }, - ]; - - let mut runbook = RunBook::new(stages.clone(), 1000).unwrap(); - assert_eq!(runbook.len(), stages.len()); - assert!(!runbook.is_empty()); - assert_eq!(runbook.max_points(), 1000); - - let mut stream = MockStream::new(stages.clone()); - let outputs = runbook.run(&mut stream).unwrap(); - - let expected_outputs: Vec = (0..stages.len()).collect(); - let non_maintenance: Vec<_> = outputs.iter().filter_map(|o| *o).collect(); - assert_eq!(non_maintenance, expected_outputs); - } -} diff --git a/diskann-utils/src/streaming/executors/bigann/validate.rs b/diskann-utils/src/streaming/executors/bigann/validate.rs deleted file mode 100644 index 05ca5933a..000000000 --- a/diskann-utils/src/streaming/executors/bigann/validate.rs +++ /dev/null @@ -1,272 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::ops::Range; - -use super::Args; -use crate::streaming; - -#[derive(Debug, Clone, Copy, PartialEq)] -struct SimpleRange { - start: usize, - end: usize, -} - -impl SimpleRange { - fn new(start: usize, end: usize) -> Self { - Self { - start, - end: end.max(start), - } - } - - fn len(&self) -> usize { - self.end - self.start - } - - fn is_empty(&self) -> bool { - self.len() == 0 - } - - fn overlaps(&self, other: SimpleRange) -> bool { - self.contains(other.start) || other.contains(self.start) - } - - fn remove(self, other: SimpleRange) -> Remaining { - if other.is_empty() { - if self.is_empty() { - return Remaining::None; - } - return Remaining::One(self); - } - - if !self.overlaps(other) { - Remaining::One(self) - } else if other.strictly_contains(self) { - Remaining::None - } else if other.start <= self.start { - Remaining::One(Self::new(other.end, self.end)) - } else if other.end >= self.end { - Remaining::One(Self::new(self.start, other.start)) - } else { - assert!(self.strictly_contains(other)); - - Remaining::Two( - Self::new(self.start, other.start), - Self::new(other.end, self.end), - ) - } - } - - fn contains(&self, i: usize) -> bool { - self.as_range().contains(&i) - } - - fn strictly_contains(&self, other: SimpleRange) -> bool { - other.start >= self.start && other.end <= self.end - } - - fn as_range(&self) -> Range { - self.start..self.end - } -} - -impl std::fmt::Display for SimpleRange { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}..{}", self.start, self.end) - } -} - -#[derive(Debug, Clone, Copy, PartialEq)] -enum Remaining { - None, - One(SimpleRange), - Two(SimpleRange, SimpleRange), -} - -impl From> for SimpleRange { - fn from(range: Range) -> Self { - Self { - start: range.start, - end: range.end, - } - } -} - -pub(super) struct Validate { - active: Vec, - max_active: usize, - currently_active: usize, - max_tag: Option, -} - -impl Validate { - pub(super) fn new() -> Self { - Self { - active: Vec::new(), - max_active: 0, - currently_active: 0, - max_tag: None, - } - } - - /// Return the maximum number of active points seen so far. - pub(super) fn max_active(&self) -> usize { - self.max_active - } - - /// Return the maximum tag seen so far. - pub(super) fn max_tag(&self) -> Option { - self.max_tag - } - - fn update_max(&mut self, tag: usize) { - match self.max_tag.as_mut() { - Some(max) => *max = (*max).max(tag), - None => self.max_tag = Some(tag), - } - } - - fn compact(&mut self) { - if self.active.len() <= 1 { - return; - } - - let mut write = 0; - for read in 1..self.active.len() { - let write_end = self.active[write].end; - let read_start = self.active[read].start; - - if write_end == read_start { - self.active[write].end = self.active[read].end; - } else { - assert!( - read_start > write_end, - "internal invariant has been violated {} > {}", - read_start, - write_end, - ); - - write += 1; - if write != read { - self.active[write] = self.active[read]; - } - } - } - - self.active.truncate(write + 1); - } - - fn find(&self, tags: SimpleRange) -> Slot { - for (i, range) in self.active.iter().enumerate() { - if tags.overlaps(*range) { - return Slot::In(i); - } - if range.start >= tags.end { - return Slot::Before(i); - } - } - Slot::Back - } - - fn insert(&mut self, tags: SimpleRange) -> anyhow::Result<()> { - match self.find(tags) { - Slot::Back => self.active.push(tags), - Slot::Before(i) => self.active.insert(i, tags), - Slot::In(_) => anyhow::bail!("tag range {} is already present for insertion", tags), - } - - self.currently_active += tags.len(); - self.max_active = self.max_active.max(self.currently_active); - if let Some(tag) = tags.end.checked_sub(1) { - self.update_max(tag); - } - - self.compact(); - Ok(()) - } - - fn replace(&mut self, tags: SimpleRange) -> anyhow::Result<()> { - match self.find(tags) { - Slot::Back | Slot::Before(_) => Err(anyhow::anyhow!( - "tag range {} not valid for replacement", - tags - )), - Slot::In(i) => { - let overlaps = self.active[i]; - if !overlaps.strictly_contains(tags) { - Err(anyhow::anyhow!("could not match the entire range {}", tags)) - } else { - Ok(()) - } - } - } - } - - fn delete(&mut self, tags: SimpleRange) -> anyhow::Result<()> { - match self.find(tags) { - Slot::Back | Slot::Before(_) => { - Err(anyhow::anyhow!("tag range {} not valid for deletion", tags)) - } - Slot::In(i) => { - let current = self.active[i]; - if !current.strictly_contains(tags) { - return Err(anyhow::anyhow!( - "could not match the entire range {} for deletion", - tags - )); - } - - match current.remove(tags) { - Remaining::None => { - self.active.remove(i); - } - Remaining::One(remaining) => self.active[i] = remaining, - Remaining::Two(first, last) => { - self.active.splice(i..(i + 1), [first, last]); - } - } - - self.currently_active -= tags.len(); - Ok(()) - } - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq)] -enum Slot { - Back, - Before(usize), - In(usize), -} - -impl streaming::Stream for Validate { - type Output = (); - - fn search(&mut self, _args: super::Search<'_>) -> anyhow::Result<()> { - Ok(()) - } - - fn insert(&mut self, args: super::Insert) -> anyhow::Result<()> { - self.insert(args.ids.into()) - } - - fn replace(&mut self, args: super::Replace) -> anyhow::Result<()> { - self.replace(args.ids.into()) - } - - fn delete(&mut self, args: super::Delete) -> anyhow::Result<()> { - self.delete(args.ids.into()) - } - - fn maintain(&mut self, _args: ()) -> anyhow::Result<()> { - Ok(()) - } - - fn needs_maintenance(&mut self) -> bool { - false - } -} diff --git a/diskann-utils/src/streaming/executors/mod.rs b/diskann-utils/src/streaming/executors/mod.rs deleted file mode 100644 index 18876a5e5..000000000 --- a/diskann-utils/src/streaming/executors/mod.rs +++ /dev/null @@ -1,12 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! # Provided Executors. -//! -//! Built-in implementations of [`Executor`] for common streaming inputs. - -#[cfg(feature = "bigann")] -#[cfg_attr(docsrs, doc(cfg(feature = "bigann")))] -pub mod bigann; diff --git a/diskann-utils/src/streaming/mod.rs b/diskann-utils/src/streaming/mod.rs deleted file mode 100644 index 9bd3b311c..000000000 --- a/diskann-utils/src/streaming/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! # Support for Streaming Operations. -//! -//! Streaming operations are defined by sequences of insertions, deletions, and replacements -//! with the goal to help study how an algorithm performs under dynamic workloads. - -mod api; -pub use api::{AnyStream, Arguments, Executor, Stream}; - -pub mod executors; diff --git a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt deleted file mode 100644 index ebc0da280..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Unrecognized key of type sequence \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml b/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml deleted file mode 100644 index 1c17b82fb..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - [1, 2, 3]: - operation: delete - start: 50 - end: 60 diff --git a/diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt b/diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt deleted file mode 100644 index 55341d4f6..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -my_dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt b/diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt deleted file mode 100644 index 4718a6dec..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt +++ /dev/null @@ -1 +0,0 @@ -dataset "my_dataset" not found in file "/runbook.yaml" - possible alternatives: [other_dataset, another_dataset] \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml b/diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml deleted file mode 100644 index 30553b966..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml +++ /dev/null @@ -1,13 +0,0 @@ -other_dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - -another_dataset: - max_pts: 200 - 0: - operation: insert - start: 0 - end: 200 diff --git a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt deleted file mode 100644 index 3eca2590b..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt +++ /dev/null @@ -1,5 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: YAML type is not a map \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml b/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml deleted file mode 100644 index a10034f9a..000000000 --- a/diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - 1: - - operation: delete - start: 50 - end: 60 diff --git a/diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt b/diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt b/diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt deleted file mode 100644 index 1bfa6bbee..000000000 --- a/diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse "delete" - 2: start (80) must be less than end (40) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml b/diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml deleted file mode 100644 index 79b218eb4..000000000 --- a/diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: delete - start: 80 - end: 40 diff --git a/diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt b/diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt b/diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt deleted file mode 100644 index 2794d487e..000000000 --- a/diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 0 - 1: trying to parse an "insert" - 2: start (50) must be less than end (25) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml b/diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml deleted file mode 100644 index e5901a5d4..000000000 --- a/diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 50 - end: 25 diff --git a/diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt b/diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt b/diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt deleted file mode 100644 index 7bd551584..000000000 --- a/diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 0 - 1: trying to parse an "insert" - 2: start (50) must be less than end (50) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml b/diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml deleted file mode 100644 index 93dfb43f1..000000000 --- a/diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 50 - end: 50 diff --git a/diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt b/diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt b/diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt deleted file mode 100644 index 08b6dac7a..000000000 --- a/diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - key "max_pts" not found \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml b/diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml deleted file mode 100644 index c95ee86f6..000000000 --- a/diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml +++ /dev/null @@ -1,5 +0,0 @@ -dataset: - 0: - operation: insert - start: 0 - end: 100 diff --git a/diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt b/diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt b/diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt deleted file mode 100644 index 44c1c4328..000000000 --- a/diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - failed to parse "max_pts" as a usize \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml b/diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml deleted file mode 100644 index 0dbab4cdd..000000000 --- a/diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: "not_an_integer" - 0: - operation: insert - start: 0 - end: 100 diff --git a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt deleted file mode 100644 index b7397a054..000000000 --- a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: ids_start (50) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml deleted file mode 100644 index 4b07f23ff..000000000 --- a/diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 50 - ids_end: 50 - tags_start: 0 - tags_end: 10 diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt deleted file mode 100644 index cb5e83c8c..000000000 --- a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: ids_start (100) must be less than ids_end (50) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml deleted file mode 100644 index a1f70e2b6..000000000 --- a/diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 100 - ids_end: 50 - tags_start: 0 - tags_end: 10 diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt deleted file mode 100644 index 014298aa0..000000000 --- a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: tags_start (80) must be less than tags_end (40) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml deleted file mode 100644 index de5f1426b..000000000 --- a/diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 0 - ids_end: 50 - tags_start: 80 - tags_end: 40 diff --git a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt deleted file mode 100644 index ffbbf0d4f..000000000 --- a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 1 - 1: trying to parse an "replace" - 2: tags_start (25) must be less than tags_end (25) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml b/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml deleted file mode 100644 index 498be2272..000000000 --- a/diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml +++ /dev/null @@ -1,12 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: insert - start: 0 - end: 100 - 1: - operation: replace - ids_start: 0 - ids_end: 50 - tags_start: 25 - tags_end: 25 diff --git a/diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt b/diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt b/diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt deleted file mode 100644 index 8781fc588..000000000 --- a/diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -while opening file "/runbook.yaml" - -Caused by: - No such file or directory (os error 2) \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt deleted file mode 100644 index e69de29bb..000000000 diff --git a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt deleted file mode 100644 index 837a2194a..000000000 --- a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt +++ /dev/null @@ -1 +0,0 @@ -invalid type: string "item1 item2", expected a YAML mapping \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml b/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml deleted file mode 100644 index de1ea531d..000000000 --- a/diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml +++ /dev/null @@ -1,2 +0,0 @@ -item1 -item2 \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt b/diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt b/diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt deleted file mode 100644 index 9143f3e00..000000000 --- a/diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Stage "1.1" must be an integer \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml b/diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml deleted file mode 100644 index a37c6f0c1..000000000 --- a/diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - 1.1: - operation: delete - start: 50 - end: 60 diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt b/diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt b/diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt deleted file mode 100644 index 8bafcb6c3..000000000 --- a/diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Unrecognized runbook key: "not_an_integer" \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml b/diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml deleted file mode 100644 index 581e83e11..000000000 --- a/diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml +++ /dev/null @@ -1,10 +0,0 @@ -dataset: - max_pts: 10 - 0: - operation: insert - start: 0 - end: 100 - not_an_integer: - operation: delete - start: 50 - end: 60 diff --git a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt deleted file mode 100644 index 527ba9779..000000000 --- a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt +++ /dev/null @@ -1,4 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - Unrecognized runbook key: "unknown_key" \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml b/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml deleted file mode 100644 index 4418781d4..000000000 --- a/diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml +++ /dev/null @@ -1,7 +0,0 @@ -dataset: - max_pts: 10 - unknown_key: 42 - 0: - operation: insert - start: 0 - end: 100 diff --git a/diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt b/diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt deleted file mode 100644 index 122af2cf4..000000000 --- a/diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt +++ /dev/null @@ -1 +0,0 @@ -dataset \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt b/diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt deleted file mode 100644 index 5c20d5c1c..000000000 --- a/diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt +++ /dev/null @@ -1,5 +0,0 @@ -parsing dataset "dataset" in file "/runbook.yaml" - -Caused by: - 0: processing stage 0 - 1: unrecognized operation: upsert \ No newline at end of file diff --git a/diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml b/diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml deleted file mode 100644 index f1b77d852..000000000 --- a/diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml +++ /dev/null @@ -1,6 +0,0 @@ -dataset: - max_pts: 100 - 0: - operation: upsert - start: 0 - end: 50 From 3b5c341a6b207c9c0d76535985beaac61a02f771 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Tue, 21 Jul 2026 16:33:46 +0000 Subject: [PATCH 06/10] fmt --- diskann-tools/src/bin/compute_streaming_groundtruth.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-tools/src/bin/compute_streaming_groundtruth.rs b/diskann-tools/src/bin/compute_streaming_groundtruth.rs index 0732b6337..8f4ffc613 100644 --- a/diskann-tools/src/bin/compute_streaming_groundtruth.rs +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -20,9 +20,9 @@ use std::{ use anyhow::Context; use clap::Parser; -use diskann_benchmark_core::streaming::executors::bigann::{FindGroundtruth, RunBook, Stage}; use diskann::neighbor::{Neighbor, NeighborPriorityQueue}; use diskann::utils::VectorRepr; +use diskann_benchmark_core::streaming::executors::bigann::{FindGroundtruth, RunBook, Stage}; use diskann_providers::storage::{FileStorageProvider, StorageReadProvider}; use diskann_tools::utils::{ init_subscriber, write_ground_truth, CMDResult, CMDToolError, DataType, From 03771f2db94918c2611106bed4c3a7db946bbbd8 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Tue, 21 Jul 2026 16:40:00 +0000 Subject: [PATCH 07/10] add back tests --- .../src/streaming/executors/bigann/parsing.rs | 307 ++++++++++++- .../src/streaming/executors/bigann/runbook.rs | 240 ++++++++++- .../streaming/executors/bigann/validate.rs | 408 +++++++++++++++++- 3 files changed, 948 insertions(+), 7 deletions(-) diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs b/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs index b98689343..3070bd47e 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/parsing.rs @@ -24,11 +24,16 @@ fn load_mapping( dataset: &str, groundtruth: &mut dyn super::FindGroundtruth, ) -> anyhow::Result { + // Multiple datasets can exist in the same YAML file. + // + // Fortunately, all datasets are keyed by their dataset name which allows us to + // quickly find whether or not it exists. let dataset_value = match mapping.get(dataset) { Some(value) => value, None => return Err(DumpKeys::new(mapping, dataset, path).into()), }; + // Try to coerce the dataset value into a `Mapping`. let dataset_mapping: &Mapping = match dataset_value.try_as() { Ok(mapping) => mapping, Err(_) => anyhow::bail!( @@ -47,6 +52,7 @@ fn load_mapping( })?; raw.stages.sort_by_key(|s| s.index); + // Translate from raw ranges into higher level steps. let context = |index: usize| { format!( "precessing stage {} of dataset \"{}\" in file \"{}\"", @@ -388,19 +394,310 @@ fn classify(value: &Value) -> &'static str { } } +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use std::{collections::HashMap, io::Write}; + + use tempfile::NamedTempFile; + + use crate::streaming::executors::bigann::Stage; + + /// A test implementation of [`super::super::FindGroundtruth`] that returns + /// pre-configured paths for each stage index. + struct MockGroundtruth { + paths: HashMap, + } + + impl MockGroundtruth { + fn new(stages: impl IntoIterator) -> Self { + Self { + paths: stages.into_iter().collect(), + } + } + } + + impl super::super::FindGroundtruth for MockGroundtruth { + fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result { + self.paths + .get(&stage) + .cloned() + .ok_or_else(|| anyhow::anyhow!("no groundtruth configured for stage {}", stage)) + } + } + + /// Helper to create a temporary YAML file with the given content. + fn create_yaml_file(content: &str) -> anyhow::Result { + let mut file = NamedTempFile::new()?; + file.write_all(content.as_bytes())?; + file.flush()?; + Ok(file) + } + + #[test] + fn test_load_simple_insert_only_runbook() { + let yaml = r#" +test_dataset: + max_pts: 1000 + 0: + operation: insert + start: 0 + end: 500 +"#; + + let file = create_yaml_file(yaml).unwrap(); + let mut groundtruth = MockGroundtruth::new([]); + + let runbook = load(file.path(), "test_dataset", &mut groundtruth).unwrap(); + + assert_eq!(runbook.max_points(), 1000); + assert_eq!(runbook.len(), 1); + + assert_eq!( + runbook.stages()[0], + Stage::Insert { + dataset_offsets_and_ids: 0..500 + } + ); + } + + #[test] + fn test_load_runbook_with_search_stage() { + let yaml = r#" +my_dataset: + max_pts: 2000 + 0: + operation: insert + start: 0 + end: 1000 + 1: + operation: search +"#; + + let file = create_yaml_file(yaml).unwrap(); + let mut groundtruth = + MockGroundtruth::new([(1, PathBuf::from("/path/to/groundtruth.bin"))]); + + let runbook = load(file.path(), "my_dataset", &mut groundtruth).unwrap(); + + assert_eq!(runbook.max_points(), 2000); + assert_eq!(runbook.len(), 2); + + assert_eq!( + runbook.stages()[1], + Stage::Search { + groundtruth: PathBuf::from("/path/to/groundtruth.bin") + } + ); + } + + #[test] + fn test_load_runbook_with_all_operation_types() { + let yaml = r#" +full_dataset: + max_pts: 5000 + 0: + operation: insert + start: 0 + end: 1000 + 1: + operation: search + 2: + operation: replace + ids_start: 1000 + ids_end: 1500 + tags_start: 0 + tags_end: 500 + 3: + operation: delete + start: 500 + end: 1000 +"#; + + let file = create_yaml_file(yaml).unwrap(); + let mut groundtruth = MockGroundtruth::new([(1, PathBuf::from("/gt/step1.bin"))]); + + let runbook = load(file.path(), "full_dataset", &mut groundtruth).unwrap(); + + assert_eq!(runbook.max_points(), 5000); + assert_eq!(runbook.len(), 4); + + // Check insert + assert_eq!( + runbook.stages()[0], + Stage::Insert { + dataset_offsets_and_ids: 0..1000 + } + ); + + // Check search + assert_eq!( + runbook.stages()[1], + Stage::Search { + groundtruth: PathBuf::from("/gt/step1.bin") + } + ); + + // Check replace + assert_eq!( + runbook.stages()[2], + Stage::Replace { + dataset_offsets: 1000..1500, + ids: 0..500 + } + ); + + // Check delete + assert_eq!(runbook.stages()[3], Stage::Delete { ids: 500..1000 }); + } + + #[test] + fn test_load_stages_out_of_order_are_sorted() { + let yaml = r#" +unordered: + max_pts: 1000 + 2: + operation: delete + start: 500 + end: 600 + 0: + operation: insert + start: 0 + end: 500 + 1: + operation: insert + start: 500 + end: 1000 +"#; + + let file = create_yaml_file(yaml).unwrap(); + let mut groundtruth = MockGroundtruth::new([]); + + let runbook = load(file.path(), "unordered", &mut groundtruth).unwrap(); + + assert_eq!(runbook.len(), 3); + + // Stages should be in order 0, 1, 2 regardless of YAML order + assert_eq!( + runbook.stages()[0], + Stage::Insert { + dataset_offsets_and_ids: 0..500 + } + ); + + assert_eq!( + runbook.stages()[1], + Stage::Insert { + dataset_offsets_and_ids: 500..1000 + } + ); + + assert_eq!(runbook.stages()[2], Stage::Delete { ids: 500..600 }); + } + + #[test] + fn test_load_multiple_datasets_in_file() { + let yaml = r#" +dataset_a: + max_pts: 100 + 0: + operation: insert + start: 0 + end: 100 + +dataset_b: + max_pts: 200 + 0: + operation: insert + start: 0 + end: 200 +"#; + + let file = create_yaml_file(yaml).unwrap(); + + // Load dataset_a + let mut groundtruth_a = MockGroundtruth::new([]); + let runbook_a = load(file.path(), "dataset_a", &mut groundtruth_a).unwrap(); + assert_eq!(runbook_a.max_points(), 100); + + // Load dataset_b + let mut groundtruth_b = MockGroundtruth::new([]); + let runbook_b = load(file.path(), "dataset_b", &mut groundtruth_b).unwrap(); + assert_eq!(runbook_b.max_points(), 200); + } + + #[test] + fn test_load_gt_url_is_ignored() { + let yaml = r#" +with_gt_url: + max_pts: 100 + gt_url: "https://example.com/groundtruth.bin" + 0: + operation: insert + start: 0 + end: 100 +"#; + + let file = create_yaml_file(yaml).unwrap(); + let mut groundtruth = MockGroundtruth::new([]); + + // Should succeed - gt_url is parsed but ignored + let runbook = load(file.path(), "with_gt_url", &mut groundtruth).unwrap(); + assert_eq!(runbook.max_points(), 100); + } +} + #[cfg(test)] mod ux_tests { use super::*; + // Exposed by the `ux-tools` feature of `diskann_benchmark_runner` use diskann_benchmark_runner::ux as runner_ux; + //---------------------------// + // File-Based UX Error Tests // + //---------------------------// + // + // These tests use checked-in YAML files and expected output files to verify error messages. + // This approach makes it easy to: + // 1. Add new test cases (just add a new directory with runbook.yaml, dataset.txt, expected.txt) + // 2. See the full error output for review + // 3. Regenerate expected output when error messages change + // + // ## Directory Structure + // + // Each test case is a directory under `tests/bigann-ux` containing: + // - `runbook.yaml` - The YAML runbook file to parse + // - `dataset.txt` - Contains the dataset name to load (single line) + // - `expected.txt` - The expected error message output + // + // ## Regenerating Expected Results + // + // Run tests with the environment variable: + // ``` + // DISKANN_TEST=overwrite + // ``` + // to regenerate the `expected.txt` files. Use `git diff` to review changes. + const TEST_DATA_DIR: &str = "bigann-ux"; const RUNBOOK_FILE: &str = "runbook.yaml"; const DATASET_FILE: &str = "dataset.txt"; const EXPECTED_FILE: &str = "expected.txt"; const PATH_PLACEHOLDER: &str = ""; + /// Replace the test directory path with a placeholder to make tests portable. + /// This handles both forward and backslash path separators. + /// + /// Additionally: + /// * All backslashes are replaced with forward slashes. + /// * Common OS-specific "file not found" error messages are normalized. fn fixup_paths_and_os_errors(s: &str, test_dir: &Path) -> String { + // Try both the native path and with normalized separators let native_path = test_dir.display().to_string(); let forward_slash_path = native_path.replace('\\', "/"); @@ -409,10 +706,11 @@ mod ux_tests { s.replace(&native_path, PATH_PLACEHOLDER) .replace(&forward_slash_path, PATH_PLACEHOLDER) - .replace('\\', "/") - .replace(NOT_FOUND_WINDOWS, NOT_FOUND_LINUX) + .replace('\\', "/") // Normalize any remaining backslashes + .replace(NOT_FOUND_WINDOWS, NOT_FOUND_LINUX) // Normalize error messages } + /// A groundtruth finder that always fails - used for error path testing. struct FailingGroundtruth; impl super::super::FindGroundtruth for FailingGroundtruth { @@ -424,15 +722,18 @@ mod ux_tests { } } + /// Run a single file-based test case. fn run_file_test(test_dir: &Path) { let runbook_path = test_dir.join(RUNBOOK_FILE); let dataset_path = test_dir.join(DATASET_FILE); let expected_path = test_dir.join(EXPECTED_FILE); + // Read the dataset name let dataset = std::fs::read_to_string(&dataset_path) .unwrap_or_else(|e| panic!("Failed to read {:?}: {}", dataset_path, e)); let dataset = dataset.trim(); + // Try to load the runbook - we expect an error let mut groundtruth = FailingGroundtruth; let result = load(&runbook_path, dataset, &mut groundtruth); @@ -444,6 +745,7 @@ mod ux_tests { Err(err) => format!("{:?}", err), }; + // Replace test directory path with placeholder for portability let actual_portable = fixup_paths_and_os_errors(&actual_output, test_dir); let actual_normalized = runner_ux::strip_backtrace(runner_ux::normalize(actual_portable)); @@ -468,6 +770,7 @@ mod ux_tests { } } + /// Run all file-based tests in the test_data directory. fn run_all_file_tests() { let test_data_path = crate::ux::test_dir().join(TEST_DATA_DIR); if !test_data_path.exists() { diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs index e9d4bada6..072b3b1f6 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs @@ -16,6 +16,9 @@ use crate::streaming::{self, Executor, Stream}; use super::{parsing, validate}; /// An executor for [BigANN-style runbooks](https://github.com/harsha-simhadri/big-ann-benchmarks/tree/main/neurips23/streaming). +/// +/// If using this struct as a [`streaming::Executor`], consider using the +/// [`super::WithData`] adaptor to provide dataset and query matrices. #[derive(Debug, Clone)] pub struct RunBook { // The individual runbook stages. @@ -297,6 +300,8 @@ pub struct Delete { } /// The argument type for a [`RunBook`]. +/// +/// See also: [`super::WithData`], [`super::DataArgs`]. #[derive(Debug, Clone, Copy)] pub struct Args; @@ -374,9 +379,10 @@ impl ScanDirectory { let files = read_dir .filter_map(|entry| { - let entry = entry.ok()?; - let file_type = entry.file_type().ok()?; - if file_type.is_file() { + if let Ok(entry) = entry + && let Ok(file_type) = entry.file_type() + && file_type.is_file() + { Some(entry.file_name().to_string_lossy().into()) } else { None @@ -418,7 +424,7 @@ impl FindGroundtruth for ScanDirectory { let mut matches = Matches::None; - for file in &self.files { + for file in self.files.iter() { if file.starts_with(&prefix) { let suffix = &file[prefix.len()..]; if suffix.chars().all(|c| c.is_ascii_digit()) { @@ -453,6 +459,10 @@ impl FindGroundtruth for ScanDirectory { mod tests { use super::*; + use std::fs::File; + + use tempfile::TempDir; + use crate::streaming::Executor; //---------// @@ -590,8 +600,230 @@ mod tests { let mut stream = MockStream::new(stages.clone()); let outputs = runbook.run(&mut stream).unwrap(); + // Verify that the outputs match the expected stage indices. let expected_outputs: Vec = (0..stages.len()).collect(); let non_maintenance: Vec<_> = outputs.iter().filter_map(|o| *o).collect(); assert_eq!(non_maintenance, expected_outputs); } + + #[test] + fn test_load_runbook_from_yaml() { + use std::io::Write; + + let temp_dir = TempDir::new().unwrap(); + + // Create groundtruth files for search stages (stages 1 and 7) + File::create(temp_dir.path().join("step1.gt100")).unwrap(); + File::create(temp_dir.path().join("step7.gt100")).unwrap(); + + // Create a YAML runbook with multiple insert, replace, and delete stages + let yaml_content = r#" +test_dataset: + max_pts: 100 + gt_url: "http://example.com/groundtruth" + 0: + operation: "insert" + start: 0 + end: 1000 + 1: + operation: "search" + 2: + operation: "insert" + start: 1000 + end: 2000 + 3: + operation: "delete" + start: 200 + end: 400 + 4: + operation: "replace" + ids_start: 2000 + ids_end: 2500 + tags_start: 400 + tags_end: 900 + 5: + operation: "insert" + start: 2500 + end: 3000 + 6: + operation: "delete" + start: 500 + end: 700 + 7: + operation: "search" +"#; + + let yaml_path = temp_dir.path().join("runbook.yaml"); + { + let mut file = File::create(&yaml_path).unwrap(); + file.write_all(yaml_content.as_bytes()).unwrap(); + } + + // Load the runbook + let mut groundtruth_finder = ScanDirectory::new(temp_dir.path()).unwrap(); + let runbook = RunBook::load(&yaml_path, "test_dataset", &mut groundtruth_finder).unwrap(); + _assert_is_clone(&runbook); + + // Verify the runbook was loaded correctly + assert_eq!(runbook.len(), 8); + assert_eq!(runbook.max_points(), 2300); + assert_eq!(runbook.max_tag(), Some(2999)); + + let stages = runbook.stages(); + + // Stage 0: Insert 0..1000 + assert_eq!( + stages[0], + Stage::Insert { + dataset_offsets_and_ids: 0..1000 + } + ); + + // Stage 1: Search + assert!( + matches!(&stages[1], Stage::Search { groundtruth } if groundtruth.file_name().unwrap() == "step1.gt100") + ); + + // Stage 2: Insert 1000..2000 + assert_eq!( + stages[2], + Stage::Insert { + dataset_offsets_and_ids: 1000..2000 + } + ); + + // Stage 3: Delete 200..400 + assert_eq!(stages[3], Stage::Delete { ids: 200..400 }); + + // Stage 4: Replace offsets 2000..2500, ids 400..900 + assert_eq!( + stages[4], + Stage::Replace { + dataset_offsets: 2000..2500, + ids: 400..900 + } + ); + + // Stage 5: Insert 2500..3000 + assert_eq!( + stages[5], + Stage::Insert { + dataset_offsets_and_ids: 2500..3000 + } + ); + + // Stage 6: Delete 500..700 + assert_eq!(stages[6], Stage::Delete { ids: 500..700 }); + + // Stage 7: Search + assert!( + matches!(&stages[7], Stage::Search { groundtruth } if groundtruth.file_name().unwrap() == "step7.gt100") + ); + } + + //---------------------// + // ScanDirectory Tests // + //---------------------// + + #[test] + fn scan_directory_finds_groundtruth_file() { + let temp_dir = TempDir::new().unwrap(); + + // Create a groundtruth file matching the expected pattern + File::create(temp_dir.path().join("step0.gt100")).unwrap(); + + let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); + let result = scanner.find_groundtruth(0).unwrap(); + assert_eq!(result, temp_dir.path().join("step0.gt100")); + } + + #[test] + fn scan_directory_finds_groundtruth_without_suffix_digits() { + let temp_dir = TempDir::new().unwrap(); + + // Create a groundtruth file with no digits after ".gt" + File::create(temp_dir.path().join("step5.gt")).unwrap(); + + let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); + let result = scanner.find_groundtruth(5).unwrap(); + assert_eq!(result, temp_dir.path().join("step5.gt")); + } + + #[test] + fn scan_directory_errors_when_no_groundtruth_found() { + let temp_dir = TempDir::new().unwrap(); + + // Create some files that don't match the pattern + File::create(temp_dir.path().join("other_file.bin")).unwrap(); + File::create(temp_dir.path().join("step0.other")).unwrap(); + + let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); + let err = scanner.find_groundtruth(0).unwrap_err(); + + let msg = err.to_string(); + assert!(msg.contains("No groundtruth found"), "Got: {}", msg); + } + + #[test] + fn scan_directory_errors_when_multiple_groundtruth_files() { + let temp_dir = TempDir::new().unwrap(); + + // Create multiple groundtruth files for the same stage + File::create(temp_dir.path().join("step0.gt100")).unwrap(); + File::create(temp_dir.path().join("step0.gt200")).unwrap(); + File::create(temp_dir.path().join("step0.gt300")).unwrap(); + + let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); + let err = scanner.find_groundtruth(0).unwrap_err(); + + let msg = err.to_string(); + assert!(msg.contains("Multiple groundtruth files"), "Got: {}", msg); + } + + #[test] + fn scan_directory_ignores_non_digit_suffix() { + let temp_dir = TempDir::new().unwrap(); + + // Create files with non-digit suffixes (should be ignored) + File::create(temp_dir.path().join("step0.gtabc")).unwrap(); + File::create(temp_dir.path().join("step0.gt100")).unwrap(); + + let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); + let result = scanner.find_groundtruth(0).unwrap(); + + // Should find only the valid one + assert_eq!(result, temp_dir.path().join("step0.gt100")); + } + + #[test] + fn scan_directory_errors_on_nonexistent_directory() { + let _ = ScanDirectory::new("/nonexistent/path/that/does/not/exist").unwrap_err(); + } + + #[test] + fn scan_directory_handles_different_stage_indices() { + let temp_dir = TempDir::new().unwrap(); + + File::create(temp_dir.path().join("step0.gt")).unwrap(); + File::create(temp_dir.path().join("step5.gt")).unwrap(); + File::create(temp_dir.path().join("step10.gt")).unwrap(); + + let mut scanner = ScanDirectory::new(temp_dir.path()).unwrap(); + + assert_eq!( + scanner.find_groundtruth(0).unwrap(), + temp_dir.path().join("step0.gt") + ); + assert_eq!( + scanner.find_groundtruth(5).unwrap(), + temp_dir.path().join("step5.gt") + ); + assert_eq!( + scanner.find_groundtruth(10).unwrap(), + temp_dir.path().join("step10.gt") + ); + + // Stage 1 doesn't exist + assert!(scanner.find_groundtruth(1).is_err()); + } } diff --git a/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs b/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs index 05ca5933a..80504111a 100644 --- a/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs +++ b/diskann-benchmark-core/src/streaming/executors/bigann/validate.rs @@ -35,13 +35,23 @@ impl SimpleRange { } fn remove(self, other: SimpleRange) -> Remaining { + // If the other range is empty - there's nothing to remove. if other.is_empty() { if self.is_empty() { return Remaining::None; + } else { + return Remaining::One(self); } - return Remaining::One(self); } + // Five cases to consider: + // + // 1. The two ranges are disjoint - there's nothing to remove. + // 2. `other` strictly contains `self - we remove all of `self`. + // 3. `other` overlaps with the just the start of `self` - return the last part of `self`. + // 4. `other` overlaps with the just the end of `self` - return the last part of `self`. + // 5. `other` is inside of `self` - we must return two pieces. + if !self.overlaps(other) { Remaining::One(self) } else if other.strictly_contains(self) { @@ -95,6 +105,7 @@ impl From> for SimpleRange { } } +/// A test `stream` that ensures a Runbook is well formed. pub(super) struct Validate { active: Vec, max_active: usize, @@ -112,6 +123,19 @@ impl Validate { } } + #[cfg(test)] + fn new_test(active: Vec) -> Self { + let currently_active = active.iter().map(|r| r.len()).sum(); + let max_tag = active.iter().filter_map(|r| r.end.checked_sub(1)).max(); + + Self { + active, + max_active: currently_active, + currently_active, + max_tag, + } + } + /// Return the maximum number of active points seen so far. pub(super) fn max_active(&self) -> usize { self.max_active @@ -129,6 +153,10 @@ impl Validate { } } + /// Merge adjacent ranges where one range's end equals another's start. + /// + /// Assumes the ranges in `active` are already disjoint and ordered by start. + /// After compaction, no two consecutive ranges will have `r1.end == r2.start`. fn compact(&mut self) { if self.active.len() <= 1 { return; @@ -140,6 +168,7 @@ impl Validate { let read_start = self.active[read].start; if write_end == read_start { + // Merge: extend the current range to include the next one self.active[write].end = self.active[read].end; } else { assert!( @@ -149,6 +178,7 @@ impl Validate { write_end, ); + // No merge: move to the next write position write += 1; if write != read { self.active[write] = self.active[read]; @@ -178,6 +208,7 @@ impl Validate { Slot::In(_) => anyhow::bail!("tag range {} is already present for insertion", tags), } + // We've successfully added tags, so this update is accurate. self.currently_active += tags.len(); self.max_active = self.max_active.max(self.currently_active); if let Some(tag) = tags.end.checked_sub(1) { @@ -195,6 +226,8 @@ impl Validate { tags )), Slot::In(i) => { + // Due to compaction, if `tags` is truly present in the list of tags, it + // cannot be split across two ranges. let overlaps = self.active[i]; if !overlaps.strictly_contains(tags) { Err(anyhow::anyhow!("could not match the entire range {}", tags)) @@ -229,6 +262,7 @@ impl Validate { } } + // We've successfully remove tags, so this update is accurate. self.currently_active -= tags.len(); Ok(()) } @@ -236,10 +270,18 @@ impl Validate { } } +/// Location where a range would be located in the sorted ranges vector. #[derive(Debug, Clone, Copy, PartialEq)] enum Slot { + /// Range would be placed in the back. It is disjoint from existing ranges. Back, + + /// Range would be before this index. This implies that the range is disjoin from + /// existing ranges. Before(usize), + + /// Range partial overlaps with this index. This should be the first such range that + /// overlaps. In(usize), } @@ -270,3 +312,367 @@ impl streaming::Stream for Validate { false } } + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + fn range(start: usize, end: usize) -> SimpleRange { + SimpleRange::new(start, end) + } + + #[test] + fn test_range() { + let r = range(0, 1); + assert_eq!(r.start, 0); + assert_eq!(r.end, 1); + assert_eq!(r.len(), 1); + + let r = range(10, 5); + assert_eq!(r.start, 10); + assert_eq!(r.end, 10); + assert_eq!(r.len(), 0); + + let r = range(5, 100); + assert_eq!(r.start, 5); + assert_eq!(r.end, 100); + assert_eq!(r.len(), 95); + } + + #[test] + fn simple_range() { + // Disjoint ranges. + { + let a = range(0, 10); + let b = range(20, 30); + + assert!(!a.overlaps(b), "a = {a}, b = {b}"); + assert!(!b.overlaps(a), "a = {a}, b = {b}"); + + assert!(!a.strictly_contains(b)); + assert!(!b.strictly_contains(a)); + + assert_eq!(a.remove(b), Remaining::One(a)); + assert_eq!(b.remove(a), Remaining::One(b)); + } + + // Disjoint ranges - b + { + let a = range(0, 10); + let b = range(10, 30); + + assert!(!a.overlaps(b), "a = {a}, b = {b}"); + assert!(!b.overlaps(a), "a = {a}, b = {b}"); + + assert!(!a.strictly_contains(b)); + assert!(!b.strictly_contains(a)); + + assert_eq!(a.remove(b), Remaining::One(a)); + assert_eq!(b.remove(a), Remaining::One(b)); + } + + // Partial overlap. + { + let a = range(10, 20); + let b = range(5, 15); + + assert!(a.overlaps(b), "a = {a}, b = {b}"); + assert!(b.overlaps(a), "a = {a}, b = {b}"); + + assert!(!a.strictly_contains(b)); + assert!(!b.strictly_contains(a)); + + assert_eq!(a.remove(b), Remaining::One(range(15, 20))); + assert_eq!(b.remove(a), Remaining::One(range(5, 10))); + } + + // Total overlap. + { + let a = range(10, 20); + let b = range(10, 20); + + assert!(a.overlaps(b), "a = {a}, b = {b}"); + assert!(a.strictly_contains(b)); + assert_eq!(a.remove(b), Remaining::None); + } + + // Totally inside. + { + let a = range(10, 20); + let b = range(12, 18); + + assert!(a.overlaps(b), "a = {a}, b = {b}"); + + assert!(a.strictly_contains(b)); + assert!(!b.strictly_contains(a)); + + assert_eq!(a.remove(b), Remaining::Two(range(10, 12), range(18, 20))); + assert_eq!(b.remove(a), Remaining::None); + } + + // Totally inside - touching left. + { + let a = range(10, 20); + let b = range(10, 15); + + assert!(a.overlaps(b), "a = {a}, b = {b}"); + + assert!(a.strictly_contains(b)); + assert!(!b.strictly_contains(a)); + + assert_eq!(a.remove(b), Remaining::One(range(15, 20))); + assert_eq!(b.remove(a), Remaining::None); + } + + // Totally inside - touching right. + { + let a = range(10, 20); + let b = range(15, 20); + + assert!(a.overlaps(b), "a = {a}, b = {b}"); + + assert!(a.strictly_contains(b)); + assert!(!b.strictly_contains(a)); + + assert_eq!(a.remove(b), Remaining::One(range(10, 15))); + assert_eq!(b.remove(a), Remaining::None); + } + + // Empty ranges. + { + let a = range(10, 10); + assert_eq!(a.remove(a), Remaining::None); + + let a = range(10, 20); + for j in [5, 10, 15, 20, 25] { + let b = range(j, j); + assert_eq!(a.remove(b), Remaining::One(a)); + } + } + } + + #[test] + fn test_compact() { + // Empty + { + let mut v = Validate::new_test(vec![]); + v.compact(); + assert_eq!(&v.active, &[]); + } + + // One + { + let start = vec![range(10, 20)]; + let mut v = Validate::new_test(start.clone()); + v.compact(); + assert_eq!(v.active, start); + } + + // No change. + { + let start = vec![range(10, 20), range(30, 40), range(50, 60)]; + + let mut v = Validate::new_test(start.clone()); + v.compact(); + + assert_eq!(v.active, start); + } + + // Merge in multiple locations with empty elements. + { + let active = vec![ + range(0, 5), // + Merged + range(5, 5), // | + range(5, 5), // | + range(5, 10), // | + range(20, 30), + range(35, 40), // + Merged + range(40, 41), // | + range(41, 41), // | + range(41, 50), // | + range(55, 60), + range(70, 70), // + Merged + range(70, 75), // | + range(75, 100), // | + ]; + + let mut v = Validate::new_test(active); + v.compact(); + + let expected = [ + range(0, 10), + range(20, 30), + range(35, 50), + range(55, 60), + range(70, 100), + ]; + + assert_eq!(&*v.active, &expected); + } + } + + #[test] + fn test_find() { + // Empty + { + let v = Validate::new_test(vec![]); + assert_eq!(v.find(range(0, 10)), Slot::Back); + } + + // One + { + let v = Validate::new_test(vec![range(10, 20)]); + + assert_eq!(v.find(range(0, 5)), Slot::Before(0)); + assert_eq!(v.find(range(5, 10)), Slot::Before(0)); + assert_eq!(v.find(range(8, 12)), Slot::In(0)); + assert_eq!(v.find(range(10, 15)), Slot::In(0)); + assert_eq!(v.find(range(15, 20)), Slot::In(0)); + assert_eq!(v.find(range(18, 22)), Slot::In(0)); + assert_eq!(v.find(range(20, 25)), Slot::Back); + } + + // Multiple + { + let v = Validate::new_test(vec![range(10, 20), range(30, 40), range(50, 60)]); + + assert_eq!(v.find(range(0, 5)), Slot::Before(0)); + assert_eq!(v.find(range(5, 10)), Slot::Before(0)); + assert_eq!(v.find(range(15, 25)), Slot::In(0)); + + assert_eq!(v.find(range(20, 25)), Slot::Before(1)); + assert_eq!(v.find(range(28, 35)), Slot::In(1)); + assert_eq!(v.find(range(35, 45)), Slot::In(1)); + + assert_eq!(v.find(range(45, 50)), Slot::Before(2)); + assert_eq!(v.find(range(55, 65)), Slot::In(2)); + assert_eq!(v.find(range(65, 70)), Slot::Back); + } + } + + #[test] + fn test_insert() { + let mut v = Validate::new_test(vec![]); + + v.insert(range(10, 20)).unwrap(); + assert_eq!(&*v.active, &[range(10, 20)]); + assert_eq!(v.max_active(), 10); + assert_eq!(v.max_tag(), Some(19)); + + v.insert(range(30, 40)).unwrap(); + assert_eq!(&*v.active, &[range(10, 20), range(30, 40)]); + assert_eq!(v.max_active(), 20); + assert_eq!(v.max_tag(), Some(39)); + + v.insert(range(20, 30)).unwrap(); + assert_eq!(&*v.active, &[range(10, 40)]); + assert!(v.max_active() == 30); + assert_eq!(v.max_tag(), Some(39)); + + let result = v.insert(range(15, 25)); + assert!(result.is_err()); + + v.insert(range(0, 5)).unwrap(); + assert_eq!(&*v.active, &[range(0, 5), range(10, 40)]); + assert_eq!(v.max_active(), 35); + assert_eq!(v.max_tag(), Some(39)); + + v.insert(range(7, 8)).unwrap(); + assert_eq!(&*v.active, &[range(0, 5), range(7, 8), range(10, 40)]); + assert_eq!(v.max_active(), 36); + assert_eq!(v.max_tag(), Some(39)); + + v.insert(range(4, 8)).unwrap_err(); + v.insert(range(50, 60)).unwrap(); + assert_eq!( + &*v.active, + &[range(0, 5), range(7, 8), range(10, 40), range(50, 60)] + ); + assert_eq!(v.max_active(), 46); + assert_eq!(v.max_tag(), Some(59)); + + v.insert(range(5, 7)).unwrap(); + assert_eq!(&*v.active, &[range(0, 8), range(10, 40), range(50, 60)]); + assert_eq!(v.max_active(), 48); + assert_eq!(v.max_tag(), Some(59)); + + v.insert(range(40, 50)).unwrap(); + assert_eq!(&*v.active, &[range(0, 8), range(10, 60)]); + assert_eq!(v.max_active(), 58); + assert_eq!(v.max_tag(), Some(59)); + + v.insert(range(8, 10)).unwrap(); + assert_eq!(&*v.active, &[range(0, 60)]); + assert_eq!(v.max_active(), 60); + assert_eq!(v.max_tag(), Some(59)); + } + + #[test] + fn test_replace() { + let mut v = Validate::new_test(vec![range(10, 20), range(30, 40)]); + + // Success + v.replace(range(12, 18)).unwrap(); + v.replace(range(10, 20)).unwrap(); + v.replace(range(15, 15)).unwrap(); + v.replace(range(30, 35)).unwrap(); + + // Failure + v.replace(range(15, 25)).unwrap_err(); + v.replace(range(5, 15)).unwrap_err(); + v.replace(range(25, 35)).unwrap_err(); + v.replace(range(40, 50)).unwrap_err(); + } + + #[test] + fn test_delete() { + let mut v = Validate::new_test(vec![range(10, 20), range(30, 40)]); + + assert_eq!(v.max_active(), 20); + + // Partial deletions + v.delete(range(15, 18)).unwrap(); + assert_eq!(&*v.active, &[range(10, 15), range(18, 20), range(30, 40)]); + + v.delete(range(35, 36)).unwrap(); + assert_eq!( + &*v.active, + &[range(10, 15), range(18, 20), range(30, 35), range(36, 40)] + ); + + // Partial at beginning + v.delete(range(10, 12)).unwrap(); + assert_eq!( + &*v.active, + &[range(12, 15), range(18, 20), range(30, 35), range(36, 40)] + ); + + // Partial at end + v.delete(range(32, 35)).unwrap(); + assert_eq!( + &*v.active, + &[range(12, 15), range(18, 20), range(30, 32), range(36, 40)] + ); + + // Full deletions + v.delete(range(18, 20)).unwrap(); + assert_eq!(&*v.active, &[range(12, 15), range(30, 32), range(36, 40)]); + + v.delete(range(36, 40)).unwrap(); + assert_eq!(&*v.active, &[range(12, 15), range(30, 32)]); + + assert_eq!(v.max_active(), 20); + assert_eq!(v.max_tag(), Some(39)); + + // Failure cases. + v.delete(range(20, 25)).unwrap_err(); + v.delete(range(0, 10)).unwrap_err(); + v.delete(range(10, 14)).unwrap_err(); + v.delete(range(31, 33)).unwrap_err(); + v.delete(range(40, 50)).unwrap_err(); + } +} From 41d06a5394e45f0f98f40a066a59f24ff7522536 Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Tue, 21 Jul 2026 20:09:48 +0000 Subject: [PATCH 08/10] make tool throw error instead of warning if not enough results --- .../src/bin/compute_streaming_groundtruth.rs | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/diskann-tools/src/bin/compute_streaming_groundtruth.rs b/diskann-tools/src/bin/compute_streaming_groundtruth.rs index 8f4ffc613..ea356265d 100644 --- a/diskann-tools/src/bin/compute_streaming_groundtruth.rs +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -119,8 +119,6 @@ fn run(args: &Args) -> CMDResult<()> { // Reverse map: external_id -> dataset_offset, needed to resolve Delete/Replace removals. let mut ext_to_offset: HashMap = HashMap::new(); - println!("Using distance function: {:?}", args.distance_function); - let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); for (stage_idx, stage) in runbook.stages().iter().enumerate() { @@ -226,7 +224,7 @@ fn run(args: &Args) -> CMDResult<()> { }) .collect(); - // Warn about queries that got fewer than K results (active set smaller than K). + // Fail if any query gets fewer than K results (active set smaller than K). let under_k: Vec = results .iter() .enumerate() @@ -239,16 +237,24 @@ fn run(args: &Args) -> CMDResult<()> { }) .collect(); if !under_k.is_empty() { - tracing::warn!( - "Stage {}: {} / {} queries have fewer than {} results (active set = {}). \ - Query indices: {:?}", - stage_idx, - under_k.len(), - n_queries, - recall_at, - n_active, - under_k, - ); + 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(CMDToolError { + details: format!( + "Stage {}: {} / {} queries have fewer than {} results (active set = {}). Query indices: {:?}{}", + stage_idx, + under_k.len(), + n_queries, + recall_at, + n_active, + preview, + suffix, + ), + }); } write_ground_truth::<()>( From 3dd5cb3024b4707264eebd34def78aa7add85b3a Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 24 Jul 2026 16:11:34 +0000 Subject: [PATCH 09/10] add stream --- .../index/bftree/full_precision_streaming.rs | 4 - .../src/bin/compute_streaming_groundtruth.rs | 342 +++++++++--------- 2 files changed, 179 insertions(+), 167 deletions(-) diff --git a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs index a87211a49..f80941ef8 100644 --- a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs +++ b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs @@ -226,9 +226,6 @@ where } } -<<<<<<< HEAD -fn bftree_streaming(input: &BfTreeDynamicRun, max_points: usize) -> BfTreeFPStreamingResult -======= type BfTreeStreamingPayload = ( bigann::WithData>, BfTreeFPIndex, @@ -238,7 +235,6 @@ fn bftree_streaming( input: &BfTreeDynamicRun, max_points: usize, ) -> anyhow::Result> ->>>>>>> 6caca35ddcfbacb2de38e57b7b53b7a7a95e1cf0 where T: bytemuck::Pod + VectorRepr + SampleableForStart, { diff --git a/diskann-tools/src/bin/compute_streaming_groundtruth.rs b/diskann-tools/src/bin/compute_streaming_groundtruth.rs index ea356265d..3cf8661ba 100644 --- a/diskann-tools/src/bin/compute_streaming_groundtruth.rs +++ b/diskann-tools/src/bin/compute_streaming_groundtruth.rs @@ -22,12 +22,14 @@ use anyhow::Context; use clap::Parser; use diskann::neighbor::{Neighbor, NeighborPriorityQueue}; use diskann::utils::VectorRepr; -use diskann_benchmark_core::streaming::executors::bigann::{FindGroundtruth, RunBook, Stage}; +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::*; @@ -100,7 +102,7 @@ fn run(args: &Args) -> CMDResult<()> { args.runbook_file, args.dataset_name ); - let runbook = RunBook::load( + let mut runbook = RunBook::load( Path::new(&args.runbook_file), &args.dataset_name, &mut finder, @@ -110,179 +112,193 @@ fn run(args: &Args) -> CMDResult<()> { })?; tracing::info!("Runbook has {} stages", runbook.len()); + let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); - // Boolean active-vector mask indexed by base offset. - let mut active: Vec = vec![false; n_base]; - // For each active dataset offset, the external ID that should appear in the groundtruth. - // For plain inserts external_id == dataset_offset; for Replace they diverge. - let mut ext_id: Vec = vec![0u32; n_base]; - // Reverse map: external_id -> dataset_offset, needed to resolve Delete/Replace removals. - let mut ext_to_offset: HashMap = HashMap::new(); + 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(), + }; - let distance_fn = V::distance(args.distance_function, Some(dataset.ncols())); + runbook + .run_with(&mut stream, |_| Ok(())) + .map_err(|e| CMDToolError { + details: e.to_string(), + })?; - for (stage_idx, stage) in runbook.stages().iter().enumerate() { - match stage { - Stage::Insert { - dataset_offsets_and_ids, - } => { - for id in dataset_offsets_and_ids.clone() { - if id < n_base { - active[id] = true; - ext_id[id] = id as u32; - ext_to_offset.insert(id as u32, id); - } + 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, + }); } - tracing::info!( - "Stage {}: insert {}..{} ({} active)", - stage_idx, - dataset_offsets_and_ids.start, - dataset_offsets_and_ids.end, - active.iter().filter(|&&b| b).count(), - ); + 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); } - Stage::Delete { ids } => { - for eid in ids.clone() { - if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { - active[offset] = false; - ext_to_offset.remove(&(eid as u32)); - } - } - tracing::info!( - "Stage {}: delete {}..{} ({} active)", - stage_idx, - ids.start, - ids.end, - active.iter().filter(|&&b| b).count(), - ); + } + 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)); } - Stage::Replace { - dataset_offsets, - ids, - } => { - // Remove old vectors by external ID. - for eid in ids.clone() { - if let Some(&offset) = ext_to_offset.get(&(eid as u32)) { - active[offset] = false; - ext_to_offset.remove(&(eid as u32)); - } - } - // Insert new vectors: they inherit the external IDs from `ids`. - for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) { - if offset < n_base { - active[offset] = true; - ext_id[offset] = eid as u32; - ext_to_offset.insert(eid as u32, offset); - } - } - tracing::info!( - "Stage {}: replace ids {}..{} with offsets {}..{} ({} active)", - stage_idx, - ids.start, - ids.end, - dataset_offsets.start, - dataset_offsets.end, - active.iter().filter(|&&b| b).count(), - ); + } + + 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); } - Stage::Search { - groundtruth: output_path, - } => { - let timer = Instant::now(); - let n_active = active.iter().filter(|&&b| b).count(); - - tracing::info!( - "Stage {}: computing top-{} groundtruth for {} active vectors against {} queries", - stage_idx, recall_at, n_active, n_queries, - ); - - let active_ids: Vec = active - .iter() - .enumerate() - .filter_map(|(i, &on)| if on { Some(i) } else { None }) - .collect(); - - // Compute KNN in parallel over queries. - // Use ext_id[offset] as the neighbor ID so that groundtruth IDs - // match external IDs (which differ from dataset offsets after Replace). - - // Using the global threadpool is fine here - #[allow(clippy::disallowed_methods)] - let results: Vec> = (0..n_queries) - .into_par_iter() - .map(|qi| { - let query = queries.row(qi); - let mut pq = NeighborPriorityQueue::new(recall_at); - for &offset in &active_ids { - let dist = distance_fn.evaluate_similarity(dataset.row(offset), query); - pq.insert(Neighbor { - id: ext_id[offset], - distance: dist, - }); - } - pq - }) - .collect(); - - // Fail if any query gets fewer than K results (active set smaller than K). - let under_k: Vec = results - .iter() - .enumerate() - .filter_map(|(qi, pq)| { - if pq.size() < recall_at { - Some(qi) - } else { - None - } - }) - .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(CMDToolError { - details: format!( - "Stage {}: {} / {} queries have fewer than {} results (active set = {}). Query indices: {:?}{}", - stage_idx, - under_k.len(), - n_queries, - recall_at, - n_active, - preview, - suffix, - ), - }); - } + } - write_ground_truth::<()>( - &storage, - output_path.to_str().ok_or_else(|| CMDToolError { - details: format!("Non-UTF8 output path: {}", output_path.display()), - })?, - n_queries, - recall_at, - results, - None, - ) - .map_err(|e| CMDToolError { - details: e.to_string(), - })?; - - tracing::info!( - "Stage {}: groundtruth written to {} in {:?}", - stage_idx, - output_path.display(), - timer.elapsed(), - ); + 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(()) } - tracing::info!("Done."); - Ok(()) + fn maintain(&mut self, _args: ()) -> anyhow::Result { + Ok(()) + } + + fn needs_maintenance(&mut self) -> bool { + false + } } #[derive(Debug, Parser)] From 89cae80ce3b6cca0c55268a97d643d83ee9059cd Mon Sep 17 00:00:00 2001 From: Magdalen Manohar Date: Fri, 24 Jul 2026 16:15:28 +0000 Subject: [PATCH 10/10] remove some intermediate states that were left by mistake --- diskann-benchmark/src/index/benchmarks.rs | 7 +------ .../src/index/bftree/full_precision_streaming.rs | 2 -- diskann-benchmark/src/index/bftree/spherical_streaming.rs | 2 -- 3 files changed, 1 insertion(+), 10 deletions(-) diff --git a/diskann-benchmark/src/index/benchmarks.rs b/diskann-benchmark/src/index/benchmarks.rs index 463c306bf..4307f2e5b 100644 --- a/diskann-benchmark/src/index/benchmarks.rs +++ b/diskann-benchmark/src/index/benchmarks.rs @@ -88,12 +88,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> )?; registry.register( "graph-index-full-precision-u8", - FullPrecision::::new() - .search(plugins::Topk) - .search(plugins::Range) - .search(plugins::TopkBetaFilter) - .search(plugins::TopkMultihopFilter) - .search(plugins::TopkInlineFilter), + FullPrecision::::new().search(plugins::Topk), )?; registry.register( "graph-index-full-precision-i8", diff --git a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs index f80941ef8..3d26773cd 100644 --- a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs +++ b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs @@ -48,8 +48,6 @@ use crate::{ //////////////////////// type BfTreeFPIndex = Arc>>; -type BfTreeFPStreamer = bigann::WithData>; -type BfTreeFPStreamingResult = anyhow::Result<(BfTreeFPStreamer, BfTreeFPIndex)>; /// The bf_tree streaming index implementation. /// diff --git a/diskann-benchmark/src/index/bftree/spherical_streaming.rs b/diskann-benchmark/src/index/bftree/spherical_streaming.rs index 5310f73e5..db773c8d9 100644 --- a/diskann-benchmark/src/index/bftree/spherical_streaming.rs +++ b/diskann-benchmark/src/index/bftree/spherical_streaming.rs @@ -52,8 +52,6 @@ use crate::{ type BfTreeSQProvider = BfTreeProvider; type BfTreeSQIndex = Arc>; -type BfTreeSQStreamer = bigann::WithData>; -type BfTreeSQStreamingResult = anyhow::Result<(BfTreeSQStreamer, BfTreeSQIndex)>; fn new_quantizer( quantizer: SphericalQuantizer,