diff --git a/Cargo.lock b/Cargo.lock index 6a88b5550..c02628a4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -596,6 +596,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 693614381..29aae9f4c 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 432c577fc..20a01949d 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}, @@ -72,6 +75,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"), @@ -122,7 +130,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 @@ -171,7 +179,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 000000000..069bb19fc --- /dev/null +++ b/diskann-disk/src/build/builder/build/pipnn.rs @@ -0,0 +1,104 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Adapter from provider-independent PiPNN adjacency to the common disk index format. +//! +//! The core crate deliberately knows nothing about dataset files, medoids, graph +//! headers, or serialization. This adapter owns that boundary: +//! +//! 1. verify on-disk dataset metadata against the requested index configuration; +//! 2. load the contiguous matrix required by batch construction; +//! 3. run PiPNN in the caller-provided Rayon pool; +//! 4. compute the production start node with the existing medoid policy; and +//! 5. serialize adjacency with the same header/layout used by Vamana. +//! +//! ```text +//! dataset file ──> metadata check ──> MatrixView ──> diskann::graph::pipnn ──> adjacency +//! │ │ +//! └──────────────────> sampled medoid ────────────────────────────┤ +//! v +//! canonical graph writer +//! ``` +//! +//! There is no PiPNN-specific disk graph format. Keeping serialization here means +//! search and loading cannot distinguish which builder produced the graph. + +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(); + // Validate metadata before allocating/loading the full matrix. A mismatch + // here otherwise turns a configuration error into a later shape failure. + let (points, dimensions) = + Metadata::read(&mut builder.storage_provider.open_reader(&data_path)?)?.into_dims(); + if dimensions != builder.index_configuration.dim { + return Err(ANNError::log_dimension_mismatch_error(format!( + "configured dimension {} does not match dataset dimension {dimensions}", + builder.index_configuration.dim + ))); + } + if points != builder.index_configuration.max_points { + return Err(ANNError::log_index_error(format!( + "configured point count {} does not match dataset point count {points}", + builder.index_configuration.max_points + ))); + } + + // PiPNN is a batch algorithm: materialize the matrix once, while all + // partition and leaf scratch stays inside the supplied pool and is released + // before the outer disk pipeline continues. + 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)?; + + // Start-node policy belongs to the persisted index, not the core graph + // constructor. Reuse the production sampled medoid implementation so the + // serialized header has the same semantics as a Vamana-built index. + 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; diff --git a/diskann-disk/src/build/builder/build/pipnn/tests.rs b/diskann-disk/src/build/builder/build/pipnn/tests.rs new file mode 100644 index 000000000..28881f706 --- /dev/null +++ b/diskann-disk/src/build/builder/build/pipnn/tests.rs @@ -0,0 +1,190 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +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 000000000..109bf34c0 --- /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}; + +/// JSON-facing PiPNN parameters. +/// +/// Graph degree, build-L, alpha, metric, threads, and memory limits remain in +/// the outer index configuration shared with Vamana. +#[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(), + k: config.k, + replicas: config.replicas, + } + } +} + +/// Selects the 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 077c5090d..19fef368d 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}; /// GB to bytes ratio. pub const BYTES_IN_GB: f64 = 1024_f64 * 1024_f64 * 1024_f64; @@ -105,9 +107,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 pipeline stages that support bounded work. + /// Explicit one-shot PiPNN selection is never silently replaced. build_memory_limit: MemoryBudget, /// Number of PQ chunks stored in-memory for search and to be generated during build. @@ -118,6 +121,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 { @@ -132,6 +138,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 pipeline. Its one-shot + /// graph build is not governed by the pipeline memory budget. + #[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), } } @@ -163,6 +189,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 25453abd0..a7e343fb5 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 20e6e4b38..27f4c124a 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 845b774df..a7d3d29a5 100644 --- a/diskann-disk/src/lib.rs +++ b/diskann-disk/src/lib.rs @@ -12,8 +12,11 @@ pub(crate) mod test_utils; 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 9b607e715..62992d82e 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,57 @@ where out.flush()?; Ok(index_size.into_usize()) } + +/// Save real-point adjacency lists in the canonical graph layout. +pub fn save_adjacency_graph

( + adjacency: &[AdjacencyList], + max_degree: u32, + provider: &P, + start_point: u32, + path: &str, +) -> ANNResult +where + P: StorageWriteProvider, +{ + save_graph( + &AdjacencyGraph { + adjacency, + max_degree, + }, + 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::log_index_error(format_args!("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) + } +} diff --git a/diskann-providers/src/storage/mod.rs b/diskann-providers/src/storage/mod.rs index 1233b11f6..f2add5ff2 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 ac5c2101d..cad31bba2 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) } + + /// Borrow the underlying pool for APIs that retain a caller-owned pool. + pub fn as_rayon(self) -> &'a rayon::ThreadPool { + self.0 + } } // Allow use of disallowed methods within this trait to provide custom