PiPNN 3/6: add core graph construction

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

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<T>`, 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.

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<f32> {
+    (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);

diskann-pipnn/src/finalization.rs

@@ -0,0 +1,165 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! Final graph-degree enforcement through the shared Vamana RobustPrune kernel.
+//!
+//! Candidate merging may produce more than `R` IDs for a point. This stage first
+//! validates every global ID, then processes rows independently in the caller's
+//! Rayon pool. Rows already within the degree bound are returned without any
+//! distance work. Overfull rows are converted to source-distance candidates,
+//! passed through RobustPrune, and rewritten from the selected output.
+//!
+//! The shared kernel owns occlusion and alpha-round semantics; this adapter owns
+//! only contiguous dataset access and distance specialization for the source
+//! representation.
+//!
+//! ```text
+//! candidate rows ──> validate row count and every global ID
+//!                                      │
+//!                         ┌────────────┴────────────┐
+//!                         v                         v
+//!                    len <= R                  len > R
+//!                  return row       source-distance candidates
+//!                                                   │
+//!                                                   v
+//!                                           shared RobustPrune
+//!                                                   │
+//!                                                   v
+//!                                           rewrite same row owner
+//! ```
+//!
+//! | Path | Distance evaluations | Allocation behavior |
+//! | --- | --- | --- |
+//! | bounded row | none | move row directly to output |
+//! | overfull row | source and occlusion distances | reuse Rayon-job workspace |
+
+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,
+    },
+}
+
+/// Per-Rayon-job state retained across rows.
+///
+/// `prune` owns candidate/state/output buffers. `cache` stores provider lookup
+/// results required by the shared kernel. Reusing both avoids per-node
+/// allocations, which would otherwise dominate finalization for millions of
+/// short rows.
+#[derive(Default)]
+struct Workspace {
+    prune: prune::Scratch<u32>,
+    cache: Vec<(f32, Option<u32>)>,
+}
+
+/// Validate candidate IDs and prune only rows whose length exceeds graph degree.
+pub(crate) fn prune_overfull<T>(
+    data: MatrixView<'_, T>,
+    candidates: Vec<AdjacencyList<u32>>,
+    graph: &Config,
+    metric: Metric,
+) -> ANNResult<Vec<AdjacencyList<u32>>>
+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)| {
+            // Candidate accumulators already enforce uniqueness. A bounded row
+            // therefore satisfies the graph policy without distance evaluation.
+            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)),
+                )
+            }));
+            // as_context sorts the active candidate prefix by source distance.
+            // The callback below is needed only for selected-to-candidate
+            // occlusion checks; dimension specialization stays in `distance`.
+            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<u32>],
+    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;