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 converges the two PQ pivot loading paths so they share a single parser for the pq_pivots.bin file layout.
Previously, load_existing_pivot_data and load_pq_pivots_bin each reimplemented the same file-reading logic for offsets, pivots, centroid data, and chunk offsets. This change introduces a shared internal helper that reads and validates those parts once, while preserving the existing public behavior.

The PR also removes a duplicate test-only PQ pivot loader and updates the tests to exercise the production PQStorage::load_pq_pivots_bin path directly. This reduces duplicated parsing logic and keeps future changes to the pivot file format or validation rules localized.

Any other comments?

Validated with:
cargo test -p diskann-providers pq_storage
cargo test -p diskann-providers fixed_chunk_pq_table

[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-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.58%. Comparing base (6caca35) to head (1d351e2).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1276      +/-   ##
==========================================
+ Coverage   91.50%   91.58%   +0.07%     
==========================================
  Files         498      497       -1     
  Lines       95522    95549      +27     
==========================================
+ Hits        87407    87505      +98     
+ Misses       8115     8044      -71     

Flag	Coverage Δ
miri	91.58% <100.00%> (+0.07%)	⬆️
unittests	91.54% <100.00%> (+0.07%)	⬆️

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

Files with missing lines	Coverage Δ
...ann-providers/src/model/pq/fixed_chunk_pq_table.rs	95.70% <100.00%> (+3.75%)	⬆️
diskann-providers/src/storage/pq_storage.rs	96.24% <100.00%> (+11.47%)	⬆️

... and 4 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)
            )));