PiPNN 4/6: integrate disk build pipeline - #1291
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.
| #[cfg(feature = "pipnn")] | ||
| let frozen_points = match self.build_algorithm { | ||
| diskann_disk::BuildAlgorithm::PiPNN(_) => NonZero::new(1).unwrap(), | ||
| _ => NonZero::new(self.start_point_strategy.count()).unwrap(), | ||
| }; |
| use std::num::NonZeroUsize; | ||
|
|
||
| use diskann::ANNError; | ||
| #[cfg(feature = "pipnn")] | ||
| use diskann::ANNResult; |
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
| /// Save real-point adjacency lists in the canonical graph layout. | ||
| pub fn save_adjacency_graph<P>( | ||
| adjacency: &[AdjacencyList<u32>], | ||
| max_degree: u32, | ||
| provider: &P, | ||
| start_point: u32, | ||
| path: &str, | ||
| ) -> ANNResult<usize> | ||
| where | ||
| P: StorageWriteProvider, | ||
| { | ||
| save_graph( | ||
| &AdjacencyGraph { | ||
| adjacency, | ||
| max_degree, | ||
| }, | ||
| provider, | ||
| start_point, | ||
| path, | ||
| ) | ||
| } |
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 ❌ Your patch status has failed because the patch coverage (23.21%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## pipnn-stack/03-core #1291 +/- ##
======================================================
Coverage ? 90.72%
======================================================
Files ? 521
Lines ? 101084
Branches ? 0
======================================================
Hits ? 91709
Misses ? 9375
Partials ? 0
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
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.
Suppressed comments (2)
diskann-providers/src/storage/bin.rs:388
save_adjacency_graphcan write a graph file thatload_graphcannot read if any row has more thanmax_degreeneighbors:load_graphallocates a buffer of lengthmax_degreeand slices it tonum_neighbors, which will panic/out-of-bounds whennum_neighbors > max_degree. The wrapper should ensure the header max degree is at least the observed maximum (or fail early).
pub fn save_adjacency_graph<P>(
adjacency: &[AdjacencyList<u32>],
max_degree: u32,
provider: &P,
start_point: u32,
diskann-providers/src/utils/rayon_util.rs:81
as_rayonis documented as a borrow, but it consumesself. WhileRayonThreadPoolRefisCopy, taking&selfbetter matches the API intent and avoids forcing a move in call sites that hold the wrapper in a non-Copycontext later.
/// Borrow the underlying pool for APIs that retain a caller-owned pool.
pub fn as_rayon(self) -> &'a rayon::ThreadPool {
self.0
}
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.
Suppressed comments (2)
diskann-providers/src/storage/bin.rs:398
save_adjacency_graphis a new public entry point but it doesn’t validate thatstart_pointis withinadjacency.len()or that any row’s degree is <=max_degree. If callers pass an out-of-range start point or an overfull row, this will write a graph whose header/contents are inconsistent and can later fail (or mis-size allocations) when loading.
save_graph(
&AdjacencyGraph {
adjacency,
max_degree,
},
diskann-disk/src/build/configuration/disk_index_build_parameter.rs:113
DiskIndexBuildParameterspreviously implementedCopy; addingbuild_algorithm: BuildAlgorithm(which can contain heap data likeVec) dropsCopy. This is an API-breaking change for downstream users that may have relied on implicit copies (e.g., passing params by value multiple times). If this is intended, it likely needs an explicit release note / versioning consideration; otherwise consider keeping the algorithm config outside this public params struct (or exposing a separate PiPNN params type) to preserve the priorCopyAPI surface.
/// 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.
7a5bbeb to
ebac06f
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.
Suppressed comments (2)
diskann-providers/src/storage/bin.rs:398
save_adjacency_graphwrites a headermax_degreethat is later trusted byload_graphto size its read buffer; if any adjacency row is longer thanmax_degree,load_graphwill 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 configuredmax_degree(and return an error if it does) before delegating tosave_graph.
save_graph(
&AdjacencyGraph {
adjacency,
max_degree,
},
diskann-disk/src/build/configuration/disk_index_build_parameter.rs:113
DiskIndexBuildParametersused to beCopy, but the newbuild_algorithm: BuildAlgorithmfield (and thePiPNNvariant’sVec) forces droppingCopyfor all builds. Sincepipnnis an optional feature and the default build still only hasBuildAlgorithm::Vamana, consider preservingCopyfor the common non-pipnnbuild viacfg_attr(not(feature = "pipnn"), derive(Copy))(and similarly makingBuildAlgorithmCopywhenpipnnis 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.
ebac06f to
2e0532f
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 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);
| )?; | ||
| save_adjacency_graph( | ||
| &adjacency, | ||
| u32_try_from(builder.index_configuration.config.pruned_degree().get())?, |
2e0532f to
449f802
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.
Suppressed comments (2)
diskann-providers/src/storage/bin.rs:387
save_adjacency_graphalways writes the caller-providedmax_degreeinto the graph header, but it does not verify that any adjacency row actually fits within that bound. If a row exceedsmax_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 callsfind_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)?)?;
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.
Suppressed comments (1)
diskann-disk/src/build/configuration/build_algorithm.rs:66
BuildAlgorithmis 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 (thePiPNNParametersstruct 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 {
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.
Suppressed comments (2)
diskann-providers/src/storage/bin.rs:388
save_adjacency_graphforces the header max_degree to the caller-providedmax_degree, but it does not validate that any adjacency row actually fits within that bound. If a caller accidentally passes a smallermax_degreethan 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 exceedsmax_degreebefore callingsave_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_graphreads the full dataset into memory viaread_bin, but then callsfind_medoid_with_samplingwhich 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-loadeddata(or add an in-memory medoid helper alongsidefind_medoid_with_samplingand 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,
Forward the disk feature to diskann/pipnn instead of depending on a separate crate. Consolidate duplicate pipeline builds and name adapter tests by the behavior they protect.
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.
Suppressed comments (4)
diskann-providers/src/storage/bin.rs:421
save_adjacency_graphpasses a configuredmax_degreeintosave_graph, butAdjacencyGraph::get_adjacency_listdoes 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_graphreturns, 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 consumepointsrows 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 viafind_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-loadedMatrixView.
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::PiPNNis selected,build_graphreturns early and never usesself.build_quantizer, butDiskIndexBuilder::newstill 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);
}
The PiPNN core returns adjacency lists, while a production DiskANN disk index also needs algorithm selection, PQ artifacts, graph headers, disk layout, start metadata, and serialization. This PR connects
diskann::graph::pipnnto the existing disk-index pipeline. It does not introduce a PiPNN-specific index format.Concepts
BuildAlgorithmdistinguishes an explicit PiPNN request from the existing automatic memory-budget path. Automatic selection may choose Vamana according to established policy. Explicit PiPNN is a user contract: it must run PiPNN or return a validation/build error.The core addresses real points only: matrix row
iproduces adjacency row/IDi. Start/frozen points, PQ, and persisted layout remain outer-pipeline responsibilities.Code map
diskann-disk/src/build/configuration/build_algorithm.rsadds JSON-facingPiPNNParametersand conversions intodiskann::graph::pipnn::PiPNNConfig.disk_index_build_parameter.rsandconfiguration/mod.rscarry explicit algorithm choice through public disk parameters.builder/build.rsvalidates explicit PiPNN before expensive work and dispatches the resolved algorithm. Only automatic strategy selection may choose Vamana.builder/build/pipnn.rsvalidates dataset metadata, loads the dense matrix, createsPiPNNBuildContext, callsbuild_graph, computes the existing sampled medoid, and writes canonical adjacency.diskann-providers/src/storage/bin.rswrites adjacency through the existing graph header/layout.diskann-diskfeaturepipnnforwards todiskann/pipnn; no implementation-crate dependency exists.End-to-end flow
Deserialize disk parameters → preserve/resolve
BuildAlgorithm→ common builder prepares dataset and graph policy → PiPNN adapter validates metadata and constructs core context → core returns adjacency rowifor real pointi→ canonical graph writer records degree/start/header/rows → common pipeline continues PQ generation, disk layout, and final serialization.Invariants and boundaries
BuildAlgorithm::PiPNNis never rewritten by memory estimation.diskann::graph::pipnndoes not depend on disk/provider crates.Review path
build_algorithm.rs.builder/build.rs.builder/build/pipnn.rsand the in-crate core call.Test architecture
Adapter tests are named for externally visible behavior:
Duplicate full builds and constructor-field tautologies were merged into these behavioral cases.
Validation
Targeted disk adapter/configuration tests pass with
diskann-disk/pipnn; Clippy compiles all targets. AArch64 and Windows cross-target checks pass. Existing canonical writer/PQ tests continue to cover shared storage policy.Stack relation
Stack 4/6. Depends on #1290 provider-independent core. #1294 invokes this production path from the disk benchmark; #1295 threads optional HashPrune configuration through the same boundary.
Stack 4/6: #1290 → #1294