mod adjacency;
mod build_kernels;
mod fastscan;
pub mod filter;
pub(crate) mod kernels;
mod projection;
mod rerank;
#[cfg(not(target_arch = "wasm32"))]
mod update;
mod vertex;
mod vnni;
#[cfg(not(target_arch = "wasm32"))]
pub use update::{
Consolidation, ConsolidationMode, ConsolidationPolicy, UpdatableVGraph, UpdateError,
};
#[cfg(feature = "experimental")]
pub mod diagnostics;
#[cfg(feature = "experimental")]
pub mod projection_diagnostic;
#[cfg(feature = "experimental")]
pub mod order_diagnostic;
#[cfg(feature = "experimental")]
pub mod ivf;
use std::cell::RefCell;
use rayon::prelude::*;
use crate::core::error::HmsError;
use adjacency::Adjacency;
use kernels::{
as_f32, as_f32_mut, as_u32, as_u32_mut, dot, dot_i8_query, dot_i8_query_x4, QueryCode,
Rotation, SplitMix,
};
#[cfg(feature = "experimental")]
pub use vertex::QueryProfile;
pub use vertex::{
huge_page_status, HugePageStatus, Residency, SearchKnobs, VGraph, VGraphMemoryBreakdown,
VMetric, VSearchParams, VSearcher,
};
pub const DEFAULT_DEGREE: usize = 32;
const NONE: u32 = u32::MAX;
fn invalid(details: String) -> HmsError {
HmsError::InvalidParam { details }
}
const UPPER_DEGREE: usize = 16;
const LAYER_RATIO: usize = 16;
const MIN_LAYER: usize = 64;
const BUILD_CODES_MIN_DIM: usize = 128;
const FIRST_PASS_EF: usize = 200;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum Metric {
#[default]
Angular,
Euclidean,
}
impl Metric {
fn code(self) -> u64 {
match self {
Self::Angular => 0,
Self::Euclidean => 1,
}
}
}
#[derive(Clone, Debug)]
pub struct BuildParams {
pub metric: Metric,
pub build_ef: usize,
pub degree: usize,
pub alpha: f32,
pub code_bits: usize,
pub seed: u64,
pub quantized_search: bool,
}
impl Default for BuildParams {
fn default() -> Self {
Self {
metric: Metric::Angular,
build_ef: 128,
degree: DEFAULT_DEGREE,
alpha: 1.0,
code_bits: 0,
seed: 0x5EED,
quantized_search: false,
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct SearchParams {
pub ef: usize,
pub max_exact: usize,
}
pub struct QGraph {
n: usize,
dim: usize,
degree: usize,
words: usize,
rotation: Rotation,
stride: usize,
codes_off: usize,
factors_off: usize,
ids_off: usize,
blocks: Vec<u64>,
entry: u32,
upper_ids: Vec<u32>,
layers: Vec<Vec<Vec<u32>>>,
}
#[derive(Clone, Copy)]
struct Candidate {
dist: f32,
id: u32,
done: bool,
}
#[cfg(feature = "experimental-build-search")]
#[inline(always)]
fn insert_build_candidate(
pool: &mut Vec<Candidate>,
candidate: Candidate,
limit: usize,
) -> Option<usize> {
debug_assert!(limit > 0 && pool.len() <= limit);
if pool.len() == limit && candidate.dist >= pool[limit - 1].dist {
return None;
}
let pos = pool.partition_point(|c| c.dist <= candidate.dist);
if pool.len() == limit {
let _ = pool.pop();
}
pool.insert(pos, candidate);
Some(pos)
}
struct Visited {
marks: Vec<u8>,
epoch: u8,
}
impl Visited {
fn new(n: usize) -> Self {
Self {
marks: vec![0; n],
epoch: 0,
}
}
fn next(&mut self) {
self.epoch = self.epoch.wrapping_add(1);
if self.epoch == 0 {
self.marks.fill(0);
self.epoch = 1;
}
}
fn insert(&mut self, id: u32) -> bool {
let m = &mut self.marks[id as usize];
let fresh = *m != self.epoch;
*m = self.epoch;
fresh
}
}
#[derive(Default)]
struct BuildVisited {
bits: Vec<u64>,
touched: Vec<u32>,
pool: Vec<Candidate>,
fresh: Vec<u32>,
dists: Vec<f32>,
}
impl BuildVisited {
fn reset(&mut self, n: usize) {
for word in self.touched.drain(..) {
self.bits[word as usize] = 0;
}
self.bits.resize(n.div_ceil(64), 0);
}
fn insert(&mut self, id: u32) -> bool {
let (w, b) = (id as usize / 64, 1u64 << (id % 64));
let old = self.bits[w];
if old & b != 0 {
return false;
}
if old == 0 {
self.touched.push(w as u32);
}
self.bits[w] = old | b;
true
}
}
#[cfg(feature = "experimental")]
#[doc(hidden)]
pub fn set_build_phase_hook(hook: fn(&str)) -> bool {
prof::HOOK.set(hook).is_ok()
}
mod prof {
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::sync::OnceLock;
use std::time::Instant;
pub(super) static GREEDY: AtomicU64 = AtomicU64::new(0);
pub(super) static PRUNE: AtomicU64 = AtomicU64::new(0);
pub(super) static REV: AtomicU64 = AtomicU64::new(0);
pub(super) static FRESH: AtomicU64 = AtomicU64::new(0);
pub(super) static EXACT: AtomicU64 = AtomicU64::new(0);
pub(super) static HOOK: OnceLock<fn(&str)> = OnceLock::new();
pub(super) fn on() -> bool {
static ON: OnceLock<bool> = OnceLock::new();
*ON.get_or_init(|| {
HOOK.get().is_some() || std::env::var_os("HMS_QGRAPH_BUILD_PROFILE").is_some()
})
}
pub(super) fn start() -> Option<Instant> {
on().then(Instant::now)
}
pub(super) fn add(c: &AtomicU64, t: Option<Instant>) {
if let Some(t) = t {
c.fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
}
}
pub(super) fn count(c: &AtomicU64, n: u64) {
if on() {
c.fetch_add(n, Relaxed);
}
}
pub(super) fn report(phase: &str, t: Option<Instant>) {
let Some(t) = t else { return };
if let Some(hook) = HOOK.get() {
hook(phase);
}
let secs = |c: &AtomicU64| c.swap(0, Relaxed) as f64 * 1e-9;
eprintln!(
"PHASE {phase} wall={:.2}s greedy_cpu={:.1}s prune_cpu={:.1}s rev_wall={:.2}s fresh={} exact={}",
t.elapsed().as_secs_f64(),
secs(&GREEDY),
secs(&PRUNE),
secs(&REV),
FRESH.swap(0, Relaxed),
EXACT.swap(0, Relaxed),
);
}
pub(super) fn graph_hash(h: impl std::hash::Hash) {
if on() {
use std::hash::Hasher;
let mut s = std::collections::hash_map::DefaultHasher::new();
h.hash(&mut s);
eprintln!("GRAPH_HASH {:016x}", s.finish());
}
}
}
#[derive(Default)]
struct InsertScratch {
cands: Vec<(f32, u32)>,
d: Vec<f32>,
}
#[derive(Default)]
struct MergeScratch {
d: Vec<f32>,
cands: Vec<(f32, u32, bool)>,
open: Vec<u32>,
out: Vec<u32>,
}
#[derive(Default)]
struct FillScratch {
ids: Vec<u32>,
skip: Vec<u32>,
bounds: Vec<(f32, f32)>,
upper: Vec<f32>,
d: Vec<f32>,
near: Vec<(f32, u32)>,
}
thread_local! {
static BUILD_VISITED: RefCell<BuildVisited> = RefCell::new(BuildVisited::default());
static INSERT_SCRATCH: RefCell<InsertScratch> = RefCell::default();
static FILL_SCRATCH: RefCell<FillScratch> = RefCell::default();
}
struct Rows<'a> {
data: &'a [f32],
dim: usize,
norms: Option<&'a [f32]>,
codes: Option<&'a BuildCodes>,
quantized: bool,
}
struct BuildCodes {
width: usize,
q: Vec<i8>,
meta: Vec<CodeMeta>,
sums: Vec<i32>,
euclidean: bool,
rel_slack: f32,
}
#[derive(Clone, Copy, Default)]
#[repr(C, align(16))]
struct CodeMeta {
scale: f32,
err: f32,
n2: f32,
n: f32,
}
impl BuildCodes {
fn new(data: &[f32], dim: usize, norms: Option<&[f32]>) -> Self {
let n = data.len() / dim;
let width = dim.next_multiple_of(16);
let mut q = vec![0i8; n * width];
let mut meta = vec![CodeMeta::default(); n];
q.par_chunks_mut(width)
.zip(meta.par_iter_mut())
.zip(data.par_chunks(dim))
.for_each(|((q, m), x)| {
let (s, e) = (&mut m.scale, &mut m.err);
let max = x.iter().fold(0f32, |m, v| m.max(v.abs()));
*s = if max > 0.0 { max / 127.0 } else { 0.0 };
let mut e2 = 0f64;
for (o, &v) in q.iter_mut().zip(x) {
let c = if *s > 0.0 {
(v / *s).round().clamp(-127.0, 127.0)
} else {
0.0
};
*o = c as i8;
let r = f64::from(v) - f64::from(*s) * f64::from(c);
e2 += r * r;
}
*e = (e2.sqrt() * (1.0 + 1e-6)) as f32 + 1e-7;
});
if let Some(norms) = norms {
for (m, &n2) in meta.iter_mut().zip(norms) {
(m.n2, m.n) = (n2, n2.sqrt());
}
}
let sums = q.par_chunks(width).map(vnni::query_sum).collect();
Self {
width,
q,
meta,
sums,
euclidean: norms.is_some(),
rel_slack: 1e-4f32.max((dim as f32 + 4.0) * 2.0 * f32::EPSILON),
}
}
fn code(&self, i: u32) -> &[i8] {
let i = i as usize;
&self.q[i * self.width..(i + 1) * self.width]
}
#[inline]
fn prefetch(&self, i: u32) {
build_kernels::prefetch_code(self.code(i));
build_kernels::prefetch_meta(&self.meta[i as usize]);
}
#[inline]
fn d2_lower(&self, p: u32, u: u32) -> f32 {
self.d2_bounds(p, u).0
}
#[inline]
fn raw(&self, p: u32, u: u32) -> i32 {
dot_i8_query(self.code(p), self.code(u), self.sums[p as usize])
}
#[inline]
fn d2_est(&self, p: u32, u: u32) -> f32 {
self.est_from_raw(p, u, self.raw(p, u))
}
#[inline]
fn d2_est4(&self, p: u32, u: [u32; 4]) -> [f32; 4] {
let raw = dot_i8_query_x4(self.code(p), u.map(|u| self.code(u)), self.sums[p as usize]);
[0, 1, 2, 3].map(|i| self.est_from_raw(p, u[i], raw[i]))
}
#[inline]
fn est_from_raw(&self, p: u32, u: u32, raw: i32) -> f32 {
let raw = raw as f32;
let (mp, mu) = (&self.meta[p as usize], &self.meta[u as usize]);
if self.euclidean {
(mp.n2 + mu.n2 - 2.0 * mp.scale * mu.scale * raw).max(0.0)
} else {
(2.0 - 2.0 * mp.scale * mu.scale * raw).max(0.0)
}
}
#[inline]
fn d2_bounds(&self, p: u32, u: u32) -> (f32, f32) {
const SLACK: f32 = 1e-4;
let (mp, mu) = (&self.meta[p as usize], &self.meta[u as usize]);
let raw = self.raw(p, u) as f32;
let ep = mp.err;
let est = mp.scale * mu.scale * raw;
if !self.euclidean {
let margin = ep + (1.0 + ep) * mu.err + SLACK;
return (2.0 - 2.0 * (est + margin), 2.0 - 2.0 * (est - margin));
}
let ((np2, np), (nu2, nu)) = ((mp.n2, mp.n), (mu.n2, mu.n));
let margin = ep * nu + (np + ep) * mu.err;
let slack = self.rel_slack * (np2 + nu2 + 2.0 * np * nu);
let base = np2 + nu2;
(
base - 2.0 * (est + margin) - slack,
base - 2.0 * (est - margin) + slack,
)
}
}
impl Rows<'_> {
fn row(&self, i: u32) -> &[f32] {
let i = i as usize;
&self.data[i * self.dim..(i + 1) * self.dim]
}
#[inline]
fn prefetch(&self, u: u32) {
build_kernels::prefetch_row(self.row(u));
if let Some(n) = self.norms {
build_kernels::prefetch_meta(&n[u as usize]);
}
}
#[inline]
fn dist(&self, a: u32, b: u32, dot: f32) -> f32 {
match self.norms {
None => (2.0 - 2.0 * dot).max(0.0),
Some(n) => (n[a as usize] + n[b as usize] - 2.0 * dot).max(0.0),
}
}
fn d2(&self, a: u32, b: u32) -> f32 {
self.dist(a, b, dot(self.row(a), self.row(b)))
}
fn d2_many(&self, p: u32, ids: &[u32], out: &mut Vec<f32>) {
out.clear();
let target = self.row(p);
let (quads, rest) = ids.as_chunks::<4>();
for q in quads {
let d = build_kernels::dot4(target, q.map(|u| self.row(u)));
out.extend(d.iter().zip(q).map(|(&x, &u)| self.dist(p, u, x)));
}
out.extend(rest.iter().map(|&u| self.d2(p, u)));
}
fn d2_stream(&self, p: u32, ids: &[u32], out: &mut Vec<f32>) {
const AHEAD: usize = 16;
out.clear();
let target = self.row(p);
for &u in ids.iter().take(AHEAD) {
self.prefetch(u);
}
let (quads, rest) = ids.as_chunks::<4>();
for (i, q) in quads.iter().enumerate() {
for &u in ids.iter().skip(AHEAD + 4 * i).take(4) {
self.prefetch(u);
}
let d = build_kernels::dot4(target, q.map(|u| self.row(u)));
out.extend(d.iter().zip(q).map(|(&x, &u)| self.dist(p, u, x)));
}
out.extend(rest.iter().map(|&u| self.d2(p, u)));
}
fn occluded(&self, c: u32, kept: &[u32], dc: f32, a2: f32) -> bool {
let rc = self.row(c);
if let Some(codes) = self.codes {
let mut open = [0u32; 4];
let mut n_open = 0;
for &s in kept {
let (lo, hi) = codes.d2_bounds(c, s);
if a2 * hi <= dc {
return true;
}
if a2 * lo > dc {
continue;
}
open[n_open] = s;
n_open += 1;
if n_open == 4 {
n_open = 0;
let d = build_kernels::dot4(rc, open.map(|u| self.row(u)));
if d.iter()
.zip(&open)
.any(|(&x, &s)| a2 * self.dist(c, s, x) <= dc)
{
return true;
}
}
}
return open[..n_open]
.iter()
.any(|&s| a2 * self.dist(s, c, dot(self.row(s), rc)) <= dc);
}
let (quads, rest) = kept.as_chunks::<4>();
for q in quads {
let d = build_kernels::dot4(rc, q.map(|u| self.row(u)));
if d.iter()
.zip(q)
.any(|(&x, &s)| a2 * self.dist(c, s, x) <= dc)
{
return true;
}
}
rest.iter()
.any(|&s| a2 * self.dist(s, c, dot(self.row(s), rc)) <= dc)
}
fn greedy(
&self,
graph: &Adjacency,
start: u32,
p: u32,
l: usize,
visited: &mut BuildVisited,
) -> Vec<(f32, u32)> {
let mut expanded = Vec::new();
self.greedy_into(graph, start, p, l, visited, &mut expanded);
expanded
}
fn greedy_into(
&self,
graph: &Adjacency,
start: u32,
p: u32,
l: usize,
visited: &mut BuildVisited,
expanded: &mut Vec<(f32, u32)>,
) {
let (mut n_fresh, mut n_exact) = (0u64, 0u64);
visited.reset(graph.len());
visited.insert(start);
let mut pool = std::mem::take(&mut visited.pool);
let mut fresh = std::mem::take(&mut visited.fresh);
let mut dists = std::mem::take(&mut visited.dists);
pool.clear();
pool.push(Candidate {
dist: self.d2(p, start),
id: start,
done: false,
});
expanded.clear();
let mut cur = 0;
while cur < pool.len() {
pool[cur].done = true;
let v = pool[cur].id;
expanded.push((pool[cur].dist, v));
fresh.clear();
match self.codes.filter(|_| self.quantized || pool.len() >= l) {
Some(codes) if self.quantized => {
for &u in graph.get(v as usize) {
if visited.insert(u) {
codes.prefetch(u);
fresh.push(u);
}
}
n_fresh += fresh.len() as u64;
}
Some(codes) => {
for &u in graph.get(v as usize) {
if visited.insert(u) {
codes.prefetch(u);
fresh.push(u);
}
}
let bound = pool[l - 1].dist;
n_fresh += fresh.len() as u64;
fresh.retain(|&u| codes.d2_lower(p, u) < bound);
for &u in &fresh {
build_kernels::prefetch_row(self.row(u));
}
}
None => {
for &u in graph.get(v as usize) {
if visited.insert(u) {
build_kernels::prefetch_row(self.row(u));
fresh.push(u);
}
}
n_fresh += fresh.len() as u64;
}
}
match self.codes.filter(|_| self.quantized) {
Some(codes) => {
dists.clear();
let (quads, rest) = fresh.as_chunks::<4>();
for &q in quads {
dists.extend(codes.d2_est4(p, q));
}
dists.extend(rest.iter().map(|&u| codes.d2_est(p, u)));
}
None => {
n_exact += fresh.len() as u64;
self.d2_many(p, &fresh, &mut dists);
}
}
let mut best = usize::MAX;
for (&u, &du) in fresh.iter().zip(&dists) {
#[cfg(feature = "experimental-build-search")]
let Some(pos) = insert_build_candidate(
&mut pool,
Candidate {
dist: du,
id: u,
done: false,
},
l,
) else {
continue;
};
#[cfg(not(feature = "experimental-build-search"))]
let pos = {
if pool.len() >= l && du >= pool[l - 1].dist {
continue;
}
let pos = pool.partition_point(|c| c.dist <= du);
pool.insert(
pos,
Candidate {
dist: du,
id: u,
done: false,
},
);
pool.truncate(l);
pos
};
best = best.min(pos);
build_kernels::prefetch_meta(graph.len_word(u as usize));
}
cur = best.min(cur + 1);
while cur < pool.len() && pool[cur].done {
cur += 1;
}
if let Some(c) = pool.get(cur) {
build_kernels::prefetch_ids(graph.get(c.id as usize));
}
}
if self.quantized {
fresh.clear();
fresh.extend(expanded.iter().map(|e| e.1));
self.d2_stream(p, &fresh, &mut dists);
n_exact += fresh.len() as u64;
for (e, &d) in expanded.iter_mut().zip(&dists) {
e.0 = d;
}
}
(visited.pool, visited.fresh, visited.dists) = (pool, fresh, dists);
prof::count(&prof::FRESH, n_fresh);
prof::count(&prof::EXACT, n_exact);
}
fn prune(&self, p: u32, mut cands: Vec<(f32, u32)>, alpha: f32, r: usize) -> Vec<u32> {
let mut out = Vec::with_capacity(r);
self.prune_into(p, &mut cands, alpha, r, &mut out);
out
}
fn prune_into(
&self,
p: u32,
cands: &mut Vec<(f32, u32)>,
alpha: f32,
r: usize,
out: &mut Vec<u32>,
) {
cands.retain(|c| c.1 != p);
cands.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
cands.dedup_by_key(|c| c.1);
let a2 = alpha * alpha;
out.clear();
for &(dc, c) in cands.iter() {
if out.len() == r {
break;
}
if !self.occluded(c, out, dc, a2) {
out.push(c);
}
}
}
#[cfg(test)]
fn prune_merged(&self, t: u32, list: &[u32], closed: usize, alpha: f32, r: usize) -> Vec<u32> {
let mut scratch = MergeScratch::default();
self.prune_merged_into(t, list, closed, alpha, r, &mut scratch);
scratch.out
}
fn prune_merged_into(
&self,
t: u32,
list: &[u32],
closed: usize,
alpha: f32,
r: usize,
scratch: &mut MergeScratch,
) {
let MergeScratch {
d,
cands,
open,
out,
} = scratch;
self.d2_stream(t, list, d);
cands.clear();
cands.extend(
d.iter()
.zip(list)
.enumerate()
.map(|(i, (&d, &u))| (d, u, i < closed)),
);
cands.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
let a2 = alpha * alpha;
out.clear();
open.clear();
for &(dc, c, was_kept) in cands.iter() {
if out.len() == r {
break;
}
if !self.occluded(c, if was_kept { open } else { out }, dc, a2) {
out.push(c);
if !was_kept {
open.push(c);
}
}
}
}
#[allow(clippy::too_many_arguments)]
fn insert_pass(
&self,
graph: &mut Adjacency,
order: &[u32],
entry: u32,
l: usize,
alpha: f32,
r: usize,
) {
let n = graph.len();
let mut closed = vec![0u32; n];
let max_batch = (n / 50).max(64);
let mut outs = Adjacency::new(max_batch.min(order.len()), r);
let (mut start, mut size) = (0, 1);
while start < order.len() {
let end = (start + size).min(order.len());
let batch = &order[start..end];
let g: &Adjacency = graph;
outs.par_set(batch.len(), |i, out| {
let p = batch[i];
INSERT_SCRATCH.with(|cell| {
let InsertScratch { cands, d } = &mut *cell.borrow_mut();
let t0 = prof::start();
BUILD_VISITED.with(|vis| {
self.greedy_into(g, entry, p, l, &mut vis.borrow_mut(), cands);
});
prof::add(&prof::GREEDY, t0);
let t0 = prof::start();
let adj = g.get(p as usize);
self.d2_stream(p, adj, d);
cands.extend(d.iter().copied().zip(adj.iter().copied()));
self.prune_into(p, cands, alpha, r, out);
prof::add(&prof::PRUNE, t0);
});
});
let trev = prof::start();
let mut rev: Vec<(u32, u32)> = batch
.iter()
.enumerate()
.flat_map(|(i, &p)| outs.get(i).iter().map(move |&u| (u, p)))
.collect();
for (i, &p) in batch.iter().enumerate() {
let out = outs.get(i);
closed[p as usize] = out.len() as u32;
graph.set(p as usize, out);
}
rev.par_sort_unstable();
let groups: Vec<_> = rev.chunk_by(|a, b| a.0 == b.0).collect();
let targets = || groups.iter().map(|grp| grp[0].0 as usize);
let closed_of = pick_mut(&mut closed, targets());
let jobs: Vec<_> = graph
.rows_mut(targets())
.zip(closed_of)
.zip(&groups)
.collect();
jobs.into_par_iter().for_each_init(
<(Vec<u32>, MergeScratch)>::default,
|(merged, scratch), (((row, len), closed), grp)| {
let t = grp[0].0;
merged.clear();
merged.extend_from_slice(&row[..*len as usize]);
for &(_, s) in grp.iter() {
if !merged.contains(&s) {
merged.push(s);
}
}
let list = if merged.len() > r {
self.prune_merged_into(t, merged, *closed as usize, alpha, r, scratch);
*closed = scratch.out.len() as u32;
&scratch.out
} else {
&*merged
};
row[..list.len()].copy_from_slice(list);
*len = list.len() as u32;
},
);
prof::add(&prof::REV, trev);
start = end;
size = (size * 2).min(max_batch);
}
}
fn fill(&self, graph: &mut Adjacency, r: usize) {
graph.par_append(|g, v, out| {
let have = g.get(v);
if have.len() >= r {
return;
}
FILL_SCRATCH.with(|cell| self.fill_one(g, v, r, &mut cell.borrow_mut(), out));
});
}
fn fill_one(&self, g: &Adjacency, v: usize, r: usize, s: &mut FillScratch, out: &mut Vec<u32>) {
let FillScratch {
ids,
skip,
bounds,
upper,
d,
near,
} = s;
let have = g.get(v);
ids.clear();
ids.extend(have.iter().flat_map(|&u| g.get(u as usize).iter().copied()));
ids.sort_unstable();
ids.dedup();
skip.clear();
skip.extend_from_slice(have);
skip.push(v as u32);
skip.sort_unstable();
let mut at = 0;
ids.retain(|&u| {
while at < skip.len() && skip[at] < u {
at += 1;
}
skip.get(at) != Some(&u)
});
let k = r - have.len();
if let Some(codes) = self.codes.filter(|_| ids.len() > k) {
const AHEAD: usize = 8;
ids.iter().take(AHEAD).for_each(|&u| codes.prefetch(u));
bounds.clear();
bounds.extend(ids.iter().enumerate().map(|(i, &u)| {
if let Some(&w) = ids.get(i + AHEAD) {
codes.prefetch(w);
}
codes.d2_bounds(v as u32, u)
}));
upper.clear();
upper.extend(bounds.iter().map(|b| b.1));
let cut = *upper.select_nth_unstable_by(k - 1, f32::total_cmp).1;
let mut b = bounds.iter();
ids.retain(|_| b.next().is_some_and(|b| b.0 <= cut));
}
self.d2_stream(v as u32, ids, d);
near.clear();
near.extend(d.iter().copied().zip(ids.iter().copied()));
let order = |a: &(f32, u32), b: &(f32, u32)| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1));
if near.len() > k {
near.select_nth_unstable_by(k - 1, order);
near.truncate(k);
}
near.sort_by(order);
out.extend(near.iter().map(|x| x.1));
}
}
fn pick_mut<T>(values: &mut [T], ids: impl Iterator<Item = usize>) -> impl Iterator<Item = &mut T> {
let mut rest = values;
let mut next = 0;
ids.map(move |v| {
assert!(v >= next, "ids must be strictly increasing");
let (x, tail) = std::mem::take(&mut rest)[v - next..]
.split_first_mut()
.expect("id in range");
rest = tail;
next = v + 1;
x
})
}
fn edge_factors(pv: &[f32], pu: &[f32], code: &mut [u64]) -> (f32, f32, f32) {
let (mut r2, mut l1, mut sv, mut vw) = (0f32, 0f32, 0f32, 0f32);
for ((a64, b64), word) in pv
.as_chunks::<64>()
.0
.iter()
.zip(pu.as_chunks::<64>().0)
.zip(code.iter_mut())
{
let mut bits = 0u64;
for (i, (&a, &b)) in a64.iter().zip(b64).enumerate() {
let w = b - a;
r2 += w * w;
l1 += w.abs();
vw += a * w;
let up = w > 0.0;
sv += if up { a } else { -a };
bits |= u64::from(up) << i;
}
*word = bits;
}
let m = if l1 > 0.0 { r2 / l1 } else { 0.0 };
let pop = code.iter().map(|c| c.count_ones()).sum::<u32>() as f32;
(vw - m * sv, m, pop)
}
impl QGraph {
pub fn build(data: &[f32], dim: usize, params: &BuildParams) -> Result<Self, HmsError> {
Self::from_graph(&Graph::build(data, dim, params)?, params)
}
pub fn from_graph(g: &Graph, params: &BuildParams) -> Result<Self, HmsError> {
let (n, dim, r) = (g.n, g.dim, g.degree);
if r != params.degree {
return Err(invalid(format!(
"graph built with degree {r}, parameters say {}",
params.degree
)));
}
if g.metric != Metric::Angular {
return Err(invalid(
"the edge index is angular only; use VGraph for Euclidean data".into(),
));
}
let rotation = Rotation::new(dim, params.code_bits, params.seed);
let words = rotation.padded() / 64;
let rows = g.rows();
let graph = &g.adj;
let codes_off = dim.div_ceil(2);
let factors_off = codes_off + r * words;
let ids_off = factors_off + 3 * r / 2;
let stride = ids_off + r / 2;
let padded = rotation.padded();
let mut rotated = vec![0f32; n * padded];
rotated
.par_chunks_mut(padded)
.enumerate()
.for_each(|(v, out)| rotation.apply(rows.row(v as u32), out));
let rot_row = |v: u32| &rotated[v as usize * padded..(v as usize + 1) * padded];
let mut blocks = vec![0u64; n * stride];
blocks
.par_chunks_mut(stride)
.enumerate()
.for_each(|(v, b)| {
let adj = graph.get(v);
let (vec_w, rest) = b.split_at_mut(codes_off);
let (codes, rest) = rest.split_at_mut(r * words);
let (fac, ids) = rest.split_at_mut(3 * r / 2);
as_f32_mut(vec_w)[..dim].copy_from_slice(rows.row(v as u32));
let fac = as_f32_mut(fac);
let ids = as_u32_mut(ids);
ids.fill(NONE);
let pv = rot_row(v as u32);
for &u in adj.iter().take(4) {
build_kernels::prefetch_row(rot_row(u));
}
for (j, &u) in adj.iter().enumerate() {
if let Some(&next) = adj.get(j + 4) {
build_kernels::prefetch_row(rot_row(next));
}
ids[j] = u;
let (k, m, pop) =
edge_factors(pv, rot_row(u), &mut codes[j * words..(j + 1) * words]);
fac[j] = k;
fac[r + j] = m;
fac[2 * r + j] = pop;
}
});
Ok(Self {
n,
dim,
degree: r,
words,
rotation,
stride,
codes_off,
factors_off,
ids_off,
blocks,
entry: g.entry,
upper_ids: g.upper_ids.clone(),
layers: g.layers.clone(),
})
}
}
pub struct Graph {
n: usize,
dim: usize,
degree: usize,
metric: Metric,
vectors: Vec<f32>,
norms: Vec<f32>,
adj: Adjacency,
entry: u32,
upper_ids: Vec<u32>,
layers: Vec<Vec<Vec<u32>>>,
build: BuildParams,
}
struct GraphStructure {
adj: Adjacency,
entry: u32,
upper_ids: Vec<u32>,
layers: Vec<Vec<Vec<u32>>>,
}
#[derive(Clone, Copy)]
struct GraphView<'a> {
n: usize,
dim: usize,
degree: usize,
metric: Metric,
vectors: &'a [f32],
norms: &'a [f32],
adj: &'a Adjacency,
entry: u32,
upper_ids: &'a [u32],
layers: &'a [Vec<Vec<u32>>],
backing: Option<&'a memmap2::Mmap>,
}
impl GraphView<'_> {
fn neighbors(&self, v: usize) -> &[u32] {
self.adj.get(v)
}
fn release_rows(&self) -> std::io::Result<()> {
if let Some(map) = self.backing {
#[cfg(target_os = "linux")]
unsafe {
map.unchecked_advise(memmap2::UncheckedAdvice::DontNeed)?
};
#[cfg(not(target_os = "linux"))]
let _ = map;
}
Ok(())
}
}
#[cfg(not(target_arch = "wasm32"))]
pub(crate) struct PreparedRows {
file: std::io::BufWriter<std::fs::File>,
values: usize,
expected: usize,
dim: usize,
}
#[cfg(not(target_arch = "wasm32"))]
impl PreparedRows {
pub(crate) fn new(count: usize, dim: usize) -> std::io::Result<Self> {
let invalid =
|| std::io::Error::new(std::io::ErrorKind::InvalidInput, "prepared row geometry");
if dim == 0 || count >= NONE as usize {
return Err(invalid());
}
let expected = count.checked_mul(dim).ok_or_else(invalid)?;
expected
.checked_mul(4)
.filter(|&n| n <= isize::MAX as usize)
.ok_or_else(invalid)?;
Ok(Self {
file: std::io::BufWriter::with_capacity(1 << 20, tempfile::tempfile()?),
values: 0,
expected,
dim,
})
}
pub(crate) fn append(&mut self, rows: &[f32]) -> std::io::Result<()> {
use std::io::Write;
let end = self
.values
.checked_add(rows.len())
.filter(|&n| n <= self.expected);
if !rows.len().is_multiple_of(self.dim) || end.is_none() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"prepared row count",
));
}
self.file.write_all(bytemuck::cast_slice(rows))?;
self.values = end.expect("checked row count");
Ok(())
}
pub(crate) fn build(self, params: &BuildParams) -> Result<PreparedGraph, HmsError> {
let dim = self.dim;
self.build_with(params, |data, norms| {
build_structure(data, norms, dim, params)
})
}
fn build_with(
self,
params: &BuildParams,
structure: impl FnOnce(&[f32], &[f32]) -> Result<GraphStructure, HmsError>,
) -> Result<PreparedGraph, HmsError> {
let io = |error: std::io::Error| HmsError::StorageFailure {
context: format!("prepared graph rows: {error}"),
};
if self.values != self.expected || self.values == 0 {
return Err(invalid("incomplete or empty prepared rows".into()));
}
let file = self.file.into_inner().map_err(|e| io(e.into_error()))?;
let mut map = unsafe { memmap2::MmapMut::map_mut(&file) }.map_err(io)?;
let data: &mut [f32] = bytemuck::cast_slice_mut(&mut map[..self.expected * 4]);
validate_build(data, self.dim, params)?;
let norms = prepare_rows(data, self.dim, params.metric);
map.flush().map_err(io)?;
let vectors = map.make_read_only().map_err(io)?;
#[cfg(unix)]
vectors.advise(memmap2::Advice::Random).map_err(io)?;
let data: &[f32] = bytemuck::cast_slice(&vectors[..self.expected * 4]);
let structure = structure(data, &norms)?;
Ok(PreparedGraph {
vectors,
norms,
dim: self.dim,
metric: params.metric,
degree: params.degree,
structure,
})
}
}
#[cfg(not(target_arch = "wasm32"))]
pub(crate) struct PreparedGraph {
vectors: memmap2::Mmap,
norms: Vec<f32>,
dim: usize,
metric: Metric,
degree: usize,
structure: GraphStructure,
}
#[cfg(not(target_arch = "wasm32"))]
impl PreparedGraph {
fn view(&self) -> GraphView<'_> {
let vectors: &[f32] = bytemuck::cast_slice(&self.vectors[..]);
GraphView {
n: vectors.len() / self.dim,
dim: self.dim,
degree: self.degree,
metric: self.metric,
vectors,
norms: &self.norms,
adj: &self.structure.adj,
entry: self.structure.entry,
upper_ids: &self.structure.upper_ids,
layers: &self.structure.layers,
backing: Some(&self.vectors),
}
}
}
fn prepare_rows(data: &mut [f32], dim: usize, metric: Metric) -> Vec<f32> {
match metric {
Metric::Angular => {
data.par_chunks_mut(dim).for_each(normalize_row);
Vec::new()
}
Metric::Euclidean => data.par_chunks(dim).map(|x| dot(x, x)).collect(),
}
}
fn normalize_row(x: &mut [f32]) {
let squared = dot(x, x);
if squared.is_finite() && squared >= f32::MIN_POSITIVE {
let norm = squared.sqrt();
x.iter_mut().for_each(|o| *o /= norm);
return;
}
let largest = x.iter().fold(0.0f32, |m, v| m.max(v.abs()));
if largest > 0.0 && largest.is_finite() {
x.iter_mut().for_each(|o| *o /= largest);
let norm = dot(x, x).sqrt();
x.iter_mut().for_each(|o| *o /= norm);
}
}
fn norms_opt(norms: &[f32]) -> Option<&[f32]> {
(!norms.is_empty()).then_some(norms)
}
impl Graph {
fn view(&self) -> GraphView<'_> {
GraphView {
n: self.n,
dim: self.dim,
degree: self.degree,
metric: self.metric,
vectors: &self.vectors,
norms: &self.norms,
adj: &self.adj,
entry: self.entry,
upper_ids: &self.upper_ids,
layers: &self.layers,
backing: None,
}
}
fn rows(&self) -> Rows<'_> {
Rows {
data: &self.vectors,
dim: self.dim,
norms: norms_opt(&self.norms),
codes: None,
quantized: false,
}
}
pub fn metric(&self) -> Metric {
self.metric
}
pub(crate) fn neighbors(&self, v: usize) -> &[u32] {
self.adj.get(v)
}
pub fn structure_hash(&self) -> u64 {
let mut h = Fnv::new();
h.word(self.entry);
for list in self.adj.iter() {
h.word(list.len() as u32);
list.iter().for_each(|&u| h.word(u));
}
for layer in &self.layers {
h.word(layer.len() as u32);
for list in layer {
h.word(list.len() as u32);
list.iter().for_each(|&u| h.word(u));
}
}
h.finish()
}
pub fn build(data: &[f32], dim: usize, params: &BuildParams) -> Result<Self, HmsError> {
Self::build_owned(data.to_vec(), dim, params)
}
pub fn build_owned(
mut data: Vec<f32>,
dim: usize,
params: &BuildParams,
) -> Result<Self, HmsError> {
validate_build(&data, dim, params)?;
let norms = prepare_rows(&mut data, dim, params.metric);
let structure = build_structure(&data, &norms, dim, params)?;
Ok(Self {
n: data.len() / dim,
dim,
degree: params.degree,
metric: params.metric,
vectors: data,
norms,
adj: structure.adj,
entry: structure.entry,
upper_ids: structure.upper_ids,
layers: structure.layers,
build: params.clone(),
})
}
}
fn validate_build(data: &[f32], dim: usize, params: &BuildParams) -> Result<(), HmsError> {
if dim == 0 || !data.len().is_multiple_of(dim) {
return Err(invalid(format!(
"ragged data: {} values are not a multiple of dimension {dim}",
data.len()
)));
}
let n = data.len() / dim;
if n == 0 {
return Err(invalid("empty index: no vectors to build from".into()));
}
if n >= NONE as usize {
return Err(invalid(format!(
"index too large: {n} vectors, the limit is {}",
NONE - 1
)));
}
if let Some(at) = data.iter().position(|v| !v.is_finite()) {
return Err(invalid(format!(
"non-finite value in row {} (coordinate {})",
at / dim,
at % dim
)));
}
if params.build_ef == 0 {
return Err(invalid("build_ef must be positive".into()));
}
let r = params.degree;
if r == 0 || !r.is_multiple_of(2) {
return Err(invalid(format!(
"degree must be positive and even, not {r}"
)));
}
Ok(())
}
fn build_structure(
vectors: &[f32],
norms: &[f32],
dim: usize,
params: &BuildParams,
) -> Result<GraphStructure, HmsError> {
let tsetup = prof::start();
let n = vectors.len() / dim;
let r = params.degree;
let metric = params.metric;
if let Some(at) = norms.iter().position(|v| !v.is_finite()) {
return Err(invalid(format!("row {at} overflows its squared norm")));
}
let codes = (params.quantized_search || dim >= BUILD_CODES_MIN_DIM)
.then(|| BuildCodes::new(vectors, dim, norms_opt(norms)));
let rows = Rows {
data: vectors,
dim,
norms: norms_opt(norms),
codes: codes.as_ref(),
quantized: params.quantized_search,
};
let entry = medoid(&rows, n, metric);
let mut rng = SplitMix(params.seed ^ 0xA5A5_A5A5);
let mut order: Vec<u32> = (0..n as u32).collect();
for i in (1..n).rev() {
order.swap(i, (rng.next_u64() % (i as u64 + 1)) as usize);
}
let mut graph = Adjacency::new(n, r);
prof::report("setup", tsetup);
let first_ef = params.build_ef.min(FIRST_PASS_EF);
for (l, alpha) in [(first_ef, 1.0), (params.build_ef, params.alpha)] {
let t = prof::start();
rows.insert_pass(&mut graph, &order, entry, l, alpha, r);
prof::report("pass", t);
}
let t = prof::start();
rows.fill(&mut graph, r);
prof::report("fill", t);
let tup = prof::start();
let layers = upper_layers(&rows, &order, params);
prof::report("upper", tup);
prof::graph_hash((&graph, &layers, entry));
let upper_ids = order[..n / LAYER_RATIO].to_vec();
Ok(GraphStructure {
adj: graph,
entry,
upper_ids,
layers,
})
}
fn medoid(rows: &Rows<'_>, n: usize, metric: Metric) -> u32 {
medoid_by(rows, n, |i| i, metric)
}
fn medoid_by(rows: &Rows<'_>, n: usize, row_at: impl Fn(u32) -> u32 + Sync, metric: Metric) -> u32 {
let dim = rows.dim;
let mut centroid = vec![0f32; dim];
for i in 0..n as u32 {
let x = rows.row(row_at(i));
centroid.iter_mut().zip(x).for_each(|(c, v)| *c += v);
}
match metric {
Metric::Angular => {
(0..n as u32)
.into_par_iter()
.map(|i| (dot(rows.row(row_at(i)), ¢roid), i))
.reduce(
|| (f32::NEG_INFINITY, NONE),
|a, b| {
if b.0 > a.0 || (b.0 == a.0 && b.1 < a.1) {
b
} else {
a
}
},
)
.1
}
Metric::Euclidean => {
centroid.iter_mut().for_each(|c| *c /= n as f32);
(0..n as u32)
.into_par_iter()
.map(|i| {
(
rows.norms.expect("Euclidean rows carry norms")[row_at(i) as usize]
- 2.0 * dot(rows.row(row_at(i)), ¢roid),
i,
)
})
.reduce(
|| (f32::INFINITY, NONE),
|a, b| {
if b.0 < a.0 || (b.0 == a.0 && b.1 < a.1) {
b
} else {
a
}
},
)
.1
}
}
}
fn upper_layers(rows: &Rows<'_>, order: &[u32], params: &BuildParams) -> Vec<Vec<Vec<u32>>> {
let dim = rows.dim;
let mut layers = Vec::new();
let mut size = order.len() / LAYER_RATIO;
while size >= MIN_LAYER {
let sub: Vec<f32> = order[..size]
.iter()
.flat_map(|&g| rows.row(g).iter().copied())
.collect();
let sub_norms: Vec<f32> = match rows.norms {
Some(n) => order[..size].iter().map(|&g| n[g as usize]).collect(),
None => Vec::new(),
};
let sub_codes = (params.quantized_search || dim >= BUILD_CODES_MIN_DIM)
.then(|| BuildCodes::new(&sub, dim, norms_opt(&sub_norms)));
let sub_rows = Rows {
data: &sub,
dim,
norms: norms_opt(&sub_norms),
codes: sub_codes.as_ref(),
quantized: params.quantized_search,
};
let local: Vec<u32> = (0..size as u32).collect();
let mut g = Adjacency::new(size, UPPER_DEGREE);
sub_rows.insert_pass(
&mut g,
&local,
0,
params.build_ef,
params.alpha,
UPPER_DEGREE,
);
layers.push(g.to_lists());
size /= LAYER_RATIO;
}
layers
}
struct Fnv(u64);
impl Fnv {
fn new() -> Self {
Self(0xcbf2_9ce4_8422_2325)
}
fn word(&mut self, w: u32) {
self.0 = (self.0 ^ u64::from(w)).wrapping_mul(0x0000_0100_0000_01b3);
}
fn finish(&self) -> u64 {
self.0
}
}
fn upper_bytes(upper_ids: &[u32], layers: &[Vec<Vec<u32>>]) -> usize {
let adj: usize = layers.iter().flatten().map(|a| a.len() * 4).sum();
upper_ids.len() * 4 + adj
}
impl QGraph {
pub fn len(&self) -> usize {
self.n
}
pub fn is_empty(&self) -> bool {
self.n == 0
}
pub fn index_bytes(&self) -> usize {
self.blocks.len() * 8 + upper_bytes(&self.upper_ids, &self.layers)
}
fn vector(&self, v: u32) -> &[f32] {
&as_f32(&self.block(v)[..self.codes_off])[..self.dim]
}
fn descend(&self, pq: &[f32]) -> (u32, usize) {
if self.layers.is_empty() {
return (self.entry, 0);
}
let mut p = 0u32;
let mut best = -dot(pq, self.vector(self.upper_ids[0]));
let mut evals = 1;
for layer in self.layers.iter().rev() {
loop {
let mut moved = false;
for &u in &layer[p as usize] {
let d = -dot(pq, self.vector(self.upper_ids[u as usize]));
evals += 1;
if d < best {
best = d;
p = u;
moved = true;
}
}
if !moved {
break;
}
}
}
(self.upper_ids[p as usize], evals)
}
pub fn searcher(&self) -> Searcher<'_> {
Searcher {
index: self,
visited: Visited::new(self.n),
pq: vec![0.0; self.rotation.padded()],
code: QueryCode::new(self.words),
raw: vec![0; self.degree],
est: vec![0.0; self.degree],
pool: Vec::new(),
results: Vec::new(),
}
}
fn block(&self, v: u32) -> &[u64] {
let v = v as usize;
&self.blocks[v * self.stride..(v + 1) * self.stride]
}
}
pub struct Searcher<'a> {
index: &'a QGraph,
visited: Visited,
pq: Vec<f32>,
code: QueryCode,
raw: Vec<u32>,
est: Vec<f32>,
pool: Vec<Candidate>,
results: Vec<(f32, u32)>,
}
impl Searcher<'_> {
pub fn search(
&mut self,
query: &[f32],
k: usize,
params: SearchParams,
out: &mut Vec<u32>,
) -> usize {
let idx = self.index;
let ef = params.ef.max(1);
idx.rotation.apply(query, &mut self.pq);
self.code.encode(&self.pq);
let (lo2, d2q, sq) = (2.0 * self.code.lo, 2.0 * self.code.delta, self.code.sum);
self.visited.next();
let (entry, mut exact) = idx.descend(query);
let mut expanded = 0;
self.visited.insert(entry);
self.pool.clear();
self.results.clear();
self.pool.push(Candidate {
dist: 0.0,
id: entry,
done: false,
});
let mut cur = 0;
while cur < self.pool.len() {
if params.max_exact > 0 && expanded >= params.max_exact {
break;
}
let v = self.pool[cur].id;
if let Some(next) = self.pool[cur + 1..].iter().find(|c| !c.done) {
kernels::prefetch(idx.block(next.id));
}
let b = idx.block(v);
let ipv = dot(query, &as_f32(&b[..idx.codes_off])[..idx.dim]);
exact += 1;
expanded += 1;
push_top(&mut self.results, k, -ipv, v);
let mut pos = self.pool.partition_point(|c| c.dist <= -ipv);
if pos > cur {
pos -= 1;
self.pool[cur..=pos].rotate_left(1);
} else {
self.pool[pos..=cur].rotate_right(1);
}
self.pool[pos] = Candidate {
dist: -ipv,
id: v,
done: true,
};
kernels::raw_estimates(
&b[idx.codes_off..idx.factors_off],
&self.code.planes,
idx.words,
&mut self.raw,
);
let fac = as_f32(&b[idx.factors_off..idx.ids_off]);
let (kf, rest) = fac.split_at(idx.degree);
let (mf, pf) = rest.split_at(idx.degree);
for ((((e, &kj), &mj), &pj), &rj) in
self.est.iter_mut().zip(kf).zip(mf).zip(pf).zip(&self.raw)
{
*e = -(ipv + kj + mj * (lo2 * pj - sq + d2q * rj as f32));
}
let mut best = usize::MAX;
for (&u, &d) in as_u32(&b[idx.ids_off..]).iter().zip(&self.est) {
if u == NONE || !self.visited.insert(u) {
continue;
}
if self.pool.len() >= ef && d >= self.pool[ef - 1].dist {
continue;
}
let pos = self.pool.partition_point(|c| c.dist <= d);
self.pool.insert(
pos,
Candidate {
dist: d,
id: u,
done: false,
},
);
self.pool.truncate(ef);
best = best.min(pos);
}
cur = best.min(cur);
while cur < self.pool.len() && self.pool[cur].done {
cur += 1;
}
}
out.clear();
out.extend(self.results.iter().map(|r| r.1));
exact
}
}
fn push_top(top: &mut Vec<(f32, u32)>, k: usize, dist: f32, id: u32) {
if top.len() >= k && top.last().is_none_or(|l| dist >= l.0) {
return;
}
let pos = top.partition_point(|c| c.0 <= dist);
top.insert(pos, (dist, id));
top.truncate(k);
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(feature = "experimental-build-search")]
#[test]
fn build_pool_retirement_matches_ordered_merge_without_growth() {
for limit in [
1usize, 2, 3, 7, 8, 9, 31, 32, 33, 63, 64, 65, 199, 200, 201, 499, 500, 501,
] {
for len in [0, limit - 1, limit] {
let original: Vec<_> = (0..len)
.map(|id| Candidate {
dist: (id / 3) as f32,
id: id as u32,
done: id % 3 == 0,
})
.collect();
let bound = original.last().map_or(0.0, |c| c.dist);
for dist in [
f32::NEG_INFINITY,
-1.0,
-0.0,
0.0,
0.25,
1.0,
bound - 0.25,
bound,
bound + 0.25,
f32::INFINITY,
f32::from_bits(0x7fc0_1234),
] {
let candidate = Candidate {
dist,
id: limit as u32,
done: false,
};
let expected_pos = if len == limit && dist >= bound {
None
} else {
Some(original.iter().take_while(|c| c.dist <= dist).count())
};
let expected: Vec<_> = match expected_pos {
Some(pos) => original[..pos]
.iter()
.copied()
.chain(std::iter::once(candidate))
.chain(original[pos..].iter().copied())
.take(limit)
.collect(),
None => original.clone(),
};
let mut pool = Vec::with_capacity(limit);
pool.extend_from_slice(&original);
let capacity = pool.capacity();
assert_eq!(
insert_build_candidate(&mut pool, candidate, limit),
expected_pos
);
assert_eq!(
pool.capacity(),
capacity,
"limit={limit}, len={len}, dist={dist}"
);
let bits = |entries: &[Candidate]| {
entries
.iter()
.map(|c| (c.dist.to_bits(), c.id, c.done))
.collect::<Vec<_>>()
};
assert_eq!(
bits(&pool),
bits(&expected),
"limit={limit}, len={len}, dist={dist}"
);
}
}
}
}
#[test]
fn angular_rows_normalize_at_extreme_magnitudes() {
let mut rows = vec![
1e20f32, 1e20, 0.0, 0.0, 1e-30, 1e-30, 1e-30, 1e-30, 3.0, 4.0, 0.0, 0.0,
];
prepare_rows(&mut rows, 4, Metric::Angular);
let half = std::f32::consts::FRAC_1_SQRT_2;
for (got, want) in rows[..8]
.iter()
.zip([half, half, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5])
{
assert!((got - want).abs() < 1e-6, "{rows:?}");
}
assert_eq!(&rows[8..], &[0.6, 0.8, 0.0, 0.0]);
}
use kernels::raw_estimates;
#[test]
fn build_visited_reuses_words_across_resets_and_graph_sizes() {
let mut visited = BuildVisited::default();
for n in [4097, 129, 128, 127, 65, 64, 63, 1, 0, 1, 63, 64, 65, 4097] {
for _ in 0..3 {
visited.reset(n);
let mut reference = std::collections::HashSet::new();
for id in (0..n as u32).rev().chain(0..n as u32) {
assert_eq!(visited.insert(id), reference.insert(id));
}
assert_eq!(visited.touched.len(), n.div_ceil(64));
let mut unique = std::collections::HashSet::new();
assert!(visited
.touched
.iter()
.all(|&word| visited.bits[word as usize] != 0 && unique.insert(word)));
visited.reset(n);
assert!(visited.touched.is_empty());
assert!(visited.bits.iter().all(|&word| word == 0));
}
}
}
#[test]
fn build_visited_records_each_sparse_word_once() {
let mut visited = BuildVisited::default();
visited.reset(8_421_866);
for id in [0, 1, 63, 64, 65, 127, 128, 8_421_865] {
assert!(visited.insert(id));
assert!(!visited.insert(id));
}
assert_eq!(visited.touched, [0, 1, 2, 131_591]);
visited.reset(0);
visited.reset(8_421_866);
for id in [0, 63, 64, 127, 128, 8_421_865] {
assert!(visited.insert(id));
}
}
#[test]
#[ignore = "synthetic visited-reset timing diagnostic; run in isolation with --release"]
fn build_visited_reset_workload() {
use std::hint::black_box;
use std::time::{Duration, Instant};
let n = 8_421_866usize;
for (visits, iterations) in [(64, 20_000), (1024, 10_000), (16_384, 1000), (n, 4)] {
let ids: Vec<_> = (0..visits).map(|i| (i * 104_729 % n) as u32).collect();
let words: std::collections::HashSet<_> = ids.iter().map(|id| id / 64).collect();
let mut visited = BuildVisited::default();
visited.reset(n);
for &id in &ids {
assert!(visited.insert(id));
}
let mut reset_time = Duration::ZERO;
let start = Instant::now();
for _ in 0..iterations {
let reset_start = Instant::now();
black_box(&mut visited).reset(black_box(n));
reset_time += reset_start.elapsed();
for &id in &ids {
black_box(visited.insert(black_box(id)));
}
}
println!(
"{{\"vertices\":{n},\"visits\":{visits},\"touched_words\":{},\"full_words\":{},\"iterations\":{iterations},\"reset_ns\":{},\"total_ns\":{}}}",
words.len(), n.div_ceil(64), reset_time.as_nanos(), start.elapsed().as_nanos()
);
}
}
pub(super) fn gaussian(rng: &mut SplitMix) -> f32 {
let u = rng.next_f32().max(1e-7);
let v = rng.next_f32();
(-2.0 * u.ln()).sqrt() * (std::f32::consts::TAU * v).cos()
}
fn unit(mut x: Vec<f32>) -> Vec<f32> {
let n = dot(&x, &x).sqrt();
x.iter_mut().for_each(|v| *v /= n);
x
}
#[test]
fn edge_estimator_is_unbiased_over_rotations() {
let d = 100;
let mut rng = SplitMix(42);
for bits in [0, 0, 0, 512, 512, 512] {
let q = unit((0..d).map(|_| gaussian(&mut rng)).collect());
let v = unit((0..d).map(|_| gaussian(&mut rng)).collect());
let u = unit(v.iter().map(|x| x + 0.15 * gaussian(&mut rng)).collect());
let truth = dot(&q, &u);
let (mut float_sum, mut quant_sum) = (0f64, 0f64);
let seeds = 400;
for seed in 0..seeds {
let rot = Rotation::new(d, bits, seed);
let p = rot.padded();
let (mut pq, mut pv, mut pu) = (vec![0.0; p], vec![0.0; p], vec![0.0; p]);
rot.apply(&q, &mut pq);
rot.apply(&v, &mut pv);
rot.apply(&u, &mut pu);
let mut code = vec![0u64; p / 64];
let (k, m, pop) = edge_factors(&pv, &pu, &mut code);
let ipv = dot(&pq, &pv);
let sq: f32 = pq
.iter()
.enumerate()
.map(|(i, x)| {
if code[i / 64] >> (i % 64) & 1 == 1 {
*x
} else {
-*x
}
})
.sum();
float_sum += f64::from(ipv + k + m * sq);
let mut qc = QueryCode::new(p / 64);
qc.encode(&pq);
let mut raw = [0u32];
raw_estimates(&code, &qc.planes, p / 64, &mut raw);
let est =
ipv + k + m * (2.0 * qc.lo * pop - qc.sum + 2.0 * qc.delta * raw[0] as f32);
quant_sum += f64::from(est);
}
let float_mean = float_sum / f64::from(seeds as u32);
let quant_mean = quant_sum / f64::from(seeds as u32);
assert!(
(float_mean - f64::from(truth)).abs() < 0.02,
"{float_mean} vs {truth}"
);
assert!(
(quant_mean - f64::from(truth)).abs() < 0.03,
"{quant_mean} vs {truth}"
);
}
}
pub(super) fn clustered(n: usize, d: usize, seed: u64) -> Vec<f32> {
let mut rng = SplitMix(seed);
let centers: Vec<Vec<f32>> = (0..20)
.map(|_| (0..d).map(|_| gaussian(&mut rng)).collect())
.collect();
(0..n)
.flat_map(|i| {
let c = ¢ers[i % centers.len()];
c.iter()
.map(|x| x + 0.4 * gaussian(&mut rng))
.collect::<Vec<_>>()
})
.collect()
}
pub(super) fn brute_force(data: &[f32], d: usize, q: &[f32], k: usize) -> Vec<u32> {
let mut s: Vec<(f32, u32)> = data
.chunks_exact(d)
.enumerate()
.map(|(i, x)| (-dot(x, q) / dot(x, x).sqrt(), i as u32))
.collect();
s.sort_by(|a, b| a.0.total_cmp(&b.0));
s.into_iter().take(k).map(|x| x.1).collect()
}
#[test]
fn recall_on_clustered_data() {
let d = 24;
let data = clustered(4000, d, 1);
let queries = clustered(100, d, 2);
let index = QGraph::build(&data, d, &BuildParams::default()).expect("build");
let mut s = index.searcher();
let mut out = Vec::new();
let mut hits = 0;
for q in queries.chunks_exact(d) {
let truth = brute_force(&data, d, q, 10);
s.search(
q,
10,
SearchParams {
ef: 64,
max_exact: 0,
},
&mut out,
);
hits += out.iter().filter(|x| truth.contains(x)).count();
}
let recall = hits as f64 / 1000.0;
assert!(recall >= 0.9, "recall {recall}");
}
#[test]
fn prune_merged_matches_prune() {
let d = 12;
let mut data = clustered(600, d, 5);
data.chunks_exact_mut(d).for_each(|x| {
let n = dot(x, x).sqrt();
x.iter_mut().for_each(|v| *v /= n);
});
let rows = Rows {
data: &data,
dim: d,
norms: None,
quantized: false,
codes: None,
};
let mut rng = SplitMix(11);
for trial in 0..200 {
let t = (rng.next_u64() % 600) as u32;
let (r, alpha) = (4 + trial % 13, [1.0, 1.2][trial % 2]);
let all = |ids: &[u32]| -> Vec<(f32, u32)> {
ids.iter().map(|&u| (rows.d2(t, u), u)).collect()
};
let pool: Vec<u32> = (0..60)
.map(|_| (rng.next_u64() % 600) as u32)
.filter(|&u| u != t)
.collect();
let mut list = rows.prune(t, all(&pool), alpha, r);
let closed = list.len();
for _ in 0..1 + trial % 9 {
let u = (rng.next_u64() % 600) as u32;
if u != t && !list.contains(&u) {
list.push(u);
}
}
assert_eq!(
rows.prune_merged(t, &list, closed, alpha, r),
rows.prune(t, all(&list), alpha, r),
"trial {trial}"
);
}
}
#[test]
fn build_codes_leave_the_graph_unchanged() {
for metric in [Metric::Angular, Metric::Euclidean] {
for d in [100, 160] {
let n = 1200;
let mut data = clustered(n, d, 3);
let mut rng = SplitMix(5);
if metric == Metric::Euclidean {
for (i, x) in data.chunks_exact_mut(d).enumerate() {
let s = 0.1 * 100f32.powf((i * 7919 % n) as f32 / n as f32);
x.iter_mut().for_each(|v| *v = s * (*v + 3.0));
}
}
for i in (0..n).step_by(7) {
let j = (rng.next_u64() % n as u64) as usize;
for k in 0..d {
data[i * d + k] = data[j * d + k] * (1.0 + 1e-6 * gaussian(&mut rng));
}
}
data[5 * d..6 * d].fill(0.0);
let norms = prepare_rows(&mut data, d, metric);
let codes = BuildCodes::new(&data, d, norms_opt(&norms));
let build = |codes: Option<&BuildCodes>| {
let rows = Rows {
data: &data,
dim: d,
norms: norms_opt(&norms),
codes,
quantized: false,
};
let order: Vec<u32> = (0..n as u32).rev().collect();
let mut graph = Adjacency::new(n, 12);
for alpha in [1.0, 1.2] {
rows.insert_pass(&mut graph, &order, 0, 40, alpha, 12);
}
rows.fill(&mut graph, 12);
graph
};
assert_eq!(build(Some(&codes)), build(None), "{metric:?}, d = {d}");
}
}
}
#[test]
fn build_and_search_are_deterministic() {
let d = 16;
let data = clustered(3000, d, 3);
let params = BuildParams {
seed: 9,
..BuildParams::default()
};
let a = QGraph::build(&data, d, ¶ms).expect("build");
let b = QGraph::build(&data, d, ¶ms).expect("build");
assert_eq!(a.entry, b.entry);
assert!(a.blocks == b.blocks);
let (mut sa, mut sb) = (a.searcher(), b.searcher());
let (mut oa, mut ob) = (Vec::new(), Vec::new());
let p = SearchParams {
ef: 32,
max_exact: 0,
};
for q in data.chunks_exact(d).take(50) {
sa.search(q, 10, p, &mut oa);
sb.search(q, 10, p, &mut ob);
assert_eq!(oa, ob);
}
}
#[test]
fn build_reproduces_the_pinned_graphs() {
let cases = [
(Metric::Angular, false, 16, 0x2b16_3721_02ba_f39b),
(Metric::Euclidean, true, 128, 0xe014_c782_9bed_e9aa),
];
for (metric, quantized_search, d, want) in cases {
let data = clustered(2500, d, 11);
let params = BuildParams {
metric,
quantized_search,
degree: 16,
build_ef: 48,
seed: 7,
..BuildParams::default()
};
let g = Graph::build(&data, d, ¶ms).expect("build");
assert_eq!(g.structure_hash(), want, "{metric:?}");
}
}
#[test]
fn exact_budget_caps_work() {
let d = 16;
let data = clustered(2000, d, 4);
let index = QGraph::build(&data, d, &BuildParams::default()).expect("build");
let mut s = index.searcher();
let mut out = Vec::new();
let q = &data[..d];
let capped = s.search(
q,
10,
SearchParams {
ef: 128,
max_exact: 20,
},
&mut out,
);
assert_eq!(out.len(), 10);
let free = s.search(
q,
10,
SearchParams {
ef: 128,
max_exact: 0,
},
&mut out,
);
assert!(capped < free);
assert_eq!(capped - index.descend(q).1, 20);
}
}