Skip to main content

nodedb_vector/quantize/
binary.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Binary Quantization (BQ): sign-bit encoding + Hamming distance.
4//!
5//! Each dimension is encoded as a single bit: 1 if positive, 0 if negative.
6//! 32x compression (D/8 bytes vs 4D bytes for FP32). Best used as a coarse
7//! pre-filter: compute Hamming distance to quickly eliminate far candidates
8//! before computing exact distances on survivors.
9//!
10//! Recall loss: 5-10% as a standalone index, but combined with a reranking
11//! step on top-K×10 candidates, effective recall loss is 1-3%.
12
13/// Encode an FP32 vector as binary: one bit per dimension.
14///
15/// Bit layout: bit 0 of byte 0 = dimension 0, bit 1 = dimension 1, etc.
16/// `output.len() = ceil(dim / 8)`.
17pub fn encode(vector: &[f32]) -> Vec<u8> {
18    let num_bytes = vector.len().div_ceil(8);
19    let mut bits = vec![0u8; num_bytes];
20    for (i, &val) in vector.iter().enumerate() {
21        if val > 0.0 {
22            bits[i / 8] |= 1 << (i % 8);
23        }
24    }
25    bits
26}
27
28/// Batch encode: encode all vectors into contiguous binary representation.
29///
30/// Returns `ceil(dim/8) * N` bytes.
31pub fn encode_batch(vectors: &[&[f32]], dim: usize) -> Vec<u8> {
32    let bytes_per = dim.div_ceil(8);
33    // no-governor: cold batch binary encode; governed at segment build call site
34    let mut out = Vec::with_capacity(bytes_per * vectors.len());
35    for v in vectors {
36        out.extend(encode(v));
37    }
38    out
39}
40
41/// Hamming distance between two binary-encoded vectors.
42///
43/// Counts the number of differing bits using `count_ones()` (hardware
44/// POPCNT on x86_64).
45#[inline]
46pub fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
47    debug_assert_eq!(a.len(), b.len());
48    let mut dist = 0u32;
49    for i in 0..a.len() {
50        dist += (a[i] ^ b[i]).count_ones();
51    }
52    dist
53}
54
55/// Hamming distance operating on u64 chunks for better throughput.
56///
57/// Processes 8 bytes at a time using u64 POPCNT. Falls back to byte-level
58/// for the remainder. ~4x faster than byte-level for vectors ≥64 dims.
59#[inline]
60pub fn hamming_distance_fast(a: &[u8], b: &[u8]) -> u32 {
61    debug_assert_eq!(a.len(), b.len());
62    let mut dist = 0u32;
63    let chunks = a.len() / 8;
64    let remainder = a.len() % 8;
65
66    // Process u64 chunks (slice is guaranteed to be 8 bytes by loop bounds).
67    for i in 0..chunks {
68        let offset = i * 8;
69        let mut a_buf = [0u8; 8];
70        let mut b_buf = [0u8; 8];
71        a_buf.copy_from_slice(&a[offset..offset + 8]);
72        b_buf.copy_from_slice(&b[offset..offset + 8]);
73        dist += (u64::from_le_bytes(a_buf) ^ u64::from_le_bytes(b_buf)).count_ones();
74    }
75
76    // Process remaining bytes.
77    let start = chunks * 8;
78    for i in 0..remainder {
79        dist += (a[start + i] ^ b[start + i]).count_ones();
80    }
81
82    dist
83}
84
85/// Binary vector size in bytes for a given dimensionality.
86pub fn binary_size(dim: usize) -> usize {
87    dim.div_ceil(8)
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn encode_positive_negative() {
96        let v = [1.0, -1.0, 1.0, -1.0, 0.0, 1.0, -0.5, 0.5];
97        let bits = encode(&v);
98        assert_eq!(bits.len(), 1);
99        // bits[0]: dim0=1, dim1=0, dim2=1, dim3=0, dim4=0, dim5=1, dim6=0, dim7=1
100        //        = 0b10100101 = 0xA5
101        assert_eq!(bits[0], 0b10100101);
102    }
103
104    #[test]
105    fn hamming_identical_is_zero() {
106        let v = [1.0, -1.0, 1.0, 0.5];
107        let a = encode(&v);
108        let b = encode(&v);
109        assert_eq!(hamming_distance(&a, &b), 0);
110    }
111
112    #[test]
113    fn hamming_opposite_is_dim() {
114        let a_vec = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
115        let b_vec = [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0];
116        let a = encode(&a_vec);
117        let b = encode(&b_vec);
118        assert_eq!(hamming_distance(&a, &b), 8);
119    }
120
121    #[test]
122    fn hamming_fast_matches_simple() {
123        // 128 dimensions = 16 bytes = 2 u64 chunks.
124        let a_vec: Vec<f32> = (0..128)
125            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
126            .collect();
127        let b_vec: Vec<f32> = (0..128)
128            .map(|i| if i % 5 == 0 { 1.0 } else { -1.0 })
129            .collect();
130        let a = encode(&a_vec);
131        let b = encode(&b_vec);
132
133        let slow = hamming_distance(&a, &b);
134        let fast = hamming_distance_fast(&a, &b);
135        assert_eq!(slow, fast);
136    }
137
138    #[test]
139    fn high_dimensional_encoding() {
140        // 768 dimensions = 96 bytes.
141        let v: Vec<f32> = (0..768).map(|i| (i as f32).sin()).collect();
142        let bits = encode(&v);
143        assert_eq!(bits.len(), 96);
144    }
145
146    #[test]
147    fn batch_encode_layout() {
148        let v1 = [1.0f32, -1.0, 1.0, -1.0];
149        let v2 = [-1.0f32, 1.0, -1.0, 1.0];
150        let batch = encode_batch(&[&v1, &v2], 4);
151        assert_eq!(batch.len(), 2); // 2 vectors × 1 byte each
152        assert_eq!(batch[0], encode(&v1)[0]);
153        assert_eq!(batch[1], encode(&v2)[0]);
154    }
155}