nodedb_vector/distance/
compute.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Top-level distance and batch-distance computation, dispatched to SIMD.
4
5use super::scalar::*;
6use super::simd;
7
8/// Compute distance between two vectors using the specified metric.
9///
10/// Dispatches to SIMD kernels (AVX-512, AVX2+FMA, NEON) where available;
11/// falls back to scalar implementations on other architectures.
12#[inline]
13pub fn distance(a: &[f32], b: &[f32], metric: DistanceMetric) -> f32 {
14    assert_eq!(
15        a.len(),
16        b.len(),
17        "distance: length mismatch (a.len()={}, b.len()={})",
18        a.len(),
19        b.len()
20    );
21    let rt = simd::runtime();
22    match metric {
23        DistanceMetric::L2 => (rt.l2_squared)(a, b),
24        DistanceMetric::Cosine => (rt.cosine_distance)(a, b),
25        DistanceMetric::InnerProduct => (rt.neg_inner_product)(a, b),
26        DistanceMetric::Manhattan => manhattan(a, b),
27        DistanceMetric::Chebyshev => chebyshev(a, b),
28        DistanceMetric::Hamming => hamming_f32(a, b),
29        DistanceMetric::Jaccard => jaccard(a, b),
30        DistanceMetric::Pearson => pearson(a, b),
31        // Unknown future metric — fall back to L2.
32        _ => (rt.l2_squared)(a, b),
33    }
34}
35
36/// Batch distance: compute distances from `query` to each candidate.
37///
38/// Returns `(index, distance)` pairs sorted ascending, truncated to `top_k`.
39pub fn batch_distances(
40    query: &[f32],
41    candidates: &[&[f32]],
42    metric: DistanceMetric,
43    top_k: usize,
44) -> Vec<(usize, f32)> {
45    let mut dists: Vec<(usize, f32)> = candidates
46        .iter()
47        .enumerate()
48        .map(|(i, c)| (i, distance(query, c, metric)))
49        .collect();
50
51    if top_k < dists.len() {
52        dists.select_nth_unstable_by(top_k, |a, b| {
53            a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
54        });
55        dists.truncate(top_k);
56    }
57    dists.sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
58    dists
59}
nodedb_vector/distance/compute.rs

nodedb_vector/distance/
compute.rs
