qql_embed/
sparse.rs

1use std::collections::HashMap;
2use std::hash::{BuildHasherDefault, Hasher};
3
4/// Sparse embedding (indices + values). Transport-neutral — not a protobuf type.
5#[derive(Debug, Clone, PartialEq, Default)]
6pub struct SparseVector {
7    pub indices: Vec<u32>,
8    pub values: Vec<f32>,
9}
10
11const OFFSET32: u32 = 2166136261;
12const PRIME32: u32 = 16777619;
13
14/// Fast Identity Hasher for u32 keys (avoids SipHash overhead).
15#[derive(Default)]
16pub struct IdentityHasher(u64);
17
18impl Hasher for IdentityHasher {
19    #[inline]
20    fn finish(&self) -> u64 {
21        self.0
22    }
23
24    #[inline]
25    fn write(&mut self, bytes: &[u8]) {
26        for &b in bytes {
27            self.0 = (self.0 << 8) | (b as u64);
28        }
29    }
30
31    #[inline]
32    fn write_u32(&mut self, i: u32) {
33        self.0 = i as u64;
34    }
35}
36
37type FastMap<K, V> = HashMap<K, V, BuildHasherDefault<IdentityHasher>>;
38
39pub fn hash_token(token: &str) -> u32 {
40    let mut h: u32 = OFFSET32;
41
42    let l = token.len() as u64;
43    for i in 0..8 {
44        h ^= (l >> (i * 8)) as u32 & 0xff;
45        h = h.wrapping_mul(PRIME32);
46    }
47
48    for &b in token.as_bytes() {
49        h ^= b as u32;
50        h = h.wrapping_mul(PRIME32);
51    }
52
53    h
54}
55
56/// Fast on-the-fly hash computation for ASCII slices with lowercase conversion.
57#[inline]
58fn hash_token_bytes(bytes: &[u8]) -> u32 {
59    let mut h: u32 = OFFSET32;
60
61    let l = bytes.len() as u64;
62    for i in 0..8 {
63        h ^= (l >> (i * 8)) as u32 & 0xff;
64        h = h.wrapping_mul(PRIME32);
65    }
66
67    for &b in bytes {
68        let lower_b = if b.is_ascii_uppercase() { b + 32 } else { b };
69        h ^= lower_b as u32;
70        h = h.wrapping_mul(PRIME32);
71    }
72
73    h
74}
75
76pub fn tokenize(text: &str) -> Vec<String> {
77    if is_ascii(text) {
78        tokenize_ascii(text)
79    } else {
80        tokenize_unicode(text)
81    }
82}
83
84fn is_ascii(s: &str) -> bool {
85    s.bytes().all(|b| b < 0x80)
86}
87
88fn is_token_byte(ch: u8) -> bool {
89    ch.is_ascii_lowercase()
90        || ch.is_ascii_uppercase()
91        || ch.is_ascii_digit()
92        || ch == b'_'
93        || ch == b'-'
94}
95
96fn is_token_rune(r: char) -> bool {
97    r.is_alphabetic() || r.is_ascii_digit() || r == '_' || r == '-'
98}
99
100fn to_lower_ascii(s: &str) -> String {
101    let bytes = s.as_bytes();
102    for (i, &b) in bytes.iter().enumerate() {
103        if b.is_ascii_uppercase() {
104            let mut buf = Vec::with_capacity(bytes.len());
105            buf.extend_from_slice(&bytes[..i]);
106            for &b in &bytes[i..] {
107                buf.push(if b.is_ascii_uppercase() { b + 32 } else { b });
108            }
109            return String::from_utf8(buf).unwrap_or_else(|_| s.to_ascii_lowercase());
110        }
111    }
112    s.to_string()
113}
114
115fn unicode_to_lower(s: &str) -> String {
116    let mut buf = String::with_capacity(s.len());
117    for c in s.chars() {
118        for lc in c.to_lowercase() {
119            buf.push(lc);
120        }
121    }
122    buf
123}
124
125fn maybe_token(s: &str) -> Option<String> {
126    if s.len() >= 2 {
127        return Some(s.to_string());
128    }
129    if s.len() == 1 {
130        let b = s.as_bytes()[0];
131        if b == b'c' {
132            return Some(s.to_string());
133        }
134    }
135    None
136}
137
138fn append_tokens(tokens: &mut Vec<String>, raw: &str) {
139    let has_hyphen = raw.contains('-');
140    if !has_hyphen {
141        if let Some(tok) = maybe_token(raw) {
142            tokens.push(tok);
143        }
144        return;
145    }
146
147    let mut start: Option<usize> = None;
148    for (i, ch) in raw.char_indices() {
149        if ch == '-' {
150            if let Some(s) = start {
151                if let Some(tok) = maybe_token(&raw[s..i]) {
152                    tokens.push(tok);
153                }
154                start = None;
155            }
156        } else {
157            if start.is_none() {
158                start = Some(i);
159            }
160        }
161    }
162    if let Some(s) = start {
163        if let Some(tok) = maybe_token(&raw[s..]) {
164            tokens.push(tok);
165        }
166    }
167}
168
169fn tokenize_ascii(text: &str) -> Vec<String> {
170    let mut tokens = Vec::new();
171    let mut start: Option<usize> = None;
172    let bytes = text.as_bytes();
173
174    for i in 0..bytes.len() {
175        let ch = bytes[i];
176        if is_token_byte(ch) {
177            if start.is_none() {
178                start = Some(i);
179            }
180        } else {
181            if let Some(s) = start {
182                append_tokens(&mut tokens, &to_lower_ascii(&text[s..i]));
183                start = None;
184            }
185        }
186    }
187    if let Some(s) = start {
188        append_tokens(&mut tokens, &to_lower_ascii(&text[s..]));
189    }
190
191    tokens
192}
193
194fn tokenize_unicode(text: &str) -> Vec<String> {
195    let lower = unicode_to_lower(text);
196    let mut tokens = Vec::new();
197    let mut start: Option<usize> = None;
198
199    for (i, ch) in lower.char_indices() {
200        if is_token_rune(ch) {
201            if start.is_none() {
202                start = Some(i);
203            }
204        } else {
205            if let Some(s) = start {
206                append_tokens(&mut tokens, &lower[s..i]);
207                start = None;
208            }
209        }
210    }
211    if let Some(s) = start {
212        append_tokens(&mut tokens, &lower[s..]);
213    }
214
215    tokens
216}
217
218#[inline]
219fn is_valid_token_len(len: usize, first_byte: u8) -> bool {
220    len >= 2 || (len == 1 && (first_byte == b'c' || first_byte == b'C'))
221}
222
223/// Zero-allocation fast pass token hashing for ASCII strings
224fn hash_tokens_ascii_fast(text: &str) -> (Vec<u32>, usize) {
225    let bytes = text.as_bytes();
226    let mut raw_hashes = Vec::with_capacity(bytes.len() / 4 + 1);
227    let mut start: Option<usize> = None;
228
229    for i in 0..bytes.len() {
230        let ch = bytes[i];
231        if is_token_byte(ch) {
232            if start.is_none() {
233                start = Some(i);
234            }
235        } else if let Some(s) = start {
236            let slice = &bytes[s..i];
237            let len = slice.len();
238            let first_b = slice[0];
239
240            if !slice.contains(&b'-') {
241                if is_valid_token_len(len, first_b) {
242                    raw_hashes.push(hash_token_bytes(slice));
243                }
244            } else {
245                let mut sub_start: Option<usize> = None;
246                for idx in 0..slice.len() {
247                    if slice[idx] == b'-' {
248                        if let Some(ss) = sub_start {
249                            let sub_slice = &slice[ss..idx];
250                            if is_valid_token_len(sub_slice.len(), sub_slice[0]) {
251                                raw_hashes.push(hash_token_bytes(sub_slice));
252                            }
253                            sub_start = None;
254                        }
255                    } else if sub_start.is_none() {
256                        sub_start = Some(idx);
257                    }
258                }
259                if let Some(ss) = sub_start {
260                    let sub_slice = &slice[ss..];
261                    if is_valid_token_len(sub_slice.len(), sub_slice[0]) {
262                        raw_hashes.push(hash_token_bytes(sub_slice));
263                    }
264                }
265            }
266            start = None;
267        }
268    }
269
270    if let Some(s) = start {
271        let slice = &bytes[s..];
272        let len = slice.len();
273        let first_b = slice[0];
274
275        if !slice.contains(&b'-') {
276            if is_valid_token_len(len, first_b) {
277                raw_hashes.push(hash_token_bytes(slice));
278            }
279        } else {
280            let mut sub_start: Option<usize> = None;
281            for idx in 0..slice.len() {
282                if slice[idx] == b'-' {
283                    if let Some(ss) = sub_start {
284                        let sub_slice = &slice[ss..idx];
285                        if is_valid_token_len(sub_slice.len(), sub_slice[0]) {
286                            raw_hashes.push(hash_token_bytes(sub_slice));
287                        }
288                        sub_start = None;
289                    }
290                } else if sub_start.is_none() {
291                    sub_start = Some(idx);
292                }
293            }
294            if let Some(ss) = sub_start {
295                let sub_slice = &slice[ss..];
296                if is_valid_token_len(sub_slice.len(), sub_slice[0]) {
297                    raw_hashes.push(hash_token_bytes(sub_slice));
298                }
299            }
300        }
301    }
302
303    let token_count = raw_hashes.len();
304    (raw_hashes, token_count)
305}
306
307pub fn build_query(text: &str) -> SparseVector {
308    let (counts, _total_tokens) = if is_ascii(text) {
309        let (hashes, total_tokens) = hash_tokens_ascii_fast(text);
310        let mut counts: FastMap<u32, f32> = FastMap::with_capacity_and_hasher(
311            hashes.len(),
312            BuildHasherDefault::<IdentityHasher>::default(),
313        );
314        for h in hashes {
315            *counts.entry(h).or_insert(0.0) += 1.0;
316        }
317        (counts, total_tokens)
318    } else {
319        let tokens = tokenize(text);
320        let mut counts: FastMap<u32, f32> = FastMap::with_capacity_and_hasher(
321            tokens.len(),
322            BuildHasherDefault::<IdentityHasher>::default(),
323        );
324        let total_tokens = tokens.len();
325        for token in &tokens {
326            *counts.entry(hash_token(token)).or_insert(0.0) += 1.0;
327        }
328        (counts, total_tokens)
329    };
330
331    if counts.is_empty() {
332        return SparseVector {
333            indices: Vec::new(),
334            values: Vec::new(),
335        };
336    }
337
338    let mut indices: Vec<u32> = counts.keys().copied().collect();
339    indices.sort_unstable();
340
341    let values: Vec<f32> = indices
342        .iter()
343        .map(|idx| 1.0 + (counts[idx] as f64).ln() as f32)
344        .collect();
345
346    SparseVector { indices, values }
347}
348
349pub fn build_document(text: &str, k1: f64, b: f64, avgdl: f64) -> SparseVector {
350    let (counts, total_tokens) = if is_ascii(text) {
351        let (hashes, total_tokens) = hash_tokens_ascii_fast(text);
352        let mut counts: FastMap<u32, f32> = FastMap::with_capacity_and_hasher(
353            hashes.len(),
354            BuildHasherDefault::<IdentityHasher>::default(),
355        );
356        for h in hashes {
357            *counts.entry(h).or_insert(0.0) += 1.0;
358        }
359        (counts, total_tokens)
360    } else {
361        let tokens = tokenize(text);
362        let mut counts: FastMap<u32, f32> = FastMap::with_capacity_and_hasher(
363            tokens.len(),
364            BuildHasherDefault::<IdentityHasher>::default(),
365        );
366        let total_tokens = tokens.len();
367        for token in &tokens {
368            *counts.entry(hash_token(token)).or_insert(0.0) += 1.0;
369        }
370        (counts, total_tokens)
371    };
372
373    if counts.is_empty() {
374        return SparseVector {
375            indices: Vec::new(),
376            values: Vec::new(),
377        };
378    }
379
380    let doc_len = total_tokens as f64;
381    let safe_avgdl = if avgdl <= 0.0 { 256.0 } else { avgdl };
382    let denom_scale = k1 * (1.0 - b + b * doc_len / safe_avgdl);
383    let k1p1 = k1 + 1.0;
384
385    let mut indices: Vec<u32> = counts.keys().copied().collect();
386    indices.sort_unstable();
387
388    let values: Vec<f32> = indices
389        .iter()
390        .map(|idx| {
391            let tf_count = counts[idx] as f64;
392            let denom = tf_count + denom_scale;
393            (tf_count * k1p1 / denom) as f32
394        })
395        .collect();
396
397    SparseVector { indices, values }
398}
399
400pub fn build_query_default(text: &str) -> SparseVector {
401    build_query(text)
402}
403
404pub fn build_document_default(text: &str) -> SparseVector {
405    build_document(text, 1.2, 0.75, 256.0)
406}
qql_embed/sparse.rs

qql_embed/
sparse.rs
