Skip to main content

velesdb_core/quantization/
mod.rs

1//! Scalar Quantization (SQ8) and Binary Quantization for memory-efficient vector storage.
2//!
3//! This module implements quantization strategies to reduce memory usage:
4//!
5//! ## Benefits
6//!
7//! | Metric | f32 | SQ8 | Binary |
8//! |--------|-----|-----|--------|
9//! | RAM/vector (768d) | 3 KB | 770 bytes | 96 bytes |
10//! | Cache efficiency | Baseline | ~4x better | ~32x better |
11//! | Recall loss | 0% | ~0.5-1% | ~5-10% |
12
13use std::io;
14
15use serde::{Deserialize, Serialize};
16
17mod binary;
18pub(crate) mod codec_helpers;
19mod pq;
20pub(crate) mod pq_kmeans;
21pub(crate) mod pq_opq;
22#[cfg(feature = "persistence")]
23mod pq_persistence;
24mod rabitq;
25pub(crate) mod rabitq_store;
26mod scalar;
27
28// Re-export binary quantization
29pub use binary::BinaryQuantizedVector;
30#[allow(unused_imports)] // Called from vector.rs search path (persistence-gated).
31pub(crate) use pq::distance_pq_l2;
32#[allow(unused_imports)] // Called from vector.rs search path (persistence-gated).
33pub(crate) use pq::pq_adc_batch_rescore;
34pub use pq::{PQCodebook, PQVector, ProductQuantizer};
35#[cfg(feature = "persistence")]
36pub use pq_opq::train_opq;
37
38// Re-export RaBitQ quantization
39#[cfg(feature = "persistence")]
40pub(crate) use rabitq::PreparedQuery;
41pub use rabitq::{RaBitQCorrection, RaBitQIndex, RaBitQVector};
42#[cfg(feature = "persistence")]
43pub(crate) use rabitq_store::RaBitQVectorStore;
44
45// Re-export scalar quantization
46pub use scalar::{
47    cosine_similarity_quantized, cosine_similarity_quantized_simd, dot_product_quantized,
48    dot_product_quantized_simd, euclidean_squared_quantized, euclidean_squared_quantized_simd,
49    QuantizedVector,
50};
51
52/// Trait for serializing and deserializing quantized vectors to/from bytes.
53///
54/// Provides a uniform interface for byte-level serialization across
55/// different quantization strategies (SQ8, Binary).
56pub trait QuantizationCodec: Sized {
57    /// Serializes the quantized vector to a byte representation.
58    fn to_bytes(&self) -> Vec<u8>;
59
60    /// Deserializes a quantized vector from bytes.
61    ///
62    /// # Errors
63    ///
64    /// Returns an error if the byte slice is too short or contains invalid data.
65    fn from_bytes(bytes: &[u8]) -> io::Result<Self>;
66}
67
68/// Storage mode for vectors.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
70#[serde(rename_all = "lowercase")]
71#[non_exhaustive]
72pub enum StorageMode {
73    /// Full precision f32 storage (default).
74    #[default]
75    Full,
76    /// 8-bit scalar quantization for 4x memory reduction.
77    SQ8,
78    /// 1-bit binary quantization for 32x memory reduction.
79    /// Best for edge/IoT devices with limited RAM.
80    Binary,
81    /// Product Quantization (PQ) for aggressive lossy compression (8x-16x typical).
82    ProductQuantization,
83    /// `RaBitQ` binary quantization for 32x compression with scalar correction.
84    RaBitQ,
85}
86
87impl StorageMode {
88    /// Returns the canonical lowercase name for this storage mode.
89    ///
90    /// This is the single source of truth for string representations,
91    /// used by [`std::fmt::Display`], [`std::str::FromStr`], and downstream crates.
92    #[must_use]
93    pub const fn canonical_name(self) -> &'static str {
94        match self {
95            Self::Full => "full",
96            Self::SQ8 => "sq8",
97            Self::Binary => "binary",
98            Self::ProductQuantization => "pq",
99            Self::RaBitQ => "rabitq",
100        }
101    }
102
103    /// Parses a storage mode string with alias support.
104    ///
105    /// Accepted aliases (case-insensitive):
106    /// - `full`, `f32` -> `Full`
107    /// - `sq8`, `int8` -> `SQ8`
108    /// - `binary`, `bit` -> `Binary`
109    /// - `pq`, `product_quantization` -> `ProductQuantization`
110    /// - `rabitq` -> `RaBitQ`
111    ///
112    /// # Examples
113    ///
114    /// ```
115    /// use velesdb_core::StorageMode;
116    ///
117    /// assert_eq!(StorageMode::parse_alias("sq8"), Some(StorageMode::SQ8));
118    /// assert_eq!(StorageMode::parse_alias("INT8"), Some(StorageMode::SQ8));
119    /// assert_eq!(StorageMode::parse_alias("unknown"), None);
120    /// ```
121    #[must_use]
122    pub fn parse_alias(value: &str) -> Option<Self> {
123        match value.trim().to_lowercase().as_str() {
124            "full" | "f32" => Some(Self::Full),
125            "sq8" | "int8" => Some(Self::SQ8),
126            "binary" | "bit" => Some(Self::Binary),
127            "pq" | "product_quantization" => Some(Self::ProductQuantization),
128            "rabitq" => Some(Self::RaBitQ),
129            _ => None,
130        }
131    }
132}
133
134impl std::fmt::Display for StorageMode {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        f.write_str(self.canonical_name())
137    }
138}
139
140impl std::str::FromStr for StorageMode {
141    type Err = String;
142
143    fn from_str(s: &str) -> Result<Self, Self::Err> {
144        Self::parse_alias(s).ok_or_else(|| {
145            format!(
146                "Unknown storage mode '{s}'. Valid options: full, f32, sq8, int8, binary, bit, pq, product_quantization, rabitq"
147            )
148        })
149    }
150}
151
152#[cfg(test)]
153mod storage_mode_parsing_tests {
154    use super::StorageMode;
155
156    #[test]
157    fn test_parse_all_canonical_names() {
158        assert_eq!("full".parse::<StorageMode>().unwrap(), StorageMode::Full);
159        assert_eq!("sq8".parse::<StorageMode>().unwrap(), StorageMode::SQ8);
160        assert_eq!(
161            "binary".parse::<StorageMode>().unwrap(),
162            StorageMode::Binary
163        );
164        assert_eq!(
165            "pq".parse::<StorageMode>().unwrap(),
166            StorageMode::ProductQuantization
167        );
168        assert_eq!(
169            "rabitq".parse::<StorageMode>().unwrap(),
170            StorageMode::RaBitQ
171        );
172    }
173
174    #[test]
175    fn test_parse_aliases() {
176        assert_eq!("f32".parse::<StorageMode>().unwrap(), StorageMode::Full);
177        assert_eq!("int8".parse::<StorageMode>().unwrap(), StorageMode::SQ8);
178        assert_eq!("bit".parse::<StorageMode>().unwrap(), StorageMode::Binary);
179        assert_eq!(
180            "product_quantization".parse::<StorageMode>().unwrap(),
181            StorageMode::ProductQuantization
182        );
183    }
184
185    #[test]
186    fn test_parse_case_insensitive() {
187        assert_eq!("SQ8".parse::<StorageMode>().unwrap(), StorageMode::SQ8);
188        assert_eq!("FULL".parse::<StorageMode>().unwrap(), StorageMode::Full);
189        assert_eq!(
190            "RaBitQ".parse::<StorageMode>().unwrap(),
191            StorageMode::RaBitQ
192        );
193    }
194
195    #[test]
196    fn test_parse_unknown_returns_error() {
197        assert!("unknown".parse::<StorageMode>().is_err());
198        assert!("".parse::<StorageMode>().is_err());
199    }
200
201    #[test]
202    fn test_canonical_name_roundtrip() {
203        for mode in [
204            StorageMode::Full,
205            StorageMode::SQ8,
206            StorageMode::Binary,
207            StorageMode::ProductQuantization,
208            StorageMode::RaBitQ,
209        ] {
210            let name = mode.canonical_name();
211            assert_eq!(name.parse::<StorageMode>().unwrap(), mode);
212        }
213    }
214
215    #[test]
216    fn test_display_uses_canonical_name() {
217        assert_eq!(format!("{}", StorageMode::Full), "full");
218        assert_eq!(format!("{}", StorageMode::SQ8), "sq8");
219        assert_eq!(format!("{}", StorageMode::Binary), "binary");
220        assert_eq!(format!("{}", StorageMode::ProductQuantization), "pq");
221        assert_eq!(format!("{}", StorageMode::RaBitQ), "rabitq");
222    }
223}