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//!   - `SoarConfig` / `MultiAssignment` - SOAR geometry-aware assignment
10//!
11//! - `quantization` - Residual product quantization
12//!   - `PQCodebook` - the index-global OPQ codebook
13//!
14//! - `index` - Segment payloads for the production ANN implementations
15//!   - `IVFPQIndex` - float vectors with residual PQ codes
16//!   - `BinaryIvfIndex` - exact packed binary vectors
17//!
18//! ## SOAR (Spilling with Orthogonality-Amplified Residuals)
19//!
20//! The IVF module includes Google's SOAR algorithm for improved recall:
21//! - Assigns vectors to multiple clusters (primary + secondary)
22//! - Secondary clusters chosen to have orthogonal residuals
23//! - Improves recall by 5-15% with ~1.3-2x storage overhead
24
25pub mod index;
26pub mod ivf;
27mod kmeans;
28pub mod quantization;
29
30/// Hard ceiling for a single decoded ANN artifact. Bincode's limit accounts
31/// for claimed container storage as well as bytes consumed, preventing a tiny
32/// corrupt payload from advertising an effectively unbounded allocation.
33///
34/// A billion-scale 2,560-bit multi-value binary segment can legitimately
35/// exceed 8 GiB. Keep the serialized-input ceiling distinct from bincode's
36/// decode budget: bincode charges both bytes consumed and requested container
37/// allocations, so a valid 9 GiB payload needs more than a 9 GiB limit.
38#[cfg(target_pointer_width = "64")]
39pub(crate) const MAX_DENSE_ANN_PAYLOAD_BYTES: usize = 16 * 1024 * 1024 * 1024;
40#[cfg(not(target_pointer_width = "64"))]
41pub(crate) const MAX_DENSE_ANN_PAYLOAD_BYTES: usize = 512 * 1024 * 1024;
42
43#[cfg(target_pointer_width = "64")]
44const MAX_DENSE_ANN_DECODE_BUDGET: usize = 32 * 1024 * 1024 * 1024;
45#[cfg(not(target_pointer_width = "64"))]
46const MAX_DENSE_ANN_DECODE_BUDGET: usize = MAX_DENSE_ANN_PAYLOAD_BYTES;
47
48const ANN_DECODE_EXPANSION_FACTOR: usize = 8;
49const ANN_DECODE_FIXED_HEADROOM: usize = 16 * 1024 * 1024;
50
51fn decode_ann_with_limit<T: serde::de::DeserializeOwned, const LIMIT: usize>(
52    data: &[u8],
53) -> std::io::Result<(T, usize)> {
54    bincode::serde::decode_from_slice(data, bincode::config::standard().with_limit::<LIMIT>())
55        .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))
56}
57
58#[cfg(target_pointer_width = "64")]
59fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
60    data: &[u8],
61    budget: usize,
62) -> std::io::Result<(T, usize)> {
63    const MIB: usize = 1024 * 1024;
64    const GIB: usize = 1024 * MIB;
65    if budget <= 32 * MIB {
66        decode_ann_with_limit::<T, { 32 * MIB }>(data)
67    } else if budget <= 128 * MIB {
68        decode_ann_with_limit::<T, { 128 * MIB }>(data)
69    } else if budget <= 512 * MIB {
70        decode_ann_with_limit::<T, { 512 * MIB }>(data)
71    } else if budget <= GIB {
72        decode_ann_with_limit::<T, GIB>(data)
73    } else if budget <= 2 * GIB {
74        decode_ann_with_limit::<T, { 2 * GIB }>(data)
75    } else if budget <= 4 * GIB {
76        decode_ann_with_limit::<T, { 4 * GIB }>(data)
77    } else if budget <= 8 * GIB {
78        decode_ann_with_limit::<T, { 8 * GIB }>(data)
79    } else if budget <= 16 * GIB {
80        decode_ann_with_limit::<T, { 16 * GIB }>(data)
81    } else {
82        decode_ann_with_limit::<T, { 32 * GIB }>(data)
83    }
84}
85
86#[cfg(not(target_pointer_width = "64"))]
87fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
88    data: &[u8],
89    budget: usize,
90) -> std::io::Result<(T, usize)> {
91    const MIB: usize = 1024 * 1024;
92    if budget <= 32 * MIB {
93        decode_ann_with_limit::<T, { 32 * MIB }>(data)
94    } else if budget <= 128 * MIB {
95        decode_ann_with_limit::<T, { 128 * MIB }>(data)
96    } else {
97        decode_ann_with_limit::<T, { 512 * MIB }>(data)
98    }
99}
100
101/// Decode exactly one bincode-backed ANN artifact under a bounded allocation
102/// budget. Trailing bytes are corruption rather than a second ignored object.
103pub(crate) fn decode_ann_bincode_exact<T: serde::de::DeserializeOwned>(
104    data: &[u8],
105    description: &str,
106) -> std::io::Result<T> {
107    if data.len() > MAX_DENSE_ANN_PAYLOAD_BYTES {
108        return Err(std::io::Error::new(
109            std::io::ErrorKind::InvalidData,
110            format!(
111                "{description} payload is {} bytes, exceeding the {}-byte decode limit",
112                data.len(),
113                MAX_DENSE_ANN_PAYLOAD_BYTES
114            ),
115        ));
116    }
117    let budget = data
118        .len()
119        .checked_mul(ANN_DECODE_EXPANSION_FACTOR)
120        .and_then(|bytes| bytes.checked_add(ANN_DECODE_FIXED_HEADROOM))
121        .unwrap_or(MAX_DENSE_ANN_DECODE_BUDGET)
122        .min(MAX_DENSE_ANN_DECODE_BUDGET);
123    let (value, consumed) = decode_ann_with_relative_limit(data, budget)?;
124    if consumed != data.len() {
125        return Err(std::io::Error::new(
126            std::io::ErrorKind::InvalidData,
127            format!(
128                "{description} payload contains {} trailing bytes",
129                data.len() - consumed
130            ),
131        ));
132    }
133    Ok(value)
134}
135
136// IVF core
137pub use ivf::{CoarseCentroids, CoarseConfig, IvfProbePlan, MultiAssignment, SoarConfig};
138
139// Quantization
140pub use quantization::{DistanceTable, PQCodebook, PQConfig};
141
142// Indexes
143pub use index::{
144    BinaryCoarseQuantizer, BinaryIvfConfig, BinaryIvfIndex, IVFPQConfig, IVFPQIndex, IvfPqQueryPlan,
145};
146
147#[cfg(test)]
148mod decode_tests {
149    use super::decode_ann_bincode_exact;
150    use super::{ANN_DECODE_EXPANSION_FACTOR, ANN_DECODE_FIXED_HEADROOM};
151    use super::{MAX_DENSE_ANN_DECODE_BUDGET, MAX_DENSE_ANN_PAYLOAD_BYTES};
152
153    #[test]
154    fn ann_bincode_decode_is_exact_and_rejects_large_claims_from_tiny_payloads() {
155        let encoded =
156            bincode::serde::encode_to_vec(vec![1u8, 2, 3], bincode::config::standard()).unwrap();
157        assert_eq!(
158            decode_ann_bincode_exact::<Vec<u8>>(&encoded, "test").unwrap(),
159            [1, 2, 3]
160        );
161
162        let mut trailing = encoded;
163        trailing.push(0);
164        assert!(decode_ann_bincode_exact::<Vec<u8>>(&trailing, "test").is_err());
165
166        // A serialized Vec starts with its usize length. This advertises a
167        // 64 MiB allocation from only a few bytes and must hit the 32 MiB tier
168        // before Vec can reserve it.
169        let oversized_claim =
170            bincode::serde::encode_to_vec(64usize * 1024 * 1024, bincode::config::standard())
171                .unwrap();
172        assert!(decode_ann_bincode_exact::<Vec<u8>>(&oversized_claim, "test").is_err());
173    }
174
175    #[cfg(target_pointer_width = "64")]
176    #[test]
177    fn billion_scale_binary_payload_fits_checked_decode_budget() {
178        // Observed production payload: 8,916,876,069 bytes. It must pass the
179        // serialized-input ceiling, while bincode receives enough accounting
180        // budget for bytes consumed plus the decoded SoA allocations.
181        const OBSERVED_BINARY_IVF_BYTES: usize = 8_916_876_069;
182        const { assert!(OBSERVED_BINARY_IVF_BYTES < MAX_DENSE_ANN_PAYLOAD_BYTES) };
183        let budget = OBSERVED_BINARY_IVF_BYTES
184            .checked_mul(ANN_DECODE_EXPANSION_FACTOR)
185            .and_then(|bytes| bytes.checked_add(ANN_DECODE_FIXED_HEADROOM))
186            .unwrap_or(MAX_DENSE_ANN_DECODE_BUDGET)
187            .min(MAX_DENSE_ANN_DECODE_BUDGET);
188        assert_eq!(budget, 32 * 1024 * 1024 * 1024);
189    }
190}