PiPNN 3/6: add core graph construction

.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",
 ]

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,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<u32>,
+    cache: Vec<(f32, Option<u32>)>,
+}
+
+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)| {
+            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<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;