Skip to content

Refactor disk index build structure #1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wuw92 wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Refactor disk index build structure #1271

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6e7bfdc
Refactor disk index build structure
wuw92 Jul 21, 2026
69ad474
Move disk build strategy into builder
wuw92 Jul 21, 2026
4d4061e
Refactor disk index builder tests and update module paths
wuw92 Jul 22, 2026
6c76dbe
Make builder test imports explicit
wuw92 Jul 22, 2026
37a839f
Restrict disk builder module visibility
wuw92 Jul 22, 2026
31d94fc
Address disk build review comments
wuw92 Jul 22, 2026
aa35711
Avoid allocations while merging neighbors
wuw92 Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
234 changes: 42 additions & 192 deletions diskann-disk/src/build/builder/build.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,46 +3,33 @@
* 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, mem};

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::{estimate_build_index_ram_usage, 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>
Expand Down Expand Up @@ -211,179 +198,42 @@ where
}
}

pub(super) async fn build_inmem_index<T, StorageProvider>(
config: IndexConfiguration,
quantizer: &BuildQuantizer,
data_path: &str,
save_path: &str,
storage_provider: &StorageProvider,
) -> ANNResult<()>
where
T: VectorRepr,
StorageProvider: StorageReadProvider + StorageWriteProvider + 'static,
<StorageProvider as StorageReadProvider>::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::<T>(index_config, provider_parameters, quantizer)?;
let medoid_id =
set_start_point_to_medoid::<T, _>(&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(())
pub(crate) enum IndexBuildStrategy {
OneShot,
Merged,
}

#[cfg(debug_assertions)]
/// Log statistics about the build process
async fn log_build_stats<T: VectorRepr>(index: &Arc<dyn InmemIndexBuilder<T>>) -> ANNResult<()> {
debug!(
"Number of points reachable in the graph: {}",
index.count_reachable_nodes().await?
pub(crate) fn determine_build_strategy<Data: GraphDataType>(
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::<Data::VectorDataType>() as u64,
index_configuration.config.max_degree().get() as u64,
build_quantization_type,
);

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
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
);
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> {
u32::try_from(value)
.map_err(|_| ANNError::log_index_error(format_args!("id {value} exceeds u32::MAX")))
}

fn set_start_point_to_medoid<T, StorageReader>(
index: &Arc<dyn InmemIndexBuilder<T>>,
path: &str,
random_seed: Option<u64>,
reader: &StorageReader,
) -> ANNResult<usize>
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::<T, _>(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<T, I>(
index: &Arc<dyn InmemIndexBuilder<T>>,
iterator: Arc<Mutex<I>>,
num_tasks: NonZeroUsize,
) -> ANNResult<()>
where
T: VectorRepr,
I: Iterator<Item = (usize, (Box<[T]>, ()))> + 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<T: VectorRepr>(
index: &Arc<dyn InmemIndexBuilder<T>>,
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"))??;
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
}

Ok(())
}
Loading
Loading