Skip to main content

hermes_core/structures/vector/
mod.rs

1//! Vector indexing data structures
2//!
3//! This module provides a modular architecture for vector search:
4//!
5//! ## Module Structure
6//!
7//! - `ivf` - Core IVF (Inverted File Index) infrastructure
8//!   - `CoarseCentroids` - k-means clustering for coarse quantization
9//!   - `ClusterData` / `ClusterStorage` - generic cluster storage
10//!   - `SoarConfig` / `MultiAssignment` - SOAR geometry-aware assignment
11//!
12//! - `quantization` - Vector quantization methods
13//!   - `RaBitQCodebook` - RaBitQ binary quantization (32x compression)
14//!   - `PQCodebook` - Product Quantization with OPQ (ScaNN-style)
15//!   - `Quantizer` trait - common interface for quantizers
16//!
17//! - `index` - Ready-to-use IVF indexes
18//!   - `IVFRaBitQIndex` - IVF + RaBitQ
19//!   - `IVFPQIndex` - IVF + PQ
20//!
21//! ## SOAR (Spilling with Orthogonality-Amplified Residuals)
22//!
23//! The IVF module includes Google's SOAR algorithm for improved recall:
24//! - Assigns vectors to multiple clusters (primary + secondary)
25//! - Secondary clusters chosen to have orthogonal residuals
26//! - Improves recall by 5-15% with ~1.3-2x storage overhead
27
28pub mod index;
29pub mod ivf;
30pub mod quantization;
31
32/// Hard ceiling for a single decoded ANN artifact. Bincode's limit accounts
33/// for claimed container storage as well as bytes consumed, preventing a tiny
34/// corrupt payload from advertising an effectively unbounded allocation.
35///
36/// The 4 GiB 64-bit ceiling still accommodates a 20M-document classic RaBitQ
37/// segment (roughly 3.5 GB decoded at 768 dimensions), while preventing one
38/// lazy decode from claiming an operationally unbounded fraction of RAM.
39#[cfg(target_pointer_width = "64")]
40pub(crate) const MAX_DENSE_ANN_DECODE_BYTES: usize = 4 * 1024 * 1024 * 1024;
41#[cfg(not(target_pointer_width = "64"))]
42pub(crate) const MAX_DENSE_ANN_DECODE_BYTES: usize = 512 * 1024 * 1024;
43
44const ANN_DECODE_EXPANSION_FACTOR: usize = 8;
45const ANN_DECODE_FIXED_HEADROOM: usize = 16 * 1024 * 1024;
46
47fn decode_ann_with_limit<T: serde::de::DeserializeOwned, const LIMIT: usize>(
48    data: &[u8],
49) -> std::io::Result<(T, usize)> {
50    bincode::serde::decode_from_slice(data, bincode::config::standard().with_limit::<LIMIT>())
51        .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))
52}
53
54#[cfg(target_pointer_width = "64")]
55fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
56    data: &[u8],
57    budget: usize,
58) -> std::io::Result<(T, usize)> {
59    const MIB: usize = 1024 * 1024;
60    const GIB: usize = 1024 * MIB;
61    if budget <= 32 * MIB {
62        decode_ann_with_limit::<T, { 32 * MIB }>(data)
63    } else if budget <= 128 * MIB {
64        decode_ann_with_limit::<T, { 128 * MIB }>(data)
65    } else if budget <= 512 * MIB {
66        decode_ann_with_limit::<T, { 512 * MIB }>(data)
67    } else if budget <= GIB {
68        decode_ann_with_limit::<T, GIB>(data)
69    } else if budget <= 2 * GIB {
70        decode_ann_with_limit::<T, { 2 * GIB }>(data)
71    } else {
72        decode_ann_with_limit::<T, { 4 * GIB }>(data)
73    }
74}
75
76#[cfg(not(target_pointer_width = "64"))]
77fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
78    data: &[u8],
79    budget: usize,
80) -> std::io::Result<(T, usize)> {
81    const MIB: usize = 1024 * 1024;
82    if budget <= 32 * MIB {
83        decode_ann_with_limit::<T, { 32 * MIB }>(data)
84    } else if budget <= 128 * MIB {
85        decode_ann_with_limit::<T, { 128 * MIB }>(data)
86    } else {
87        decode_ann_with_limit::<T, { 512 * MIB }>(data)
88    }
89}
90
91/// Decode exactly one bincode-backed ANN artifact under a bounded allocation
92/// budget. Trailing bytes are corruption rather than a second ignored object.
93pub(crate) fn decode_ann_bincode_exact<T: serde::de::DeserializeOwned>(
94    data: &[u8],
95    description: &str,
96) -> std::io::Result<T> {
97    if data.len() > MAX_DENSE_ANN_DECODE_BYTES {
98        return Err(std::io::Error::new(
99            std::io::ErrorKind::InvalidData,
100            format!(
101                "{description} payload is {} bytes, exceeding the {}-byte decode limit",
102                data.len(),
103                MAX_DENSE_ANN_DECODE_BYTES
104            ),
105        ));
106    }
107    let budget = data
108        .len()
109        .checked_mul(ANN_DECODE_EXPANSION_FACTOR)
110        .and_then(|bytes| bytes.checked_add(ANN_DECODE_FIXED_HEADROOM))
111        .unwrap_or(MAX_DENSE_ANN_DECODE_BYTES)
112        .min(MAX_DENSE_ANN_DECODE_BYTES);
113    let (value, consumed) = decode_ann_with_relative_limit(data, budget)?;
114    if consumed != data.len() {
115        return Err(std::io::Error::new(
116            std::io::ErrorKind::InvalidData,
117            format!(
118                "{description} payload contains {} trailing bytes",
119                data.len() - consumed
120            ),
121        ));
122    }
123    Ok(value)
124}
125
126// IVF core
127pub use ivf::{
128    ClusterData, ClusterStorage, CoarseCentroids, CoarseConfig, MultiAssignment, QuantizedCode,
129    SoarConfig,
130};
131
132// Quantization
133pub use quantization::{
134    DistanceTable, PQCodebook, PQConfig, PQVector, QuantizedQuery, QuantizedVector, Quantizer,
135    RaBitQCodebook, RaBitQConfig,
136};
137
138// Indexes
139pub use index::{
140    BinaryIvfConfig, BinaryIvfIndex, IVFPQConfig, IVFPQIndex, IVFRaBitQConfig, IVFRaBitQIndex,
141    RaBitQIndex,
142};
143
144#[cfg(test)]
145mod decode_tests {
146    use super::decode_ann_bincode_exact;
147
148    #[test]
149    fn ann_bincode_decode_is_exact_and_rejects_large_claims_from_tiny_payloads() {
150        let encoded =
151            bincode::serde::encode_to_vec(vec![1u8, 2, 3], bincode::config::standard()).unwrap();
152        assert_eq!(
153            decode_ann_bincode_exact::<Vec<u8>>(&encoded, "test").unwrap(),
154            [1, 2, 3]
155        );
156
157        let mut trailing = encoded;
158        trailing.push(0);
159        assert!(decode_ann_bincode_exact::<Vec<u8>>(&trailing, "test").is_err());
160
161        // A serialized Vec starts with its usize length. This advertises a
162        // 64 MiB allocation from only a few bytes and must hit the 32 MiB tier
163        // before Vec can reserve it.
164        let oversized_claim =
165            bincode::serde::encode_to_vec(64usize * 1024 * 1024, bincode::config::standard())
166                .unwrap();
167        assert!(decode_ann_bincode_exact::<Vec<u8>>(&oversized_claim, "test").is_err());
168    }
169}