Repository navigation
PiPNN 4/6: integrate disk build pipeline - #1291
weiyaoluo (SeliMeli) wants to merge 359 commits into
Conversation
There was a problem hiding this comment.
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
BuildAlgorithmselector (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.
8134eb3 to
ce5f054
Compare
There was a problem hiding this comment.
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_memoryusessize_of::<...>()unqualified, but this module doesn’t importstd::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
DiskIndexBuildParametersis publicly re-exported (viadiskann-disk’s public API) and this change removesCopyfrom 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 {
ce5f054 to
3893460
Compare
3893460 to
f836f69
Compare
There was a problem hiding this comment.
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.rsenables#![warn(missing_docs)], and the newpub(crate) fn pipnn_confighas 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_graphreads the full dataset into memory (read_bin) and then callsfind_medoid_with_sampling, which performs another full pass over the dataset viaVectorDataIterator(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_graphis re-exported as a public API, but its signature exposesdiskann::graph::AdjacencyList<u32>, forcing downstream callers to take a dependency ondiskanngraph types just to usediskann-providersserialization. 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 leakingdiskanninternals intodiskann-providers’ public surface.
pub fn save_adjacency_graph<P>(
adjacency: &[AdjacencyList<u32>],
max_degree: u32,
provider: &P,
start_point: u32,
f836f69 to
2181da2
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ 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
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
2181da2 to
19c3b5a
Compare
…date/04-integration
# 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>
This reverts commit 535a213.
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>
- 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 reverts commit 18dcbca.
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
BuildAlgorithmaccepts explicit PiPNN parameters.diskann::graph::pipnn::build_graph.The PiPNN core returns rows for real points only. The disk pipeline still owns frozen points and persisted metadata.
Review order
build/configuration/build_algorithm.rs.builder/build.rs.builder/build/pipnn.rs.storage/bin.rsgraph validation and serialization.Validation
Stack
Stack 4/6. Depends on #1290. #1294 adds benchmark routes.