From 6e7bfdc8e7286b71fa82f8a45debe198fdd58f48 Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Tue, 21 Jul 2026 11:08:47 +0800 Subject: [PATCH 1/7] Refactor disk index build structure Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ecf6179f-7467-44be-8a77-bf7a96e319af --- diskann-disk/src/build/builder/build.rs | 210 +----------------- .../src/build/builder/inmem_builder.rs | 202 ++++++++++++++++- .../builder/{core.rs => merged_index.rs} | 2 +- diskann-disk/src/build/builder/mod.rs | 4 +- diskann-disk/src/build/builder/quantizer.rs | 2 +- diskann-disk/src/build/builder/tests.rs | 2 +- .../src/search/provider/disk_provider.rs | 2 +- 7 files changed, 210 insertions(+), 214 deletions(-) rename diskann-disk/src/build/builder/{core.rs => merged_index.rs} (99%) diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index 432c577fc..4e0098c6c 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -3,37 +3,23 @@ * Licensed under the MIT license. */ -//! Async disk index builder implementation. -use std::{ - marker::PhantomData, - num::NonZeroUsize, - sync::{Arc, Mutex}, -}; +//! Disk index builder implementation. +use std::marker::PhantomData; use crate::data_model::GraphDataType; -use diskann::{ - utils::{async_tools, VectorRepr, ONE}, - ANNError, ANNResult, -}; +use diskann::{utils::VectorRepr, ANNResult}; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ - model::{ - graph::provider::async_::inmem::DefaultProviderParameters, IndexConfiguration, - MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS, - }, - storage::{DiskGraphOnly, PQStorage}, - utils::{ - create_thread_pool, find_medoid_with_sampling, RayonThreadPoolRef, VectorDataIterator, - MAX_MEDOID_SAMPLE_SIZE, - }, + model::{IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS}, + storage::PQStorage, + utils::{create_thread_pool, RayonThreadPoolRef}, }; -use tokio::task::JoinSet; -use tracing::{debug, info}; +use tracing::info; use crate::{ build::builder::{ - core::{determine_build_strategy, IndexBuildStrategy, MergedVamanaIndexBuilder}, - inmem_builder::{new_inmem_index_builder, InmemIndexBuilder}, + inmem_builder::build_inmem_index, + merged_index::{determine_build_strategy, IndexBuildStrategy, MergedVamanaIndexBuilder}, quantizer::BuildQuantizer, tokio::create_runtime, }, @@ -58,7 +44,6 @@ where build_quantizer: BuildQuantizer, _phantom: PhantomData, } - impl<'a, Data, StorageProvider> DiskIndexBuilder<'a, Data, StorageProvider> where Data: GraphDataType, @@ -210,180 +195,3 @@ where Ok(()) } } - -pub(super) async fn build_inmem_index( - config: IndexConfiguration, - quantizer: &BuildQuantizer, - data_path: &str, - save_path: &str, - storage_provider: &StorageProvider, -) -> ANNResult<()> -where - T: VectorRepr, - StorageProvider: StorageReadProvider + StorageWriteProvider + 'static, - ::Reader: std::marker::Send, -{ - // use either user-specified number of threads or default to available parallelism - let num_tasks = NonZeroUsize::new(config.num_threads) - .or_else(|| std::thread::available_parallelism().ok()) - .ok_or_else(|| ANNError::log_index_error("Failed to determine number of threads"))?; - - // Associated data will only be used in the write_disk_layout function which only requires the none-partitioned associated data stream. - let dataset_iter = Arc::new(Mutex::new({ - let iter = VectorDataIterator::<_, T>::new(data_path, Option::None, storage_provider)?; - iter.enumerate() - })); - - let index_config = config.config.clone(); - let provider_parameters = DefaultProviderParameters { - max_points: config.max_points, - frozen_points: ONE, - metric: config.dist_metric, - dim: config.dim, - max_degree: index_config.max_degree_u32().get(), - prefetch_lookahead: config.prefetch_lookahead.map(|x| x.get()), - prefetch_cache_line_level: config.prefetch_cache_line_level, - }; - let index = new_inmem_index_builder::(index_config, provider_parameters, quantizer)?; - let medoid_id = - set_start_point_to_medoid::(&index, data_path, config.random_seed, storage_provider)?; - let start_point = u32_try_from(medoid_id)?; - - run_build(&index, dataset_iter, num_tasks).await?; - - #[cfg(debug_assertions)] - log_build_stats::<_>(&index).await?; - - run_final_prune(&index, num_tasks).await?; - index - .save_graph( - storage_provider, - &(start_point, DiskGraphOnly::new(save_path)), - ) - .await?; - - Ok(()) -} - -#[cfg(debug_assertions)] -/// Log statistics about the build process -async fn log_build_stats(index: &Arc>) -> ANNResult<()> { - debug!( - "Number of points reachable in the graph: {}", - index.count_reachable_nodes().await? - ); - - let (full_vector, quant_vector) = index.counts_for_get_vector(); - let capacity = index.capacity(); - debug!( - "Number of get vector calls per insert: {}", - full_vector as f32 / capacity as f32 - ); - debug!( - "Number of get quantized vector calls per insert: {}", - quant_vector as f32 / capacity as f32 - ); - - Ok(()) -} - -/// Convert a `usize` index into the `u32` internal id type, erroring if it does not fit. -/// -/// The async index uses `u32` internal ids, so positions in the dataset must not exceed -/// `u32::MAX`. -fn u32_try_from(value: usize) -> ANNResult { - u32::try_from(value) - .map_err(|_| ANNError::log_index_error(format_args!("id {value} exceeds u32::MAX"))) -} - -fn set_start_point_to_medoid( - index: &Arc>, - path: &str, - random_seed: Option, - reader: &StorageReader, -) -> ANNResult -where - T: VectorRepr, - StorageReader: StorageReadProvider, -{ - let mut rng = diskann_providers::utils::create_rnd_from_optional_seed(random_seed); - let (medoid, medoid_id) = - find_medoid_with_sampling::(path, reader, MAX_MEDOID_SAMPLE_SIZE, &mut rng)?; - - index.set_start_point(medoid.as_slice())?; - - debug!("Set start point to medoid ID: {}", medoid_id); - - Ok(medoid_id) -} - -async fn run_build( - index: &Arc>, - iterator: Arc>, - num_tasks: NonZeroUsize, -) -> ANNResult<()> -where - T: VectorRepr, - I: Iterator, ()))> + Send + 'static, -{ - let total_points = index.capacity(); - let partitions = async_tools::PartitionIter::new(total_points, num_tasks); - - let mut tasks = JoinSet::new(); - - for partition in partitions { - let index_clone = index.clone(); - let iterator_clone = iterator.clone(); - tasks.spawn(async move { - for _ in partition { - let vector_data = { - let mut guard = iterator_clone.lock().map_err(|_| { - ANNError::log_index_error("Poisoned mutex during construction") - })?; - guard.next() - }; - - match vector_data { - Some((i, (vector, _))) => { - let id = u32_try_from(i)?; - index_clone.insert_vector(id, vector.as_ref()).await?; - } - None => break, - } - } - ANNResult::Ok(()) - }); - } - - // Wait for all tasks to complete. - while let Some(res) = tasks.join_next().await { - res.map_err(|_| ANNError::log_index_error("A spawned insert task failed"))??; - } - - info!("Linked all points. Num points: #{}", total_points); - Ok(()) -} - -async fn run_final_prune( - index: &Arc>, - num_tasks: NonZeroUsize, -) -> ANNResult<()> { - let partitions = async_tools::PartitionIter::new(index.total_points(), num_tasks); - - let mut tasks = JoinSet::new(); - - for partition in partitions { - let index_clone = index.clone(); - tasks.spawn(async move { - let range = u32_try_from(partition.start)?..u32_try_from(partition.end)?; - index_clone.final_prune(range).await - }); - } - - // Wait for all final prune tasks to complete - while let Some(res) = tasks.join_next().await { - res.map_err(|_| ANNError::log_index_error("A spawned final prune task failed"))??; - } - - Ok(()) -} diff --git a/diskann-disk/src/build/builder/inmem_builder.rs b/diskann-disk/src/build/builder/inmem_builder.rs index 2eeab9780..ae3d67aed 100644 --- a/diskann-disk/src/build/builder/inmem_builder.rs +++ b/diskann-disk/src/build/builder/inmem_builder.rs @@ -3,7 +3,12 @@ * Licensed under the MIT license. */ -use std::{marker::PhantomData, pin::Pin, sync::Arc}; +use std::{ + marker::PhantomData, + num::NonZeroUsize, + pin::Pin, + sync::{Arc, Mutex}, +}; use diskann::{ graph::{ @@ -11,10 +16,12 @@ use diskann::{ Config, DiskANNIndex, }, provider::DefaultContext, - utils::VectorRepr, + utils::{async_tools, VectorRepr, ONE}, ANNError, ANNResult, }; -use diskann_providers::storage::{DynWriteProvider, WriteProviderWrapper}; +use diskann_providers::storage::{ + DynWriteProvider, StorageReadProvider, StorageWriteProvider, WriteProviderWrapper, +}; use diskann_providers::{ index::diskann_async, model::graph::provider::async_::{ @@ -23,9 +30,13 @@ use diskann_providers::{ DefaultProvider, DefaultProviderParameters, FullPrecisionProvider, SetStartPoints, }, }, + model::IndexConfiguration, storage::{DiskGraphOnly, SaveWith}, + utils::{find_medoid_with_sampling, VectorDataIterator, MAX_MEDOID_SAMPLE_SIZE}, }; use diskann_utils::future::{AsyncFriendly, SendFuture}; +use tokio::task::JoinSet; +use tracing::{debug, info}; use super::quantizer::BuildQuantizer; @@ -33,7 +44,7 @@ use super::quantizer::BuildQuantizer; /// /// Thread safety: /// Implementors must be `Send` and `Sync`. Methods can be called from many tasks. -pub(super) trait InmemIndexBuilder: Send + Sync { +trait InmemIndexBuilder: Send + Sync { /// Return the total capacity of the provider, **excluding** start points. fn capacity(&self) -> usize; @@ -153,7 +164,7 @@ where // Quant Implementation // ////////////////////////// -pub(super) struct QuantInMemBuilder +struct QuantInMemBuilder where Q: AsyncFriendly, { @@ -165,7 +176,7 @@ impl QuantInMemBuilder where Q: AsyncFriendly, { - pub fn new(index: DiskANNIndex>) -> Self { + fn new(index: DiskANNIndex>) -> Self { Self { index, _vector_data_type: PhantomData, @@ -266,7 +277,7 @@ where /// /// # Errors /// Returns an error if the underlying index creation fails. -pub(super) fn new_inmem_index_builder( +fn new_inmem_index_builder( config: Config, params: DefaultProviderParameters, build_quantizer: &BuildQuantizer, @@ -288,3 +299,180 @@ where } } } + +pub(super) async fn build_inmem_index( + config: IndexConfiguration, + quantizer: &BuildQuantizer, + data_path: &str, + save_path: &str, + storage_provider: &StorageProvider, +) -> ANNResult<()> +where + T: VectorRepr, + StorageProvider: StorageReadProvider + StorageWriteProvider + 'static, + ::Reader: std::marker::Send, +{ + // use either user-specified number of threads or default to available parallelism + let num_tasks = NonZeroUsize::new(config.num_threads) + .or_else(|| std::thread::available_parallelism().ok()) + .ok_or_else(|| ANNError::log_index_error("Failed to determine number of threads"))?; + + // Associated data will only be used in the write_disk_layout function which only requires the none-partitioned associated data stream. + let dataset_iter = Arc::new(Mutex::new({ + let iter = VectorDataIterator::<_, T>::new(data_path, Option::None, storage_provider)?; + iter.enumerate() + })); + + let index_config = config.config.clone(); + let provider_parameters = DefaultProviderParameters { + max_points: config.max_points, + frozen_points: ONE, + metric: config.dist_metric, + dim: config.dim, + max_degree: index_config.max_degree_u32().get(), + prefetch_lookahead: config.prefetch_lookahead.map(|x| x.get()), + prefetch_cache_line_level: config.prefetch_cache_line_level, + }; + let index = new_inmem_index_builder::(index_config, provider_parameters, quantizer)?; + let medoid_id = + set_start_point_to_medoid::(&index, data_path, config.random_seed, storage_provider)?; + let start_point = u32_try_from(medoid_id)?; + + run_build(&index, dataset_iter, num_tasks).await?; + + #[cfg(debug_assertions)] + log_build_stats::<_>(&index).await?; + + run_final_prune(&index, num_tasks).await?; + index + .save_graph( + storage_provider, + &(start_point, DiskGraphOnly::new(save_path)), + ) + .await?; + + Ok(()) +} + +#[cfg(debug_assertions)] +/// Log statistics about the build process +async fn log_build_stats(index: &Arc>) -> ANNResult<()> { + debug!( + "Number of points reachable in the graph: {}", + index.count_reachable_nodes().await? + ); + + let (full_vector, quant_vector) = index.counts_for_get_vector(); + let capacity = index.capacity(); + debug!( + "Number of get vector calls per insert: {}", + full_vector as f32 / capacity as f32 + ); + debug!( + "Number of get quantized vector calls per insert: {}", + quant_vector as f32 / capacity as f32 + ); + + Ok(()) +} + +/// Convert a `usize` index into the `u32` internal id type, erroring if it does not fit. +/// +/// The in-memory index uses `u32` internal ids, so positions in the dataset must not exceed +/// `u32::MAX`. +fn u32_try_from(value: usize) -> ANNResult { + u32::try_from(value) + .map_err(|_| ANNError::log_index_error(format_args!("id {value} exceeds u32::MAX"))) +} + +fn set_start_point_to_medoid( + index: &Arc>, + path: &str, + random_seed: Option, + reader: &StorageReader, +) -> ANNResult +where + T: VectorRepr, + StorageReader: StorageReadProvider, +{ + let mut rng = diskann_providers::utils::create_rnd_from_optional_seed(random_seed); + let (medoid, medoid_id) = + find_medoid_with_sampling::(path, reader, MAX_MEDOID_SAMPLE_SIZE, &mut rng)?; + + index.set_start_point(medoid.as_slice())?; + + debug!("Set start point to medoid ID: {}", medoid_id); + + Ok(medoid_id) +} + +async fn run_build( + index: &Arc>, + iterator: Arc>, + num_tasks: NonZeroUsize, +) -> ANNResult<()> +where + T: VectorRepr, + I: Iterator, ()))> + Send + 'static, +{ + let total_points = index.capacity(); + let partitions = async_tools::PartitionIter::new(total_points, num_tasks); + + let mut tasks = JoinSet::new(); + + for partition in partitions { + let index_clone = index.clone(); + let iterator_clone = iterator.clone(); + tasks.spawn(async move { + for _ in partition { + let vector_data = { + let mut guard = iterator_clone.lock().map_err(|_| { + ANNError::log_index_error("Poisoned mutex during construction") + })?; + guard.next() + }; + + match vector_data { + Some((i, (vector, _))) => { + let id = u32_try_from(i)?; + index_clone.insert_vector(id, vector.as_ref()).await?; + } + None => break, + } + } + ANNResult::Ok(()) + }); + } + + // Wait for all tasks to complete. + while let Some(res) = tasks.join_next().await { + res.map_err(|_| ANNError::log_index_error("A spawned insert task failed"))??; + } + + info!("Linked all points. Num points: #{}", total_points); + Ok(()) +} + +async fn run_final_prune( + index: &Arc>, + num_tasks: NonZeroUsize, +) -> ANNResult<()> { + let partitions = async_tools::PartitionIter::new(index.total_points(), num_tasks); + + let mut tasks = JoinSet::new(); + + for partition in partitions { + let index_clone = index.clone(); + tasks.spawn(async move { + let range = u32_try_from(partition.start)?..u32_try_from(partition.end)?; + index_clone.final_prune(range).await + }); + } + + // Wait for all final prune tasks to complete + while let Some(res) = tasks.join_next().await { + res.map_err(|_| ANNError::log_index_error("A spawned final prune task failed"))??; + } + + Ok(()) +} diff --git a/diskann-disk/src/build/builder/core.rs b/diskann-disk/src/build/builder/merged_index.rs similarity index 99% rename from diskann-disk/src/build/builder/core.rs rename to diskann-disk/src/build/builder/merged_index.rs index 6f7e157d2..1c346606d 100644 --- a/diskann-disk/src/build/builder/core.rs +++ b/diskann-disk/src/build/builder/merged_index.rs @@ -22,7 +22,7 @@ use rand::seq::SliceRandom; use tracing::info; use crate::{ - build::builder::{build::build_inmem_index, quantizer::BuildQuantizer}, + build::builder::{inmem_builder::build_inmem_index, quantizer::BuildQuantizer}, disk_index_build_parameter::BYTES_IN_GB, storage::{CachedReader, CachedWriter, DiskIndexWriter}, utils::instrumentation::{BuildMergedVamanaIndexCheckpoint, PerfLogger}, diff --git a/diskann-disk/src/build/builder/mod.rs b/diskann-disk/src/build/builder/mod.rs index 9c22a0766..126e262d3 100644 --- a/diskann-disk/src/build/builder/mod.rs +++ b/diskann-disk/src/build/builder/mod.rs @@ -5,11 +5,11 @@ //! Disk index builders and related functionality. pub mod build; -pub mod core; +mod merged_index; pub mod quantizer; pub mod inmem_builder; pub mod tokio; #[cfg(test)] -mod tests; +pub(crate) mod tests; diff --git a/diskann-disk/src/build/builder/quantizer.rs b/diskann-disk/src/build/builder/quantizer.rs index 7bbfacf3c..d98fa48e7 100644 --- a/diskann-disk/src/build/builder/quantizer.rs +++ b/diskann-disk/src/build/builder/quantizer.rs @@ -21,7 +21,7 @@ use tracing::info; use crate::QuantizationType; -/// Quantizer types used specifically for async disk index building. +/// Quantizer types used for disk index building. #[derive(Clone)] pub enum BuildQuantizer { NoQuant(NoStore), diff --git a/diskann-disk/src/build/builder/tests.rs b/diskann-disk/src/build/builder/tests.rs index 8ce51a3c5..0d88de292 100644 --- a/diskann-disk/src/build/builder/tests.rs +++ b/diskann-disk/src/build/builder/tests.rs @@ -10,7 +10,7 @@ mod disk_index_build_tests { use rstest::rstest; use crate::{ - build::builder::core::disk_index_builder_tests::{ + build::builder::merged_index::disk_index_builder_tests::{ new_vfs, verify_search_result_with_ground_truth, IndexBuildFixture, TestParams, }, QuantizationType, diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index c497414a6..ffc78238f 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -1249,7 +1249,7 @@ mod disk_provider_tests { use super::*; use crate::{ - build::builder::core::disk_index_builder_tests::{IndexBuildFixture, TestParams}, + build::builder::merged_index::disk_index_builder_tests::{IndexBuildFixture, TestParams}, search::provider::aligned_file_reader::VirtualAlignedReaderFactory, utils::QueryStatistics, }; From 69ad47499d30de3f3b60bc534ed674db3925ab5f Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Tue, 21 Jul 2026 14:23:09 +0800 Subject: [PATCH 2/7] Move disk build strategy into builder Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ecf6179f-7467-44be-8a77-bf7a96e319af --- diskann-disk/src/build/builder/build.rs | 130 +++++++++++++++++- .../src/build/builder/merged_index.rs | 117 +--------------- 2 files changed, 126 insertions(+), 121 deletions(-) diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index 4e0098c6c..cecd09812 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -4,13 +4,16 @@ */ //! Disk index builder implementation. -use std::marker::PhantomData; +use std::{marker::PhantomData, mem}; use crate::data_model::GraphDataType; use diskann::{utils::VectorRepr, ANNResult}; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ - model::{IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS}, + model::{ + IndexConfiguration, GRAPH_SLACK_FACTOR, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, + NUM_PQ_CENTROIDS, + }, storage::PQStorage, utils::{create_thread_pool, RayonThreadPoolRef}, }; @@ -18,17 +21,16 @@ use tracing::info; use crate::{ build::builder::{ - inmem_builder::build_inmem_index, - merged_index::{determine_build_strategy, IndexBuildStrategy, MergedVamanaIndexBuilder}, - quantizer::BuildQuantizer, - tokio::create_runtime, + inmem_builder::build_inmem_index, merged_index::MergedVamanaIndexBuilder, + quantizer::BuildQuantizer, tokio::create_runtime, }, + disk_index_build_parameter::BYTES_IN_GB, storage::{ quant::{PQGeneration, PQGenerationContext, QuantDataGenerator}, DiskIndexWriter, }, utils::instrumentation::{DiskIndexBuildCheckpoint, PerfLogger}, - DiskIndexBuildParameters, + DiskIndexBuildParameters, QuantizationType, }; pub struct DiskIndexBuilder<'a, Data, StorageProvider> @@ -44,6 +46,7 @@ where build_quantizer: BuildQuantizer, _phantom: PhantomData, } + impl<'a, Data, StorageProvider> DiskIndexBuilder<'a, Data, StorageProvider> where Data: GraphDataType, @@ -195,3 +198,116 @@ where Ok(()) } } + +pub(crate) enum IndexBuildStrategy { + OneShot, + Merged, +} + +pub(crate) fn determine_build_strategy( + index_configuration: &IndexConfiguration, + index_build_ram_limit_in_bytes: f64, + build_quantization_type: &QuantizationType, +) -> IndexBuildStrategy { + let estimated_index_ram_in_bytes = estimate_build_index_ram_usage( + index_configuration.max_points as u64, + index_configuration.dim as u64, + mem::size_of::() as u64, + index_configuration.config.max_degree().get() as u64, + build_quantization_type, + ); + + info!( + "Estimated index RAM usage: {} GB, index_build_ram_limit={} GB", + estimated_index_ram_in_bytes / BYTES_IN_GB, + index_build_ram_limit_in_bytes / BYTES_IN_GB + ); + + if estimated_index_ram_in_bytes >= index_build_ram_limit_in_bytes { + info!( + "Insufficient memory budget for index build in one shot, index_build_ram_limit={} GB estimated_index_ram={} GB", + index_build_ram_limit_in_bytes / BYTES_IN_GB, + estimated_index_ram_in_bytes / BYTES_IN_GB, + ); + IndexBuildStrategy::Merged + } else { + info!( + "Full index fits in RAM budget, should consume at most {} GBs, so building in one shot", + estimated_index_ram_in_bytes / BYTES_IN_GB + ); + IndexBuildStrategy::OneShot + } +} + +/// Overhead factor for RAM estimation during index build (10% buffer). +const OVERHEAD_FACTOR: f64 = 1.1f64; + +/// Estimate RAM usage in bytes for building an index. +#[inline] +fn estimate_build_index_ram_usage( + num_points: u64, + dim: u64, + datasize: u64, + graph_degree: u64, + build_quantization_type: &QuantizationType, +) -> f64 { + let graph_size = (num_points * graph_degree * mem::size_of::() as u64) as f64 + * GRAPH_SLACK_FACTOR; + + let single_vec_size = match *build_quantization_type { + QuantizationType::FP => dim.next_multiple_of(8u64) * datasize, + // We can skip PQ pivots data as it is very small(~3MB) for even large datasets like OAI-3072. + QuantizationType::PQ { num_chunks } => num_chunks as u64, + // `+ std::mem::size_of::()` for f32 compensation metadata for the scalar quantizer. + QuantizationType::SQ { nbits, .. } => { + (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 + } + }; + + OVERHEAD_FACTOR * (graph_size + (single_vec_size * num_points) as f64) +} + +#[cfg(test)] +mod ram_estimation_tests { + use rstest::rstest; + + use super::*; + use crate::QuantizationType; + + #[rstest] + #[case(QuantizationType::FP)] + #[case(QuantizationType::PQ { num_chunks: 15 })] + #[case(QuantizationType::SQ { nbits: 1, standard_deviation: None })] + fn test_estimate_build_index_ram_usage( + #[case] build_quantization_type: QuantizationType, + ) { + let num_points = 1000; + let dim = 128; + let size_of_t = std::mem::size_of::() as u64; + let graph_degree = 50; + + let single_vec_size = match build_quantization_type { + QuantizationType::FP => dim * size_of_t, + QuantizationType::PQ { num_chunks } => num_chunks as u64, + QuantizationType::SQ { nbits, .. } => { + (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 + } + }; + let mut expected_ram_usage = (num_points as f64) + * (graph_degree as f64) + * (std::mem::size_of::() as f64) + * GRAPH_SLACK_FACTOR + + (num_points * single_vec_size) as f64; + expected_ram_usage *= OVERHEAD_FACTOR; + + let actual_ram_usage = estimate_build_index_ram_usage( + num_points, + dim, + size_of_t, + graph_degree, + &build_quantization_type, + ); + + assert_eq!(actual_ram_usage, expected_ram_usage); + } +} diff --git a/diskann-disk/src/build/builder/merged_index.rs b/diskann-disk/src/build/builder/merged_index.rs index 1c346606d..1048b15f6 100644 --- a/diskann-disk/src/build/builder/merged_index.rs +++ b/diskann-disk/src/build/builder/merged_index.rs @@ -11,7 +11,7 @@ use crate::data_model::GraphDataType; use diskann::{utils::VectorRepr, ANNResult}; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ - model::{IndexConfiguration, GRAPH_SLACK_FACTOR, MAX_PQ_TRAINING_SET_SIZE}, + model::{IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE}, utils::{ load_metadata_from_file, RayonThreadPoolRef, SampleVectorReader, SamplingDensity, READ_WRITE_BLOCK_SIZE, @@ -23,44 +23,15 @@ use tracing::info; use crate::{ build::builder::{inmem_builder::build_inmem_index, quantizer::BuildQuantizer}, - disk_index_build_parameter::BYTES_IN_GB, storage::{CachedReader, CachedWriter, DiskIndexWriter}, utils::instrumentation::{BuildMergedVamanaIndexCheckpoint, PerfLogger}, utils::partition_with_ram_budget, - DiskIndexBuildParameters, QuantizationType, + DiskIndexBuildParameters, }; -/// Overhead factor for RAM estimation during index build (10% buffer). -const OVERHEAD_FACTOR: f64 = 1.1f64; - /// Number of nearest shards each vector is assigned to during partitioning. const PARTITION_ASSIGNMENTS_PER_VECTOR: usize = 2; -/// Estimate RAM usage in bytes for building an index. -#[inline] -fn estimate_build_index_ram_usage( - num_points: u64, - dim: u64, - datasize: u64, - graph_degree: u64, - build_quantization_type: &QuantizationType, -) -> f64 { - let graph_size = - (num_points * graph_degree * mem::size_of::() as u64) as f64 * GRAPH_SLACK_FACTOR; - - let single_vec_size = match *build_quantization_type { - QuantizationType::FP => dim.next_multiple_of(8u64) * datasize, - // We can skip PQ pivots data as it is very small(~3MB) for even large datasets like OAI-3072. - QuantizationType::PQ { num_chunks } => num_chunks as u64, - // `+ std::mem::size_of::()` for f32 compensation metadata for the scalar quantizer. - QuantizationType::SQ { nbits, .. } => { - (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 - } - }; - - OVERHEAD_FACTOR * (graph_size + (single_vec_size * num_points) as f64) -} - /// Builds a merged Vamana index from overlapping dataset shards. pub(super) struct MergedVamanaIndexBuilder<'a, Data, StorageProvider> where @@ -495,46 +466,6 @@ where } } -pub(crate) enum IndexBuildStrategy { - OneShot, - Merged, -} - -pub(crate) fn determine_build_strategy( - index_configuration: &IndexConfiguration, - index_build_ram_limit_in_bytes: f64, - build_quantization_type: &QuantizationType, -) -> IndexBuildStrategy { - let estimated_index_ram_in_bytes = estimate_build_index_ram_usage( - index_configuration.max_points as u64, - index_configuration.dim as u64, - mem::size_of::() as u64, - index_configuration.config.max_degree().get() as u64, - build_quantization_type, - ); - - info!( - "Estimated index RAM usage: {} GB, index_build_ram_limit={} GB", - estimated_index_ram_in_bytes / BYTES_IN_GB, - index_build_ram_limit_in_bytes / BYTES_IN_GB - ); - - if estimated_index_ram_in_bytes >= index_build_ram_limit_in_bytes { - info!( - "Insufficient memory budget for index build in one shot, index_build_ram_limit={} GB estimated_index_ram={} GB", - index_build_ram_limit_in_bytes / BYTES_IN_GB, - estimated_index_ram_in_bytes / BYTES_IN_GB, - ); - IndexBuildStrategy::Merged - } else { - info!( - "Full index fits in RAM budget, should consume at most {} GBs, so building in one shot", - estimated_index_ram_in_bytes / BYTES_IN_GB - ); - IndexBuildStrategy::OneShot - } -} - #[cfg(test)] pub(crate) mod disk_index_builder_tests { use std::{io::Read, sync::Arc, time::Instant}; @@ -569,6 +500,7 @@ pub(crate) mod disk_index_builder_tests { }, storage::disk_index_reader::DiskIndexReader, utils::QueryStatistics, + QuantizationType, }; const DEFAULT_DISK_SECTOR_LEN: usize = 4096; pub const TEST_DATA_FILE: &str = "/sift/siftsmall_learn_256pts.fbin"; @@ -1155,46 +1087,3 @@ pub(crate) mod disk_index_builder_tests { ) } } - -#[cfg(test)] -mod ram_estimation_tests { - use rstest::rstest; - - use super::*; - use crate::QuantizationType; - - #[rstest] - #[case(QuantizationType::FP)] - #[case(QuantizationType::PQ { num_chunks: 15 })] - #[case(QuantizationType::SQ { nbits: 1, standard_deviation: None })] - fn test_estimate_build_index_ram_usage(#[case] build_quantization_type: QuantizationType) { - let num_points = 1000; - let dim = 128; - let size_of_t = std::mem::size_of::() as u64; - let graph_degree = 50; - - let single_vec_size = match build_quantization_type { - QuantizationType::FP => dim * size_of_t, - QuantizationType::PQ { num_chunks } => num_chunks as u64, - QuantizationType::SQ { nbits, .. } => { - (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 - } - }; - let mut expected_ram_usage = (num_points as f64) - * (graph_degree as f64) - * (std::mem::size_of::() as f64) - * GRAPH_SLACK_FACTOR - + (num_points * single_vec_size) as f64; - expected_ram_usage *= OVERHEAD_FACTOR; - - let actual_ram_usage = estimate_build_index_ram_usage( - num_points, - dim, - size_of_t, - graph_degree, - &build_quantization_type, - ); - - assert_eq!(actual_ram_usage, expected_ram_usage); - } -} From 4d4061eac58a57c8ea64eb844c99b604b016f109 Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Wed, 22 Jul 2026 13:47:20 +0800 Subject: [PATCH 3/7] Refactor disk index builder tests and update module paths - Moved disk index builder tests to a new module structure for better organization. - Updated import paths in `disk_provider.rs` to reflect the new location of `IndexBuildFixture` and `TestParams`. - Enhanced the test suite for disk index building, including additional test cases and improved parameter handling. --- diskann-disk/src/build/builder/build.rs | 84 +-- .../src/build/builder/merged_index.rs | 671 ++---------------- diskann-disk/src/build/builder/tests.rs | 627 +++++++++++++++- .../src/search/provider/disk_provider.rs | 2 +- 4 files changed, 689 insertions(+), 695 deletions(-) diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index cecd09812..d898e7c9d 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -10,10 +10,7 @@ use crate::data_model::GraphDataType; use diskann::{utils::VectorRepr, ANNResult}; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ - model::{ - IndexConfiguration, GRAPH_SLACK_FACTOR, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, - NUM_PQ_CENTROIDS, - }, + model::{IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS}, storage::PQStorage, utils::{create_thread_pool, RayonThreadPoolRef}, }; @@ -21,8 +18,10 @@ use tracing::info; use crate::{ build::builder::{ - inmem_builder::build_inmem_index, merged_index::MergedVamanaIndexBuilder, - quantizer::BuildQuantizer, tokio::create_runtime, + inmem_builder::build_inmem_index, + merged_index::{estimate_build_index_ram_usage, MergedVamanaIndexBuilder}, + quantizer::BuildQuantizer, + tokio::create_runtime, }, disk_index_build_parameter::BYTES_IN_GB, storage::{ @@ -238,76 +237,3 @@ pub(crate) fn determine_build_strategy( IndexBuildStrategy::OneShot } } - -/// Overhead factor for RAM estimation during index build (10% buffer). -const OVERHEAD_FACTOR: f64 = 1.1f64; - -/// Estimate RAM usage in bytes for building an index. -#[inline] -fn estimate_build_index_ram_usage( - num_points: u64, - dim: u64, - datasize: u64, - graph_degree: u64, - build_quantization_type: &QuantizationType, -) -> f64 { - let graph_size = (num_points * graph_degree * mem::size_of::() as u64) as f64 - * GRAPH_SLACK_FACTOR; - - let single_vec_size = match *build_quantization_type { - QuantizationType::FP => dim.next_multiple_of(8u64) * datasize, - // We can skip PQ pivots data as it is very small(~3MB) for even large datasets like OAI-3072. - QuantizationType::PQ { num_chunks } => num_chunks as u64, - // `+ std::mem::size_of::()` for f32 compensation metadata for the scalar quantizer. - QuantizationType::SQ { nbits, .. } => { - (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 - } - }; - - OVERHEAD_FACTOR * (graph_size + (single_vec_size * num_points) as f64) -} - -#[cfg(test)] -mod ram_estimation_tests { - use rstest::rstest; - - use super::*; - use crate::QuantizationType; - - #[rstest] - #[case(QuantizationType::FP)] - #[case(QuantizationType::PQ { num_chunks: 15 })] - #[case(QuantizationType::SQ { nbits: 1, standard_deviation: None })] - fn test_estimate_build_index_ram_usage( - #[case] build_quantization_type: QuantizationType, - ) { - let num_points = 1000; - let dim = 128; - let size_of_t = std::mem::size_of::() as u64; - let graph_degree = 50; - - let single_vec_size = match build_quantization_type { - QuantizationType::FP => dim * size_of_t, - QuantizationType::PQ { num_chunks } => num_chunks as u64, - QuantizationType::SQ { nbits, .. } => { - (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 - } - }; - let mut expected_ram_usage = (num_points as f64) - * (graph_degree as f64) - * (std::mem::size_of::() as f64) - * GRAPH_SLACK_FACTOR - + (num_points * single_vec_size) as f64; - expected_ram_usage *= OVERHEAD_FACTOR; - - let actual_ram_usage = estimate_build_index_ram_usage( - num_points, - dim, - size_of_t, - graph_degree, - &build_quantization_type, - ); - - assert_eq!(actual_ram_usage, expected_ram_usage); - } -} diff --git a/diskann-disk/src/build/builder/merged_index.rs b/diskann-disk/src/build/builder/merged_index.rs index 1048b15f6..0a4cf84d9 100644 --- a/diskann-disk/src/build/builder/merged_index.rs +++ b/diskann-disk/src/build/builder/merged_index.rs @@ -11,7 +11,7 @@ use crate::data_model::GraphDataType; use diskann::{utils::VectorRepr, ANNResult}; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ - model::{IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE}, + model::{IndexConfiguration, GRAPH_SLACK_FACTOR, MAX_PQ_TRAINING_SET_SIZE}, utils::{ load_metadata_from_file, RayonThreadPoolRef, SampleVectorReader, SamplingDensity, READ_WRITE_BLOCK_SIZE, @@ -26,12 +26,40 @@ use crate::{ storage::{CachedReader, CachedWriter, DiskIndexWriter}, utils::instrumentation::{BuildMergedVamanaIndexCheckpoint, PerfLogger}, utils::partition_with_ram_budget, - DiskIndexBuildParameters, + DiskIndexBuildParameters, QuantizationType, }; +/// Overhead factor for RAM estimation during index build (10% buffer). +const OVERHEAD_FACTOR: f64 = 1.1f64; + /// Number of nearest shards each vector is assigned to during partitioning. const PARTITION_ASSIGNMENTS_PER_VECTOR: usize = 2; +/// Estimate RAM usage in bytes for building an index. +#[inline] +pub(super) fn estimate_build_index_ram_usage( + num_points: u64, + dim: u64, + datasize: u64, + graph_degree: u64, + build_quantization_type: &QuantizationType, +) -> f64 { + let graph_size = + (num_points * graph_degree * mem::size_of::() as u64) as f64 * GRAPH_SLACK_FACTOR; + + let single_vec_size = match *build_quantization_type { + QuantizationType::FP => dim.next_multiple_of(8u64) * datasize, + // We can skip PQ pivots data as it is very small(~3MB) for even large datasets like OAI-3072. + QuantizationType::PQ { num_chunks } => num_chunks as u64, + // `+ std::mem::size_of::()` for f32 compensation metadata for the scalar quantizer. + QuantizationType::SQ { nbits, .. } => { + (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 + } + }; + + OVERHEAD_FACTOR * (graph_size + (single_vec_size * num_points) as f64) +} + /// Builds a merged Vamana index from overlapping dataset shards. pub(super) struct MergedVamanaIndexBuilder<'a, Data, StorageProvider> where @@ -467,623 +495,44 @@ where } #[cfg(test)] -pub(crate) mod disk_index_builder_tests { - use std::{io::Read, sync::Arc, time::Instant}; - - use crate::test_utils::{GraphDataF32VectorU32Data, GraphDataF32VectorUnitData}; - use diskann::{ - graph::config, - utils::{IntoUsize, VectorRepr, ONE}, - ANNResult, - }; - use diskann_providers::storage::VirtualStorageProvider; - use diskann_providers::storage::{ - get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, - }; - - use diskann_utils::test_data_root; - use diskann_vector::{ - distance::Metric::{self, L2}, - DistanceFunction, - }; +mod ram_estimation_tests { use rstest::rstest; - use vfs::OverlayFS; use super::*; - use crate::{ - build::builder::build::DiskIndexBuilder, - data_model::{CachingStrategy, GraphHeader}, - disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}, - search::provider::{ - aligned_file_reader::VirtualAlignedReaderFactory, disk_provider::DiskIndexSearcher, - disk_vertex_provider_factory::DiskVertexProviderFactory, - }, - storage::disk_index_reader::DiskIndexReader, - utils::QueryStatistics, - QuantizationType, - }; - const DEFAULT_DISK_SECTOR_LEN: usize = 4096; - pub const TEST_DATA_FILE: &str = "/sift/siftsmall_learn_256pts.fbin"; - /// We can use the same index prefix for all tests since we use virtual storage provider - const INDEX_PATH_PREFIX: &str = "/disk_index_build/sift_learn_test_disk_index_build"; - const TRUTH_INDEX_PATH_PREFIX_R4_L50: &str = "/disk_index_build/truth_sift_learn_R4_L50"; - - pub struct TestParams { - pub dim: usize, - pub full_dim: usize, - pub max_degree: u32, - pub num_pq_chunks: usize, - pub build_quantization_type: QuantizationType, - pub l_build: u32, - pub data_path: String, - pub index_path_prefix: String, - pub associated_data_path: Option, - pub index_build_ram_gb: f64, - pub data_compression_chunk_vector_count: Option, - pub num_threads: usize, - pub metric: Metric, - pub block_size: usize, - } - - impl Default for TestParams { - fn default() -> Self { - Self { - dim: 128, // D - full_dim: 128, - max_degree: 4, // R - num_pq_chunks: 128, - build_quantization_type: QuantizationType::FP, // No quantization, i.e. QuantizationType::FP - l_build: 50, - data_path: TEST_DATA_FILE.to_string(), - index_path_prefix: INDEX_PATH_PREFIX.to_string(), - associated_data_path: None, - index_build_ram_gb: 1.0, - data_compression_chunk_vector_count: None, - num_threads: 1, - metric: L2, - block_size: DEFAULT_DISK_SECTOR_LEN, - } - } - } - - impl TestParams { - /// Returns the appropriate truth index path prefix for build comparison. - fn truth_index_path_prefix(&self) -> &str { - match (self.max_degree, self.l_build, self.index_build_ram_gb) { - (4, 50, 1.0) => TRUTH_INDEX_PATH_PREFIX_R4_L50, - (max_degree, l_build, index_build_ram_gb) => panic!( - "Truth index path not found for max_degree={}, l_build={}, index_build_ram_gb={}", - max_degree, l_build, index_build_ram_gb - ), - } - } - pub fn truth_pq_compressed_path(&self) -> String { - let prefix = match self.num_pq_chunks { - 128 => TRUTH_INDEX_PATH_PREFIX_R4_L50, - num_pq_chunks => panic!( - "Truth pq compressed path not found for num_pq_chunks={}", - num_pq_chunks, - ), - }; - get_compressed_pq_file(prefix) - } - - pub fn pq_compressed_path(&self) -> String { - get_compressed_pq_file(&self.index_path_prefix) - } - } - - pub fn new_vfs() -> VirtualStorageProvider { - VirtualStorageProvider::new_overlay(test_data_root()) - } - - pub struct IndexBuildFixture { - pub storage_provider: Arc, - pub params: TestParams, - } - - impl - IndexBuildFixture - { - pub fn new(storage_provider: StorageProvider, params: TestParams) -> ANNResult { - Ok(Self { - storage_provider: Arc::new(storage_provider), - params, - }) - } - - pub fn build(&self) -> ANNResult<()> - where - T: GraphDataType, - StorageProvider::Reader: std::marker::Send + Read, - { - // Create disk index build parameters - let mut disk_index_build_parameters = DiskIndexBuildParameters::new( - MemoryBudget::try_from_gb(self.params.index_build_ram_gb)?, - self.params.build_quantization_type, - NumPQChunks::new_with(self.params.num_pq_chunks, self.params.full_dim)?, - ); - if let Some(chunk_vector_count) = self.params.data_compression_chunk_vector_count { - disk_index_build_parameters = disk_index_build_parameters - .with_data_compression_chunk_vector_count(chunk_vector_count); - } - - let config = config::Builder::new_with( - self.params.max_degree.into_usize(), - config::MaxDegree::default_slack(), - self.params.l_build.into_usize(), - self.params.metric.into(), - |b| { - b.saturate_after_prune(true); - }, - ) - .build()?; - - let metadata = - load_metadata_from_file(self.storage_provider.as_ref(), &self.params.data_path) - .unwrap(); - - assert_eq!( - self.params.dim, - metadata.ndims(), - "Parameters dimension {} and data dimension {} are not equal", - self.params.dim, - metadata.ndims(), - ); - - let config = IndexConfiguration::new( - self.params.metric, - self.params.dim, - metadata.npoints(), - ONE, - self.params.num_threads, - config, - ) - .with_pseudo_rng_from_seed(100); - - let disk_index_writer = DiskIndexWriter::new( - self.params.data_path.clone(), - self.params.index_path_prefix.clone(), - self.params.associated_data_path.clone(), - self.params.block_size, - )?; - - let mut disk_index = DiskIndexBuilder::::new( - self.storage_provider.as_ref(), - disk_index_build_parameters, - config, - disk_index_writer, - )?; - - let timer = Instant::now(); - disk_index.build()?; - println!("Indexing time: {} seconds", timer.elapsed().as_secs_f64()); - - Ok(()) - } - - pub fn compare_pq_compressed_files(&self) { - self.compare_files( - &self.params.pq_compressed_path(), - &self.params.truth_pq_compressed_path(), - ); - } - - pub fn assert_index_max_degree(&self) -> ANNResult<()> { - let index_file_path = get_disk_index_file(&self.params.index_path_prefix); - let file_data = load_file_to_vec(self.storage_provider.as_ref(), &index_file_path); - let graph_header = GraphHeader::try_from(&file_data[8..])?; - let max_degree = graph_header.max_degree::()?; - assert_eq!( - max_degree, self.params.max_degree as usize, - "Max degree mismatch: expected {}, got {}", - self.params.max_degree, max_degree - ); - - Ok(()) - } - - fn compare_disk_index_with_associated_data( - &self, - pivot_file_prefix_test: &str, - pivot_file_prefix_expected: &str, - index_file_suffix: &str, - ) { - let pq_pivot_path = pivot_file_prefix_test.to_string() + index_file_suffix; - let pq_pivot_path_truth = pivot_file_prefix_expected.to_string() + index_file_suffix; - let file1 = load_file_to_vec(self.storage_provider.as_ref(), &pq_pivot_path); - let file2 = load_file_to_vec(self.storage_provider.as_ref(), &pq_pivot_path_truth); - compare_disk_index_graphs(&file1, &file2) - } - - pub fn compare_files(&self, file_path1: &str, file_path2: &str) { - let file1 = load_file_to_vec(self.storage_provider.as_ref(), file_path1); - let file2 = load_file_to_vec(self.storage_provider.as_ref(), file_path2); - - assert_eq!(file1.len(), file2.len()); - assert_eq!(file1, file2) - } - } - - /// Common helper function for one-shot async index build tests - fn run_one_shot_test(index_path_prefix: String, params_customizer: F) - where - F: FnOnce(TestParams) -> TestParams, - { - let l_build = 64; - let max_degree = 16; - let top_k = 10; - let search_l = 32; - - let base_params = TestParams { - l_build, - max_degree, - index_path_prefix, - ..TestParams::default() - }; - - let params = params_customizer(base_params); - - let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); - fixture.build::().unwrap(); - - // Validate search recall against ground truth for async tests - verify_search_result_with_ground_truth::( - &fixture.params, - top_k, - search_l, - &fixture.storage_provider, - ) - .unwrap(); - - fixture - .assert_index_max_degree::() - .unwrap(); - - // Assert that all data was kept in memory and no files were written to the disk. - let mem_index_file_path = format!("{}_mem.index.data", fixture.params.index_path_prefix); - assert!(!fixture.storage_provider.exists(&mem_index_file_path)); - } + use crate::QuantizationType; #[rstest] - fn test_build_from_iter_one_shot_with_metric( - #[values(Metric::L2, Metric::InnerProduct, Metric::Cosine)] metric: Metric, - ) { - let index_path_prefix = format!("{}_metric_{:?}", INDEX_PATH_PREFIX, metric); - - run_one_shot_test(index_path_prefix, |params| TestParams { metric, ..params }); - } - - /// Forces the multi-sector-per-node layout: with a 512-byte block, a 128-d f32 vector - /// (512 B) plus its neighbor list exceeds one sector, so each node spans multiple - /// sectors (`node_len > block_size`). The default 4096-byte sector packs these small - /// nodes many-per-sector, leaving the crossing path otherwise untested; 512 is one of - /// the two block sizes the disk format supports. - #[test] - fn test_build_multi_sector_per_node() { - let params = TestParams { - block_size: 512, - max_degree: 16, - l_build: 64, - index_path_prefix: format!("{}_multi_sector", INDEX_PATH_PREFIX), - ..TestParams::default() - }; - let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); - fixture.build::().unwrap(); - verify_search_result_with_ground_truth::( - &fixture.params, - 10, - 32, - &fixture.storage_provider, - ) - .unwrap(); - } - - #[test] - fn test_build_from_iter_one_shot_with_associated_data() { - // Set up test data - let params = TestParams { - associated_data_path: Some( - "/sift/siftsmall_learn_256pts_u32_associated_data.fbin".to_string(), - ), - ..TestParams::default() - }; - - // Create fixture with virtual storage provider - let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); - - // Build the index with the associated data - fixture.build::().unwrap(); - - // Assert that all data was kept in memory and no files were written to the disk. - let mem_index_file_path = format!("{}_mem.index.data", fixture.params.index_path_prefix); - let mem_index_associated_data_path = format!( - "{}_mem.index.associated_data", - fixture.params.index_path_prefix - ); - assert!(!fixture.storage_provider.exists(&mem_index_file_path)); - assert!(!fixture - .storage_provider - .exists(&mem_index_associated_data_path)); - - // assert index files are expected. - fixture.compare_disk_index_with_associated_data( - &fixture.params.index_path_prefix, - fixture.params.truth_index_path_prefix(), - "_disk.index", - ); - } - - #[test] - fn test_build_from_iter_merged_index() { - // Use the same parameters from [test_sift_build_and_search] in diskann_index - let l_build = 64; - let max_degree = 16; - let top_k = 10; - let search_l = 32; - - let index_path_prefix = - "/disk_index_build/disk_index_sift_learn_test_disk_index_build_merged".to_string(); - let params = TestParams { - l_build, - max_degree, - index_path_prefix, - index_build_ram_gb: 0.0001, // small enough to trigger merged index build - ..TestParams::default() + #[case(QuantizationType::FP)] + #[case(QuantizationType::PQ { num_chunks: 15 })] + #[case(QuantizationType::SQ { nbits: 1, standard_deviation: None })] + fn test_estimate_build_index_ram_usage(#[case] build_quantization_type: QuantizationType) { + let num_points = 1000; + let dim = 128; + let size_of_t = std::mem::size_of::() as u64; + let graph_degree = 50; + + let single_vec_size = match build_quantization_type { + QuantizationType::FP => dim * size_of_t, + QuantizationType::PQ { num_chunks } => num_chunks as u64, + QuantizationType::SQ { nbits, .. } => { + (nbits as u64 * dim).div_ceil(8) + std::mem::size_of::() as u64 + } }; - - let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); - - fixture.build::().unwrap(); - - verify_search_result_with_ground_truth::( - &fixture.params, - top_k, - search_l, - &fixture.storage_provider, - ) - .unwrap(); - - fixture - .assert_index_max_degree::() - .unwrap(); - } - - #[rstest] - #[case(QuantizationType::SQ { nbits: 2, standard_deviation: None }, "SQ quantization is only supported for 1 bit")] - fn test_build_quantization_type_failure_cases( - #[case] build_quantization_type: QuantizationType, - #[case] error_message: &str, - ) { - let storage_provider = VirtualStorageProvider::new_overlay(test_data_root()); - let disk_index_builder = create_disk_index_builder( - 1000, // num_points - 128, // dim - 128, // num_pq_chunks - &storage_provider, - build_quantization_type, - ); - - let err = disk_index_builder.err().unwrap(); - assert!(err.to_string().contains(error_message)); - } - - fn load_file_to_vec( - storage_provider: &StorageType, - file_path: &str, - ) -> Vec { - let mut file = storage_provider.open_reader(file_path).unwrap(); - let mut buffer = vec![]; - file.read_to_end(&mut buffer).unwrap(); - buffer - } - - /// Verifies that search results exactly match the ground truth of nearest neighbors - /// - /// This function performs validation of search results by: - /// 1. Running searches on the index using actual data points from the dataset as queries - /// 2. Computing the exact ground truth results using direct distance calculations - /// 3. Verifying that the search engine returns precisely the same results as the ground truth - pub(crate) fn verify_search_result_with_ground_truth< - G: GraphDataType, - >( - params: &TestParams, - top_k: usize, - search_l: u32, - storage_provider: &Arc>, - ) -> ANNResult<()> { - let pq_pivot_path = get_pq_pivot_file(¶ms.index_path_prefix); - let pq_compressed_path = get_compressed_pq_file(¶ms.index_path_prefix); - let index_file_path = get_disk_index_file(¶ms.index_path_prefix); - - let index_reader = - DiskIndexReader::new(pq_pivot_path, pq_compressed_path, storage_provider.as_ref())?; - - let vertex_provider_factory = DiskVertexProviderFactory::new( - VirtualAlignedReaderFactory::new(index_file_path, Arc::clone(storage_provider)), - CachingStrategy::None, - )?; - - let search_engine = DiskIndexSearcher::>::new( - 1, - u32::MAX as usize, - &index_reader, - vertex_provider_factory, - params.metric, - None, - )?; - - let data = - read_bin::(&mut storage_provider.open_reader(¶ms.data_path)?)?; - let dim = data.ncols(); - let distance = ::distance(params.metric, Some(dim)); - - // Here, we use elements of the dataset to search the dataset itself. - // - // We do this for each query, computing the expected ground truth and verifying - // that our simple graph search matches. - // - // Because this dataset is small, we can expect exact equality. - for (q, query_data) in data.row_iter().enumerate() { - let gt = - diskann_providers::test_utils::groundtruth(data.as_view(), query_data, |a, b| { - distance.evaluate_similarity(a, b) - }); - - let mut query_stats = QueryStatistics::default(); - - let mut indices = vec![0u32; top_k]; - let mut distances = vec![0f32; top_k]; - let mut associated_data = vec![(); top_k]; - - _ = search_engine.search_internal( - query_data, - top_k, - search_l, - None, // beam_width - &mut query_stats, - &mut indices, - &mut distances, - &mut associated_data, - &crate::search::search_mode::SearchMode::graph(), - ); - - diskann_providers::test_utils::assert_top_k_exactly_match( - q, >, &indices, &distances, top_k, - ); - } - - Ok(()) - } - - // Compare that the index built in test is the same as the truth index. The truth index doesn't have associated data, we are only comparing the vector and neighbor data. - pub fn compare_disk_index_graphs(graph_data: &[u8], truth_graph_data: &[u8]) { - let graph_header = GraphHeader::try_from(&graph_data[8..]).unwrap(); - let truth_graph_header = GraphHeader::try_from(&truth_graph_data[8..]).unwrap(); - - let test_node_per_block = graph_header.metadata().num_nodes_per_block; - let test_max_node_length = graph_header.metadata().node_len; - - let truth_node_per_block = truth_graph_header.metadata().num_nodes_per_block; - let truth_max_node_length = truth_graph_header.metadata().node_len; - - assert_eq!( - graph_header.metadata().num_pts, - truth_graph_header.metadata().num_pts - ); - - assert_eq!( - graph_header.metadata().dims, - truth_graph_header.metadata().dims - ); - - let num_pts = graph_header.metadata().num_pts as usize; - let dim = graph_header.metadata().dims; - - for idx in 0..num_pts { - let test_node_id_offset = node_data_offset( - idx, - test_max_node_length as usize, - test_node_per_block as usize, - DEFAULT_DISK_SECTOR_LEN, - ); - - let truth_node_id_offset = node_data_offset( - idx, - truth_max_node_length as usize, - truth_node_per_block as usize, - DEFAULT_DISK_SECTOR_LEN, - ); - - // Assert that the vector data is the same between the test and truth graphs for this node. - assert_eq!( - &graph_data - [test_node_id_offset..test_node_id_offset + dim * std::mem::size_of::()], - &truth_graph_data - [truth_node_id_offset..truth_node_id_offset + dim * std::mem::size_of::()] - ); - - // Assert that the neighbor count is the same between the test and truth graphs for this node. - let test_nbr_cnt_offset = test_node_id_offset + dim * std::mem::size_of::(); - let truth_nbr_cnt_offset = truth_node_id_offset + dim * std::mem::size_of::(); - - let test_nbr_count = u32::from_le_bytes([ - graph_data[test_nbr_cnt_offset], - graph_data[test_nbr_cnt_offset + 1], - graph_data[test_nbr_cnt_offset + 2], - graph_data[test_nbr_cnt_offset + 3], - ]); - - let truth_nbr_count = u32::from_le_bytes([ - truth_graph_data[truth_nbr_cnt_offset], - truth_graph_data[truth_nbr_cnt_offset + 1], - truth_graph_data[truth_nbr_cnt_offset + 2], - truth_graph_data[truth_nbr_cnt_offset + 3], - ]); - - assert_eq!(test_nbr_count, truth_nbr_count); - - // Assert the neighbors (u32) are the same between the test and truth graphs for this node. - let test_nbr_offset = test_nbr_cnt_offset + 4; - let truth_nbr_offset = truth_nbr_cnt_offset + 4; - assert_eq!( - graph_data[test_nbr_offset..test_nbr_offset + test_nbr_count as usize * 4], - truth_graph_data[truth_nbr_offset..truth_nbr_offset + truth_nbr_count as usize * 4] - ); - } - } - - pub fn node_data_offset( - node_id: usize, - node_length: usize, - nodes_per_block: usize, - block_size: usize, - ) -> usize { - let block_id = node_id / nodes_per_block; - let node_id_in_block = node_id % nodes_per_block; - let offset = block_id * block_size + node_id_in_block * node_length; - offset + block_size - } - - fn create_disk_index_builder( - num_points: usize, - dim: usize, - num_of_pq_chunks: usize, - storage_provider: &VirtualStorageProvider, - build_quantization_type: QuantizationType, - ) -> ANNResult< - DiskIndexBuilder<'_, GraphDataF32VectorUnitData, VirtualStorageProvider>, - > { - let memory_budget = MemoryBudget::try_from_gb(1.0)?; - let num_pq_chunks = NumPQChunks::new_with(num_of_pq_chunks, dim)?; - - let build_parameters = - DiskIndexBuildParameters::new(memory_budget, build_quantization_type, num_pq_chunks); - - let index_configuration = IndexConfiguration::new( - L2, - dim, + let mut expected_ram_usage = (num_points as f64) + * (graph_degree as f64) + * (std::mem::size_of::() as f64) + * GRAPH_SLACK_FACTOR + + (num_points * single_vec_size) as f64; + expected_ram_usage *= OVERHEAD_FACTOR; + + let actual_ram_usage = estimate_build_index_ram_usage( num_points, - ONE, - 1, - config::Builder::new_with(4, config::MaxDegree::default_slack(), 50, L2.into(), |b| { - b.saturate_after_prune(true); - }) - .build()?, + dim, + size_of_t, + graph_degree, + &build_quantization_type, ); - let disk_index_writer = DiskIndexWriter::new( - "data_path".to_string(), - "index_path_prefix".to_string(), - None, - DEFAULT_DISK_SECTOR_LEN, - )?; - - DiskIndexBuilder::>::new( - storage_provider, - build_parameters, - index_configuration, - disk_index_writer, - ) + assert_eq!(actual_ram_usage, expected_ram_usage); } } diff --git a/diskann-disk/src/build/builder/tests.rs b/diskann-disk/src/build/builder/tests.rs index 0d88de292..f01957caf 100644 --- a/diskann-disk/src/build/builder/tests.rs +++ b/diskann-disk/src/build/builder/tests.rs @@ -4,17 +4,636 @@ */ //! Disk index builder tests. +use crate::{data_model::GraphDataType, storage::DiskIndexWriter}; +use diskann_providers::{ + model::IndexConfiguration, + storage::{StorageReadProvider, StorageWriteProvider}, + utils::load_metadata_from_file, +}; +use diskann_utils::io::read_bin; + #[cfg(test)] -mod disk_index_build_tests { - use crate::test_utils::{GraphDataF32VectorUnitData, GraphDataMinMaxVectorUnitData}; +pub(crate) mod disk_index_builder_tests { + use std::{io::Read, sync::Arc, time::Instant}; + + use crate::test_utils::{ + GraphDataF32VectorU32Data, GraphDataF32VectorUnitData, GraphDataMinMaxVectorUnitData, + }; + use diskann::{ + graph::config, + utils::{IntoUsize, VectorRepr, ONE}, + ANNResult, + }; + use diskann_providers::storage::VirtualStorageProvider; + use diskann_providers::storage::{ + get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, + }; + + use diskann_utils::test_data_root; + use diskann_vector::{ + distance::Metric::{self, L2}, + DistanceFunction, + }; use rstest::rstest; + use vfs::OverlayFS; + use super::*; use crate::{ - build::builder::merged_index::disk_index_builder_tests::{ - new_vfs, verify_search_result_with_ground_truth, IndexBuildFixture, TestParams, + build::builder::build::DiskIndexBuilder, + data_model::{CachingStrategy, GraphHeader}, + disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}, + search::provider::{ + aligned_file_reader::VirtualAlignedReaderFactory, disk_provider::DiskIndexSearcher, + disk_vertex_provider_factory::DiskVertexProviderFactory, }, + storage::disk_index_reader::DiskIndexReader, + utils::QueryStatistics, QuantizationType, }; + const DEFAULT_DISK_SECTOR_LEN: usize = 4096; + pub const TEST_DATA_FILE: &str = "/sift/siftsmall_learn_256pts.fbin"; + /// We can use the same index prefix for all tests since we use virtual storage provider + const INDEX_PATH_PREFIX: &str = "/disk_index_build/sift_learn_test_disk_index_build"; + const TRUTH_INDEX_PATH_PREFIX_R4_L50: &str = "/disk_index_build/truth_sift_learn_R4_L50"; + + pub struct TestParams { + pub dim: usize, + pub full_dim: usize, + pub max_degree: u32, + pub num_pq_chunks: usize, + pub build_quantization_type: QuantizationType, + pub l_build: u32, + pub data_path: String, + pub index_path_prefix: String, + pub associated_data_path: Option, + pub index_build_ram_gb: f64, + pub data_compression_chunk_vector_count: Option, + pub num_threads: usize, + pub metric: Metric, + pub block_size: usize, + } + + impl Default for TestParams { + fn default() -> Self { + Self { + dim: 128, // D + full_dim: 128, + max_degree: 4, // R + num_pq_chunks: 128, + build_quantization_type: QuantizationType::FP, // No quantization, i.e. QuantizationType::FP + l_build: 50, + data_path: TEST_DATA_FILE.to_string(), + index_path_prefix: INDEX_PATH_PREFIX.to_string(), + associated_data_path: None, + index_build_ram_gb: 1.0, + data_compression_chunk_vector_count: None, + num_threads: 1, + metric: L2, + block_size: DEFAULT_DISK_SECTOR_LEN, + } + } + } + + impl TestParams { + /// Returns the appropriate truth index path prefix for build comparison. + fn truth_index_path_prefix(&self) -> &str { + match (self.max_degree, self.l_build, self.index_build_ram_gb) { + (4, 50, 1.0) => TRUTH_INDEX_PATH_PREFIX_R4_L50, + (max_degree, l_build, index_build_ram_gb) => panic!( + "Truth index path not found for max_degree={}, l_build={}, index_build_ram_gb={}", + max_degree, l_build, index_build_ram_gb + ), + } + } + pub fn truth_pq_compressed_path(&self) -> String { + let prefix = match self.num_pq_chunks { + 128 => TRUTH_INDEX_PATH_PREFIX_R4_L50, + num_pq_chunks => panic!( + "Truth pq compressed path not found for num_pq_chunks={}", + num_pq_chunks, + ), + }; + get_compressed_pq_file(prefix) + } + + pub fn pq_compressed_path(&self) -> String { + get_compressed_pq_file(&self.index_path_prefix) + } + } + + pub fn new_vfs() -> VirtualStorageProvider { + VirtualStorageProvider::new_overlay(test_data_root()) + } + + pub struct IndexBuildFixture { + pub storage_provider: Arc, + pub params: TestParams, + } + + impl + IndexBuildFixture + { + pub fn new(storage_provider: StorageProvider, params: TestParams) -> ANNResult { + Ok(Self { + storage_provider: Arc::new(storage_provider), + params, + }) + } + + pub fn build(&self) -> ANNResult<()> + where + T: GraphDataType, + StorageProvider::Reader: std::marker::Send + Read, + { + // Create disk index build parameters + let mut disk_index_build_parameters = DiskIndexBuildParameters::new( + MemoryBudget::try_from_gb(self.params.index_build_ram_gb)?, + self.params.build_quantization_type, + NumPQChunks::new_with(self.params.num_pq_chunks, self.params.full_dim)?, + ); + if let Some(chunk_vector_count) = self.params.data_compression_chunk_vector_count { + disk_index_build_parameters = disk_index_build_parameters + .with_data_compression_chunk_vector_count(chunk_vector_count); + } + + let config = config::Builder::new_with( + self.params.max_degree.into_usize(), + config::MaxDegree::default_slack(), + self.params.l_build.into_usize(), + self.params.metric.into(), + |b| { + b.saturate_after_prune(true); + }, + ) + .build()?; + + let metadata = + load_metadata_from_file(self.storage_provider.as_ref(), &self.params.data_path) + .unwrap(); + + assert_eq!( + self.params.dim, + metadata.ndims(), + "Parameters dimension {} and data dimension {} are not equal", + self.params.dim, + metadata.ndims(), + ); + + let config = IndexConfiguration::new( + self.params.metric, + self.params.dim, + metadata.npoints(), + ONE, + self.params.num_threads, + config, + ) + .with_pseudo_rng_from_seed(100); + + let disk_index_writer = DiskIndexWriter::new( + self.params.data_path.clone(), + self.params.index_path_prefix.clone(), + self.params.associated_data_path.clone(), + self.params.block_size, + )?; + + let mut disk_index = DiskIndexBuilder::::new( + self.storage_provider.as_ref(), + disk_index_build_parameters, + config, + disk_index_writer, + )?; + + let timer = Instant::now(); + disk_index.build()?; + println!("Indexing time: {} seconds", timer.elapsed().as_secs_f64()); + + Ok(()) + } + + pub fn compare_pq_compressed_files(&self) { + self.compare_files( + &self.params.pq_compressed_path(), + &self.params.truth_pq_compressed_path(), + ); + } + + pub fn assert_index_max_degree(&self) -> ANNResult<()> { + let index_file_path = get_disk_index_file(&self.params.index_path_prefix); + let file_data = load_file_to_vec(self.storage_provider.as_ref(), &index_file_path); + let graph_header = GraphHeader::try_from(&file_data[8..])?; + let max_degree = graph_header.max_degree::()?; + assert_eq!( + max_degree, self.params.max_degree as usize, + "Max degree mismatch: expected {}, got {}", + self.params.max_degree, max_degree + ); + + Ok(()) + } + + fn compare_disk_index_with_associated_data( + &self, + pivot_file_prefix_test: &str, + pivot_file_prefix_expected: &str, + index_file_suffix: &str, + ) { + let pq_pivot_path = pivot_file_prefix_test.to_string() + index_file_suffix; + let pq_pivot_path_truth = pivot_file_prefix_expected.to_string() + index_file_suffix; + let file1 = load_file_to_vec(self.storage_provider.as_ref(), &pq_pivot_path); + let file2 = load_file_to_vec(self.storage_provider.as_ref(), &pq_pivot_path_truth); + compare_disk_index_graphs(&file1, &file2) + } + + pub fn compare_files(&self, file_path1: &str, file_path2: &str) { + let file1 = load_file_to_vec(self.storage_provider.as_ref(), file_path1); + let file2 = load_file_to_vec(self.storage_provider.as_ref(), file_path2); + + assert_eq!(file1.len(), file2.len()); + assert_eq!(file1, file2) + } + } + + /// Common helper function for one-shot async index build tests + fn run_one_shot_test(index_path_prefix: String, params_customizer: F) + where + F: FnOnce(TestParams) -> TestParams, + { + let l_build = 64; + let max_degree = 16; + let top_k = 10; + let search_l = 32; + + let base_params = TestParams { + l_build, + max_degree, + index_path_prefix, + ..TestParams::default() + }; + + let params = params_customizer(base_params); + + let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); + fixture.build::().unwrap(); + + // Validate search recall against ground truth for async tests + verify_search_result_with_ground_truth::( + &fixture.params, + top_k, + search_l, + &fixture.storage_provider, + ) + .unwrap(); + + fixture + .assert_index_max_degree::() + .unwrap(); + + // Assert that all data was kept in memory and no files were written to the disk. + let mem_index_file_path = format!("{}_mem.index.data", fixture.params.index_path_prefix); + assert!(!fixture.storage_provider.exists(&mem_index_file_path)); + } + + #[rstest] + fn test_build_from_iter_one_shot_with_metric( + #[values(Metric::L2, Metric::InnerProduct, Metric::Cosine)] metric: Metric, + ) { + let index_path_prefix = format!("{}_metric_{:?}", INDEX_PATH_PREFIX, metric); + + run_one_shot_test(index_path_prefix, |params| TestParams { metric, ..params }); + } + + /// Forces the multi-sector-per-node layout: with a 512-byte block, a 128-d f32 vector + /// (512 B) plus its neighbor list exceeds one sector, so each node spans multiple + /// sectors (`node_len > block_size`). The default 4096-byte sector packs these small + /// nodes many-per-sector, leaving the crossing path otherwise untested; 512 is one of + /// the two block sizes the disk format supports. + #[test] + fn test_build_multi_sector_per_node() { + let params = TestParams { + block_size: 512, + max_degree: 16, + l_build: 64, + index_path_prefix: format!("{}_multi_sector", INDEX_PATH_PREFIX), + ..TestParams::default() + }; + let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); + fixture.build::().unwrap(); + verify_search_result_with_ground_truth::( + &fixture.params, + 10, + 32, + &fixture.storage_provider, + ) + .unwrap(); + } + + #[test] + fn test_build_from_iter_one_shot_with_associated_data() { + // Set up test data + let params = TestParams { + associated_data_path: Some( + "/sift/siftsmall_learn_256pts_u32_associated_data.fbin".to_string(), + ), + ..TestParams::default() + }; + + // Create fixture with virtual storage provider + let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); + + // Build the index with the associated data + fixture.build::().unwrap(); + + // Assert that all data was kept in memory and no files were written to the disk. + let mem_index_file_path = format!("{}_mem.index.data", fixture.params.index_path_prefix); + let mem_index_associated_data_path = format!( + "{}_mem.index.associated_data", + fixture.params.index_path_prefix + ); + assert!(!fixture.storage_provider.exists(&mem_index_file_path)); + assert!(!fixture + .storage_provider + .exists(&mem_index_associated_data_path)); + + // assert index files are expected. + fixture.compare_disk_index_with_associated_data( + &fixture.params.index_path_prefix, + fixture.params.truth_index_path_prefix(), + "_disk.index", + ); + } + + #[test] + fn test_build_from_iter_merged_index() { + // Use the same parameters from [test_sift_build_and_search] in diskann_index + let l_build = 64; + let max_degree = 16; + let top_k = 10; + let search_l = 32; + + let index_path_prefix = + "/disk_index_build/disk_index_sift_learn_test_disk_index_build_merged".to_string(); + let params = TestParams { + l_build, + max_degree, + index_path_prefix, + index_build_ram_gb: 0.0001, // small enough to trigger merged index build + ..TestParams::default() + }; + + let fixture = IndexBuildFixture::new(new_vfs(), params).unwrap(); + + fixture.build::().unwrap(); + + verify_search_result_with_ground_truth::( + &fixture.params, + top_k, + search_l, + &fixture.storage_provider, + ) + .unwrap(); + + fixture + .assert_index_max_degree::() + .unwrap(); + } + + #[rstest] + #[case(QuantizationType::SQ { nbits: 2, standard_deviation: None }, "SQ quantization is only supported for 1 bit")] + fn test_build_quantization_type_failure_cases( + #[case] build_quantization_type: QuantizationType, + #[case] error_message: &str, + ) { + let storage_provider = VirtualStorageProvider::new_overlay(test_data_root()); + let disk_index_builder = create_disk_index_builder( + 1000, // num_points + 128, // dim + 128, // num_pq_chunks + &storage_provider, + build_quantization_type, + ); + + let err = disk_index_builder.err().unwrap(); + assert!(err.to_string().contains(error_message)); + } + + fn load_file_to_vec( + storage_provider: &StorageType, + file_path: &str, + ) -> Vec { + let mut file = storage_provider.open_reader(file_path).unwrap(); + let mut buffer = vec![]; + file.read_to_end(&mut buffer).unwrap(); + buffer + } + + /// Verifies that search results exactly match the ground truth of nearest neighbors + /// + /// This function performs validation of search results by: + /// 1. Running searches on the index using actual data points from the dataset as queries + /// 2. Computing the exact ground truth results using direct distance calculations + /// 3. Verifying that the search engine returns precisely the same results as the ground truth + pub(crate) fn verify_search_result_with_ground_truth< + G: GraphDataType, + >( + params: &TestParams, + top_k: usize, + search_l: u32, + storage_provider: &Arc>, + ) -> ANNResult<()> { + let pq_pivot_path = get_pq_pivot_file(¶ms.index_path_prefix); + let pq_compressed_path = get_compressed_pq_file(¶ms.index_path_prefix); + let index_file_path = get_disk_index_file(¶ms.index_path_prefix); + + let index_reader = + DiskIndexReader::new(pq_pivot_path, pq_compressed_path, storage_provider.as_ref())?; + + let vertex_provider_factory = DiskVertexProviderFactory::new( + VirtualAlignedReaderFactory::new(index_file_path, Arc::clone(storage_provider)), + CachingStrategy::None, + )?; + + let search_engine = DiskIndexSearcher::>::new( + 1, + u32::MAX as usize, + &index_reader, + vertex_provider_factory, + params.metric, + None, + )?; + + let data = + read_bin::(&mut storage_provider.open_reader(¶ms.data_path)?)?; + let dim = data.ncols(); + let distance = ::distance(params.metric, Some(dim)); + + // Here, we use elements of the dataset to search the dataset itself. + // + // We do this for each query, computing the expected ground truth and verifying + // that our simple graph search matches. + // + // Because this dataset is small, we can expect exact equality. + for (q, query_data) in data.row_iter().enumerate() { + let gt = + diskann_providers::test_utils::groundtruth(data.as_view(), query_data, |a, b| { + distance.evaluate_similarity(a, b) + }); + + let mut query_stats = QueryStatistics::default(); + + let mut indices = vec![0u32; top_k]; + let mut distances = vec![0f32; top_k]; + let mut associated_data = vec![(); top_k]; + + _ = search_engine.search_internal( + query_data, + top_k, + search_l, + None, // beam_width + &mut query_stats, + &mut indices, + &mut distances, + &mut associated_data, + &crate::search::search_mode::SearchMode::graph(), + ); + + diskann_providers::test_utils::assert_top_k_exactly_match( + q, >, &indices, &distances, top_k, + ); + } + + Ok(()) + } + + // Compare that the index built in test is the same as the truth index. The truth index doesn't have associated data, we are only comparing the vector and neighbor data. + pub fn compare_disk_index_graphs(graph_data: &[u8], truth_graph_data: &[u8]) { + let graph_header = GraphHeader::try_from(&graph_data[8..]).unwrap(); + let truth_graph_header = GraphHeader::try_from(&truth_graph_data[8..]).unwrap(); + + let test_node_per_block = graph_header.metadata().num_nodes_per_block; + let test_max_node_length = graph_header.metadata().node_len; + + let truth_node_per_block = truth_graph_header.metadata().num_nodes_per_block; + let truth_max_node_length = truth_graph_header.metadata().node_len; + + assert_eq!( + graph_header.metadata().num_pts, + truth_graph_header.metadata().num_pts + ); + + assert_eq!( + graph_header.metadata().dims, + truth_graph_header.metadata().dims + ); + + let num_pts = graph_header.metadata().num_pts as usize; + let dim = graph_header.metadata().dims; + + for idx in 0..num_pts { + let test_node_id_offset = node_data_offset( + idx, + test_max_node_length as usize, + test_node_per_block as usize, + DEFAULT_DISK_SECTOR_LEN, + ); + + let truth_node_id_offset = node_data_offset( + idx, + truth_max_node_length as usize, + truth_node_per_block as usize, + DEFAULT_DISK_SECTOR_LEN, + ); + + // Assert that the vector data is the same between the test and truth graphs for this node. + assert_eq!( + &graph_data + [test_node_id_offset..test_node_id_offset + dim * std::mem::size_of::()], + &truth_graph_data + [truth_node_id_offset..truth_node_id_offset + dim * std::mem::size_of::()] + ); + + // Assert that the neighbor count is the same between the test and truth graphs for this node. + let test_nbr_cnt_offset = test_node_id_offset + dim * std::mem::size_of::(); + let truth_nbr_cnt_offset = truth_node_id_offset + dim * std::mem::size_of::(); + + let test_nbr_count = u32::from_le_bytes([ + graph_data[test_nbr_cnt_offset], + graph_data[test_nbr_cnt_offset + 1], + graph_data[test_nbr_cnt_offset + 2], + graph_data[test_nbr_cnt_offset + 3], + ]); + + let truth_nbr_count = u32::from_le_bytes([ + truth_graph_data[truth_nbr_cnt_offset], + truth_graph_data[truth_nbr_cnt_offset + 1], + truth_graph_data[truth_nbr_cnt_offset + 2], + truth_graph_data[truth_nbr_cnt_offset + 3], + ]); + + assert_eq!(test_nbr_count, truth_nbr_count); + + // Assert the neighbors (u32) are the same between the test and truth graphs for this node. + let test_nbr_offset = test_nbr_cnt_offset + 4; + let truth_nbr_offset = truth_nbr_cnt_offset + 4; + assert_eq!( + graph_data[test_nbr_offset..test_nbr_offset + test_nbr_count as usize * 4], + truth_graph_data[truth_nbr_offset..truth_nbr_offset + truth_nbr_count as usize * 4] + ); + } + } + + pub fn node_data_offset( + node_id: usize, + node_length: usize, + nodes_per_block: usize, + block_size: usize, + ) -> usize { + let block_id = node_id / nodes_per_block; + let node_id_in_block = node_id % nodes_per_block; + let offset = block_id * block_size + node_id_in_block * node_length; + offset + block_size + } + + fn create_disk_index_builder( + num_points: usize, + dim: usize, + num_of_pq_chunks: usize, + storage_provider: &VirtualStorageProvider, + build_quantization_type: QuantizationType, + ) -> ANNResult< + DiskIndexBuilder<'_, GraphDataF32VectorUnitData, VirtualStorageProvider>, + > { + let memory_budget = MemoryBudget::try_from_gb(1.0)?; + let num_pq_chunks = NumPQChunks::new_with(num_of_pq_chunks, dim)?; + + let build_parameters = + DiskIndexBuildParameters::new(memory_budget, build_quantization_type, num_pq_chunks); + + let index_configuration = IndexConfiguration::new( + L2, + dim, + num_points, + ONE, + 1, + config::Builder::new_with(4, config::MaxDegree::default_slack(), 50, L2.into(), |b| { + b.saturate_after_prune(true); + }) + .build()?, + ); + + let disk_index_writer = DiskIndexWriter::new( + "data_path".to_string(), + "index_path_prefix".to_string(), + None, + DEFAULT_DISK_SECTOR_LEN, + )?; + + DiskIndexBuilder::>::new( + storage_provider, + build_parameters, + index_configuration, + disk_index_writer, + ) + } #[derive(PartialEq)] enum BuildType { diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index ffc78238f..3e191db4e 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -1249,7 +1249,7 @@ mod disk_provider_tests { use super::*; use crate::{ - build::builder::merged_index::disk_index_builder_tests::{IndexBuildFixture, TestParams}, + build::builder::tests::disk_index_builder_tests::{IndexBuildFixture, TestParams}, search::provider::aligned_file_reader::VirtualAlignedReaderFactory, utils::QueryStatistics, }; From 6c76dbeb727fd2f644466679f2f4565b9ec60b3b Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Wed, 22 Jul 2026 15:36:55 +0800 Subject: [PATCH 4/7] Make builder test imports explicit Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ecf6179f-7467-44be-8a77-bf7a96e319af --- diskann-disk/src/build/builder/tests.rs | 26 ++++++++++--------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/diskann-disk/src/build/builder/tests.rs b/diskann-disk/src/build/builder/tests.rs index f01957caf..49fac2f76 100644 --- a/diskann-disk/src/build/builder/tests.rs +++ b/diskann-disk/src/build/builder/tests.rs @@ -4,14 +4,6 @@ */ //! Disk index builder tests. -use crate::{data_model::GraphDataType, storage::DiskIndexWriter}; -use diskann_providers::{ - model::IndexConfiguration, - storage::{StorageReadProvider, StorageWriteProvider}, - utils::load_metadata_from_file, -}; -use diskann_utils::io::read_bin; - #[cfg(test)] pub(crate) mod disk_index_builder_tests { use std::{io::Read, sync::Arc, time::Instant}; @@ -24,12 +16,15 @@ pub(crate) mod disk_index_builder_tests { utils::{IntoUsize, VectorRepr, ONE}, ANNResult, }; - use diskann_providers::storage::VirtualStorageProvider; - use diskann_providers::storage::{ - get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, + use diskann_providers::{ + model::IndexConfiguration, + storage::{ + get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, StorageReadProvider, + StorageWriteProvider, VirtualStorageProvider, + }, + utils::load_metadata_from_file, }; - - use diskann_utils::test_data_root; + use diskann_utils::{io::read_bin, test_data_root}; use diskann_vector::{ distance::Metric::{self, L2}, DistanceFunction, @@ -37,16 +32,15 @@ pub(crate) mod disk_index_builder_tests { use rstest::rstest; use vfs::OverlayFS; - use super::*; use crate::{ build::builder::build::DiskIndexBuilder, - data_model::{CachingStrategy, GraphHeader}, + data_model::{CachingStrategy, GraphDataType, GraphHeader}, disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}, search::provider::{ aligned_file_reader::VirtualAlignedReaderFactory, disk_provider::DiskIndexSearcher, disk_vertex_provider_factory::DiskVertexProviderFactory, }, - storage::disk_index_reader::DiskIndexReader, + storage::{disk_index_reader::DiskIndexReader, DiskIndexWriter}, utils::QueryStatistics, QuantizationType, }; From 37a839fdb055724948499c824841deb35f162acc Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Wed, 22 Jul 2026 15:46:09 +0800 Subject: [PATCH 5/7] Restrict disk builder module visibility Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ecf6179f-7467-44be-8a77-bf7a96e319af --- diskann-disk/src/build/builder/mod.rs | 11 +++++++---- diskann-disk/src/search/provider/disk_provider.rs | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/diskann-disk/src/build/builder/mod.rs b/diskann-disk/src/build/builder/mod.rs index 126e262d3..38dd4b796 100644 --- a/diskann-disk/src/build/builder/mod.rs +++ b/diskann-disk/src/build/builder/mod.rs @@ -5,11 +5,14 @@ //! Disk index builders and related functionality. pub mod build; + +mod inmem_builder; mod merged_index; -pub mod quantizer; +mod quantizer; +mod tokio; -pub mod inmem_builder; -pub mod tokio; +#[cfg(test)] +mod tests; #[cfg(test)] -pub(crate) mod tests; +pub(crate) use tests::disk_index_builder_tests::{IndexBuildFixture, TestParams}; diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 3e191db4e..6f19df31f 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -1249,7 +1249,7 @@ mod disk_provider_tests { use super::*; use crate::{ - build::builder::tests::disk_index_builder_tests::{IndexBuildFixture, TestParams}, + build::builder::{IndexBuildFixture, TestParams}, search::provider::aligned_file_reader::VirtualAlignedReaderFactory, utils::QueryStatistics, }; From 31d94fc94889237484105755714636f74bf5f794 Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Wed, 22 Jul 2026 17:35:25 +0800 Subject: [PATCH 6/7] Address disk build review comments Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ecf6179f-7467-44be-8a77-bf7a96e319af --- diskann-disk/src/build/builder/quantizer.rs | 4 ++-- diskann-disk/src/build/builder/tests.rs | 6 ++++-- diskann-disk/src/build/builder/tokio.rs | 2 +- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/diskann-disk/src/build/builder/quantizer.rs b/diskann-disk/src/build/builder/quantizer.rs index d98fa48e7..b1bf86830 100644 --- a/diskann-disk/src/build/builder/quantizer.rs +++ b/diskann-disk/src/build/builder/quantizer.rs @@ -23,7 +23,7 @@ use crate::QuantizationType; /// Quantizer types used for disk index building. #[derive(Clone)] -pub enum BuildQuantizer { +pub(super) enum BuildQuantizer { NoQuant(NoStore), Scalar1Bit(WithBits<1>), PQ(FixedChunkPQTable), @@ -31,7 +31,7 @@ pub enum BuildQuantizer { impl BuildQuantizer { /// Train a new quantizer from scratch. - pub fn train( + pub(super) fn train( build_quantization_type: &QuantizationType, index_path_prefix: &str, index_configuration: &IndexConfiguration, diff --git a/diskann-disk/src/build/builder/tests.rs b/diskann-disk/src/build/builder/tests.rs index 49fac2f76..ddea80d5c 100644 --- a/diskann-disk/src/build/builder/tests.rs +++ b/diskann-disk/src/build/builder/tests.rs @@ -505,9 +505,11 @@ pub(crate) mod disk_index_builder_tests { let test_node_per_block = graph_header.metadata().num_nodes_per_block; let test_max_node_length = graph_header.metadata().node_len; + let test_block_size = graph_header.block_size() as usize; let truth_node_per_block = truth_graph_header.metadata().num_nodes_per_block; let truth_max_node_length = truth_graph_header.metadata().node_len; + let truth_block_size = truth_graph_header.block_size() as usize; assert_eq!( graph_header.metadata().num_pts, @@ -527,14 +529,14 @@ pub(crate) mod disk_index_builder_tests { idx, test_max_node_length as usize, test_node_per_block as usize, - DEFAULT_DISK_SECTOR_LEN, + test_block_size, ); let truth_node_id_offset = node_data_offset( idx, truth_max_node_length as usize, truth_node_per_block as usize, - DEFAULT_DISK_SECTOR_LEN, + truth_block_size, ); // Assert that the vector data is the same between the test and truth graphs for this node. diff --git a/diskann-disk/src/build/builder/tokio.rs b/diskann-disk/src/build/builder/tokio.rs index 4b66c37ce..7e807e7ba 100644 --- a/diskann-disk/src/build/builder/tokio.rs +++ b/diskann-disk/src/build/builder/tokio.rs @@ -7,7 +7,7 @@ use diskann::{ANNError, ANNResult}; /// Creates a new multi-threaded tokio runtime with the specified number of worker threads. /// If `num_threads` is 0, it defaults to the number of logical CPUs. -pub fn create_runtime(num_threads: usize) -> ANNResult { +pub(super) fn create_runtime(num_threads: usize) -> ANNResult { let mut builder = tokio::runtime::Builder::new_multi_thread(); if num_threads != 0 { From aa357111942555266340912649728e7a11040b70 Mon Sep 17 00:00:00 2001 From: Wei Wu Date: Wed, 22 Jul 2026 19:20:38 +0800 Subject: [PATCH 7/7] Avoid allocations while merging neighbors Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ecf6179f-7467-44be-8a77-bf7a96e319af --- .../src/build/builder/merged_index.rs | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/diskann-disk/src/build/builder/merged_index.rs b/diskann-disk/src/build/builder/merged_index.rs index 0a4cf84d9..cfca29910 100644 --- a/diskann-disk/src/build/builder/merged_index.rs +++ b/diskann-disk/src/build/builder/merged_index.rs @@ -359,17 +359,10 @@ where let nnbrs: u32 = std::cmp::min(final_nbrs.len() as u32, max_degree); merged_vamana_cached_writer.write(&nnbrs.to_le_bytes())?; - let bytes = final_nbrs - .iter() - .take(nnbrs as usize) - .flat_map(|x| x.to_le_bytes()) - .collect::>(); - merged_vamana_cached_writer.write(&bytes)?; + merged_vamana_cached_writer + .write(bytemuck::must_cast_slice(&final_nbrs[..nnbrs as usize]))?; merged_index_size += (size_of::() + nnbrs as usize * size_of::()) as u64; - if cur_id % 499999 == 1 { - print!("."); - } cur_id = node_id; final_nbrs.iter().for_each(|p| nbr_set[*p as usize] = false); @@ -407,12 +400,8 @@ where let nnbrs: u32 = std::cmp::min(final_nbrs.len() as u32, max_degree); merged_vamana_cached_writer.write(&nnbrs.to_le_bytes())?; - let bytes = final_nbrs - .iter() - .take(nnbrs as usize) - .flat_map(|x| x.to_le_bytes()) - .collect::>(); - merged_vamana_cached_writer.write(&bytes)?; + merged_vamana_cached_writer + .write(bytemuck::must_cast_slice(&final_nbrs[..nnbrs as usize]))?; merged_index_size += (size_of::() + nnbrs as usize * size_of::()) as u64;