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) | ||
| } |
|
@SeliMeli please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
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::*;
Adds provider-independent PiPNN partition, leaf-build, direct-candidate, and finalization stages.
Review key points
MatrixView+ graph policy + caller-owned Rayon pool.RefCell, providers, start points, or serialization.Stack 3/6: #1288 → #1291