diff --git a/.cargo/mutants.toml b/.cargo/mutants.toml index 2d7811d4c..baa95fe77 100644 --- a/.cargo/mutants.toml +++ b/.cargo/mutants.toml @@ -9,4 +9,8 @@ exclude_re = [ "diskann-pipnn/src/partition_kernel\\.rs:.*replace \\* with / in process_(unary|binary)", "diskann-pipnn/src/leaf_kernel\\.rs:.*replace > with >= in .*run_simd", "diskann/src/graph/prune\\.rs:[0-9]+:17: replace < with <= in robust_prune", + "diskann-pipnn/src/(lib|leaf_build|partitioning)\\.rs:.*replace > with (==|>=) in (build_graph_inner|build_leaf_candidates|partition)", + "diskann-pipnn/src/partitioning\\.rs:.*replace < with (==|>|<=) in scatter_assignments", + "diskann-pipnn/src/partitioning\\.rs:.*replace (>|<) with (>=|<=) in global_merge_small", + "diskann-pipnn/src/partitioning\\.rs:.*assignment_stripe_rows", ] diff --git a/Cargo.lock b/Cargo.lock index fd7eeadd3..98915fca3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -661,10 +661,16 @@ name = "diskann-pipnn" version = "0.55.0" dependencies = [ "criterion", + "diskann", "diskann-linalg", + "diskann-utils", "diskann-vector", "diskann-wide", + "half", + "rand", + "rayon", "thiserror 2.0.17", + "tracing", ] [[package]] diff --git a/diskann-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml index 1ff7c3bfd..3f706b02d 100644 --- a/diskann-pipnn/Cargo.toml +++ b/diskann-pipnn/Cargo.toml @@ -11,17 +11,27 @@ license.workspace = true edition.workspace = true [dependencies] +diskann.workspace = true +diskann-linalg.workspace = true +diskann-utils.workspace = true diskann-vector.workspace = true diskann-wide.workspace = true +rand.workspace = true +rayon.workspace = true thiserror.workspace = true +tracing.workspace = true [dev-dependencies] criterion.workspace = true -diskann-linalg.workspace = true +half.workspace = true [[bench]] name = "kernels" harness = false +[[bench]] +name = "core" +harness = false + [lints] workspace = true diff --git a/diskann-pipnn/README.md b/diskann-pipnn/README.md new file mode 100644 index 000000000..4ac1618f0 --- /dev/null +++ b/diskann-pipnn/README.md @@ -0,0 +1,19 @@ +# PiPNN graph construction + +This crate implements the graph-construction stages from [PiPNN: Pick in Partitions for Fast and Accurate ANN Graph Construction](https://arxiv.org/html/2602.21247v1). + +## Boundary + +PiPNN core consumes a dense `MatrixView`, graph policy, and a caller-owned Rayon pool, then returns adjacency lists for the dataset's real point IDs. It does not own start or frozen points, vector or neighbor providers, PQ, disk headers, serialization, or search. Those concerns remain in the outer in-memory and disk pipelines. + +The dense view is intentional: partition assignment and leaf all-pairs kernels operate over the whole source matrix. Materializing provider state inside the algorithm would couple numerical graph construction to storage lifecycle and would require a second dataset copy. Integrations should finish PiPNN scratch before allocating or populating their searchable provider. + +## Policy ownership + +- `PiPNNConfig` owns partition and leaf-selection parameters: leaf bounds, sampling fraction, fanout levels, leaf `k`, and replicas. +- DiskANN graph configuration owns metric, output degree, build-L, alpha, and prune policy. +- Candidate-merging policies are separate validated options; they must not make graph policy fields redundant or silently cap the requested degree. + +## Execution + +A build runs partitioning, leaf construction, candidate merging, then graph finalization. All parallel work executes in the supplied pool. Per-job scratch is initialized through Rayon and is released through normal ownership when its stage completes; the core has no global thread-local buffers or cleanup broadcasts. diff --git a/diskann-pipnn/benches/core.rs b/diskann-pipnn/benches/core.rs new file mode 100644 index 000000000..7a142c063 --- /dev/null +++ b/diskann-pipnn/benches/core.rs @@ -0,0 +1,117 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{hint::black_box, time::Duration}; + +use criterion::{criterion_group, criterion_main, Criterion, Throughput}; +use diskann::graph::{ + config::{self, MaxDegree}, + Config, +}; +use diskann_pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; + +const DIMENSIONS: usize = 128; +const POINTS: usize = 1_024; +const DEGREE: usize = 64; + +fn fixed_data(rows: usize, columns: usize) -> Vec { + (0..rows * columns) + .map(|index| { + ((index.wrapping_mul(1_664_525).wrapping_add(1_013_904_223) % 2_003) as f32 - 1_001.0) + / 1_001.0 + }) + .collect() +} + +fn graph_config(degree: usize) -> Config { + config::Builder::new_with( + degree, + MaxDegree::same(), + 72, + Metric::L2.into(), + |builder| { + builder.alpha(1.2); + }, + ) + .build() + .unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap() +} + +fn benchmark_stage_focused_builds(c: &mut Criterion) { + let data = fixed_data(POINTS, DIMENSIONS); + let view = MatrixView::try_from(data.as_slice(), POINTS, DIMENSIONS).unwrap(); + let pool = pool(); + let mut group = c.benchmark_group("pipnn/core"); + group.throughput(Throughput::Elements(POINTS as u64)); + + // These use the public build boundary so the benchmark does not widen the + // production API. Each workload suppresses unrelated work where possible. + let scenarios = [ + ( + "partition-heavy-full-build", + PiPNNConfig { + c_max: 64, + c_min: 16, + p_samp: 0.1, + fanout: vec![3, 2], + k: 1, + replicas: 1, + }, + POINTS, + ), + ( + "single-leaf-candidates-full-build", + PiPNNConfig { + c_max: POINTS, + c_min: 1, + p_samp: 0.01, + fanout: vec![1], + k: 2, + replicas: 1, + }, + POINTS, + ), + ( + "overfull-finalization-full-build", + PiPNNConfig { + c_max: POINTS, + c_min: 1, + p_samp: 0.01, + fanout: vec![1], + k: 96, + replicas: 1, + }, + DEGREE, + ), + ]; + + for (name, config, degree) in scenarios { + let graph = graph_config(degree); + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + group.bench_function(name, |bencher| { + bencher.iter(|| black_box(build_graph(view, &context).unwrap())); + }); + } + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(3)); + targets = benchmark_stage_focused_builds +} +criterion_main!(benches); diff --git a/diskann-pipnn/src/finalization.rs b/diskann-pipnn/src/finalization.rs new file mode 100644 index 000000000..3aac61fdd --- /dev/null +++ b/diskann-pipnn/src/finalization.rs @@ -0,0 +1,123 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Orders complete PiPNN candidate rows and applies shared Vamana RobustPrune. + +use std::convert::Infallible; + +use diskann::{ + graph::{prune, AdjacencyList, Config}, + neighbor::Neighbor, + utils::VectorRepr, + ANNError, ANNResult, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::{distance::Metric, DistanceFunction}; +use rayon::prelude::*; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum FinalizationError { + #[error("candidate row count {rows} does not match the dataset point count {points}")] + RowCountMismatch { rows: usize, points: usize }, + #[error("candidate ID {candidate} in row {row} is outside a {points}-point dataset")] + InvalidCandidateId { + row: usize, + candidate: u32, + points: usize, + }, +} + +#[derive(Default)] +struct Workspace { + prune: prune::Scratch, + cache: Vec<(f32, Option)>, +} + +pub(crate) fn prune_overfull( + data: MatrixView<'_, T>, + candidates: Vec>, + graph: &Config, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + validate_candidates(&candidates, data.nrows()).map_err(ANNError::opaque)?; + + let degree = graph.pruned_degree().get(); + let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); + let distance = T::distance(metric, Some(data.ncols())); + + // build_graph installs the complete call tree in the caller-owned pool. + #[allow(clippy::disallowed_methods)] + candidates + .into_par_iter() + .enumerate() + .map_init(Workspace::default, |workspace, (source, mut row)| { + if row.len() <= degree { + return Ok(row); + } + + let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_vector = data.row(source); + let pool = workspace.prune.candidates_mut(); + pool.clear(); + pool.try_reserve(row.len()).map_err(ANNError::opaque)?; + pool.extend(row.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + let candidate_count = pool.len(); + let mut context = workspace.prune.as_context(candidate_count); + prune::robust_prune( + &mut context, + policy, + &mut workspace.cache, + Some, + |left, right| { + Ok::<_, Infallible>( + distance.evaluate_similarity( + data.row(*left as usize), + data.row(*right as usize), + ), + ) + }, + |id| id == source_id, + ) + .map_err(ANNError::opaque)?; + + row.clear(); + row.extend_from_slice(workspace.prune.neighbors()); + Ok(row) + }) + .collect() +} + +fn validate_candidates( + candidates: &[AdjacencyList], + points: usize, +) -> Result<(), FinalizationError> { + if candidates.len() != points { + return Err(FinalizationError::RowCountMismatch { + rows: candidates.len(), + points, + }); + } + for (row_id, row) in candidates.iter().enumerate() { + if let Some(&candidate) = row.iter().find(|&&id| id as usize >= points) { + return Err(FinalizationError::InvalidCandidateId { + row: row_id, + candidate, + points, + }); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/finalization/tests.rs b/diskann-pipnn/src/finalization/tests.rs new file mode 100644 index 000000000..62b317d36 --- /dev/null +++ b/diskann-pipnn/src/finalization/tests.rs @@ -0,0 +1,103 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann::graph::{ + config::{self, MaxDegree}, + AdjacencyList, +}; +use diskann_utils::views::MatrixView; + +use super::*; + +fn graph_config(degree: usize) -> Config { + config::Builder::new_with( + degree, + MaxDegree::same(), + degree, + Metric::L2.into(), + |builder| { + builder.alpha(1.2); + }, + ) + .build() + .unwrap() +} + +fn row(ids: impl IntoIterator) -> AdjacencyList { + AdjacencyList::from_iter_untrusted(ids) +} + +#[test] +fn preserves_rows_within_the_degree_bound() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![row([3, 1]), row([]), row([]), row([])]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); +} + +#[test] +fn prunes_an_overfull_row_with_the_vamana_kernel() { + let data = [0.0_f32, 1.0, 2.0, -3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![row([3, 2, 1]), row([]), row([]), row([])]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert!(actual[0].len() <= 2); + assert!(actual[0].contains(1)); +} + +#[test] +fn rejects_invalid_candidate_ids_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![row([1, 3]), row([]), row([])]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::InvalidCandidateId { + row: 0, + candidate: 3, + points: 3, + }) + )); +} + +#[test] +fn rejects_candidate_row_count_mismatch_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![row([]), row([]), row([]), row([])]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::RowCountMismatch { rows: 4, points: 3 }) + )); +} + +#[test] +fn rejects_more_candidates_than_the_shared_position_type_can_represent() { + let count = u16::MAX as usize + 1; + let data = vec![0.0_f32; count + 1]; + let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); + let mut candidates = Vec::with_capacity(count + 1); + candidates.push(row(1..=count as u32)); + candidates.resize_with(count + 1, AdjacencyList::new); + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::>(), + Some(prune::RobustPruneError::TooManyCandidates { actual, max }) + if *actual == count && *max == u16::MAX as usize + )); +} diff --git a/diskann-pipnn/src/leaf_build.rs b/diskann-pipnn/src/leaf_build.rs new file mode 100644 index 000000000..1647ba35a --- /dev/null +++ b/diskann-pipnn/src/leaf_build.rs @@ -0,0 +1,341 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Leaf construction and direct candidate accumulation. + +use std::{ + collections::{HashSet, TryReserveError}, + sync::Mutex, +}; + +use diskann::{graph::AdjacencyList, utils::VectorRepr}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use rayon::prelude::*; + +use crate::leaf_kernel::{ + nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, +}; + +/// Failure while converting leaves into direct graph candidates. +#[derive(Debug, thiserror::Error)] +pub(crate) enum LeafBuildError { + #[error("leaf build requires at least one dimension")] + EmptyDimensions, + #[error("dataset point count {0} exceeds the u32 ID limit")] + TooManyPoints(usize), + #[error("leaf {leaf} is empty")] + EmptyLeaf { leaf: usize }, + #[error("point ID {point} in leaf {leaf} is outside a {points}-point dataset")] + InvalidPointId { + leaf: usize, + point: u32, + points: usize, + }, + #[error("point ID {point} appears more than once in leaf {leaf}")] + DuplicatePointId { leaf: usize, point: u32 }, + #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] + ShapeOverflow { + leaf: usize, + rows: usize, + columns: usize, + }, + #[error("failed to reserve {additional} values for {buffer}")] + Allocation { + buffer: &'static str, + additional: usize, + #[source] + source: TryReserveError, + }, + #[error("failed to convert point {point} in leaf {leaf}")] + Conversion { + leaf: usize, + point: u32, + #[source] + source: diskann::ANNError, + }, + #[error("lower-AAT failed for leaf {leaf}")] + LowerAat { + leaf: usize, + #[source] + source: diskann_linalg::SgemmError, + }, + #[error("nearest-neighbor selection failed for leaf {leaf}")] + Kernel { + leaf: usize, + #[source] + source: LeafKernelError, + }, + #[error("leaf kernel returned local position {position} for a {points}-point leaf")] + InvalidLocalPosition { position: u32, points: usize }, + #[error("candidate row {point} is poisoned")] + PoisonedCandidateRow { point: u32 }, +} + +#[derive(Default)] +struct LeafBuffers { + points: Vec, + dots: Vec, + nearest: Vec, + local_graph: Vec>, + top_k: LeafTopKWorkspace, + seen_ids: HashSet, +} + +impl LeafBuffers { + fn prepare( + &mut self, + leaf: usize, + points: usize, + dimensions: usize, + k: usize, + ) -> Result { + let point_values = points + .checked_mul(dimensions) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: dimensions, + })?; + let dot_values = points + .checked_mul(points) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: points, + })?; + let actual_k = k.min(points.saturating_sub(1)); + let nearest_values = points + .checked_mul(actual_k) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: actual_k, + })?; + + resize("leaf points", &mut self.points, point_values, 0.0)?; + resize("leaf dot products", &mut self.dots, dot_values, 0.0)?; + resize( + "leaf nearest neighbors", + &mut self.nearest, + nearest_values, + LeafNeighbor::default(), + )?; + let additional = points.saturating_sub(self.local_graph.len()); + self.local_graph + .try_reserve(additional) + .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; + self.local_graph.resize_with(points, AdjacencyList::new); + self.local_graph[..points] + .iter_mut() + .for_each(AdjacencyList::clear); + Ok(actual_k) + } +} + +struct DirectCandidates { + rows: Vec>>, +} + +impl DirectCandidates { + fn new(points: usize) -> Result { + let mut rows = Vec::new(); + rows.try_reserve_exact(points) + .map_err(|source| allocation_error("candidate rows", points, source))?; + rows.resize_with(points, || Mutex::new(AdjacencyList::new())); + Ok(Self { rows }) + } + + fn add_leaf( + &self, + point_ids: &[u32], + local_graph: &[AdjacencyList], + ) -> Result<(), LeafBuildError> { + for (&source, additions) in point_ids.iter().zip(local_graph) { + // Every point ID is validated before leaf-local work begins. + let row = &self.rows[source as usize]; + let mut row = row.lock().map_err(|_| poisoned_row(source))?; + row.extend_from_slice(additions); + } + Ok(()) + } + + fn into_rows(self) -> Result>, LeafBuildError> { + let mut output = Vec::new(); + output + .try_reserve_exact(self.rows.len()) + .map_err(|source| allocation_error("candidate output", self.rows.len(), source))?; + for (point, row) in self.rows.into_iter().enumerate() { + let mut row = row.into_inner().map_err(|_| poisoned_row(point as u32))?; + row.sort(); + output.push(row); + } + Ok(output) + } +} + +/// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +pub(crate) fn build_leaf_candidates( + data: MatrixView<'_, T>, + leaves: Vec>, + k: usize, + metric: Metric, +) -> Result>, LeafBuildError> +where + T: VectorRepr + 'static, +{ + if data.ncols() == 0 { + return Err(LeafBuildError::EmptyDimensions); + } + if data.nrows() > u32::MAX as usize { + return Err(LeafBuildError::TooManyPoints(data.nrows())); + } + + let candidates = DirectCandidates::new(data.nrows())?; + leaves.par_iter().enumerate().try_for_each_init( + LeafBuffers::default, + |buffers, (leaf, point_ids)| { + build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + }, + )?; + candidates.into_rows() +} + +fn build_leaf( + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + k: usize, + metric: Metric, + buffers: &mut LeafBuffers, + candidates: &DirectCandidates, +) -> Result<(), LeafBuildError> +where + T: VectorRepr + 'static, +{ + if point_ids.is_empty() { + return Err(LeafBuildError::EmptyLeaf { leaf }); + } + buffers.seen_ids.clear(); + buffers + .seen_ids + .try_reserve(point_ids.len()) + .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; + for &point in point_ids { + if point as usize >= data.nrows() { + return Err(LeafBuildError::InvalidPointId { + leaf, + point, + points: data.nrows(), + }); + } + if !buffers.seen_ids.insert(point) { + return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } + } + let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; + if actual_k == 0 { + return Ok(()); + } + + for (&point, output) in point_ids + .iter() + .zip(buffers.points.chunks_exact_mut(data.ncols())) + { + // Point IDs were validated before the zero-k/singleton fast path. + let row = data.row(point as usize); + T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { + leaf, + point, + source: source.into(), + })?; + } + + diskann_linalg::sgemm_aat_lower( + &buffers.points, + point_ids.len(), + data.ncols(), + &mut buffers.dots, + ) + .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; + nearest_leaf_neighbors( + LeafTopK { + dots: &buffers.dots, + points: point_ids.len(), + metric, + }, + k, + &mut buffers.nearest, + &mut buffers.top_k, + ) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + + add_symmetric_edges( + point_ids, + actual_k, + &buffers.nearest, + &mut buffers.local_graph[..point_ids.len()], + )?; + candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) +} + +fn add_symmetric_edges( + point_ids: &[u32], + k: usize, + nearest: &[LeafNeighbor], + local_graph: &mut [AdjacencyList], +) -> Result<(), LeafBuildError> { + for (source, nearest) in nearest.chunks_exact(k).enumerate() { + for neighbor in nearest { + let target = neighbor.position as usize; + let Some(&target_id) = point_ids.get(target) else { + return Err(LeafBuildError::InvalidLocalPosition { + position: neighbor.position, + points: point_ids.len(), + }); + }; + let source_id = point_ids[source]; + if source_id != target_id { + local_graph[source].push(target_id); + local_graph[target].push(source_id); + } + } + } + Ok(()) +} + +fn resize( + buffer: &'static str, + values: &mut Vec, + len: usize, + value: T, +) -> Result<(), LeafBuildError> { + let additional = len.saturating_sub(values.len()); + values + .try_reserve(additional) + .map_err(|source| allocation_error(buffer, additional, source))?; + values.resize(len, value); + Ok(()) +} + +fn allocation_error( + buffer: &'static str, + additional: usize, + source: TryReserveError, +) -> LeafBuildError { + LeafBuildError::Allocation { + buffer, + additional, + source, + } +} + +fn poisoned_row(point: u32) -> LeafBuildError { + LeafBuildError::PoisonedCandidateRow { point } +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/leaf_build/tests.rs b/diskann-pipnn/src/leaf_build/tests.rs new file mode 100644 index 000000000..efa597dd7 --- /dev/null +++ b/diskann-pipnn/src/leaf_build/tests.rs @@ -0,0 +1,372 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use half::f16; +use std::collections::BTreeSet; + +use super::{ + add_symmetric_edges, allocation_error, build_leaf_candidates, DirectCandidates, LeafBuffers, + LeafBuildError, +}; + +fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() +} + +fn build( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, +) -> Result>, LeafBuildError> +where + T: diskann::utils::VectorRepr + 'static, +{ + pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) +} + +fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() +} + +fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { + let mut graph = vec![BTreeSet::new(); data.len()]; + for (source, left) in data.iter().enumerate() { + let mut nearest: Vec<_> = data + .iter() + .enumerate() + .filter(|(target, _)| *target != source) + .map(|(target, right)| { + let distance = left + .iter() + .zip(right) + .map(|(x, y)| (x - y) * (x - y)) + .sum::(); + (target, distance) + }) + .collect(); + nearest.sort_by(|left, right| { + left.1 + .total_cmp(&right.1) + .then_with(|| left.0.cmp(&right.0)) + }); + for &(target, _) in nearest.iter().take(k) { + graph[source].insert(target as u32); + graph[target].insert(source as u32); + } + } + graph + .into_iter() + .map(|neighbors| neighbors.into_iter().collect()) + .collect() +} + +#[test] +fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + let points = [ + [0.0_f32, 0.0], + [1.0, 0.2], + [3.1, 0.5], + [7.8, 1.4], + [-2.3, 4.1], + [6.7, -3.2], + ]; + let flat: Vec<_> = points.into_iter().flatten().collect(); + + let actual = rows( + build( + view(&flat, points.len(), 2), + &[(0..points.len() as u32).collect()], + 2, + Metric::L2, + ) + .unwrap(), + ); + + assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); +} + +#[test] +fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + + let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + rows(graph), + [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] + ); +} + +#[test] +fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + let dimensions = 9; + let mut data = vec![0.0_f32; 10 * dimensions]; + for row in 1..10 { + data[row * dimensions + row - 1] = 1.0; + } + + let graph = build( + view(&data, 10, dimensions), + &[(0..10).collect()], + 1, + Metric::L2, + ) + .unwrap(); + + assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); + assert!(graph.iter().enumerate().all(|(source, neighbors)| { + neighbors.iter().all(|&target| target as usize != source) + && neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + })); +} + +#[test] +fn global_id_translation_is_independent_of_leaf_order() { + let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; + let leaves = vec![vec![4, 1, 3]]; + + let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + rows(graph), + [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] + ); +} + +fn assert_source_type(data: &[T]) +where + T: diskann::utils::VectorRepr + 'static, +{ + let leaves = vec![vec![0, 1, 2, 3]]; + let graph = build(view(data, 4, 2), &leaves, 1, Metric::L2).unwrap(); + assert_eq!(rows(graph), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); +} + +#[test] +fn gathers_every_supported_source_type_without_full_dataset_conversion() { + assert_source_type(&[0.0_f32, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0]); + assert_source_type(&[0_i8, 0, 1, 0, 2, 0, 3, 0]); + assert_source_type(&[0_u8, 0, 1, 0, 2, 0, 3, 0]); + assert_source_type(&[ + f16::from_f32(0.0), + f16::from_f32(0.0), + f16::from_f32(1.0), + f16::from_f32(0.0), + f16::from_f32(2.0), + f16::from_f32(0.0), + f16::from_f32(3.0), + f16::from_f32(0.0), + ]); +} + +#[test] +fn all_metrics_produce_symmetric_unique_non_self_candidates() { + let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; + let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); + for (source, neighbors) in graph.iter().enumerate() { + assert!(neighbors.iter().all(|&target| target as usize != source)); + assert!(neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32))); + assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); + } + } +} + +#[test] +fn parallel_leaf_schedule_does_not_change_candidate_order() { + let data: Vec = (0..64).map(|value| value as f32).collect(); + let leaves: Vec> = (0..32) + .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) + .collect(); + let pool = pool(); + pool.install(|| { + let expected = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); + assert_eq!(actual, expected); + } + }); +} + +#[test] +fn rejects_invalid_shape_inputs_without_panicking() { + let data = [0.0_f32, 1.0]; + let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); + assert!(matches!( + build(no_dimensions, &[], 1, Metric::L2), + Err(LeafBuildError::EmptyDimensions) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), + Err(LeafBuildError::EmptyLeaf { leaf: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { + leaf: 0, + point: 2, + points: 2 + }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) + )); +} + +#[test] +fn singleton_and_zero_k_leaves_add_no_candidates() { + let data = [0.0_f32, 1.0, 2.0]; + let singleton = build( + view(&data, 3, 1), + &[vec![0], vec![1], vec![2]], + 1, + Metric::L2, + ) + .unwrap(); + let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); + assert!(singleton.iter().chain(&zero_k).all(|row| row.is_empty())); +} + +#[test] +fn reuses_worker_buffers_for_smaller_leaves() { + let mut buffers = LeafBuffers::default(); + buffers.prepare(0, 64, 128, 2).unwrap(); + let points = buffers.points.as_ptr(); + let dots = buffers.dots.as_ptr(); + let nearest = buffers.nearest.as_ptr(); + + buffers.prepare(1, 8, 128, 2).unwrap(); + + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.dots.as_ptr(), dots); + assert_eq!(buffers.nearest.as_ptr(), nearest); +} + +#[test] +fn reports_shape_overflow_before_allocating() { + let mut buffers = LeafBuffers::default(); + assert!(matches!( + buffers.prepare(7, usize::MAX, 2, 1), + Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) + )); +} + +#[test] +fn rejects_an_invalid_kernel_position() { + let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let error = add_symmetric_edges( + &[10, 20], + 1, + &[ + crate::leaf_kernel::LeafNeighbor::new(9, 1.0), + crate::leaf_kernel::LeafNeighbor::new(0, 1.0), + ], + &mut graph, + ) + .unwrap_err(); + assert!(matches!( + error, + LeafBuildError::InvalidLocalPosition { + position: 9, + points: 2 + } + )); +} + +#[test] +fn skips_duplicate_global_ids_without_self_edges() { + let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + add_symmetric_edges( + &[7, 7], + 1, + &[ + crate::leaf_kernel::LeafNeighbor::new(1, 0.0), + crate::leaf_kernel::LeafNeighbor::new(0, 0.0), + ], + &mut graph, + ) + .unwrap(); + assert!(graph.iter().all(|row| row.is_empty())); +} + +#[test] +fn poisoned_candidate_rows_return_errors() { + let candidates = DirectCandidates::new(1).unwrap(); + let _ = std::panic::catch_unwind(|| { + let _guard = candidates.rows[0].lock().unwrap(); + panic!("poison candidate row"); + }); + assert!(matches!( + candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), + Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + )); + assert!(matches!( + candidates.into_rows(), + Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + )); +} + +#[test] +fn allocation_errors_preserve_buffer_context() { + let mut values = Vec::::new(); + let source = values.try_reserve(usize::MAX).unwrap_err(); + let error = allocation_error("test", 1, source); + assert!(matches!( + error, + LeafBuildError::Allocation { + buffer: "test", + additional: 1, + .. + } + )); +} + +#[test] +fn direct_candidate_accumulator_keeps_unique_sorted_rows() { + let candidates = DirectCandidates::new(2).unwrap(); + candidates + .add_leaf( + &[0, 1], + &[ + diskann::graph::AdjacencyList::from_iter_untrusted([1, 1]), + diskann::graph::AdjacencyList::from_iter_untrusted([0]), + ], + ) + .unwrap(); + assert_eq!(rows(candidates.into_rows().unwrap()), [vec![1], vec![0]]); +} diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index 198434b72..6ed06f8d7 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -5,5 +5,200 @@ //! PiPNN graph construction. +mod finalization; +mod leaf_build; pub mod leaf_kernel; pub mod partition_kernel; +mod partitioning; + +use diskann::{ + graph::{AdjacencyList, Config}, + utils::VectorRepr, + ANNError, ANNResult, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use rayon::ThreadPool; + +/// Configuration of PiPNN's partitioning and local-neighbor algorithm. +/// +/// Graph degree, pruning policy, and alpha belong to DiskANN's graph +/// configuration and are supplied separately through [`PiPNNBuildContext`]. +#[derive(Clone, Debug, PartialEq)] +pub struct PiPNNConfig { + /// 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 partition leaders. + pub p_samp: f64, + /// Number of nearest leaders retained at each overlapping partition level. + pub fanout: Vec, + /// Number of nearest neighbors selected within each leaf. + pub k: usize, + /// Number of independent partition passes over the dataset. + pub replicas: usize, +} + +impl PiPNNConfig { + /// Validate the algorithm-specific partition and leaf-build parameters. + pub fn validate(&self) -> ANNResult<()> { + if self.c_max == 0 { + return Err(config_error("c_max must be greater than zero")); + } + if self.c_min == 0 { + return Err(config_error("c_min must be greater than zero")); + } + if self.c_min > self.c_max { + return Err(config_error(format!( + "c_min ({}) must not exceed c_max ({})", + self.c_min, self.c_max + ))); + } + if !self.p_samp.is_finite() || !(0.0..=1.0).contains(&self.p_samp) || self.p_samp == 0.0 { + return Err(config_error(format!( + "p_samp ({}) must be finite and in (0, 1]", + self.p_samp + ))); + } + if self.fanout.is_empty() { + return Err(config_error("fanout must not be empty")); + } + if let Some(&fanout) = self + .fanout + .iter() + .find(|&&fanout| !(1..=partition_kernel::MAX_PARTITION_FANOUT).contains(&fanout)) + { + return Err(config_error(format!( + "fanout ({fanout}) must be in [1, {}]", + partition_kernel::MAX_PARTITION_FANOUT + ))); + } + if self.k == 0 { + return Err(config_error("k must be greater than zero")); + } + if self.replicas == 0 { + return Err(config_error("replicas must be greater than zero")); + } + Ok(()) + } +} + +/// Validated, borrowed policy and execution context for one PiPNN graph build. +#[derive(Debug)] +pub struct PiPNNBuildContext<'a> { + pub(crate) config: PiPNNConfig, + pub(crate) graph: &'a Config, + pub(crate) metric: Metric, + pub(crate) pool: &'a ThreadPool, +} + +impl<'a> PiPNNBuildContext<'a> { + /// Validate and combine PiPNN configuration with outer graph policy. + pub fn new( + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, + ) -> ANNResult { + config.validate()?; + if !graph.alpha().is_finite() || graph.alpha() < 1.0 { + return Err(config_error(format!( + "graph alpha ({}) must be finite and at least 1", + graph.alpha() + ))); + } + if graph.prune_kind() != metric.into() { + return Err(config_error(format!( + "graph prune kind {:?} is incompatible with metric {metric:?}", + graph.prune_kind() + ))); + } + + Ok(Self { + config, + graph, + metric, + pool, + }) + } +} + +/// Build PiPNN adjacency for real rows in `data`. +/// +/// This is the core algorithm boundary. Search entry-point selection, frozen nodes, +/// providers, serialization, and index writers belong to the outer build pipelines. +/// For raw `u8` and `i8` rows, `CosineNormalized` is evaluated as `Cosine` because +/// those representations are converted to f32 scratch but are not unit-normalized. +pub fn build_graph( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + context.pool.install(|| build_graph_inner(data, context)) +} + +fn build_graph_inner( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + if data.nrows() == 0 { + return Err(ANNError::log_dimension_mismatch_error( + "PiPNN requires at least one data row".into(), + )); + } + if data.ncols() == 0 { + return Err(ANNError::log_dimension_mismatch_error( + "PiPNN requires at least one data dimension".into(), + )); + } + if data.nrows() > u32::MAX as usize { + return Err(config_error(format!( + "dataset row count ({}) exceeds the u32 graph ID limit", + data.nrows() + ))); + } + data.nrows().checked_mul(data.ncols()).ok_or_else(|| { + ANNError::log_dimension_mismatch_error(format!( + "PiPNN dataset shape {} x {} overflows usize", + data.nrows(), + data.ncols() + )) + })?; + let metric = effective_metric::(context.metric); + + let partition = partitioning::PartitionConfig::from(&context.config); + let leaves = tracing::info_span!("pipnn.partition") + .in_scope(|| partitioning::partition(data, partition, metric))?; + let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) + .map_err(ANNError::opaque) + })?; + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) +} + +fn effective_metric(metric: Metric) -> Metric { + use std::any::TypeId; + + if metric == Metric::CosineNormalized + && (TypeId::of::() == TypeId::of::() || TypeId::of::() == TypeId::of::()) + { + Metric::Cosine + } else { + metric + } +} + +#[track_caller] +fn config_error(message: impl std::fmt::Display) -> ANNError { + ANNError::log_index_config_error("PiPNN".into(), message.to_string()) +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/partitioning.rs b/diskann-pipnn/src/partitioning.rs new file mode 100644 index 000000000..a30c329ac --- /dev/null +++ b/diskann-pipnn/src/partitioning.rs @@ -0,0 +1,640 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Deterministic overlapping partition construction for PiPNN. +//! +//! The stage maps real dataset rows to bounded leaf ID lists. Numerical work +//! reuses the partition kernel and dense GEMM; scratch belongs to the Rayon +//! iterator that uses it, so no thread-local cleanup protocol is required. + +use std::collections::HashSet; + +use diskann::{utils::VectorRepr, ANNError, ANNResult}; +use diskann_linalg::Transpose; +use diskann_utils::views::MatrixView; +use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; +use rand::{prelude::IndexedRandom, SeedableRng}; +use rayon::prelude::*; + +use crate::{ + partition_kernel::{nearest_leaders, PartitionTopK}, + PiPNNConfig, +}; + +// Private algorithm and batching constants live together. None are user policy. +const PARTITION_SEED: u64 = 1_000; +const REPLICA_SEED_STEP: u64 = 7_919; +const LEADER_CAP: usize = 1_000; +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; +const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; +const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; +const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; +const SCATTER_STRIPE_ROWS: usize = 65_536; +const MAX_PARTITION_ITERATIONS: usize = 30; + +/// Policy owned by the partition stage. Leaf-neighbor and merge settings do +/// not cross this boundary. +#[derive(Clone, Debug)] +pub(crate) struct PartitionConfig { + c_max: usize, + c_min: usize, + p_samp: f64, + fanout: Vec, + replicas: usize, +} + +impl From<&PiPNNConfig> for PartitionConfig { + fn from(config: &PiPNNConfig) -> Self { + Self { + c_max: config.c_max, + c_min: config.c_min, + p_samp: config.p_samp, + fanout: config.fanout.clone(), + replicas: config.replicas, + } + } +} + +/// A partition failure with enough context to diagnose non-progressing input. +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub(crate) enum PartitionError { + #[error("PiPNN cannot partition an empty dataset")] + EmptyDataset, + #[error("PiPNN cannot partition vectors with zero dimensions")] + EmptyDimensions, + #[error("dataset has {0} rows, which exceeds the u32 ID limit")] + TooManyPoints(usize), + #[error("{buffer} shape {rows} x {cols} overflows usize")] + ShapeOverflow { + buffer: &'static str, + rows: usize, + cols: usize, + }, + #[error( + "partition stopped after {limit} iterations with an oversized cluster of size \ + {size} at level {level}" + )] + IterationLimit { + size: usize, + level: usize, + limit: usize, + }, + #[error("partition produced an invalid leaf of size {size}; expected 1..={limit}")] + InvalidLeaf { size: usize, limit: usize }, + #[error("invalid {buffer} length: expected {expected}, got {actual}")] + InvalidBufferLength { + buffer: &'static str, + expected: usize, + actual: usize, + }, + #[error("partition worker did not publish its result")] + MissingWorkerResult, +} + +struct WorkItem { + indices: Vec, + level: usize, + seed: u64, +} + +#[derive(Default)] +struct StripeBuffers { + points: Vec, + dots: Vec, + row_scales: Vec, +} + +/// Partition every configured replica into overlapping bounded leaves. +/// +/// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped +/// to the private leader bound), assigns every point to its nearest `fanout` +/// leaders for the current level, and recurses only on oversized clusters. +/// Levels beyond `fanout.len()` retain one leader assignment. Completed small +/// leaves are merged without exceeding `c_max`; every input point must remain +/// covered once per replica. The caller installs the operation in its pool. +pub(crate) fn partition( + data: MatrixView<'_, T>, + config: PartitionConfig, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let points = data.nrows(); + if points == 0 { + return Err(ANNError::opaque(PartitionError::EmptyDataset)); + } + if data.ncols() == 0 { + return Err(ANNError::opaque(PartitionError::EmptyDimensions)); + } + if points > u32::MAX as usize { + return Err(ANNError::opaque(PartitionError::TooManyPoints(points))); + } + + let mut leaves = Vec::new(); + for replica in 0..config.replicas { + let seed = replica_seed(replica); + let mut replica_leaves = partition_replica(data, &config, metric, seed)?; + leaves + .try_reserve(replica_leaves.len()) + .map_err(ANNError::opaque)?; + leaves.append(&mut replica_leaves); + } + validate_leaves(&leaves, config.c_max)?; + Ok(leaves) +} + +fn partition_replica( + data: MatrixView<'_, T>, + config: &PartitionConfig, + metric: Metric, + seed: u64, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let initial_indices = point_ids(data.nrows())?; + if data.nrows() <= config.c_max { + let mut leaves = Vec::new(); + leaves.try_reserve_exact(1).map_err(ANNError::opaque)?; + leaves.push(initial_indices); + return Ok(leaves); + } + + let mut leaves = Vec::new(); + let mut work = Vec::new(); + work.try_reserve_exact(1).map_err(ANNError::opaque)?; + work.push(WorkItem { + indices: initial_indices, + level: 0, + seed, + }); + + for _ in 0..MAX_PARTITION_ITERATIONS { + if work.is_empty() { + return global_merge_small(leaves, config.c_min, config.c_max); + } + + let mut results = Vec::new(); + results + .try_reserve_exact(work.len()) + .map_err(ANNError::opaque)?; + results.resize_with(work.len(), || None); + // build_graph installs this complete private call tree into the + // caller-owned pool; the indexed fill cannot escape that pool. + #[allow(clippy::disallowed_methods)] + results + .par_iter_mut() + .zip(work.into_par_iter()) + .for_each(|(slot, item)| { + *slot = Some(partition_one_level(data, config, metric, item)); + }); + + let mut next_work = Vec::new(); + for result in results { + let (mut pending, mut finished) = + result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??; + next_work + .try_reserve(pending.len()) + .map_err(ANNError::opaque)?; + leaves + .try_reserve(finished.len()) + .map_err(ANNError::opaque)?; + next_work.append(&mut pending); + leaves.append(&mut finished); + } + work = next_work; + } + + if work.is_empty() { + return global_merge_small(leaves, config.c_min, config.c_max); + } + let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { + return global_merge_small(leaves, config.c_min, config.c_max); + }; + Err(ANNError::opaque(PartitionError::IterationLimit { + size: largest.indices.len(), + level: largest.level, + limit: MAX_PARTITION_ITERATIONS, + })) +} + +fn partition_one_level( + data: MatrixView<'_, T>, + config: &PartitionConfig, + metric: Metric, + item: WorkItem, +) -> ANNResult<(Vec, Vec>)> +where + T: VectorRepr + Send + Sync, +{ + let points = item.indices.len(); + let fanout = config.fanout.get(item.level).copied().unwrap_or(1); + let leaders = sample_leaders( + &item.indices, + config.p_samp, + mix_seed(item.seed, points as u64), + )?; + let clusters = assign_to_leaders(data, &item.indices, &leaders, fanout, metric)?; + + let mut pending = Vec::new(); + let mut finished = Vec::new(); + pending + .try_reserve(clusters.len()) + .map_err(ANNError::opaque)?; + finished + .try_reserve(clusters.len()) + .map_err(ANNError::opaque)?; + let child_seed = mix_seed(item.seed, points as u64); + for cluster in clusters { + if cluster.is_empty() { + continue; + } + if cluster.len() <= config.c_max { + finished.push(cluster); + } else { + pending.push(WorkItem { + indices: cluster, + level: item.level + 1, + seed: child_seed, + }); + } + } + Ok((pending, finished)) +} + +fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { + let count = sample_num_leaders(points.len(), sampling_fraction); + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut leaders = Vec::new(); + leaders.try_reserve_exact(count).map_err(ANNError::opaque)?; + leaders.extend(points.choose_multiple(&mut rng, count).copied()); + Ok(leaders) +} + +fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { + ((points as f64 * sampling_fraction).ceil() as usize) + .clamp(2, LEADER_CAP) + .min(points) +} + +fn replica_seed(replica: usize) -> u64 { + PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) +} + +// A single LCG mixer derives recursive seeds. Wrapping makes the mapping stable +// across debug/release builds and supported platforms. +fn mix_seed(seed: u64, salt: u64) -> u64 { + seed.wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(salt) +} + +fn assign_to_leaders( + data: MatrixView<'_, T>, + points: &[u32], + leaders: &[u32], + fanout: usize, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let dimensions = data.ncols(); + let leader_values_len = checked_area("leader data", leaders.len(), dimensions)?; + let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; + gather_rows(data, leaders, &mut leader_values)?; + + let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { + filled_vec(leaders.len(), 0.0f32)? + } else { + Vec::new() + }; + for (scale, row) in leader_scales + .iter_mut() + .zip(leader_values.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + if metric == Metric::Cosine { + *scale = scale.sqrt(); + } + } + + let fanout = fanout.min(leaders.len()); + let assignment_len = checked_area("partition assignments", points.len(), fanout)?; + let mut assignments = filled_vec(assignment_len, 0u32)?; + let stripe_rows = assignment_stripe_rows(leaders.len()); + let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; + + // build_graph pins this terminal operation to the caller-owned pool. + #[allow(clippy::disallowed_methods)] + assignments + .par_chunks_mut(assignment_stripe) + .enumerate() + .try_for_each_init(StripeBuffers::default, |buffers, (stripe, output)| { + let first = stripe * stripe_rows; + let rows = output.len() / fanout; + let point_values_len = checked_area("point stripe", rows, dimensions)?; + let dots_len = checked_area("dot-product stripe", rows, leaders.len())?; + resize_fallible(&mut buffers.points, point_values_len, 0.0)?; + resize_fallible(&mut buffers.dots, dots_len, 0.0)?; + gather_rows(data, &points[first..first + rows], &mut buffers.points)?; + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + rows, + leaders.len(), + dimensions, + 1.0, + &buffers.points, + &leader_values, + None, + &mut buffers.dots, + ) + .map_err(ANNError::opaque)?; + + let row_scales = if metric == Metric::Cosine { + resize_fallible(&mut buffers.row_scales, rows, 0.0)?; + for (scale, row) in buffers + .row_scales + .iter_mut() + .zip(buffers.points.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + } + buffers.row_scales.as_slice() + } else { + &[] + }; + nearest_leaders( + PartitionTopK { + dots: &buffers.dots, + rows, + leaders: leaders.len(), + row_scales, + leader_scales: &leader_scales, + metric, + }, + fanout, + output, + ) + .map_err(ANNError::opaque) + })?; + + scatter_assignments(points, &assignments, fanout, leaders.len()) +} + +fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> +where + T: VectorRepr, +{ + let expected = checked_area("gather output", indices.len(), data.ncols())?; + if output.len() != expected { + return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "gather output", + expected, + actual: output.len(), + })); + } + for (&index, row) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { + T::as_f32_into(data.row(index as usize), row).map_err(Into::::into)?; + } + Ok(()) +} + +fn scatter_assignments( + points: &[u32], + assignments: &[u32], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + if points.len() < PARALLEL_SCATTER_MIN_POINTS { + return scatter_serial(points, assignments, fanout, leaders); + } + + let assignment_stripe = checked_area("scatter assignment stripe", SCATTER_STRIPE_ROWS, fanout)?; + let stripes = points.len().div_ceil(SCATTER_STRIPE_ROWS); + let mut partials = Vec::new(); + partials + .try_reserve_exact(stripes) + .map_err(ANNError::opaque)?; + partials.resize_with(stripes, || None); + // See the pool invariant at the other partition terminal operations. + #[allow(clippy::disallowed_methods)] + partials + .par_iter_mut() + .zip( + points + .par_chunks(SCATTER_STRIPE_ROWS) + .zip(assignments.par_chunks(assignment_stripe)), + ) + .for_each(|(slot, (points, assignments))| { + *slot = Some(scatter_serial(points, assignments, fanout, leaders)); + }); + + let mut locals = Vec::new(); + locals + .try_reserve_exact(stripes) + .map_err(ANNError::opaque)?; + for result in partials { + locals.push(result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??); + } + + let mut sizes = filled_vec(leaders, 0usize)?; + for local in &locals { + for (size, cluster) in sizes.iter_mut().zip(local) { + *size = size.checked_add(cluster.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "cluster size", + rows: *size, + cols: cluster.len(), + }) + })?; + } + } + + let mut clusters = clusters_with_capacities(&sizes)?; + for local in locals { + for (cluster, part) in clusters.iter_mut().zip(local) { + debug_assert!(cluster.capacity().saturating_sub(cluster.len()) >= part.len()); + cluster.extend(part); + } + } + Ok(clusters) +} + +fn scatter_serial( + points: &[u32], + assignments: &[u32], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + let mut sizes = filled_vec(leaders, 0usize)?; + for &leader in assignments { + let Some(size) = sizes.get_mut(leader as usize) else { + return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: leaders, + actual: leader as usize + 1, + })); + }; + *size = size.checked_add(1).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "cluster size", + rows: *size, + cols: 1, + }) + })?; + } + let mut clusters = clusters_with_capacities(&sizes)?; + for (&point, row) in points.iter().zip(assignments.chunks_exact(fanout)) { + for &leader in row { + clusters[leader as usize].push(point); + } + } + Ok(clusters) +} + +fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { + let mut clusters = Vec::new(); + clusters + .try_reserve_exact(sizes.len()) + .map_err(ANNError::opaque)?; + for &size in sizes { + let mut cluster = Vec::new(); + cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + clusters.push(cluster); + } + Ok(clusters) +} + +fn global_merge_small( + leaves: Vec>, + c_min: usize, + c_max: usize, +) -> ANNResult>> { + let mut merged = Vec::new(); + let mut small_leaves = Vec::new(); + merged.try_reserve(leaves.len()).map_err(ANNError::opaque)?; + small_leaves + .try_reserve(leaves.len()) + .map_err(ANNError::opaque)?; + for leaf in leaves { + if leaf.len() >= c_min { + merged.push(leaf); + } else { + small_leaves.push(leaf); + } + } + if small_leaves.is_empty() { + return Ok(merged); + } + + let mut small = HashSet::new(); + small.try_reserve(c_max).map_err(ANNError::opaque)?; + + for leaf in small_leaves { + let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf merge", + rows: small.len(), + cols: leaf.len(), + }) + })?; + if combined > c_max { + merged.push(drain_sorted(&mut small)?); + } + small.try_reserve(leaf.len()).map_err(ANNError::opaque)?; + small.extend(leaf); + if small.len() >= c_min { + merged.push(drain_sorted(&mut small)?); + } + } + + if !small.is_empty() { + let mut remainder = drain_sorted(&mut small)?; + if remainder.len() < c_min { + if let Some(last) = merged.last_mut() { + remainder.retain(|id| !last.contains(id)); + let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf tail merge", + rows: last.len(), + cols: remainder.len(), + }) + })?; + if combined <= c_max { + last.try_reserve(remainder.len()) + .map_err(ANNError::opaque)?; + last.append(&mut remainder); + last.sort_unstable(); + } + } + } + if !remainder.is_empty() { + merged.push(remainder); + } + } + + validate_leaves(&merged, c_max)?; + Ok(merged) +} + +fn drain_sorted(set: &mut HashSet) -> ANNResult> { + let mut values = Vec::new(); + values + .try_reserve_exact(set.len()) + .map_err(ANNError::opaque)?; + values.extend(set.drain()); + values.sort_unstable(); + Ok(values) +} + +fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { + if let Some(leaf) = leaves + .iter() + .find(|leaf| leaf.is_empty() || leaf.len() > c_max) + { + return Err(ANNError::opaque(PartitionError::InvalidLeaf { + size: leaf.len(), + limit: c_max, + })); + } + Ok(()) +} + +fn point_ids(points: usize) -> ANNResult> { + let mut ids = Vec::new(); + ids.try_reserve_exact(points).map_err(ANNError::opaque)?; + ids.extend(0..points as u32); + Ok(ids) +} + +fn filled_vec(len: usize, value: T) -> ANNResult> { + let mut values = Vec::new(); + values.try_reserve_exact(len).map_err(ANNError::opaque)?; + values.resize(len, value); + Ok(values) +} + +fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { + values + .try_reserve(len.saturating_sub(values.len())) + .map_err(ANNError::opaque)?; + values.resize(len, value); + Ok(()) +} + +fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { + rows.checked_mul(cols) + .ok_or_else(|| ANNError::opaque(PartitionError::ShapeOverflow { buffer, rows, cols })) +} + +fn assignment_stripe_rows(leaders: usize) -> usize { + (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::())) + .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/partitioning/tests.rs b/diskann-pipnn/src/partitioning/tests.rs new file mode 100644 index 000000000..ef7b4b517 --- /dev/null +++ b/diskann-pipnn/src/partitioning/tests.rs @@ -0,0 +1,342 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::{Matrix, MatrixView}; +use diskann_vector::{distance::Metric, Half}; + +use super::*; + +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PartitionConfig { + PartitionConfig { + c_max, + c_min, + p_samp: 0.25, + fanout, + replicas, + } +} + +fn clustered_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let row = position / dimensions; + let column = position % dimensions; + position += 1; + (row / 8) as f32 * 10.0 + column as f32 * 0.01 + row as f32 * 0.001 + } + }), + points, + dimensions, + ) +} + +fn directional_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let row = position / dimensions; + let column = position % dimensions; + position += 1; + let angle = std::f32::consts::TAU * row as f32 / points as f32; + match column { + 0 => angle.cos(), + 1 => angle.sin(), + _ => 0.0, + } + } + }), + points, + dimensions, + ) +} + +fn sorted_memberships(leaves: &[Vec]) -> Vec> { + let mut memberships: Vec> = leaves + .iter() + .map(|leaf| { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids + }) + .collect(); + memberships.sort(); + memberships +} + +fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + assert!(leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max)); + let mut counts = vec![0usize; points]; + for leaf in leaves { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); + for &id in leaf { + assert!((id as usize) < points); + counts[id as usize] += 1; + } + } + assert!(counts.iter().all(|&count| count >= replicas)); +} + +#[test] +fn returns_one_leaf_at_and_below_c_max() { + for points in [7, 8] { + let data = clustered_data(points, 3); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + assert_eq!(leaves, vec![(0..points as u32).collect::>()]); + } +} + +#[test] +fn partition_is_fixed_seed_deterministic_and_bounded() { + let data = clustered_data(96, 8); + let config = config(4, 16, vec![3, 2], 2); + + let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); + let second = partition(data.as_view(), config, Metric::L2).unwrap(); + + assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); + assert_valid_partition(&first, 96, 16, 2); + assert!(first.iter().map(Vec::len).sum::() > 96 * 2); +} + +#[test] +fn recursion_after_fanout_levels_falls_back_to_one() { + let data = clustered_data(80, 4); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + + assert_valid_partition(&leaves, 80, 8, 1); +} + +#[test] +fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { + let data = Matrix::new(1.0f32, 24, 4); + let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = error.downcast::().unwrap(); + + assert!(matches!( + error, + PartitionError::IterationLimit { + size: 24, + limit: MAX_PARTITION_ITERATIONS, + .. + } + )); +} + +#[test] +fn global_merge_canonicalizes_small_leaf_membership() { + let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + + let merged = global_merge_small(leaves, 4, 8).unwrap(); + + assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); +} + +#[test] +fn global_merge_never_overfills_before_reaching_c_min() { + let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + + let merged = global_merge_small(leaves, 11, 11).unwrap(); + + assert_eq!( + merged, + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] + ); +} + +#[test] +fn global_merge_fills_exact_capacity_before_flushing() { + let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + + assert_eq!(merged, vec![vec![0, 1, 2, 3]]); +} + +#[test] +fn replicas_cover_every_point_once_or_more_per_replica() { + let data = directional_data(72, 5); + let leaves = partition( + data.as_view(), + config(3, 12, vec![3, 2], 3), + Metric::CosineNormalized, + ) + .unwrap(); + + assert_valid_partition(&leaves, 72, 12, 3); +} + +#[test] +fn supported_source_types_share_partition_contract() { + let f32_data: Vec = (0..64 * 4).map(|value| (value % 23) as f32).collect(); + let half_data: Vec = f32_data.iter().copied().map(Half::from_f32).collect(); + let u8_data: Vec = f32_data.iter().map(|value| *value as u8).collect(); + let i8_data: Vec = u8_data.iter().map(|value| *value as i8 - 11).collect(); + let config = config(2, 16, vec![2, 1], 1); + + let f32_leaves = partition( + MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let half_leaves = partition( + MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let u8_leaves = partition( + MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let i8_leaves = partition( + MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), + config, + Metric::L2, + ) + .unwrap(); + + for leaves in [&f32_leaves, &half_leaves, &u8_leaves, &i8_leaves] { + assert_valid_partition(leaves, 64, 16, 1); + } + assert_eq!( + sorted_memberships(&f32_leaves), + sorted_memberships(&half_leaves) + ); + assert_eq!( + sorted_memberships(&f32_leaves), + sorted_memberships(&u8_leaves) + ); + assert_eq!( + sorted_memberships(&u8_leaves), + sorted_memberships(&i8_leaves) + ); +} + +#[test] +fn all_metrics_produce_valid_partitions() { + let data = directional_data(64, 8); + let config = config(2, 20, vec![2], 1); + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); + assert_valid_partition(&leaves, 64, 20, 1); + } +} + +#[test] +fn leader_count_is_bounded() { + assert_eq!(sample_num_leaders(1, 1.0), 1); + assert_eq!(sample_num_leaders(10, 0.01), 2); + assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); +} + +#[test] +fn replica_seed_derivation_is_stable_and_distinct() { + assert_eq!(replica_seed(0), 1_000); + assert_eq!(replica_seed(1), 8_919); +} + +#[test] +fn leader_assignment_handles_multiple_stripes() { + let points = 2_048; + let data: Vec = (0..points).map(|point| point as f32).collect(); + let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); + let point_ids: Vec = (0..points as u32).collect(); + + let clusters = assign_to_leaders(data, &point_ids, &[0, 2_047], 1, Metric::L2).unwrap(); + + assert_eq!(clusters[0], (0..1_024).collect::>()); + assert_eq!(clusters[1], (1_024..2_048).collect::>()); +} + +#[test] +fn parallel_scatter_matches_serial_order() { + let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); + let assignments: Vec = points + .iter() + .flat_map(|point| [point % 7, (point + 3) % 7]) + .collect(); + + let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + + assert_eq!(actual, expected); +} + +#[test] +fn rejects_empty_dataset() { + let data = Matrix::::new(0.0, 0, 4); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDataset + ); +} + +#[test] +fn rejects_zero_dimensions() { + let data = Matrix::::new(0.0, 4, 0); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDimensions + ); +} + +#[test] +fn rejects_invalid_gather_output_length() { + let data = Matrix::::new(0.0, 2, 2); + let error = gather_rows(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "gather output", + expected: 4, + actual: 3, + } + ); +} + +#[test] +fn rejects_assignment_to_an_unknown_leader() { + let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: 2, + actual: 3, + } + ); +} + +#[test] +fn rejects_empty_and_oversized_leaves() { + for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { + let error = validate_leaves(&leaves, 2).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidLeaf { size, limit: 2 } + ); + } +} diff --git a/diskann-pipnn/src/tests.rs b/diskann-pipnn/src/tests.rs new file mode 100644 index 000000000..d5689eede --- /dev/null +++ b/diskann-pipnn/src/tests.rs @@ -0,0 +1,27 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; +use half::f16; + +#[test] +fn integer_normalized_cosine_uses_unnormalized_cosine() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let expected = if metric == Metric::CosineNormalized { + Metric::Cosine + } else { + metric + }; + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); + } +} diff --git a/diskann-pipnn/tests/build_graph.rs b/diskann-pipnn/tests/build_graph.rs new file mode 100644 index 000000000..04086be8e --- /dev/null +++ b/diskann-pipnn/tests/build_graph.rs @@ -0,0 +1,225 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann::graph::config::{self, MaxDegree}; +use diskann_pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use half::f16; +use rand::{rngs::StdRng, Rng, SeedableRng}; + +fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 4, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + } +} + +fn graph_config(metric: Metric, degree: usize) -> diskann::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { + builder.alpha(1.2); + }) + .build() + .unwrap() +} + +fn pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() +} + +fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() +} + +fn assert_graph_invariants( + graph: &[diskann::graph::AdjacencyList], + points: usize, + degree: usize, +) { + assert_eq!(graph.len(), points); + for (source, row) in graph.iter().enumerate() { + assert!(row.len() <= degree); + let mut sorted = row.to_vec(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), row.len()); + assert!(row + .iter() + .all(|&id| (id as usize) < points && id as usize != source)); + } +} + +#[test] +fn builds_a_single_leaf_graph_for_real_dataset_ids() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let graph = graph_config(Metric::L2, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + + let graph = graph_config(Metric::L2, 1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let pruned = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&pruned, 4, 1); + for (source, neighbors) in pruned.iter().enumerate() { + assert_eq!(source.abs_diff(neighbors[0] as usize), 1); + } +} + +#[test] +fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let graph = graph_config(Metric::L2, 1); + let pool = pool(2); + let config = PiPNNConfig { + c_max: 5, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 4, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 5, 1); + assert!(actual.iter().all(|row| row.len() == 1)); +} + +#[test] +fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + let graph = graph_config(Metric::L2, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); + let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + + assert!(build_graph(no_rows, &context).is_err()); + assert!(build_graph(no_columns, &context).is_err()); +} + +#[test] +fn supports_every_source_type_and_metric() { + fn build(values: &[T], metric: Metric) { + let data = MatrixView::try_from(values, 6, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + let actual = build_graph(data, &context).unwrap(); + assert_graph_invariants(&actual, 6, 2); + } + + let values = [ + 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, + ]; + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + build(&values, metric); + } + build(&values.map(f16::from_f32), Metric::L2); + build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); + build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); +} + +#[test] +fn integer_normalized_cosine_matches_cosine() { + fn assert_match(values: &[T]) { + let data = MatrixView::try_from(values, 8, 2).unwrap(); + let pool = pool(2); + let build = |metric| { + let graph = graph_config(metric, 2); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + rows(build_graph(data, &context).unwrap()) + }; + assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); + } + + assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); + assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); +} + +#[test] +fn is_deterministic_for_a_fixed_pool_size() { + let data: Vec = (0..96 * 4) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let first = build_graph(data, &context).unwrap(); + let second = build_graph(data, &context).unwrap(); + + assert_eq!(first, second); + assert_graph_invariants(&first, 96, 8); +} + +#[test] +fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { + let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); + for case in 0..24 { + let points = rng.random_range(4..=32); + let dimensions = rng.random_range(1..=8); + let c_max = rng.random_range(4..=points.min(12)); + let c_min = rng.random_range(1..=c_max); + let degree = rng.random_range(1..=points.min(8)); + let values: Vec = (0..points * dimensions) + .map(|_| rng.random_range(-10.0..10.0)) + .collect(); + let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, degree); + let pool = pool(2); + let config = PiPNNConfig { + c_max, + c_min, + p_samp: 0.5, + fanout: vec![2], + k: rng.random_range(1..=3), + replicas: rng.random_range(1..=2), + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context) + .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); + assert_graph_invariants(&actual, points, degree); + } +} diff --git a/diskann-pipnn/tests/config.rs b/diskann-pipnn/tests/config.rs new file mode 100644 index 000000000..923faee9e --- /dev/null +++ b/diskann-pipnn/tests/config.rs @@ -0,0 +1,127 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann::graph::config::{self, MaxDegree}; +use diskann_pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann_vector::distance::Metric; + +fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 512, + c_min: 64, + p_samp: 0.01, + fanout: vec![10, 3], + k: 2, + replicas: 1, + } +} + +fn graph_config(metric: Metric, alpha: f32) -> diskann::graph::Config { + config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + builder.alpha(alpha); + }) + .build() + .unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap() +} + +#[test] +fn accepts_the_six_algorithm_parameters_with_outer_graph_policy() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); +} + +#[test] +fn rejects_each_invalid_algorithm_parameter() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + let mut cases = [ + PiPNNConfig { + c_max: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 513, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 0.0, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: -0.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 1.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: f64::NAN, + ..pipnn_config() + }, + PiPNNConfig { + fanout: Vec::new(), + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![1, 0], + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![17], + ..pipnn_config() + }, + PiPNNConfig { + k: 0, + ..pipnn_config() + }, + PiPNNConfig { + replicas: 0, + ..pipnn_config() + }, + ]; + + for config in &mut cases { + let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + } +} + +#[test] +fn rejects_graph_policy_for_a_different_metric() { + let graph = graph_config(Metric::InnerProduct, 1.2); + let pool = pool(); + + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("prune kind")); +} + +#[test] +fn rejects_invalid_outer_alpha() { + let pool = pool(); + for alpha in [0.9, f32::NAN, f32::INFINITY] { + let graph = graph_config(Metric::L2, alpha); + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("alpha")); + } +}