PiPNN 3/6: add core graph construction - #1290
Conversation
There was a problem hiding this comment.
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
PiPNNConfigvalidation and aPiPNNBuildContextthat 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 (nouse 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.
| *scale = FastL2NormSquared.evaluate(row); | ||
| if metric == Metric::Cosine { | ||
| *scale = scale.sqrt(); | ||
| } | ||
| } |
5be0c9c to
50047c6
Compare
| 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) | ||
| } |
There was a problem hiding this comment.
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_ofis 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
Normis imported but never used in this module, which will tripunused_importswarnings (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::*;
There was a problem hiding this comment.
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_rowsusesTypeId::of::<T>(), which implicitly requiresT: 'static. Making that bound explicit here avoids surprising/indirect trait-bound errors later and matches the publicbuild_graphboundary (which already requires'static).
fn gather_rows<T>(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()>
where
T: VectorRepr,
1324668 to
857e200
Compare
There was a problem hiding this comment.
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 withstd::mem::size_of(or add an explicit import).
let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>());
| use diskann::{utils::VectorRepr, ANNError, ANNResult}; | ||
| use diskann_linalg::Transpose; | ||
| use diskann_utils::views::MatrixView; | ||
| use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; |
857e200 to
f642204
Compare
There was a problem hiding this comment.
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_leafis executed from a Rayon parallel context (viabuild_leaf_candidates), so it should also explicitly requireT: Send + Syncto reflect the actual thread-safety requirement.
where
T: VectorRepr + 'static,
{
diskann-pipnn/src/leaf_build/tests.rs:154
assert_source_typeforwardsTinto the parallel leaf build path, so it should also includeSend + Syncbounds 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_candidatesuses Rayon parallel iteration overdata, soTmust beSend + 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) requiresT: 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,
{
| fanout: usize, | ||
| leaders: usize, | ||
| ) -> ANNResult<Vec<Vec<u32>>> { | ||
| let mut sizes = filled_vec(leaders, 0usize)?; |
f642204 to
b1181d6
Compare
| 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() { | ||
| point_count | ||
| } else { | ||
| point_count.next_power_of_two() / 2 | ||
| }; | ||
| point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS) | ||
| } |
There was a problem hiding this comment.
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_errblock is not rustfmt-formatted (indentation inside the struct literal is off). This is likely to failcargo fmt --checkin 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_leafcallsbuffers.prepare(...)(which can grow the dot-product buffer topoint_count^2) before checking whether any neighbors will be produced. Forrequested_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)?;
Move core construction under diskann::graph::pipnn so finalization can reuse private RobustPrune state. Remove standalone PiPNN Cargo benchmarks, group public tests by PiPNN behavior, and delete duplicate or non-discriminating cases.
Restore the finalization adapter overflow oracle while consolidating stage-focused builds in the shared diskann IAI target.
Keep sorting, workspace allocation, source exclusion, and adjacency rewriting in PiPNN finalization. The shared internal kernel now sees only prepared candidates and state.
Address the renamed internal modules and preserve graph Config alpha behavior without adding PiPNN-specific validation.
There was a problem hiding this comment.
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 adduse 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() {
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
fanoutleaders, so child partitions overlap.klocal companions per point; these are candidates, not final graph degreeR.PiPNNConfigowns partition/leaf policy; DiskANNgraph::Configowns metric compatibility,R, alpha, prune kind, and saturation.Code map
diskann/src/graph/pipnn/mod.rsPiPNNConfig::validatechecks leaf bounds, sampling, fanout,k, and replicas.PiPNNBuildContext::newcombines algorithm policy with graph policy, metric, and Rayon pool.build_graphinstalls the full call tree in that pool and orchestrates all stages.partitioning.rsleaf_build.rsA · Aᵀ, prepared leaf kernel, local/global ID translation, symmetric candidate insertion, sort/dedup.finalization.rsRunchanged and sends only overfull rows through private shared RobustPrune.diskann/benches/benchmarks_iai/pipnn_core.rsadds partition-heavy, single-leaf, and overfull-finalization cases to the shared IAI target.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 globalu32IDs → merge/sort/deduplicate candidates per source → preserve bounded rows and RobustPrune only overfull rows → return one adjacency list per input row.Invariants and boundaries
usize; row count must fitu32graph IDs.0 < c_min <= c_max; sampling is finite in(0,1]; fanout is nonempty and bounded;kand replicas are nonzero.graph::Configcontract.c_max.map_initstate. No TLS or cleanup broadcast exists.Review path
PiPNNConfig,PiPNNBuildContext, andbuild_graphto establish policy ownership.Test architecture
The suite keeps one independent differential oracle per numerical kernel. Private tests target branch-heavy state; integration tests exercise only public composition. Removed cases were duplicate formula-sharing differentials, constructor-only success checks, or weaker invariant-only wrappers. Exact retained oracles cover:
u16candidate overflow;Rust test names are behavior phrases (
rejects_*,preserves_*,*_is_*); no mechanicaltest_orshould_prefix is used for new PiPNN tests.Validation
graph::Config.cargo test -p diskann --features pipnnand Clippy with all targets pass with and without default features.Stack relation
Stack 3/6. Depends on #1287 numerical kernels and #1288 private RobustPrune. #1291 adds production disk integration; #1294 adds benchmark entry points; #1295 adds optional HashPrune merging.
Stack 3/6: #1288 → #1291