[diskann-utils] Matrix Update 4 of 4 - #1457
Conversation
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a workspace-wide unsafe foundational API and still has an unresolved mismatch between the documented and implemented MatrixMut contract.
Review effort: Balanced
Findings: 1
Open (3)
What changed in this PR
Migrates the workspace to the new row-major matrix API, replacing legacy matrix types with Owned, Ref, and Mut plus the Matrix/MatrixMut traits.
Changes:
- Introduces the row-major matrix traits, views, ownership types, and iterators.
- Updates matrix construction, traversal, mutation, and parallel APIs workspace-wide.
- Adapts storage, quantization, indexing, benchmarks, tests, and documentation.
| File | Description |
|---|---|
diskann/src/graph/workingset/map.rs |
Migrates working-set matrices. |
diskann/src/graph/test/synthetic.rs |
Migrates synthetic test data. |
diskann/src/graph/test/provider.rs |
Updates provider batch types. |
diskann/src/graph/test/cases/helpers.rs |
Updates row iteration. |
diskann/src/graph/test/cases/grid_insert.rs |
Updates matrix views and ownership. |
diskann/src/graph/start_point.rs |
Migrates start-point generation. |
diskann/src/graph/pipnn/topk.rs |
Migrates top-k matrix operations. |
diskann/src/graph/pipnn/partition_kernel.rs |
Updates partition-kernel views. |
diskann/src/graph/pipnn/mod.rs |
Updates distance scratch views. |
diskann/src/graph/pipnn/leaf_kernel.rs |
Updates leaf-kernel views. |
diskann/src/graph/glue.rs |
Updates batch matrix integration. |
diskann/src/flat/test/provider.rs |
Migrates flat test provider. |
diskann/src/flat/test/harness.rs |
Updates brute-force row iteration. |
diskann/src/error/ann_error.rs |
Updates matrix error conversions. |
diskann-utils/src/views/mod.rs |
Retains dense-data traits and exports row-major APIs. |
diskann-utils/src/views/rowmajor.rs |
Defines the new matrix API. |
diskann-utils/src/views/rowmajor/iter.rs |
Implements row/window iterators. |
diskann-utils/src/strided.rs |
Integrates row-major views with strided matrices. |
diskann-utils/src/sampling/latin_hypercube.rs |
Migrates sampling matrices. |
diskann-utils/src/sampling/medoid.rs |
Migrates medoid calculations. |
diskann-utils/src/io.rs |
Updates binary matrix I/O. |
diskann-utils/src/internal.rs |
Adds pointer ownership helpers. |
diskann-tools/src/utils/relative_contrast.rs |
Updates matrix traversal. |
diskann-tools/src/utils/ground_truth.rs |
Migrates ground-truth matrices. |
diskann-tools/src/bin/generate_minmax.rs |
Imports the new matrix trait. |
diskann-tools/src/bin/compute_streaming_groundtruth.rs |
Updates streaming ground-truth matrices. |
diskann-quantization/src/utils.rs |
Migrates statistical matrix utilities. |
diskann-quantization/src/test_util.rs |
Updates quantization test matrices. |
diskann-quantization/src/spherical/quantizer.rs |
Migrates spherical training inputs. |
diskann-quantization/src/scalar/train.rs |
Updates scalar training views. |
diskann-quantization/src/scalar/quantizer.rs |
Updates scalar quantizer examples/tests. |
diskann-quantization/src/scalar/mod.rs |
Updates scalar module documentation. |
diskann-quantization/src/product/tables/transposed/pivots.rs |
Migrates pivot tests. |
diskann-quantization/src/product/tables/basic.rs |
Generalizes pivots over Matrix. |
diskann-quantization/src/multi_vector/matrix.rs |
Updates row-major view conversion. |
diskann-quantization/src/multi_vector/distance/factory.rs |
Imports matrix operations. |
diskann-quantization/src/matrix_kernels/test_util.rs |
Updates generated test matrices. |
diskann-quantization/src/matrix_kernels/maxsim/test.rs |
Migrates max-sim fixtures. |
diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f32.rs |
Updates test imports. |
diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f16.rs |
Updates test imports. |
diskann-quantization/src/matrix_kernels/blocks/unpacked.rs |
Migrates unpacked block views. |
diskann-quantization/src/matrix_kernels/blocks/packed.rs |
Migrates packed block tests. |
diskann-quantization/src/binary/quantizer.rs |
Updates binary quantizer tests. |
diskann-quantization/src/algorithms/transforms/random_rotation.rs |
Migrates rotation matrices. |
diskann-quantization/src/algorithms/hadamard.rs |
Migrates Hadamard test matrices. |
diskann-providers/src/utils/storage_utils.rs |
Updates storage matrix I/O. |
diskann-providers/src/utils/file_util.rs |
Migrates multivector loading. |
diskann-providers/src/test_utils/search_utils.rs |
Updates ground-truth views. |
diskann-providers/src/storage/pq_storage.rs |
Migrates PQ persistence matrices. |
diskann-providers/src/storage/index_storage.rs |
Updates index-storage tests. |
diskann-providers/src/model/pq/views.rs |
Updates bridged matrix errors. |
diskann-providers/src/model/pq/pq_construction.rs |
Migrates PQ construction views. |
diskann-providers/src/model/pq/fixed_chunk_pq_table.rs |
Updates PQ table views. |
diskann-providers/src/model/pq/debug.rs |
Migrates PQ comparison inputs. |
diskann-providers/src/model/graph/provider/determinant_diversity.rs |
Updates mutable candidate matrices. |
diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs |
Migrates spherical provider tests. |
diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs |
Updates scalar training view. |
diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs |
Migrates reranking matrices. |
diskann-providers/src/index/wrapped_async.rs |
Updates row iteration. |
diskann-inmem/src/repr/spherical.rs |
Migrates spherical representation matrices. |
diskann-inmem/src/repr/full.rs |
Migrates full-precision matrices. |
diskann-inmem/src/provider.rs |
Updates provider tests. |
diskann-inmem/integration/support/io.rs |
Migrates converted datasets. |
diskann-inmem/integration/support/datatype.rs |
Updates dataset matrix variants. |
diskann-inmem/integration/index/tests.rs |
Migrates result matrices. |
diskann-inmem/integration/index/runner.rs |
Updates integration data bundles. |
diskann-garnet/src/quantization.rs |
Migrates quantizer training views. |
diskann-garnet/src/provider.rs |
Updates provider training matrices. |
diskann-disk/src/storage/quant/pq/pq_generation.rs |
Migrates PQ compression matrices. |
diskann-disk/src/storage/quant/pq/pq_dataset.rs |
Updates compressed PQ storage. |
diskann-disk/src/storage/quant/generator.rs |
Migrates quantization generator views. |
diskann-disk/src/storage/quant/compressor.rs |
Updates compressor interface. |
diskann-disk/src/search/provider/disk_provider.rs |
Migrates search candidate matrices. |
diskann-disk/src/search/pq/quantizer_preprocess.rs |
Updates PQ scratch view. |
diskann-disk/src/build/builder/quantizer.rs |
Migrates quantizer training matrices. |
diskann-disk/src/build/builder/core.rs |
Updates build-test traversal. |
diskann-bftree/src/quant.rs |
Migrates test quantizer data. |
diskann-benchmark/src/utils/datafiles.rs |
Updates matrix-loading APIs. |
diskann-benchmark/src/index/streaming/managed.rs |
Migrates managed stream interfaces. |
diskann-benchmark/src/index/streaming/full_precision.rs |
Updates streaming insertion matrices. |
diskann-benchmark/src/index/inmem2.rs |
Migrates in-memory benchmark data. |
diskann-benchmark/src/index/inmem/spherical.rs |
Updates spherical benchmark matrices. |
diskann-benchmark/src/index/inmem/scalar.rs |
Updates scalar benchmark views. |
diskann-benchmark/src/index/inmem/product.rs |
Updates product benchmark views. |
diskann-benchmark/src/index/build.rs |
Migrates build helpers. |
diskann-benchmark/src/index/bftree/spherical.rs |
Updates BfTree spherical data. |
diskann-benchmark/src/index/bftree/spherical_streaming.rs |
Migrates spherical streaming views. |
diskann-benchmark/src/index/bftree/full_precision.rs |
Imports matrix operations. |
diskann-benchmark/src/index/bftree/full_precision_streaming.rs |
Migrates full-precision streaming views. |
diskann-benchmark/src/index/benchmarks.rs |
Updates benchmark matrix interfaces. |
diskann-benchmark/src/flat/search.rs |
Migrates flat-search matrices. |
diskann-benchmark/src/exhaustive/spherical.rs |
Updates spherical exhaustive store. |
diskann-benchmark/src/exhaustive/product.rs |
Updates product exhaustive store. |
diskann-benchmark/src/exhaustive/minmax.rs |
Updates min-max exhaustive store. |
diskann-benchmark/src/exhaustive/algos.rs |
Migrates linear-search matrices. |
diskann-benchmark/src/disk_index/search.rs |
Updates disk-search query matrices. |
diskann-benchmark-simd/src/lib.rs |
Migrates SIMD benchmark matrices. |
diskann-benchmark-core/src/streaming/graph/test.rs |
Imports matrix operations. |
diskann-benchmark-core/src/streaming/executors/bigann/withdata.rs |
Migrates streaming datasets. |
diskann-benchmark-core/src/search/ids.rs |
Updates result-ID matrices. |
diskann-benchmark-core/src/search/graph/range.rs |
Migrates range-search queries. |
diskann-benchmark-core/src/search/graph/multihop.rs |
Migrates multihop queries. |
diskann-benchmark-core/src/search/graph/knn.rs |
Migrates KNN queries. |
diskann-benchmark-core/src/search/graph/inline.rs |
Migrates inline-filter queries. |
diskann-benchmark-core/src/search/graph/filtered_range.rs |
Migrates filtered-range queries. |
diskann-benchmark-core/src/search/api.rs |
Updates search output matrices. |
diskann-benchmark-core/src/build/graph/single.rs |
Migrates single-insert data. |
diskann-benchmark-core/src/build/graph/multi.rs |
Migrates multi-insert batches. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Suryansh Gupta (suri-kumkaran)
left a comment
There was a problem hiding this comment.
nit: Old type names left in comments: strided.rs:187, 420 , generator.rs:139 , tables/test.rs:290 , multi_vector/matrix.rs:891 , agents.md:50 , diskann-utils/README.md:7 .
Another reminder: Nightly miri doesn't run on diskann-utils, which now holds most of this unsafe code.
Thanks - I updated the stale references I could find, and added |
# DiskANN v0.60.0 ## Major Changes ### `diskann-utils`: Matrix API changes There have been significant changes to the matrix APIs in `diskann-utils`. The previous contents of `diskann_utils::views` have been moved to `diskann_utils::views::rowmajor`. The following type changes have been made | v0.59 | v0.60 | Remarks | |-|-|-| | `views::Matrix<T>` | `views::rowmajor::Owned<T>` | Size changed from 32 bytes to 24 bytes | | `views::MatrixView<'a, T>` | `views::rowmajor::Ref<'a, T>` | Size changed from 32 bytes to 24 bytes | | `views::MatrixViewMut<'a, T>` | `views::rowmajor::Mut<'a, T>` | Size changed from 32 bytes to 24 bytes | | `MatrixBase` inherent methods | `rowmajor::Matrix` and `rowmajor::MatrixMut` trait methods | | The `rowmajor::Matrix` and `rowmajor::MatrixMut` trait methods need to be imported to access most matrix functionality. In addition, the following methods have been renamed and changed: | v0.59 | v0.60 | Remark | |-|-|-| | `row_iter` | `rows` | No longer panics when `ncols == 0` | | `row_iter_mut` | `rows_mut` | No longer panics when `ncols == 0` | | `par_row_iter` | `par_rows` | No longer panics when `ncols == 0` | | `par_row_iter_mut` | `par_rows_mut` | | | `get_row_unchecked` | `row_unchecked` | | | `get_row_unchecked_mut` | `row_unchecked_mut` | | | `as_mut_view` | `as_view_mut` | | | `try_from` constructors | `try_from_data` | | In addition, * The `Index` and `IndexMut` implementations have been removed. Use `Matrix::element` and `Matrix::element_mut` instead. * The `Init` and `Generator` style constructors for `Matrix::new` have been removed. Instead, `rowmajor::Owned::from_fn` and `rowmajor::Owned::from_element` can be used instead. The former's closure has been changed to take an argument `rowmajor::RowCol` indicated the row and column being initialized. Both of these constructors take now take the initializer as the last argument. * Matrix constructors validate against integer overflow and `isize::MAX` allocations via a `rowmajor::Layout<T>` layout. Constructors panic if these validations fail. `try_*` variations of these constructors which return `rowmajor::LayoutError` instead. ### `diskann-utils`: Strided API changes * Structs `StridedBase`, `StridedView`, and `MutStridedView` were collapsed into a single, immutable `Strided<'a, T>`. * `strided::linear_length` was removed, use `strided::Layout::linear_length` instead. * `strided::TryFromErrorLight` was removed. * `strided::TryFromError` changed from a struct to an enum, and `into_inner`/`as_static` were removed. ### `diskann` `InsertStrategy` is no longer a `SearchStrategy` subtrait. Implementations must now provide ```rust type SearchAccessor; type SearchAccessorError; fn insert_search_accessor(...) -> Result<Self::SearchAccessor, Self::SearchAccessorError>; ``` Previously, `insert_search_accessor` had a provided implementation forwarding to `SearchStrategy::search_accessor`. That default is now gone. This change was done to decouple insert element types from search element types. ### `diskann-disk` The publicly re-exported types below were removed in favor of standard buffered I/O: ``` diskann_disk::storage::CachedReader diskann_disk::storage::CachedWriter ``` Consumers should migrate to `std::io::BufReader`/`BufWriter` and the relevant storage-provider reader/writer. `READ_WRITE_BLOCK_SIZE` also changed from `u64` to `usize`. ### `diskann-quantization` These public traits now require `Debug`: ``` spherical::iface::Quantizer spherical::iface::DynQueryComputer spherical::iface::DynDistanceComputer ``` ### `diskann-inmem` (experimental) The old `layers` abstraction was removed and replaced by the `Representation`/`Store`/`Slots` architecture. And support for a spherically quantized backed was added. # Full Change Log * [diskann-inmem] Prepare code for quantization and beyond by @hildebrandmw in #1352 * Add Neon inner product U8xU1, U8xU2 and U8xU4 kernels by @pfoxARM in #1372 * [CI] Check all workspace features and private docs by @wuw92 in #1400 * [CI] Check Rust license headers by @wuw92 in #1401 * [disk] Replace custom cached I/O with standard buffered readers and writers by @wuw92 in #1396 * [dikkann_wide] Integer SIMDMinMax by @suri-kumkaran in #1407 * Add Neon Squared L2 kernels for USlice2 and USlice4 by @pfoxARM in #1398 * Bump the github-actions group with 3 updates by @dependabot[bot] in #1406 * [CI] Check workspace direct dependency versions by @wuw92 in #1410 * Train fresh disk-search PQ codebooks for each index build by @xwj-ox in #1299 * [diskann-quantization] Make most things in `diskann-quantization/spherical/iface.rs` debuggable by @hildebrandmw in #1414 * [diskann] Make `InsertStrategy` independent of `SearchStrategy` by @hildebrandmw in #1393 * Use #[expect] to detect unnecessary lint suppressions by @wuw92 in #1409 * [diskann-garnet] Store attributes during insert, not after by @metajack in #1412 * [diskann-wide] Fix `alias!` by @hildebrandmw in #1417 * [diskann-utils] Simplify `StridedView`. by @hildebrandmw in #1376 * [diskann-wide] More Neon Zip/Unzip impls by @hildebrandmw in #1405 * [diskann-garnet] Deserialize max id for quantization on load by @metajack in #1419 * More ergonomic `spherical::iface::Quantizer` constructors by @hildebrandmw in #1418 * [diskann-bftree] Clamp number of concurrency test parallel workers by @hildebrandmw in #1433 * Tweak behavior of adaptive L by @magdalendobson in #1354 * Bump the github-actions group with 2 updates by @dependabot[bot] in #1436 * Reform Range Search by @magdalendobson in #1337 * PiPNN 2/6: add numerical kernels by @SeliMeli in #1287 * Matrix overhaul 1 of N - Data Accessors by @hildebrandmw in #1415 * Add `assert_contains!` helper by @hildebrandmw in #1421 * Matrix overhaul 2 of N - `Matrix` invariants by @hildebrandmw in #1420 * [CI] Run SDE tests for diskann with all features by @SeliMeli in #1440 * [diskann-wide] Add Neon `vabdq_u8` and `vabdq_s8`. by @hildebrandmw in #1437 * [diskann-inmem] Add support for spherical quantization by @hildebrandmw in #1432 * [benchmark-runner] Bump edition to 2024 by @hildebrandmw in #1444 * [diskann-wide] Add byte-vector reinterpretations for u32x8 and u64x8 by @partychen in #1445 * Propagate invalid matrix layout errors in production paths by @arrayka in #1453 * Add sparse vector distance kernels (f32/f16 L2, inner product, cosine) by @mana-agarwal in #1434 * Bump the github-actions group with 3 updates by @dependabot[bot] in #1454 * [diskann-utils] Matrix Update 3 of N (constructors) by @hildebrandmw in #1451 * [diskann-utils] Matrix Update 4 of 4 by @hildebrandmw in #1457 ## New Contributors * @mana-agarwal made their first contribution in #1434 **Full Changelog**: v0.59.0...v0.60.0


This is the last big change in the core Matrix API (I hope). The main changes are as follows:
The module
diskann_utils::viewsis moved todiskann_utils::views::rowmajor. DiskANN already supports multiple matrix types scattered about, withStrided(indiskann_utils::strided) and the block-transpose matrix indiskann_quantization::multi_vector::block_transpose. I would eventually like to consolidate these somewhat. Differentiating by storage philosophy seems reasonable.The API for working with matrices is changed from inherent methods on
MatrixBaseto two traits:Matrixfor non-mutating operations andMatrixMutfor mutating operations. The core requirement for these traits is:The trait is
unsafedue to idempotency and lifetime requirements. Please read the safety documentation for the relevant traits.The types
Matrix,MatrixView, andMatrixViewMutare replaced byOwned,Ref, andMutrespectively. As these are fairly common names, most code using these types has been updated to include therowmajormodule name as part of the type identity. So, for example,Matrixwould becomerowmajor::Owned.One key factor behind these types is that they are all just 3 words in size (24-bytes on a 64-bit system), consisting of just a pointer, the number of columns, and the number of rows. This is 8-bytes less than the existing matrices, which have slices for their base storage.
API changes:
Matrix::row_iterbecomesMatrix::rowsand returns aniter::Rowscustom iterator. As a side-effect, this method no longer panics when the number of columns is zero.MatrixMut::row_iter_mutbecomesMatrixMut::rows_mutand returns aniter::RowsMutcustom iterator. Likerows, this iterator works correctly when the number of columns is zero.MatrixMut::window_iteralso returns a customiter::Windowsthat works correctly when the number of columns is zero.Matrix::get_row_uncheckedbecomesMatrix::row_uncheckedfor symmetry withMatrix::element_unchecked. Similarlly,MatrixMut::get_row_unchecked_mutbecomesMatrixMut::row_unchecked_mut.Matrix::par_row_iterbecomesMatrix::par_rowsand works correctly when the number of columns is zero.Matrix::par_window_iterworks correctly when the number of columns is zero.MatrixMut::par_row_iter_mutbecomesMatrix::par_rows_mut. It still panics when the number of columns is zero, but the message is a little clearer.MatrixMut::par_window_iter_mutstill panics when the number of columns is zero, but has a better error message. It also fixes an overflow condition that previously existed when the batchsize was large.MatrixMut::as_mut_viewbecomesMatrixMut::as_view_mutfor consistency.try_frombecometry_from_datato differentiate fromTryFromThe bulk of the meaningful changes are in
diskann-utils/src/view/*. Everything else is the tedious process of updating all the uses of this fundamental data type.Thoughts on the API
I feel this is a good middle ground between the current matrix in
diskann-utilsand the more flexible but tedious matrix indiskann-quantization.The problem with the current
diskann-utilsmatrix is that while it's fairly simple and doesn't require a trait to call its inherent methods, it's harder to extend. The underlying storage can be abstracted withDenseDataandDenseDataMut, but this can include some redundant information. For example,MatrixView<'_, T>is 32 bytes because it contains a&[T]at its base. The length tracked by the slice is redundant since we can compute it from the already stored number of rows and columns. Further, it's impossible to resize. With this approach, we can do something likeThe problem with the matrix in
diskann-quantizationis that it still runs into the problems of non-extendable inherent methods and is writing newReprimplementations is extremely tedious. Additionally, writing new matrix types (e.g. block transpose) requires forwarding a large number of inherent methods. This design (small core with everything layered on top) adds a lot of flexibility at the cost of needing to pull in a trait or two.Next Steps
As a follow-up to this, we can begin removing the
Standardrepr to thediskann-quantizationmatrix. Eventually,BlockTransposeshould get the same treatment, which would hopefully simplify it some. The min-max reprs can also be rewritten to be generic overM: Matrix<Element = u8>.