PQ: tighten dim contract; right-size scratch buffer

[wuw92]

Background

Historically, query buffers came from AlignedBoxWithSlice, which silently rounded length up to a multiple of 8 for SIMD alignment. Downstream populate functions therefore had to accept query.len() >= dim instead of query.len() == dim — pre-PR comment in TableL2::populate:

Alignment means that the size of query gets increased ...
This makes is VERY hard to do error checking on dimension propagation.

With #960 removing AlignedBoxWithSlice from the PQ path, the subtree can refactor dim handling along the boundary convert / internal trust idiom.

Three-layer dim contract

Layer	Where	Action	Failure
Boundary (inmem)	QueryComputer::new, MultiQueryComputer::new, DistanceComputer::evaluate_similarity	Validate len == dim	Result::Err / assert_eq!
Boundary (disk)	PQScratch::set	Validate len >= dim, slice [..dim]	Result::Err on undersize
Internal	TableL2/IP/Cosine::{new, populate}	Trusted, no re-validation	—
Inner kernel	preprocess_query, populate_chunk_distances_impl, direct_distance_impl	Contract boundary debug_assert_eq!	dev/CI panic, release zero-cost; also fail-loud for direct OSS callers

Boundary methods take &[f32]. Quantized inputs are decoded via VectorRepr::as_f32 once at the caller boundary; the PQ subtree is f32-only internally.

Why entries take &[f32]

Into<f32> is per-element. MinMaxElement<8> is a single byte that can't decode without the full slice's trailing metadata — it cannot implement Into<f32>. Production callers supporting MinMax were therefore always pre-decoding via VectorRepr::as_f32 and passing &[f32] upstream; the previous <U: Into<f32>> generic on PQ entries was effectively orphan-only.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 99.30070% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 89.48%. Comparing base (d516da1) to head (6b8ade0).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-disk/src/search/pq/quantizer_preprocess.rs	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1044      +/-   ##
==========================================
- Coverage   90.60%   89.48%   -1.13%     
==========================================
  Files         461      461              
  Lines       85559    85742     +183     
==========================================
- Hits        77525    76723     -802     
- Misses       8034     9019     +985     

Flag	Coverage Δ
miri	89.48% <99.30%> (-1.13%)	⬇️
unittests	89.32% <99.30%> (-1.25%)	⬇️

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

Files with missing lines	Coverage Δ
diskann-disk/src/search/pq/pq_scratch.rs	89.18% <100.00%> (+9.64%)	⬆️
diskann-disk/src/search/provider/disk_provider.rs	90.88% <100.00%> (-0.01%)	⬇️
diskann-disk/src/storage/quant/pq/pq_dataset.rs	97.33% <100.00%> (+0.19%)	⬆️
...aph/provider/async_/experimental/multi_pq_async.rs	96.40% <100.00%> (ø)
...ovider/async_/fast_memory_quant_vector_provider.rs	98.46% <100.00%> (ø)
...ph/provider/async_/memory_quant_vector_provider.rs	98.26% <100.00%> (ø)
diskann-providers/src/model/pq/distance/cosine.rs	98.76% <100.00%> (-1.24%)	⬇️
diskann-providers/src/model/pq/distance/dynamic.rs	96.81% <100.00%> (+2.70%)	⬆️
...nn-providers/src/model/pq/distance/innerproduct.rs	100.00% <100.00%> (ø)
diskann-providers/src/model/pq/distance/l2.rs	100.00% <100.00%> (ø)
... and 4 more

... and 39 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

This PR tightens and clarifies the PQ query-dimension contract by moving conversion/validation to boundary APIs, making PQ internals f32-only, and sizing scratch buffers to the PQ table’s logical dimension (rather than alignment-padded lengths). It also hides previously public PQ internals that were effectively low-level/FFI-oriented.

Changes:

• Switch PQ “entry points” (e.g., QueryComputer / MultiQueryComputer) to accept &[f32], validate len >= dim, and internally slice to [..dim].
• Add debug-only dimension assertions inside inner kernels and right-size PQ scratch/query buffers to the PQ table’s logical dimension.
• Remove public exposure of some PQ internals (direct_distance_impl, populate_chunk_distances_impl, inner_product_raw) and adjust callers/tests accordingly.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
diskann-providers/src/model/pq/mod.rs	Stops re-exporting direct_distance_impl from the PQ module surface.
diskann-providers/src/model/pq/fixed_chunk_pq_table.rs	Makes several PQ internals private and adds debug-only dim assertions in kernels.
diskann-providers/src/model/pq/distance/test_utils.rs	Updates PQ distance test helpers to generate/use f32 queries directly.
diskann-providers/src/model/pq/distance/multi.rs	Updates MultiQueryComputer::new to take &[f32] and adjusts tests.
diskann-providers/src/model/pq/distance/l2.rs	Changes TableL2 constructor/populate path to accept &[f32]; updates tests.
diskann-providers/src/model/pq/distance/innerproduct.rs	Changes TableIP constructor/populate path to accept &[f32]; updates tests.
diskann-providers/src/model/pq/distance/dynamic.rs	Enforces boundary dim checks/slicing in QueryComputer::new; slices FP input in DistanceComputer. Adds tests for undersized query handling.
diskann-providers/src/model/pq/distance/cosine.rs	Changes cosine query handling to copy from &[f32]; updates tests.
diskann-providers/src/model/mod.rs	Removes re-export of direct_distance_impl from the top-level model API.
diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs	Moves query decoding to the provider boundary via VectorRepr::as_f32.
diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs	Aligns query constraint to T: VectorRepr (dropping Copy).
diskann-providers/src/model/graph/provider/async_/experimental/multi_pq_async.rs	Decodes queries to f32 at the boundary and passes &[f32] into PQ code.
diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs	Aligns query constraint to T: VectorRepr (dropping Copy).
diskann-disk/src/storage/quant/pq/pq_dataset.rs	Adds PQData::get_dim() to expose PQ logical dimension for sizing buffers.
diskann-disk/src/search/provider/disk_provider.rs	Uses PQ logical dim and decodes query to f32 before PQScratch::set.
diskann-disk/src/search/pq/quantizer_preprocess.rs	Removes manual query slicing and relies on right-sized rotated_query.
diskann-disk/src/search/pq/pq_scratch.rs	Updates PQScratch::set to accept &[f32] and copy exactly dim elements into a right-sized buffer.

Comments suppressed due to low confidence (1)

diskann-providers/src/model/pq/distance/multi.rs:377

• MultiQueryComputer::new now requires &[f32] rather than &[U: Into<f32> + Copy], which is a breaking change for downstreams passing non-f32 query types directly. If this is intended, consider a transition path (deprecated overload or helper) consistent with the new boundary decode pattern.

    /// Construct a new `MultiQueryComputer` with the requested metric and query.
    pub fn new(table: MultiTable<T, I>, metric: Metric, query: &[f32]) -> ANNResult<Self> {
        let s = match table {
            MultiTable::One { table, version } => Self::One {
                computer: { QueryComputer::new(table, metric, query, None)? },
                version,

---

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