From 29b35f199dd5ee0f26a8b5d07814841cb9dd9640 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Mon, 15 Jun 2026 12:58:16 -0700 Subject: [PATCH 1/7] Refactor streaming benchmarks and unify bftree inputs Behavioral changes on top of the module reorganization: - Introduce shared StreamRunner/Managed/build_streamer abstractions for streaming benchmarks, replacing duplicated per-backend logic - Rename Dynamic -> Streaming for runbook params and input types - Unify bftree build/streaming inputs with QuantConfig enum (none or spherical) instead of separate input structs - Add StreamingSearchParams with scalar num_threads/search_l, removing the vestigial groundtruth file requirement and unnecessary parameter sweep arrays from streaming search - Extract quantizer_util for shared spherical quantizer construction - Add InmemMaintainer for inmem streaming maintenance (DropDeleted + Release) - Privatize input struct fields behind accessor methods Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- diskann-benchmark/README.md | 4 +- .../example/graph-index-bftree-spherical.json | 35 +- .../graph-index-bftree-stream-spherical.json | 84 ++++ .../example/graph-index-bftree-stream.json | 26 +- .../example/graph-index-bftree.json | 5 +- ...x-dynamic.json => graph-index-stream.json} | 23 +- diskann-benchmark/src/index/benchmarks.rs | 219 ++++----- .../src/index/bftree/full_precision.rs | 34 +- .../index/bftree/full_precision_streaming.rs | 228 +++------ diskann-benchmark/src/index/bftree/mod.rs | 18 +- .../src/index/bftree/quantizer_util.rs | 67 +++ .../src/index/bftree/spherical.rs | 71 +-- .../src/index/bftree/spherical_streaming.rs | 257 +++------- .../src/index/bftree/streaming_utils.rs | 80 ---- diskann-benchmark/src/index/build.rs | 22 +- diskann-benchmark/src/index/inmem/mod.rs | 3 + diskann-benchmark/src/index/inmem/product.rs | 28 +- diskann-benchmark/src/index/inmem/scalar.rs | 42 +- .../src/index/inmem/spherical.rs | 46 +- .../src/index/inmem/streaming.rs | 122 +++++ diskann-benchmark/src/index/result.rs | 58 +-- .../src/index/streaming/full_precision.rs | 223 --------- diskann-benchmark/src/index/streaming/mod.rs | 121 ++++- .../src/index/streaming/runner.rs | 215 +++++++++ diskann-benchmark/src/inputs/bftree.rs | 443 +++++------------- diskann-benchmark/src/inputs/exhaustive.rs | 2 +- diskann-benchmark/src/inputs/graph_index.rs | 238 ++++++++-- diskann-benchmark/src/main.rs | 14 +- 28 files changed, 1315 insertions(+), 1413 deletions(-) create mode 100644 diskann-benchmark/example/graph-index-bftree-stream-spherical.json rename diskann-benchmark/example/{graph-index-dynamic.json => graph-index-stream.json} (69%) create mode 100644 diskann-benchmark/src/index/bftree/quantizer_util.rs delete mode 100644 diskann-benchmark/src/index/bftree/streaming_utils.rs create mode 100644 diskann-benchmark/src/index/inmem/streaming.rs delete mode 100644 diskann-benchmark/src/index/streaming/full_precision.rs create mode 100644 diskann-benchmark/src/index/streaming/runner.rs diff --git a/diskann-benchmark/README.md b/diskann-benchmark/README.md index b48fdf18a..6f756a434 100644 --- a/diskann-benchmark/README.md +++ b/diskann-benchmark/README.md @@ -269,9 +269,9 @@ First, set up the runbook and ground truth for the desired workload. Refer to th Benchmarks are run with ```sh -cargo run --release --package diskann-benchmark -- run --input-file ./diskann-benchmark/example/graph-index-dynamic.json --output-file dynamic-output.json +cargo run --release --package diskann-benchmark -- run --input-file ./diskann-benchmark/example/graph-index-stream.json --output-file stream-output.json ``` -Note in the example json that the benchmark is registered under `graph-index-dynamic-run`, +Note in the example json that the benchmark is registered under `graph-index-stream-run`, instead of `graph-index-build` etc.. A streaming run happens in several phases. diff --git a/diskann-benchmark/example/graph-index-bftree-spherical.json b/diskann-benchmark/example/graph-index-bftree-spherical.json index 064fe8047..198511484 100644 --- a/diskann-benchmark/example/graph-index-bftree-spherical.json +++ b/diskann-benchmark/example/graph-index-bftree-spherical.json @@ -4,7 +4,7 @@ ], "jobs": [ { - "type": "graph-index-build-bftree-spherical-quantization", + "type": "graph-index-bftree", "content": { "build": { "data_type": "float32", @@ -40,10 +40,24 @@ } ] }, - "seed": 42, - "transform_kind": "null", - "num_bits": 2, - "pre_scale": "reciprocal_mean_norm", + "quantization": { + "kind": "spherical", + "seed": 42, + "transform_kind": "null", + "num_bits": 2, + "pre_scale": "reciprocal_mean_norm", + "quant_store_config": { + "cb_size_byte": 67108864, + "leaf_page_size": 4096, + "cb_max_record_size": null, + "cb_min_record_size": null, + "read_promotion_rate": null, + "scan_promotion_rate": null, + "cb_copy_on_access_ratio": null, + "read_record_cache": null, + "cache_only": null + } + }, "vector_store_config": { "cb_size_byte": 67108864, "leaf_page_size": 4096, @@ -65,17 +79,6 @@ "cb_copy_on_access_ratio": null, "read_record_cache": null, "cache_only": null - }, - "quant_store_config": { - "cb_size_byte": 67108864, - "leaf_page_size": 4096, - "cb_max_record_size": null, - "cb_min_record_size": null, - "read_promotion_rate": null, - "scan_promotion_rate": null, - "cb_copy_on_access_ratio": null, - "read_record_cache": null, - "cache_only": null } } } diff --git a/diskann-benchmark/example/graph-index-bftree-stream-spherical.json b/diskann-benchmark/example/graph-index-bftree-stream-spherical.json new file mode 100644 index 000000000..26a79c08d --- /dev/null +++ b/diskann-benchmark/example/graph-index-bftree-stream-spherical.json @@ -0,0 +1,84 @@ +{ + "search_directories": [ + "test_data/disk_index_search" + ], + "jobs": [ + { + "type": "graph-index-stream-bftree", + "content": { + "build": { + "data_type": "float32", + "data": "disk_index_siftsmall_learn_256pts_data.fbin", + "distance": "squared_l2", + "max_degree": 32, + "l_build": 50, + "start_point_strategy": "medoid", + "alpha": 1.2, + "backedge_ratio": 1.0, + "num_threads": 1 + }, + "search": { + "queries": "disk_index_sample_query_10pts.fbin", + "reps": 1, + "num_threads": 1, + "search_l": 40, + "search_n": 10, + "recall_k": 10 + }, + "runbook_params": { + "runbook_path": "example_runbook.yaml", + "dataset_name": "sift-small-256", + "gt_directory": "example_runbook_gt", + "ip_delete_method": { + "method": "visited_and_top_k", + "params": { + "k_value": 10, + "l_value": 40 + } + }, + "ip_delete_num_to_replace": 3 + }, + "quantization": { + "kind": "spherical", + "seed": 42, + "transform_kind": "null", + "num_bits": 2, + "pre_scale": "reciprocal_mean_norm", + "quant_store_config": { + "cb_size_byte": 67108864, + "leaf_page_size": 4096, + "cb_max_record_size": null, + "cb_min_record_size": null, + "read_promotion_rate": null, + "scan_promotion_rate": null, + "cb_copy_on_access_ratio": null, + "read_record_cache": null, + "cache_only": null + } + }, + "vector_store_config": { + "cb_size_byte": 67108864, + "leaf_page_size": 4096, + "cb_max_record_size": null, + "cb_min_record_size": null, + "read_promotion_rate": null, + "scan_promotion_rate": null, + "cb_copy_on_access_ratio": null, + "read_record_cache": null, + "cache_only": null + }, + "neighbor_store_config": { + "cb_size_byte": 67108864, + "leaf_page_size": 4096, + "cb_max_record_size": null, + "cb_min_record_size": null, + "read_promotion_rate": null, + "scan_promotion_rate": null, + "cb_copy_on_access_ratio": null, + "read_record_cache": null, + "cache_only": null + } + } + } + ] +} diff --git a/diskann-benchmark/example/graph-index-bftree-stream.json b/diskann-benchmark/example/graph-index-bftree-stream.json index dbcfd9677..ec42978fb 100644 --- a/diskann-benchmark/example/graph-index-bftree-stream.json +++ b/diskann-benchmark/example/graph-index-bftree-stream.json @@ -4,7 +4,7 @@ ], "jobs": [ { - "type": "graph-index-stream-bftree-full-precision", + "type": "graph-index-stream-bftree", "content": { "build": { "data_type": "float32", @@ -17,24 +17,13 @@ "backedge_ratio": 1.0, "num_threads": 1 }, - "search_phase": { - "search-type": "topk", + "search": { "queries": "disk_index_sample_query_10pts.fbin", - "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin", "reps": 1, - "num_threads": [ - 1 - ], - "runs": [ - { - "search_n": 10, - "search_l": [ - 20, - 40 - ], - "recall_k": 10 - } - ] + "num_threads": 1, + "search_l": 40, + "search_n": 10, + "recall_k": 10 }, "runbook_params": { "runbook_path": "example_runbook.yaml", @@ -49,6 +38,9 @@ }, "ip_delete_num_to_replace": 3 }, + "quantization": { + "kind": "none" + }, "vector_store_config": { "cb_size_byte": 67108864, "leaf_page_size": 4096, diff --git a/diskann-benchmark/example/graph-index-bftree.json b/diskann-benchmark/example/graph-index-bftree.json index 69634ef03..6a0cc8ef8 100644 --- a/diskann-benchmark/example/graph-index-bftree.json +++ b/diskann-benchmark/example/graph-index-bftree.json @@ -4,7 +4,7 @@ ], "jobs": [ { - "type": "graph-index-build-bftree-full-precision", + "type": "graph-index-bftree", "content": { "build": { "data_type": "float32", @@ -37,6 +37,9 @@ } ] }, + "quantization": { + "kind": "none" + }, "vector_store_config": { "cb_size_byte": 67108864, "leaf_page_size": 4096, diff --git a/diskann-benchmark/example/graph-index-dynamic.json b/diskann-benchmark/example/graph-index-stream.json similarity index 69% rename from diskann-benchmark/example/graph-index-dynamic.json rename to diskann-benchmark/example/graph-index-stream.json index b18629dbd..3434eb2c3 100644 --- a/diskann-benchmark/example/graph-index-dynamic.json +++ b/diskann-benchmark/example/graph-index-stream.json @@ -4,7 +4,7 @@ ], "jobs": [ { - "type": "graph-index-dynamic-run", + "type": "graph-index-stream-run", "content": { "build": { "data_type": "float32", @@ -17,24 +17,13 @@ "num_threads": 4, "start_point_strategy": "medoid" }, - "search_phase": { - "search-type": "topk", + "search": { "queries": "disk_index_sample_query_10pts.fbin", - "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin", "reps": 1, - "num_threads": [ - 2 - ], - "runs": [ - { - "search_n": 10, - "search_l": [ - 20, - 40 - ], - "recall_k": 10 - } - ] + "num_threads": 2, + "search_l": 40, + "search_n": 10, + "recall_k": 10 }, "runbook_params": { "runbook_path": "example_runbook.yaml", diff --git a/diskann-benchmark/src/index/benchmarks.rs b/diskann-benchmark/src/index/benchmarks.rs index 0a66576a5..43f389461 100644 --- a/diskann-benchmark/src/index/benchmarks.rs +++ b/diskann-benchmark/src/index/benchmarks.rs @@ -11,10 +11,7 @@ use diskann::{ provider::{self, DataProvider, DefaultContext}, utils::VectorRepr, }; -use diskann_benchmark_core::{ - self as benchmark_core, - streaming::{executors::bigann, Executor}, -}; +use diskann_benchmark_core::{self as benchmark_core, streaming::executors::bigann}; use diskann_benchmark_runner::{ benchmark::{FailureScore, MatchScore}, output::Output, @@ -42,11 +39,14 @@ use super::{ }; use crate::{ index::{ + inmem::InmemMaintainer, result::{AggregatedSearchResults, BuildResult}, search::plugins, - streaming::{self, managed, stats::StreamStats, FullPrecisionStream, Managed}, + streaming::{self, managed, stats::StreamStats, Managed, StreamRunner}, + }, + inputs::graph_index::{ + IndexBuild, IndexOperation, IndexSource, SearchPhase, StreamingIndexRun, }, - inputs::graph_index::{DynamicIndexRun, IndexBuild, IndexOperation, IndexSource, SearchPhase}, utils::{ self, datafiles::{self}, @@ -58,7 +58,7 @@ use crate::{ // Benchmark Registration // //////////////////////////// -pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> { +pub(super) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> { // Notes on registration: // // We register all supported search types for `f32`, but intentionally limit the number @@ -95,22 +95,22 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> FullPrecision::::new().search(plugins::Topk), )?; - // Dynamic Full Precision + // Streaming Full Precision registry.register( - "graph-index-dynamic-full-precision-f32", - DynamicFullPrecision::::new(), + "graph-index-stream-full-precision-f32", + StreamingFullPrecision::::new(), )?; registry.register( - "graph-index-dynamic-full-precision-f16", - DynamicFullPrecision::::new(), + "graph-index-stream-full-precision-f16", + StreamingFullPrecision::::new(), )?; registry.register( - "graph-index-dynamic-full-precision-u8", - DynamicFullPrecision::::new(), + "graph-index-stream-full-precision-u8", + StreamingFullPrecision::::new(), )?; registry.register( - "graph-index-dynamic-full-precision-i8", - DynamicFullPrecision::::new(), + "graph-index-stream-full-precision-i8", + StreamingFullPrecision::::new(), )?; product::register_benchmarks(registry)?; @@ -129,7 +129,7 @@ type FullPrecisionProvider = inmem::DefaultProvider< /// Associate a type (usually a [`diskann::provider::DataProvider`]) with a full-precision /// element type. This is used in implementations of [`plugins::Plugin`] to derive the /// correct query types to load. -pub(crate) trait QueryType { +pub(super) trait QueryType { type Element: VectorRepr; } @@ -180,8 +180,8 @@ where type Output = BuildResult; fn try_match(&self, input: &IndexOperation) -> Result { - let score = utils::match_data_type::(*input.source.data_type()); - if self.plugins.is_match(&input.search_phase) { + let score = utils::match_data_type::(*input.source().data_type()); + if self.plugins.is_match(input.search_phase()) { score } else { match score { @@ -198,16 +198,16 @@ where ) -> std::fmt::Result { match input { Some(arg) => { - let desc = T::describe(*arg.source.data_type()); + let desc = T::describe(*arg.source().data_type()); if !desc.is_match() { writeln!(f, "Data/Query Type: {}", desc)?; } - if !self.plugins.is_match(&arg.search_phase) { + if !self.plugins.is_match(arg.search_phase()) { writeln!( f, "Unsupported search phase: \"{}\" - expected one of {}", - arg.search_phase.kind(), + arg.search_phase().kind(), self.plugins.format_kinds(), )?; } @@ -227,7 +227,7 @@ where mut output: &mut dyn Output, ) -> anyhow::Result { writeln!(output, "{}", input)?; - let (index, build_stats) = match &input.source { + let (index, build_stats) = match input.source() { IndexSource::Build(build) => { let (index, build_stats) = run_build( build, @@ -272,7 +272,7 @@ where let search_results = self.plugins.run( index, - &input.search_phase, + input.search_phase(), &Strategy::new(common::FullPrecision), )?; @@ -283,12 +283,12 @@ where } } -// Graph Index Dynamic Run -pub(crate) struct DynamicFullPrecision { +// Graph Index Streaming Run +pub(super) struct StreamingFullPrecision { _type: std::marker::PhantomData, } -impl DynamicFullPrecision { +impl StreamingFullPrecision { fn new() -> Self { Self { _type: std::marker::PhantomData, @@ -296,97 +296,51 @@ impl DynamicFullPrecision { } } -impl Benchmark for DynamicFullPrecision +impl Benchmark for StreamingFullPrecision where T: VectorRepr + diskann_utils::sampling::WithApproximateNorm + diskann::graph::SampleableForStart + AsDataType, { - type Input = DynamicIndexRun; + type Input = StreamingIndexRun; type Output = Vec>; - fn try_match(&self, input: &DynamicIndexRun) -> Result { - utils::match_data_type::(input.build.data_type()) + fn try_match(&self, input: &StreamingIndexRun) -> Result { + utils::match_data_type::(input.build().data_type()) } fn description( &self, f: &mut std::fmt::Formatter<'_>, - input: Option<&DynamicIndexRun>, + input: Option<&StreamingIndexRun>, ) -> std::fmt::Result { match input { - Some(i) => write!(f, "{}", T::describe(i.build.data_type())), + Some(i) => write!(f, "{}", T::describe(i.build().data_type())), None => write!(f, "{}", T::DATA_TYPE), } } fn run( &self, - input: &DynamicIndexRun, + input: &StreamingIndexRun, _checkpoint: Checkpoint<'_>, mut output: &mut dyn Output, ) -> anyhow::Result>> { writeln!(output, "{}", input)?; - let groundtruth_directory = input - .runbook_params - .resolved_gt_directory - .as_ref() - .ok_or_else(|| { - anyhow::anyhow!("Ground truth directory path was not resolved during validation") - })?; - - let mut runbook = bigann::RunBook::load( - &input.runbook_params.runbook_path, - &input.runbook_params.dataset_name, - &mut bigann::ScanDirectory::new(groundtruth_directory)?, - )?; - - let mut streamer = full_precision_streaming::(input, runbook.max_points())?; - - let mut results = Vec::new(); - let stages = runbook.len(); - let mut i = 1; - - runbook.run_with( - &mut streamer, - |o: managed::Stats| -> anyhow::Result<()> { - if o.inner().is_maintain() { - let message = format!("Ran maintenance before stage {}", i); - write!(output, "{}", crate::utils::SmallBanner(&message))?; - } else { - let message = - format!("Finished stage {} of {}: {}", i, stages, o.inner().kind()); - write!(output, "{}", crate::utils::SmallBanner(&message))?; - i += 1; - } - writeln!(output, "{}", o)?; - results.push(o); - Ok(()) - }, - )?; - - write!( - output, - "{}", - crate::utils::SmallBanner("End of Run Summary") - )?; - - writeln!( + streaming::run_streaming::( + input.runbook_params(), + |max_points| full_precision_streaming::(input, max_points), output, - "{}", - streaming::stats::Summary::new(results.iter().map(|r| r.inner())) - )?; - - Ok(results) + ) } } // Simplify reasoning about this rather hefty type. type Index = Arc>; -pub(crate) fn run_build( +pub(super) fn run_build( input: &IndexBuild, build_strategy: B, data: Option>>, @@ -424,14 +378,14 @@ where /// This exists so we can implement [`search::Plugin`] for a raw generic `DP` without /// forming a blanket implementation for all `DP`/parameter `P` pairs. #[derive(Debug, Clone, Copy)] -pub(crate) struct Strategy(S); +pub(super) struct Strategy(S); impl Strategy { - pub(crate) fn new(strategy: S) -> Self { + pub(super) fn new(strategy: S) -> Self { Self(strategy) } - pub(crate) fn inner(&self) -> S + pub(super) fn inner(&self) -> S where S: Clone, { @@ -766,7 +720,7 @@ where /// The stack looks like this: /// -/// - Bottom: [`FullPrecisionStream`]: The core streaming index implementation. +/// - Bottom: [`StreamRunner`]: The core streaming index implementation. /// - Middle: [`Managed`]: Since the in-mem index currently does not split internal and external /// IDs, the [`Managed`] layer is introduced as a temporary measure. This is responsible /// for ID mapping. @@ -774,64 +728,49 @@ where /// /// This function constructs the entire stack. fn full_precision_streaming( - input: &DynamicIndexRun, + input: &StreamingIndexRun, max_points: usize, ) -> anyhow::Result>> where T: bytemuck::Pod + VectorRepr + WithApproximateNorm + SampleableForStart, { - let topk = input.search_phase.as_topk()?; + let search = input.search(); let consolidate_threshold: f32 = input - .runbook_params + .runbook_params() .consolidate_threshold .ok_or_else(|| anyhow::anyhow!("consolidate_threshold is required for inmem streaming"))?; - - let data = datafiles::load_dataset::(datafiles::BinFile(input.build.data()))?; - let queries = Arc::new(datafiles::load_dataset::(datafiles::BinFile( - &topk.queries, - ))?); - - // Create a little extra headroom to handle deferred maintenance. - let max_points = ((max_points as f32) * (1.0 + 2.0 * consolidate_threshold)).ceil() as usize; - - let index = diskann_async::new_index::( - input.try_as_config(input.build.l_build())?.build()?, - input.inmem_parameters(max_points, data.ncols()), - common::TableBasedDeletes, - )?; - - build::set_start_points( - index.provider(), - data.as_view(), - *input.build.start_point_strategy(), - )?; - - let num_threads_and_tasks = NonZeroUsize::new(input.build.num_threads()).unwrap(); - let managed_stream = FullPrecisionStream { - index, - search: topk.clone(), - runtime: benchmark_core::tokio::runtime(num_threads_and_tasks.get())?, - ntasks: num_threads_and_tasks, - inplace_delete_num_to_replace: input.runbook_params.ip_delete_num_to_replace, - inplace_delete_method: input.runbook_params.ip_delete_method.into(), - }; - - let managed = Managed::new( - max_points, - managed::SlotReclaim::Deferred(consolidate_threshold), - managed_stream, - ); - - // compute the maximum value of k used in any search - let max_k = topk.max_k(); - - let layered = bigann::WithData::new(managed, data, queries, move |path| { - Ok(Box::new(datafiles::load_groundtruth( - datafiles::BinFile(path), - Some(max_k), - )?)) - }); - - Ok(layered) + let capacity = ((max_points as f32) * (1.0 + 2.0 * consolidate_threshold)).ceil() as usize; + + streaming::build_streamer( + input.build().data(), + search, + streaming::managed::SlotReclaim::Deferred(consolidate_threshold), + capacity, + |data, capacity| { + let index = diskann_async::new_index::( + input.try_as_config(input.build().l_build())?.build()?, + input.inmem_parameters(capacity, data.ncols()), + common::TableBasedDeletes, + )?; + + build::set_start_points( + index.provider(), + data.as_view(), + *input.build().start_point_strategy(), + )?; + + let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap(); + Ok(StreamRunner::new( + index, + common::FullPrecision, + search.clone(), + benchmark_core::tokio::runtime(num_threads_and_tasks.get())?, + num_threads_and_tasks, + input.runbook_params().ip_delete_num_to_replace, + input.runbook_params().ip_delete_method.into(), + InmemMaintainer, + )) + }, + ) } diff --git a/diskann-benchmark/src/index/bftree/full_precision.rs b/diskann-benchmark/src/index/bftree/full_precision.rs index d6f2d1f54..0acf849ce 100644 --- a/diskann-benchmark/src/index/bftree/full_precision.rs +++ b/diskann-benchmark/src/index/bftree/full_precision.rs @@ -25,7 +25,10 @@ use crate::{ result::BuildResult, search::plugins::{Plugin, Plugins}, }, - inputs::{bftree::BfTreeFullPrecisionBuild, graph_index::SearchPhase}, + inputs::{ + bftree::{BfTreeBuild, QuantConfig}, + graph_index::SearchPhase, + }, utils::{self}, }; @@ -68,25 +71,32 @@ impl Benchmark for BfTreeFullPrecision where T: VectorRepr + AsDataType + SampleableForStart + WithApproximateNorm + 'static, { - type Input = BfTreeFullPrecisionBuild; + type Input = BfTreeBuild; type Output = BuildResult; fn try_match(&self, input: &Self::Input) -> Result { - let score = utils::match_data_type::(input.data_type()); - if self.plugins.is_match(input.search_phase()) { - score - } else { - match score { - Ok(_) => Err(FailureScore(0)), - Err(s) => Err(s), - } + let mut failure_score: Option = None; + + if !matches!(input.quantization(), QuantConfig::None) { + *failure_score.get_or_insert(0) += 1; + } + if let Err(s) = utils::match_data_type::(input.data_type()) { + *failure_score.get_or_insert(0) += s.0; + } + if !self.plugins.is_match(input.search_phase()) { + *failure_score.get_or_insert(0) += 1; + } + + match failure_score { + None => Ok(MatchScore(0)), + Some(score) => Err(FailureScore(score)), } } fn description( &self, f: &mut std::fmt::Formatter<'_>, - input: Option<&Self::Input>, + input: Option<&BfTreeBuild>, ) -> std::fmt::Result { match input { Some(arg) => { @@ -113,7 +123,7 @@ where fn run( &self, - input: &Self::Input, + input: &BfTreeBuild, checkpoint: Checkpoint<'_>, mut output: &mut dyn Output, ) -> anyhow::Result { diff --git a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs index e3d460a78..73019fa02 100644 --- a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs +++ b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs @@ -2,13 +2,13 @@ * Copyright (c) Microsoft Corporation. * Licensed under the MIT license. */ -use std::{borrow::Cow, io::Write, num::NonZeroUsize, sync::Arc}; +use std::{io::Write, num::NonZeroUsize, sync::Arc}; use diskann::{ - graph::{DiskANNIndex, InplaceDeleteMethod, SampleableForStart}, - utils::{VectorRepr, ONE}, + graph::{DiskANNIndex, SampleableForStart}, + utils::VectorRepr, }; -use diskann_benchmark_core::{self as benchmark_core, recall::Rows, streaming::executors::bigann}; +use diskann_benchmark_core::{self as benchmark_core, streaming::executors::bigann}; use diskann_benchmark_runner::{ benchmark::{FailureScore, MatchScore}, output::Output, @@ -17,133 +17,24 @@ use diskann_benchmark_runner::{ }; use diskann_bftree::{BfTreeProvider, NoStore}; use diskann_providers::model::graph::provider::async_::common::FullPrecision; -use diskann_utils::{ - sampling::WithApproximateNorm, - views::{Matrix, MatrixView}, -}; +use diskann_utils::sampling::WithApproximateNorm; use crate::{ - index::{ - build::{BuildKind, BuildStats}, - search::knn, - streaming::{ - managed::{self, Managed}, - stats::{GenericStats, StreamStats}, - ManagedStream, - }, - }, - inputs::{ - bftree::BfTreeDynamicRun, - graph_index::{SearchPhase, TopkSearchPhase}, + index::streaming::{ + managed::{self, Managed}, + runner::BfTreeMaintainer, + stats::StreamStats, + StreamRunner, }, - utils::{self, datafiles}, + inputs::bftree::{BfTreeStreamingRun, QuantConfig}, + utils, }; //////////////////////// // Streaming BfTree // //////////////////////// -type BfTreeFPIndex = Arc>>; - -/// The bf_tree streaming index implementation. -/// -/// Mirrors the in-memory `FullPrecisionStream` but targets `BfTreeProvider`. -struct BfTreeStream -where - T: VectorRepr, -{ - index: BfTreeFPIndex, - search: TopkSearchPhase, - runtime: tokio::runtime::Runtime, - ntasks: NonZeroUsize, - inplace_delete_num_to_replace: usize, - inplace_delete_method: InplaceDeleteMethod, -} - -impl BfTreeStream -where - T: VectorRepr, -{ - fn insert_(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { - let runner = benchmark_core::build::graph::SingleInsert::new( - self.index.clone(), - Arc::new(data.to_owned()), - FullPrecision, - benchmark_core::build::ids::Slice::new(slots.into()), - ); - - let results = benchmark_core::build::build( - runner, - benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - BuildStats::new(BuildKind::SingleInsert, results) - } -} - -impl ManagedStream for BfTreeStream -where - T: VectorRepr, -{ - type Output = StreamStats; - - fn search( - &self, - queries: Arc>, - groundtruth: &dyn Rows, - ) -> anyhow::Result { - let knn = benchmark_core::search::graph::KNN::new( - self.index.clone(), - queries, - benchmark_core::search::graph::Strategy::broadcast(FullPrecision), - )?; - - let steps = knn::SearchSteps::new( - self.search.reps, - &self.search.num_threads, - &self.search.runs, - ); - let results = knn::run(&knn, groundtruth, steps)?; - Ok(StreamStats::Search(results)) - } - - fn insert(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { - Ok(StreamStats::Insert(self.insert_(data, slots)?)) - } - - fn replace(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { - Ok(StreamStats::Replace(self.insert_(data, slots)?)) - } - - fn delete(&self, slots: &[u32]) -> anyhow::Result { - let runner = benchmark_core::streaming::graph::InplaceDelete::new( - self.index.clone(), - FullPrecision, - self.inplace_delete_num_to_replace, - self.inplace_delete_method, - benchmark_core::build::ids::Slice::new(slots.into()), - ); - - let results = benchmark_core::build::build( - runner, - benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - Ok(StreamStats::Delete(GenericStats::new( - Cow::Borrowed("Delete"), - results, - )?)) - } - - fn maintain(&self) -> anyhow::Result { - // bf-tree uses hard deletes — no deferred cleanup needed. - Ok(StreamStats::Maintain(Vec::new())) - } -} - -/// The dynamic/streaming benchmark for bf_tree full precision. +/// The streaming benchmark for bf_tree full precision. pub(super) struct StreamingFullPrecision { _type: std::marker::PhantomData, } @@ -160,19 +51,18 @@ impl Benchmark for StreamingFullPrecision where T: VectorRepr + WithApproximateNorm + SampleableForStart + AsDataType + bytemuck::Pod, { - type Input = BfTreeDynamicRun; + type Input = BfTreeStreamingRun; type Output = Vec>; fn try_match(&self, input: &Self::Input) -> Result { let mut failure_score: Option = None; - if let Err(s) = utils::match_data_type::(input.data_type()) { - failure_score = Some(s.0); - } - - if !matches!(input.search_phase(), SearchPhase::Topk(_)) { + if !matches!(input.quantization(), QuantConfig::None) { *failure_score.get_or_insert(0) += 1; } + if let Err(s) = utils::match_data_type::(input.data_type()) { + *failure_score.get_or_insert(0) += s.0; + } match failure_score { None => Ok(MatchScore(0)), @@ -186,7 +76,9 @@ where input: Option<&Self::Input>, ) -> std::fmt::Result { match input { - Some(i) => write!(f, "{}", T::describe(i.build().data_type())), + Some(i) => { + write!(f, "{}", T::describe(i.build().data_type())) + } None => write!(f, "{}", T::DATA_TYPE), } } @@ -199,7 +91,7 @@ where ) -> anyhow::Result { writeln!(output, "{}", input)?; - super::streaming_utils::run_streaming::( + crate::index::streaming::run_streaming::( input.runbook_params(), |max_points| bftree_streaming::(input, max_points), output, @@ -208,55 +100,43 @@ where } fn bftree_streaming( - input: &BfTreeDynamicRun, + input: &BfTreeStreamingRun, max_points: usize, ) -> anyhow::Result>> where T: bytemuck::Pod + VectorRepr + WithApproximateNorm + SampleableForStart, { - let topk = match &input.search_phase() { - SearchPhase::Topk(topk) => topk, - _ => anyhow::bail!("Only TopK is currently supported by the streaming index"), - }; - - let data = datafiles::load_dataset::(datafiles::BinFile(input.build().data()))?; - let queries = Arc::new(datafiles::load_dataset::(datafiles::BinFile( - &topk.queries, - ))?); - - let config = input.try_as_config()?.build()?; - let params = input.bftree_parameters(max_points, data.ncols()); - let start_points = input - .build() - .start_point_strategy() - .compute(data.as_view())?; - let provider = BfTreeProvider::new(params, start_points.as_view(), NoStore)?; - let index = Arc::new(DiskANNIndex::new(config, provider, None)); - - let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap(); - let managed_stream = BfTreeStream { - index, - search: topk.clone(), - runtime: benchmark_core::tokio::runtime(num_threads_and_tasks.get())?, - ntasks: num_threads_and_tasks, - inplace_delete_num_to_replace: input.runbook_params().ip_delete_num_to_replace, - inplace_delete_method: input.runbook_params().ip_delete_method.into(), - }; + let search = input.search(); let num_start_points = input.build().start_point_strategy().count(); - let managed = Managed::new( - max_points + num_start_points, - managed::SlotReclaim::Immediate, - managed_stream, - ); - - let max_k = topk.max_k(); - let layered = bigann::WithData::new(managed, data, queries, move |path| { - Ok(Box::new(datafiles::load_groundtruth( - datafiles::BinFile(path), - Some(max_k), - )?)) - }); - - Ok(layered) + let capacity = max_points + num_start_points; + + crate::index::streaming::build_streamer( + input.build().data(), + search, + crate::index::streaming::managed::SlotReclaim::Immediate, + capacity, + |data, capacity| { + let config = input.try_as_config()?.build()?; + let params = input.bftree_parameters(capacity, data.ncols()); + let start_points = input + .build() + .start_point_strategy() + .compute(data.as_view())?; + let provider = BfTreeProvider::new(params, start_points.as_view(), NoStore)?; + let index = Arc::new(DiskANNIndex::new(config, provider, None)); + + let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap(); + Ok(StreamRunner::new( + index, + FullPrecision, + search.clone(), + benchmark_core::tokio::runtime(num_threads_and_tasks.get())?, + num_threads_and_tasks, + input.runbook_params().ip_delete_num_to_replace, + input.runbook_params().ip_delete_method.into(), + BfTreeMaintainer, + )) + }, + ) } diff --git a/diskann-benchmark/src/index/bftree/mod.rs b/diskann-benchmark/src/index/bftree/mod.rs index 600d94d50..7619a0815 100644 --- a/diskann-benchmark/src/index/bftree/mod.rs +++ b/diskann-benchmark/src/index/bftree/mod.rs @@ -10,19 +10,17 @@ //! once, search) and streaming (insert/delete/search interleaved) modes. //! //! Registered tags: -//! - `graph-index-bftree-full-precision-f32` — static FP build + search -//! - `graph-index-bftree-stream-full-precision-f32` — streaming FP -//! - `graph-index-build-bftree-spherical-quantization` — static spherical (1/2/4-bit) -//! - `graph-index-stream-bftree-spherical-quantization` — streaming spherical (1/2/4-bit) +//! - `graph-index-bftree` — static build + search (full-precision or spherical) +//! - `graph-index-stream-bftree` — streaming (full-precision or spherical) -use crate::index::search::plugins::Topk; +use super::search::plugins::Topk; use diskann_benchmark_runner::Registry; mod full_precision; mod full_precision_streaming; +mod quantizer_util; mod spherical; mod spherical_streaming; -mod streaming_utils; pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> { registry.register( @@ -31,13 +29,13 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> )?; registry.register( - "graph-index-bftree-stream-full-precision-f32", - full_precision_streaming::StreamingFullPrecision::::new(), + "graph-index-bftree-spherical-quantization", + spherical::BfTreeSpherical::new().search(Topk), )?; registry.register( - "graph-index-bftree-spherical-quantization", - spherical::BfTreeSpherical::new().search(Topk), + "graph-index-stream-bftree-full-precision-f32", + full_precision_streaming::StreamingFullPrecision::::new(), )?; registry.register( diff --git a/diskann-benchmark/src/index/bftree/quantizer_util.rs b/diskann-benchmark/src/index/bftree/quantizer_util.rs new file mode 100644 index 000000000..6f6a9d6e6 --- /dev/null +++ b/diskann-benchmark/src/index/bftree/quantizer_util.rs @@ -0,0 +1,67 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_quantization::{ + alloc::{AllocatorError, GlobalAllocator, Poly}, + spherical::{ + iface::{self as spherical_iface, Quantizer}, + SphericalQuantizer, + }, +}; +use diskann_utils::views::MatrixView; +use rand::SeedableRng; + +use crate::{inputs::bftree::QuantConfig, utils::SimilarityMeasure}; + +fn new_quantizer( + quantizer: SphericalQuantizer, +) -> Result, AllocatorError> +where + spherical_iface::Impl: spherical_iface::Constructible + Quantizer, +{ + let imp = spherical_iface::Impl::::new(quantizer)?; + diskann_quantization::poly!(Quantizer, imp, GlobalAllocator) +} + +pub(super) fn build_quantizer( + quantization: &QuantConfig, + data: MatrixView<'_, f32>, + distance: SimilarityMeasure, +) -> anyhow::Result>> { + match quantization { + QuantConfig::None => Ok(None), + QuantConfig::Spherical { + seed, + transform_kind, + num_bits, + pre_scale, + .. + } => { + let m: diskann_vector::distance::Metric = distance.into(); + let pre_scale = match pre_scale { + Some(v) => (*v).try_into()?, + None => diskann_quantization::spherical::PreScale::None, + }; + + let quantizer = SphericalQuantizer::train( + data, + transform_kind.into(), + m.try_into()?, + pre_scale, + &mut rand::rngs::StdRng::seed_from_u64(*seed), + GlobalAllocator, + )?; + + let poly = match num_bits.get() { + 1 => new_quantizer::<1>(quantizer)?, + 2 => new_quantizer::<2>(quantizer)?, + 4 => new_quantizer::<4>(quantizer)?, + n => anyhow::bail!("{n} bits not supported for spherical quantization"), + }; + + Ok(Some(poly)) + } + } +} diff --git a/diskann-benchmark/src/index/bftree/spherical.rs b/diskann-benchmark/src/index/bftree/spherical.rs index 0e9ae9a70..119321798 100644 --- a/diskann-benchmark/src/index/bftree/spherical.rs +++ b/diskann-benchmark/src/index/bftree/spherical.rs @@ -14,15 +14,7 @@ use diskann_benchmark_runner::{ }; use diskann_bftree::{quant::QuantVectorProvider, BfTreeProvider}; use diskann_providers::model::graph::provider::async_::common::Quantized; -use diskann_quantization::{ - alloc::{AllocatorError, GlobalAllocator, Poly}, - spherical::{ - iface::{self as spherical_iface, Quantizer}, - SphericalQuantizer, - }, -}; use diskann_utils::views::Matrix; -use rand::SeedableRng; use crate::{ index::{ @@ -31,7 +23,10 @@ use crate::{ result::BuildResult, search::plugins::{Plugin, Plugins}, }, - inputs::{bftree::BfTreeSphericalBuild, graph_index::SearchPhase}, + inputs::{ + bftree::{BfTreeBuild, QuantConfig}, + graph_index::SearchPhase, + }, utils::{self, datafiles}, }; @@ -65,29 +60,19 @@ impl BfTreeSpherical { } } -fn new_quantizer( - quantizer: SphericalQuantizer, -) -> Result, AllocatorError> -where - spherical_iface::Impl: spherical_iface::Constructible + Quantizer, -{ - let imp = spherical_iface::Impl::::new(quantizer)?; - diskann_quantization::poly!(Quantizer, imp, GlobalAllocator) -} - impl Benchmark for BfTreeSpherical { - type Input = BfTreeSphericalBuild; + type Input = BfTreeBuild; type Output = BuildResult; - fn try_match(&self, input: &BfTreeSphericalBuild) -> Result { + fn try_match(&self, input: &BfTreeBuild) -> Result { let mut failure_score: Option = None; - if let Err(s) = utils::match_data_type::(input.data_type()) { - failure_score = Some(s.0); - } - if !matches!(input.num_bits().get(), 1 | 2 | 4) { + if !matches!(input.quantization(), QuantConfig::Spherical { .. }) { *failure_score.get_or_insert(0) += 1; } + if let Err(s) = utils::match_data_type::(input.data_type()) { + *failure_score.get_or_insert(0) += s.0; + } if !self.search.is_match(input.search_phase()) { *failure_score.get_or_insert(0) += 1; } @@ -101,7 +86,7 @@ impl Benchmark for BfTreeSpherical { fn description( &self, f: &mut std::fmt::Formatter<'_>, - input: Option<&BfTreeSphericalBuild>, + input: Option<&BfTreeBuild>, ) -> std::fmt::Result { match input { None => { @@ -136,7 +121,7 @@ impl Benchmark for BfTreeSpherical { fn run( &self, - input: &BfTreeSphericalBuild, + input: &BfTreeBuild, checkpoint: Checkpoint<'_>, mut output: &mut dyn Output, ) -> anyhow::Result { @@ -148,33 +133,17 @@ impl Benchmark for BfTreeSpherical { // 1. Train the spherical quantizer. let start = std::time::Instant::now(); - let m: diskann_vector::distance::Metric = build.distance().into(); - let pre_scale = match input.pre_scale() { - Some(&v) => v.try_into()?, - None => diskann_quantization::spherical::PreScale::None, - }; - - let quantizer = diskann_quantization::spherical::SphericalQuantizer::train( + let quantizer_poly = super::quantizer_util::build_quantizer( + input.quantization(), data.as_view(), - (input.transform_kind()).into(), - m.try_into()?, - pre_scale, - &mut rand::rngs::StdRng::seed_from_u64(input.seed()), - GlobalAllocator, - )?; + build.distance(), + )? + .expect("spherical quantization config guaranteed by try_match"); let training_time = start.elapsed(); writeln!(output, "Training time: {:.2}s", training_time.as_secs_f64())?; - // 2. Dispatch on num_bits to create the type-erased quantizer. - let quantizer_poly = match input.num_bits().get() { - 1 => new_quantizer::<1>(quantizer)?, - 2 => new_quantizer::<2>(quantizer)?, - 4 => new_quantizer::<4>(quantizer)?, - _ => unreachable!("try_match handles bit validation"), - }; - - // 3. Build the bf_tree provider with quantization. + // 2. Build the bf_tree provider with quantization. let config = input.try_as_config()?.build()?; let params = input.bftree_parameters(data.nrows(), data.ncols()); let start_points = input @@ -184,13 +153,13 @@ impl Benchmark for BfTreeSpherical { let provider = BfTreeProvider::new(params, start_points.as_view(), quantizer_poly)?; let index = Arc::new(DiskANNIndex::new(config, provider, None)); - // 4. Insert all vectors using Quantized strategy. + // 3. Insert all vectors using Quantized strategy. let build_stats = single_or_multi_insert(index.clone(), Quantized, data.clone(), build, output)?; checkpoint.checkpoint(&build_stats)?; - // 5. Search using Quantized strategy. + // 4. Search using Quantized strategy. let search_results = self.search .run(index, input.search_phase(), &Strategy::new(Quantized))?; diff --git a/diskann-benchmark/src/index/bftree/spherical_streaming.rs b/diskann-benchmark/src/index/bftree/spherical_streaming.rs index 8dfe54833..31ae2fa04 100644 --- a/diskann-benchmark/src/index/bftree/spherical_streaming.rs +++ b/diskann-benchmark/src/index/bftree/spherical_streaming.rs @@ -6,144 +6,30 @@ // Streaming BfTree SQ // //////////////////////// -use std::{borrow::Cow, io::Write, num::NonZeroUsize, sync::Arc}; +use std::{io::Write, num::NonZeroUsize, sync::Arc}; -use diskann::graph::{DiskANNIndex, InplaceDeleteMethod}; -use diskann::utils::ONE; +use diskann::graph::DiskANNIndex; use diskann_benchmark_core as benchmark_core; -use diskann_benchmark_core::{recall::Rows, streaming::executors::bigann}; +use diskann_benchmark_core::streaming::executors::bigann; use diskann_benchmark_runner::{ benchmark::{FailureScore, MatchScore}, output::Output, utils::datatype::AsDataType, Benchmark, Checkpoint, }; -use diskann_bftree::{quant::QuantVectorProvider, BfTreeProvider}; +use diskann_bftree::BfTreeProvider; use diskann_providers::model::graph::provider::async_::common::Quantized; -use diskann_quantization::alloc::{AllocatorError, GlobalAllocator, Poly}; -use diskann_quantization::spherical::{ - iface::{self as spherical_iface, Quantizer}, - SphericalQuantizer, -}; -use diskann_utils::views::{Matrix, MatrixView}; -use rand::SeedableRng; use crate::{ - index::{ - build::{BuildKind, BuildStats}, - search::knn, - streaming::{ - managed::{self, Managed}, - stats::{GenericStats, StreamStats}, - ManagedStream, - }, - }, - inputs::{ - bftree::BfTreeSphericalDynamicRun, - graph_index::{SearchPhase, TopkSearchPhase}, + index::streaming::{ + managed::{self, Managed}, + stats::StreamStats, + BfTreeMaintainer, StreamRunner, }, - utils::{self, datafiles}, + inputs::bftree::{BfTreeStreamingRun, QuantConfig}, + utils, }; -type BfTreeSQProvider = BfTreeProvider; -type BfTreeSQIndex = Arc>; - -fn new_quantizer( - quantizer: SphericalQuantizer, -) -> Result, AllocatorError> -where - spherical_iface::Impl: spherical_iface::Constructible + Quantizer, -{ - let imp = spherical_iface::Impl::::new(quantizer)?; - diskann_quantization::poly!(Quantizer, imp, GlobalAllocator) -} - -struct BfTreeSQStream { - index: BfTreeSQIndex, - search: TopkSearchPhase, - runtime: tokio::runtime::Runtime, - ntasks: NonZeroUsize, - inplace_delete_num_to_replace: usize, - inplace_delete_method: InplaceDeleteMethod, -} - -impl BfTreeSQStream { - fn insert_(&self, data: MatrixView<'_, f32>, slots: &[u32]) -> anyhow::Result { - let runner = benchmark_core::build::graph::SingleInsert::new( - self.index.clone(), - Arc::new(data.to_owned()), - Quantized, - benchmark_core::build::ids::Slice::new(slots.into()), - ); - - let results = benchmark_core::build::build( - runner, - benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - BuildStats::new(BuildKind::SingleInsert, results) - } -} - -impl ManagedStream for BfTreeSQStream { - type Output = StreamStats; - - fn search( - &self, - queries: Arc>, - groundtruth: &dyn Rows, - ) -> anyhow::Result { - let knn = benchmark_core::search::graph::KNN::new( - self.index.clone(), - queries, - benchmark_core::search::graph::Strategy::broadcast(Quantized), - )?; - - let steps = knn::SearchSteps::new( - self.search.reps, - &self.search.num_threads, - &self.search.runs, - ); - let results = knn::run(&knn, groundtruth, steps)?; - Ok(StreamStats::Search(results)) - } - - fn insert(&self, data: MatrixView<'_, f32>, slots: &[u32]) -> anyhow::Result { - Ok(StreamStats::Insert(self.insert_(data, slots)?)) - } - - fn replace(&self, data: MatrixView<'_, f32>, slots: &[u32]) -> anyhow::Result { - Ok(StreamStats::Replace(self.insert_(data, slots)?)) - } - - fn delete(&self, slots: &[u32]) -> anyhow::Result { - let runner = benchmark_core::streaming::graph::InplaceDelete::new( - self.index.clone(), - Quantized, - self.inplace_delete_num_to_replace, - self.inplace_delete_method, - benchmark_core::build::ids::Slice::new(slots.into()), - ); - - let results = benchmark_core::build::build( - runner, - benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - Ok(StreamStats::Delete(GenericStats::new( - Cow::Borrowed("Delete"), - results, - )?)) - } - - fn maintain(&self) -> anyhow::Result { - // bf-tree uses hard deletes — no deferred cleanup needed. - Ok(StreamStats::Maintain(Vec::new())) - } -} - /// The streaming benchmark for bf_tree spherical quantization. /// /// Dispatches `num_bits` at runtime to avoid const-generic monomorphization. @@ -156,20 +42,17 @@ impl StreamingSpherical { } impl Benchmark for StreamingSpherical { - type Input = BfTreeSphericalDynamicRun; + type Input = BfTreeStreamingRun; type Output = Vec>; fn try_match(&self, input: &Self::Input) -> Result { let mut failure_score: Option = None; - if let Err(s) = utils::match_data_type::(input.data_type()) { - failure_score = Some(s.0); - } - if !matches!(input.num_bits().get(), 1 | 2 | 4) { + if !matches!(input.quantization(), QuantConfig::Spherical { .. }) { *failure_score.get_or_insert(0) += 1; } - if !matches!(input.search_phase(), SearchPhase::Topk(_)) { - *failure_score.get_or_insert(0) += 1; + if let Err(s) = utils::match_data_type::(input.data_type()) { + *failure_score.get_or_insert(0) += s.0; } match failure_score { @@ -181,7 +64,7 @@ impl Benchmark for StreamingSpherical { fn description( &self, f: &mut std::fmt::Formatter<'_>, - input: Option<&Self::Input>, + input: Option<&BfTreeStreamingRun>, ) -> std::fmt::Result { match input { None => { @@ -198,13 +81,13 @@ impl Benchmark for StreamingSpherical { fn run( &self, - input: &Self::Input, + input: &BfTreeStreamingRun, _checkpoint: Checkpoint<'_>, mut output: &mut dyn Output, ) -> anyhow::Result { writeln!(output, "{}", input)?; - super::streaming_utils::run_streaming::( + crate::index::streaming::run_streaming::( input.runbook_params(), |max_points| bftree_sq_streaming_impl(input, max_points), output, @@ -213,75 +96,47 @@ impl Benchmark for StreamingSpherical { } fn bftree_sq_streaming_impl( - input: &BfTreeSphericalDynamicRun, + input: &BfTreeStreamingRun, max_points: usize, ) -> anyhow::Result>> { - let topk = match input.search_phase() { - SearchPhase::Topk(topk) => topk, - _ => anyhow::bail!("Only TopK is currently supported by the streaming index"), - }; - - let data = datafiles::load_dataset::(datafiles::BinFile(input.build().data()))?; - let queries = Arc::new(datafiles::load_dataset::(datafiles::BinFile( - &topk.queries, - ))?); - - // Train the spherical quantizer. - let m: diskann_vector::distance::Metric = input.build().distance().into(); - let pre_scale = match input.pre_scale() { - Some(&v) => v.try_into()?, - None => diskann_quantization::spherical::PreScale::None, - }; - - let quantizer = diskann_quantization::spherical::SphericalQuantizer::train( - data.as_view(), - (input.transform_kind()).into(), - m.try_into()?, - pre_scale, - &mut rand::rngs::StdRng::seed_from_u64(input.seed()), - GlobalAllocator, - )?; - - let quantizer_poly = match input.num_bits().get() { - 1 => new_quantizer::<1>(quantizer)?, - 2 => new_quantizer::<2>(quantizer)?, - 4 => new_quantizer::<4>(quantizer)?, - _ => unreachable!("try_match handles bit validation"), - }; - - let config = input.try_as_config()?.build()?; - let params = input.bftree_parameters(max_points, data.ncols()); - let start_points = input - .build() - .start_point_strategy() - .compute(data.as_view())?; - let provider = BfTreeProvider::new(params, start_points.as_view(), quantizer_poly)?; - let index = Arc::new(DiskANNIndex::new(config, provider, None)); - - let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap(); - let managed_stream = BfTreeSQStream { - index, - search: topk.clone(), - runtime: benchmark_core::tokio::runtime(num_threads_and_tasks.get())?, - ntasks: num_threads_and_tasks, - inplace_delete_num_to_replace: input.runbook_params().ip_delete_num_to_replace, - inplace_delete_method: input.runbook_params().ip_delete_method.into(), - }; + let search = input.search(); let num_start_points = input.build().start_point_strategy().count(); - let managed = Managed::new( - max_points + num_start_points, - managed::SlotReclaim::Immediate, - managed_stream, - ); - - let max_k = topk.max_k(); - let layered = bigann::WithData::new(managed, data, queries, move |path| { - Ok(Box::new(datafiles::load_groundtruth( - datafiles::BinFile(path), - Some(max_k), - )?)) - }); - - Ok(layered) + let capacity = max_points + num_start_points; + + crate::index::streaming::build_streamer( + input.build().data(), + search, + crate::index::streaming::managed::SlotReclaim::Immediate, + capacity, + |data, capacity| { + let quantizer_poly = super::quantizer_util::build_quantizer( + input.quantization(), + data.as_view(), + input.build().distance(), + )? + .expect("spherical quantization config guaranteed by try_match"); + + let config = input.try_as_config()?.build()?; + let params = input.bftree_parameters(capacity, data.ncols()); + let start_points = input + .build() + .start_point_strategy() + .compute(data.as_view())?; + let provider = BfTreeProvider::new(params, start_points.as_view(), quantizer_poly)?; + let index = Arc::new(DiskANNIndex::new(config, provider, None)); + + let num_threads_and_tasks = NonZeroUsize::new(input.build().num_threads()).unwrap(); + Ok(StreamRunner::new( + index, + Quantized, + search.clone(), + benchmark_core::tokio::runtime(num_threads_and_tasks.get())?, + num_threads_and_tasks, + input.runbook_params().ip_delete_num_to_replace, + input.runbook_params().ip_delete_method.into(), + BfTreeMaintainer, + )) + }, + ) } diff --git a/diskann-benchmark/src/index/bftree/streaming_utils.rs b/diskann-benchmark/src/index/bftree/streaming_utils.rs deleted file mode 100644 index d505eeae7..000000000 --- a/diskann-benchmark/src/index/bftree/streaming_utils.rs +++ /dev/null @@ -1,80 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::io::Write; - -use diskann_benchmark_core::streaming::{executors::bigann, Executor}; -use diskann_benchmark_runner::output::Output; - -use crate::index::streaming::{ - managed::{self, Managed}, - stats::StreamStats, -}; -use crate::inputs::graph_index::DynamicRunbookParams; - -/// Run a streaming benchmark using the given runbook parameters. -/// -/// `make_streamer` receives `max_points` from the loaded runbook and returns the -/// constructed streamer. This is shared between full-precision and spherical streaming -/// benchmarks to avoid duplicating the runbook load → run_with → stage banner → summary logic. -pub(super) fn run_streaming( - runbook_params: &DynamicRunbookParams, - make_streamer: F, - mut output: &mut dyn Output, -) -> anyhow::Result>> -where - T: 'static, - F: FnOnce(usize) -> anyhow::Result>>, -{ - let groundtruth_directory = runbook_params - .resolved_gt_directory - .as_ref() - .ok_or_else(|| { - anyhow::anyhow!("Ground truth directory path was not resolved during validation") - })?; - - let mut runbook = bigann::RunBook::load( - &runbook_params.runbook_path, - &runbook_params.dataset_name, - &mut bigann::ScanDirectory::new(groundtruth_directory)?, - )?; - - let mut streamer = make_streamer(runbook.max_points())?; - - let mut results = Vec::new(); - let stages = runbook.len(); - let mut i = 1; - - runbook.run_with( - &mut streamer, - |o: managed::Stats| -> anyhow::Result<()> { - if o.inner().is_maintain() { - let message = format!("Ran maintenance before stage {}", i); - write!(output, "{}", crate::utils::SmallBanner(&message))?; - } else { - let message = format!("Finished stage {} of {}: {}", i, stages, o.inner().kind()); - write!(output, "{}", crate::utils::SmallBanner(&message))?; - i += 1; - } - writeln!(output, "{}", o)?; - results.push(o); - Ok(()) - }, - )?; - - write!( - output, - "{}", - crate::utils::SmallBanner("End of Run Summary") - )?; - - writeln!( - output, - "{}", - crate::index::streaming::stats::Summary::new(results.iter().map(|r| r.inner())) - )?; - - Ok(results) -} diff --git a/diskann-benchmark/src/index/build.rs b/diskann-benchmark/src/index/build.rs index 2e2fbd84d..b624c9b70 100644 --- a/diskann-benchmark/src/index/build.rs +++ b/diskann-benchmark/src/index/build.rs @@ -34,7 +34,7 @@ use crate::inputs::graph_index::IndexBuild; // Start Point Configuration // /////////////////////////////// -pub(crate) fn set_start_points( +pub(super) fn set_start_points( provider: &DP, data: MatrixView<'_, T>, start_strategy: StartPointStrategy, @@ -55,7 +55,7 @@ where // Build // /////////// -pub(crate) fn single_or_multi_insert( +pub(super) fn single_or_multi_insert( index: Arc>, strategy: S, data: Arc>, @@ -111,7 +111,7 @@ where } #[cfg(any(feature = "scalar-quantization", feature = "spherical-quantization"))] -pub(crate) fn only_single_insert( +pub(super) fn only_single_insert( index: Arc>, strategy: S, data: Arc>, @@ -155,7 +155,7 @@ where #[derive(Debug, Clone, Copy, Serialize)] #[serde(rename = "kebab-case")] -pub(crate) enum BuildKind { +pub(super) enum BuildKind { SingleInsert, MultiInsert, } @@ -170,11 +170,11 @@ impl std::fmt::Display for BuildKind { } #[derive(Debug, Serialize)] -pub(crate) struct BuildStats { - pub(crate) kind: BuildKind, - pub(crate) total_time: MicroSeconds, - pub(crate) vectors_inserted: usize, - pub(crate) insert_latencies: percentiles::Percentiles, +pub(super) struct BuildStats { + pub(super) kind: BuildKind, + pub(super) total_time: MicroSeconds, + pub(super) vectors_inserted: usize, + pub(super) insert_latencies: percentiles::Percentiles, } impl BuildStats { @@ -253,7 +253,7 @@ impl build_core::Progress for Meter { // Save and Load API /// //////////////////////// -pub(crate) async fn save_index( +pub(super) async fn save_index( index: Arc>, save_path: &str, ) -> anyhow::Result<()> @@ -273,7 +273,7 @@ where } // for now, this only works with full-precision indices -pub(crate) async fn load_index<'a, DP>( +pub(super) async fn load_index<'a, DP>( load_path: &'a str, index_config: &IndexConfiguration, ) -> anyhow::Result> diff --git a/diskann-benchmark/src/index/inmem/mod.rs b/diskann-benchmark/src/index/inmem/mod.rs index bb64f6cdf..f979fc107 100644 --- a/diskann-benchmark/src/index/inmem/mod.rs +++ b/diskann-benchmark/src/index/inmem/mod.rs @@ -6,3 +6,6 @@ pub(crate) mod product; pub(crate) mod scalar; pub(crate) mod spherical; +pub(crate) mod streaming; + +pub(crate) use streaming::InmemMaintainer; diff --git a/diskann-benchmark/src/index/inmem/product.rs b/diskann-benchmark/src/index/inmem/product.rs index c26e5838b..44c11b8b8 100644 --- a/diskann-benchmark/src/index/inmem/product.rs +++ b/diskann-benchmark/src/index/inmem/product.rs @@ -90,7 +90,7 @@ mod imp { /// types. /// /// The kinds of quantized and full-precision searches are kept in-sync. - pub(crate) struct ProductQuantized + pub(super) struct ProductQuantized where T: VectorRepr, { @@ -102,14 +102,14 @@ mod imp { where T: VectorRepr, { - pub(crate) fn new() -> Self { + pub(super) fn new() -> Self { Self { quant_search: plugins::Plugins::new(), full_search: plugins::Plugins::new(), } } - pub(crate) fn search

(mut self, plugin: P) -> Self + pub(super) fn search

(mut self, plugin: P) -> Self where P: plugins::Plugin, SearchPhase, Strategy> + plugins::Plugin, SearchPhase, Strategy> @@ -133,10 +133,10 @@ mod imp { type Output = QuantBuildResult; fn try_match(&self, input: &IndexPQOperation) -> Result { - let score = utils::match_data_type::(*input.index_operation.source.data_type()); + let score = utils::match_data_type::(*input.index_operation().source().data_type()); if self .quant_search - .is_match(&input.index_operation.search_phase) + .is_match(input.index_operation().search_phase()) { score } else { @@ -154,19 +154,19 @@ mod imp { ) -> std::fmt::Result { match input { Some(arg) => { - let desc = T::describe(*arg.index_operation.source.data_type()); + let desc = T::describe(*arg.index_operation().source().data_type()); if !desc.is_match() { writeln!(f, "Data/Query Type: {}", desc,)?; } if !self .quant_search - .is_match(&arg.index_operation.search_phase) + .is_match(arg.index_operation().search_phase()) { writeln!( f, "Unsupported search phase: \"{}\" - expected one of {}", - arg.index_operation.search_phase.kind(), + arg.index_operation().search_phase().kind(), self.quant_search.format_kinds(), )?; } @@ -188,9 +188,9 @@ mod imp { ) -> anyhow::Result { writeln!(output, "{}", input)?; - let hybrid = common::Hybrid::new(input.max_fp_vecs_per_prune); + let hybrid = common::Hybrid::new(input.max_fp_vecs_per_prune()); - let (index, build_stats, quant_training_time) = match &input.index_operation.source { + let (index, build_stats, quant_training_time) = match input.index_operation().source() { IndexSource::Load(load) => { let index_config: &IndexConfiguration = &input.to_config()?; @@ -213,8 +213,8 @@ mod imp { diskann_async::train_pq( train_data.as_view(), - input.num_pq_chunks, - &mut StdRng::seed_from_u64(input.seed), + input.num_pq_chunks(), + &mut StdRng::seed_from_u64(input.seed()), diskann_providers::utils::create_thread_pool(build.num_threads())? .as_ref(), )? @@ -257,8 +257,8 @@ mod imp { // Save construction stats before running queries. checkpoint.checkpoint(&build_stats)?; - let search_phase = &input.index_operation.search_phase; - let search = if input.use_fp_for_search { + let search_phase = input.index_operation().search_phase(); + let search = if input.use_fp_for_search() { self.full_search .run(index, search_phase, &Strategy::new(common::FullPrecision))? } else { diff --git a/diskann-benchmark/src/index/inmem/scalar.rs b/diskann-benchmark/src/index/inmem/scalar.rs index ab26638dd..569959c1b 100644 --- a/diskann-benchmark/src/index/inmem/scalar.rs +++ b/diskann-benchmark/src/index/inmem/scalar.rs @@ -100,7 +100,7 @@ mod imp { /// types. /// /// The kinds of quantized and full-precision searches are kept in-sync. - pub(crate) struct ScalarQuantized + pub(super) struct ScalarQuantized where T: VectorRepr, { @@ -114,14 +114,14 @@ mod imp { where T: VectorRepr, { - pub(crate) fn new() -> Self { + pub(super) fn new() -> Self { Self { quant_search: plugins::Plugins::new(), full_search: plugins::Plugins::new(), } } - pub(crate) fn search

(mut self, plugin: P) -> Self + pub(super) fn search

(mut self, plugin: P) -> Self where P: plugins::Plugin, SearchPhase, Strategy> + plugins::Plugin, SearchPhase, Strategy> @@ -142,25 +142,25 @@ mod imp { fn try_match(&self, input: &IndexSQOperation) -> Result { let mut failure_score: Option = None; - match input.index_operation.source { + match input.index_operation().source() { IndexSource::Load(_) => {} - IndexSource::Build(ref build) => { + IndexSource::Build(build) => { if build.multi_insert().is_some() { failure_score = Some(1); } } } - if !<$T>::is_match(*input.index_operation.source.data_type()) { + if !<$T>::is_match(*input.index_operation().source().data_type()) { *failure_score.get_or_insert(0) += 1; } - if !self.quant_search.is_match(&input.index_operation.search_phase) { + if !self.quant_search.is_match(input.index_operation().search_phase()) { *failure_score.get_or_insert(0) += 1; } - if input.num_bits != $N { - *failure_score.get_or_insert(0) += 10 + ($N as usize).abs_diff(input.num_bits) as u32; + if input.num_bits() != $N { + *failure_score.get_or_insert(0) += 10 + ($N as usize).abs_diff(input.num_bits()) as u32; } match failure_score { @@ -191,16 +191,16 @@ mod imp { writeln!(f, "- Search Kinds: {}", self.quant_search.format_kinds())?; } Some(input) => { - if input.num_bits != $N { + if input.num_bits() != $N { writeln!( f, "- Expected {} bits, instead got {}", $N, - input.num_bits + input.num_bits() )?; } - let data_type = *input.index_operation.source.data_type(); + let data_type = *input.index_operation().source().data_type(); if !<$T>::is_match(data_type) { writeln!( f, @@ -210,7 +210,7 @@ mod imp { )?; } - if let IndexSource::Build(ref build) = input.index_operation.source { + if let IndexSource::Build(build) = input.index_operation().source() { if build.multi_insert().is_some() { writeln!( f, @@ -219,11 +219,11 @@ mod imp { } } - if !self.quant_search.is_match(&input.index_operation.search_phase) { + if !self.quant_search.is_match(input.index_operation().search_phase()) { writeln!( f, "- Unsupported search phase: \"{}\" - expected one of {}", - input.index_operation.search_phase.kind(), + input.index_operation().search_phase().kind(), self.quant_search.format_kinds(), )?; } @@ -239,14 +239,14 @@ mod imp { mut output: &mut dyn Output, ) -> anyhow::Result { assert_eq!( - input.num_bits, + input.num_bits(), $N, "INTERNAL ERROR: this should not have passed the match check" ); writeln!(output, "{}", input)?; - let (index, build_stats, quant_training_time) = match &input.index_operation.source { + let (index, build_stats, quant_training_time) = match input.index_operation().source() { IndexSource::Load(load) => { let index_config: &IndexConfiguration = &load.to_config()?; @@ -263,7 +263,7 @@ mod imp { let start = std::time::Instant::now(); let quantizer = diskann_quantization::scalar::train::ScalarQuantizationParameters::new( - diskann_quantization::num::Positive::new(input.standard_deviations).context( + diskann_quantization::num::Positive::new(input.standard_deviations()).context( "please file a bug report, this should not have made it past the\ front end", )?, @@ -305,16 +305,16 @@ mod imp { // Save construction stats before running queries. checkpoint.checkpoint(&build_stats)?; - let search = if input.use_fp_for_search { + let search = if input.use_fp_for_search() { self.full_search.run( index, - &input.index_operation.search_phase, + input.index_operation().search_phase(), &Strategy::new(common::FullPrecision), )? } else { self.quant_search.run( index, - &input.index_operation.search_phase, + input.index_operation().search_phase(), &Strategy::new(common::Quantized), )? }; diff --git a/diskann-benchmark/src/index/inmem/spherical.rs b/diskann-benchmark/src/index/inmem/spherical.rs index 7a9994648..c62a83b1c 100644 --- a/diskann-benchmark/src/index/inmem/spherical.rs +++ b/diskann-benchmark/src/index/inmem/spherical.rs @@ -112,18 +112,18 @@ mod imp { /// A [`Benchmark`] for spherical-quantized searches containing a dynamic list of search /// types. - pub(crate) struct SphericalQ { + pub(super) struct SphericalQ { search: search::plugins::Plugins, } impl SphericalQ { - pub(crate) fn new() -> Self { + pub(super) fn new() -> Self { Self { search: search::plugins::Plugins::new(), } } - pub(crate) fn search

(mut self, plugin: P) -> Self + pub(super) fn search

(mut self, plugin: P) -> Self where P: search::plugins::Plugin + 'static, @@ -190,19 +190,19 @@ mod imp { input: &SphericalQuantBuild, ) -> Result { let mut failure_score: Option = None; - if input.build.multi_insert().is_some() { + if input.build().multi_insert().is_some() { failure_score = Some(1); } - if !f32::is_match(input.build.data_type()) { + if !f32::is_match(input.build().data_type()) { *failure_score.get_or_insert(0) += 1; } - if !self.search.is_match(&input.search_phase) { + if !self.search.is_match(input.search_phase()) { *failure_score.get_or_insert(0) += 1; } - let num_bits = input.num_bits.get(); + let num_bits = input.num_bits().get(); if num_bits != $N { *failure_score.get_or_insert(0) += ($N as usize) .abs_diff(num_bits) @@ -234,31 +234,31 @@ mod imp { writeln!(f, "- Search Kinds: {}", self.search.format_kinds())?; } Some(input) => { - let num_bits = input.num_bits.get(); + let num_bits = input.num_bits().get(); if num_bits != $N { writeln!(f, "- Expected {} bits, got {}", $N, num_bits)?; } - if input.build.multi_insert().is_some() { + if input.build().multi_insert().is_some() { writeln!( f, "- Spherical Quantization does not support multi-insert" )?; } - if !f32::is_match(input.build.data_type()) { + if !f32::is_match(input.build().data_type()) { writeln!( f, "- Only `float32` data type is supported. Instead, got {}", - input.build.data_type() + input.build().data_type() )?; } - if !self.search.is_match(&input.search_phase) { + if !self.search.is_match(input.search_phase()) { writeln!( f, "- Unsupported search phase: \"{}\" - expected one of {}", - input.search_phase.kind(), + input.search_phase().kind(), self.search.format_kinds(), )?; } @@ -274,31 +274,31 @@ mod imp { mut output: &mut dyn Output, ) -> anyhow::Result { assert_eq!( - input.num_bits.get(), + input.num_bits().get(), $N, "INTERNAL ERROR: this should not have passed the match check" ); writeln!(output, "{}", input)?; - let build = &input.build; + let build = input.build(); let data: Arc> = Arc::new(datafiles::load_dataset(datafiles::BinFile(build.data()))?); let start = std::time::Instant::now(); let m: diskann_vector::distance::Metric = build.distance().into(); - let pre_scale = match input.pre_scale { - Some(v) => v.try_into()?, + let pre_scale = match input.pre_scale() { + Some(&v) => v.try_into()?, None => diskann_quantization::spherical::PreScale::None, }; let quantizer = diskann_quantization::spherical::SphericalQuantizer::train( data.as_view(), - (&input.transform_kind).into(), + (input.transform_kind()).into(), m.try_into()?, pre_scale, - &mut rand::rngs::StdRng::seed_from_u64(input.seed), + &mut rand::rngs::StdRng::seed_from_u64(input.seed()), GlobalAllocator, )?; @@ -340,10 +340,10 @@ mod imp { // Save construction stats before running queries. checkpoint.checkpoint(&result)?; - for layout in input.query_layouts.iter() { - let search = self - .search - .run(index.clone(), &input.search_phase, layout)?; + for layout in input.query_layouts().iter() { + let search = + self.search + .run(index.clone(), input.search_phase(), layout)?; result.append(SearchRun { layout: *layout, diff --git a/diskann-benchmark/src/index/inmem/streaming.rs b/diskann-benchmark/src/index/inmem/streaming.rs new file mode 100644 index 000000000..8153d9254 --- /dev/null +++ b/diskann-benchmark/src/index/inmem/streaming.rs @@ -0,0 +1,122 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{borrow::Cow, num::NonZeroUsize, sync::Arc}; + +use diskann::{ + graph::DiskANNIndex, + provider::{self, Delete}, + utils::ONE, + ANNError, ANNErrorKind, ANNResult, +}; +use diskann_benchmark_core::{ + build::{self, ids::Range, Parallelism}, + streaming::graph::DropDeleted, +}; +use diskann_providers::model::graph::provider::async_::{ + inmem::DefaultProvider, TableDeleteProviderAsync, +}; +use diskann_utils::future::AsyncFriendly; +use tokio::runtime::Runtime; + +use crate::index::streaming::{ + runner::Maintainer, + stats::{GenericStats, StreamStats}, +}; + +///////////// +// Helpers // +///////////// + +struct Release +where + U: AsyncFriendly, + V: AsyncFriendly, +{ + index: Arc>>, +} + +impl Release +where + U: AsyncFriendly, + V: AsyncFriendly, +{ + fn new(index: Arc>>) -> Arc { + Arc::new(Self { index }) + } +} + +/// NOTE: The implementation here strictly targets the implementation of [`provider::Delete`] +/// for [`DefaultProvider`] with the [`TableDeleteProviderAsync`] delete provider. +/// +/// Trying to make this generic over the delete provider is not a recipe for a good time. +impl diskann_benchmark_core::build::Build for Release +where + U: AsyncFriendly, + V: AsyncFriendly, +{ + type Output = (); + + fn num_data(&self) -> usize { + self.index.provider().total_points() + } + + async fn build(&self, range: std::ops::Range) -> ANNResult<()> { + let provider = self.index.provider(); + let ctx = &provider::DefaultContext; + + for internal_id in range { + let internal_id: u32 = internal_id + .try_into() + .map_err(|_| ANNError::message(ANNErrorKind::Opaque, "invalid id provided"))?; + if provider + .status_by_external_id(ctx, &internal_id) + .await? + .is_deleted() + { + provider.release(ctx, internal_id).await?; + } + } + Ok(()) + } +} + +////////////////////// +// Inmem Maintainer // +////////////////////// + +/// Inmem maintenance: runs `DropDeleted` to unlink deleted neighbors, then `Release` +/// to free deleted slots for reuse. +pub(crate) struct InmemMaintainer; + +impl Maintainer> for InmemMaintainer +where + U: AsyncFriendly, + V: AsyncFriendly, +{ + fn maintain( + &self, + index: &Arc>>, + runtime: &Runtime, + ntasks: NonZeroUsize, + ) -> anyhow::Result { + let range = index.provider().iter(); + + let runner = DropDeleted::new(index.clone(), false, Range::new(range)); + + let drop_deleted = build::build(runner, Parallelism::fixed(Some(ONE), ntasks), runtime)?; + + let release = build::build( + Release::new(index.clone()), + Parallelism::fixed(Some(ONE), ntasks), + runtime, + )?; + + Ok(StreamStats::Maintain(vec![ + GenericStats::new(Cow::Borrowed("Drop Deleted"), drop_deleted)?, + GenericStats::new(Cow::Borrowed("Release"), release)?, + ])) + } +} diff --git a/diskann-benchmark/src/index/result.rs b/diskann-benchmark/src/index/result.rs index 8bd90b33e..d12ffd2b9 100644 --- a/diskann-benchmark/src/index/result.rs +++ b/diskann-benchmark/src/index/result.rs @@ -16,13 +16,13 @@ use crate::{ // BuildResult // ////////////////// #[derive(Debug, Serialize)] -pub(crate) struct BuildResult { - pub(crate) build: Option, - pub(crate) search: AggregatedSearchResults, +pub(super) struct BuildResult { + pub(super) build: Option, + pub(super) search: AggregatedSearchResults, } impl BuildResult { - pub(crate) fn new(build: Option, search: AggregatedSearchResults) -> Self { + pub(super) fn new(build: Option, search: AggregatedSearchResults) -> Self { Self { build, search } } } @@ -45,9 +45,9 @@ impl std::fmt::Display for BuildResult { #[cfg(any(feature = "product-quantization", feature = "scalar-quantization",))] #[derive(Debug, Serialize)] -pub(crate) struct QuantBuildResult { - pub(crate) quant_training_time: MicroSeconds, - pub(crate) build: BuildResult, +pub(super) struct QuantBuildResult { + pub(super) quant_training_time: MicroSeconds, + pub(super) build: BuildResult, } #[cfg(any(feature = "product-quantization", feature = "scalar-quantization",))] @@ -67,7 +67,7 @@ impl std::fmt::Display for QuantBuildResult { /////////////////// #[derive(Debug, Serialize)] -pub(crate) enum AggregatedSearchResults { +pub(super) enum AggregatedSearchResults { Topk(Vec), Range(Vec), } @@ -83,18 +83,18 @@ impl std::fmt::Display for AggregatedSearchResults { } #[derive(Debug, Serialize)] -pub(crate) struct SearchResults { - pub(crate) num_tasks: usize, - pub(crate) search_n: usize, - pub(crate) search_l: usize, - pub(crate) qps: Vec, - pub(crate) search_latencies: Vec, - pub(crate) mean_latencies: Vec, - pub(crate) p90_latencies: Vec, - pub(crate) p99_latencies: Vec, - pub(crate) recall: utils::recall::RecallMetrics, - pub(crate) mean_cmps: f32, - pub(crate) mean_hops: f32, +pub(super) struct SearchResults { + pub(super) num_tasks: usize, + pub(super) search_n: usize, + pub(super) search_l: usize, + pub(super) qps: Vec, + pub(super) search_latencies: Vec, + pub(super) mean_latencies: Vec, + pub(super) p90_latencies: Vec, + pub(super) p99_latencies: Vec, + pub(super) recall: utils::recall::RecallMetrics, + pub(super) mean_cmps: f32, + pub(super) mean_hops: f32, } impl SearchResults { @@ -229,15 +229,15 @@ impl std::fmt::Display for DisplayWrapper<'_, [SearchResults]> { //////////////////////// #[derive(Debug, Serialize)] -pub(crate) struct RangeSearchResults { - pub(crate) num_tasks: usize, - pub(crate) initial_l: usize, - pub(crate) qps: Vec, - pub(crate) search_latencies: Vec, - pub(crate) mean_latencies: Vec, - pub(crate) p90_latencies: Vec, - pub(crate) p99_latencies: Vec, - pub(crate) average_precision: utils::recall::AveragePrecisionMetrics, +pub(super) struct RangeSearchResults { + pub(super) num_tasks: usize, + pub(super) initial_l: usize, + pub(super) qps: Vec, + pub(super) search_latencies: Vec, + pub(super) mean_latencies: Vec, + pub(super) p90_latencies: Vec, + pub(super) p99_latencies: Vec, + pub(super) average_precision: utils::recall::AveragePrecisionMetrics, } impl RangeSearchResults { diff --git a/diskann-benchmark/src/index/streaming/full_precision.rs b/diskann-benchmark/src/index/streaming/full_precision.rs deleted file mode 100644 index 3f7d73e76..000000000 --- a/diskann-benchmark/src/index/streaming/full_precision.rs +++ /dev/null @@ -1,223 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::{borrow::Cow, num::NonZeroUsize, sync::Arc}; - -use diskann::{ - graph::{DiskANNIndex, InplaceDeleteMethod}, - provider::{self, Delete}, - utils::{VectorRepr, ONE}, - ANNError, ANNErrorKind, ANNResult, -}; -use diskann_benchmark_core::recall::Rows; -use diskann_providers::model::graph::provider::async_::{ - common, - inmem::{self, DefaultProvider}, - TableDeleteProviderAsync, -}; -use diskann_utils::{ - future::AsyncFriendly, - views::{Matrix, MatrixView}, -}; - -use super::{ - stats::{GenericStats, StreamStats}, - ManagedStream, -}; -use crate::{ - index::{ - build::{BuildKind, BuildStats}, - search::knn, - }, - inputs::graph_index::TopkSearchPhase, -}; - -type FullPrecisionIndex = Arc< - DiskANNIndex< - DefaultProvider, common::NoStore, TableDeleteProviderAsync>, - >, ->; - -/// Full-Precision Streaming Index Implementation. -/// -/// ## Behavior with Deletes -/// -/// Deletes are processed by using `inplace_delete` to soft-delete data. Slots deleted this way -/// are not reused until maintenance is run, which drops deleted neighbors and releases the -/// deleted data points. -pub(crate) struct FullPrecisionStream -where - T: VectorRepr, -{ - pub(crate) index: FullPrecisionIndex, - pub(crate) search: TopkSearchPhase, - pub(crate) runtime: tokio::runtime::Runtime, - pub(crate) ntasks: NonZeroUsize, - pub(crate) inplace_delete_num_to_replace: usize, - pub(crate) inplace_delete_method: InplaceDeleteMethod, -} - -impl FullPrecisionStream -where - T: VectorRepr, -{ - // Common code-path for both inserts and replace. - fn insert_(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { - let runner = diskann_benchmark_core::build::graph::SingleInsert::new( - self.index.clone(), - Arc::new(data.to_owned()), - common::FullPrecision, - diskann_benchmark_core::build::ids::Slice::new(slots.into()), - ); - - let results = diskann_benchmark_core::build::build( - runner, - diskann_benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - BuildStats::new(BuildKind::SingleInsert, results) - } -} - -impl ManagedStream for FullPrecisionStream -where - T: VectorRepr, -{ - type Output = StreamStats; - - fn search( - &self, - queries: Arc>, - groundtruth: &dyn Rows, - ) -> anyhow::Result { - let knn = diskann_benchmark_core::search::graph::KNN::new( - self.index.clone(), - queries, - diskann_benchmark_core::search::graph::Strategy::broadcast(common::FullPrecision), - )?; - - let steps = knn::SearchSteps::new( - self.search.reps, - &self.search.num_threads, - &self.search.runs, - ); - let results = knn::run(&knn, groundtruth, steps)?; - Ok(StreamStats::Search(results)) - } - - fn insert(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { - Ok(StreamStats::Insert(self.insert_(data, slots)?)) - } - - fn replace(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { - Ok(StreamStats::Replace(self.insert_(data, slots)?)) - } - - fn delete(&self, slots: &[u32]) -> anyhow::Result { - let runner = diskann_benchmark_core::streaming::graph::InplaceDelete::new( - self.index.clone(), - common::FullPrecision, - self.inplace_delete_num_to_replace, - self.inplace_delete_method, - diskann_benchmark_core::build::ids::Slice::new(slots.into()), - ); - - let results = diskann_benchmark_core::build::build( - runner, - diskann_benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - Ok(StreamStats::Delete(GenericStats::new( - Cow::Borrowed("Delete"), - results, - )?)) - } - - fn maintain(&self) -> anyhow::Result { - let range = self.index.provider().iter(); - - let runner = diskann_benchmark_core::streaming::graph::DropDeleted::new( - self.index.clone(), - false, - diskann_benchmark_core::build::ids::Range::new(range), - ); - - let drop_deleted = diskann_benchmark_core::build::build( - runner, - diskann_benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - let release = diskann_benchmark_core::build::build( - Release::new(self.index.clone()), - diskann_benchmark_core::build::Parallelism::fixed(Some(ONE), self.ntasks), - &self.runtime, - )?; - - Ok(StreamStats::Maintain(vec![ - GenericStats::new(Cow::Borrowed("Drop Deleted"), drop_deleted)?, - GenericStats::new(Cow::Borrowed("Release"), release)?, - ])) - } -} - -///////////// -// Helpers // -///////////// - -struct Release -where - U: AsyncFriendly, - V: AsyncFriendly, -{ - index: Arc>>, -} - -impl Release -where - U: AsyncFriendly, - V: AsyncFriendly, -{ - fn new(index: Arc>>) -> Arc { - Arc::new(Self { index }) - } -} - -/// NOTE: The implementation here strictly targets the implementation of [`provider::Delete`] -/// for [`DefaultProvider`] with the [`TableDeleteProviderAsync`] delete provider. -/// -/// Trying to make this generic over the delete provider is not a recipe for a good time. -impl diskann_benchmark_core::build::Build for Release -where - U: AsyncFriendly, - V: AsyncFriendly, -{ - type Output = (); - - fn num_data(&self) -> usize { - self.index.provider().total_points() - } - - async fn build(&self, range: std::ops::Range) -> ANNResult<()> { - let provider = self.index.provider(); - let ctx = &provider::DefaultContext; - - for internal_id in range { - let internal_id: u32 = internal_id - .try_into() - .map_err(|_| ANNError::message(ANNErrorKind::Opaque, "invalid id provided"))?; - if provider - .status_by_external_id(ctx, &internal_id) - .await? - .is_deleted() - { - provider.release(ctx, internal_id).await?; - } - } - Ok(()) - } -} diff --git a/diskann-benchmark/src/index/streaming/mod.rs b/diskann-benchmark/src/index/streaming/mod.rs index 3b4814549..b4e3bc834 100644 --- a/diskann-benchmark/src/index/streaming/mod.rs +++ b/diskann-benchmark/src/index/streaming/mod.rs @@ -3,9 +3,126 @@ * Licensed under the MIT license. */ -pub(crate) mod full_precision; +use std::{io::Write, sync::Arc}; + +use diskann::utils::VectorRepr; +use diskann_benchmark_core::streaming::{executors::bigann, Executor}; +use diskann_benchmark_runner::output::Output; +use diskann_utils::views::Matrix; + pub(crate) mod managed; +pub(crate) mod runner; pub(crate) mod stats; -pub(crate) use full_precision::FullPrecisionStream; pub(crate) use managed::{Managed, ManagedStream}; +#[cfg(feature = "bftree")] +pub(crate) use runner::BfTreeMaintainer; +pub(crate) use runner::StreamRunner; + +use crate::{ + inputs::graph_index::{StreamingRunbookParams, StreamingSearchParams}, + utils::datafiles, +}; + +/// Construct the streaming stack: load data/queries, create the managed stream via the +/// closure, then wrap in [`Managed`] and [`bigann::WithData`]. +/// +/// `capacity` is the pre-computed slot count passed to [`Managed::new`]. Each backend +/// computes this differently (inmem applies headroom, bf_tree adds start points). +/// +/// The closure receives `(&data, capacity)` so it can use `data.ncols()` for provider params. +pub(crate) fn build_streamer( + data_path: &diskann_benchmark_runner::files::InputFile, + search: &StreamingSearchParams, + reclaim: managed::SlotReclaim, + capacity: usize, + make_stream: F, +) -> anyhow::Result>> +where + T: bytemuck::Pod + VectorRepr + 'static, + M: ManagedStream + 'static, + F: FnOnce(&Matrix, usize) -> anyhow::Result, +{ + let data = datafiles::load_dataset::(datafiles::BinFile(data_path))?; + let queries = Arc::new(datafiles::load_dataset::(datafiles::BinFile( + &search.queries, + ))?); + + let managed_stream = make_stream(&data, capacity)?; + let managed = Managed::new(capacity, reclaim, managed_stream); + + let max_k = search.max_k(); + let layered = bigann::WithData::new(managed, data, queries, move |path| { + Ok(Box::new(datafiles::load_groundtruth( + datafiles::BinFile(path), + Some(max_k), + )?)) + }); + + Ok(layered) +} + +/// Run a streaming benchmark using the given runbook parameters. +/// +/// `make_streamer` receives `max_points` from the loaded runbook and returns the +/// constructed streamer. This is shared across all streaming benchmarks (inmem, bftree) +/// to avoid duplicating the runbook load → run_with → stage banner → summary logic. +pub(crate) fn run_streaming( + runbook_params: &StreamingRunbookParams, + make_streamer: F, + mut output: &mut dyn Output, +) -> anyhow::Result>> +where + T: 'static, + F: FnOnce(usize) -> anyhow::Result>>, +{ + let groundtruth_directory = runbook_params + .resolved_gt_directory + .as_ref() + .ok_or_else(|| { + anyhow::anyhow!("Ground truth directory path was not resolved during validation") + })?; + + let mut runbook = bigann::RunBook::load( + &runbook_params.runbook_path, + &runbook_params.dataset_name, + &mut bigann::ScanDirectory::new(groundtruth_directory)?, + )?; + + let mut streamer = make_streamer(runbook.max_points())?; + + let mut results = Vec::new(); + let stages = runbook.len(); + let mut i = 1; + + runbook.run_with( + &mut streamer, + |o: managed::Stats| -> anyhow::Result<()> { + if o.inner().is_maintain() { + let message = format!("Ran maintenance before stage {}", i); + write!(output, "{}", crate::utils::SmallBanner(&message))?; + } else { + let message = format!("Finished stage {} of {}: {}", i, stages, o.inner().kind()); + write!(output, "{}", crate::utils::SmallBanner(&message))?; + i += 1; + } + writeln!(output, "{}", o)?; + results.push(o); + Ok(()) + }, + )?; + + write!( + output, + "{}", + crate::utils::SmallBanner("End of Run Summary") + )?; + + writeln!( + output, + "{}", + stats::Summary::new(results.iter().map(|r| r.inner())) + )?; + + Ok(results) +} diff --git a/diskann-benchmark/src/index/streaming/runner.rs b/diskann-benchmark/src/index/streaming/runner.rs new file mode 100644 index 000000000..3a271c8b6 --- /dev/null +++ b/diskann-benchmark/src/index/streaming/runner.rs @@ -0,0 +1,215 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{borrow::Cow, marker::PhantomData, num::NonZeroUsize, sync::Arc}; + +use diskann::{ + graph::{ + glue::{DefaultSearchStrategy, InplaceDeleteStrategy, InsertStrategy}, + DiskANNIndex, InplaceDeleteMethod, + }, + provider::{DataProvider, Delete, SetElement}, + utils::ONE, +}; +use diskann_benchmark_core::{ + build::{self, graph::SingleInsert, ids::Slice, Parallelism}, + recall::Rows, + search::{ + self, + graph::{Strategy, KNN}, + }, + streaming::graph::InplaceDelete, +}; +use diskann_utils::{ + future::AsyncFriendly, + views::{Matrix, MatrixView}, +}; +use tokio::runtime::Runtime; + +use crate::{ + index::{ + build::{BuildKind, BuildStats}, + search::knn, + streaming::{ + stats::{GenericStats, StreamStats}, + ManagedStream, + }, + }, + inputs::graph_index::{GraphSearch, StreamingSearchParams}, +}; + +pub(crate) trait Maintainer { + fn maintain( + &self, + index: &Arc>, + runtime: &Runtime, + ntasks: NonZeroUsize, + ) -> anyhow::Result; +} + +pub(crate) struct StreamRunner +where + DP: DataProvider, +{ + index: Arc>, + strategy: S, + search: StreamingSearchParams, + runtime: Runtime, + ntasks: NonZeroUsize, + inplace_delete_num_to_replace: usize, + inplace_delete_method: InplaceDeleteMethod, + maintainer: M, + _marker: PhantomData T>, +} + +impl StreamRunner +where + DP: DataProvider, +{ + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + index: Arc>, + strategy: S, + search: StreamingSearchParams, + runtime: Runtime, + ntasks: NonZeroUsize, + inplace_delete_num_to_replace: usize, + inplace_delete_method: InplaceDeleteMethod, + maintainer: M, + ) -> Self { + Self { + index, + strategy, + search, + runtime, + ntasks, + inplace_delete_num_to_replace, + inplace_delete_method, + maintainer, + _marker: PhantomData, + } + } +} + +impl StreamRunner +where + DP: DataProvider + for<'a> SetElement<&'a [T]>, + S: for<'a> InsertStrategy<'a, DP, &'a [T]> + Clone + AsyncFriendly, + T: AsyncFriendly + Clone, +{ + fn build(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + let runner = SingleInsert::new( + self.index.clone(), + Arc::new(data.to_owned()), + self.strategy.clone(), + Slice::new(slots.into()), + ); + + let results = build::build( + runner, + Parallelism::fixed(Some(ONE), self.ntasks), + &self.runtime, + )?; + + BuildStats::new(BuildKind::SingleInsert, results) + } +} + +impl ManagedStream for StreamRunner +where + DP: DataProvider + + for<'a> SetElement<&'a [T]> + + Delete, + DP::ExternalId: search::Id, + S: for<'a> InsertStrategy<'a, DP, &'a [T]> + + for<'a> DefaultSearchStrategy<'a, DP, &'a [T], DP::ExternalId> + + InplaceDeleteStrategy + + Clone + + AsyncFriendly, + T: AsyncFriendly + Clone, + M: Maintainer, +{ + type Output = StreamStats; + + fn search( + &self, + queries: Arc>, + groundtruth: &dyn Rows, + ) -> anyhow::Result { + let knn = KNN::new( + self.index.clone(), + queries, + Strategy::broadcast(self.strategy.clone()), + )?; + + let run = GraphSearch { + search_n: self.search.search_n, + search_l: vec![self.search.search_l], + recall_k: self.search.recall_k, + }; + let num_threads = [self.search.num_threads]; + let runs = [run]; + let steps = knn::SearchSteps::new(self.search.reps, &num_threads, &runs); + let results = knn::run(&knn, groundtruth, steps)?; + Ok(StreamStats::Search(results)) + } + + fn insert(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + Ok(StreamStats::Insert(self.build(data, slots)?)) + } + + fn replace(&self, data: MatrixView<'_, T>, slots: &[u32]) -> anyhow::Result { + Ok(StreamStats::Replace(self.build(data, slots)?)) + } + + fn delete(&self, slots: &[u32]) -> anyhow::Result { + let runner = InplaceDelete::new( + self.index.clone(), + self.strategy.clone(), + self.inplace_delete_num_to_replace, + self.inplace_delete_method, + Slice::new(slots.into()), + ); + + let results = build::build( + runner, + Parallelism::fixed(Some(ONE), self.ntasks), + &self.runtime, + )?; + + Ok(StreamStats::Delete(GenericStats::new( + Cow::Borrowed("Delete"), + results, + )?)) + } + + fn maintain(&self) -> anyhow::Result { + self.maintainer + .maintain(&self.index, &self.runtime, self.ntasks) + } +} + +///////////////////// +// BfTree Maintain // +///////////////////// + +/// BfTree maintenance: bf-tree uses hard deletes, no deferred cleanup needed. +#[cfg(feature = "bftree")] +pub(crate) struct BfTreeMaintainer; + +#[cfg(feature = "bftree")] +impl Maintainer for BfTreeMaintainer +where + DP: DataProvider, +{ + fn maintain( + &self, + _index: &Arc>, + _runtime: &Runtime, + _ntasks: NonZeroUsize, + ) -> anyhow::Result { + Ok(StreamStats::Maintain(Vec::new())) + } +} diff --git a/diskann-benchmark/src/inputs/bftree.rs b/diskann-benchmark/src/inputs/bftree.rs index da4156ccf..77f5c2c2b 100644 --- a/diskann-benchmark/src/inputs/bftree.rs +++ b/diskann-benchmark/src/inputs/bftree.rs @@ -6,7 +6,9 @@ use std::num::{NonZero, NonZeroUsize}; use crate::inputs::{ as_input, exhaustive, - graph_index::{DynamicRunbookParams, IndexBuild, SearchPhase, TopkSearchPhase}, + graph_index::{ + IndexBuild, SearchPhase, StreamingRunbookParams, StreamingSearchParams, TopkSearchPhase, + }, write_field, Example, PRINT_WIDTH, }; use diskann::graph::config; @@ -209,129 +211,65 @@ fn bftree_parameters_from( } } -as_input!(BfTreeFullPrecisionBuild); - -#[derive(Debug, Serialize, Deserialize)] -pub(crate) struct BfTreeFullPrecisionBuild { - build: IndexBuild, - search_phase: SearchPhase, - #[serde(deserialize_with = "Deserialize::deserialize")] - vector_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - neighbor_store_config: Option, -} - -impl BfTreeFullPrecisionBuild { - pub(crate) const fn tag() -> &'static str { - "graph-index-build-bftree-full-precision" - } - - pub(crate) fn try_as_config(&self) -> anyhow::Result { - // Delegate to IndexBuild's try_as_config which uses the default - // MaxDegree::Value(exact_max_degree) with 1.3x slack. The bf_tree - // neighbor pages are sized to exact_max_degree to accommodate this. - self.build.try_as_config() - } - - pub(crate) fn data_type(&self) -> DataType { - self.build.data_type() - } - - pub(crate) fn search_phase(&self) -> &SearchPhase { - &self.search_phase - } - - pub(crate) fn build(&self) -> &IndexBuild { - &self.build - } - - pub(crate) fn bftree_parameters( - &self, - num_points: usize, - dim: usize, - ) -> BfTreeProviderParameters { - bftree_parameters_from( - &self.build, - num_points, - dim, - &self.vector_store_config, - &self.neighbor_store_config, - &None, - ) - } - pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { - self.build.validate(checker)?; - self.search_phase.validate(checker)?; - if let Some(cfg) = &mut self.neighbor_store_config { - cfg.fill_defaults(); - cfg.validate()?; - } - if let Some(cfg) = &mut self.vector_store_config { - cfg.fill_defaults(); - cfg.validate()?; - } - Ok(()) - } -} - -impl Example for BfTreeFullPrecisionBuild { - fn example() -> Self { - let build = IndexBuild::example(); - - Self { - build, - search_phase: SearchPhase::Topk(TopkSearchPhase::example()), - vector_store_config: None, - neighbor_store_config: None, - } - } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "kind")] +pub(crate) enum QuantConfig { + #[serde(rename = "none")] + None, + #[serde(rename = "spherical")] + Spherical { + seed: u64, + transform_kind: exhaustive::TransformKind, + num_bits: NonZeroUsize, + #[serde(deserialize_with = "Deserialize::deserialize")] + pre_scale: Option, + #[serde(deserialize_with = "Deserialize::deserialize")] + quant_store_config: Option, + }, } -impl std::fmt::Display for BfTreeFullPrecisionBuild { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!(f, "Graph Index Bf_Tree Full-Precision Build")?; - if cfg!(not(feature = "bftree")) { - writeln!(f, "Requires the `bftree` feature")?; +impl QuantConfig { + pub(crate) fn validate(&mut self) -> anyhow::Result<()> { + match self { + Self::None => Ok(()), + Self::Spherical { + num_bits, + quant_store_config, + .. + } => { + match num_bits.get() { + 1 | 2 | 4 => {} + n => anyhow::bail!("{n} bits are not supported for spherical quantization"), + } + if let Some(cfg) = quant_store_config { + cfg.fill_defaults(); + cfg.validate()?; + } + Ok(()) + } } - write_field!(f, "tag", Self::tag())?; - - writeln!(f)?; - self.build.summarize_fields(f)?; - - if let Some(ref cfg) = self.vector_store_config { - writeln!(f, "\n Vector Store:")?; - write!(f, "{}", cfg)?; - } - if let Some(ref cfg) = self.neighbor_store_config { - writeln!(f, "\n Neighbor Store:")?; - write!(f, "{}", cfg)?; - } - - Ok(()) } } -as_input!(BfTreeDynamicRun); - #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct BfTreeDynamicRun { +pub(crate) struct BfTreeBuild { build: IndexBuild, search_phase: SearchPhase, - runbook_params: DynamicRunbookParams, + quantization: QuantConfig, #[serde(deserialize_with = "Deserialize::deserialize")] vector_store_config: Option, #[serde(deserialize_with = "Deserialize::deserialize")] neighbor_store_config: Option, } - -impl BfTreeDynamicRun { +impl BfTreeBuild { pub(crate) const fn tag() -> &'static str { - "graph-index-stream-bftree-full-precision" + "graph-index-bftree" } pub(crate) fn try_as_config(&self) -> anyhow::Result { // Delegate to IndexBuild's try_as_config which uses the default - // MaxDegree::Value(exact_max_degree) with 1.3x slack. + // MaxDegree::Value(exact_max_degree) with 1.3x slack. The bf_tree + // neighbor pages are sized to exact_max_degree to accommodate this. self.build.try_as_config() } @@ -347,201 +285,81 @@ impl BfTreeDynamicRun { &self.build } - pub(crate) fn runbook_params(&self) -> &DynamicRunbookParams { - &self.runbook_params - } - pub(crate) fn bftree_parameters( &self, num_points: usize, dim: usize, ) -> BfTreeProviderParameters { + let quant_store_config = match &self.quantization { + QuantConfig::None => &None, + QuantConfig::Spherical { + quant_store_config, .. + } => quant_store_config, + }; bftree_parameters_from( &self.build, num_points, dim, &self.vector_store_config, &self.neighbor_store_config, - &None, + quant_store_config, ) } + pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.build.validate(checker)?; self.search_phase.validate(checker)?; - self.runbook_params.validate(checker)?; - if let Some(cfg) = &mut self.vector_store_config { - cfg.fill_defaults(); - cfg.validate()?; - } if let Some(cfg) = &mut self.neighbor_store_config { cfg.fill_defaults(); cfg.validate()?; } - Ok(()) - } -} - -impl Example for BfTreeDynamicRun { - fn example() -> Self { - let build = IndexBuild::example(); - - Self { - build, - search_phase: SearchPhase::Topk(TopkSearchPhase::example()), - runbook_params: DynamicRunbookParams::example_immediate(), - vector_store_config: None, - neighbor_store_config: None, - } - } -} - -impl std::fmt::Display for BfTreeDynamicRun { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!(f, "Graph Index Bf_Tree Full-Precision Streaming")?; - if cfg!(not(feature = "bftree")) { - writeln!(f, "Requires the `bftree` feature")?; - } - write_field!(f, "tag", Self::tag())?; - - writeln!(f)?; - self.build.summarize_fields(f)?; - - if let Some(ref cfg) = self.vector_store_config { - writeln!(f, "\n Vector Store:")?; - write!(f, "{}", cfg)?; - } - if let Some(ref cfg) = self.neighbor_store_config { - writeln!(f, "\n Neighbor Store:")?; - write!(f, "{}", cfg)?; - } - - Ok(()) - } -} - -// ─── Spherical Quantization ─────────────────────────────────────────────────── - -as_input!(BfTreeSphericalBuild); - -#[derive(Debug, Serialize, Deserialize)] -pub(crate) struct BfTreeSphericalBuild { - build: IndexBuild, - search_phase: SearchPhase, - seed: u64, - transform_kind: exhaustive::TransformKind, - num_bits: NonZeroUsize, - pre_scale: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - vector_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - neighbor_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - quant_store_config: Option, -} - -impl BfTreeSphericalBuild { - pub(crate) const fn tag() -> &'static str { - "graph-index-build-bftree-spherical-quantization" - } - - pub(crate) fn try_as_config(&self) -> anyhow::Result { - self.build.try_as_config() - } - - pub(crate) fn bftree_parameters( - &self, - num_points: usize, - dim: usize, - ) -> BfTreeProviderParameters { - bftree_parameters_from( - &self.build, - num_points, - dim, - &self.vector_store_config, - &self.neighbor_store_config, - &self.quant_store_config, - ) - } - - pub(crate) fn data_type(&self) -> DataType { - self.build.data_type() - } - - pub(crate) fn search_phase(&self) -> &SearchPhase { - &self.search_phase - } - - pub(crate) fn build(&self) -> &IndexBuild { - &self.build - } - - pub(crate) fn seed(&self) -> u64 { - self.seed - } - - pub(crate) fn transform_kind(&self) -> &exhaustive::TransformKind { - &self.transform_kind - } - - pub(crate) fn num_bits(&self) -> NonZeroUsize { - self.num_bits - } - - pub(crate) fn pre_scale(&self) -> Option<&exhaustive::PreScale> { - self.pre_scale.as_ref() - } - - pub(crate) fn validate(&mut self, checker: &mut Checker) -> anyhow::Result<()> { - self.build.validate(checker)?; - self.search_phase.validate(checker)?; if let Some(cfg) = &mut self.vector_store_config { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.neighbor_store_config { - cfg.fill_defaults(); - cfg.validate()?; - } - if let Some(cfg) = &mut self.quant_store_config { - cfg.fill_defaults(); - cfg.validate()?; - } - - match self.num_bits.get() { - 1 | 2 | 4 => {} - n => anyhow::bail!("{n} bits are not supported for spherical quantization"), - } - + self.quantization.validate()?; Ok(()) } + + pub(crate) fn quantization(&self) -> &QuantConfig { + &self.quantization + } } -impl Example for BfTreeSphericalBuild { +impl Example for BfTreeBuild { fn example() -> Self { Self { build: IndexBuild::example(), search_phase: SearchPhase::Topk(TopkSearchPhase::example()), - seed: 42, - transform_kind: exhaustive::TransformKind::Null, - num_bits: NonZeroUsize::new(1).unwrap(), - pre_scale: None, + quantization: QuantConfig::None, vector_store_config: None, neighbor_store_config: None, - quant_store_config: None, } } } -impl std::fmt::Display for BfTreeSphericalBuild { +impl std::fmt::Display for BfTreeBuild { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!(f, "Graph Index Bf_Tree Spherical Quantization Build")?; + writeln!(f, "Graph Index Bf_Tree Build")?; if cfg!(not(feature = "bftree")) { writeln!(f, "Requires the `bftree` feature")?; } write_field!(f, "tag", Self::tag())?; - write_field!(f, "num_bits", self.num_bits)?; - write_field!(f, "seed", self.seed)?; - write_field!(f, "transform_kind", self.transform_kind)?; + + if let QuantConfig::Spherical { + seed, + transform_kind, + num_bits, + .. + } = &self.quantization + { + write_field!(f, "quantization", "spherical")?; + write_field!(f, "num_bits", num_bits)?; + write_field!(f, "seed", seed)?; + write_field!(f, "transform_kind", transform_kind)?; + } else { + write_field!(f, "quantization", "none")?; + } writeln!(f)?; self.build.summarize_fields(f)?; @@ -554,7 +372,11 @@ impl std::fmt::Display for BfTreeSphericalBuild { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.quant_store_config { + if let QuantConfig::Spherical { + quant_store_config: Some(ref cfg), + .. + } = self.quantization + { writeln!(f, "\n Quant Store:")?; write!(f, "{}", cfg)?; } @@ -563,30 +385,24 @@ impl std::fmt::Display for BfTreeSphericalBuild { } } -// ─── Spherical Streaming ────────────────────────────────────────────────────── - -as_input!(BfTreeSphericalDynamicRun); +as_input!(BfTreeBuild); +as_input!(BfTreeStreamingRun); #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct BfTreeSphericalDynamicRun { +pub(crate) struct BfTreeStreamingRun { build: IndexBuild, - search_phase: SearchPhase, - runbook_params: DynamicRunbookParams, - seed: u64, - transform_kind: exhaustive::TransformKind, - num_bits: NonZeroUsize, - pre_scale: Option, + search: StreamingSearchParams, + runbook_params: StreamingRunbookParams, + quantization: QuantConfig, #[serde(deserialize_with = "Deserialize::deserialize")] vector_store_config: Option, #[serde(deserialize_with = "Deserialize::deserialize")] neighbor_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - quant_store_config: Option, } -impl BfTreeSphericalDynamicRun { +impl BfTreeStreamingRun { pub(crate) const fn tag() -> &'static str { - "graph-index-stream-bftree-spherical-quantization" + "graph-index-stream-bftree" } pub(crate) fn try_as_config(&self) -> anyhow::Result { @@ -597,32 +413,20 @@ impl BfTreeSphericalDynamicRun { self.build.data_type() } - pub(crate) fn search_phase(&self) -> &SearchPhase { - &self.search_phase + pub(crate) fn search(&self) -> &StreamingSearchParams { + &self.search } pub(crate) fn build(&self) -> &IndexBuild { &self.build } - pub(crate) fn runbook_params(&self) -> &DynamicRunbookParams { + pub(crate) fn runbook_params(&self) -> &StreamingRunbookParams { &self.runbook_params } - pub(crate) fn seed(&self) -> u64 { - self.seed - } - - pub(crate) fn transform_kind(&self) -> &exhaustive::TransformKind { - &self.transform_kind - } - - pub(crate) fn num_bits(&self) -> NonZeroUsize { - self.num_bits - } - - pub(crate) fn pre_scale(&self) -> Option<&exhaustive::PreScale> { - self.pre_scale.as_ref() + pub(crate) fn quantization(&self) -> &QuantConfig { + &self.quantization } pub(crate) fn bftree_parameters( @@ -630,19 +434,25 @@ impl BfTreeSphericalDynamicRun { num_points: usize, dim: usize, ) -> BfTreeProviderParameters { + let quant_store_config = match &self.quantization { + QuantConfig::None => &None, + QuantConfig::Spherical { + quant_store_config, .. + } => quant_store_config, + }; bftree_parameters_from( &self.build, num_points, dim, &self.vector_store_config, &self.neighbor_store_config, - &self.quant_store_config, + quant_store_config, ) } pub(crate) fn validate(&mut self, checker: &mut Checker) -> anyhow::Result<()> { self.build.validate(checker)?; - self.search_phase.validate(checker)?; + self.search.validate(checker)?; self.runbook_params.validate(checker)?; if let Some(cfg) = &mut self.vector_store_config { cfg.fill_defaults(); @@ -652,47 +462,46 @@ impl BfTreeSphericalDynamicRun { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.quant_store_config { - cfg.fill_defaults(); - cfg.validate()?; - } - - match self.num_bits.get() { - 1 | 2 | 4 => {} - n => anyhow::bail!("{n} bits are not supported for spherical quantization"), - } - + self.quantization.validate()?; Ok(()) } } -impl Example for BfTreeSphericalDynamicRun { +impl Example for BfTreeStreamingRun { fn example() -> Self { Self { build: IndexBuild::example(), - search_phase: SearchPhase::Topk(TopkSearchPhase::example()), - runbook_params: DynamicRunbookParams::example_immediate(), - seed: 42, - transform_kind: exhaustive::TransformKind::Null, - num_bits: NonZeroUsize::new(1).unwrap(), - pre_scale: None, + search: StreamingSearchParams::example(), + runbook_params: StreamingRunbookParams::example_immediate(), + quantization: QuantConfig::None, vector_store_config: None, neighbor_store_config: None, - quant_store_config: None, } } } -impl std::fmt::Display for BfTreeSphericalDynamicRun { +impl std::fmt::Display for BfTreeStreamingRun { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!(f, "Graph Index Bf_Tree Spherical Quantization Streaming")?; + writeln!(f, "Graph Index Bf_Tree Streaming")?; if cfg!(not(feature = "bftree")) { writeln!(f, "Requires the `bftree` feature")?; } write_field!(f, "tag", Self::tag())?; - write_field!(f, "num_bits", self.num_bits)?; - write_field!(f, "seed", self.seed)?; - write_field!(f, "transform_kind", self.transform_kind)?; + + if let QuantConfig::Spherical { + seed, + transform_kind, + num_bits, + .. + } = &self.quantization + { + write_field!(f, "quantization", "spherical")?; + write_field!(f, "num_bits", num_bits)?; + write_field!(f, "seed", seed)?; + write_field!(f, "transform_kind", transform_kind)?; + } else { + write_field!(f, "quantization", "none")?; + } writeln!(f)?; self.build.summarize_fields(f)?; @@ -705,7 +514,11 @@ impl std::fmt::Display for BfTreeSphericalDynamicRun { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.quant_store_config { + if let QuantConfig::Spherical { + quant_store_config: Some(ref cfg), + .. + } = self.quantization + { writeln!(f, "\n Quant Store:")?; write!(f, "{}", cfg)?; } diff --git a/diskann-benchmark/src/inputs/exhaustive.rs b/diskann-benchmark/src/inputs/exhaustive.rs index 20583de85..ccd6215c3 100644 --- a/diskann-benchmark/src/inputs/exhaustive.rs +++ b/diskann-benchmark/src/inputs/exhaustive.rs @@ -152,7 +152,7 @@ impl From for diskann_quantization::algorithms::transforms::TargetDim } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub(crate) enum TransformKind { PaddingHadamard(TargetDim), diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index 0506c8083..2bd8556b8 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -35,7 +35,7 @@ as_input!(IndexOperation); as_input!(IndexPQOperation); as_input!(IndexSQOperation); as_input!(SphericalQuantBuild); -as_input!(DynamicIndexRun); +as_input!(StreamingIndexRun); //////////// // Search // @@ -218,7 +218,7 @@ impl BetaSearchPhase { } #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct MultihopFilterSearchPhase { +pub(crate) struct MultiHopSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, pub(crate) groundtruth: InputFile, @@ -229,7 +229,7 @@ pub(crate) struct MultihopFilterSearchPhase { pub(crate) runs: Vec, } -impl MultihopFilterSearchPhase { +impl MultiHopSearchPhase { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; self.query_predicates.resolve(checker)?; @@ -259,7 +259,7 @@ impl AdaptiveL { } #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct InlineFilterSearchPhase { +pub(crate) struct InlineSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, pub(crate) groundtruth: InputFile, @@ -271,7 +271,7 @@ pub(crate) struct InlineFilterSearchPhase { pub(crate) adaptive_l: Option, } -impl InlineFilterSearchPhase { +impl InlineSearchPhase { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; self.query_predicates.resolve(checker)?; @@ -407,8 +407,8 @@ pub(crate) enum SearchPhase { Topk(TopkSearchPhase), Range(RangeSearchPhase), TopkBetaFilter(BetaSearchPhase), - TopkMultihopFilter(MultihopFilterSearchPhase), - TopkInlineFilter(InlineFilterSearchPhase), + TopkMultihopFilter(MultiHopSearchPhase), + TopkInlineFilter(InlineSearchPhase), TopkDeterminantDiversity(TopkDeterminantDiversityPhase), } @@ -473,7 +473,7 @@ impl SearchPhase { pub(crate) fn as_topk_multihop_filter( &self, - ) -> Result<&MultihopFilterSearchPhase, WrongSearchPhaseKind> { + ) -> Result<&MultiHopSearchPhase, WrongSearchPhaseKind> { match self { Self::TopkMultihopFilter(phase) => Ok(phase), _ => Err(WrongSearchPhaseKind::new( @@ -483,9 +483,7 @@ impl SearchPhase { } } - pub(crate) fn as_topk_inline_filter( - &self, - ) -> Result<&InlineFilterSearchPhase, WrongSearchPhaseKind> { + pub(crate) fn as_topk_inline_filter(&self) -> Result<&InlineSearchPhase, WrongSearchPhaseKind> { match self { Self::TopkInlineFilter(phase) => Ok(phase), _ => Err(WrongSearchPhaseKind::new( @@ -886,8 +884,8 @@ impl IndexSource { #[derive(Debug, Serialize, Deserialize)] pub(crate) struct IndexOperation { - pub(crate) source: IndexSource, // either load or build - pub(crate) search_phase: SearchPhase, + source: IndexSource, // either load or build + search_phase: SearchPhase, } impl IndexOperation { @@ -895,6 +893,14 @@ impl IndexOperation { "graph-index-build" } + pub(crate) fn source(&self) -> &IndexSource { + &self.source + } + + pub(crate) fn search_phase(&self) -> &SearchPhase { + &self.search_phase + } + pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.source.validate(checker)?; self.search_phase.validate(checker)?; @@ -928,11 +934,11 @@ impl std::fmt::Display for IndexOperation { #[derive(Debug, Serialize, Deserialize)] pub(crate) struct IndexPQOperation { - pub(crate) index_operation: IndexOperation, // either load or build - pub(crate) num_pq_chunks: usize, - pub(crate) seed: u64, - pub(crate) max_fp_vecs_per_prune: Option, - pub(crate) use_fp_for_search: bool, + index_operation: IndexOperation, // either load or build + num_pq_chunks: usize, + seed: u64, + max_fp_vecs_per_prune: Option, + use_fp_for_search: bool, } impl IndexPQOperation { @@ -940,6 +946,31 @@ impl IndexPQOperation { "graph-index-build-pq" } + #[cfg(feature = "product-quantization")] + pub(crate) fn index_operation(&self) -> &IndexOperation { + &self.index_operation + } + + #[cfg(feature = "product-quantization")] + pub(crate) fn num_pq_chunks(&self) -> usize { + self.num_pq_chunks + } + + #[cfg(feature = "product-quantization")] + pub(crate) fn seed(&self) -> u64 { + self.seed + } + + #[cfg(feature = "product-quantization")] + pub(crate) fn max_fp_vecs_per_prune(&self) -> Option { + self.max_fp_vecs_per_prune + } + + #[cfg(feature = "product-quantization")] + pub(crate) fn use_fp_for_search(&self) -> bool { + self.use_fp_for_search + } + #[cfg(feature = "product-quantization")] pub(crate) fn to_config(&self) -> Result { match &self.index_operation.source { @@ -1018,10 +1049,10 @@ impl std::fmt::Display for IndexPQOperation { #[derive(Debug, Serialize, Deserialize)] pub(crate) struct IndexSQOperation { - pub(crate) index_operation: IndexOperation, - pub(crate) num_bits: usize, - pub(crate) standard_deviations: f64, - pub(crate) use_fp_for_search: bool, + index_operation: IndexOperation, + num_bits: usize, + standard_deviations: f64, + use_fp_for_search: bool, } impl IndexSQOperation { @@ -1029,6 +1060,26 @@ impl IndexSQOperation { "graph-index-build-sq" } + #[cfg(feature = "scalar-quantization")] + pub(crate) fn index_operation(&self) -> &IndexOperation { + &self.index_operation + } + + #[cfg(feature = "scalar-quantization")] + pub(crate) fn num_bits(&self) -> usize { + self.num_bits + } + + #[cfg(feature = "scalar-quantization")] + pub(crate) fn standard_deviations(&self) -> f64 { + self.standard_deviations + } + + #[cfg(feature = "scalar-quantization")] + pub(crate) fn use_fp_for_search(&self) -> bool { + self.use_fp_for_search + } + #[cfg(feature = "scalar-quantization")] pub(crate) fn try_as_config(&self) -> anyhow::Result { match &self.index_operation.source { @@ -1108,13 +1159,13 @@ impl std::fmt::Display for IndexSQOperation { #[derive(Debug, Serialize, Deserialize)] pub(crate) struct SphericalQuantBuild { - pub(crate) build: IndexBuild, // spherical does not support saving and loading - pub(crate) search_phase: SearchPhase, - pub(crate) seed: u64, - pub(crate) transform_kind: inputs::exhaustive::TransformKind, - pub(crate) query_layouts: Vec, - pub(crate) num_bits: NonZeroUsize, - pub(crate) pre_scale: Option, + build: IndexBuild, // spherical does not support saving and loading + search_phase: SearchPhase, + seed: u64, + transform_kind: inputs::exhaustive::TransformKind, + query_layouts: Vec, + num_bits: NonZeroUsize, + pre_scale: Option, } impl SphericalQuantBuild { @@ -1122,6 +1173,41 @@ impl SphericalQuantBuild { "graph-index-build-spherical-quantization" } + #[cfg(feature = "spherical-quantization")] + pub(crate) fn build(&self) -> &IndexBuild { + &self.build + } + + #[cfg(feature = "spherical-quantization")] + pub(crate) fn search_phase(&self) -> &SearchPhase { + &self.search_phase + } + + #[cfg(feature = "spherical-quantization")] + pub(crate) fn seed(&self) -> u64 { + self.seed + } + + #[cfg(feature = "spherical-quantization")] + pub(crate) fn transform_kind(&self) -> &inputs::exhaustive::TransformKind { + &self.transform_kind + } + + #[cfg(feature = "spherical-quantization")] + pub(crate) fn query_layouts(&self) -> &[inputs::exhaustive::SphericalQuery] { + &self.query_layouts + } + + #[cfg(feature = "spherical-quantization")] + pub(crate) fn num_bits(&self) -> NonZeroUsize { + self.num_bits + } + + #[cfg(feature = "spherical-quantization")] + pub(crate) fn pre_scale(&self) -> Option<&inputs::exhaustive::PreScale> { + self.pre_scale.as_ref() + } + #[cfg(feature = "spherical-quantization")] pub(crate) fn try_as_config(&self) -> anyhow::Result { self.build.try_as_config() @@ -1221,7 +1307,7 @@ impl std::fmt::Display for SphericalQuantBuild { } //////////////////////////// -// Dynamic Runbook Params // +// Streaming Runbook Params // //////////////////////////// #[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize)] @@ -1249,7 +1335,7 @@ impl From for graph::InplaceDeleteMethod { /// Runbook loading and phase type definitions are in utils.datafiles #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct DynamicRunbookParams { +pub(crate) struct StreamingRunbookParams { pub(crate) runbook_path: InputFile, pub(crate) dataset_name: String, pub(crate) gt_directory: String, @@ -1267,7 +1353,7 @@ pub(crate) struct DynamicRunbookParams { // 1. The runbook file can be parsed // 2. The dataset_name exists in the runbook // 3. All required ground truth files exist in gt_directory -impl DynamicRunbookParams { +impl StreamingRunbookParams { pub(crate) fn validate(&mut self, checker: &mut Checker) -> anyhow::Result<()> { self.runbook_path.resolve(checker)?; @@ -1334,7 +1420,7 @@ impl DynamicRunbookParams { } } -impl Example for DynamicRunbookParams { +impl Example for StreamingRunbookParams { fn example() -> Self { Self { runbook_path: InputFile::new("path/to/runbook"), @@ -1351,7 +1437,7 @@ impl Example for DynamicRunbookParams { } } -impl DynamicRunbookParams { +impl StreamingRunbookParams { /// Example for hard-delete providers that don't use consolidation. #[cfg(feature = "bftree")] pub(crate) fn example_immediate() -> Self { @@ -1362,9 +1448,9 @@ impl DynamicRunbookParams { } } -impl std::fmt::Display for DynamicRunbookParams { +impl std::fmt::Display for StreamingRunbookParams { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!(f, "Dynamic Runbook Parameters")?; + writeln!(f, "Streaming Runbook Parameters")?; write_field!(f, "Runbook Path", self.runbook_path.display())?; write_field!(f, "Dataset Name", self.dataset_name)?; @@ -1398,26 +1484,78 @@ impl std::fmt::Display for DynamicRunbookParams { } /////////////////////////// -// Graph Index Dynamic // +// Graph Index Streaming // /////////////////////////// +/// Search parameters for streaming benchmarks. +/// +/// Unlike [`TopkSearchPhase`], this does not include a groundtruth file because streaming +/// benchmarks load per-stage groundtruth from the runbook's `gt_directory`. +/// +/// Streaming runs a single search configuration between each runbook stage rather than +/// sweeping a parameter matrix, so thread count and search_l are scalars. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct StreamingSearchParams { + pub(crate) queries: InputFile, + pub(crate) reps: NonZeroUsize, + pub(crate) num_threads: NonZeroUsize, + pub(crate) search_l: usize, + pub(crate) search_n: usize, + pub(crate) recall_k: usize, +} + +impl StreamingSearchParams { + pub(crate) fn max_k(&self) -> usize { + self.recall_k + } + + pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { + self.queries.resolve(checker)?; + Ok(()) + } +} + +impl Example for StreamingSearchParams { + fn example() -> Self { + Self { + queries: InputFile::new("path/to/queries"), + reps: NonZeroUsize::new(5).unwrap(), + num_threads: NonZeroUsize::new(4).unwrap(), + search_l: 40, + search_n: 10, + recall_k: 10, + } + } +} + #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct DynamicIndexRun { - pub(crate) build: IndexBuild, - pub(crate) search_phase: SearchPhase, - pub(crate) runbook_params: DynamicRunbookParams, +pub(crate) struct StreamingIndexRun { + build: IndexBuild, + search: StreamingSearchParams, + runbook_params: StreamingRunbookParams, } -impl DynamicIndexRun { +impl StreamingIndexRun { pub(crate) const fn tag() -> &'static str { - "graph-index-dynamic-run" + "graph-index-stream-run" + } + + pub(crate) fn build(&self) -> &IndexBuild { + &self.build + } + + pub(crate) fn search(&self) -> &StreamingSearchParams { + &self.search + } + + pub(crate) fn runbook_params(&self) -> &StreamingRunbookParams { + &self.runbook_params } pub(crate) fn validate(&mut self, checker: &mut Checker) -> anyhow::Result<()> { self.build.validate(checker)?; self.runbook_params.validate(checker)?; - self.search_phase.validate(checker)?; - + self.search.validate(checker)?; Ok(()) } @@ -1436,21 +1574,21 @@ impl DynamicIndexRun { } } -impl Example for DynamicIndexRun { +impl Example for StreamingIndexRun { fn example() -> Self { let build = IndexBuild::example(); Self { build, - search_phase: SearchPhase::Topk(TopkSearchPhase::example()), - runbook_params: DynamicRunbookParams::example(), + search: StreamingSearchParams::example(), + runbook_params: StreamingRunbookParams::example(), } } } -impl std::fmt::Display for DynamicIndexRun { +impl std::fmt::Display for StreamingIndexRun { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!(f, "Graph Index Dynamic Run")?; + writeln!(f, "Graph Index Streaming Run")?; write_field!(f, "tag", Self::tag())?; writeln!(f, "Runbook Parameters:")?; write!(f, "{}", self.runbook_params)?; diff --git a/diskann-benchmark/src/main.rs b/diskann-benchmark/src/main.rs index 873b99a54..e3ceea6f6 100644 --- a/diskann-benchmark/src/main.rs +++ b/diskann-benchmark/src/main.rs @@ -309,12 +309,12 @@ mod tests { } //////////////////////////// - // Dynamic Index // + // Streaming Index // //////////////////////////// #[test] - fn graph_index_dynamic_integration() { - let raw = value_from_file(&example_directory().join("graph-index-dynamic.json")); + fn graph_index_stream_integration() { + let raw = value_from_file(&example_directory().join("graph-index-stream.json")); run_integration_test(raw); } @@ -343,6 +343,14 @@ mod tests { run_integration_test(raw); } + #[test] + #[cfg(feature = "bftree")] + fn graph_index_bftree_stream_spherical_integration() { + let raw = + value_from_file(&example_directory().join("graph-index-bftree-stream-spherical.json")); + run_integration_test(raw); + } + //////////////////////////// // MinMax Quantization // //////////////////////////// From a0f481c8038521ad4e505fb82567e69ae2aa38b7 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Thu, 18 Jun 2026 10:56:05 -0700 Subject: [PATCH 2/7] undo name rename --- diskann-benchmark/src/inputs/graph_index.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index 2bd8556b8..0dbf8139b 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -218,7 +218,7 @@ impl BetaSearchPhase { } #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct MultiHopSearchPhase { +pub(crate) struct MultihopFilterSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, pub(crate) groundtruth: InputFile, @@ -229,7 +229,7 @@ pub(crate) struct MultiHopSearchPhase { pub(crate) runs: Vec, } -impl MultiHopSearchPhase { +impl MultihopFilterSearchPhase { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; self.query_predicates.resolve(checker)?; @@ -407,7 +407,7 @@ pub(crate) enum SearchPhase { Topk(TopkSearchPhase), Range(RangeSearchPhase), TopkBetaFilter(BetaSearchPhase), - TopkMultihopFilter(MultiHopSearchPhase), + TopkMultihopFilter(MultihopFilterSearchPhase), TopkInlineFilter(InlineSearchPhase), TopkDeterminantDiversity(TopkDeterminantDiversityPhase), } @@ -473,7 +473,7 @@ impl SearchPhase { pub(crate) fn as_topk_multihop_filter( &self, - ) -> Result<&MultiHopSearchPhase, WrongSearchPhaseKind> { + ) -> Result<&MultihopFilterSearchPhase, WrongSearchPhaseKind> { match self { Self::TopkMultihopFilter(phase) => Ok(phase), _ => Err(WrongSearchPhaseKind::new( From b37c21fb56fad082164a175991e18779b6807e7d Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Thu, 18 Jun 2026 13:10:50 -0700 Subject: [PATCH 3/7] enum --- .../src/index/bftree/quantizer_util.rs | 14 +++-- diskann-benchmark/src/inputs/bftree.rs | 60 ++++++++++++++++--- 2 files changed, 59 insertions(+), 15 deletions(-) diff --git a/diskann-benchmark/src/index/bftree/quantizer_util.rs b/diskann-benchmark/src/index/bftree/quantizer_util.rs index 6f6a9d6e6..7cff874c5 100644 --- a/diskann-benchmark/src/index/bftree/quantizer_util.rs +++ b/diskann-benchmark/src/index/bftree/quantizer_util.rs @@ -13,7 +13,10 @@ use diskann_quantization::{ use diskann_utils::views::MatrixView; use rand::SeedableRng; -use crate::{inputs::bftree::QuantConfig, utils::SimilarityMeasure}; +use crate::{ + inputs::bftree::{QuantConfig, SphericalBits}, + utils::SimilarityMeasure, +}; fn new_quantizer( quantizer: SphericalQuantizer, @@ -54,11 +57,10 @@ pub(super) fn build_quantizer( GlobalAllocator, )?; - let poly = match num_bits.get() { - 1 => new_quantizer::<1>(quantizer)?, - 2 => new_quantizer::<2>(quantizer)?, - 4 => new_quantizer::<4>(quantizer)?, - n => anyhow::bail!("{n} bits not supported for spherical quantization"), + let poly = match *num_bits { + SphericalBits::One => new_quantizer::<1>(quantizer)?, + SphericalBits::Two => new_quantizer::<2>(quantizer)?, + SphericalBits::Four => new_quantizer::<4>(quantizer)?, }; Ok(Some(poly)) diff --git a/diskann-benchmark/src/inputs/bftree.rs b/diskann-benchmark/src/inputs/bftree.rs index 77f5c2c2b..a46dc290c 100644 --- a/diskann-benchmark/src/inputs/bftree.rs +++ b/diskann-benchmark/src/inputs/bftree.rs @@ -2,7 +2,7 @@ * Copyright (c) Microsoft Corporation. * Licensed under the MIT license. */ -use std::num::{NonZero, NonZeroUsize}; +use std::num::NonZero; use crate::inputs::{ as_input, exhaustive, @@ -211,6 +211,54 @@ fn bftree_parameters_from( } } +/// Supported bit widths for spherical quantization. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub(crate) enum SphericalBits { + #[serde(rename = "1")] + One = 1, + #[serde(rename = "2")] + Two = 2, + #[serde(rename = "4")] + Four = 4, +} + +impl SphericalBits { + pub(crate) const fn get(self) -> usize { + self as usize + } +} + +impl std::fmt::Display for SphericalBits { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.get()) + } +} + +impl TryFrom for SphericalBits { + type Error = String; + + fn try_from(value: usize) -> Result { + match value { + 1 => Ok(Self::One), + 2 => Ok(Self::Two), + 4 => Ok(Self::Four), + n => Err(format!( + "{n} bits not supported for spherical quantization; expected 1, 2, or 4" + )), + } + } +} + +impl<'de> Deserialize<'de> for SphericalBits { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let n = usize::deserialize(deserializer)?; + SphericalBits::try_from(n).map_err(serde::de::Error::custom) + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "kind")] pub(crate) enum QuantConfig { @@ -220,7 +268,7 @@ pub(crate) enum QuantConfig { Spherical { seed: u64, transform_kind: exhaustive::TransformKind, - num_bits: NonZeroUsize, + num_bits: SphericalBits, #[serde(deserialize_with = "Deserialize::deserialize")] pre_scale: Option, #[serde(deserialize_with = "Deserialize::deserialize")] @@ -233,14 +281,8 @@ impl QuantConfig { match self { Self::None => Ok(()), Self::Spherical { - num_bits, - quant_store_config, - .. + quant_store_config, .. } => { - match num_bits.get() { - 1 | 2 | 4 => {} - n => anyhow::bail!("{n} bits are not supported for spherical quantization"), - } if let Some(cfg) = quant_store_config { cfg.fill_defaults(); cfg.validate()?; From 634b25c3f219ec3cf798bd151a2039a5e8a33bb0 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Thu, 18 Jun 2026 13:38:02 -0700 Subject: [PATCH 4/7] QuantConfig implements display --- diskann-benchmark/src/inputs/bftree.rs | 76 +++++++++++--------------- 1 file changed, 32 insertions(+), 44 deletions(-) diff --git a/diskann-benchmark/src/inputs/bftree.rs b/diskann-benchmark/src/inputs/bftree.rs index a46dc290c..84b6c9877 100644 --- a/diskann-benchmark/src/inputs/bftree.rs +++ b/diskann-benchmark/src/inputs/bftree.rs @@ -293,6 +293,36 @@ impl QuantConfig { } } +impl std::fmt::Display for QuantConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Spherical { + seed, + transform_kind, + num_bits, + pre_scale, + quant_store_config, + } => { + write_field!(f, "quantization", "spherical")?; + write_field!(f, "num_bits", num_bits)?; + write_field!(f, "seed", seed)?; + write_field!(f, "transform_kind", transform_kind)?; + if let Some(pre_scale) = pre_scale { + write_field!(f, "pre_scale", pre_scale)?; + } + if let Some(cfg) = quant_store_config { + writeln!(f, "\n Quant Store:")?; + write!(f, "{}", cfg)?; + } + } + Self::None => { + write_field!(f, "quantization", "none")?; + } + } + Ok(()) + } +} + #[derive(Debug, Serialize, Deserialize)] pub(crate) struct BfTreeBuild { build: IndexBuild, @@ -388,20 +418,7 @@ impl std::fmt::Display for BfTreeBuild { } write_field!(f, "tag", Self::tag())?; - if let QuantConfig::Spherical { - seed, - transform_kind, - num_bits, - .. - } = &self.quantization - { - write_field!(f, "quantization", "spherical")?; - write_field!(f, "num_bits", num_bits)?; - write_field!(f, "seed", seed)?; - write_field!(f, "transform_kind", transform_kind)?; - } else { - write_field!(f, "quantization", "none")?; - } + write!(f, "{}", self.quantization)?; writeln!(f)?; self.build.summarize_fields(f)?; @@ -414,14 +431,6 @@ impl std::fmt::Display for BfTreeBuild { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } - if let QuantConfig::Spherical { - quant_store_config: Some(ref cfg), - .. - } = self.quantization - { - writeln!(f, "\n Quant Store:")?; - write!(f, "{}", cfg)?; - } Ok(()) } @@ -530,20 +539,7 @@ impl std::fmt::Display for BfTreeStreamingRun { } write_field!(f, "tag", Self::tag())?; - if let QuantConfig::Spherical { - seed, - transform_kind, - num_bits, - .. - } = &self.quantization - { - write_field!(f, "quantization", "spherical")?; - write_field!(f, "num_bits", num_bits)?; - write_field!(f, "seed", seed)?; - write_field!(f, "transform_kind", transform_kind)?; - } else { - write_field!(f, "quantization", "none")?; - } + write!(f, "{}", self.quantization)?; writeln!(f)?; self.build.summarize_fields(f)?; @@ -556,14 +552,6 @@ impl std::fmt::Display for BfTreeStreamingRun { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } - if let QuantConfig::Spherical { - quant_store_config: Some(ref cfg), - .. - } = self.quantization - { - writeln!(f, "\n Quant Store:")?; - write!(f, "{}", cfg)?; - } Ok(()) } From b306e94d93654cc053b6404f8abc2a240025a7d8 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Mon, 22 Jun 2026 08:48:54 -0700 Subject: [PATCH 5/7] Validate search_l >= search_n in StreamingSearchParams Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- diskann-benchmark/src/inputs/graph_index.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index 0dbf8139b..dacae8716 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -1511,6 +1511,13 @@ impl StreamingSearchParams { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; + if self.search_l < self.search_n { + return Err(anyhow!( + "search_l {} is less than search_n: {}", + self.search_l, + self.search_n + )); + } Ok(()) } } From f9a9213b54ef87975ee31d7295d1e5b4a9b162e0 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Mon, 22 Jun 2026 15:03:56 -0700 Subject: [PATCH 6/7] Address review: restore InlineFilterSearchPhase name, add recall_k <= search_n validation - Restore InlineFilterSearchPhase naming that was accidentally overwritten - Add recall_k > search_n validation to GraphSearch::validate and StreamingSearchParams::validate to fail fast before running benchmarks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- diskann-benchmark/src/inputs/graph_index.rs | 25 +++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index dacae8716..62b9194c5 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -50,6 +50,14 @@ pub(crate) struct GraphSearch { impl GraphSearch { pub(crate) fn validate(&mut self, _checker: &mut Checker) -> Result<(), anyhow::Error> { + if self.recall_k > self.search_n { + return Err(anyhow!( + "recall_k {} is greater than search_n: {}", + self.recall_k, + self.search_n + )); + } + for (i, l) in self.search_l.iter().enumerate() { if *l < self.search_n { return Err(anyhow!( @@ -259,7 +267,7 @@ impl AdaptiveL { } #[derive(Debug, Serialize, Deserialize)] -pub(crate) struct InlineSearchPhase { +pub(crate) struct InlineFilterSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, pub(crate) groundtruth: InputFile, @@ -271,7 +279,7 @@ pub(crate) struct InlineSearchPhase { pub(crate) adaptive_l: Option, } -impl InlineSearchPhase { +impl InlineFilterSearchPhase { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; self.query_predicates.resolve(checker)?; @@ -408,7 +416,7 @@ pub(crate) enum SearchPhase { Range(RangeSearchPhase), TopkBetaFilter(BetaSearchPhase), TopkMultihopFilter(MultihopFilterSearchPhase), - TopkInlineFilter(InlineSearchPhase), + TopkInlineFilter(InlineFilterSearchPhase), TopkDeterminantDiversity(TopkDeterminantDiversityPhase), } @@ -483,7 +491,9 @@ impl SearchPhase { } } - pub(crate) fn as_topk_inline_filter(&self) -> Result<&InlineSearchPhase, WrongSearchPhaseKind> { + pub(crate) fn as_topk_inline_filter( + &self, + ) -> Result<&InlineFilterSearchPhase, WrongSearchPhaseKind> { match self { Self::TopkInlineFilter(phase) => Ok(phase), _ => Err(WrongSearchPhaseKind::new( @@ -1511,6 +1521,13 @@ impl StreamingSearchParams { pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.queries.resolve(checker)?; + if self.recall_k > self.search_n { + return Err(anyhow!( + "recall_k {} is greater than search_n: {}", + self.recall_k, + self.search_n + )); + } if self.search_l < self.search_n { return Err(anyhow!( "search_l {} is less than search_n: {}", From ff5cff6bd8523a838353f08f146889dd3e571d4e Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Thu, 25 Jun 2026 10:55:31 -0700 Subject: [PATCH 7/7] Make Managed opt-in: bftree streams bypass ID management layer Introduce StreamingOutput trait and make run_streaming generic over the stream type. This allows: - bftree: implements Stream directly on StreamRunner, using tags as slot IDs without translation (no Managed wrapper) - inmem: continues using Managed for tag-to-slot ID translation When inmem 2.0 lands with native ID management, Managed can be deleted entirely with no infrastructure changes needed. Also: - Remove SlotReclaim::Immediate and recycle_tags (bftree-only paths) - Simplify PhantomData T> to PhantomData on StreamRunner Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- diskann-benchmark/src/index/benchmarks.rs | 4 +- .../index/bftree/full_precision_streaming.rs | 19 ++-- .../src/index/bftree/spherical_streaming.rs | 22 ++-- .../src/index/streaming/managed.rs | 9 -- diskann-benchmark/src/index/streaming/mod.rs | 104 ++++++++++++++---- .../src/index/streaming/runner.rs | 100 ++++++++++++++++- diskann-benchmark/src/utils/streaming.rs | 19 ---- 7 files changed, 201 insertions(+), 76 deletions(-) diff --git a/diskann-benchmark/src/index/benchmarks.rs b/diskann-benchmark/src/index/benchmarks.rs index 43f389461..98399cc1f 100644 --- a/diskann-benchmark/src/index/benchmarks.rs +++ b/diskann-benchmark/src/index/benchmarks.rs @@ -329,7 +329,7 @@ where ) -> anyhow::Result>> { writeln!(output, "{}", input)?; - streaming::run_streaming::( + streaming::run_streaming::, _>( input.runbook_params(), |max_points| full_precision_streaming::(input, max_points), output, @@ -742,7 +742,7 @@ where .ok_or_else(|| anyhow::anyhow!("consolidate_threshold is required for inmem streaming"))?; let capacity = ((max_points as f32) * (1.0 + 2.0 * consolidate_threshold)).ceil() as usize; - streaming::build_streamer( + streaming::build_managed_streamer( input.build().data(), search, streaming::managed::SlotReclaim::Deferred(consolidate_threshold), diff --git a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs index 73019fa02..317ef8235 100644 --- a/diskann-benchmark/src/index/bftree/full_precision_streaming.rs +++ b/diskann-benchmark/src/index/bftree/full_precision_streaming.rs @@ -20,12 +20,7 @@ use diskann_providers::model::graph::provider::async_::common::FullPrecision; use diskann_utils::sampling::WithApproximateNorm; use crate::{ - index::streaming::{ - managed::{self, Managed}, - runner::BfTreeMaintainer, - stats::StreamStats, - StreamRunner, - }, + index::streaming::{runner::BfTreeMaintainer, stats::StreamStats, StreamRunner}, inputs::bftree::{BfTreeStreamingRun, QuantConfig}, utils, }; @@ -52,7 +47,7 @@ where T: VectorRepr + WithApproximateNorm + SampleableForStart + AsDataType + bytemuck::Pod, { type Input = BfTreeStreamingRun; - type Output = Vec>; + type Output = Vec; fn try_match(&self, input: &Self::Input) -> Result { let mut failure_score: Option = None; @@ -91,7 +86,7 @@ where ) -> anyhow::Result { writeln!(output, "{}", input)?; - crate::index::streaming::run_streaming::( + crate::index::streaming::run_streaming::, _>( input.runbook_params(), |max_points| bftree_streaming::(input, max_points), output, @@ -99,10 +94,13 @@ where } } +type BfTreeFullPrecisionStream = + StreamRunner, T, FullPrecision, BfTreeMaintainer>; + fn bftree_streaming( input: &BfTreeStreamingRun, max_points: usize, -) -> anyhow::Result>> +) -> anyhow::Result>> where T: bytemuck::Pod + VectorRepr + WithApproximateNorm + SampleableForStart, { @@ -111,10 +109,9 @@ where let num_start_points = input.build().start_point_strategy().count(); let capacity = max_points + num_start_points; - crate::index::streaming::build_streamer( + crate::index::streaming::build_direct_streamer( input.build().data(), search, - crate::index::streaming::managed::SlotReclaim::Immediate, capacity, |data, capacity| { let config = input.try_as_config()?.build()?; diff --git a/diskann-benchmark/src/index/bftree/spherical_streaming.rs b/diskann-benchmark/src/index/bftree/spherical_streaming.rs index 31ae2fa04..826141397 100644 --- a/diskann-benchmark/src/index/bftree/spherical_streaming.rs +++ b/diskann-benchmark/src/index/bftree/spherical_streaming.rs @@ -21,11 +21,7 @@ use diskann_bftree::BfTreeProvider; use diskann_providers::model::graph::provider::async_::common::Quantized; use crate::{ - index::streaming::{ - managed::{self, Managed}, - stats::StreamStats, - BfTreeMaintainer, StreamRunner, - }, + index::streaming::{stats::StreamStats, BfTreeMaintainer, StreamRunner}, inputs::bftree::{BfTreeStreamingRun, QuantConfig}, utils, }; @@ -43,7 +39,7 @@ impl StreamingSpherical { impl Benchmark for StreamingSpherical { type Input = BfTreeStreamingRun; - type Output = Vec>; + type Output = Vec; fn try_match(&self, input: &Self::Input) -> Result { let mut failure_score: Option = None; @@ -87,7 +83,7 @@ impl Benchmark for StreamingSpherical { ) -> anyhow::Result { writeln!(output, "{}", input)?; - crate::index::streaming::run_streaming::( + crate::index::streaming::run_streaming::( input.runbook_params(), |max_points| bftree_sq_streaming_impl(input, max_points), output, @@ -95,19 +91,25 @@ impl Benchmark for StreamingSpherical { } } +type BfTreeSphericalStream = StreamRunner< + BfTreeProvider, + f32, + Quantized, + BfTreeMaintainer, +>; + fn bftree_sq_streaming_impl( input: &BfTreeStreamingRun, max_points: usize, -) -> anyhow::Result>> { +) -> anyhow::Result> { let search = input.search(); let num_start_points = input.build().start_point_strategy().count(); let capacity = max_points + num_start_points; - crate::index::streaming::build_streamer( + crate::index::streaming::build_direct_streamer( input.build().data(), search, - crate::index::streaming::managed::SlotReclaim::Immediate, capacity, |data, capacity| { let quantizer_poly = super::quantizer_util::build_quantizer( diff --git a/diskann-benchmark/src/index/streaming/managed.rs b/diskann-benchmark/src/index/streaming/managed.rs index cbe7188d9..5d23fc4c5 100644 --- a/diskann-benchmark/src/index/streaming/managed.rs +++ b/diskann-benchmark/src/index/streaming/managed.rs @@ -21,11 +21,6 @@ use crate::utils::streaming::TagSlotManager; /// while soft-delete providers need a deferred consolidation pass. #[derive(Debug, Clone)] pub(crate) enum SlotReclaim { - /// Slots are recycled to `empty_slots` immediately during delete. - /// No maintenance pass is needed. - #[cfg(feature = "bftree")] - Immediate, - /// Slots are held in `deleted_slots` until maintenance fires. /// Maintenance triggers when `num_deleted > num_active * threshold`. Deferred(f32), @@ -148,8 +143,6 @@ where let (overhead_slots, slots) = timed!(self.book_keeping.find_slots_by_tags(tags.clone())?); let output = self.stream.delete(&slots)?; let (overhead_reclaim, _) = timed!(match &self.reclaim { - #[cfg(feature = "bftree")] - SlotReclaim::Immediate => self.book_keeping.recycle_tags(tags)?, SlotReclaim::Deferred(_) => self.book_keeping.mark_tags_deleted(tags)?, }); Ok(Stats::new(overhead_slots + overhead_reclaim, output)) @@ -163,8 +156,6 @@ where fn needs_maintenance(&mut self) -> bool { match &self.reclaim { - #[cfg(feature = "bftree")] - SlotReclaim::Immediate => false, SlotReclaim::Deferred(threshold) => { let num_active = self.book_keeping.num_active(); let limit = (num_active as f32 * threshold) as usize; diff --git a/diskann-benchmark/src/index/streaming/mod.rs b/diskann-benchmark/src/index/streaming/mod.rs index b4e3bc834..4cd756a1a 100644 --- a/diskann-benchmark/src/index/streaming/mod.rs +++ b/diskann-benchmark/src/index/streaming/mod.rs @@ -6,7 +6,7 @@ use std::{io::Write, sync::Arc}; use diskann::utils::VectorRepr; -use diskann_benchmark_core::streaming::{executors::bigann, Executor}; +use diskann_benchmark_core::streaming::{self, executors::bigann, Executor}; use diskann_benchmark_runner::output::Output; use diskann_utils::views::Matrix; @@ -24,6 +24,26 @@ use crate::{ utils::datafiles, }; +/// Trait for streaming benchmark outputs that wrap or produce [`stats::StreamStats`]. +/// +/// This allows [`run_streaming`] to work with both direct streams (which produce +/// `StreamStats` directly) and managed streams (which produce `managed::Stats`). +pub(crate) trait StreamingOutput: std::fmt::Display + 'static { + fn stream_stats(&self) -> &stats::StreamStats; +} + +impl StreamingOutput for stats::StreamStats { + fn stream_stats(&self) -> &stats::StreamStats { + self + } +} + +impl StreamingOutput for managed::Stats { + fn stream_stats(&self) -> &stats::StreamStats { + self.inner() + } +} + /// Construct the streaming stack: load data/queries, create the managed stream via the /// closure, then wrap in [`Managed`] and [`bigann::WithData`]. /// @@ -31,7 +51,7 @@ use crate::{ /// computes this differently (inmem applies headroom, bf_tree adds start points). /// /// The closure receives `(&data, capacity)` so it can use `data.ncols()` for provider params. -pub(crate) fn build_streamer( +pub(crate) fn build_managed_streamer( data_path: &diskann_benchmark_runner::files::InputFile, search: &StreamingSearchParams, reclaim: managed::SlotReclaim, @@ -62,19 +82,57 @@ where Ok(layered) } +/// Construct a direct streaming stack (no ID management layer). +/// +/// For providers where external IDs match internal slots (e.g., bf-tree), the +/// [`Managed`] layer is unnecessary. This function creates the stack without it. +/// +/// The closure receives `(&data, capacity)` so it can use `data.ncols()` for provider params. +#[cfg(feature = "bftree")] +pub(crate) fn build_direct_streamer( + data_path: &diskann_benchmark_runner::files::InputFile, + search: &StreamingSearchParams, + capacity: usize, + make_stream: F, +) -> anyhow::Result> +where + T: bytemuck::Pod + VectorRepr + 'static, + S: streaming::Stream> + 'static, + F: FnOnce(&Matrix, usize) -> anyhow::Result, +{ + let data = datafiles::load_dataset::(datafiles::BinFile(data_path))?; + let queries = Arc::new(datafiles::load_dataset::(datafiles::BinFile( + &search.queries, + ))?); + + let stream = make_stream(&data, capacity)?; + + let max_k = search.max_k(); + let layered = bigann::WithData::new(stream, data, queries, move |path| { + Ok(Box::new(datafiles::load_groundtruth( + datafiles::BinFile(path), + Some(max_k), + )?)) + }); + + Ok(layered) +} + /// Run a streaming benchmark using the given runbook parameters. /// /// `make_streamer` receives `max_points` from the loaded runbook and returns the /// constructed streamer. This is shared across all streaming benchmarks (inmem, bftree) /// to avoid duplicating the runbook load → run_with → stage banner → summary logic. -pub(crate) fn run_streaming( +pub(crate) fn run_streaming( runbook_params: &StreamingRunbookParams, make_streamer: F, mut output: &mut dyn Output, -) -> anyhow::Result>> +) -> anyhow::Result> where T: 'static, - F: FnOnce(usize) -> anyhow::Result>>, + S: streaming::Stream>, + S::Output: StreamingOutput, + F: FnOnce(usize) -> anyhow::Result>, { let groundtruth_directory = runbook_params .resolved_gt_directory @@ -95,22 +153,24 @@ where let stages = runbook.len(); let mut i = 1; - runbook.run_with( - &mut streamer, - |o: managed::Stats| -> anyhow::Result<()> { - if o.inner().is_maintain() { - let message = format!("Ran maintenance before stage {}", i); - write!(output, "{}", crate::utils::SmallBanner(&message))?; - } else { - let message = format!("Finished stage {} of {}: {}", i, stages, o.inner().kind()); - write!(output, "{}", crate::utils::SmallBanner(&message))?; - i += 1; - } - writeln!(output, "{}", o)?; - results.push(o); - Ok(()) - }, - )?; + runbook.run_with(&mut streamer, |o: S::Output| -> anyhow::Result<()> { + if o.stream_stats().is_maintain() { + let message = format!("Ran maintenance before stage {}", i); + write!(output, "{}", crate::utils::SmallBanner(&message))?; + } else { + let message = format!( + "Finished stage {} of {}: {}", + i, + stages, + o.stream_stats().kind() + ); + write!(output, "{}", crate::utils::SmallBanner(&message))?; + i += 1; + } + writeln!(output, "{}", o)?; + results.push(o); + Ok(()) + })?; write!( output, @@ -121,7 +181,7 @@ where writeln!( output, "{}", - stats::Summary::new(results.iter().map(|r| r.inner())) + stats::Summary::new(results.iter().map(|r| r.stream_stats())) )?; Ok(results) diff --git a/diskann-benchmark/src/index/streaming/runner.rs b/diskann-benchmark/src/index/streaming/runner.rs index 3a271c8b6..871f2f153 100644 --- a/diskann-benchmark/src/index/streaming/runner.rs +++ b/diskann-benchmark/src/index/streaming/runner.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use std::{borrow::Cow, marker::PhantomData, num::NonZeroUsize, sync::Arc}; +use std::{borrow::Cow, marker::PhantomData, num::NonZeroUsize, ops::Range, sync::Arc}; use diskann::{ graph::{ @@ -20,7 +20,7 @@ use diskann_benchmark_core::{ self, graph::{Strategy, KNN}, }, - streaming::graph::InplaceDelete, + streaming::{self, executors::bigann, graph::InplaceDelete}, }; use diskann_utils::{ future::AsyncFriendly, @@ -61,7 +61,7 @@ where inplace_delete_num_to_replace: usize, inplace_delete_method: InplaceDeleteMethod, maintainer: M, - _marker: PhantomData T>, + _marker: PhantomData, } impl StreamRunner @@ -213,3 +213,97 @@ where Ok(StreamStats::Maintain(Vec::new())) } } + +///////////////////////////////////// +// Direct Stream (no ID management) // +///////////////////////////////////// + +/// Direct [`streaming::Stream`] implementation for providers that manage their own IDs. +/// +/// In this mode, the external tag IDs from the runbook are used directly as slot IDs +/// (cast to `u32`). No [`super::Managed`] layer is needed. +impl streaming::Stream> for StreamRunner +where + DP: DataProvider + + for<'a> SetElement<&'a [T]> + + Delete, + DP::ExternalId: search::Id, + S: for<'a> InsertStrategy<'a, DP, &'a [T]> + + for<'a> DefaultSearchStrategy<'a, DP, &'a [T], DP::ExternalId> + + InplaceDeleteStrategy + + Clone + + AsyncFriendly, + T: AsyncFriendly + Clone, + M: Maintainer, +{ + type Output = StreamStats; + + fn search( + &mut self, + (queries, groundtruth): (Arc>, &dyn Rows), + ) -> anyhow::Result { + let knn = KNN::new( + self.index.clone(), + queries, + Strategy::broadcast(self.strategy.clone()), + )?; + + let run = GraphSearch { + search_n: self.search.search_n, + search_l: vec![self.search.search_l], + recall_k: self.search.recall_k, + }; + let num_threads = [self.search.num_threads]; + let runs = [run]; + let steps = knn::SearchSteps::new(self.search.reps, &num_threads, &runs); + let results = knn::run(&knn, groundtruth, steps)?; + Ok(StreamStats::Search(results)) + } + + fn insert( + &mut self, + (data, tags): (MatrixView<'_, T>, Range), + ) -> anyhow::Result { + let slots: Vec = tags.map(|t| t as u32).collect(); + Ok(StreamStats::Insert(self.build(data, &slots)?)) + } + + fn replace( + &mut self, + (data, tags): (MatrixView<'_, T>, Range), + ) -> anyhow::Result { + let slots: Vec = tags.map(|t| t as u32).collect(); + Ok(StreamStats::Replace(self.build(data, &slots)?)) + } + + fn delete(&mut self, tags: Range) -> anyhow::Result { + let slots: Vec = tags.map(|t| t as u32).collect(); + let runner = InplaceDelete::new( + self.index.clone(), + self.strategy.clone(), + self.inplace_delete_num_to_replace, + self.inplace_delete_method, + Slice::new(slots.into()), + ); + + let results = build::build( + runner, + Parallelism::fixed(Some(ONE), self.ntasks), + &self.runtime, + )?; + + Ok(StreamStats::Delete(GenericStats::new( + Cow::Borrowed("Delete"), + results, + )?)) + } + + fn maintain(&mut self, _: ()) -> anyhow::Result { + self.maintainer + .maintain(&self.index, &self.runtime, self.ntasks) + } + + fn needs_maintenance(&mut self) -> bool { + false + } +} diff --git a/diskann-benchmark/src/utils/streaming.rs b/diskann-benchmark/src/utils/streaming.rs index ee9456c83..b1382a261 100644 --- a/diskann-benchmark/src/utils/streaming.rs +++ b/diskann-benchmark/src/utils/streaming.rs @@ -132,25 +132,6 @@ impl TagSlotManager { } /// Immediately recycle tags back to empty slots, removing all mappings. - /// - /// Used by hard-delete providers where slots can be reused immediately - /// without waiting for a consolidation pass. - #[cfg(feature = "bftree")] - pub fn recycle_tags(&mut self, tag_range: std::ops::Range) -> anyhow::Result<()> { - for tag in tag_range { - if let Some(slot) = self.tag_to_slot.remove(&tag) { - self.slot_to_tag.remove(&slot); - self.empty_slots.push_back(slot); - } else { - return Err(anyhow::anyhow!( - "Tag {} not found in slot mapping for recycle", - tag - )); - } - } - Ok(()) - } - /// Consolidate deleted slots back to empty slots after background cleaning /// /// This method processes all deleted slots by: