reddb_server/storage/engine/turboquant/
codec.rs

1//! TurboQuant per-vector encode + per-collection score driver.
2//!
3//! `EncodedVector` is a [`BlockedCodeStorage`] handle — `(block_idx,
4//! lane)` — per ADR 0024. The per-vector `Vec<u8>` of the rejected
5//! layout is gone; encoded codes live in collection-owned blocked
6//! buffers and are scored block-at-a-time through a
7//! [`PerBlockScorer`].
8
9use super::codebook::Codebook;
10use super::rotation::RotationMatrix;
11use super::scoring::{select_scorer, QueryLut};
12use super::storage::{BlockHandle, BlockedCodeStorage, BLOCK_LANES};
13use crate::storage::engine::distance;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct EncodedVector {
17    pub block_idx: u32,
18    pub lane: u8,
19}
20
21impl From<BlockHandle> for EncodedVector {
22    fn from(h: BlockHandle) -> Self {
23        Self {
24            block_idx: h.block_idx,
25            lane: h.lane,
26        }
27    }
28}
29
30#[derive(Debug, Clone)]
31pub struct Codec {
32    dim: usize,
33    rotation: RotationMatrix,
34    codebook: Codebook,
35}
36
37impl Codec {
38    pub fn new(dim: usize, seed: u64) -> Self {
39        Self {
40            dim,
41            rotation: RotationMatrix::new(dim, seed),
42            codebook: Codebook::for_dim_bits(dim, 4),
43        }
44    }
45
46    pub fn dim(&self) -> usize {
47        self.dim
48    }
49
50    pub fn n_byte_groups(&self) -> usize {
51        self.dim.div_ceil(2)
52    }
53
54    /// Encode `vector` into the packed 4-bit code groups +
55    /// pre-rotation L2 norm without appending to any storage. This is
56    /// the persistence-layer entry point used by
57    /// `TurboExtent::append` in the durable write path (#693). The
58    /// `BlockedCodeStorage::append` view of the same bytes is what
59    /// the in-memory scorer reads.
60    pub fn encode_packed(&self, vector: &[f32]) -> (Vec<u8>, f32) {
61        assert_eq!(vector.len(), self.dim, "encode dimension must match codec",);
62        let scale = distance::l2_norm(vector);
63        let normalized = if scale > 0.0 {
64            vector.iter().map(|v| *v / scale).collect::<Vec<_>>()
65        } else {
66            vec![0.0; vector.len()]
67        };
68        let rotated = self.rotation.rotate(&normalized);
69        let mut packed = vec![0u8; self.n_byte_groups()];
70        for (i, pair) in rotated.chunks(2).enumerate() {
71            let lo = self.codebook.quantize(pair[0]) & 0x0f;
72            let hi = pair
73                .get(1)
74                .map(|value| self.codebook.quantize(*value) & 0x0f)
75                .unwrap_or(0);
76            packed[i] = lo | (hi << 4);
77        }
78        (packed, scale)
79    }
80
81    /// Encode `vector` and append it to `storage`. Returns a handle the
82    /// caller can pass back to [`Self::score_many`] (indirectly through
83    /// the storage itself).
84    pub fn encode_into(&self, storage: &mut BlockedCodeStorage, vector: &[f32]) -> EncodedVector {
85        let (packed, scale) = self.encode_packed(vector);
86        storage.append(&packed, scale).into()
87    }
88
89    pub fn scalar_score(
90        &self,
91        query: &[f32],
92        candidate: &[f32],
93        metric: distance::DistanceMetric,
94    ) -> f32 {
95        let raw = distance::distance(query, candidate, metric);
96        match metric {
97            distance::DistanceMetric::Cosine => 1.0 - raw,
98            distance::DistanceMetric::InnerProduct | distance::DistanceMetric::L2 => -raw,
99        }
100    }
101
102    /// Score `query` against every vector currently in `storage`.
103    /// Returns a `Vec<f32>` of length `storage.n_blocks() * BLOCK_LANES`
104    /// indexed by `block_idx * BLOCK_LANES + lane`; entries in unused
105    /// trailing-block lanes are filled with `f32::NEG_INFINITY`.
106    pub fn score_many(
107        &self,
108        query: &[f32],
109        storage: &BlockedCodeStorage,
110        metric: distance::DistanceMetric,
111    ) -> Vec<f32> {
112        match self.score_many_checked(query, storage, metric, |_| {
113            Ok::<(), std::convert::Infallible>(())
114        }) {
115            Ok(scores) => scores,
116            Err(never) => match never {},
117        }
118    }
119
120    /// Poll before each packed scoring block; failed checks never return partial scores.
121    pub(crate) fn score_many_checked<E>(
122        &self,
123        query: &[f32],
124        storage: &BlockedCodeStorage,
125        metric: distance::DistanceMetric,
126        mut check: impl FnMut(u64) -> Result<(), E>,
127    ) -> Result<Vec<f32>, E> {
128        assert_eq!(query.len(), self.dim, "Vector dimensions must match");
129
130        check(0)?;
131        let n_blocks = storage.n_blocks();
132        let mut scores = vec![f32::NEG_INFINITY; n_blocks * BLOCK_LANES];
133
134        let query_norm = distance::l2_norm(query);
135        if query_norm == 0.0 {
136            for b in 0..n_blocks {
137                let filled = storage.block_lanes_filled(b);
138                check(u64::try_from(filled).expect("block holds at most 32 lanes"))?;
139                for lane in 0..filled {
140                    let s = storage.lane_scale(b, lane);
141                    scores[b * BLOCK_LANES + lane] = match metric {
142                        distance::DistanceMetric::L2 => -(s * s),
143                        _ => 0.0,
144                    };
145                }
146            }
147            check(0)?;
148            return Ok(scores);
149        }
150
151        let normalized: Vec<f32> = query.iter().map(|v| *v / query_norm).collect();
152        let rotated = self.rotation.rotate(&normalized);
153        let lut = QueryLut::build(&rotated, self.codebook.centroids());
154        let scorer = select_scorer();
155
156        let n_byte_groups = storage.n_byte_groups();
157        let mut block_scores = [0.0f32; BLOCK_LANES];
158        for b in 0..n_blocks {
159            let filled = storage.block_lanes_filled(b);
160            check(u64::try_from(filled).expect("block holds at most 32 lanes"))?;
161            scorer.score_block(
162                &lut,
163                storage.block_codes(b),
164                n_byte_groups,
165                filled,
166                &mut block_scores,
167            );
168            for lane in 0..filled {
169                let unit_dot = block_scores[lane];
170                let lane_scale = storage.lane_scale(b, lane);
171                let raw_dot = unit_dot * query_norm * lane_scale;
172                let metric_score = match metric {
173                    distance::DistanceMetric::Cosine => {
174                        if lane_scale > 0.0 {
175                            unit_dot
176                        } else {
177                            0.0
178                        }
179                    }
180                    distance::DistanceMetric::InnerProduct => raw_dot,
181                    distance::DistanceMetric::L2 => {
182                        -(query_norm * query_norm + lane_scale * lane_scale - 2.0 * raw_dot)
183                    }
184                };
185                scores[b * BLOCK_LANES + lane] = metric_score;
186            }
187        }
188        check(0)?;
189        Ok(scores)
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn checked_scoring_stops_between_blocks_for_zero_and_nonzero_queries() {
199        let codec = Codec::new(2, 11);
200        let mut storage = BlockedCodeStorage::new(codec.n_byte_groups());
201        for _ in 0..70 {
202            codec.encode_into(&mut storage, &[1.0, 0.0]);
203        }
204        for query in [[0.0, 0.0], [1.0, 0.0]] {
205            let mut charged = 0;
206            let result = codec.score_many_checked(
207                &query,
208                &storage,
209                distance::DistanceMetric::Cosine,
210                |work| {
211                    if charged + work > 32 {
212                        return Err("exhausted");
213                    }
214                    charged += work;
215                    Ok(())
216                },
217            );
218            assert_eq!(
219                result.expect_err("second block exceeds budget"),
220                "exhausted"
221            );
222            assert_eq!(charged, 32, "only the first block was admitted");
223        }
224    }
225
226    #[test]
227    fn encode_is_bit_exact_for_frozen_vectors() {
228        let codec = Codec::new(4, 7);
229        let mut a = BlockedCodeStorage::new(codec.n_byte_groups());
230        let mut b = BlockedCodeStorage::new(codec.n_byte_groups());
231        let ha = codec.encode_into(&mut a, &[1.0, 0.0, -1.0, 0.5]);
232        let hb = codec.encode_into(&mut b, &[1.0, 0.0, -1.0, 0.5]);
233        assert_eq!(ha, hb);
234        assert_eq!(
235            a.decode_lane(ha.block_idx as usize, ha.lane as usize),
236            b.decode_lane(hb.block_idx as usize, hb.lane as usize),
237        );
238    }
239
240    #[test]
241    fn score_many_layout_indexes_by_block_lane() {
242        let codec = Codec::new(2, 11);
243        let mut storage = BlockedCodeStorage::new(codec.n_byte_groups());
244        let h0 = codec.encode_into(&mut storage, &[1.0, 0.0]);
245        let h1 = codec.encode_into(&mut storage, &[0.0, 1.0]);
246        let scores = codec.score_many(&[1.0, 0.0], &storage, distance::DistanceMetric::Cosine);
247        let s0 = scores[h0.block_idx as usize * BLOCK_LANES + h0.lane as usize];
248        let s1 = scores[h1.block_idx as usize * BLOCK_LANES + h1.lane as usize];
249        // Self-similarity dominates the orthogonal entry under cosine.
250        assert!(
251            s0 >= s1,
252            "expected vector aligned with query to outrank orthogonal one: s0={s0}, s1={s1}",
253        );
254    }
255}
reddb_server/storage/engine/turboquant/codec.rs

reddb_server/storage/engine/turboquant/
codec.rs
