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
Wei Wu (wuw92) wants to merge 10 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
10 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
bfce216
Remove tautological RAM estimate test
wuw92 Jul 27, 2026
9cb8148
Merge main into disk build refactor
wuw92 Jul 31, 2026
9b1de20
Preserve little-endian neighbor serialization
wuw92 Jul 31, 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
242 changes: 42 additions & 200 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,47 +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},
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,

@partychen juchen-ms (partychen) Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The name build_inmem_index is now used at two different abstraction levels. DiskIndexBuilder::build_inmem_index selects either the one-shot or merged strategy, while this imported helper builds and persists one complete Vamana graph (for either the full dataset or one shard). The outer operation is not necessarily in-memory, and the inner helper is specifically the one-shot leaf operation.

Could we distinguish these responsibilities explicitly - for example, build_vamana_index for the outer phase and build_one_shot_vamana or build_vamana_graph for the leaf helper? That would make the call graph understandable without relying on the historical _mem.index terminology.

merged_index::{estimate_build_index_ram_usage, MergedVamanaIndexBuilder},
quantizer::BuildQuantizer,
tokio::create_runtime,
},
error::{diskann_error, ErrorKind},
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 @@ -212,186 +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(|| {
diskann_error!(
ErrorKind::IndexError,
"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>(
Comment on lines +201 to +206

@partychen juchen-ms (partychen) Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IndexBuildStrategy sounds like it controls the complete disk-index build, but it only selects how the Vamana graph is constructed. Its current placement also leaves the decision split across modules: build.rs owns the strategy, while the RAM estimate it requires remains in merged_index.rs, even though that estimate is not specific to the merged path.

Could the strategy and its estimate be named and owned together with the Vamana builders (for example, VamanaBuildStrategy)? That would leave build.rs focused on the top-level PQ, Vamana, and disk-layout orchestration.

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(|_| diskann_error!(ErrorKind::IndexError, "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(|_| {
diskann_error!(ErrorKind::IndexError, "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(|_| diskann_error!(ErrorKind::IndexError, "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(|_| {
diskann_error!(ErrorKind::IndexError, "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