[weiyaoluo (SeliMeli)]

Summary

Add PiPNN graph construction under diskann::graph::pipnn, behind the opt-in diskann/pipnn feature.

build_graph borrows one contiguous MatrixView of f32, f16, u8 or i8 vectors and returns one adjacency list for each point, in three steps. The build context carries the f32 row length, which differs from the stored row length for compressed types such as MinMax.

• Partition: split the dataset into overlapping leaves of at most c_max points with the partition kernel from PiPNN 2/6: add numerical kernels #1287.
• Leaf build: select the leaf_k nearest neighbors of each point inside each leaf with the leaf kernel from PiPNN 2/6: add numerical kernels #1287, and merge the pairs into one candidate list per point.
• Finalization: prune each list longer than the graph degree with the RobustPrune shared with Vamana (PiPNN 1/6: extract shared RobustPrune #1315).

The core does not load providers, select start points or serialize indexes. This is part 3 of the PiPNN stack; #1291 integrates it into the disk index build. Review in pipeline order: mod.rs, partitioning.rs, leaf_build.rs, finalization.rs.

Design

Take the graph policy from the DiskANN Config. Degree, alpha, max_occlusion_size and the prune kind come from Config, as in Vamana. PiPNNConfig holds the partition and leaf parameters, with NonZeroUsize counts and a PiPNNConfigError from validate.

Dispatch once per build. build_graph selects the SIMD architecture and the metric, so every step compiles for one of each.

Split clusters one level at a time. A split samples ceil(p_samp * n) leaders, clamped to [2, max_leaders] (default 1000, as in the paper), assigns each point to its fanout nearest leaders, and groups the points by leader. Children of at most c_max points become leaves. Each split seeds its children from its own random stream, and children keep the input order of their points, so the leaves do not depend on the thread count.

End partitioning at max_depth. A level whose fanout is at least the leader count would copy the cluster into every child, so the cluster moves to the next level unsplit. A cluster that is still larger than c_max at max_depth (default 30), for example copies of one vector, is cut into equal leaves in point order, and the build logs one warning. Small leaves are kept as they are: on BigANN-10M, merging leaves below 64 points changed recall@10 by at most 0.0006, the run-to-run spread, and cost 3.7% to 10.9% of build time.

Assign points in stripes of 256. A stripe gathers its points as f32 and ranks them against all leaders in one kernel call. Sweeps on four CPU families kept this height within 1.4% of the best on average. A cache-size rule did no better, because a VM reports a last-level cache that other tenants share.

Convert f16 with the instructions of the running CPU. Crates that depend on diskann do not inherit its x86-64-v3 target flag, so the generic VectorRepr::as_f32_into converts f16 in software for them, which made gathers 4.6 to 22 times slower on an EPYC 7763.

Lock each candidate list once per leaf. A leaf job collects its pairs in both directions locally, then appends them to each point's shared list under one lock. The lists stay unsorted, because final pruning ranks candidates by distance; sorting them cost 10 to 12 s of serial time per BigANN-10M build with no recall change. List order follows job completion and only breaks distance ties, so leaves are reproducible and the graph is not.

Prune only overfull lists. A longer list is ranked by full-precision distance to its source, cut to the max_occlusion_size nearest candidates as in Vamana, and pruned with the shared RobustPrune.

How the build works, and the branches that tests cover

build_graph builds the graph in three steps:

1. Partition: split the points into small leaves that overlap.
2. Leaf build: inside each leaf, connect each point to its nearest points.
3. Finalization: prune each point's neighbor list to the graph degree.

Kernels, top-k and RobustPrune are lower layers with their own tests. build_graph has an end-to-end test on the 256-point SIFT sample.

build_graph (mod.rs)

• No points or no dimensions: returns an error before any row is read.
• Metric: L2, cosine, normalized cosine and inner product each select their own kernels.

Partition (partitioning.rs)

A split samples some points as leaders and puts each point into the groups of its fanout nearest leaders. A point is in several groups, so groups overlap. A group of at most c_max points is a leaf; a larger group is split again at the next level. fanout[level] sets the fanout of each level, and max_depth limits the number of levels.

• partition: splits the dataset level by level until every group is a leaf.
  • Data within c_max: one leaf with every point, no split.
• split_cluster: one split of one group.
  • Level past the end of fanout: fanout 1, so a point joins only one group.
  • Leader count is ceil(points * p_samp), at least 2 and at most max_leaders.
  • Fanout > 1 and fanout >= leader count: every point would join every group, so the group moves to the next level unchanged.
  • Level max_depth: the group is cut into equal leaves in point order, and partition logs one warning. A group can stay larger than c_max when its points are copies of one vector, or under inner product when one leader wins every point.
  • A child group with no points (a leader that no point chose, for example a NaN leader): dropped.
• assign_to_leaders: finds the fanout nearest leaders of each point, in batches (stripes) of 256 points. Unit tests use stripes of 8.
  • More points than one stripe: each stripe, including the partial last one, writes the results of its own points.
  • A smaller stripe after a larger one reuses the buffers and reads only its own part.
• group_by_leader: turns each point's leader list into one point group per leader.
  • A point whose distances to all leaders are NaN has no leader and joins no group.

Leaf build (leaf_build.rs)

In each leaf, each point picks its leaf_k nearest points of that leaf, and each picked pair becomes an edge in both directions. A point is in several leaves, so its edges from all leaves merge into one candidate list.

• add_leaf_candidates: one leaf.
  • leaf_k at or above the leaf size: each point picks every other point.
  • Single-point leaf: no edges.
  • A pair that the kernel cannot rank (NaN or infinite distance): no edge.
  • Overlapping leaves find the same pair: the candidate list holds it once.
  • The next leaf on the same thread reuses the buffers and gets none of the earlier leaf's edges.

Finalization (finalization.rs)

Merged lists can hold more candidates than the graph degree. Vamana RobustPrune cuts such a list: it keeps the nearest candidate first, and drops a candidate when a kept neighbor is more than alpha times nearer to it than the source is.

• PruneWorkspace::prune: one list.
  • List within the degree: unchanged, in its order.
  • Longer list: keeps the max_occlusion_size nearest candidates, then runs RobustPrune.
  • The next list reuses the workspace and keeps no state of the previous one.
• prune_overfull: every list. Each result stays with its source point.

Conversion (conversion.rs)

The kernels compute in f32. gather_as_f32 copies the rows of one stripe or leaf into one f32 matrix. The f32 row length comes from the build context: compressed types such as MinMax store more elements per row.

• f16: uses the CPU's FP16 instructions. Other types use VectorRepr::as_f32_into.
• A row length that does not match the context: an error.

[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,
{

[Mark Hildebrand (hildebrandmw)]

From a high level architectural perspective, the design here forces monomorphization of "glue level" code based on the architecture type, the metric type, and the data type.

Almost all of these can be moved behind an enum and/or trait objects I think without really sacrificing anything (see my comment in leaf_build.rs). Granted, doing this in the leaf_build stage is a little easier than in the partitioning stage due to the associated type on PartitionMetric. But even in that context, there are coarsed grained operations on concrete types that live behind a generic - meaning it's a good candidate for type erasure.

My other large architectural concern for this review pass is the state of configuration. Having configuration object dedicated solely to the algorithm (and not pulling in an arbitrary subset of graph::Config) will make this code far easier to maintain and use independently from the graph code.

[Mark Hildebrand (hildebrandmw)]

Thanks! Love to see the consolidation into well defined trait boundaries. I think others have reviewed the mechanics of the algorithm enough, I'm limiting my reviews to the higher level components and user facing side.

My last requests are in two buckets:

1. Please make rowmajor::Buffer a little more general purpose so it can be used more readily in other places (and rename to rowmajor::Resizeable).
2. I have some suggestions for how to decouple the implementation from rigid use of diskann_vector::distance::Metric without necessarily requiring a whole slew of user facing traits. I think it would be a good idea to do this preemptively so we can be more easily experiment with different kernel types in the future.