diskann_quantization/multi_vector/
mod.rs

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

diskann_quantization/multi_vector/
mod.rs
