PiPNN 3/6: add core graph construction

[weiyaoluo (SeliMeli)]

PiPNN constructs an ANN graph without beam-searching a partially built graph for every inserted point. It creates overlapping randomized partitions, performs dense leaf-local distance work, merges symmetric candidates, then applies DiskANN graph policy. This PR adds that provider-independent core under diskann::graph::pipnn.

The core consumes a borrowed dense MatrixView<T>, validated PiPNN/graph policy, and a caller-owned Rayon pool. It returns adjacency for real dataset IDs. Providers, start/frozen points, PQ, disk headers, serialization, and search remain outside.

Concepts

• A sampled leader names a child partition.
• Each point keeps its nearest fanout leaders, so child partitions overlap.
• Oversized children recurse until bounded leaves remain.
• Each leaf selects k local companions per point; these are candidates, not final graph degree R.
• Independent replicas repeat partitioning with deterministic derived seeds.
• PiPNNConfig owns partition/leaf policy; DiskANN graph::Config owns metric compatibility, R, alpha, prune kind, and saturation.

Code map

1. diskann/src/graph/pipnn/mod.rs
   • PiPNNConfig::validate checks leaf bounds, sampling, fanout, k, and replicas.
   • PiPNNBuildContext::new combines algorithm policy with graph policy, metric, and Rayon pool.
   • build_graph installs the full call tree in that pool and orchestrates all stages.
2. partitioning.rs
   • deterministic replicas, leader sampling, striped GEMM assignment, ordered scatter, recursion, small-leaf merge, and final leaf validation.
3. leaf_build.rs
   • per-job reusable buffers, lower A · Aᵀ, prepared leaf kernel, local/global ID translation, symmetric candidate insertion, sort/dedup.
4. finalization.rs
   • moves rows at or below R unchanged and sends only overfull rows through private shared RobustPrune.
5. Tests are co-located by production seam: kernel, partitioning, leaf assembly/CSR, finalization, configuration, and complete graph construction.

End-to-end flow

Row-major MatrixView<T> → validate dataset/configuration → create deterministic overlapping leaves for each replica → gather each leaf and compute lower-triangular Gram matrix → select symmetric leaf-local neighbors → translate local positions to global u32 IDs → merge/sort/deduplicate candidates per source → preserve bounded rows and RobustPrune only overfull rows → return one adjacency list per input row.

Invariants and boundaries

• Input has at least one row and dimension; area must not overflow usize; row count must fit u32 graph IDs.
• 0 < c_min <= c_max; sampling is finite in (0,1]; fanout is nonempty and bounded; k and replicas are nonzero.
• Graph prune kind matches metric; alpha remains owned and validated (or not) by the existing graph::Config contract.
• Every replica covers every point; leaves may overlap but none is empty or exceeds c_max.
• Seed derivation, scatter order, and final candidate sets are deterministic for fixed input/pool policy.
• Grow-only scratch may retain a larger high-water allocation; all reads use the active prefix.
• Partition pool locks cover lease/return only, never GEMM. Leaf jobs own their map_init state. No TLS or cleanup broadcast exists.
• Core output uses real point IDs only and owns no storage/search lifecycle.

Review path

1. Start with PiPNNConfig, PiPNNBuildContext, and build_graph to establish policy ownership.
2. Follow one oversized partition through leader sampling, stripes, ordered scatter, recursion, and small-leaf merge.
3. Follow one leaf through gather, lower-triangle kernel, symmetric CSR, and global candidate insertion.
4. Finish in finalization; compare its private RobustPrune policy with graph configuration.
5. Review ownership/reuse and allocation-error paths after algorithm flow.

Test architecture

The suite keeps one independent differential oracle per numerical kernel. Private seam and module-level composition tests live beside their implementation. Removed cases were duplicate formula-sharing differentials, constructor-only success checks, or weaker invariant-only wrappers. Exact retained oracles cover:

• partition seed/order/coverage, fanout exhaustion, merge flush boundaries, stripes, conversion, and validation;
• leaf all-pairs reference, symmetric candidate/CSR assembly, duplicate/global IDs, conversion, buffer reuse, and poisoned/allocation errors;
• finalization exact bounded/overfull output, invalid IDs/list counts, and propagated u16 candidate overflow;
• public graph shape, IDs, degree, metrics/types, determinism, and fixed-seed randomized invariants.

Rust test names are behavior phrases (rejects_*, preserves_*, *_is_*); no mechanical test_ or should_ prefix is used for new PiPNN tests.

Validation

• Private seam tests plus co-located configuration/build composition cases cover this layer, including a regression that PiPNN adds no alpha validation beyond graph::Config.
• cargo test -p diskann --features pipnn and Clippy with all targets pass with and without default features.

Stack relation

Stack 3/6. Depends on #1315 private RobustPrune and #1287 numerical kernels. #1291 adds production disk integration; #1294 adds benchmark entry points; #1295 adds optional HashPrune merging.

Stack 3/6: #1287 → #1291

[Copilot]

Pull request overview

This PR introduces the “core” PiPNN build pipeline in the diskann-pipnn crate, wiring together deterministic partitioning, leaf-local candidate construction, and final pruning into a public build_graph API with a validated build context.

Changes:

• Adds PiPNNConfig validation and a PiPNNBuildContext that binds PiPNN policy to DiskANN graph pruning policy and a caller-owned Rayon thread pool.
• Implements the three main stages: partitioning (partitioning.rs), leaf candidate construction (leaf_build.rs), and final pruning via shared Vamana robust prune (finalization.rs).
• Adds comprehensive unit/integration tests and a Criterion benchmark for core scenarios; updates dependencies, lockfile, and mutation-test exclusions.

Reviewed changes

Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
diskann-pipnn/src/lib.rs	Adds public PiPNN API (PiPNNConfig, PiPNNBuildContext, build_graph) and stage orchestration.
diskann-pipnn/src/partitioning.rs	Implements deterministic overlapping partition construction and leader assignment/scatter.
diskann-pipnn/src/partitioning/tests.rs	Adds unit tests covering partition determinism, invariants, error cases, and helpers.
diskann-pipnn/src/leaf_build.rs	Builds leaf-local symmetric k-NN candidates and accumulates global candidates safely in parallel.
diskann-pipnn/src/leaf_build/tests.rs	Adds unit tests for candidate correctness, invariants, type support, and error handling.
diskann-pipnn/src/finalization.rs	Orders/prunes candidate rows using shared robust_prune and validates candidate IDs/shape.
diskann-pipnn/src/finalization/tests.rs	Adds unit tests for pruning behavior and candidate validation failures.
diskann-pipnn/src/tests.rs	Tests effective_metric behavior for integer cosine-normalized handling.
diskann-pipnn/tests/config.rs	Integration tests for config validation and graph-policy compatibility checks.
diskann-pipnn/tests/build_graph.rs	Integration tests for end-to-end graph building, invariants, determinism, and type/metric support.
diskann-pipnn/benches/core.rs	Adds a Criterion benchmark for stage-focused core build scenarios.
diskann-pipnn/Cargo.toml	Updates crate dependencies/dev-dependencies and registers the new core benchmark target.
Cargo.lock	Records dependency graph changes for the updated diskann-pipnn crate dependencies.
.cargo/mutants.toml	Adds mutation-test exclusions for key PiPNN public boundary checks and partitioning invariants.

Comments suppressed due to low confidence (1)

diskann-pipnn/src/partitioning.rs:604

• size_of::<f32>() is used without being in scope (no use std::mem::size_of; and not qualified), so this function won’t compile as written.

fn assignment_stripe_rows(leaders: usize) -> usize {
    (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
        .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

diskann-pipnn/src/partitioning.rs:637

• size_of is used without being in scope (std::mem::size_of), which will not compile. Qualify the call or import it.

fn assignment_stripe_rows(leaders: usize) -> usize {
    (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
        .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}

diskann-pipnn/src/partitioning.rs:19

• Norm is imported but never used in this module, which will trip unused_imports warnings (and can become CI failures under -D warnings). Remove it from the import list.

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::*;

[Copilot]

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

diskann-pipnn/src/partitioning.rs:390

• gather_rows uses TypeId::of::<T>(), which implicitly requires T: 'static. Making that bound explicit here avoids surprising/indirect trait-bound errors later and matches the public build_graph boundary (which already requires 'static).

fn gather_rows<T>(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()>
where
    T: VectorRepr,

[Copilot]

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

diskann-pipnn/src/partitioning.rs:641

• size_of::<f32>() is used without being imported or qualified, which will fail to compile. Qualify it with std::mem::size_of (or add an explicit import).

    let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>());

[Copilot]

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (4)

diskann-pipnn/src/leaf_build.rs:222

• build_leaf is executed from a Rayon parallel context (via build_leaf_candidates), so it should also explicitly require T: Send + Sync to reflect the actual thread-safety requirement.

where
    T: VectorRepr + 'static,
{

diskann-pipnn/src/leaf_build/tests.rs:154

• assert_source_type forwards T into the parallel leaf build path, so it should also include Send + Sync bounds to match the production requirements.

fn assert_source_type<T>(data: &[T])
where
    T: diskann::utils::VectorRepr + 'static,
{

diskann-pipnn/src/leaf_build.rs:193

• build_leaf_candidates uses Rayon parallel iteration over data, so T must be Send + Sync. Making this explicit in the signature avoids confusing trait-bound errors at call sites and documents the thread-safety requirement.

This issue also appears on line 220 of the same file.

where
    T: VectorRepr + 'static,
{

diskann-pipnn/src/leaf_build/tests.rs:35

• This test helper calls build_leaf_candidates, which (via Rayon) requires T: Send + Sync. Add the bounds here so the test continues to compile once the production signature is tightened.

This issue also appears on line 151 of the same file.

where
    T: diskann::utils::VectorRepr + 'static,
{

[Copilot]

Pull request overview

Copilot reviewed 15 out of 16 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.

Suppressed comments (2)

diskann/src/graph/pipnn/leaf_build.rs:344

• This map_err block is not rustfmt-formatted (indentation inside the struct literal is off). This is likely to fail cargo fmt --check in CI; reformat the block to match the rest of the file.

    let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len())
        .map_err(|error| LeafBuildError::Kernel {
        leaf,
        source: LeafKernelError::InvalidBufferLength {
            buffer: "leaf dot-product matrix",
            expected: dot_count,
            actual: error.into_inner().len(),
        },
    })?;

diskann/src/graph/pipnn/leaf_build.rs:305

• build_leaf calls buffers.prepare(...) (which can grow the dot-product buffer to point_count^2) before checking whether any neighbors will be produced. For requested_k == 0 (and trivially for singleton leaves), we can return early after ID validation and avoid unnecessary allocation/work.

    let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?;

[Copilot]

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann/src/graph/pipnn/partitioning.rs:794

• size_of::<f32>() isn’t in the Rust prelude; this will fail to compile unless you import it. Qualify it or add use std::mem::size_of;.

fn assignment_stripe_point_count(leader_count: usize) -> usize {
    let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count.max(1) * size_of::<f32>());
    let point_count = if point_count.is_power_of_two() {

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (pipnn-stack/01-kernels@84dff03). Learn more about missing BASE report.

Additional details and impacted files

@@                    Coverage Diff                    @@
##             pipnn-stack/01-kernels    #1290   +/-   ##
=========================================================
  Coverage                          ?   91.28%           
=========================================================
  Files                             ?      517           
  Lines                             ?    98609           
  Branches                          ?        0           
=========================================================
  Hits                              ?    90013           
  Misses                            ?     8596           
  Partials                          ?        0           

Flag	Coverage Δ
miri	91.28% <ø> (?)
unittests	90.95% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.