velesdb_core/quantization/
mod.rs1use 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
28pub use binary::BinaryQuantizedVector;
30#[allow(unused_imports)] pub(crate) use pq::distance_pq_l2;
32#[allow(unused_imports)] pub(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#[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
45pub 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
52pub trait QuantizationCodec: Sized {
57 fn to_bytes(&self) -> Vec<u8>;
59
60 fn from_bytes(bytes: &[u8]) -> io::Result<Self>;
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
70#[serde(rename_all = "lowercase")]
71#[non_exhaustive]
72pub enum StorageMode {
73 #[default]
75 Full,
76 SQ8,
78 Binary,
81 ProductQuantization,
83 RaBitQ,
85}
86
87impl StorageMode {
88 #[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 #[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}