Converge PQ pivot loading

[xinyuwen2]

• Does this PR have a descriptive title that could go in our release notes?
• Does this PR add any new dependencies?
• Does this PR modify any existing APIs?
• Is the change to the API backwards compatible?
• Should this result in any changes to our documentation, either updating existing docs or adding new ones?

Reference Issues/PRs

Fixes #1015.

What does this implement/fix? Briefly explain your changes.

This PR consolidates PQ pivot loading around PQStorage::load_pivots, which uses PQStorage::pivot_data_path as the single source of truth and returns a fully validated BasicTable.

• Route existing PQ pivot loading through a shared parser for offsets, pivots, centroid data, and chunk offsets.
• Fold legacy centroid data into the pivot matrix before constructing BasicTable, so the loaded table is fully formed at construction time.
• Use BasicTable / ChunkOffsetsBase to validate structural table invariants in the quantization layer.
• Remove the raw pivot-data loading API and the load_pq_pivots_bin compatibility wrapper.
• Remove the explicit pivot path parameter from pivot loading APIs and remove get_pivot_data_path.
• Add TryFrom<BasicTable> for FixedChunkPQTable so callers validate fixed-table center-count bounds when converting loaded tables.
• Preserve the previous fixed PQ table loading behavior that accepts up to NUM_PQ_CENTROIDS centers while rejecting tables with more centers.
• Move configuration-specific compatibility checks, such as expected chunk count and dimensionality, to the call sites that own those expectations.
• Update PQ construction, PQ generation, disk index loading, and memory quant vector providers to use load_pivots.

Any other comments?

Validated with:

• cargo test -p diskann-providers pq_storage
• cargo test -p diskann-providers pq_construction
• cargo test -p diskann-providers fixed_chunk_pq_table_test
• cargo test -p diskann-providers fast_memory_quant_vector_provider::tests
• cargo test -p diskann-providers memory_quant_vector_provider::tests

[Copilot]

Pull request overview

This PR consolidates the two PQ pivot loading paths in diskann-providers to share a single parser for the pq_pivots.bin layout, reducing duplicated file parsing/validation logic and updating tests to exercise the production loader.

Changes:

• Introduces an internal load_pivot_file_parts helper (via PivotFileParts) to parse offsets, pivots, centroid, and chunk offsets once.
• Refactors load_existing_pivot_data and load_pq_pivots_bin to use the shared helper, preserving centroid-folding behavior only for the load_pq_pivots_bin path.
• Removes the duplicate test-only pivot loader and updates tests to call PQStorage::load_pq_pivots_bin directly.

Reviewed changes

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

File	Description
diskann-providers/src/storage/pq_storage.rs	Adds shared pivot-file parsing helper and routes existing loaders through it.
diskann-providers/src/model/pq/fixed_chunk_pq_table.rs	Updates tests to use the production PQStorage::load_pq_pivots_bin instead of a test-only parser.

Comments suppressed due to low confidence (3)

diskann-providers/src/storage/pq_storage.rs:301

• The offset-table validation checks only nrows() == 4, but the documented layout is 4x1. A malformed offsets matrix with nrows == 4 and ncols != 1 would currently pass validation and likely produce confusing downstream failures.

This issue also appears on line 377 of the same file.

        let offsets = read_bin_from::<u64>(&mut reader, 0)?;
        if offsets.nrows() != 4 {
            return Err(ANNError::log_pq_error(format_args!(
                "Error reading pq_pivots file {}. Offsets don't contain correct metadata, \
                 # offsets = {}, but expecting 4.",

diskann-providers/src/storage/pq_storage.rs:356

• When expected_num_pq_chunks is None (inference mode), the current chunk-offsets error message still prints an "expected nr=0" value and mentions "pass 0" even though the helper now takes an Option. Splitting the validation into explicit branches will keep the error message accurate for both inferred and fixed chunk counts.

        let chunk_offsets_m = read_bin_from::<u32>(&mut reader, file_offset_data[(2, 0)])?;
        if expected_num_pq_chunks
            .is_some_and(|num_pq_chunks| chunk_offsets_m.nrows() != num_pq_chunks + 1)
            || chunk_offsets_m.ncols() != 1
        {

diskann-providers/src/storage/pq_storage.rs:380

• This newly added chunk-offsets validation error omits the file path, which makes debugging harder when multiple PQ pivot files are involved. Including pq_pivots in the message would align with the other errors in this function.

            return Err(ANNError::log_pq_error(format_args!(
                "Error reading pq_pivots file at chunk offsets; chunk offsets must start at 0, end at dim {}, and contain at least two entries.",
                parts.dim()
            )));

---

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

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

diskann-providers/src/storage/pq_storage.rs:365

• The chunk-offsets validation combines “wrong row count” and “wrong column count” into one error message, and when expected_num_pq_chunks is None (infer case) the message prints expecting nr=0, which is misleading (the real requirement in that branch is only nc=1). Consider splitting the checks so the error message accurately reflects what was expected in each case.

        let chunk_offsets_m = read_bin_from::<u32>(&mut reader, file_offset_data[(2, 0)])?;
        if expected_num_pq_chunks
            .is_some_and(|num_pq_chunks| chunk_offsets_m.nrows() != num_pq_chunks + 1)
            || chunk_offsets_m.ncols() != 1
        {
            return Err(ANNError::log_pq_error(format_args!(
                "Error reading pq_pivots file at chunk offsets; file has nr={}, nc={} \
                 but expecting nr={} and nc=1. The expected num_pq_chunks should be \
                 passed as 0 if we want to infer.",
                chunk_offsets_m.nrows(),
                chunk_offsets_m.ncols(),
                expected_num_pq_chunks.map_or(0, |num_pq_chunks| num_pq_chunks + 1)
            )));
        }

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 90.00000% with 25 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.51%. Comparing base (c1f8c7d) to head (ce57c16).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-providers/src/model/pq/pq_construction.rs	68.96%	9 Missing ⚠️
diskann-providers/src/storage/pq_storage.rs	95.94%	6 Missing ⚠️
...ovider/async_/fast_memory_quant_vector_provider.rs	44.44%	5 Missing ⚠️
...ph/provider/async_/memory_quant_vector_provider.rs	44.44%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1276      +/-   ##
==========================================
- Coverage   91.66%   91.51%   -0.15%     
==========================================
  Files         514      516       +2     
  Lines       98732    98254     -478     
==========================================
- Hits        90505    89920     -585     
- Misses       8227     8334     +107     

Flag	Coverage Δ
miri	91.51% <90.00%> (-0.15%)	⬇️
unittests	91.19% <90.00%> (-0.45%)	⬇️

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

Files with missing lines	Coverage Δ
diskann-disk/src/storage/disk_index_reader.rs	93.33% <100.00%> (+1.26%)	⬆️
diskann-disk/src/storage/quant/pq/pq_generation.rs	92.59% <100.00%> (+0.77%)	⬆️
...aph/provider/async_/experimental/multi_pq_async.rs	97.53% <100.00%> (+0.33%)	⬆️
...ann-providers/src/model/pq/fixed_chunk_pq_table.rs	96.82% <100.00%> (+4.39%)	⬆️
...ovider/async_/fast_memory_quant_vector_provider.rs	97.84% <44.44%> (+0.36%)	⬆️
...ph/provider/async_/memory_quant_vector_provider.rs	97.43% <44.44%> (-0.83%)	⬇️
diskann-providers/src/storage/pq_storage.rs	96.51% <95.94%> (+10.97%)	⬆️
diskann-providers/src/model/pq/pq_construction.rs	93.97% <68.96%> (-1.71%)	⬇️

... and 250 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]

Pull request overview

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

Comments suppressed due to low confidence (2)

diskann-providers/src/storage/pq_storage.rs:380

• This chunk-offset invariant error omits the file path and the observed start/end values, which makes it harder to diagnose malformed pivot files in logs (especially when multiple files are loaded). Including pq_pivots and the observed bounds would make the message actionable.

        if parts.chunk_offsets.nrows() < 2
            || parts.chunk_offsets[(0, 0)] != 0
            || parts.chunk_offsets[(parts.chunk_offsets.nrows() - 1, 0)] != parts.dim()
        {
            return Err(ANNError::log_pq_error(format_args!(
                "Error reading pq_pivots file at chunk offsets; chunk offsets must start at 0, end at dim {}, and contain at least two entries.",
                parts.dim()
            )));

diskann-providers/src/storage/pq_storage.rs:198

• num_pq_chunks == 0 is now interpreted as "infer chunk count from file" (via then_some), which is a non-obvious semantic for a public parameter. Adding an inline comment here would help future maintainers and prevent accidental behavior changes.

        let parts = self.load_pivot_file_parts(
            &self.pivot_data_path,
            (*num_pq_chunks != 0).then_some(*num_pq_chunks),
            Some(*num_centers),

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

diskann-providers/src/storage/pq_storage.rs:364

• When inferring the number of PQ chunks (i.e., expected_num_pq_chunks == None), the error path formats the expected row count as 0 via map_or(0, ...), producing confusing messages like “expecting nr=0”. This makes the diagnostic misleading, especially when the error is actually due to nc != 1.

                 passed as 0 if we want to infer.",
                chunk_offsets_m.nrows(),
                chunk_offsets_m.ncols(),
                expected_num_pq_chunks.map_or(0, |num_pq_chunks| num_pq_chunks + 1)
            )));

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

diskann-providers/src/storage/pq_storage.rs:315

• load_pivot_file_parts validates that the offset table has 4 rows but never validates the expected 1-column layout (the comment above says offset table(4*1)). A malformed file with 4 rows and >1 column would currently be accepted and only the first column would be used via (0,0)/(1,0)/(2,0), which can mask corruption and lead to confusing downstream read errors.

        if offsets.nrows() != 4 {
            return Err(ANNError::log_pq_error(format_args!(
                "Error reading pq_pivots file {}. Offsets don't contain correct metadata, \
                 # offsets = {}, but expecting 4.",
                pq_pivots,

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (3)

diskann-providers/src/storage/pq_storage.rs:272

• load_pq_pivots_bin now rejects num_pq_chunks == 0, but the previous implementation explicitly documented 0 as the sentinel to infer chunk count. This is a backwards-incompatible behavioral change and conflicts with the PR metadata checkbox claiming API compatibility. Consider preserving the old behavior by delegating 0 to load_pq_pivots_bin_infer_chunks (and keeping the new explicit inference API as a clearer alternative).

        if num_pq_chunks == 0 {
            return Err(ANNError::log_pq_error(
                "num_pq_chunks must be non-zero; use load_pq_pivots_bin_infer_chunks to infer from the file.",
            ));
        }

diskann-providers/src/storage/pq_storage.rs:364

• This error message says "expecting {NUM_PQ_CENTROIDS} centers", but the condition is pivots.nrows() > NUM_PQ_CENTROIDS (i.e., the file is allowed to have fewer centers). The message should reflect that it's an upper bound to avoid confusing callers.

                "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.",
                pq_pivots,
                pivots.nrows(),
                NUM_PQ_CENTROIDS
            )));

diskann-providers/src/storage/pq_storage.rs:223

• load_existing_pivot_table clones the centroid matrix just to return it alongside the constructed table. This introduces an extra allocation/copy on every load. Consider restructuring PivotFileParts / pivot_file_parts_into_basic_table so the centroid can be moved out (or returned separately) without cloning.

        let centroid = parts.centroid.clone();
        let table = Self::pivot_file_parts_into_basic_table(&self.pivot_data_path, parts)?;
        Ok((table, centroid))

[Wei Wu (wuw92)]

Thanks for working through this convergence.
One point I am unsure about is the compatibility boundary. This PR adds or retains several pub APIs primarily to support existing in-workspace patterns, Although these are technically available to downstream users, they appear intended mainly for other DiskANN crates rather than as a stable external API.

Should we preserve these conservatively as public compatibility surfaces, or can we treat them as workspace-internal implementation details and simplify them directly? If compatibility is required, deprecating the legacy APIs explicitly may be clearer. Otherwise, this PR could leave load_pivots as the single loader, remove the raw-data compatibility path. Mark Hildebrand (@hildebrandmw) Aditya Krishnan (@arkrishn94)

[Mark Hildebrand (hildebrandmw)]

Thanks for working through this convergence. One point I am unsure about is the compatibility boundary. This PR adds or retains several pub APIs primarily to support existing in-workspace patterns, Although these are technically available to downstream users, they appear intended mainly for other DiskANN crates rather than as a stable external API.

Should we preserve these conservatively as public compatibility surfaces, or can we treat them as workspace-internal implementation details and simplify them directly? If compatibility is required, deprecating the legacy APIs explicitly may be clearer. Otherwise, this PR could leave load_pivots as the single loader, remove the raw-data compatibility path. Mark Hildebrand (@hildebrandmw) Aditya Krishnan (@arkrishn94)

Thanks Wei Wu (@wuw92) - while we're getting closer, DiskANN is not quite at a place (in most parts of the code) where I think trying to provide exact API backwards compatibility is realistic or desirable. In situations like this where implementations and API surface get worse as a result, I'd rather take the "L" on compatibility and do the full simplification. Especially since there is a clear path to go from before this change to after this change.

In other words, I fully support removing the older methods.

[Copilot]

Pull request overview

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

[Wei Wu (wuw92)]

Thanks for the update! left a few small comments

[Copilot]

Pull request overview

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

Suppressed comments (2)

diskann-providers/src/storage/pq_storage.rs:248

• read_pivot_file trusts the on-disk offset table values (after only checking the 4x1 shape). If the offsets are corrupt/non-monotonic or point outside the file, the loader can seek to the wrong locations and mis-parse the pivot/centroid/chunk-offset matrices, producing confusing errors or potentially accepting an invalid file layout. Add basic sanity checks that the offsets are strictly increasing, start at METADATA_SIZE, and that the final offset does not exceed the file size before using them.

        let file_offset_data = offsets.map(|x| x.into_usize());

        info!(" Offset data: {:?}", file_offset_data.as_slice());

        let pivots = read_bin_from::<f32>(&mut reader, file_offset_data[(0, 0)])?;

diskann-providers/src/model/pq/pq_construction.rs:558

• generate_pq_data_from_pivots_from_membuf_batch allocates a fresh Vec<f32> for every vector inside the parallel loop. This adds significant allocation overhead and extra copies on large datasets. Prefer reusing a per-thread scratch buffer via try_for_each_init_in_pool to avoid per-item allocations.

        .try_for_each_in_pool(pool, |(pq_slice, vector)| {
            let data = vector.iter().map(|x| (*x).into()).collect::<Vec<f32>>();
            table
                .compress_into(data.as_slice(), pq_slice)
                .map_err(ANNError::new)

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

Suppressed comments (1)

diskann-providers/src/storage/pq_storage.rs:247

• read_pivot_file trusts the on-disk file offset table after only checking its shape. A corrupted (non-monotonic / wrong start) offset table can cause read_bin_from to read the wrong regions while still producing a structurally valid BasicTable, which defeats the “file layout” validation claim in load_pivots. Add basic monotonicity + expected-start checks before using the offsets.

        let file_offset_data = offsets.map(|x| x.into_usize());

        info!(" Offset data: {:?}", file_offset_data.as_slice());

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

Suppressed comments (3)

diskann-providers/src/storage/pq_storage.rs:247

• read_pivot_file trusts the 4-entry offset table without validating that offsets are sane (start at METADATA_SIZE, strictly increasing, and within the file length). If the offset table is corrupted but still points into the file, the loader can parse structurally-valid matrices from the wrong regions and return a silently corrupted pivot table.

        let file_offset_data = offsets.map(|x| x.into_usize());

        info!(" Offset data: {:?}", file_offset_data.as_slice());

diskann-providers/src/storage/pq_storage.rs:255

• The error message for too many pivot rows says the file is "expecting {NUM_PQ_CENTROIDS} centers", but the check is actually enforcing an upper bound (pivots.nrows() > NUM_PQ_CENTROIDS). This is misleading now that load_pivots supports loading fewer than NUM_PQ_CENTROIDS centers (with call sites validating their expectations).

        if pivots.nrows() > NUM_PQ_CENTROIDS {
            return Err(ANNError::message(format!(
                "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.",
                pq_pivots,
                pivots.nrows(),
                NUM_PQ_CENTROIDS
            )));

diskann-providers/src/model/pq/fixed_chunk_pq_table.rs:447

• PR description says an impl From<BasicTable> for FixedChunkPQTable was added so callers can convert with .into(), but the code adds TryFrom<BasicTable> (fallible) and call sites use .try_into()?. Either the PR description should be updated, or the conversion semantics should be aligned with what the PR claims.

impl TryFrom<BasicTable> for FixedChunkPQTable {
    type Error = ANNError;

    fn try_from(table: BasicTable) -> Result<Self, Self::Error> {
        if table.ncenters() > NUM_PQ_CENTROIDS {
            return Err(ANNError::message(format!(
                "PQ pivot table mismatch: file has {} centers but supports at most {} centers.",
                table.ncenters(),
                NUM_PQ_CENTROIDS
            )));
        }

        Ok(Self { table })
    }
}

[Mark Hildebrand (hildebrandmw)]

Awesome thanks for your patience throughout this process. The PR looks great now!