diff --git a/Cargo.lock b/Cargo.lock index 72af940896..f756d1dc63 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -597,6 +597,7 @@ dependencies = [ "rayon", "rstest", "serde", + "serde_json", "tempfile", "thiserror 2.0.17", "tokio", diff --git a/diskann-disk/Cargo.toml b/diskann-disk/Cargo.toml index 693614381d..29aae9f4c3 100644 --- a/diskann-disk/Cargo.toml +++ b/diskann-disk/Cargo.toml @@ -68,6 +68,7 @@ features = [ rstest.workspace = true tempfile.workspace = true vfs.workspace = true +serde_json.workspace = true diskann-providers = { workspace = true, default-features = false, features = [ "testing", "virtual_storage", @@ -82,6 +83,8 @@ proptest.workspace = true [features] default = [] perf_test = ["dep:opentelemetry"] +pipnn = ["diskann/pipnn"] +virtual_storage = ["diskann-providers/virtual_storage"] experimental_diversity_search = [ "diskann/experimental_diversity_search", "diskann-providers/experimental_diversity_search", diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index c0df5d6b75..7b7e13e551 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -30,6 +30,9 @@ use diskann_providers::{ use tokio::task::JoinSet; use tracing::{debug, info}; +#[cfg(feature = "pipnn")] +mod pipnn; + use crate::{ build::builder::{ core::{determine_build_strategy, IndexBuildStrategy, MergedVamanaIndexBuilder}, @@ -73,6 +76,11 @@ where index_configuration: IndexConfiguration, index_writer: DiskIndexWriter, ) -> ANNResult { + #[cfg(feature = "pipnn")] + if let Some(config) = disk_build_param.pipnn_config() { + config.validate()?; + } + let pq_storage = PQStorage::new( &(index_writer.get_index_path_prefix() + "_pq_pivots.bin"), &(index_writer.get_index_path_prefix() + "_pq_compressed.bin"), @@ -123,7 +131,7 @@ where self.generate_compressed_data(pool.as_ref())?; logger.log_checkpoint(DiskIndexBuildCheckpoint::PqConstruction); - self.build_inmem_index(pool.as_ref()).await?; + self.build_graph(pool.as_ref()).await?; logger.log_checkpoint(DiskIndexBuildCheckpoint::InmemIndexBuild); // Use physical file to pass the memory index to the disk writer @@ -172,7 +180,12 @@ where ) } - async fn build_inmem_index(&mut self, pool: RayonThreadPoolRef<'_>) -> ANNResult<()> { + async fn build_graph(&mut self, pool: RayonThreadPoolRef<'_>) -> ANNResult<()> { + #[cfg(feature = "pipnn")] + if let Some(config) = self.disk_build_param.pipnn_config() { + return pipnn::build_graph(self, pool, config); + } + match determine_build_strategy::( &self.index_configuration, self.disk_build_param.build_memory_limit().in_bytes() as f64, diff --git a/diskann-disk/src/build/builder/build/pipnn.rs b/diskann-disk/src/build/builder/build/pipnn.rs new file mode 100644 index 0000000000..488ea8ae7e --- /dev/null +++ b/diskann-disk/src/build/builder/build/pipnn.rs @@ -0,0 +1,277 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Write a PiPNN graph in the DiskANN disk-index format. +//! +//! This adapter checks dataset metadata and loads one contiguous matrix. It runs +//! PiPNN in the supplied Rayon pool. It computes the start point with the sampled +//! medoid policy. It then writes the common graph header and adjacency layout. +//! +//! PiPNN and Vamana use the same disk graph format. + +use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann::{utils::VectorRepr, ANNError, ANNResult}; +use diskann_providers::{ + storage::{save_adjacency_graph, StorageReadProvider, StorageWriteProvider}, + utils::{find_medoid_with_sampling, RayonThreadPoolRef, MAX_MEDOID_SAMPLE_SIZE}, +}; +use diskann_utils::io::{read_bin, Metadata}; + +use super::{u32_try_from, DiskIndexBuilder}; +use crate::data_model::GraphDataType; + +/// Build PiPNN adjacency and persist it through the canonical disk graph writer. +pub(super) fn build_graph( + builder: &DiskIndexBuilder<'_, Data, StorageProvider>, + pool: RayonThreadPoolRef<'_>, + config: PiPNNConfig, +) -> ANNResult<()> +where + Data: GraphDataType, + Data::VectorDataType: VectorRepr, + StorageProvider: StorageReadProvider + StorageWriteProvider, +{ + let data_path = builder.index_writer.get_dataset_file(); + // Check metadata before the code loads the full matrix. Report a configuration + // error instead of a matrix-shape error. + let (points, dimensions) = + Metadata::read(&mut builder.storage_provider.open_reader(&data_path)?)?.into_dims(); + if dimensions != builder.index_configuration.dim { + return Err(ANNError::message(format!( + "configured dimension {} does not match dataset dimension {dimensions}", + builder.index_configuration.dim + ))); + } + if points != builder.index_configuration.max_points { + return Err(ANNError::message(format!( + "configured point count {} does not match dataset point count {points}", + builder.index_configuration.max_points + ))); + } + + // PiPNN requires one contiguous matrix. Partition and leaf work use the + // supplied Rayon pool. + let data = + read_bin::(&mut builder.storage_provider.open_reader(&data_path)?)?; + let context = PiPNNBuildContext::new( + config, + &builder.index_configuration.config, + builder.index_configuration.dist_metric, + pool.as_rayon(), + )?; + let adjacency = diskann::graph::pipnn::build_graph(data.as_view(), &context)?; + + // The disk header requires a start point. Use the same sampled medoid policy + // as the Vamana disk builder. + let mut rng = diskann_providers::utils::create_rnd_from_optional_seed( + builder.index_configuration.random_seed, + ); + let (_, start_id) = find_medoid_with_sampling::( + &data_path, + builder.storage_provider, + MAX_MEDOID_SAMPLE_SIZE, + &mut rng, + )?; + save_adjacency_graph( + &adjacency, + u32_try_from(builder.index_configuration.config.pruned_degree().get())?, + builder.storage_provider, + u32_try_from(start_id)?, + &builder.index_writer.get_mem_index_file(), + )?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use diskann::{graph::config, utils::ONE}; + use diskann_providers::utils::create_thread_pool; + use diskann_providers::{ + model::IndexConfiguration, + storage::{ + get_disk_index_file, StorageReadProvider, StorageWriteProvider, VirtualStorageProvider, + }, + }; + use diskann_utils::{io::write_bin, views::MatrixView}; + use diskann_vector::distance::Metric; + use vfs::MemoryFS; + + use crate::{ + build::{ + builder::build::DiskIndexBuilder, + configuration::{MemoryBudget, NumPQChunks, PiPNNParameters}, + }, + data_model::AdHoc, + storage::DiskIndexWriter, + DiskIndexBuildParameters, + }; + + fn pipnn() -> PiPNNParameters { + PiPNNParameters { + c_max: 512, + c_min: 64, + p_samp: 0.01, + fanout: vec![10, 3], + k: 2, + replicas: 1, + } + } + + fn write_data(storage: &VirtualStorageProvider, points: usize, dimensions: usize) { + let data: Vec = (0..points * dimensions) + .map(|index| ((index * 17) % 251) as f32) + .collect(); + write_bin( + MatrixView::try_from(data.as_slice(), points, dimensions).unwrap(), + &mut storage.create_for_write("/data.fbin").unwrap(), + ) + .unwrap(); + } + + fn graph_config(degree: usize, alpha: f32) -> diskann::graph::Config { + config::Builder::new_with( + degree, + config::MaxDegree::default_slack(), + 50, + Metric::L2.into(), + |builder| { + builder.alpha(alpha); + }, + ) + .build() + .unwrap() + } + + fn builder<'a>( + storage: &'a VirtualStorageProvider, + points: usize, + dimensions: usize, + budget_gib: f64, + alpha: f32, + parameters: PiPNNParameters, + ) -> DiskIndexBuilder<'a, AdHoc, VirtualStorageProvider> { + let params = DiskIndexBuildParameters::new_pipnn( + MemoryBudget::try_from_gb(budget_gib).unwrap(), + NumPQChunks::new_with(dimensions, dimensions).unwrap(), + parameters, + ); + let config = IndexConfiguration::new( + Metric::L2, + dimensions, + points, + ONE, + 1, + graph_config(32, alpha), + ) + .with_pseudo_rng_from_seed(42); + let writer = + DiskIndexWriter::new("/data.fbin".into(), "/index".into(), None, 4096).unwrap(); + DiskIndexBuilder::new(storage, params, config, writer).unwrap() + } + + #[test] + fn disk_build_rejects_dataset_shape_mismatch() { + let storage = VirtualStorageProvider::new_memory(); + write_data(&storage, 2, 8); + let params = DiskIndexBuildParameters::new_pipnn( + MemoryBudget::try_from_gb(10_000.0).unwrap(), + NumPQChunks::new_with(4, 4).unwrap(), + PiPNNParameters::default(), + ); + let config = IndexConfiguration::new(Metric::L2, 4, 3, ONE, 1, graph_config(4, 1.2)); + let writer = + DiskIndexWriter::new("/data.fbin".into(), "/index".into(), None, 4096).unwrap(); + let mut builder = + DiskIndexBuilder::, _>::new(&storage, params, config, writer).unwrap(); + + let error = builder.build().unwrap_err(); + assert!(format!("{error:?}").contains("configured dimension 4")); + assert!(storage.exists("/index_pq_compressed.bin")); + } + + #[test] + fn graph_adapter_rejects_point_count_mismatch() { + let storage = VirtualStorageProvider::new_memory(); + write_data(&storage, 2, 8); + let parameters = pipnn(); + let builder = builder(&storage, 3, 8, 1.0, 1.2, parameters.clone()); + let pool = create_thread_pool(1).unwrap(); + + let error = super::build_graph(&builder, pool.as_ref(), (¶meters).into()).unwrap_err(); + assert!(format!("{error:?}").contains("configured point count 3")); + assert!(!storage.exists(&builder.index_writer.get_mem_index_file())); + } + + #[test] + fn graph_adapter_writes_degree_medoid_and_frozen_count() { + let storage = VirtualStorageProvider::new_memory(); + let (points, dimensions) = (256, 8); + write_data(&storage, points, dimensions); + let parameters = pipnn(); + let builder = builder(&storage, points, dimensions, 1.0, 1.2, parameters.clone()); + let pool = create_thread_pool(1).unwrap(); + + super::build_graph(&builder, pool.as_ref(), (¶meters).into()).unwrap(); + + let mut header = [0_u8; 24]; + std::io::Read::read_exact( + &mut storage + .open_reader(&builder.index_writer.get_mem_index_file()) + .unwrap(), + &mut header, + ) + .unwrap(); + assert_eq!(u32::from_le_bytes(header[8..12].try_into().unwrap()), 32); + assert!(u32::from_le_bytes(header[12..16].try_into().unwrap()) < points as u32); + assert_eq!(u64::from_le_bytes(header[16..24].try_into().unwrap()), 0); + } + + #[test] + fn explicit_selection_ignores_the_vamana_memory_strategy() { + let storage = VirtualStorageProvider::new_memory(); + let (points, dimensions) = (256, 8); + write_data(&storage, points, dimensions); + let mut builder = builder(&storage, points, dimensions, 0.000001, 1.3, pipnn()); + + assert!(matches!( + builder.disk_build_param.build_algorithm(), + crate::BuildAlgorithm::PiPNN(_) + )); + assert_eq!( + builder.disk_build_param.build_quantization(), + &crate::QuantizationType::FP + ); + assert_eq!(builder.index_configuration.config.pruned_degree().get(), 32); + assert_eq!(builder.index_configuration.config.l_build().get(), 50); + assert_eq!(builder.index_configuration.config.alpha(), 1.3); + builder.build().unwrap(); + assert!(storage.exists(&get_disk_index_file("/index"))); + assert!(storage.exists("/index_pq_compressed.bin")); + } + + #[test] + fn builder_rejects_invalid_pipnn_config() { + let storage = VirtualStorageProvider::new_memory(); + let invalid = PiPNNParameters { + c_max: 0, + ..PiPNNParameters::default() + }; + let params = DiskIndexBuildParameters::new_pipnn( + MemoryBudget::try_from_gb(0.0001).unwrap(), + NumPQChunks::new_with(1, 1).unwrap(), + invalid, + ); + let config = IndexConfiguration::new(Metric::L2, 1, 1, ONE, 1, graph_config(4, 1.2)); + let writer = + DiskIndexWriter::new("/data.fbin".into(), "/index".into(), None, 4096).unwrap(); + + let error = match DiskIndexBuilder::, _>::new(&storage, params, config, writer) { + Ok(_) => panic!("invalid PiPNN config must be rejected"), + Err(error) => error, + }; + + assert!(format!("{error:?}").contains("c_max must be greater than zero")); + } +} diff --git a/diskann-disk/src/build/configuration/build_algorithm.rs b/diskann-disk/src/build/configuration/build_algorithm.rs new file mode 100644 index 0000000000..d33e6b0b89 --- /dev/null +++ b/diskann-disk/src/build/configuration/build_algorithm.rs @@ -0,0 +1,123 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Graph-build algorithm selection and its JSON-facing configuration. + +use std::fmt; + +use serde::{Deserialize, Serialize}; + +/// PiPNN parameters in the JSON build configuration. +/// +/// The common index configuration supplies graph degree, alpha, metric, thread +/// count, and memory limit. +#[cfg(feature = "pipnn")] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct PiPNNParameters { + /// Maximum number of points in a leaf. + pub c_max: usize, + /// Minimum leaf size used by global small-leaf merging. + pub c_min: usize, + /// Fraction of a cluster sampled as leaders. + pub p_samp: f64, + /// Number of nearest leaders retained at each partition level. + pub fanout: Vec, + /// Number of nearest neighbors selected within each leaf. + pub k: usize, + /// Number of independent partition passes. + pub replicas: usize, +} + +#[cfg(feature = "pipnn")] +impl Default for PiPNNParameters { + fn default() -> Self { + Self { + c_max: 256, + c_min: 16, + p_samp: 0.005, + fanout: vec![8, 3], + k: 2, + replicas: 1, + } + } +} + +#[cfg(feature = "pipnn")] +impl From<&PiPNNParameters> for diskann::graph::pipnn::PiPNNConfig { + fn from(config: &PiPNNParameters) -> Self { + Self { + c_max: config.c_max, + c_min: config.c_min, + p_samp: config.p_samp, + fanout: config.fanout.clone(), + leaf_k: config.k, + replicas: config.replicas, + } + } +} + +/// Graph construction algorithm for index building. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(tag = "algorithm")] +#[non_exhaustive] +pub enum BuildAlgorithm { + /// Default Vamana graph construction. + #[default] + Vamana, + + /// PiPNN one-shot partition-based graph construction. + #[cfg(feature = "pipnn")] + PiPNN(PiPNNParameters), +} + +impl fmt::Display for BuildAlgorithm { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Vamana => write!(f, "Vamana"), + #[cfg(feature = "pipnn")] + Self::PiPNN(config) => write!(f, "PiPNN({config:?})"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_is_vamana() { + assert_eq!(BuildAlgorithm::default(), BuildAlgorithm::Vamana); + } + + #[test] + fn vamana_serde_roundtrip() { + let json = serde_json::to_string(&BuildAlgorithm::Vamana).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + BuildAlgorithm::Vamana + ); + } + + #[cfg(feature = "pipnn")] + #[test] + fn pipnn_serde_uses_inline_defaults_and_rejects_unknown_fields() { + let algorithm: BuildAlgorithm = serde_json::from_str( + r#"{"algorithm":"PiPNN","c_max":512,"c_min":64,"fanout":[10,3],"k":3}"#, + ) + .unwrap(); + let BuildAlgorithm::PiPNN(config) = algorithm else { + panic!("expected PiPNN"); + }; + assert_eq!(config.c_max, 512); + assert_eq!(config.c_min, 64); + assert_eq!(config.fanout, [10, 3]); + assert_eq!(config.k, 3); + assert_eq!(config.replicas, 1); + assert!( + serde_json::from_str::(r#"{"algorithm":"PiPNN","l_max":72}"#).is_err() + ); + } +} diff --git a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs index b4ba343b18..a8b4017529 100644 --- a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs +++ b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs @@ -10,7 +10,9 @@ use std::num::NonZeroUsize; use diskann::ANNError; use thiserror::Error; -use super::QuantizationType; +#[cfg(feature = "pipnn")] +use super::PiPNNParameters; +use super::{BuildAlgorithm, QuantizationType}; use crate::error::{diskann_error, ErrorKind}; @@ -107,9 +109,10 @@ impl NumPQChunks { } /// Parameters specific for disk index construction. -#[derive(Clone, Copy, PartialEq, Debug)] +#[derive(Clone, PartialEq, Debug)] pub struct DiskIndexBuildParameters { - /// Limit on the memory allowed for building the index. + /// Memory budget for disk-index stages that support bounded work. + /// PiPNN always uses its one-shot graph build. build_memory_limit: MemoryBudget, /// Number of PQ chunks stored in-memory for search and to be generated during build. @@ -120,6 +123,9 @@ pub struct DiskIndexBuildParameters { /// Number of vectors processed per data-compression chunk. data_compression_chunk_vector_count: usize, + + /// Which graph construction algorithm to use. + build_algorithm: BuildAlgorithm, } impl DiskIndexBuildParameters { @@ -134,6 +140,26 @@ impl DiskIndexBuildParameters { search_pq_chunks, build_quantization, data_compression_chunk_vector_count: DEFAULT_DATA_COMPRESSION_CHUNK_VECTOR_COUNT, + build_algorithm: BuildAlgorithm::default(), + } + } + + /// Create parameters for one-shot PiPNN graph construction. + /// + /// PiPNN uses the common search-PQ and disk layout. The memory budget does + /// not limit its one-shot graph build. + #[cfg(feature = "pipnn")] + pub fn new_pipnn( + build_memory_limit: MemoryBudget, + search_pq_chunks: NumPQChunks, + config: PiPNNParameters, + ) -> Self { + Self { + build_memory_limit, + search_pq_chunks, + build_quantization: QuantizationType::FP, + data_compression_chunk_vector_count: DEFAULT_DATA_COMPRESSION_CHUNK_VECTOR_COUNT, + build_algorithm: BuildAlgorithm::PiPNN(config), } } @@ -165,6 +191,19 @@ impl DiskIndexBuildParameters { pub fn data_compression_chunk_vector_count(&self) -> usize { self.data_compression_chunk_vector_count } + + /// Get the graph-construction algorithm. + pub fn build_algorithm(&self) -> &BuildAlgorithm { + &self.build_algorithm + } + + #[cfg(feature = "pipnn")] + pub(crate) fn pipnn_config(&self) -> Option { + match &self.build_algorithm { + BuildAlgorithm::PiPNN(config) => Some(config.into()), + BuildAlgorithm::Vamana => None, + } + } } #[cfg(test)] diff --git a/diskann-disk/src/build/configuration/mod.rs b/diskann-disk/src/build/configuration/mod.rs index 25453abd09..a7e343fb57 100644 --- a/diskann-disk/src/build/configuration/mod.rs +++ b/diskann-disk/src/build/configuration/mod.rs @@ -2,6 +2,11 @@ * Copyright (c) Microsoft Corporation. * Licensed under the MIT license. */ +pub mod build_algorithm; +pub use build_algorithm::BuildAlgorithm; +#[cfg(feature = "pipnn")] +pub use build_algorithm::PiPNNParameters; + pub mod disk_index_build_parameter; pub use disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}; diff --git a/diskann-disk/src/build/mod.rs b/diskann-disk/src/build/mod.rs index 20e6e4b389..27f4c124aa 100644 --- a/diskann-disk/src/build/mod.rs +++ b/diskann-disk/src/build/mod.rs @@ -12,6 +12,9 @@ pub mod builder; pub mod configuration; // Re-export key types for convenience +#[cfg(feature = "pipnn")] +pub use configuration::PiPNNParameters; pub use configuration::{ - disk_index_build_parameter, filter_parameter, DiskIndexBuildParameters, QuantizationType, + disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters, + QuantizationType, }; diff --git a/diskann-disk/src/lib.rs b/diskann-disk/src/lib.rs index facaebf094..5d9e6c368d 100644 --- a/diskann-disk/src/lib.rs +++ b/diskann-disk/src/lib.rs @@ -14,8 +14,11 @@ pub(crate) mod test_utils; pub mod error; pub mod build; +#[cfg(feature = "pipnn")] +pub use build::PiPNNParameters; pub use build::{ - disk_index_build_parameter, filter_parameter, DiskIndexBuildParameters, QuantizationType, + disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters, + QuantizationType, }; pub mod data_model; diff --git a/diskann-providers/src/storage/bin.rs b/diskann-providers/src/storage/bin.rs index 358356265d..c4f24dc7c5 100644 --- a/diskann-providers/src/storage/bin.rs +++ b/diskann-providers/src/storage/bin.rs @@ -9,6 +9,7 @@ use super::{StorageReadProvider, StorageWriteProvider}; use byteorder::{LittleEndian, ReadBytesExt}; use diskann::{ ANNError, ANNResult, + graph::AdjacencyList, utils::{IntoUsize, VectorRepr}, }; use diskann_utils::io::Metadata; @@ -378,3 +379,152 @@ where out.flush()?; Ok(index_size.into_usize()) } + +/// Save real-point adjacency lists in the canonical graph layout. +/// +/// # Errors +/// +/// Returns an error before file creation if an ID is outside `adjacency`. It also +/// returns an error if a row exceeds `max_degree`. The function returns storage +/// creation and write errors from `provider`. +pub fn save_adjacency_graph

( + adjacency: &[AdjacencyList], + max_degree: u32, + provider: &P, + start_point: u32, + path: &str, +) -> ANNResult +where + P: StorageWriteProvider, +{ + let points = adjacency.len(); + if start_point.into_usize() >= points { + return Err(ANNError::message(format!( + "graph start point {start_point} is outside {points} rows" + ))); + } + let max_degree = max_degree.into_usize(); + for (source, neighbors) in adjacency.iter().enumerate() { + if neighbors.len() > max_degree { + return Err(ANNError::message(format!( + "graph row {source} has degree {}, exceeding configured max degree {max_degree}", + neighbors.len() + ))); + } + if let Some(&neighbor) = neighbors.iter().find(|&&id| id.into_usize() >= points) { + return Err(ANNError::message(format!( + "graph row {source} has neighbor {neighbor} outside {points} rows" + ))); + } + } + + save_graph( + &AdjacencyGraph { + adjacency, + max_degree: max_degree as u32, + }, + provider, + start_point, + path, + ) +} + +struct AdjacencyGraph<'a> { + adjacency: &'a [AdjacencyList], + max_degree: u32, +} + +impl GetAdjacencyList for AdjacencyGraph<'_> { + type Element = u32; + type Item<'a> + = &'a [u32] + where + Self: 'a; + + fn get_adjacency_list(&self, index: usize) -> ANNResult> { + self.adjacency + .get(index) + .map(|row| &**row) + .ok_or_else(|| ANNError::message(format!("missing graph row {index}"))) + } + + fn total(&self) -> usize { + self.adjacency.len() + } + + fn additional_points(&self) -> u64 { + 0 + } + + fn max_degree(&self) -> Option { + Some(self.max_degree) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::VirtualStorageProvider; + use vfs::MemoryFS; + + #[derive(Debug)] + struct LoadedGraph { + rows: Vec>, + max_degree: usize, + start_points: usize, + } + + impl SetAdjacencyList for LoadedGraph { + type Item = u32; + + fn set_adjacency_list(&mut self, index: usize, neighbors: &[u32]) -> ANNResult<()> { + self.rows[index].extend_from_slice(neighbors); + Ok(()) + } + } + + fn adjacency(rows: &[&[u32]]) -> Vec> { + rows.iter() + .map(|row| AdjacencyList::from_iter_untrusted(row.iter().copied())) + .collect() + } + + #[test] + fn adjacency_graph_roundtrips_through_the_canonical_loader() { + let storage = VirtualStorageProvider::::new_memory(); + let expected = adjacency(&[&[1, 2], &[0], &[]]); + + save_adjacency_graph(&expected, 2, &storage, 1, "/graph").unwrap(); + let actual = load_graph(&storage, "/graph", |points, max_degree, start_points| { + Ok(LoadedGraph { + rows: vec![Vec::new(); points], + max_degree, + start_points, + }) + }) + .unwrap(); + + assert_eq!(actual.rows, [vec![1, 2], vec![0], vec![]]); + assert_eq!(actual.max_degree, 2); + assert_eq!(actual.start_points, 0); + } + + #[test] + fn adjacency_graph_rejects_inconsistent_header_and_ids_before_writing() { + let cases = [ + (adjacency(&[&[1, 2], &[], &[]]), 1, 0, "degree"), + (adjacency(&[&[3], &[], &[]]), 1, 0, "neighbor"), + (adjacency(&[&[], &[]]), 1, 2, "start point"), + ]; + + for (index, (graph, max_degree, start_point, message)) in cases.into_iter().enumerate() { + let storage = VirtualStorageProvider::::new_memory(); + let path = format!("/graph-{index}"); + let error = + save_adjacency_graph(&graph, max_degree, &storage, start_point, &path).unwrap_err(); + + assert!(error.to_string().contains(message), "{error}"); + assert!(!storage.exists(&path)); + } + } +} diff --git a/diskann-providers/src/storage/mod.rs b/diskann-providers/src/storage/mod.rs index 1233b11f61..f2add5ff25 100644 --- a/diskann-providers/src/storage/mod.rs +++ b/diskann-providers/src/storage/mod.rs @@ -17,6 +17,7 @@ mod api; pub use api::{AsyncIndexMetadata, AsyncQuantLoadContext, DiskGraphOnly, LoadWith, SaveWith}; pub(crate) mod bin; +pub use bin::save_adjacency_graph; pub(crate) mod file_storage_provider; // Use VirtualStorageProvider in tests to avoid filesystem side-effects diff --git a/diskann-providers/src/utils/rayon_util.rs b/diskann-providers/src/utils/rayon_util.rs index 744cbd6ccf..d60fde72fb 100644 --- a/diskann-providers/src/utils/rayon_util.rs +++ b/diskann-providers/src/utils/rayon_util.rs @@ -74,6 +74,11 @@ impl<'a> RayonThreadPoolRef<'a> { { self.0.install(op) } + + /// Return a reference to the underlying Rayon pool. + pub fn as_rayon(self) -> &'a rayon::ThreadPool { + self.0 + } } // Allow use of disallowed methods within this trait to provide custom