Unify the two matrix types#1269
Conversation
There was a problem hiding this comment.
Pull request overview
This PR unifies DiskANN’s matrix abstractions by moving the more general Repr-based Mat framework into diskann-utils::matrix and rebuilding the dense Matrix/MatrixView APIs as thin row-major aliases (Matrix<T> = Mat<RowMajor<T>>). Call sites across the workspace are updated from views::* to matrix::*, and Standard is renamed to RowMajor for consistency with existing terminology.
Changes:
- Introduces
diskann-utils::matrixas the shared home forMat+Reprand the dense row-majorMatrixalias layer; removes/rewires the previousviewsusage. - Renames
Standard→RowMajorand updates constructors (Mat::new(repr, init)→Mat::from_repr(...), denseMatrix::new(gen, rows, cols)usage). - Updates quantization/table code to depend on the shared
diskann_utils::matrixframework and adjusts a small number of non-mechanical invariants/paths (e.g., pivot-dimension checks).
Reviewed changes
Copilot reviewed 126 out of 126 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| diskann/src/graph/workingset/map.rs | Updates doctests/tests to import Matrix from diskann_utils::matrix. |
| diskann/src/graph/test/synthetic.rs | Switches graph test helpers to diskann_utils::matrix::{Init, Matrix}. |
| diskann/src/graph/test/provider.rs | Updates matrix imports to the new module. |
| diskann/src/graph/test/cases/grid_insert.rs | Updates MatrixView import path to diskann_utils::matrix. |
| diskann/src/graph/start_point.rs | Rewires Matrix/MatrixView imports from views → matrix. |
| diskann/src/graph/glue.rs | Updates Batch docs/impl to target diskann_utils::matrix::Matrix. |
| diskann/src/flat/test/provider.rs | Switches test provider to diskann_utils::matrix::Matrix. |
| diskann/src/error/ann_error.rs | Updates conversion to new diskann_utils::matrix::TryFromError type. |
| diskann-utils/src/strided.rs | Adapts strided matrix utilities to new dense view types and DenseData traits under matrix. |
| diskann-utils/src/sampling/medoid.rs | Updates sampling code to use crate::matrix::MatrixView (and tests to Init/Matrix). |
| diskann-utils/src/sampling/latin_hypercube.rs | Switches LHS sampling to crate::matrix::{Matrix, MatrixView} (and tests). |
| diskann-utils/src/matrix.rs | Adds the unified Mat/Repr framework + dense row-major aliases and constructors. |
| diskann-utils/src/lib.rs | Removes views module export and adds pub mod matrix. |
| diskann-utils/src/io.rs | Rewires IO helpers to use crate::matrix::{Matrix, MatrixView} (and Init in tests). |
| diskann-tools/src/utils/ground_truth.rs | Updates ground-truth utility to consume diskann_utils::matrix::Matrix. |
| diskann-quantization/tests/compile-fail/multi/matref_rows.stderr | Mechanical Standard → RowMajor update in expected compiler output. |
| diskann-quantization/tests/compile-fail/multi/matref_rows.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matref_get_row.stderr | Mechanical Standard → RowMajor update in expected compiler output. |
| diskann-quantization/tests/compile-fail/multi/matref_get_row.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_rows.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_reborrow.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_invariant.stderr | Mechanical Standard → RowMajor update in expected compiler output. |
| diskann-quantization/tests/compile-fail/multi/matmut_invariant.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_get_row.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_rows.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_rows_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_reborrow.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_invariant.stderr | Mechanical Standard → RowMajor update in expected compiler output. |
| diskann-quantization/tests/compile-fail/multi/mat_invariant.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_get_row.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_as_view.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.rs | Updates compile-fail harness to new constructor naming and RowMajor. |
| diskann-quantization/src/views.rs | Switches quantization views to diskann_utils::matrix::{DenseData, MutDenseData}. |
| diskann-quantization/src/utils.rs | Updates helper utilities to use diskann_utils::matrix::MatrixView (and tests). |
| diskann-quantization/src/test_util.rs | Switches test utilities to diskann_utils::matrix::Matrix. |
| diskann-quantization/src/spherical/quantizer.rs | Updates matrix imports/usage and adapts tests to matrix::{self, Matrix}. |
| diskann-quantization/src/spherical/iface.rs | Updates docs/tests to use diskann_utils::matrix::{Matrix, MatrixView}. |
| diskann-quantization/src/scalar/train.rs | Updates scalar training signature to consume matrix::MatrixView. |
| diskann-quantization/src/scalar/quantizer.rs | Updates docs/tests to new matrix module paths. |
| diskann-quantization/src/scalar/mod.rs | Updates module docs to use diskann_utils::matrix::Matrix. |
| diskann-quantization/src/product/train.rs | Updates product-quantizer training signatures to matrix::MatrixView. |
| diskann-quantization/src/product/tables/transposed/table.rs | Replaces MutMatrixView with MatrixViewMut and rewires table compression APIs to new types. |
| diskann-quantization/src/product/tables/transposed/pivots.rs | Updates tests/helpers to use diskann_utils::matrix. |
| diskann-quantization/src/product/tables/test.rs | Updates test helper trait bounds from MutMatrixView → MatrixViewMut. |
| diskann-quantization/src/product/tables/basic.rs | Refactors pivot storage to owned/borrowed via a PivotTable abstraction and fixes dimension invariant to ncols. |
| diskann-quantization/src/multi_vector/mod.rs | Switches multi_vector to re-export diskann_utils::matrix types and RowMajor. |
| diskann-quantization/src/multi_vector/distance/mod.rs | Updates docs to use RowMajor and MatRef::from_repr. |
| diskann-quantization/src/multi_vector/distance/max_sim.rs | Updates docs to refer to MatRef<RowMajor<_>>. |
| diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs | Rewires tests and constructors to RowMajor + from_repr. |
| diskann-quantization/src/multi_vector/distance/kernels/layouts.rs | Updates layout binding impl to MatRef<RowMajor<T>>. |
| diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs | Updates kernel signatures from Standard → RowMajor. |
| diskann-quantization/src/multi_vector/distance/kernels/f16.rs | Updates kernel signatures from Standard → RowMajor. |
| diskann-quantization/src/multi_vector/distance/kernel.rs | Updates object-safe kernel trait to accept MatRef<RowMajor<T>>. |
| diskann-quantization/src/multi_vector/distance/fallback.rs | Updates fallback kernels and docs to use RowMajor and from_repr. |
| diskann-quantization/src/multi_vector/distance/factory.rs | Updates arch-dispatch paths to use RowMajor matrix views. |
| diskann-quantization/src/multi_vector/block_transposed.rs | Updates block-transposed impls to use matrix::{MatrixView, MatrixViewMut} and adapts NewOwned impl shape. |
| diskann-quantization/src/minmax/multi/mod.rs | Updates docs/examples to RowMajor and Mat::from_repr. |
| diskann-quantization/src/minmax/multi/meta.rs | Updates matrix type parameters (Standard → RowMajor) and tweaks allocation strategy. |
| diskann-quantization/src/minmax/multi/max_sim.rs | Updates tests to construct RowMajor input views and use Mat::from_repr. |
| diskann-quantization/src/binary/quantizer.rs | Updates tests to import Matrix from diskann_utils::matrix. |
| diskann-quantization/src/algorithms/transforms/random_rotation.rs | Rewires stored transform matrix type and constructors to diskann_utils::matrix::Matrix. |
| diskann-quantization/src/algorithms/kmeans/plusplus.rs | Switches MutMatrixView APIs to MatrixViewMut and updates tests. |
| diskann-quantization/src/algorithms/kmeans/lloyds.rs | Switches MutMatrixView APIs to MatrixViewMut and updates imports/tests. |
| diskann-quantization/src/algorithms/hadamard.rs | Updates tests to use diskann_utils::matrix::{self, Matrix, MatrixView}. |
| diskann-providers/src/utils/storage_utils.rs | Updates provider utilities to use diskann_utils::matrix::{Matrix, MatrixView}. |
| diskann-providers/src/utils/file_util.rs | Updates metadata/file utils to use diskann_utils::matrix::Matrix. |
| diskann-providers/src/test_utils/search_utils.rs | Updates test utils to diskann_utils::matrix::MatrixView. |
| diskann-providers/src/storage/pq_storage.rs | Updates PQ storage and tests to diskann_utils::matrix view types. |
| diskann-providers/src/storage/index_storage.rs | Updates tests to diskann_utils::matrix::MatrixView. |
| diskann-providers/src/model/pq/views.rs | Updates PQ model view conversions to new matrix::TryFromError. |
| diskann-providers/src/model/pq/strided.rs | Updates strided PQ conversion bounds to matrix::DenseData. |
| diskann-providers/src/model/pq/pq_construction.rs | Switches mutable matrix views to MatrixViewMut and updates call sites. |
| diskann-providers/src/model/pq/fixed_chunk_pq_table.rs | Updates PQ table construction to use Matrix::try_from and rewires view return types. |
| diskann-providers/src/model/pq/debug.rs | Updates PQ comparison utilities to matrix::MatrixView types. |
| diskann-providers/src/model/graph/provider/determinant_diversity.rs | Updates determinant-diversity to accept MatrixViewMut candidates. |
| diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs | Updates tests and dataset construction to diskann_utils::matrix. |
| diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs | Updates scalar in-mem tests to matrix::MatrixView. |
| diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs | Updates provider implementation to use diskann_utils::matrix::Matrix. |
| diskann-providers/src/index/diskann_async.rs | Updates training/test helpers to use diskann_utils::matrix and rewires signatures. |
| diskann-garnet/src/quantization.rs | Updates Garnet quantization interfaces/tests to diskann_utils::matrix. |
| diskann-garnet/src/provider.rs | Updates provider matrix imports and a MatrixView::try_from call site to new slice reference forms. |
| diskann-disk/src/storage/quant/pq/pq_generation.rs | Switches compressor interfaces from MatrixBase to MatrixView/MatrixViewMut. |
| diskann-disk/src/storage/quant/pq/pq_dataset.rs | Updates PQ dataset storage to use diskann_utils::matrix::Matrix. |
| diskann-disk/src/storage/quant/generator.rs | Updates generator logic to use matrix::MatrixView/MatrixViewMut. |
| diskann-disk/src/storage/quant/compressor.rs | Updates compressor trait to accept MatrixViewMut<u8>. |
| diskann-disk/src/search/provider/disk_provider.rs | Updates disk provider to import Matrix from diskann_utils::matrix. |
| diskann-disk/src/search/pq/quantizer_preprocess.rs | Updates scratch view construction to matrix::MatrixViewMut. |
| diskann-disk/src/build/builder/quantizer.rs | Updates build-time quantizer training input type to matrix::MatrixView. |
| diskann-bftree/src/quant.rs | Updates quantizer test setup to diskann_utils::matrix::{Init, Matrix}. |
| diskann-bftree/src/provider.rs | Updates provider docs/tests to diskann_utils::matrix types. |
| diskann-benchmark/src/utils/datafiles.rs | Updates benchmark dataset loading and conversions to diskann_utils::matrix. |
| diskann-benchmark/src/multi_vector/driver.rs | Updates multi-vector benchmark types to RowMajor (but still uses Mat::from_fn). |
| diskann-benchmark/src/index/streaming/managed.rs | Updates streaming managed index benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/streaming/full_precision.rs | Updates full-precision streaming benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/inmem/spherical.rs | Updates in-mem spherical benchmark to matrix::Matrix. |
| diskann-benchmark/src/index/inmem/scalar.rs | Updates in-mem scalar benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/inmem/product.rs | Updates in-mem product benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/build.rs | Updates index build benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/bftree/spherical.rs | Updates bftree spherical benchmark to matrix::Matrix. |
| diskann-benchmark/src/index/bftree/spherical_streaming.rs | Updates bftree spherical streaming benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/bftree/full_precision_streaming.rs | Updates bftree full-precision streaming benchmark to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/index/benchmarks.rs | Updates benchmark harness to matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/flat/search.rs | Updates flat search benchmark to matrix::Matrix. |
| diskann-benchmark/src/exhaustive/spherical.rs | Updates exhaustive spherical store to matrix::Matrix and input to matrix::MatrixView. |
| diskann-benchmark/src/exhaustive/product.rs | Updates exhaustive product store to matrix::Matrix and input to matrix::MatrixView. |
| diskann-benchmark/src/exhaustive/minmax.rs | Updates exhaustive minmax store to matrix::Matrix and input to matrix::MatrixView. |
| diskann-benchmark/src/exhaustive/algos.rs | Updates exhaustive algorithms to use matrix::{Matrix, MatrixView}. |
| diskann-benchmark/src/disk_index/search.rs | Updates disk-index benchmark search to matrix::Matrix. |
| diskann-benchmark-simd/src/lib.rs | Updates SIMD benchmark library to use matrix::{Matrix, MatrixView} and matrix::Init. |
| diskann-benchmark-core/src/streaming/executors/bigann/withdata.rs | Updates BigANN streaming executor to matrix::{Matrix, MatrixView}. |
| diskann-benchmark-core/src/search/ids.rs | Updates ID utilities to matrix::{self, Matrix} and matrix::Init. |
| diskann-benchmark-core/src/search/graph/range.rs | Updates range search to matrix::Matrix. |
| diskann-benchmark-core/src/search/graph/multihop.rs | Updates multihop search to matrix::Matrix. |
| diskann-benchmark-core/src/search/graph/knn.rs | Updates KNN search to matrix::Matrix. |
| diskann-benchmark-core/src/search/graph/inline.rs | Updates inline search to matrix::Matrix. |
| diskann-benchmark-core/src/search/api.rs | Updates search API internals to matrix::{self, Matrix} and matrix::Init. |
| diskann-benchmark-core/src/recall.rs | Updates recall computation/tests to matrix::{Matrix, MatrixView} and matrix::Init. |
| diskann-benchmark-core/src/build/graph/single.rs | Updates single build path to matrix::Matrix. |
| diskann-benchmark-core/src/build/graph/multi.rs | Updates multi build path to matrix::Matrix. |
Comments suppressed due to low confidence (7)
diskann-utils/src/matrix.rs:753
Mat::from_fn(...)is still used in the repo (e.g. diskann-benchmark/src/multi_vector/driver.rs) butdiskann_utils::matrix::Mat<RowMajor<T>>no longer defines this constructor, so the workspace won't compile. Reintroduce a dense-onlyfrom_fnthat takes an already-validatedRowMajor<T>and fills elements from a closure (similar to the previous API).
diskann-utils/src/matrix.rs:1226MatrixView::rowcomputesi * ncolswithout first assertingi < nrows. In release buildsusizemultiplication can wrap, so a very largeican return an in-bounds slice from the wrong row instead of panicking as documented. Add an explicit bounds check before doing the offset math.
diskann-utils/src/matrix.rs:1231MatrixView::row_iteruseschunks_exact(self.ncols()), which panics whenncols == 0. Zero-column matrices are explicitly allowed byRowMajor::new, so iterating rows should remain well-defined (yieldingnrowsempty slices). Using the existingrows()iterator avoids thechunks_exact(0)panic.
diskann-utils/src/matrix.rs:1243MatrixView::window_iterwill panic or misbehave whenncols == 0(division by zero ind.len() / ncols, andchunks(0)panics). It also risksncols * batchsizeoverflow in release builds. Iterating by row-index windows and usingsubviewavoids these edge cases and preserves the intended semantics.
diskann-utils/src/matrix.rs:1273MatrixView::par_window_iterhas the samencols == 0hazards aswindow_iter(par_chunks(0)panics, andd.len() / ncolsdivides by zero), plus potentialncols * batchsizeoverflow. Consider switching this to an index-window parallel iterator (usingstep_by) andsubview, mirroring the sequential implementation.
diskann-utils/src/matrix.rs:1609- The
PartialEqimpls only comparencolsand the flattened slice. Forncols == 0, all matrices/views have an empty backing slice regardless ofnrows, so shapes like5x0and10x0will compare equal. Includenrows(or comparerepr) to make equality shape-correct for zero-column matrices.
diskann-utils/src/matrix.rs:36 - The
rayon::iter::ParallelIteratorimport is currently unused in this module (all public APIs referencerayon::iter::IndexedParallelIteratorby fully-qualified path, and method bodies import the needed slice traits locally). With-D warningsin CI, this unused import will fail builds when therayonfeature is enabled. Either remove it, or import it as_if you only need it for trait resolution.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1269 +/- ##
==========================================
+ Coverage 90.05% 90.30% +0.24%
==========================================
Files 500 497 -3
Lines 95975 94773 -1202
==========================================
- Hits 86435 85588 -847
+ Misses 9540 9185 -355
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
| /// [`MatrixView`] so a [`BasicTableBase`] reads its dimensions and rows from one place. | ||
| pub trait PivotTable { | ||
| fn view(&self) -> MatrixView<'_, f32>; | ||
| } |
There was a problem hiding this comment.
I'm not sure how I feel about this. I'll need to think about it more ...
There was a problem hiding this comment.
Do you have any better suggestion?
arkrishn94
left a comment
There was a problem hiding this comment.
Thanks for taking this on Suryansh. I left some minor comments re errors.
Apart from that, I notice there is quite a large increase in the number of constructors and accessors for Mat and its friends. I recognize that most of this is needed because of the switch from the old view type but can you -
- Do an audit and confirm that each of these accessors/constructors in
matrix.rsis needed and can't be consolidated. - That these are also tested correctly? I don't notice any extra testing there.
| //! ).unwrap(); | ||
| //! | ||
| //! // Create owned matrices for quantized output using Mat::new | ||
| //! // Create owned matrices for quantized output using Mat::from_repr |
| } | ||
|
|
||
| /// Return a mutable base pointer for the [`Mat`]. | ||
| pub fn as_raw_mut_ptr(&mut self) -> *mut u8 { |
We carry two matrix types for overlapping jobs: the dense
MatrixBaseindiskann-utils(theMatrix/MatrixVieweverything reaches for), anddiskann-quantization'sRepr-genericMat<Standard<T>>, which also covers the exotic layouts (block-transposed, quantized). This PR folds them into one.The
Reprframework is the more capable of the two, so it stays — moving intodiskann-utils, where the dense types already live. Dense becomes a thin alias layer (Matrix<T> = Mat<RowMajor<T>>) that keeps every existing name, so callers barely notice, and the oldMatrixBaseis deleted.A few things worth flagging:
Standardis renamed toRowMajor— the codebase already says "row-major" everywhere;Standardwas the odd one out.Mat::new(repr, init)becomesMat::from_repr, freeing plainnewfor the denseMatrix::new(gen, rows, cols)most callers use.impl diskann_utils::matrix::Reprinstead of owning the framework.diskann-utils— it's already a grab-bag, so the types stay underdiskann_utils::matrix.How to review
Most of the diff (~110 files) is mechanical and safe to skim:
Standard→RowMajor,MatrixBase→Matrix, the constructor renames,views::*→matrix::*. The parts that want real attention:diskann-utils/src/matrix.rs— the heart of it. Git tracks it as a rename ofdiskann-quantization/src/multi_vector/matrix.rs, so read it as "that file,Standard→RowMajor, plus the dense aliases/API andDenseDatafolded in."lib.rs, the deletedviews.rs,strided.rs— the module wiring, and whereMatrixBaseused to live.multi_vector/mod.rs, the deletedmulti_vector/matrix.rs,block_transposed.rs— the quantization side now building on the shared framework.product/tables/basic.rsandminmax/multi/meta.rs— the only spots with real logic changes, not mechanical ones.