PiPNN 4/6: integrate disk build pipeline

[weiyaoluo (SeliMeli)]

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]

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]

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]

Pull request overview

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

[Copilot]

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]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated no new comments.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 91.22807% with 10 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (pipnn-stack/03-core@61a03bc). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...nn-disk/src/build/configuration/build_algorithm.rs	66.66%	4 Missing ⚠️
.../build/configuration/disk_index_build_parameter.rs	50.00%	3 Missing ⚠️
diskann-providers/src/utils/rayon_util.rs	0.00%	3 Missing ⚠️

Additional details and impacted files

@@                  Coverage Diff                   @@
##             pipnn-stack/03-core    #1291   +/-   ##
======================================================
  Coverage                       ?   91.32%           
======================================================
  Files                          ?      518           
  Lines                          ?    98701           
  Branches                       ?        0           
======================================================
  Hits                           ?    90138           
  Misses                         ?     8563           
  Partials                       ?        0           

Flag	Coverage Δ
miri	91.32% <91.22%> (?)
unittests	91.00% <91.22%> (?)

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

Files with missing lines	Coverage Δ
diskann-disk/src/build/builder/build.rs	92.82% <100.00%> (ø)
diskann-providers/src/storage/bin.rs	96.72% <100.00%> (ø)
.../build/configuration/disk_index_build_parameter.rs	94.73% <50.00%> (ø)
diskann-providers/src/utils/rayon_util.rs	96.31% <0.00%> (ø)
...nn-disk/src/build/configuration/build_algorithm.rs	66.66% <66.66%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (2)

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

• save_adjacency_graph writes a header max_degree that is later trusted by load_graph to size its read buffer; if any adjacency row is longer than max_degree, load_graph will panic when slicing the buffer. Since this is now a public entry point, it should defensively validate that the observed max row length does not exceed the configured max_degree (and return an error if it does) before delegating to save_graph.

    save_graph(
        &AdjacencyGraph {
            adjacency,
            max_degree,
        },

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

• DiskIndexBuildParameters used to be Copy, but the new build_algorithm: BuildAlgorithm field (and the PiPNN variant’s Vec) forces dropping Copy for all builds. Since pipnn is an optional feature and the default build still only has BuildAlgorithm::Vamana, consider preserving Copy for the common non-pipnn build via cfg_attr(not(feature = "pipnn"), derive(Copy)) (and similarly making BuildAlgorithm Copy when pipnn is disabled). This reduces downstream breakage for consumers not enabling PiPNN.

/// Parameters specific for disk index construction.
#[derive(Clone, PartialEq, Debug)]
pub struct DiskIndexBuildParameters {
    /// Memory budget for disk-index pipeline stages that support bounded work.
    /// Explicit one-shot PiPNN selection is never silently replaced.

[Copilot]

Pull request overview

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

Suppressed comments (2)

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

• save_adjacency_graph() forwards a caller-supplied max_degree into the file header without validating it against the actual adjacency row lengths. If any row is longer than max_degree, load_graph() will panic because it allocates a buffer of size max_degree and then slices it to num_neighbors. Consider validating and returning an ANNError when observed_max_degree > max_degree.

    save_graph(
        &AdjacencyGraph {
            adjacency,
            max_degree,
        },

diskann-disk/src/build/builder/build/pipnn/tests.rs:154

• This test currently hard-codes the serialized graph header max_degree to the pruned degree (32). If the adapter writes the canonical max_degree (Config::max_degree_u32) like the Vamana pipeline, this assertion will fail even though the header is correct.

    assert_eq!(u32::from_le_bytes(header[8..12].try_into().unwrap()), 32);

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (2)

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

• save_adjacency_graph always writes the caller-provided max_degree into the graph header, but it does not verify that any adjacency row actually fits within that bound. If a row exceeds max_degree, the file will contain per-row degrees larger than the header’s max degree, which can corrupt downstream assumptions (and is hard to debug).

Consider validating the observed maximum degree in adjacency and returning an error if it exceeds max_degree before calling save_graph.

/// Save real-point adjacency lists in the canonical graph layout.
pub fn save_adjacency_graph<P>(
    adjacency: &[AdjacencyList<u32>],
    max_degree: u32,
    provider: &P,

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

• This adapter reads the full dataset into memory via read_bin(...), and then later calls find_medoid_with_sampling(...), which re-reads metadata and scans the dataset again to compute the start node. For large datasets and non-local storage providers this doubles I/O and can dominate build time.

Consider deriving the medoid/start ID from the already-loaded data (or adding a find_medoid_* helper that accepts a MatrixView/slice) so the PiPNN path only performs one full read.

    // PiPNN is a batch algorithm: materialize the matrix once, while all
    // partition and leaf scratch stays inside the supplied pool and is released
    // before the outer disk pipeline continues.
    let data =
        read_bin::<Data::VectorDataType>(&mut builder.storage_provider.open_reader(&data_path)?)?;

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann-disk/src/build/configuration/build_algorithm.rs:66

• BuildAlgorithm is deserialized from JSON with an internally-tagged enum, but it does not deny unknown fields. That means inputs like { "algorithm": "Vamana", "c_max": 512 } can be accepted while silently ignoring the extra keys, which makes misconfiguration/typos hard to detect for this user-facing config surface. Consider rejecting unknown fields at the enum level as well (the PiPNNParameters struct already does this).

/// Selects the graph construction algorithm for index building.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(tag = "algorithm")]
#[non_exhaustive]
pub enum BuildAlgorithm {

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (2)

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

• save_adjacency_graph forces the header max_degree to the caller-provided max_degree, but it does not validate that any adjacency row actually fits within that bound. If a caller accidentally passes a smaller max_degree than the data contains, the file will serialize successfully but later consumers will see an inconsistent header (and may allocate based on it). Consider rejecting graphs whose observed degree exceeds max_degree before calling save_graph.

pub fn save_adjacency_graph<P>(
    adjacency: &[AdjacencyList<u32>],
    max_degree: u32,
    provider: &P,
    start_point: u32,

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

• build_graph reads the full dataset into memory via read_bin, but then calls find_medoid_with_sampling which re-opens and re-reads the dataset file to compute the medoid/start point. For large datasets this doubles IO work in the PiPNN path. If possible, derive the start point from the already-loaded data (or add an in-memory medoid helper alongside find_medoid_with_sampling and use it here).

    let mut rng = diskann_providers::utils::create_rnd_from_optional_seed(
        builder.index_configuration.random_seed,
    );
    let (_, start_id) = find_medoid_with_sampling::<Data::VectorDataType, _>(
        &data_path,

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (4)

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

• save_adjacency_graph passes a configured max_degree into save_graph, but AdjacencyGraph::get_adjacency_list does not enforce that each row length is <= max_degree. If any row exceeds the configured degree, the file header will claim a smaller degree than the data contains, and downstream readers that size buffers from the header can misbehave.

Consider validating per-row degree here and returning an error that includes the offending row index/length.

    fn get_adjacency_list(&self, index: usize) -> ANNResult<Self::Item<'_>> {
        self.adjacency
            .get(index)
            .map(|row| &**row)
            .ok_or_else(|| ANNError::log_index_error(format_args!("missing graph row {index}")))

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

• After diskann::graph::pipnn::build_graph returns, the adapter writes the adjacency to the canonical graph file but never checks that the returned row count matches the dataset point count validated from metadata. If the core ever returns a mismatched length (or if a future refactor changes its contract), the disk layout writer will later consume points rows from the file and hit EOF / corrupt the index.

Add a defensive check that adjacency.len() == points before serializing.

    let adjacency = diskann::graph::pipnn::build_graph(data.as_view(), &context)?;

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

• PiPNN builds already load the entire dataset into memory (read_bin), but start node selection re-reads the dataset from storage via find_medoid_with_sampling (which itself does at least one full scan after centroid computation). For large datasets this adds substantial extra I/O and wall time compared to computing the sampled medoid from the already-loaded MatrixView.

Consider adding an in-memory medoid helper (same sampling semantics) and using it here to avoid re-reading the dataset file.

    let mut rng = diskann_providers::utils::create_rnd_from_optional_seed(
        builder.index_configuration.random_seed,
    );
    let (_, start_id) = find_medoid_with_sampling::<Data::VectorDataType, _>(
        &data_path,
        builder.storage_provider,
        MAX_MEDOID_SAMPLE_SIZE,
        &mut rng,
    )?;

diskann-disk/src/build/builder/build.rs:186

• When BuildAlgorithm::PiPNN is selected, build_graph returns early and never uses self.build_quantizer, but DiskIndexBuilder::new still trains it unconditionally. Training the build quantizer can be expensive and becomes pure overhead for PiPNN builds.

Consider deferring BuildQuantizer::train until after algorithm selection, or making it conditional on BuildAlgorithm::Vamana (e.g., store Option<BuildQuantizer> and initialize it only in the Vamana branches).

    async fn build_graph(&mut self, pool: RayonThreadPoolRef<'_>) -> ANNResult<()> {
        #[cfg(feature = "pipnn")]
        if let Some(config) = self.disk_build_param.pipnn_config() {
            return pipnn::build_graph(self, pool, config);
        }

[juchen-ms (partychen)]

The algorithm itself looks sound. My concern is mostly that the current code structure does not map cleanly to the algorithm’s stages.

CandidateMerge currently determines not only candidate aggregation, but also candidate retention and final neighbor selection. Because the branch happens before leaf construction, both paths separately wrap the same validation, parallel traversal, and compute_leaf() logic. It also makes final_prune meaningful only for the HashPrune path.

Would it make sense to separate these responsibilities into explicit stages?

Diagram

This would make the stable pipeline directly visible:
LeafBuilder → CandidateAggregation → NeighborSelection → FinalAdjacency

Leaf generation would remain shared, while aggregation and neighbor selection could vary independently. This should preserve the existing algorithms while reducing duplicated traversal logic, clarifying the role of final_prune, and making future strategies easier to add.