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

velesdb_core/quantization/
mod.rs
