core_api/
exact_knn.rs

1//! Exact cosine kernel: pack L2-normalised rows, GEMV / gram via one `dgemm`.
2//!
3//! Used by brute `find_similar` and `pairwise_similar`. Not used by HNSW.
4
5use core_query::GraphView;
6use core_storage::Value;
7use std::borrow::Cow;
8use std::cell::Cell;
9
10/// Above this packed n, `pairwise_similar` uses n `gemv`s instead of `gram`.
11pub const PAIRWISE_GRAM_MAX: usize = 4_096;
12/// Hard refuse for `pairwise_similar` on resolved unique key count.
13pub const PAIRWISE_MAX_N: usize = 8_192;
14
15thread_local! {
16    static PAIRWISE_GRAM_MAX_OVERRIDE: Cell<Option<usize>> = const { Cell::new(None) };
17    static PAIRWISE_MAX_N_OVERRIDE: Cell<Option<usize>> = const { Cell::new(None) };
18}
19
20pub(crate) fn pairwise_gram_max() -> usize {
21    PAIRWISE_GRAM_MAX_OVERRIDE.with(|c| c.get().unwrap_or(PAIRWISE_GRAM_MAX))
22}
23
24pub(crate) fn pairwise_max_n() -> usize {
25    PAIRWISE_MAX_N_OVERRIDE.with(|c| c.get().unwrap_or(PAIRWISE_MAX_N))
26}
27
28/// Run `f` with temporary pairwise n caps. Restores the previous overrides
29/// (including across panics). Thread-local, same shape as `HNSW_BUILD_BATCH`.
30pub fn with_pairwise_caps<R>(gram_max: usize, max_n: usize, f: impl FnOnce() -> R) -> R {
31    let prev_gram = PAIRWISE_GRAM_MAX_OVERRIDE.with(|c| c.replace(Some(gram_max)));
32    let prev_n = PAIRWISE_MAX_N_OVERRIDE.with(|c| c.replace(Some(max_n)));
33    let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
34    PAIRWISE_GRAM_MAX_OVERRIDE.with(|c| c.set(prev_gram));
35    PAIRWISE_MAX_N_OVERRIDE.with(|c| c.set(prev_n));
36    match out {
37        Ok(v) => v,
38        Err(p) => std::panic::resume_unwind(p),
39    }
40}
41
42/// Row-major packed, L2-normalised f64 matrix. Not persisted.
43pub struct PackedVectors {
44    pub ids: Vec<u32>, // row i is node ids[i]
45    pub dim: usize,
46    pub data: Vec<f64>, // len == ids.len() * dim, each row unit-length
47}
48
49/// Cosine of two already-unit vectors (dot). Test helper only — mixed-dim
50/// candidates are skipped, not tailed through this.
51#[allow(dead_code)]
52pub fn cosine_unit(a: &[f64], b: &[f64]) -> f64 {
53    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
54}
55
56/// Pack candidates whose `len() == dim` and whose L2 norm is non-zero.
57/// Other candidates are omitted (not scored).
58///
59/// A row whose L2² is within `PACK_UNIT_L2SQ_EPS` of 1 is already unit:
60/// copy it and skip the second L2 (the scale pass).
61pub fn pack<'a, I>(rows: I, dim: usize) -> PackedVectors
62where
63    I: IntoIterator<Item = (u32, &'a [f64])>,
64{
65    let mut ids = Vec::new();
66    let mut data = Vec::new();
67    for (id, row) in rows {
68        if row.len() != dim {
69            continue;
70        }
71        let n2 = pack_l2sq(row);
72        if n2 == 0.0 {
73            continue;
74        }
75        ids.push(id);
76        if (n2 - 1.0).abs() <= PACK_UNIT_L2SQ_EPS {
77            data.extend_from_slice(row);
78        } else {
79            let norm = n2.sqrt();
80            pack_scale_unit(row, norm, &mut data);
81        }
82    }
83    PackedVectors { ids, dim, data }
84}
85
86/// |‖x‖² − 1| at or below this → already unit. Same domain as the zero-row
87/// check (`n2 == 0.0`): squared L2, not ‖x‖. Tight enough that the GEMM-vs-
88/// scalar pin (1e-9) is unaffected.
89const PACK_UNIT_L2SQ_EPS: f64 = 1e-12;
90
91#[cfg(test)]
92thread_local! {
93    static PACK_L2_CALLS: Cell<u64> = const { Cell::new(0) };
94}
95
96#[inline]
97fn note_pack_l2() {
98    #[cfg(test)]
99    PACK_L2_CALLS.with(|c| c.set(c.get().saturating_add(1)));
100}
101
102#[inline]
103fn pack_l2sq(row: &[f64]) -> f64 {
104    note_pack_l2();
105    row.iter().map(|x| x * x).sum()
106}
107
108#[inline]
109fn pack_scale_unit(row: &[f64], norm: f64, data: &mut Vec<f64>) {
110    note_pack_l2();
111    data.extend(row.iter().map(|x| x / norm));
112}
113
114/// `out[i] = row(i) · q_unit`. `q_unit.len() == packed.dim`.
115pub fn gemv(packed: &PackedVectors, q_unit: &[f64]) -> Vec<f64> {
116    debug_assert_eq!(q_unit.len(), packed.dim);
117    let n = packed.ids.len();
118    let dim = packed.dim;
119    let mut out = vec![0.0; n];
120    if q_unit.len() != dim {
121        return out;
122    }
123    dgemm_f64(
124        n,
125        dim,
126        1,
127        &packed.data,
128        dim as isize,
129        1,
130        q_unit,
131        1,
132        1,
133        &mut out,
134        1,
135        1,
136    );
137    out
138}
139
140/// `out` is n×n row-major `A Aᵀ`. Diagonal is ~1 for unit rows.
141/// Callers must not invoke this for `n > PAIRWISE_GRAM_MAX`.
142pub fn gram(packed: &PackedVectors) -> Vec<f64> {
143    let n = packed.ids.len();
144    let dim = packed.dim;
145    let mut out = vec![0.0; n.saturating_mul(n)];
146    dgemm_f64(
147        n,
148        dim,
149        n,
150        &packed.data,
151        dim as isize,
152        1,
153        &packed.data,
154        1,
155        dim as isize,
156        &mut out,
157        n as isize,
158        1,
159    );
160    out
161}
162
163/// Overlay-aware f64 view of node `id`'s `field`.
164/// Base `ColumnData::Vector` and no overlay/tombstone → `ColumnsView::vector`
165/// (`Cow::Borrowed` when aligned). Overlay `Value::List` → owned f64s via
166/// `value_as_float_list`. Missing / non-list → `None`.
167pub fn vector_f64<'a>(view: &'a GraphView<'_>, id: u32, field: &str) -> Option<Cow<'a, [f64]>> {
168    if let Some(v) = view.props.vector(id, field) {
169        return Some(v);
170    }
171    let vr = view.prop(id, field)?;
172    let xs = value_as_float_list(vr.as_value())?;
173    Some(Cow::Owned(xs))
174}
175
176fn value_as_float_list(v: &Value) -> Option<Vec<f64>> {
177    match v {
178        Value::List(items) => items
179            .iter()
180            .map(|item| match item {
181                Value::Float(f) => Some(*f),
182                Value::Int(i) => Some(*i as f64),
183                _ => None,
184            })
185            .collect(),
186        _ => None,
187    }
188}
189
190/// C ← A B. Empty `m`/`k`/`n` returns without calling dgemm (`c` left as-is).
191#[allow(clippy::too_many_arguments)]
192fn dgemm_f64(
193    m: usize,
194    k: usize,
195    n: usize,
196    a: &[f64],
197    rsa: isize,
198    csa: isize,
199    b: &[f64],
200    rsb: isize,
201    csb: isize,
202    c: &mut [f64],
203    rsc: isize,
204    csc: isize,
205) {
206    if m == 0 || k == 0 || n == 0 {
207        return;
208    }
209    debug_assert!(c.len() >= m.saturating_mul(n));
210    // SAFETY: m, k, n are non-zero. `a` is the m×k matrix at (`rsa`, `csa`);
211    // `b` is the k×n matrix at (`rsb`, `csb`); `c` is the m×n output at
212    // (`rsc`, `csc`) and does not alias `a` or `b`. `rsc`/`csc` are non-zero
213    // at every call site, so C elements do not alias each other. β = 0 so C
214    // need not be initialized; the Vec is zeroed anyway.
215    unsafe {
216        matrixmultiply::dgemm(
217            m,
218            k,
219            n,
220            1.0,
221            a.as_ptr(),
222            rsa,
223            csa,
224            b.as_ptr(),
225            rsb,
226            csb,
227            0.0,
228            c.as_mut_ptr(),
229            rsc,
230            csc,
231        );
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn gram_2x2_orthonormal() {
241        let r0: [f64; 2] = [1.0, 0.0];
242        let r1: [f64; 2] = [0.0, 1.0];
243        let packed = pack([(0, r0.as_slice()), (1, r1.as_slice())], 2);
244        let g = gram(&packed);
245        assert_eq!(g.len(), 4);
246        assert!((g[0] - 1.0).abs() < 1e-12, "g00={}", g[0]);
247        assert!(g[1].abs() < 1e-12, "g01={}", g[1]);
248        assert!(g[2].abs() < 1e-12, "g10={}", g[2]);
249        assert!((g[3] - 1.0).abs() < 1e-12, "g11={}", g[3]);
250    }
251
252    fn ulps(a: f64, b: f64) -> u64 {
253        if a == b {
254            return 0;
255        }
256        let mut ai = a.to_bits() as i64;
257        let mut bi = b.to_bits() as i64;
258        if ai < 0 {
259            ai = i64::MIN - ai;
260        }
261        if bi < 0 {
262            bi = i64::MIN - bi;
263        }
264        ai.abs_diff(bi)
265    }
266
267    fn pack_l2_calls() -> u64 {
268        PACK_L2_CALLS.with(|c| c.get())
269    }
270
271    fn pack_l2_calls_reset() {
272        PACK_L2_CALLS.with(|c| c.set(0));
273    }
274
275    /// Already-unit rows must match the always-normalise pack within 1 ulp
276    /// and must not run the scale pass (second L2).
277    #[test]
278    fn pack_skips_second_l2_on_unit() {
279        let unit = [0.6_f64, 0.8];
280        let n2 = unit[0] * unit[0] + unit[1] * unit[1];
281        assert!(
282            (n2 - 1.0).abs() <= PACK_UNIT_L2SQ_EPS,
283            "fixture must be unit in the packer's epsilon, n2={n2}"
284        );
285
286        pack_l2_calls_reset();
287        let packed = pack([(7, unit.as_slice())], 2);
288        assert_eq!(packed.ids, vec![7]);
289        assert_eq!(packed.dim, 2);
290        assert_eq!(packed.data.len(), 2);
291
292        let norm = n2.sqrt();
293        let oracle = [unit[0] / norm, unit[1] / norm];
294        for (i, (&got, &expect)) in packed.data.iter().zip(oracle.iter()).enumerate() {
295            assert!(
296                ulps(got, expect) <= 1,
297                "unit[{i}]: packed {got} vs always-L2 {expect} ulps={}",
298                ulps(got, expect)
299            );
300        }
301        assert_eq!(pack_l2_calls(), 1, "already-unit row must not L2 twice");
302    }
303
304    #[test]
305    fn gemv_matches_cosine_unit() {
306        let a = [3.0, 4.0];
307        let b = [1.0, 0.0];
308        let c = [0.0, 2.0];
309        let packed = pack([(0, a.as_slice()), (1, b.as_slice()), (2, c.as_slice())], 2);
310        let q = [1.0, 1.0];
311        let qn = q.iter().map(|x| x * x).sum::<f64>().sqrt();
312        let q_unit = [q[0] / qn, q[1] / qn];
313        let scores = gemv(&packed, &q_unit);
314        assert_eq!(scores.len(), packed.ids.len());
315        for (i, row) in packed.data.chunks(packed.dim).enumerate() {
316            let expected = cosine_unit(row, &q_unit);
317            assert!(
318                (scores[i] - expected).abs() < 1e-12,
319                "row {i}: gemv {} vs cosine_unit {expected}",
320                scores[i]
321            );
322        }
323    }
324}
core_api/exact_knn.rs

core_api/
exact_knn.rs
