Skip to content

PiPNN 3/6: add core graph construction - #1290

Open
weiyaoluo (SeliMeli) wants to merge 52 commits into
pipnn-stack/01-kernelsfrom
pipnn-stack/03-core
Open

PiPNN 3/6: add core graph construction#1290
weiyaoluo (SeliMeli) wants to merge 52 commits into
pipnn-stack/01-kernelsfrom
pipnn-stack/03-core

Conversation

@SeliMeli

@SeliMeli weiyaoluo (SeliMeli) commented Jul 29, 2026

Copy link
Copy Markdown

Purpose

This PR adds provider-independent PiPNN graph construction under diskann::graph::pipnn.

The core accepts a MatrixView, graph policy, and Rayon pool. It returns one adjacency list for each data point.

Main changes

  • PiPNNConfig defines partition and leaf policy.
  • build_graph selects one architecture and one metric marker.
  • partitioning creates overlapping bounded leaves.
  • Partition workers call PartitionMetric to prepare point and leader norms.
  • Partition workers pass active point and dot matrix views to the numerical stage.
  • leaf_build gathers vectors and computes lower-triangular Gram matrices.
  • Leaf workers call LeafMetric to prepare metric-specific norms.
  • finalization applies shared RobustPrune only to overfull rows.
  • Production callers remove the three temporary PR2 dead-code attributes.

The core does not load providers. It does not select start points. It does not serialize indexes.

Required invariants

  • Point IDs fit in u32.
  • 0 < c_min <= c_max.
  • Each fanout value is positive.
  • leaf_k is in the range 1 through 3.
  • Leaf IDs are sorted and unique.
  • Each worker reads only the active buffer prefix.

Review order

  1. Review PiPNNConfig, PiPNNBuildContext, and build_graph in mod.rs.
  2. Review recursive partition flow in partitioning.rs.
  3. Review leaf construction in leaf_build.rs.
  4. Review final graph pruning in finalization.rs.
  5. Confirm that PR2 dead-code attributes disappear in this layer.

Validation

  • Partition tests cover fanout boundaries, recursion, replicas, malformed assignments, and buffer reuse.
  • Leaf tests cover all source types, ID conversion, symmetry, duplicate removal, and allocation errors.
  • Graph tests cover all metrics, deterministic builds, valid IDs, and degree bounds.
  • All-target Clippy passes with -Dwarnings.

Stack

Stack 3/6. Depends on #1287. #1291 adds disk-index integration.

@SeliMeli
weiyaoluo (SeliMeli) requested review from a team and a lite review from Copilot July 29, 2026 11:52

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread diskann-pipnn/src/partitioning.rs Outdated
Copilot AI review requested due to automatic review settings July 29, 2026 13:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Comment thread diskann-pipnn/src/partitioning.rs Outdated
Copilot AI review requested due to automatic review settings July 29, 2026 16:38
@SeliMeli weiyaoluo (SeliMeli) changed the title Pipnn stack/03 core PiPNN 3/6: add core graph construction Jul 29, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 AI review requested due to automatic review settings July 30, 2026 07:47

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 AI review requested due to automatic review settings July 30, 2026 08:26

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>());

Comment thread diskann-pipnn/src/partitioning.rs Outdated
Copilot AI review requested due to automatic review settings July 30, 2026 08:55

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread diskann/src/graph/pipnn/partitioning.rs
Copilot AI review requested due to automatic review settings July 30, 2026 11:26
Pass the existing SortedNeighbors witness into internal prune so source-distance ordering is enforced by type rather than caller documentation.
Reject k outside 1..=3 at context construction so production never reaches an unsupported leaf-kernel width.
Partitioning is the only production source and emits sorted unique IDs. Reject unsorted input linearly and remove the HashSet fallback.
Rely on validated MatrixView shape and avoid cloning PiPNNConfig and its fanout vector.
Run partition orchestration under one architecture/metric specialization and reuse a runtime-sized tracker instead of imposing a fanout cap.
Keep architecture and metric concrete across the Rayon leaf pass and call the generic kernel directly.
Transpose::Ordinary,
point_count,
leader_count,
dimensions,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe I'm misunderstanding something, but it seems that this part of the computation scales linearly with the vector dimension. Given that PiPNN is mainly targeting faster index construction, do we expect this to become a bottleneck on high-dimensional datasets, or has that not been an issue in practice?


diskann_linalg::sgemm_aat_lower(
point_ids.len(),
data.ncols(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar to my understanding that the partitioning step is affected by dimensionality, the leaf-building stage here also seems to have computational complexity that scales roughly linearly with the dimension.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants