diskann_quantization/multi_vector/
mod.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Multi-vector matrix types and distance functions.
7//!
8//! Row-major matrix abstractions for multi-vector representations, where each
9//! entity is encoded as multiple embedding vectors (e.g., per-token embeddings).
10//!
11//! # Core Types
12//!
13//! | Type | Description |
14//! |------|-------------|
15//! | [`Mat`] | Owning matrix that manages its own memory |
16//! | [`MatRef`] | Immutable borrowed view  |
17//! | [`MatMut`] | Mutable borrowed view |
18//! | [`Repr`] | Trait defining row layout (e.g., [`Standard`]) |
19//! | [`BlockTransposed`] | Owning block-transposed matrix |
20//! | [`BlockTransposedRef`] | Immutable view of a block-transposed matrix |
21//! | [`BlockTransposedMut`] | Mutable view of a block-transposed matrix |
22//! | [`QueryMatRef`] | Query wrapper for asymmetric distances |
23//! | [`MaxSim`] | Per-query-vector max similarity computation |
24//! | [`Chamfer`] | Asymmetric Chamfer distance (sum of MaxSim) |
25//!
26//! # Example
27//!
28//! ```
29//! use diskann_quantization::multi_vector::{
30//!     distance::QueryMatRef,
31//!     Chamfer, Mat, MatMut, MatRef, MaxSim, Standard,
32//! };
33//! use diskann_utils::ReborrowMut;
34//! use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
35//!
36//! // Create an owned matrix (2 vectors, dim 3, initialized to 0.0)
37//! let mut owned = Mat::new(Standard::new(2, 3).unwrap(), 0.0f32).unwrap();
38//! assert_eq!(owned.num_vectors(), 2);
39//!
40//! // Modify via mutable view
41//! let mut view = owned.reborrow_mut();
42//! if let Some(row) = view.get_row_mut(0) {
43//!    row[0] = 1.0;
44//! }
45//!
46//! // Create views from slices
47//! let query_data = [1.0f32, 0.0, 0.0, 1.0];
48//! let doc_data = [1.0f32, 0.0, 0.0, 1.0];
49//!
50//! // Wrap query as QueryMatRef for type-safe asymmetric distance
51//! let query: QueryMatRef<_> = MatRef::new(
52//!     Standard::new(2, 2).unwrap(),
53//!     &query_data,
54//! ).unwrap().into();
55//! let doc = MatRef::new(Standard::new(2, 2).unwrap(), &doc_data).unwrap();
56//!
57//! // Chamfer distance (sum of max similarities)
58//! let distance = Chamfer::evaluate(query, doc);
59//! assert_eq!(distance, -2.0); // Perfect match: -1.0 per vector
60//!
61//! // MaxSim (per-query-vector scores)
62//! let mut scores = vec![0.0f32; 2];
63//! let mut max_sim = MaxSim::new(&mut scores).unwrap();
64//! max_sim.evaluate(query, doc);
65//! assert_eq!(scores[0], -1.0);
66//! assert_eq!(scores[1], -1.0);
67//! ```
68
69pub mod block_transposed;
70pub mod distance;
71pub(crate) mod matrix;
72
73pub use block_transposed::{BlockTransposed, BlockTransposedMut, BlockTransposedRef};
74pub use distance::{Chamfer, MaxSim, MaxSimError, QueryMatRef};
75pub use matrix::{
76    Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,
77    Repr, ReprMut, ReprOwned, SliceError, Standard,
78};
diskann_quantization/multi_vector/mod.rs

diskann_quantization/multi_vector/
mod.rs
