Skip to content

PiPNN 4/6: integrate disk build pipeline - #1291

Open
weiyaoluo (SeliMeli) wants to merge 359 commits into
mainfrom
pipnn-stack/04-integration
Open

weiyaoluo (SeliMeli) wants to merge 359 commits into
mainfrom
pipnn-stack/04-integration

Conversation

@SeliMeli

@SeliMeli weiyaoluo (SeliMeli) commented Jul 29, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

This PR connects PiPNN graph construction to the production disk-index builder.

It uses the existing DiskANN graph format. It does not add a PiPNN-specific index format.

Main changes

  • BuildAlgorithm accepts explicit PiPNN parameters.
  • The common builder validates PiPNN before expensive work.
  • Explicit PiPNN selection never falls back to Vamana.
  • The PiPNN adapter loads the dense matrix and calls diskann::graph::pipnn::build_graph.
  • The adapter selects the existing sampled medoid.
  • The canonical graph writer stores degree, start point, and adjacency rows.
  • The common pipeline continues PQ training and disk layout.

The PiPNN core returns rows for real points only. The disk pipeline still owns frozen points and persisted metadata.

Review order

  1. Review build/configuration/build_algorithm.rs.
  2. Review algorithm selection in builder/build.rs.
  3. Review the PiPNN adapter in builder/build/pipnn.rs.
  4. Review storage/bin.rs graph validation and serialization.
  5. Confirm that PQ and disk assembly use the common pipeline.

Validation

  • Tests reject invalid PiPNN configuration before artifact creation.
  • Tests reject dataset shape and point-count mismatches.
  • Tests check degree, medoid, and frozen-count headers.
  • Tests confirm explicit PiPNN selection under a low memory budget.
  • Writer tests reject invalid degrees and IDs before file creation.
  • Windows GNU and AArch64 checks pass.

Stack

Stack 4/6. Depends on #1290. #1294 adds benchmark routes.

Copilot AI lite review requested due to automatic review settings July 29, 2026 11:53

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 integrates PiPNN graph construction into the DiskANN disk-index build pipeline and exposes it through the benchmark CLI inputs, including automated fixture-based tests to validate the new JSON configuration and execution paths.

Changes:

  • Add a BuildAlgorithm selector (Vamana vs PiPNN) with JSON-facing PiPNN parameters, plus memory-estimate-based fallback to Vamana for disk builds.
  • Implement PiPNN graph building adapters for disk builds and in-memory benchmark builds, including vector-store “flat prefix” access for dense layouts.
  • Add a shared cached CLI fixture runner and new PiPNN benchmark fixtures (disk + in-memory) to validate end-to-end CLI wiring.

Reviewed changes

Copilot reviewed 33 out of 37 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
diskann-providers/src/utils/rayon_util.rs Exposes underlying Rayon pool reference for APIs that need &rayon::ThreadPool.
diskann-providers/src/storage/mod.rs Re-exports new adjacency-graph save helper.
diskann-providers/src/storage/bin.rs Adds save_adjacency_graph for canonical adjacency graph serialization.
diskann-providers/src/model/graph/provider/async_/fast_memory_vector_provider.rs Exposes unsafe dense-prefix slice access to support PiPNN build path.
diskann-providers/src/model/graph/provider/async_/common.rs Implements AlignedMemoryVectorStore::flat_prefix and adds a contiguous-prefix test.
diskann-disk/src/lib.rs Re-exports PiPNN parameters (feature-gated) and exposes build-algorithm selection.
diskann-disk/src/build/mod.rs Re-exports BuildAlgorithm and (feature-gated) PiPNN parameters.
diskann-disk/src/build/configuration/mod.rs Introduces build-algorithm configuration module and re-exports its types.
diskann-disk/src/build/configuration/disk_index_build_parameter.rs Adds build_algorithm, PiPNN constructor, PiPNN config conversion, and memory-estimate-based fallback.
diskann-disk/src/build/configuration/build_algorithm.rs New: serde-driven BuildAlgorithm enum and PiPNN JSON parameter struct.
diskann-disk/src/build/builder/build/pipnn.rs New: PiPNN graph build adapter producing a canonical on-disk adjacency graph.
diskann-disk/src/build/builder/build/pipnn/tests.rs New: disk-build integration tests for PiPNN selection, fallback, and header correctness.
diskann-disk/src/build/builder/build.rs Routes graph construction to PiPNN (when selected) while keeping existing Vamana pipeline as default.
diskann-disk/Cargo.toml Adds optional diskann-pipnn dependency + feature, plus test deps for new tests.
diskann-benchmark/src/main.rs Refactors registry creation and adds PiPNN fixture test hook (feature-gated).
diskann-benchmark/src/inputs/graph_index.rs Adds PiPNN build-algorithm selection for graph-index build inputs (feature-gated).
diskann-benchmark/src/inputs/disk.rs Adds alpha, makes quantization_type optional, and adds disk build-algorithm selection + validation.
diskann-benchmark/src/index/build.rs Adds PiPNN in-memory build path leveraging dense flat-prefix when possible.
diskann-benchmark/src/index/benchmarks.rs Wires PiPNN build path selection into benchmark build flow.
diskann-benchmark/src/disk_index/build.rs Wires disk BuildAlgorithm into disk-index build parameter construction.
diskann-benchmark/Cargo.toml Adds optional PiPNN dependency/feature and enables fixture runner for dev-tests.
diskann-benchmark/tests/pipnn/inmemory/stdout.txt New: expected output for in-memory PiPNN CLI fixture.
diskann-benchmark/tests/pipnn/inmemory/stdin.txt New: commands for in-memory PiPNN CLI fixture.
diskann-benchmark/tests/pipnn/inmemory/input.json New: PiPNN in-memory input JSON fixture.
diskann-benchmark/tests/pipnn/disk/stdout.txt New: expected output for disk PiPNN CLI fixture.
diskann-benchmark/tests/pipnn/disk/stdin.txt New: commands for disk PiPNN CLI fixture.
diskann-benchmark/tests/pipnn/disk/input.json New: PiPNN disk input JSON fixture.
diskann-benchmark-runner/src/lib.rs Exposes new fixture runner module for tests/downstream crates (feature-gated).
diskann-benchmark-runner/src/fixture.rs New: shared cached CLI fixture runner implementation.
diskann-benchmark-runner/src/fixture/tests.rs New: unit tests for fixture runner behavior.
diskann-benchmark-runner/src/app.rs Updates integration-test documentation and switches tests to the shared fixture runner.
diskann-benchmark-runner/Cargo.toml Adds test-fixtures feature + optional tempfile dep for downstream fixture usage.
Cargo.lock Adds/records diskann-pipnn and serde_json dependency changes.
.cargo/mutants.toml Excludes new fixture overwrite behavior from mutation testing.

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

Copilot AI review requested due to automatic review settings July 29, 2026 13:10
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/04-integration branch from 8134eb3 to ce5f054 Compare July 29, 2026 13:10
@SeliMeli
weiyaoluo (SeliMeli) requested a review from a team 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 12 out of 13 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

diskann-disk/src/build/configuration/disk_index_build_parameter.rs:17

  • estimate_pipnn_peak_memory uses size_of::<...>() unqualified, but this module doesn’t import std::mem::size_of. This will fail to compile under the repo’s Rust toolchain.
use std::num::NonZeroUsize;

use diskann::ANNError;
#[cfg(feature = "pipnn")]
use diskann::ANNResult;
use thiserror::Error;

#[cfg(feature = "pipnn")]
use super::PiPNNParameters;
use super::{BuildAlgorithm, QuantizationType};

diskann-disk/src/build/configuration/disk_index_build_parameter.rs:113

  • DiskIndexBuildParameters is publicly re-exported (via diskann-disk’s public API) and this change removes Copy from its derives. That is a breaking API change for downstream users that relied on implicit copies; if this is intentional, it should be called out explicitly (and potentially gated or versioned accordingly).
/// Parameters specific for disk index construction.
#[derive(Clone, PartialEq, Debug)]
pub struct DiskIndexBuildParameters {

Copilot AI review requested due to automatic review settings July 29, 2026 16:38
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/04-integration branch from ce5f054 to 3893460 Compare July 29, 2026 16:38
@SeliMeli weiyaoluo (SeliMeli) changed the title Pipnn stack/04 integration PiPNN 4/6: integrate disk build pipeline 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 12 out of 13 changed files in this pull request and generated 1 comment.

Comment thread diskann-providers/src/storage/bin.rs Resolved
Copilot AI review requested due to automatic review settings July 30, 2026 07:47
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/04-integration branch from 3893460 to f836f69 Compare 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 12 out of 13 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

diskann-disk/src/build/configuration/disk_index_build_parameter.rs:204

  • disk_index_build_parameter.rs enables #![warn(missing_docs)], and the new pub(crate) fn pipnn_config has no doc comment. If CI promotes warnings to errors, this will fail builds; even if not, it adds new lint noise.
    #[cfg(feature = "pipnn")]
    pub(crate) fn pipnn_config(&self) -> Option<diskann_pipnn::PiPNNConfig> {
        match &self.build_algorithm {
            BuildAlgorithm::PiPNN(config) => Some(config.into()),
            BuildAlgorithm::Vamana => None,
        }
    }

diskann-disk/src/build/builder/build/pipnn.rs:49

  • pipnn::build_graph reads the full dataset into memory (read_bin) and then calls find_medoid_with_sampling, which performs another full pass over the dataset via VectorDataIterator (plus a sampled centroid pass). On large datasets this doubles I/O and parsing work even though the vectors are already in memory.
    let data =
        read_bin::<Data::VectorDataType>(&mut builder.storage_provider.open_reader(&data_path)?)?;
    let context = PiPNNBuildContext::new(
        config,
        &builder.index_configuration.config,

diskann-providers/src/storage/bin.rs:388

  • save_adjacency_graph is re-exported as a public API, but its signature exposes diskann::graph::AdjacencyList<u32>, forcing downstream callers to take a dependency on diskann graph types just to use diskann-providers serialization. If the goal is a narrow cross-crate adjacency entry point, consider changing this API to accept adjacency rows as plain slices (or any type that derefs/borrows to &[u32]) to avoid leaking diskann internals into diskann-providers’ public surface.
pub fn save_adjacency_graph<P>(
    adjacency: &[AdjacencyList<u32>],
    max_degree: u32,
    provider: &P,
    start_point: u32,

Copilot AI review requested due to automatic review settings July 30, 2026 08:26
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/04-integration branch from f836f69 to 2181da2 Compare 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 12 out of 13 changed files in this pull request and generated no new comments.

@codecov-commenter

Codecov Comments Bot (codecov-commenter) commented Jul 30, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.60479% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.16%. Comparing base (dd896a3) to head (fec1146).

Files with missing lines Patch % Lines
diskann-disk/src/build/builder/pipnn.rs 97.35% 4 Missing ⚠️
...src/build/configuration/disk_index_build_config.rs 93.84% 4 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1291      +/-   ##
==========================================
+ Coverage   92.14%   92.16%   +0.01%     
==========================================
  Files         621      623       +2     
  Lines      130986   131240     +254     
==========================================
+ Hits       120694   120953     +259     
+ Misses      10292    10287       -5     
Flag Coverage Δ
miri 92.16% <97.60%> (+0.01%) ⬆️
unittests 92.11% <97.60%> (+0.01%) ⬆️

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

Files with missing lines Coverage Δ
diskann-benchmark/src/disk_index/build.rs 93.33% <100.00%> (+1.12%) ⬆️
diskann-disk/src/build/builder/backend.rs 95.34% <ø> (ø)
diskann-disk/src/build/builder/build.rs 94.92% <100.00%> (+2.47%) ⬆️
diskann-disk/src/build/builder/core.rs 97.71% <100.00%> (+0.14%) ⬆️
.../build/configuration/disk_index_build_parameter.rs 97.27% <ø> (ø)
...n-disk/src/search/provider/disk_vertex_provider.rs 86.50% <100.00%> (-0.11%) ⬇️
diskann-providers/src/common/minmax_repr.rs 100.00% <100.00%> (ø)
diskann-disk/src/build/builder/pipnn.rs 97.35% <97.35%> (ø)
...src/build/configuration/disk_index_build_config.rs 93.84% <93.84%> (ø)

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot AI review requested due to automatic review settings July 30, 2026 08:55
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/04-integration branch from 2181da2 to 19c3b5a Compare July 30, 2026 08:55
weiyaoluo (SeliMeli) and others added 18 commits November 2, 2026 11:03
# Conflicts:
#	diskann/src/graph/pipnn/mod.rs
Read in this order: mod.rs (config, build_graph), kernels.rs,
conversion.rs, then partitioning.rs, leaf_build.rs, finalization.rs, then
the kernel files, then rowmajor.rs and Cargo.toml.

Port to the rowmajor views API. main moved diskann-utils views to
views::rowmajor: MatrixView is rowmajor::Ref, MutMatrixView is
rowmajor::Mut, the owned Matrix is rowmajor::Owned, and the row accessors
live on the Matrix and MatrixMut traits.

Configuration. PiPNNConfig gets private fields and a Builder that checks
every invariant once, so the build steps rely on the configuration
instead of a validate call higher in the stack. The builder owns the four
prune parameters that PiPNN reads from graph::Config (degree, alpha,
max_occlusion_size as NonZeroU16, prune_kind); Builder::from_graph copies
them. The metric moves into the configuration, and the f32 row length
comes from VectorRepr::full_dimension of the first row.

Thread pool. PiPNNBuildContext and its &rayon::ThreadPool go away. The
build runs as Rayon jobs in the pool of the caller: an enclosing
ThreadPool::install, or else the global pool.

Type erasure. build_graph takes the dataset as a Dataset trait object:
rows gathered as f32 for the partition and leaf steps, and the distance
between two rows for final pruning. Rows<T> implements it for a matrix of
one Element type; Element selects each type's row conversion, and f16
overrides it with the run-time FP16 dispatch that the TypeId branch used
to select. kernels::select runs the diskann-wide dispatch once per build
and boxes the leaf and partition kernels of the metric behind the Kernels
and LeaderSet objects; a call covers one stripe or leaf. No build step
carries the architecture, metric or element type as a generic.

Scratch. rowmajor::Buffer is a Vec-backed Matrix/MatrixMut whose shape
changes in place and whose allocation only grows. The kernel workspaces
and the leaf and stripe job buffers use it instead of a Vec plus a
hand-cut prefix view, and the distance_scratch helper goes away.

Logging. The three info_span! scopes become one timed tracked_debug!
event per build step, and the max_depth warning uses tracked_warn!
without a structured field, which the no-tracing fallback cannot format.
The pipnn feature no longer forces tracing.
Group the build parameters by the step that reads them. `Config` now holds
a private `PartitionConfig`, the leaf `k`, and a private `PruneConfig`;
`partition` and `prune_overfull` take their own group. The public
`Builder` and the `Config` getters stay flat.

The sampling seed is a plain `u64` with a fixed default. A caller that
wants fresh leaves on each build passes a random value.

`partition` owns the level loop: it reads the fanout of each level, splits
the pending clusters, and cuts the clusters that are still pending after
`max_depth` levels. `Cluster` no longer carries a level and `ClusterSplit`
no longer counts cut points. `split_cluster` only splits.

Tests: the medoid helper uses `ComputeMedoid`; one helper builds, indexes
and asserts a SIFT case; the vector-type test drops the f32 arm that the
metric test already runs.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
`Builder` keeps the seed as `Option<u64>` and `build` resolves it once:
the given seed, or a random one. `Config::seed` returns the resolved
value, so a caller can log it to repeat a build. A default build no
longer samples the same leaves every time.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The library build no longer uses `VectorRepr` in `pipnn/mod.rs`, so the
crate-level import failed `-D warnings` in the clippy-features and docs
jobs. Only the construction tests call `f32::distance`.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The SDE jobs run the `diskann` tests with all features, which now read the
SIFT fixtures in `test_data/sift` through LFS, as the other test jobs do.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
`add_leaf_candidates` computes the effective `k` and returns early for a
leaf without pairs; `LeafBuffers::prepare` only shapes storage, and now
also clears the edge lists. The leaf loop consumes the leaves instead of
borrowing them.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The two end-to-end PiPNN tests read the SIFT fixtures from LFS, which the
SDE jobs do not check out. The PiPNN kernel tests still run under SDE and
cover the SIMD paths.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The core PR replaced `PiPNNBuildContext` and `PiPNNConfig` with
`pipnn::Builder`, a `Dataset` trait object, and the caller's Rayon pool.

- `PiPNNParameters::builder(&graph)` starts a `pipnn::Builder` from the
  JSON parameters and the graph prune policy; the disk builder adds the
  seed and `build()` validates, so `validate()` is gone.
- `build_adjacency` owns the dataset: it reads the rows, builds the graph
  under `pool.install`, and returns the adjacency lists, so the dataset is
  dropped before the start point and the graph file are written.
- `BuildVector` names the element bound of the disk build: `VectorRepr`,
  or the PiPNN `Element` gather with the `pipnn` feature. `MinMaxElement`
  implements `Element` under a new `pipnn` feature of diskann-providers.
- `RayonThreadPoolRef::as_rayon` is no longer needed.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The benchmark starts its PiPNN configuration from the same JSON
parameters and needs the same element bound as the disk build.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
- `rowmajor::Buffer` becomes `rowmajor::Resizable`, a general resizable
  matrix: `resize` (needs only `T: Default`), `resize_with` for any `T`,
  `with_capacity`, `capacity`, `clear`, `shrink_to_fit`, conversions to
  and from `Owned`, and `PartialEq`. The values after a resize are
  unspecified, so the implementation may reuse storage freely.
- `Element` no longer extends `VectorRepr`. It asks only for what PiPNN
  reads: the converted row length, a row conversion, and a distance.
  `f32`, `u8` and `i8` share one macro and never fail; a compressed type
  such as MinMax returns its decode errors from `convert`. `f16` keeps
  its batched runtime-dispatched gather.
- The build steps read rows through a private `inner::Dataset` that an
  opaque `Backend` holds, following `into_backend` in diskann-inmem. The
  inner trait hands out the kernels and the prune kind instead of a
  `Metric`, so `Metric` stays at the public boundary. `Dataset` gains
  `Debug` and a provided `backend()`. `build_graph` takes any
  `IntoBackend`: a `&Dataset` or a `Backend`.
- Public docs describe the contract, including the documented panics for
  out-of-range IDs and wrong output shapes; implementation notes moved to
  comments. The lane-count check in `topk` is now a compile-time assert.
- The SDE jobs check out the LFS fixtures instead of skipping the SIFT
  tests.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`Element` no longer extends `VectorRepr`, so:

- `MinMaxElement` implements `Element` directly: `dimensions` reads the
  row metadata, `convert` decodes with its `as_f32_into` and returns the
  decode error, and distances use its code-domain `FnPtr`.
- `BuildVector` and `build_adjacency` require `VectorRepr` and `Element`
  explicitly, since the Vamana path still needs `VectorRepr`.

`build_graph` takes `&rows` through `IntoBackend`, so the adapter call is
unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- PiPNNParameters fields are optional, so the pipnn::Builder defaults are
  defined once. `k` is renamed to `leaf_k` to match the builder.
- DiskIndexBuildParameters selects the algorithm with
  `with_build_algorithm`, like the other optional parameters.
  `new_pipnn` is removed.
- DiskIndexBuilder::new parses the PiPNN configuration once, before the
  quantizer is trained, and rejects a build quantization other than FP.
- The PiPNN adapter takes explicit arguments like the Vamana
  `build_inmem_index`, reuses `start_point_from_medoid`, and reports a
  dataset shape mismatch with `diskann_error!`.
- save_adjacency_graph documents its input contract instead of checking
  each row.
- The benchmark disk build is bounded on BuildVector, which fixes the
  `--all-features` clippy job.
- The unused `virtual_storage` feature of diskann-disk is removed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
weiyaoluo (SeliMeli) and others added 11 commits November 11, 2026 07:32
- DiskIndexBuilder::build_with_pipnn takes a pipnn::Config. Only this entry
  point requires pipnn::Element, so Vamana builds and their callers keep
  their bounds, and the `pipnn` feature stays additive. The BuildVector
  alias, BuildAlgorithm and PiPNNParameters are removed.
- build() and build_with_pipnn() share the PQ and disk layout steps.
- The PiPNN adapter moves to builder/pipnn.rs and writes the in-memory
  graph file itself, as the merged Vamana build does. The providers
  storage module is unchanged.
- The benchmark is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
BREAKING CHANGE: DiskIndexBuilder::new takes DiskIndexBuildConfig instead of separate disk and index parameters. Select PiPNN in config and call build().
BREAKING CHANGE: DiskIndexBuildConfig holds shared inputs. GraphBuildConfig contains only algorithm-specific parameters.
BREAKING CHANGE: with pipnn enabled, BuildElement requires Element. Keep MinMax row support and remove the duplicate InputRows adapter.
BREAKING CHANGE: new now stores config and returns Self. Graph preparation errors are returned by build, before output generation.

This branch has not been deployed

No deployments
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.

Integrate PiPNN with disk index builds

4 participants