iqdb_quantize/code.rs
1//! Owned, immutable code types: [`Sq8Code`], [`BqCode`], and [`PqCode`].
2//!
3//! Each is a thin newtype over the underlying storage with no public
4//! mutators. Codes are produced exclusively by their owning quantizer; a
5//! caller cannot construct a `Sq8Code`, `BqCode`, or `PqCode` outside this
6//! crate, which keeps the contract that a code's contents match the
7//! calibrated quantizer.
8
9/// A scalar-quantized (SQ8) code: one `u8` per dimension of the trained
10/// vector space.
11///
12/// Produced by [`Quantizer::quantize`](crate::Quantizer::quantize) on a
13/// trained [`ScalarQuantizer`](crate::ScalarQuantizer). The byte at
14/// position `i` is the linear `u8` encoding of the original `f32`
15/// component under that dimension's affine calibration; it is not useful
16/// on its own. Decode it with
17/// [`Quantizer::dequantize`](crate::Quantizer::dequantize) or compare it
18/// against a query through
19/// [`Quantizer::distance`](crate::Quantizer::distance).
20///
21/// # Examples
22///
23/// ```
24/// use iqdb_quantize::{Quantizer, ScalarQuantizer};
25///
26/// let mut sq = ScalarQuantizer::new();
27/// sq.train(&[&[0.0_f32, 1.0][..], &[1.0_f32, 0.0][..]]).expect("ok");
28/// let code = sq.quantize(&[0.5_f32, 0.5]).expect("ok");
29/// assert_eq!(code.len(), 2);
30/// ```
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct Sq8Code {
33 pub(crate) bytes: Vec<u8>,
34}
35
36impl Sq8Code {
37 /// Returns the dimension of the encoded vector (one byte per dimension).
38 ///
39 /// # Examples
40 ///
41 /// ```
42 /// use iqdb_quantize::{Quantizer, ScalarQuantizer};
43 ///
44 /// let mut sq = ScalarQuantizer::new();
45 /// sq.train(&[&[0.0_f32, 1.0, 2.0][..]]).expect("ok");
46 /// let code = sq.quantize(&[0.5_f32, 0.5, 0.5]).expect("ok");
47 /// assert_eq!(code.len(), 3);
48 /// ```
49 #[must_use]
50 pub fn len(&self) -> usize {
51 self.bytes.len()
52 }
53
54 /// Returns `true` if the code holds no bytes.
55 ///
56 /// A `Sq8Code` produced by [`Quantizer::quantize`](crate::Quantizer::quantize)
57 /// on a trained [`ScalarQuantizer`](crate::ScalarQuantizer) is never
58 /// empty (empty inputs are rejected at the boundary); this method
59 /// exists for API symmetry with [`Sq8Code::len`].
60 ///
61 /// # Examples
62 ///
63 /// ```
64 /// use iqdb_quantize::{Quantizer, ScalarQuantizer};
65 ///
66 /// let mut sq = ScalarQuantizer::new();
67 /// sq.train(&[&[0.0_f32, 1.0][..]]).expect("ok");
68 /// let code = sq.quantize(&[0.5_f32, 0.5]).expect("ok");
69 /// assert!(!code.is_empty());
70 /// ```
71 #[must_use]
72 pub fn is_empty(&self) -> bool {
73 self.bytes.is_empty()
74 }
75
76 /// Borrow the raw `u8` code bytes.
77 ///
78 /// # Examples
79 ///
80 /// ```
81 /// use iqdb_quantize::{Quantizer, ScalarQuantizer};
82 ///
83 /// let mut sq = ScalarQuantizer::new();
84 /// sq.train(&[&[0.0_f32, 1.0][..]]).expect("ok");
85 /// let code = sq.quantize(&[0.5_f32, 0.5]).expect("ok");
86 /// assert_eq!(code.as_bytes().len(), 2);
87 /// ```
88 #[must_use]
89 pub fn as_bytes(&self) -> &[u8] {
90 &self.bytes
91 }
92}
93
94/// A binary-quantized (BQ) code: one bit per dimension, packed into
95/// `u64` words.
96///
97/// Produced by [`Quantizer::quantize`](crate::Quantizer::quantize) on a
98/// trained [`BinaryQuantizer`](crate::BinaryQuantizer). Each bit is `1`
99/// when the corresponding `f32` component is at or above the trained
100/// per-dimension mean, `0` otherwise. When the dimension is not a
101/// multiple of 64 the trailing word has unused high bits; those bits
102/// are always `0`, so they cannot contribute to Hamming distance.
103///
104/// # Examples
105///
106/// ```
107/// use iqdb_quantize::{BinaryQuantizer, Quantizer};
108///
109/// let mut bq = BinaryQuantizer::new();
110/// bq.train(&[&[0.0_f32, 1.0, 2.0][..], &[2.0_f32, 1.0, 0.0][..]]).expect("ok");
111/// let code = bq.quantize(&[0.5_f32, 1.5, 2.5]).expect("ok");
112/// assert_eq!(code.dim(), 3);
113/// // dim 3 fits in a single u64 word.
114/// assert_eq!(code.as_words().len(), 1);
115/// ```
116#[derive(Debug, Clone, PartialEq, Eq)]
117pub struct BqCode {
118 pub(crate) words: Vec<u64>,
119 pub(crate) dim: usize,
120}
121
122impl BqCode {
123 /// Returns the original vector dimension this code was produced from.
124 ///
125 /// This is the number of meaningful bits in the packed words; the
126 /// trailing word may have unused high bits, which are always zero.
127 ///
128 /// # Examples
129 ///
130 /// ```
131 /// use iqdb_quantize::{BinaryQuantizer, Quantizer};
132 ///
133 /// let mut bq = BinaryQuantizer::new();
134 /// bq.train(&[&[0.0_f32; 65][..], &[1.0_f32; 65][..]]).expect("ok");
135 /// let code = bq.quantize(&[0.5_f32; 65]).expect("ok");
136 /// assert_eq!(code.dim(), 65);
137 /// // 65 bits requires two u64 words.
138 /// assert_eq!(code.as_words().len(), 2);
139 /// ```
140 #[must_use]
141 pub fn dim(&self) -> usize {
142 self.dim
143 }
144
145 /// Returns `true` if the code encodes a zero-dimensional vector.
146 ///
147 /// A `BqCode` produced by [`Quantizer::quantize`](crate::Quantizer::quantize)
148 /// on a trained [`BinaryQuantizer`](crate::BinaryQuantizer) is never
149 /// empty (empty inputs are rejected at the boundary); this method
150 /// exists for API symmetry with [`BqCode::dim`].
151 ///
152 /// # Examples
153 ///
154 /// ```
155 /// use iqdb_quantize::{BinaryQuantizer, Quantizer};
156 ///
157 /// let mut bq = BinaryQuantizer::new();
158 /// bq.train(&[&[0.0_f32, 1.0][..]]).expect("ok");
159 /// let code = bq.quantize(&[0.5_f32, 0.5]).expect("ok");
160 /// assert!(!code.is_empty());
161 /// ```
162 #[must_use]
163 pub fn is_empty(&self) -> bool {
164 self.dim == 0
165 }
166
167 /// Borrow the raw packed `u64` words.
168 ///
169 /// # Examples
170 ///
171 /// ```
172 /// use iqdb_quantize::{BinaryQuantizer, Quantizer};
173 ///
174 /// let mut bq = BinaryQuantizer::new();
175 /// bq.train(&[&[0.0_f32; 4][..], &[1.0_f32; 4][..]]).expect("ok");
176 /// let code = bq.quantize(&[0.5_f32; 4]).expect("ok");
177 /// assert_eq!(code.as_words().len(), 1);
178 /// ```
179 #[must_use]
180 pub fn as_words(&self) -> &[u64] {
181 &self.words
182 }
183}
184
185/// A product-quantized (PQ) code: one `u8` centroid index per subvector.
186///
187/// Produced by [`Quantizer::quantize`](crate::Quantizer::quantize) on a
188/// trained [`ProductQuantizer`](crate::ProductQuantizer). The byte at
189/// position `m` is the index (in `0..n_centroids`, where `n_centroids
190/// <= 256`) of the centroid in subvector codebook `m` that best
191/// approximates the `m`-th subvector of the encoded vector. Decode it
192/// with [`Quantizer::dequantize`](crate::Quantizer::dequantize) (lossy)
193/// or compare it against a query through
194/// [`Quantizer::distance`](crate::Quantizer::distance).
195///
196/// # Examples
197///
198/// ```
199/// use iqdb_quantize::{ProductQuantizer, Quantizer};
200///
201/// let mut pq = ProductQuantizer::with_config(2, 4, 42);
202/// let training: Vec<Vec<f32>> = (0..8)
203/// .map(|i| vec![i as f32, (i * 2) as f32, (i * 3) as f32, (i * 4) as f32])
204/// .collect();
205/// let refs: Vec<&[f32]> = training.iter().map(Vec::as_slice).collect();
206/// pq.train(&refs).expect("training succeeds");
207///
208/// let code = pq.quantize(&[1.0_f32, 2.0, 3.0, 4.0]).expect("quantize");
209/// assert_eq!(code.n_subvectors(), 2);
210/// assert_eq!(code.dim(), 4);
211/// assert_eq!(code.as_bytes().len(), 2);
212/// ```
213#[derive(Debug, Clone, PartialEq, Eq)]
214pub struct PqCode {
215 pub(crate) codes: Vec<u8>,
216 pub(crate) dim: usize,
217 pub(crate) n_subvectors: usize,
218}
219
220impl PqCode {
221 /// Returns the original vector dimension this code was produced from.
222 ///
223 /// # Examples
224 ///
225 /// ```
226 /// use iqdb_quantize::{ProductQuantizer, Quantizer};
227 ///
228 /// let mut pq = ProductQuantizer::with_config(2, 4, 7);
229 /// let training: Vec<Vec<f32>> = (0..8)
230 /// .map(|i| vec![i as f32, (i + 1) as f32, (i + 2) as f32, (i + 3) as f32])
231 /// .collect();
232 /// let refs: Vec<&[f32]> = training.iter().map(Vec::as_slice).collect();
233 /// pq.train(&refs).expect("training succeeds");
234 /// let code = pq.quantize(&[1.0_f32, 2.0, 3.0, 4.0]).expect("quantize");
235 /// assert_eq!(code.dim(), 4);
236 /// ```
237 #[must_use]
238 pub fn dim(&self) -> usize {
239 self.dim
240 }
241
242 /// Returns the number of subvectors `M` this code was produced under.
243 ///
244 /// Equal to `self.as_bytes().len()` and to the
245 /// [`ProductQuantizer::n_subvectors`](crate::ProductQuantizer::n_subvectors)
246 /// the code was produced with.
247 ///
248 /// # Examples
249 ///
250 /// ```
251 /// use iqdb_quantize::{ProductQuantizer, Quantizer};
252 ///
253 /// let mut pq = ProductQuantizer::with_config(2, 4, 7);
254 /// let training: Vec<Vec<f32>> = (0..8)
255 /// .map(|i| vec![i as f32, (i + 1) as f32, (i + 2) as f32, (i + 3) as f32])
256 /// .collect();
257 /// let refs: Vec<&[f32]> = training.iter().map(Vec::as_slice).collect();
258 /// pq.train(&refs).expect("training succeeds");
259 /// let code = pq.quantize(&[1.0_f32, 2.0, 3.0, 4.0]).expect("quantize");
260 /// assert_eq!(code.n_subvectors(), 2);
261 /// ```
262 #[must_use]
263 pub fn n_subvectors(&self) -> usize {
264 self.n_subvectors
265 }
266
267 /// Returns the number of bytes in the code (equal to
268 /// [`PqCode::n_subvectors`]).
269 ///
270 /// # Examples
271 ///
272 /// ```
273 /// use iqdb_quantize::{ProductQuantizer, Quantizer};
274 ///
275 /// let mut pq = ProductQuantizer::with_config(2, 4, 7);
276 /// let training: Vec<Vec<f32>> = (0..8)
277 /// .map(|i| vec![i as f32, (i + 1) as f32, (i + 2) as f32, (i + 3) as f32])
278 /// .collect();
279 /// let refs: Vec<&[f32]> = training.iter().map(Vec::as_slice).collect();
280 /// pq.train(&refs).expect("training succeeds");
281 /// let code = pq.quantize(&[1.0_f32, 2.0, 3.0, 4.0]).expect("quantize");
282 /// assert_eq!(code.len(), 2);
283 /// ```
284 #[must_use]
285 pub fn len(&self) -> usize {
286 self.codes.len()
287 }
288
289 /// Returns `true` if the code holds no centroid indices.
290 ///
291 /// A `PqCode` produced by [`Quantizer::quantize`](crate::Quantizer::quantize)
292 /// on a trained [`ProductQuantizer`](crate::ProductQuantizer) is never
293 /// empty (empty inputs and `n_subvectors == 0` are rejected at the
294 /// boundary); this method exists for API symmetry with [`PqCode::len`].
295 ///
296 /// # Examples
297 ///
298 /// ```
299 /// use iqdb_quantize::{ProductQuantizer, Quantizer};
300 ///
301 /// let mut pq = ProductQuantizer::with_config(2, 4, 7);
302 /// let training: Vec<Vec<f32>> = (0..8)
303 /// .map(|i| vec![i as f32, (i + 1) as f32, (i + 2) as f32, (i + 3) as f32])
304 /// .collect();
305 /// let refs: Vec<&[f32]> = training.iter().map(Vec::as_slice).collect();
306 /// pq.train(&refs).expect("training succeeds");
307 /// let code = pq.quantize(&[1.0_f32, 2.0, 3.0, 4.0]).expect("quantize");
308 /// assert!(!code.is_empty());
309 /// ```
310 #[must_use]
311 pub fn is_empty(&self) -> bool {
312 self.codes.is_empty()
313 }
314
315 /// Borrow the raw centroid-index bytes.
316 ///
317 /// # Examples
318 ///
319 /// ```
320 /// use iqdb_quantize::{ProductQuantizer, Quantizer};
321 ///
322 /// let mut pq = ProductQuantizer::with_config(2, 4, 7);
323 /// let training: Vec<Vec<f32>> = (0..8)
324 /// .map(|i| vec![i as f32, (i + 1) as f32, (i + 2) as f32, (i + 3) as f32])
325 /// .collect();
326 /// let refs: Vec<&[f32]> = training.iter().map(Vec::as_slice).collect();
327 /// pq.train(&refs).expect("training succeeds");
328 /// let code = pq.quantize(&[1.0_f32, 2.0, 3.0, 4.0]).expect("quantize");
329 /// assert_eq!(code.as_bytes().len(), 2);
330 /// ```
331 #[must_use]
332 pub fn as_bytes(&self) -> &[u8] {
333 &self.codes
334 }
335}