Batch distance: compute distances from `query` to each candidate. Returns `(index, distance)` pairs sorted ascending, truncated to `top_k`.
(
query: &[f32],
candidates: &[&[f32]],
metric: DistanceMetric,
top_k: usize,
)
| 41 | /// |
| 42 | /// Returns `(index, distance)` pairs sorted ascending, truncated to `top_k`. |
| 43 | pub fn batch_distances( |
| 44 | query: &[f32], |
| 45 | candidates: &[&[f32]], |
| 46 | metric: DistanceMetric, |
| 47 | top_k: usize, |
| 48 | ) -> Vec<(usize, f32)> { |
| 49 | let mut dists: Vec<(usize, f32)> = candidates |
| 50 | .iter() |
| 51 | .enumerate() |
| 52 | .map(|(i, c)| (i, distance(query, c, metric))) |
| 53 | .collect(); |
| 54 | |
| 55 | if top_k < dists.len() { |
| 56 | dists.select_nth_unstable_by(top_k, |a, b| { |
| 57 | a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal) |
| 58 | }); |
| 59 | dists.truncate(top_k); |
| 60 | } |
| 61 | dists.sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); |
| 62 | dists |
| 63 | } |