osobhandClaude Opus 5
97e65f2adf
The int8-quantised HNSW index compared vectors with a scalar loop that the compiler vectorised for the x86-64 baseline (SSE2), while the f32 path it was measured against goes through `clawhdf5-accel` and runs AVX2. So the ~13% throughput cost recorded for `quantized_index` was a missing kernel rather than a property of int8. `clawhdf5_accel::dot_i8` adds a scalar fallback and an AVX2 path: sign-extend each 16-byte half to i16, then `madd_epi16`, which multiplies and sums adjacent pairs straight into i32 lanes. It is dispatched through the same detected backend as the f32 kernels, and the index now calls it. Integer arithmetic, so the SIMD path must agree with scalar bit for bit — tested at lengths that are and are not multiples of the block, and at the -128 extreme for overflow. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
2426 lines
86 KiB
Rust
2426 lines
86 KiB
Rust
//! HNSW index implementation with HDF5 serialization.
|
|
|
|
use std::collections::BinaryHeap;
|
|
|
|
use clawhdf5_format::attribute::extract_attributes_full;
|
|
use clawhdf5_format::data_layout::DataLayout;
|
|
use clawhdf5_format::data_read::{read_as_f32, read_as_i32, read_raw_data_full};
|
|
use clawhdf5_format::dataspace::Dataspace;
|
|
use clawhdf5_format::datatype::Datatype;
|
|
use clawhdf5_format::error::FormatError;
|
|
use clawhdf5_format::file_writer::{AttrValue, FileWriter as FmtWriter};
|
|
use clawhdf5_format::filter_pipeline::FilterPipeline;
|
|
use clawhdf5_format::group_v2::resolve_path_any;
|
|
use clawhdf5_format::message_type::MessageType;
|
|
use clawhdf5_format::object_header::ObjectHeader;
|
|
use clawhdf5_format::signature::find_signature;
|
|
use clawhdf5_format::superblock::Superblock;
|
|
use clawhdf5_io::FileWriter as IoFileWriter;
|
|
|
|
/// Distance metric for the HNSW index.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum DistanceMetric {
|
|
/// L2 (Euclidean) distance.
|
|
L2,
|
|
/// Cosine distance (1 - cosine_similarity).
|
|
Cosine,
|
|
}
|
|
|
|
impl DistanceMetric {
|
|
fn as_str(self) -> &'static str {
|
|
match self {
|
|
DistanceMetric::L2 => "l2",
|
|
DistanceMetric::Cosine => "cosine",
|
|
}
|
|
}
|
|
|
|
fn from_str(s: &str) -> Option<Self> {
|
|
match s {
|
|
"l2" => Some(DistanceMetric::L2),
|
|
"cosine" => Some(DistanceMetric::Cosine),
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compute distance between two vectors using the given metric.
|
|
///
|
|
/// Delegates to `clawhdf5-accel`'s runtime-dispatched SIMD kernels (AVX2 on
|
|
/// x86_64, NEON on aarch64, portable scalar fallback elsewhere) — this is
|
|
/// the hottest loop in both HNSW build and every `hybrid_search` query.
|
|
fn compute_distance(a: &[f32], b: &[f32], metric: DistanceMetric) -> f32 {
|
|
match metric {
|
|
DistanceMetric::L2 => clawhdf5_accel::l2_distance(a, b),
|
|
// Both sides are unit length (see `prepare`), so cosine similarity is
|
|
// the plain dot product. Computing it as dot / (|a| * |b|) re-derived
|
|
// both norms on every call — three reductions instead of one, in the
|
|
// innermost loop of both build and search.
|
|
DistanceMetric::Cosine => 1.0 - clawhdf5_accel::dot_product(a, b),
|
|
}
|
|
}
|
|
|
|
/// Put a vector in the form the index stores and compares: unit length for the
|
|
/// cosine metric, unchanged for L2. A zero vector stays zero, giving distance 1
|
|
/// to everything — what the cosine kernel reports for a degenerate input.
|
|
fn prepare(mut v: Vec<f32>, metric: DistanceMetric) -> Vec<f32> {
|
|
if metric == DistanceMetric::Cosine {
|
|
let norm = clawhdf5_accel::vector_norm(&v);
|
|
if norm > f32::EPSILON {
|
|
let inv = 1.0 / norm;
|
|
v.iter_mut().for_each(|x| *x *= inv);
|
|
} else {
|
|
v.iter_mut().for_each(|x| *x = 0.0);
|
|
}
|
|
}
|
|
v
|
|
}
|
|
|
|
/// Assign a random level to a new node based on the HNSW probability distribution.
|
|
///
|
|
/// Uses a deterministic approach based on the node index for reproducibility.
|
|
fn assign_level(node_id: usize, m: usize) -> usize {
|
|
let ml = 1.0 / (m as f64).ln();
|
|
// Use a simple hash-based pseudo-random for reproducibility
|
|
let hash = splitmix64(node_id as u64);
|
|
let uniform = (hash >> 11) as f64 / (1u64 << 53) as f64;
|
|
(-uniform.ln() * ml).floor() as usize
|
|
}
|
|
|
|
/// Simple splitmix64 hash for deterministic level assignment.
|
|
fn splitmix64(mut x: u64) -> u64 {
|
|
x = x.wrapping_add(0x9e3779b97f4a7c15);
|
|
x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
|
|
x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
|
|
x ^ (x >> 31)
|
|
}
|
|
|
|
/// Candidate neighbor for priority queue operations.
|
|
#[derive(Debug, Clone)]
|
|
struct Candidate {
|
|
id: usize,
|
|
distance: f32,
|
|
}
|
|
|
|
impl PartialEq for Candidate {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.distance.to_bits() == other.distance.to_bits() && self.id == other.id
|
|
}
|
|
}
|
|
|
|
impl Eq for Candidate {}
|
|
|
|
impl PartialOrd for Candidate {
|
|
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl Ord for Candidate {
|
|
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
|
// Reverse ordering for min-heap behavior
|
|
other
|
|
.distance
|
|
.partial_cmp(&self.distance)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
}
|
|
}
|
|
|
|
/// Max-heap candidate (furthest first).
|
|
#[derive(Debug, Clone)]
|
|
struct FarCandidate {
|
|
id: usize,
|
|
distance: f32,
|
|
}
|
|
|
|
impl PartialEq for FarCandidate {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.distance.to_bits() == other.distance.to_bits() && self.id == other.id
|
|
}
|
|
}
|
|
|
|
impl Eq for FarCandidate {}
|
|
|
|
impl PartialOrd for FarCandidate {
|
|
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl Ord for FarCandidate {
|
|
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
|
self.distance
|
|
.partial_cmp(&other.distance)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
}
|
|
}
|
|
|
|
/// How the index keeps its copy of the vectors.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
|
pub enum Storage {
|
|
/// Exactly as given: `dim * 4` bytes per vector.
|
|
#[default]
|
|
Float32,
|
|
/// Each component scaled to an `i8`: `dim` bytes per vector, a quarter of
|
|
/// the space, at some cost in precision.
|
|
///
|
|
/// Only meaningful for [`DistanceMetric::Cosine`]: rows are stored
|
|
/// unit-length, so a quantised dot product reconstructs the similarity
|
|
/// directly. Requesting it for `L2` keeps `Float32`, because an L2
|
|
/// distance cannot be recovered from a dot product alone.
|
|
Int8,
|
|
}
|
|
|
|
/// The index's copy of the vectors, flat and row-major.
|
|
#[derive(Debug, Clone)]
|
|
enum Vectors {
|
|
F32 {
|
|
dim: usize,
|
|
flat: Vec<f32>,
|
|
},
|
|
/// `flat[i * dim + j]` is component `j` of vector `i` divided by
|
|
/// `scales[i]`; multiplying back recovers it.
|
|
///
|
|
/// The scale is per row rather than global. A unit-length row in `d`
|
|
/// dimensions has components around `1/sqrt(d)`, so a fixed `[-1, 1]`
|
|
/// scale spends fewer than 12 of the 255 levels on a 128-dimensional
|
|
/// vector and the reconstruction error swamps the gaps between near
|
|
/// neighbours — measured at 0.35 top-10 overlap with the exact ranking.
|
|
/// Scaling each row by its own largest component uses the full range.
|
|
Int8 {
|
|
dim: usize,
|
|
flat: Vec<i8>,
|
|
scales: Vec<f32>,
|
|
},
|
|
}
|
|
|
|
/// Levels either side of zero. 127, not 128, so the range is symmetric.
|
|
const INT8_LEVELS: f32 = 127.0;
|
|
|
|
/// Quantise one row, returning the codes and the scale that inverts them.
|
|
fn quantise_row(v: &[f32], out: &mut Vec<i8>) -> f32 {
|
|
let max_abs = v.iter().fold(0.0f32, |m, x| m.max(x.abs()));
|
|
if max_abs <= f32::MIN_POSITIVE {
|
|
out.extend(core::iter::repeat_n(0i8, v.len()));
|
|
return 0.0;
|
|
}
|
|
let inv = INT8_LEVELS / max_abs;
|
|
out.extend(
|
|
v.iter()
|
|
.map(|x| (x * inv).round().clamp(-INT8_LEVELS, INT8_LEVELS) as i8),
|
|
);
|
|
max_abs / INT8_LEVELS
|
|
}
|
|
|
|
impl Vectors {
|
|
fn new(dim: usize, storage: Storage, metric: DistanceMetric) -> Self {
|
|
match storage {
|
|
Storage::Int8 if metric == DistanceMetric::Cosine => Vectors::Int8 {
|
|
dim,
|
|
flat: Vec::new(),
|
|
scales: Vec::new(),
|
|
},
|
|
_ => Vectors::F32 {
|
|
dim,
|
|
flat: Vec::new(),
|
|
},
|
|
}
|
|
}
|
|
|
|
fn dim(&self) -> usize {
|
|
match self {
|
|
Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim,
|
|
}
|
|
}
|
|
|
|
fn storage(&self) -> Storage {
|
|
match self {
|
|
Vectors::F32 { .. } => Storage::Float32,
|
|
Vectors::Int8 { .. } => Storage::Int8,
|
|
}
|
|
}
|
|
|
|
fn len(&self) -> usize {
|
|
let dim = self.dim();
|
|
if dim == 0 {
|
|
return 0;
|
|
}
|
|
match self {
|
|
Vectors::F32 { flat, .. } => flat.len() / dim,
|
|
Vectors::Int8 { flat, .. } => flat.len() / dim,
|
|
}
|
|
}
|
|
|
|
/// Set the row width, for a store seeded empty by `new`.
|
|
fn set_dim(&mut self, new_dim: usize) {
|
|
match self {
|
|
Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim = new_dim,
|
|
}
|
|
}
|
|
|
|
fn push(&mut self, vector: &[f32]) {
|
|
match self {
|
|
Vectors::F32 { flat, .. } => flat.extend_from_slice(vector),
|
|
Vectors::Int8 { flat, scales, .. } => scales.push(quantise_row(vector, flat)),
|
|
}
|
|
}
|
|
|
|
/// Row `i` as `f32`, for callers that need the values back (serialization,
|
|
/// and the f32 fast paths). Quantised rows are reconstructed, so this is
|
|
/// lossy in exactly the way the storage is.
|
|
fn row(&self, i: usize) -> Vec<f32> {
|
|
let dim = self.dim();
|
|
let start = i * dim;
|
|
match self {
|
|
Vectors::F32 { flat, .. } => flat[start..start + dim].to_vec(),
|
|
Vectors::Int8 { flat, scales, .. } => flat[start..start + dim]
|
|
.iter()
|
|
.map(|&q| f32::from(q) * scales[i])
|
|
.collect(),
|
|
}
|
|
}
|
|
|
|
/// Distance between two stored vectors.
|
|
fn dist(&self, a: usize, b: usize, metric: DistanceMetric) -> f32 {
|
|
let dim = self.dim();
|
|
match self {
|
|
Vectors::F32 { flat, .. } => {
|
|
let (x, y) = (a * dim, b * dim);
|
|
compute_distance(&flat[x..x + dim], &flat[y..y + dim], metric)
|
|
}
|
|
Vectors::Int8 { flat, scales, .. } => {
|
|
let (x, y) = (a * dim, b * dim);
|
|
let dot = dot_i8(&flat[x..x + dim], &flat[y..y + dim]);
|
|
1.0 - dot as f32 * scales[a] * scales[b]
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Distance from a prepared query to stored vector `i`.
|
|
fn dist_query(&self, query: &Query, i: usize, metric: DistanceMetric) -> f32 {
|
|
let dim = self.dim();
|
|
let start = i * dim;
|
|
match (self, query) {
|
|
(Vectors::F32 { flat, .. }, Query::F32(q)) => {
|
|
compute_distance(q, &flat[start..start + dim], metric)
|
|
}
|
|
(Vectors::Int8 { flat, scales, .. }, Query::Int8(q, q_scale)) => {
|
|
let dot = dot_i8(q, &flat[start..start + dim]);
|
|
1.0 - dot as f32 * q_scale * scales[i]
|
|
}
|
|
// Mixed forms cannot occur: `Query` is built from the same storage.
|
|
_ => f32::MAX,
|
|
}
|
|
}
|
|
|
|
/// Build a store from prepared rows.
|
|
fn from_rows(rows: &[Vec<f32>], storage: Storage, metric: DistanceMetric) -> Self {
|
|
let dim = rows.first().map_or(0, Vec::len);
|
|
let mut out = Vectors::new(dim, storage, metric);
|
|
for row in rows {
|
|
out.push(row);
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Prepare `query` for comparison against this store.
|
|
fn query(&self, query: Vec<f32>) -> Query {
|
|
match self {
|
|
Vectors::F32 { .. } => Query::F32(query),
|
|
Vectors::Int8 { .. } => {
|
|
let mut codes = Vec::with_capacity(query.len());
|
|
let scale = quantise_row(&query, &mut codes);
|
|
Query::Int8(codes, scale)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// What a layer search is measuring distance *to*: an incoming query, or a
|
|
/// node already in the index (which is what insertion compares against).
|
|
enum Target<'a> {
|
|
Query(&'a Query),
|
|
Node(usize),
|
|
}
|
|
|
|
impl Vectors {
|
|
fn dist_to(&self, target: &Target<'_>, i: usize, metric: DistanceMetric) -> f32 {
|
|
match target {
|
|
Target::Query(q) => self.dist_query(q, i, metric),
|
|
Target::Node(n) => self.dist(*n, i, metric),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A search query in whichever form the store compares against.
|
|
enum Query {
|
|
F32(Vec<f32>),
|
|
/// Codes and the scale that inverts them, as in [`Vectors::Int8`].
|
|
Int8(Vec<i8>, f32),
|
|
}
|
|
|
|
/// Sum of products, widened so it cannot overflow. Runtime-dispatched to the
|
|
/// same SIMD backend as the f32 kernels, so the two storages are compared on
|
|
/// equal terms.
|
|
#[inline]
|
|
fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
|
clawhdf5_accel::dot_i8(a, b)
|
|
}
|
|
|
|
/// Magic for [`HnswIndex::graph_to_bytes`].
|
|
const GRAPH_MAGIC: &[u8; 4] = b"CHG1";
|
|
|
|
/// On-disk format version for the serialized HNSW index.
|
|
///
|
|
/// - Version 1: original layout (`vectors`, `graph_layer_*`, `config`), no
|
|
/// deletion support and no explicit version tag.
|
|
/// - Version 2: adds a `format_version` attribute and a `deleted` bitset dataset
|
|
/// so live insert/delete state survives a save/load round-trip.
|
|
///
|
|
/// Files written before this constant existed are treated as version 1 on load.
|
|
pub const HNSW_FORMAT_VERSION: i64 = 2;
|
|
|
|
/// HNSW (Hierarchical Navigable Small World) approximate nearest neighbor index.
|
|
///
|
|
/// Supports building an index from vectors, incremental insertion and soft
|
|
/// deletion, searching for nearest neighbors, and serializing/deserializing to
|
|
/// HDF5 format.
|
|
#[derive(Debug, Clone)]
|
|
pub struct HnswIndex {
|
|
/// All vectors in the index, flat and row-major.
|
|
vectors: Vectors,
|
|
/// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs.
|
|
graph: Vec<Vec<Vec<usize>>>,
|
|
/// Soft-deletion flags, one per node. Deleted nodes remain in the graph for
|
|
/// connectivity but are never returned from [`HnswIndex::search`].
|
|
deleted: Vec<bool>,
|
|
/// Entry point node ID.
|
|
entry_point: usize,
|
|
/// Maximum number of connections per node (per layer).
|
|
m: usize,
|
|
/// Maximum connections for layer 0 (typically 2*m).
|
|
m_max0: usize,
|
|
/// ef parameter used during construction.
|
|
ef_construction: usize,
|
|
/// Maximum layer assigned to each node.
|
|
node_levels: Vec<usize>,
|
|
/// Distance metric.
|
|
metric: DistanceMetric,
|
|
}
|
|
|
|
impl HnswIndex {
|
|
/// Build an HNSW index from a set of vectors.
|
|
///
|
|
/// # Parameters
|
|
/// - `vectors`: The vectors to index. All must have the same dimension.
|
|
/// - `m`: Maximum number of connections per node (higher = more accurate, more memory).
|
|
/// - `ef_construction`: Size of the dynamic candidate list during construction.
|
|
///
|
|
/// Uses L2 distance by default. Use [`build_with_metric`] to specify the metric.
|
|
pub fn build(vectors: &[Vec<f32>], m: usize, ef_construction: usize) -> Self {
|
|
Self::build_with_metric(vectors, m, ef_construction, DistanceMetric::L2)
|
|
}
|
|
|
|
/// Build an HNSW index with a specific distance metric.
|
|
pub fn build_with_metric(
|
|
vectors: &[Vec<f32>],
|
|
m: usize,
|
|
ef_construction: usize,
|
|
metric: DistanceMetric,
|
|
) -> Self {
|
|
Self::build_with(vectors, m, ef_construction, metric, Storage::default())
|
|
}
|
|
|
|
/// Build an index, choosing how the vectors are stored.
|
|
///
|
|
/// [`Storage::Int8`] keeps them at a quarter of the size; see its docs for
|
|
/// what that costs and when it applies.
|
|
pub fn build_with(
|
|
vectors: &[Vec<f32>],
|
|
m: usize,
|
|
ef_construction: usize,
|
|
metric: DistanceMetric,
|
|
storage: Storage,
|
|
) -> Self {
|
|
assert!(!vectors.is_empty(), "cannot build index from empty vectors");
|
|
assert!(m >= 2, "m must be at least 2");
|
|
let dim = vectors[0].len();
|
|
for v in vectors {
|
|
assert_eq!(v.len(), dim, "all vectors must have the same dimension");
|
|
}
|
|
|
|
let m_max0 = m * 2;
|
|
let n = vectors.len();
|
|
let mut prepared = Vectors::new(dim, storage, metric);
|
|
for v in vectors {
|
|
prepared.push(&prepare(v.clone(), metric));
|
|
}
|
|
let vectors = &prepared;
|
|
|
|
// Assign levels to all nodes
|
|
let mut node_levels = Vec::with_capacity(n);
|
|
let mut max_level = 0;
|
|
for i in 0..n {
|
|
let level = assign_level(i, m);
|
|
if level > max_level {
|
|
max_level = level;
|
|
}
|
|
node_levels.push(level);
|
|
}
|
|
|
|
// Initialize graph layers
|
|
let num_layers = max_level + 1;
|
|
let mut graph: Vec<Vec<Vec<usize>>> = Vec::with_capacity(num_layers);
|
|
for _ in 0..num_layers {
|
|
graph.push(vec![Vec::new(); n]);
|
|
}
|
|
|
|
let mut entry_point = 0;
|
|
let mut ep_level = node_levels[0];
|
|
|
|
// Insert in batches. Each batch is planned against the graph as it
|
|
// stood when the batch began (read-only, so the plans are independent
|
|
// and run in parallel with the `parallel` feature), then linked, then
|
|
// every neighbour list that overflowed is pruned once. Pruning is ~90%
|
|
// of a build's distance evaluations, and a node that gains several
|
|
// back-links in one batch is pruned once instead of once per link.
|
|
//
|
|
// Nodes in the same batch cannot see each other while planning, so
|
|
// batches start at one node and grow only as the graph does — a batch
|
|
// is never more than a small fraction of what is already linked. The
|
|
// result is deterministic and identical with or without `parallel`.
|
|
let mut next = 1;
|
|
while next < n {
|
|
let mut end = (next + batch_len(next)).min(n);
|
|
// A node that raises the top layer becomes the new entry point and
|
|
// changes how every later node descends: give it a batch alone.
|
|
if let Some(tall) = (next..end).find(|&i| node_levels[i] > ep_level) {
|
|
end = if tall == next { next + 1 } else { tall };
|
|
}
|
|
|
|
let plans = plan_batch(
|
|
vectors,
|
|
&graph,
|
|
&node_levels,
|
|
next..end,
|
|
entry_point,
|
|
ep_level,
|
|
(m, m_max0, ef_construction),
|
|
metric,
|
|
);
|
|
|
|
let mut overflowed: Vec<(usize, usize)> = Vec::new();
|
|
for (offset, plan) in plans.into_iter().enumerate() {
|
|
let node = next + offset;
|
|
for (layer, selected) in plan {
|
|
let max_conn = if layer == 0 { m_max0 } else { m };
|
|
for &neighbor in &selected {
|
|
let list = &mut graph[layer][neighbor];
|
|
list.push(node);
|
|
if list.len() == max_conn + 1 {
|
|
overflowed.push((layer, neighbor));
|
|
}
|
|
}
|
|
graph[layer][node] = selected;
|
|
}
|
|
}
|
|
prune_overflowed(vectors, &mut graph, overflowed, (m, m_max0), metric);
|
|
|
|
for (i, &level) in node_levels.iter().enumerate().take(end).skip(next) {
|
|
if level > ep_level {
|
|
entry_point = i;
|
|
ep_level = level;
|
|
}
|
|
}
|
|
next = end;
|
|
}
|
|
|
|
Self {
|
|
vectors: prepared,
|
|
graph,
|
|
deleted: vec![false; n],
|
|
entry_point,
|
|
m,
|
|
m_max0,
|
|
ef_construction,
|
|
node_levels,
|
|
metric,
|
|
}
|
|
}
|
|
|
|
/// Create an empty index with the given parameters. Used as the starting
|
|
/// point for incremental [`HnswIndex::insert`] and as the result of
|
|
/// [`HnswIndex::compact`] when every vector has been deleted.
|
|
pub fn new(m: usize, ef_construction: usize, metric: DistanceMetric) -> Self {
|
|
Self::new_with(m, ef_construction, metric, Storage::default())
|
|
}
|
|
|
|
/// [`HnswIndex::new`], choosing how the vectors are stored.
|
|
pub fn new_with(
|
|
m: usize,
|
|
ef_construction: usize,
|
|
metric: DistanceMetric,
|
|
storage: Storage,
|
|
) -> Self {
|
|
assert!(m >= 2, "m must be at least 2");
|
|
Self {
|
|
// The dimension is set by the first insert.
|
|
vectors: Vectors::new(0, storage, metric),
|
|
graph: Vec::new(),
|
|
deleted: Vec::new(),
|
|
entry_point: 0,
|
|
m,
|
|
m_max0: m * 2,
|
|
ef_construction,
|
|
node_levels: Vec::new(),
|
|
metric,
|
|
}
|
|
}
|
|
|
|
/// Insert a single vector into the index incrementally and return its id.
|
|
///
|
|
/// The id is the vector's position in insertion order and is stable for the
|
|
/// life of the index (until [`HnswIndex::compact`] renumbers survivors).
|
|
/// Inserting into an empty index seeds the entry point.
|
|
///
|
|
/// # Panics
|
|
/// Panics if `vector`'s dimension does not match the existing vectors.
|
|
pub fn insert(&mut self, vector: Vec<f32>) -> usize {
|
|
let vector = prepare(vector, self.metric);
|
|
let id = self.vectors.len();
|
|
|
|
// Seed an empty index.
|
|
if id == 0 {
|
|
let node_level = assign_level(0, self.m);
|
|
self.vectors.set_dim(vector.len());
|
|
self.vectors.push(&vector);
|
|
self.deleted.push(false);
|
|
self.node_levels.push(node_level);
|
|
self.graph = (0..=node_level).map(|_| vec![Vec::new(); 1]).collect();
|
|
self.entry_point = 0;
|
|
return 0;
|
|
}
|
|
|
|
assert_eq!(
|
|
vector.len(),
|
|
self.vectors.dim(),
|
|
"insert dimension mismatch"
|
|
);
|
|
|
|
let node_level = assign_level(id, self.m);
|
|
self.vectors.push(&vector);
|
|
self.deleted.push(false);
|
|
self.node_levels.push(node_level);
|
|
|
|
// Grow every existing layer with an empty adjacency slot for `id`, and
|
|
// add any brand-new top layers this node introduces.
|
|
for layer in self.graph.iter_mut() {
|
|
layer.push(Vec::new());
|
|
}
|
|
while self.graph.len() <= node_level {
|
|
self.graph.push(vec![Vec::new(); id + 1]);
|
|
}
|
|
|
|
let ep_level = self.node_levels[self.entry_point];
|
|
let mut ep = self.entry_point;
|
|
|
|
// Phase 1: greedy descent from the top down to node_level + 1.
|
|
for layer in (node_level + 1..=ep_level).rev() {
|
|
ep = greedy_closest(
|
|
&self.vectors,
|
|
&self.graph[layer],
|
|
&Target::Node(id),
|
|
ep,
|
|
self.metric,
|
|
);
|
|
}
|
|
|
|
// Phase 2: search and connect from min(node_level, ep_level) down to 0.
|
|
let bottom = node_level.min(ep_level);
|
|
for layer in (0..=bottom).rev() {
|
|
let max_conn = if layer == 0 { self.m_max0 } else { self.m };
|
|
let neighbors = search_layer(
|
|
&self.vectors,
|
|
&self.graph[layer],
|
|
&Target::Node(id),
|
|
ep,
|
|
self.ef_construction,
|
|
self.metric,
|
|
None,
|
|
);
|
|
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect();
|
|
let selected = select_neighbors(&self.vectors, &scored, max_conn, self.metric);
|
|
self.graph[layer][id] = selected.clone();
|
|
link_back(
|
|
&self.vectors,
|
|
&mut self.graph[layer],
|
|
id,
|
|
&selected,
|
|
max_conn,
|
|
self.metric,
|
|
);
|
|
if !selected.is_empty() {
|
|
ep = selected[0];
|
|
}
|
|
}
|
|
|
|
// Promote the entry point if this node sits on a higher layer.
|
|
if node_level > ep_level {
|
|
self.entry_point = id;
|
|
}
|
|
|
|
id
|
|
}
|
|
|
|
/// Soft-delete the vector with the given id. The node stays in the graph so
|
|
/// traversal/connectivity is preserved, but it will never be returned from
|
|
/// [`HnswIndex::search`]. Idempotent; out-of-range ids are ignored.
|
|
///
|
|
/// Returns `true` if the id existed and was not already deleted.
|
|
pub fn mark_deleted(&mut self, id: usize) -> bool {
|
|
if id >= self.deleted.len() || self.deleted[id] {
|
|
return false;
|
|
}
|
|
self.deleted[id] = true;
|
|
true
|
|
}
|
|
|
|
/// Returns whether the vector with the given id is soft-deleted.
|
|
pub fn is_deleted(&self, id: usize) -> bool {
|
|
self.deleted.get(id).copied().unwrap_or(false)
|
|
}
|
|
|
|
/// Number of soft-deleted vectors still occupying the index.
|
|
pub fn deleted_count(&self) -> usize {
|
|
self.deleted.iter().filter(|&&d| d).count()
|
|
}
|
|
|
|
/// Number of live (non-deleted) vectors.
|
|
pub fn active_len(&self) -> usize {
|
|
self.vectors.len() - self.deleted_count()
|
|
}
|
|
|
|
/// Rebuild the index from scratch, dropping all soft-deleted vectors and
|
|
/// renumbering the survivors into a compact `0..active_len` id space.
|
|
///
|
|
/// Returns a mapping from old id to new id (`None` for dropped vectors) so
|
|
/// callers can rewrite any external id references they keep.
|
|
pub fn compact(&mut self) -> Vec<Option<usize>> {
|
|
let mut mapping = vec![None; self.vectors.len()];
|
|
let mut surviving: Vec<Vec<f32>> = Vec::with_capacity(self.active_len());
|
|
for (old, slot) in mapping.iter_mut().enumerate() {
|
|
if !self.deleted[old] {
|
|
*slot = Some(surviving.len());
|
|
surviving.push(self.vectors.row(old));
|
|
}
|
|
}
|
|
// Rebuilding must keep the storage the caller chose; a compaction is
|
|
// not the place to silently quadruple the index's memory.
|
|
let storage = self.vectors.storage();
|
|
*self = if surviving.is_empty() {
|
|
Self::new_with(self.m, self.ef_construction, self.metric, storage)
|
|
} else {
|
|
Self::build_with(
|
|
&surviving,
|
|
self.m,
|
|
self.ef_construction,
|
|
self.metric,
|
|
storage,
|
|
)
|
|
};
|
|
mapping
|
|
}
|
|
|
|
/// Search the index for the `k` nearest neighbors to the query vector.
|
|
///
|
|
/// # Parameters
|
|
/// - `query`: The query vector.
|
|
/// - `k`: Number of nearest neighbors to return.
|
|
/// - `ef`: Size of the dynamic candidate list during search (must be >= k).
|
|
///
|
|
/// # Returns
|
|
/// A vector of `(id, distance)` pairs sorted by distance (closest first).
|
|
pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(usize, f32)> {
|
|
if self.vectors.len() == 0 {
|
|
return Vec::new();
|
|
}
|
|
assert_eq!(query.len(), self.vectors.dim(), "query dimension mismatch");
|
|
let ef = ef.max(k);
|
|
// Prepared and, for a quantised store, quantised once per search
|
|
// rather than once per comparison.
|
|
let prepared = self.vectors.query(prepare(query.to_vec(), self.metric));
|
|
let target = Target::Query(&prepared);
|
|
|
|
let mut ep = self.entry_point;
|
|
let top_layer = self.graph.len().saturating_sub(1);
|
|
|
|
// Greedy search from top layer down to layer 1
|
|
for layer in (1..=top_layer).rev() {
|
|
ep = greedy_closest(&self.vectors, &self.graph[layer], &target, ep, self.metric);
|
|
}
|
|
|
|
// Search layer 0 for the ef nearest *live* nodes. Deleted nodes are
|
|
// still traversed (they remain valid graph waypoints) but take no
|
|
// result slot, so deletions near the query don't shrink the answer.
|
|
let candidates = search_layer(
|
|
&self.vectors,
|
|
&self.graph[0],
|
|
&target,
|
|
ep,
|
|
ef,
|
|
self.metric,
|
|
Some(&self.deleted),
|
|
);
|
|
|
|
candidates
|
|
.into_iter()
|
|
.take(k)
|
|
.map(|c| (c.id, c.distance))
|
|
.collect()
|
|
}
|
|
|
|
/// Save the index to an HDF5 file via the given writer.
|
|
pub fn save_to_hdf5(&self, writer: &mut IoFileWriter) -> Result<(), FormatError> {
|
|
let bytes = self.to_hdf5_bytes()?;
|
|
writer
|
|
.write_bytes_owned(bytes)
|
|
.map_err(|e| FormatError::SerializationError(e.to_string()))?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Serialize the index to HDF5 bytes.
|
|
pub fn to_hdf5_bytes(&self) -> Result<Vec<u8>, FormatError> {
|
|
let mut fw = FmtWriter::new();
|
|
let n = self.vectors.len();
|
|
let dim = self.vectors.dim();
|
|
|
|
// Flatten vectors into a 1D array for storage
|
|
let mut flat_vectors: Vec<f32> = Vec::with_capacity(n * dim);
|
|
for i in 0..n {
|
|
flat_vectors.extend_from_slice(&self.vectors.row(i));
|
|
}
|
|
|
|
let mut group = fw.create_group("ann");
|
|
|
|
group
|
|
.create_dataset("vectors")
|
|
.with_f32_data(&flat_vectors)
|
|
.with_shape(&[n as u64, dim as u64])
|
|
.set_attr("rows", AttrValue::I64(n as i64))
|
|
.set_attr("cols", AttrValue::I64(dim as i64));
|
|
|
|
// Serialize graph layers: store as flat i32 arrays with metadata
|
|
let num_layers = self.graph.len();
|
|
for (layer_idx, layer) in self.graph.iter().enumerate() {
|
|
// Flatten: for each node, store [count, neighbor1, neighbor2, ...]
|
|
let mut flat: Vec<i32> = Vec::new();
|
|
for neighbors in layer {
|
|
flat.push(neighbors.len() as i32);
|
|
for &n_id in neighbors {
|
|
flat.push(n_id as i32);
|
|
}
|
|
}
|
|
|
|
let ds_name = format!("graph_layer_{layer_idx}");
|
|
group
|
|
.create_dataset(&ds_name)
|
|
.with_i32_data(&flat)
|
|
.set_attr("layer", AttrValue::I64(layer_idx as i64));
|
|
}
|
|
|
|
// Soft-deletion bitset (format version 2+): 0 = live, 1 = deleted.
|
|
let deleted_i32: Vec<i32> = self.deleted.iter().map(|&d| d as i32).collect();
|
|
group.create_dataset("deleted").with_i32_data(&deleted_i32);
|
|
|
|
// Store config as attributes on a small dataset
|
|
let node_levels_i32: Vec<i32> = self.node_levels.iter().map(|&l| l as i32).collect();
|
|
group
|
|
.create_dataset("config")
|
|
.with_i32_data(&node_levels_i32)
|
|
.set_attr("format_version", AttrValue::I64(HNSW_FORMAT_VERSION))
|
|
.set_attr("m", AttrValue::I64(self.m as i64))
|
|
.set_attr(
|
|
"ef_construction",
|
|
AttrValue::I64(self.ef_construction as i64),
|
|
)
|
|
.set_attr("entry_point", AttrValue::I64(self.entry_point as i64))
|
|
.set_attr("num_layers", AttrValue::I64(num_layers as i64))
|
|
.set_attr(
|
|
"metric",
|
|
AttrValue::String(self.metric.as_str().to_string()),
|
|
)
|
|
.set_attr("num_vectors", AttrValue::I64(n as i64))
|
|
.set_attr("dimension", AttrValue::I64(dim as i64));
|
|
|
|
let finished = group.finish();
|
|
fw.add_group(finished);
|
|
fw.finish()
|
|
}
|
|
|
|
/// Load an HNSW index from HDF5 bytes.
|
|
///
|
|
/// The HDF5 data must contain the `/ann/vectors`, `/ann/graph_layer_*`,
|
|
/// and `/ann/config` datasets as produced by [`to_hdf5_bytes`].
|
|
pub fn load_from_hdf5(data: &[u8]) -> Result<Self, FormatError> {
|
|
let sig_offset = find_signature(data)?;
|
|
let sb = Superblock::parse(data, sig_offset)?;
|
|
|
|
// Read config dataset and its attributes
|
|
let config_attrs = read_dataset_attrs(data, &sb, "ann/config")?;
|
|
let config_raw = read_dataset_raw(data, &sb, "ann/config")?;
|
|
let config_dt = read_dataset_datatype(data, &sb, "ann/config")?;
|
|
let node_levels_i32 = read_as_i32(&config_raw, &config_dt)?;
|
|
|
|
// Files written before format version 2 have no version attribute; treat
|
|
// them as version 1. Reject anything newer than we understand.
|
|
let format_version = get_attr_i64_opt(&config_attrs, "format_version").unwrap_or(1);
|
|
if format_version > HNSW_FORMAT_VERSION {
|
|
return Err(FormatError::SerializationError(format!(
|
|
"unsupported HNSW format version {format_version} (this build understands up to {HNSW_FORMAT_VERSION})"
|
|
)));
|
|
}
|
|
|
|
let m = get_attr_i64(&config_attrs, "m")? as usize;
|
|
let ef_construction = get_attr_i64(&config_attrs, "ef_construction")? as usize;
|
|
let entry_point = get_attr_i64(&config_attrs, "entry_point")? as usize;
|
|
let num_layers = get_attr_i64(&config_attrs, "num_layers")? as usize;
|
|
let n = get_attr_i64(&config_attrs, "num_vectors")? as usize;
|
|
let dim = get_attr_i64(&config_attrs, "dimension")? as usize;
|
|
let metric_str = get_attr_string(&config_attrs, "metric")?;
|
|
let metric = DistanceMetric::from_str(&metric_str).ok_or_else(|| {
|
|
FormatError::SerializationError(format!("unknown metric: {metric_str}"))
|
|
})?;
|
|
|
|
let node_levels: Vec<usize> = node_levels_i32.iter().map(|&l| l as usize).collect();
|
|
|
|
// Read vectors
|
|
let vectors_raw = read_dataset_raw(data, &sb, "ann/vectors")?;
|
|
let vectors_dt = read_dataset_datatype(data, &sb, "ann/vectors")?;
|
|
let flat_vectors = read_as_f32(&vectors_raw, &vectors_dt)?;
|
|
let mut vectors = Vec::with_capacity(n);
|
|
for i in 0..n {
|
|
let start = i * dim;
|
|
let end = start + dim;
|
|
if end > flat_vectors.len() {
|
|
return Err(FormatError::DataSizeMismatch {
|
|
expected: end,
|
|
actual: flat_vectors.len(),
|
|
});
|
|
}
|
|
// Files written before vectors were stored unit-length hold the
|
|
// raw ones; preparing is idempotent, so this handles both.
|
|
vectors.push(prepare(flat_vectors[start..end].to_vec(), metric));
|
|
}
|
|
|
|
// Read graph layers
|
|
let mut graph = Vec::with_capacity(num_layers);
|
|
for layer_idx in 0..num_layers {
|
|
let ds_name = format!("ann/graph_layer_{layer_idx}");
|
|
let layer_raw = read_dataset_raw(data, &sb, &ds_name)?;
|
|
let layer_dt = read_dataset_datatype(data, &sb, &ds_name)?;
|
|
let flat = read_as_i32(&layer_raw, &layer_dt)?;
|
|
|
|
let mut layer_graph = Vec::with_capacity(n);
|
|
let mut pos = 0;
|
|
while pos < flat.len() {
|
|
let count = flat[pos] as usize;
|
|
pos += 1;
|
|
let mut neighbors = Vec::with_capacity(count);
|
|
for _ in 0..count {
|
|
if pos >= flat.len() {
|
|
return Err(FormatError::SerializationError(
|
|
"truncated graph data".into(),
|
|
));
|
|
}
|
|
neighbors.push(flat[pos] as usize);
|
|
pos += 1;
|
|
}
|
|
layer_graph.push(neighbors);
|
|
}
|
|
// Pad with empty if needed (nodes not present at this layer)
|
|
while layer_graph.len() < n {
|
|
layer_graph.push(Vec::new());
|
|
}
|
|
graph.push(layer_graph);
|
|
}
|
|
|
|
// Deleted bitset (version 2+). Older files default every node to live.
|
|
let deleted = if format_version >= 2 {
|
|
let deleted_raw = read_dataset_raw(data, &sb, "ann/deleted")?;
|
|
let deleted_dt = read_dataset_datatype(data, &sb, "ann/deleted")?;
|
|
let deleted_i32 = read_as_i32(&deleted_raw, &deleted_dt)?;
|
|
let mut deleted: Vec<bool> = deleted_i32.iter().map(|&d| d != 0).collect();
|
|
deleted.resize(n, false);
|
|
deleted
|
|
} else {
|
|
vec![false; n]
|
|
};
|
|
|
|
Ok(Self {
|
|
// Serialized files carry f32 vectors and no storage tag: a
|
|
// quantised index is rebuilt, not loaded.
|
|
vectors: Vectors::from_rows(&vectors, Storage::Float32, metric),
|
|
graph,
|
|
deleted,
|
|
entry_point,
|
|
m,
|
|
m_max0: m * 2,
|
|
ef_construction,
|
|
node_levels,
|
|
metric,
|
|
})
|
|
}
|
|
|
|
/// Serialize the **graph only** — levels, tombstones and adjacency, not the
|
|
/// vectors — for a caller that already stores the vectors elsewhere (the
|
|
/// agent's record cache). [`HnswIndex::to_hdf5_bytes`] writes a complete,
|
|
/// self-contained index including a full copy of every vector, which would
|
|
/// double such a store's size. Reattach with
|
|
/// [`HnswIndex::from_graph_bytes`].
|
|
///
|
|
/// Layout (little endian): magic `CHG1`, then u32 fields `n`, `m`,
|
|
/// `m_max0`, `ef_construction`, `entry_point`, `num_layers`, `metric`;
|
|
/// `n` level bytes; `n` tombstone bytes; per layer, per node that exists on
|
|
/// that layer: u32 neighbour count + u32 ids; trailing CRC32 of all of it.
|
|
pub fn graph_to_bytes(&self) -> Vec<u8> {
|
|
let n = self.vectors.len();
|
|
let mut out = Vec::with_capacity(32 + n * 2 + n * self.m_max0 * 4);
|
|
out.extend_from_slice(GRAPH_MAGIC);
|
|
for field in [
|
|
n,
|
|
self.m,
|
|
self.m_max0,
|
|
self.ef_construction,
|
|
self.entry_point,
|
|
self.graph.len(),
|
|
match self.metric {
|
|
DistanceMetric::L2 => 0,
|
|
DistanceMetric::Cosine => 1,
|
|
},
|
|
] {
|
|
out.extend_from_slice(&(field as u32).to_le_bytes());
|
|
}
|
|
out.extend(self.node_levels.iter().map(|&l| l.min(255) as u8));
|
|
out.extend(self.deleted.iter().map(|&d| u8::from(d)));
|
|
for (layer, adjacency) in self.graph.iter().enumerate() {
|
|
for (node, neighbors) in adjacency.iter().enumerate() {
|
|
if self.node_levels[node] < layer {
|
|
continue; // node does not exist on this layer
|
|
}
|
|
out.extend_from_slice(&(neighbors.len() as u32).to_le_bytes());
|
|
for &id in neighbors {
|
|
out.extend_from_slice(&(id as u32).to_le_bytes());
|
|
}
|
|
}
|
|
}
|
|
let crc = clawhdf5_format::checksum::crc32(&out);
|
|
out.extend_from_slice(&crc.to_le_bytes());
|
|
out
|
|
}
|
|
|
|
/// Rebuild an index from [`HnswIndex::graph_to_bytes`] output and the
|
|
/// vectors it was built over (same order). Every structural claim in
|
|
/// `bytes` is validated — a corrupt or mismatched graph is an error, never
|
|
/// an index that panics or walks out of bounds during a search.
|
|
pub fn from_graph_bytes(bytes: &[u8], vectors: Vec<Vec<f32>>) -> Result<Self, FormatError> {
|
|
Self::from_graph_bytes_with(bytes, vectors, Storage::default())
|
|
}
|
|
|
|
/// As [`from_graph_bytes`](Self::from_graph_bytes), choosing how the
|
|
/// rehydrated vectors are stored.
|
|
pub fn from_graph_bytes_with(
|
|
bytes: &[u8],
|
|
vectors: Vec<Vec<f32>>,
|
|
storage: Storage,
|
|
) -> Result<Self, FormatError> {
|
|
let bad = |what: &str| FormatError::SerializationError(format!("HNSW graph: {what}"));
|
|
let body_len = bytes
|
|
.len()
|
|
.checked_sub(4)
|
|
.filter(|&l| l >= GRAPH_MAGIC.len() + 7 * 4)
|
|
.ok_or_else(|| bad("truncated"))?;
|
|
let (body, crc_bytes) = bytes.split_at(body_len);
|
|
if &body[..4] != GRAPH_MAGIC {
|
|
return Err(bad("bad magic"));
|
|
}
|
|
let stored_crc =
|
|
u32::from_le_bytes([crc_bytes[0], crc_bytes[1], crc_bytes[2], crc_bytes[3]]);
|
|
if clawhdf5_format::checksum::crc32(body) != stored_crc {
|
|
return Err(bad("checksum mismatch"));
|
|
}
|
|
|
|
let mut pos = 4;
|
|
let next_u32 = |pos: &mut usize| -> Result<usize, FormatError> {
|
|
let b = body.get(*pos..*pos + 4).ok_or_else(|| bad("truncated"))?;
|
|
*pos += 4;
|
|
Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize)
|
|
};
|
|
let n = next_u32(&mut pos)?;
|
|
let m = next_u32(&mut pos)?;
|
|
let m_max0 = next_u32(&mut pos)?;
|
|
let ef_construction = next_u32(&mut pos)?;
|
|
let entry_point = next_u32(&mut pos)?;
|
|
let num_layers = next_u32(&mut pos)?;
|
|
let metric = match next_u32(&mut pos)? {
|
|
0 => DistanceMetric::L2,
|
|
1 => DistanceMetric::Cosine,
|
|
_ => return Err(bad("unknown metric")),
|
|
};
|
|
if n != vectors.len() {
|
|
return Err(bad("vector count does not match the graph"));
|
|
}
|
|
if n == 0 || entry_point >= n || m < 2 || num_layers == 0 || num_layers > 256 {
|
|
return Err(bad("invalid header"));
|
|
}
|
|
let dim = vectors[0].len();
|
|
if vectors.iter().any(|v| v.len() != dim) {
|
|
return Err(bad("vectors have mixed dimensions"));
|
|
}
|
|
|
|
let levels = body.get(pos..pos + n).ok_or_else(|| bad("truncated"))?;
|
|
pos += n;
|
|
let node_levels: Vec<usize> = levels.iter().map(|&l| l as usize).collect();
|
|
if node_levels.iter().any(|&l| l >= num_layers)
|
|
|| node_levels[entry_point] + 1 != num_layers
|
|
{
|
|
return Err(bad("levels inconsistent with layer count"));
|
|
}
|
|
let deleted: Vec<bool> = body
|
|
.get(pos..pos + n)
|
|
.ok_or_else(|| bad("truncated"))?
|
|
.iter()
|
|
.map(|&d| d != 0)
|
|
.collect();
|
|
pos += n;
|
|
|
|
let mut graph: Vec<Vec<Vec<usize>>> = Vec::with_capacity(num_layers);
|
|
for layer in 0..num_layers {
|
|
let max_conn = if layer == 0 { m_max0 } else { m };
|
|
let mut adjacency = vec![Vec::new(); n];
|
|
for (node, slot) in adjacency.iter_mut().enumerate() {
|
|
if node_levels[node] < layer {
|
|
continue;
|
|
}
|
|
let count = next_u32(&mut pos)?;
|
|
if count > max_conn {
|
|
return Err(bad("neighbour list exceeds the connection limit"));
|
|
}
|
|
let mut neighbors = Vec::with_capacity(count);
|
|
for _ in 0..count {
|
|
let id = next_u32(&mut pos)?;
|
|
// A neighbour must exist, and exist on this layer.
|
|
if id >= n || node_levels[id] < layer {
|
|
return Err(bad("neighbour id out of range for its layer"));
|
|
}
|
|
neighbors.push(id);
|
|
}
|
|
*slot = neighbors;
|
|
}
|
|
graph.push(adjacency);
|
|
}
|
|
if pos != body.len() {
|
|
return Err(bad("trailing bytes"));
|
|
}
|
|
|
|
Ok(Self {
|
|
vectors: Vectors::from_rows(
|
|
&vectors
|
|
.into_iter()
|
|
.map(|v| prepare(v, metric))
|
|
.collect::<Vec<_>>(),
|
|
storage,
|
|
metric,
|
|
),
|
|
graph,
|
|
deleted,
|
|
entry_point,
|
|
m,
|
|
m_max0,
|
|
ef_construction,
|
|
node_levels,
|
|
metric,
|
|
})
|
|
}
|
|
|
|
/// Returns the number of vectors in the index.
|
|
pub fn len(&self) -> usize {
|
|
self.vectors.len()
|
|
}
|
|
|
|
/// Returns true if the index is empty.
|
|
pub fn is_empty(&self) -> bool {
|
|
self.vectors.len() == 0
|
|
}
|
|
|
|
/// How this index stores its copy of the vectors.
|
|
pub fn storage(&self) -> Storage {
|
|
self.vectors.storage()
|
|
}
|
|
|
|
/// Returns the dimension of vectors in the index.
|
|
pub fn dimension(&self) -> usize {
|
|
self.vectors.dim()
|
|
}
|
|
|
|
/// Returns the number of layers in the graph.
|
|
pub fn num_layers(&self) -> usize {
|
|
self.graph.len()
|
|
}
|
|
|
|
/// Returns the distance metric used by this index.
|
|
pub fn metric(&self) -> DistanceMetric {
|
|
self.metric
|
|
}
|
|
|
|
/// Returns the maximum number of connections at layer 0.
|
|
pub fn m_max0(&self) -> usize {
|
|
self.m_max0
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Internal HNSW algorithms
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Greedy search: find the single closest node to `query` starting from `ep`.
|
|
fn greedy_closest(
|
|
vectors: &Vectors,
|
|
layer: &[Vec<usize>],
|
|
target: &Target<'_>,
|
|
mut ep: usize,
|
|
metric: DistanceMetric,
|
|
) -> usize {
|
|
let mut best_dist = vectors.dist_to(target, ep, metric);
|
|
loop {
|
|
let mut changed = false;
|
|
for &neighbor in &layer[ep] {
|
|
let d = vectors.dist_to(target, neighbor, metric);
|
|
if d < best_dist {
|
|
best_dist = d;
|
|
ep = neighbor;
|
|
changed = true;
|
|
}
|
|
}
|
|
if !changed {
|
|
break;
|
|
}
|
|
}
|
|
ep
|
|
}
|
|
|
|
/// Search a single layer for the ef closest nodes to `query`.
|
|
/// Best-first search of one layer, returning up to `ef` nodes by ascending
|
|
/// distance.
|
|
///
|
|
/// `skip` marks nodes that must not be *returned* (soft-deleted ones). They
|
|
/// are still traversed — a tombstone is a perfectly good waypoint — but they
|
|
/// never occupy one of the `ef` result slots. Filtering them out afterwards
|
|
/// instead meant a query whose neighbourhood had been deleted got back fewer
|
|
/// than `k` results, or none, however many live records were nearby.
|
|
fn search_layer(
|
|
vectors: &Vectors,
|
|
layer: &[Vec<usize>],
|
|
target: &Target<'_>,
|
|
ep: usize,
|
|
ef: usize,
|
|
metric: DistanceMetric,
|
|
skip: Option<&[bool]>,
|
|
) -> Vec<Candidate> {
|
|
let ep_dist = vectors.dist_to(target, ep, metric);
|
|
|
|
// Min-heap of candidates to explore
|
|
let mut candidates = BinaryHeap::new();
|
|
candidates.push(Candidate {
|
|
id: ep,
|
|
distance: ep_dist,
|
|
});
|
|
|
|
// Max-heap of current results (furthest first)
|
|
let mut results = BinaryHeap::new();
|
|
if !skip.is_some_and(|s| s[ep]) {
|
|
results.push(FarCandidate {
|
|
id: ep,
|
|
distance: ep_dist,
|
|
});
|
|
}
|
|
|
|
VISITED.with_borrow_mut(|visited| {
|
|
visited.begin(vectors.len());
|
|
visited.insert(ep);
|
|
search_layer_visit(
|
|
vectors, layer, target, ef, metric, skip, visited, candidates, results,
|
|
)
|
|
})
|
|
}
|
|
|
|
/// Which nodes a layer search has already seen. A `HashSet` allocated per call
|
|
/// was the hottest non-arithmetic cost in both build and query; this is one
|
|
/// `u32` stamp per node, reused across calls: a node is visited iff its stamp
|
|
/// equals the current epoch, so "clearing" is just bumping the epoch.
|
|
#[derive(Default)]
|
|
struct Visited {
|
|
stamps: Vec<u32>,
|
|
epoch: u32,
|
|
}
|
|
|
|
impl Visited {
|
|
fn begin(&mut self, n: usize) {
|
|
if self.stamps.len() < n {
|
|
self.stamps.resize(n, 0);
|
|
}
|
|
self.epoch = self.epoch.wrapping_add(1);
|
|
if self.epoch == 0 {
|
|
// Wrapped: stale stamps could collide with the new epoch.
|
|
self.stamps.iter_mut().for_each(|s| *s = 0);
|
|
self.epoch = 1;
|
|
}
|
|
}
|
|
|
|
/// Mark `id` visited; `true` if it was not already.
|
|
fn insert(&mut self, id: usize) -> bool {
|
|
let seen = self.stamps[id] == self.epoch;
|
|
self.stamps[id] = self.epoch;
|
|
!seen
|
|
}
|
|
}
|
|
|
|
thread_local! {
|
|
/// Per-thread scratch, so `search(&self)` stays shareable across threads.
|
|
static VISITED: std::cell::RefCell<Visited> = std::cell::RefCell::new(Visited::default());
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn search_layer_visit(
|
|
vectors: &Vectors,
|
|
layer: &[Vec<usize>],
|
|
target: &Target<'_>,
|
|
ef: usize,
|
|
metric: DistanceMetric,
|
|
skip: Option<&[bool]>,
|
|
visited: &mut Visited,
|
|
mut candidates: BinaryHeap<Candidate>,
|
|
mut results: BinaryHeap<FarCandidate>,
|
|
) -> Vec<Candidate> {
|
|
while let Some(closest) = candidates.pop() {
|
|
let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance);
|
|
if closest.distance > furthest_dist && results.len() >= ef {
|
|
break;
|
|
}
|
|
|
|
for &neighbor in &layer[closest.id] {
|
|
if !visited.insert(neighbor) {
|
|
continue;
|
|
}
|
|
|
|
let d = vectors.dist_to(target, neighbor, metric);
|
|
let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance);
|
|
|
|
if d < furthest_dist || results.len() < ef {
|
|
candidates.push(Candidate {
|
|
id: neighbor,
|
|
distance: d,
|
|
});
|
|
if skip.is_some_and(|s| s[neighbor]) {
|
|
continue; // explore through it, but never return it
|
|
}
|
|
results.push(FarCandidate {
|
|
id: neighbor,
|
|
distance: d,
|
|
});
|
|
if results.len() > ef {
|
|
results.pop();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Convert to sorted vec
|
|
let mut result: Vec<Candidate> = results
|
|
.into_iter()
|
|
.map(|f| Candidate {
|
|
id: f.id,
|
|
distance: f.distance,
|
|
})
|
|
.collect();
|
|
result.sort_by(|a, b| {
|
|
a.distance
|
|
.partial_cmp(&b.distance)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
result
|
|
}
|
|
|
|
/// Choose up to `max_conn` neighbours for a node from `candidates` (sorted by
|
|
/// ascending distance to that node) — the HNSW paper's Algorithm 4 with
|
|
/// `keepPrunedConnections`.
|
|
///
|
|
/// Taking the plain `max_conn` closest is what breaks the graph on clustered
|
|
/// data: every link of a node inside a tight cluster goes to that same cluster,
|
|
/// so clusters become islands that a search entering elsewhere can never
|
|
/// reach, however large `ef` is. Instead a candidate is accepted only if it is
|
|
/// closer to the node than to every neighbour already accepted, which spreads
|
|
/// links across directions and keeps the long edges that join clusters. Any
|
|
/// remaining slots are then filled with the closest rejected candidates, so a
|
|
/// node is never left under-connected.
|
|
fn select_neighbors(
|
|
vectors: &Vectors,
|
|
candidates: &[(usize, f32)],
|
|
max_conn: usize,
|
|
metric: DistanceMetric,
|
|
) -> Vec<usize> {
|
|
if candidates.len() <= max_conn {
|
|
return candidates.iter().map(|&(id, _)| id).collect();
|
|
}
|
|
let mut selected: Vec<usize> = Vec::with_capacity(max_conn);
|
|
let mut rejected: Vec<usize> = Vec::new();
|
|
for &(id, dist_to_node) in candidates {
|
|
if selected.len() >= max_conn {
|
|
break;
|
|
}
|
|
let diverse = selected
|
|
.iter()
|
|
.all(|&s| vectors.dist(id, s, metric) > dist_to_node);
|
|
if diverse {
|
|
selected.push(id);
|
|
} else {
|
|
rejected.push(id);
|
|
}
|
|
}
|
|
for id in rejected {
|
|
if selected.len() >= max_conn {
|
|
break;
|
|
}
|
|
selected.push(id);
|
|
}
|
|
selected
|
|
}
|
|
|
|
/// How many nodes to plan together once `linked` nodes are in the graph.
|
|
fn batch_len(linked: usize) -> usize {
|
|
(linked / 16).clamp(1, 512)
|
|
}
|
|
|
|
/// For each node in `batch`: the neighbours to link it to on each of its
|
|
/// layers, found by searching the graph as it currently stands.
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn plan_batch(
|
|
vectors: &Vectors,
|
|
graph: &[Vec<Vec<usize>>],
|
|
node_levels: &[usize],
|
|
batch: std::ops::Range<usize>,
|
|
entry_point: usize,
|
|
ep_level: usize,
|
|
(m, m_max0, ef_construction): (usize, usize, usize),
|
|
metric: DistanceMetric,
|
|
) -> Vec<Vec<(usize, Vec<usize>)>> {
|
|
let plan_one = |i: usize| -> Vec<(usize, Vec<usize>)> {
|
|
let node_level = node_levels[i];
|
|
let mut ep = entry_point;
|
|
// Phase 1: greedy descent from the top layer down to node_level + 1.
|
|
for layer in (node_level + 1..=ep_level).rev() {
|
|
ep = greedy_closest(vectors, &graph[layer], &Target::Node(i), ep, metric);
|
|
}
|
|
// Phase 2: search and select on every layer the node lives on.
|
|
let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1);
|
|
for layer in (0..=node_level.min(ep_level)).rev() {
|
|
let max_conn = if layer == 0 { m_max0 } else { m };
|
|
let neighbors = search_layer(
|
|
vectors,
|
|
&graph[layer],
|
|
&Target::Node(i),
|
|
ep,
|
|
ef_construction,
|
|
metric,
|
|
None,
|
|
);
|
|
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect();
|
|
let selected = select_neighbors(vectors, &scored, max_conn, metric);
|
|
if let Some(&closest) = selected.first() {
|
|
ep = closest;
|
|
}
|
|
plan.push((layer, selected));
|
|
}
|
|
plan
|
|
};
|
|
|
|
#[cfg(feature = "parallel")]
|
|
if batch.len() >= PARALLEL_MIN {
|
|
use rayon::prelude::*;
|
|
return batch.into_par_iter().map(plan_one).collect();
|
|
}
|
|
batch.map(plan_one).collect()
|
|
}
|
|
|
|
/// Prune every `(layer, node)` neighbour list in `overflowed` back to its
|
|
/// limit. Each list belongs to a different node, so they are independent.
|
|
fn prune_overflowed(
|
|
vectors: &Vectors,
|
|
graph: &mut [Vec<Vec<usize>>],
|
|
overflowed: Vec<(usize, usize)>,
|
|
(m, m_max0): (usize, usize),
|
|
metric: DistanceMetric,
|
|
) {
|
|
let limit = |layer: usize| if layer == 0 { m_max0 } else { m };
|
|
|
|
#[cfg(feature = "parallel")]
|
|
if overflowed.len() >= PARALLEL_MIN {
|
|
use rayon::prelude::*;
|
|
let mut work: Vec<(usize, usize, Vec<usize>)> = overflowed
|
|
.into_iter()
|
|
.map(|(layer, node)| (layer, node, std::mem::take(&mut graph[layer][node])))
|
|
.collect();
|
|
work.par_iter_mut().for_each(|(layer, node, list)| {
|
|
prune_connections(vectors, list, *node, limit(*layer), metric);
|
|
});
|
|
for (layer, node, list) in work {
|
|
graph[layer][node] = list;
|
|
}
|
|
return;
|
|
}
|
|
|
|
for (layer, node) in overflowed {
|
|
prune_connections(vectors, &mut graph[layer][node], node, limit(layer), metric);
|
|
}
|
|
}
|
|
|
|
/// Fewest independent tasks worth handing to the thread pool.
|
|
#[cfg(feature = "parallel")]
|
|
const PARALLEL_MIN: usize = 8;
|
|
|
|
/// Add the back-link `neighbor -> new_id` for every selected neighbour, pruning
|
|
/// each list that overflows.
|
|
///
|
|
/// Used by incremental [`HnswIndex::insert`]. (A single insert's handful of
|
|
/// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16
|
|
/// cores; bulk builds batch their pruning instead, see `prune_overflowed`.)
|
|
fn link_back(
|
|
vectors: &Vectors,
|
|
layer: &mut [Vec<usize>],
|
|
new_id: usize,
|
|
selected: &[usize],
|
|
max_conn: usize,
|
|
metric: DistanceMetric,
|
|
) {
|
|
let mut overflowed: Vec<usize> = Vec::new();
|
|
for &neighbor in selected {
|
|
layer[neighbor].push(new_id);
|
|
if layer[neighbor].len() > max_conn {
|
|
overflowed.push(neighbor);
|
|
}
|
|
}
|
|
|
|
for node in overflowed {
|
|
prune_connections(vectors, &mut layer[node], node, max_conn, metric);
|
|
}
|
|
}
|
|
|
|
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
|
|
fn prune_connections(
|
|
vectors: &Vectors,
|
|
neighbors: &mut Vec<usize>,
|
|
node: usize,
|
|
max_conn: usize,
|
|
metric: DistanceMetric,
|
|
) {
|
|
if neighbors.len() <= max_conn {
|
|
return;
|
|
}
|
|
let mut scored: Vec<(usize, f32)> = neighbors
|
|
.iter()
|
|
.map(|&n| (n, vectors.dist(node, n, metric)))
|
|
.collect();
|
|
scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
|
|
*neighbors = select_neighbors(vectors, &scored, max_conn, metric);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// HDF5 reading helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn read_dataset_raw(data: &[u8], sb: &Superblock, path: &str) -> Result<Vec<u8>, FormatError> {
|
|
let addr = resolve_path_any(data, sb, path)?;
|
|
let header = ObjectHeader::parse(data, addr as usize, sb.offset_size, sb.length_size)?;
|
|
|
|
let dt_msg = header
|
|
.messages
|
|
.iter()
|
|
.find(|m| m.msg_type == MessageType::Datatype)
|
|
.ok_or(FormatError::DatasetMissingData)?;
|
|
let (datatype, _) = Datatype::parse(&dt_msg.data)?;
|
|
|
|
let ds_msg = header
|
|
.messages
|
|
.iter()
|
|
.find(|m| m.msg_type == MessageType::Dataspace)
|
|
.ok_or(FormatError::DatasetMissingShape)?;
|
|
let dataspace = Dataspace::parse(&ds_msg.data, sb.length_size)?;
|
|
|
|
let dl_msg = header
|
|
.messages
|
|
.iter()
|
|
.find(|m| m.msg_type == MessageType::DataLayout)
|
|
.ok_or(FormatError::DatasetMissingData)?;
|
|
let layout = DataLayout::parse(&dl_msg.data, sb.offset_size, sb.length_size)?;
|
|
|
|
let pipeline = header
|
|
.messages
|
|
.iter()
|
|
.find(|m| m.msg_type == MessageType::FilterPipeline)
|
|
.and_then(|msg| FilterPipeline::parse(&msg.data).ok());
|
|
|
|
read_raw_data_full(
|
|
data,
|
|
&layout,
|
|
&dataspace,
|
|
&datatype,
|
|
pipeline.as_ref(),
|
|
sb.offset_size,
|
|
sb.length_size,
|
|
)
|
|
}
|
|
|
|
fn read_dataset_datatype(
|
|
data: &[u8],
|
|
sb: &Superblock,
|
|
path: &str,
|
|
) -> Result<Datatype, FormatError> {
|
|
let addr = resolve_path_any(data, sb, path)?;
|
|
let header = ObjectHeader::parse(data, addr as usize, sb.offset_size, sb.length_size)?;
|
|
let dt_msg = header
|
|
.messages
|
|
.iter()
|
|
.find(|m| m.msg_type == MessageType::Datatype)
|
|
.ok_or(FormatError::DatasetMissingData)?;
|
|
let (datatype, _) = Datatype::parse(&dt_msg.data)?;
|
|
Ok(datatype)
|
|
}
|
|
|
|
fn read_dataset_attrs(
|
|
data: &[u8],
|
|
sb: &Superblock,
|
|
path: &str,
|
|
) -> Result<Vec<(String, AttrValue)>, FormatError> {
|
|
let addr = resolve_path_any(data, sb, path)?;
|
|
let header = ObjectHeader::parse(data, addr as usize, sb.offset_size, sb.length_size)?;
|
|
let attr_msgs = extract_attributes_full(data, &header, sb.offset_size, sb.length_size)?;
|
|
|
|
let mut result = Vec::new();
|
|
for attr in &attr_msgs {
|
|
let name = attr.name.clone();
|
|
if let Some(val) = decode_simple_attr(attr) {
|
|
result.push((name, val));
|
|
}
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
fn decode_simple_attr(attr: &clawhdf5_format::attribute::AttributeMessage) -> Option<AttrValue> {
|
|
let raw = &attr.raw_data;
|
|
match &attr.datatype {
|
|
Datatype::FixedPoint { size, signed, .. } => {
|
|
if *signed {
|
|
match size {
|
|
8 => {
|
|
if raw.len() >= 8 {
|
|
let val = i64::from_le_bytes(raw[..8].try_into().ok()?);
|
|
Some(AttrValue::I64(val))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
4 => {
|
|
if raw.len() >= 4 {
|
|
let val = i32::from_le_bytes(raw[..4].try_into().ok()?) as i64;
|
|
Some(AttrValue::I64(val))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
_ => None,
|
|
}
|
|
} else if raw.len() >= *size as usize {
|
|
let val = u64::from_le_bytes({
|
|
let mut buf = [0u8; 8];
|
|
buf[..*size as usize].copy_from_slice(&raw[..*size as usize]);
|
|
buf
|
|
});
|
|
Some(AttrValue::U64(val))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Datatype::FloatingPoint { size: 8, .. } => {
|
|
if raw.len() >= 8 {
|
|
let val = f64::from_le_bytes(raw[..8].try_into().ok()?);
|
|
Some(AttrValue::F64(val))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Datatype::String { size, .. } => {
|
|
let s = std::str::from_utf8(&raw[..*size as usize]).ok()?;
|
|
Some(AttrValue::String(s.trim_end_matches('\0').to_string()))
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
fn get_attr_i64(attrs: &[(String, AttrValue)], name: &str) -> Result<i64, FormatError> {
|
|
for (n, v) in attrs {
|
|
if n == name {
|
|
return match v {
|
|
AttrValue::I64(val) => Ok(*val),
|
|
AttrValue::U64(val) => Ok(*val as i64),
|
|
_ => Err(FormatError::SerializationError(format!(
|
|
"attribute {name} is not an integer"
|
|
))),
|
|
};
|
|
}
|
|
}
|
|
Err(FormatError::SerializationError(format!(
|
|
"missing attribute: {name}"
|
|
)))
|
|
}
|
|
|
|
/// Like [`get_attr_i64`] but returns `None` when the attribute is absent or not
|
|
/// an integer, instead of erroring. Used for optional/back-compat attributes.
|
|
fn get_attr_i64_opt(attrs: &[(String, AttrValue)], name: &str) -> Option<i64> {
|
|
attrs
|
|
.iter()
|
|
.find(|(n, _)| n == name)
|
|
.and_then(|(_, v)| match v {
|
|
AttrValue::I64(val) => Some(*val),
|
|
AttrValue::U64(val) => Some(*val as i64),
|
|
_ => None,
|
|
})
|
|
}
|
|
|
|
fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String, FormatError> {
|
|
for (n, v) in attrs {
|
|
if n == name {
|
|
return match v {
|
|
AttrValue::String(s) => Ok(s.clone()),
|
|
_ => Err(FormatError::SerializationError(format!(
|
|
"attribute {name} is not a string"
|
|
))),
|
|
};
|
|
}
|
|
}
|
|
Err(FormatError::SerializationError(format!(
|
|
"missing attribute: {name}"
|
|
)))
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::collections::HashSet;
|
|
|
|
/// Tight, well-separated clusters — the shape real embeddings have, and
|
|
/// the case plain closest-M neighbour selection fails on: each cluster
|
|
/// becomes an island, so recall is capped no matter how large `ef` is.
|
|
fn clustered(n: usize, dim: usize, clusters: usize, seed: u64) -> Vec<Vec<f32>> {
|
|
let mut state = seed;
|
|
let mut next = move || {
|
|
state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
|
let mut z = state;
|
|
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
|
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
|
((z ^ (z >> 31)) >> 40) as f32 / (1u64 << 24) as f32 - 0.5
|
|
};
|
|
let centres: Vec<Vec<f32>> = (0..clusters)
|
|
.map(|_| (0..dim).map(|_| next() * 10.0).collect())
|
|
.collect();
|
|
(0..n)
|
|
.map(|i| {
|
|
centres[i % clusters]
|
|
.iter()
|
|
.map(|c| c + next() * 0.5)
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn recall_at_10(
|
|
index: &HnswIndex,
|
|
vectors: &[Vec<f32>],
|
|
queries: &[Vec<f32>],
|
|
ef: usize,
|
|
) -> f64 {
|
|
let mut hits = 0;
|
|
for q in queries {
|
|
let mut exact: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::L2)))
|
|
.collect();
|
|
exact.sort_by(|a, b| a.1.total_cmp(&b.1));
|
|
let want: Vec<usize> = exact[..10].iter().map(|e| e.0).collect();
|
|
hits += index
|
|
.search(q, 10, ef)
|
|
.iter()
|
|
.filter(|(id, _)| want.contains(id))
|
|
.count();
|
|
}
|
|
hits as f64 / (10 * queries.len()) as f64
|
|
}
|
|
|
|
#[test]
|
|
fn clustered_data_keeps_high_recall() {
|
|
// Data and queries come from the same clusters: one draw, split.
|
|
let mut vectors = clustered(3060, 24, 30, 1);
|
|
let queries = vectors.split_off(3000);
|
|
let built = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2);
|
|
let recall = recall_at_10(&built, &vectors, &queries, 64);
|
|
assert!(recall >= 0.95, "bulk build recall@10 = {recall}");
|
|
|
|
// Incremental inserts go through the same neighbour selection.
|
|
let mut incremental = HnswIndex::new(8, 40, DistanceMetric::L2);
|
|
for v in &vectors {
|
|
incremental.insert(v.clone());
|
|
}
|
|
let recall = recall_at_10(&incremental, &vectors, &queries, 64);
|
|
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
|
|
}
|
|
|
|
#[test]
|
|
fn int8_storage_needs_an_exact_re_score_to_match_f32() {
|
|
// Cosine only: rows are unit-length, so a quantised dot product
|
|
// reconstructs the similarity directly.
|
|
//
|
|
// Not the `clustered` generator: its clusters are far tighter than any
|
|
// real embedding, so neighbours sit closer together than the
|
|
// quantisation error and top-10 identity there is noise — that would
|
|
// measure the fixture, not the storage.
|
|
let mut vectors = make_random_vectors(3060, 128, 5);
|
|
let queries = vectors.split_off(3000);
|
|
let f32_index =
|
|
HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Float32);
|
|
let quantised =
|
|
HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Int8);
|
|
assert_eq!(quantised.storage(), Storage::Int8);
|
|
|
|
// Ground truth, not the f32 index's answers: re-scoring can beat that
|
|
// index, and measuring against it would score being right as drift.
|
|
let truth: Vec<Vec<usize>> = queries
|
|
.iter()
|
|
.map(|q| {
|
|
let mut d: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::Cosine)))
|
|
.collect();
|
|
d.sort_by(|a, b| a.1.total_cmp(&b.1));
|
|
d[..10].iter().map(|x| x.0).collect()
|
|
})
|
|
.collect();
|
|
|
|
let recall = |got: &dyn Fn(&[f32]) -> Vec<usize>| -> f64 {
|
|
let mut hits = 0;
|
|
for (q, want) in queries.iter().zip(&truth) {
|
|
hits += got(q).iter().filter(|id| want.contains(id)).count();
|
|
}
|
|
hits as f64 / (10 * queries.len()) as f64
|
|
};
|
|
|
|
let exact_recall = recall(&|q| f32_index.search(q, 10, 64).iter().map(|r| r.0).collect());
|
|
let raw_recall = recall(&|q| quantised.search(q, 10, 64).iter().map(|r| r.0).collect());
|
|
// Quantised distances alone cost recall, and `ef` cannot buy it back:
|
|
// the loss is in the distances, not in the graph.
|
|
assert!(
|
|
raw_recall < exact_recall,
|
|
"int8 alone should cost recall: {raw_recall} vs {exact_recall}"
|
|
);
|
|
|
|
// Re-scoring a wider candidate pool against the exact vectors — what a
|
|
// caller holding them (the agent's embedding cache) does — puts it
|
|
// back, because only the *ordering* was approximate.
|
|
let rescored_recall = recall(&|q| {
|
|
let mut pool: Vec<(usize, f32)> = quantised
|
|
.search(q, 40, 64)
|
|
.into_iter()
|
|
.map(|(id, _)| {
|
|
(
|
|
id,
|
|
compute_distance(q, &vectors[id], DistanceMetric::Cosine),
|
|
)
|
|
})
|
|
.collect();
|
|
pool.sort_by(|a, b| a.1.total_cmp(&b.1));
|
|
pool.truncate(10);
|
|
pool.into_iter().map(|p| p.0).collect()
|
|
});
|
|
assert!(
|
|
rescored_recall >= exact_recall - 0.01,
|
|
"int8 + exact re-score should match f32: {rescored_recall} vs {exact_recall} (raw {raw_recall})"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn int8_storage_falls_back_to_f32_for_non_cosine_metrics() {
|
|
// L2 distance is not recoverable from a quantised dot product, so the
|
|
// store silently stays f32 rather than returning wrong distances.
|
|
let vectors = clustered(100, 8, 5, 3);
|
|
let index = HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::L2, Storage::Int8);
|
|
assert_eq!(index.storage(), Storage::Float32);
|
|
}
|
|
|
|
#[test]
|
|
fn deletions_near_the_query_do_not_shrink_or_degrade_results() {
|
|
let mut vectors = clustered(2040, 16, 20, 11);
|
|
let queries = vectors.split_off(2000);
|
|
let mut index = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2);
|
|
|
|
let mut short = 0;
|
|
let mut hits = 0;
|
|
for q in &queries {
|
|
// Delete this query's 40 nearest neighbours: more than ef, so every
|
|
// candidate a plain search collects is a tombstone.
|
|
let mut exact: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(i, _)| !index.is_deleted(*i))
|
|
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::L2)))
|
|
.collect();
|
|
exact.sort_by(|a, b| a.1.total_cmp(&b.1));
|
|
for &(id, _) in &exact[..40] {
|
|
index.mark_deleted(id);
|
|
}
|
|
let want: Vec<usize> = exact[40..50].iter().map(|e| e.0).collect();
|
|
|
|
let got = index.search(q, 10, 32);
|
|
assert!(got.iter().all(|(id, _)| !index.is_deleted(*id)));
|
|
short += usize::from(got.len() < 10);
|
|
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
|
}
|
|
assert_eq!(short, 0, "searches returned fewer than k live results");
|
|
let recall = hits as f64 / (10 * queries.len()) as f64;
|
|
assert!(recall >= 0.9, "recall@10 among live records = {recall}");
|
|
}
|
|
|
|
#[test]
|
|
fn bulk_build_is_deterministic() {
|
|
// Batched planning runs on a thread pool with the `parallel` feature;
|
|
// the graph must not depend on scheduling. (It is also the same graph
|
|
// with and without the feature: both take this exact code path.)
|
|
let vectors = clustered(2500, 16, 20, 21);
|
|
let a = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::Cosine);
|
|
let b = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::Cosine);
|
|
assert_eq!(a.graph_to_bytes(), b.graph_to_bytes());
|
|
}
|
|
|
|
#[test]
|
|
fn batches_stay_a_small_fraction_of_the_graph() {
|
|
assert_eq!(batch_len(1), 1);
|
|
assert_eq!(batch_len(15), 1);
|
|
assert_eq!(batch_len(160), 10);
|
|
assert_eq!(batch_len(1_000_000), 512);
|
|
}
|
|
|
|
#[test]
|
|
fn graph_bytes_round_trip_gives_identical_searches() {
|
|
let mut vectors = clustered(1260, 16, 12, 9);
|
|
let queries = vectors.split_off(1200);
|
|
let mut index = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2);
|
|
index.mark_deleted(3);
|
|
index.mark_deleted(700);
|
|
|
|
let bytes = index.graph_to_bytes();
|
|
// The graph is a small fraction of the vectors it indexes... not
|
|
// necessarily at dim 16, but it must not embed them.
|
|
assert!(bytes.len() < 1200 * (16 * 2 + 2) * 4);
|
|
let restored = HnswIndex::from_graph_bytes(&bytes, vectors.clone()).unwrap();
|
|
assert_eq!(restored.deleted_count(), 2);
|
|
for q in &queries {
|
|
assert_eq!(restored.search(q, 10, 50), index.search(q, 10, 50));
|
|
}
|
|
// A restored index keeps working incrementally.
|
|
let mut restored = restored;
|
|
let id = restored.insert(queries[0].clone());
|
|
assert_eq!(restored.search(&queries[0], 1, 50)[0].0, id);
|
|
}
|
|
|
|
#[test]
|
|
fn damaged_or_mismatched_graph_bytes_are_errors() {
|
|
let vectors = clustered(300, 8, 6, 4);
|
|
let index = HnswIndex::build_with_metric(&vectors, 6, 30, DistanceMetric::Cosine);
|
|
let bytes = index.graph_to_bytes();
|
|
|
|
// Wrong vector set.
|
|
assert!(HnswIndex::from_graph_bytes(&bytes, vectors[..299].to_vec()).is_err());
|
|
// Every truncation.
|
|
for len in 0..bytes.len() {
|
|
assert!(
|
|
HnswIndex::from_graph_bytes(&bytes[..len], vectors.clone()).is_err(),
|
|
"truncated to {len}"
|
|
);
|
|
}
|
|
// A flipped bit anywhere.
|
|
for i in (0..bytes.len()).step_by(7) {
|
|
let mut damaged = bytes.clone();
|
|
damaged[i] ^= 0x10;
|
|
assert!(
|
|
HnswIndex::from_graph_bytes(&damaged, vectors.clone()).is_err(),
|
|
"bit flip at {i}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn structurally_invalid_graph_with_a_valid_checksum_is_rejected() {
|
|
// The CRC only proves the bytes are what was written; a hostile or
|
|
// buggy writer can checksum nonsense. Out-of-range neighbour ids must
|
|
// still be caught, or search would index out of bounds.
|
|
let vectors = clustered(50, 4, 3, 5);
|
|
let index = HnswIndex::build_with_metric(&vectors, 4, 20, DistanceMetric::L2);
|
|
let mut bytes = index.graph_to_bytes();
|
|
let body_len = bytes.len() - 4;
|
|
// First neighbour id of node 0 on layer 0 sits right after the header,
|
|
// levels, tombstones and node 0's count.
|
|
let at = 4 + 7 * 4 + 50 + 50 + 4;
|
|
bytes[at..at + 4].copy_from_slice(&9999u32.to_le_bytes());
|
|
let crc = clawhdf5_format::checksum::crc32(&bytes[..body_len]);
|
|
bytes[body_len..].copy_from_slice(&crc.to_le_bytes());
|
|
assert!(HnswIndex::from_graph_bytes(&bytes, vectors).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn select_neighbors_prefers_diverse_directions_and_fills_up() {
|
|
// Node at the origin. Three candidates bunched together on the right,
|
|
// one on the left. With room for two, plain closest-M would take two
|
|
// from the bunch and lose the only link leftwards.
|
|
let vectors = vec![
|
|
vec![0.0, 0.0], // 0: the node
|
|
vec![1.0, 0.0], // 1
|
|
vec![1.1, 0.0], // 2
|
|
vec![1.2, 0.0], // 3
|
|
vec![-2.0, 0.0], // 4
|
|
];
|
|
let store = Vectors::from_rows(&vectors, Storage::Float32, DistanceMetric::L2);
|
|
let scored: Vec<(usize, f32)> = (1..5)
|
|
.map(|i| {
|
|
(
|
|
i,
|
|
compute_distance(&vectors[0], &vectors[i], DistanceMetric::L2),
|
|
)
|
|
})
|
|
.collect();
|
|
assert_eq!(
|
|
select_neighbors(&store, &scored, 2, DistanceMetric::L2),
|
|
[1, 4]
|
|
);
|
|
// Spare capacity is filled with the closest rejected candidates.
|
|
assert_eq!(
|
|
select_neighbors(&store, &scored, 3, DistanceMetric::L2),
|
|
[1, 4, 2]
|
|
);
|
|
}
|
|
|
|
fn make_random_vectors(n: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
|
|
let mut vectors = Vec::with_capacity(n);
|
|
let mut state = seed;
|
|
for _ in 0..n {
|
|
let mut v = Vec::with_capacity(dim);
|
|
for _ in 0..dim {
|
|
state = splitmix64(state);
|
|
let val = (state >> 40) as f32 / 16777216.0 - 0.5;
|
|
v.push(val);
|
|
}
|
|
vectors.push(v);
|
|
}
|
|
vectors
|
|
}
|
|
|
|
#[test]
|
|
fn build_small_index() {
|
|
let vectors = vec![
|
|
vec![1.0, 0.0, 0.0],
|
|
vec![0.0, 1.0, 0.0],
|
|
vec![0.0, 0.0, 1.0],
|
|
vec![1.0, 1.0, 0.0],
|
|
];
|
|
let index = HnswIndex::build(&vectors, 4, 16);
|
|
assert_eq!(index.len(), 4);
|
|
assert_eq!(index.dimension(), 3);
|
|
assert!(!index.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn search_exact_match() {
|
|
let vectors = vec![
|
|
vec![1.0, 0.0],
|
|
vec![0.0, 1.0],
|
|
vec![1.0, 1.0],
|
|
vec![0.5, 0.5],
|
|
];
|
|
let index = HnswIndex::build(&vectors, 4, 16);
|
|
let results = index.search(&[1.0, 0.0], 1, 16);
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, 0); // exact match
|
|
assert!(results[0].1 < 1e-6); // distance ~0
|
|
}
|
|
|
|
#[test]
|
|
fn search_k_neighbors() {
|
|
let vectors = make_random_vectors(50, 8, 42);
|
|
let index = HnswIndex::build(&vectors, 8, 32);
|
|
let results = index.search(&vectors[0], 5, 32);
|
|
assert_eq!(results.len(), 5);
|
|
// First result should be the query itself
|
|
assert_eq!(results[0].0, 0);
|
|
assert!(results[0].1 < 1e-6);
|
|
// Results should be sorted by distance
|
|
for i in 1..results.len() {
|
|
assert!(results[i].1 >= results[i - 1].1);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn search_returns_correct_count() {
|
|
let vectors = make_random_vectors(20, 4, 123);
|
|
let index = HnswIndex::build(&vectors, 4, 16);
|
|
let results = index.search(&vectors[5], 3, 16);
|
|
assert_eq!(results.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn cosine_distance() {
|
|
let a = vec![1.0, 0.0];
|
|
let b = vec![0.0, 1.0];
|
|
let d = compute_distance(&a, &b, DistanceMetric::Cosine);
|
|
assert!((d - 1.0).abs() < 1e-6); // orthogonal vectors = cosine distance 1
|
|
|
|
let c = vec![1.0, 0.0];
|
|
let d_same = compute_distance(&a, &c, DistanceMetric::Cosine);
|
|
assert!(d_same < 1e-6); // same direction = cosine distance 0
|
|
}
|
|
|
|
#[test]
|
|
fn l2_distance() {
|
|
let a = vec![0.0, 0.0];
|
|
let b = vec![3.0, 4.0];
|
|
let d = compute_distance(&a, &b, DistanceMetric::L2);
|
|
assert!((d - 5.0).abs() < 1e-6); // 3-4-5 triangle
|
|
}
|
|
|
|
#[test]
|
|
fn cosine_index_build_and_search() {
|
|
let vectors = vec![
|
|
vec![1.0, 0.0, 0.0],
|
|
vec![0.9, 0.1, 0.0],
|
|
vec![0.0, 1.0, 0.0],
|
|
vec![0.0, 0.0, 1.0],
|
|
];
|
|
let index = HnswIndex::build_with_metric(&vectors, 4, 16, DistanceMetric::Cosine);
|
|
let results = index.search(&[1.0, 0.0, 0.0], 2, 16);
|
|
assert_eq!(results.len(), 2);
|
|
// The closest should be vector 0 (exact) or vector 1 (very similar)
|
|
assert!(results[0].0 == 0 || results[0].0 == 1);
|
|
}
|
|
|
|
#[test]
|
|
fn save_and_load_roundtrip() {
|
|
let vectors = make_random_vectors(30, 4, 999);
|
|
let index = HnswIndex::build(&vectors, 4, 16);
|
|
|
|
let bytes = index.to_hdf5_bytes().unwrap();
|
|
assert!(!bytes.is_empty());
|
|
assert_eq!(&bytes[..8], b"\x89HDF\r\n\x1a\n");
|
|
|
|
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
|
assert_eq!(loaded.len(), index.len());
|
|
assert_eq!(loaded.dimension(), index.dimension());
|
|
assert_eq!(loaded.metric(), DistanceMetric::L2);
|
|
assert_eq!(loaded.m, index.m);
|
|
assert_eq!(loaded.ef_construction, index.ef_construction);
|
|
assert_eq!(loaded.entry_point, index.entry_point);
|
|
|
|
// Verify vectors match
|
|
for i in 0..loaded.len() {
|
|
assert_eq!(loaded.vectors.row(i), index.vectors.row(i));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn save_and_load_preserves_search_results() {
|
|
let vectors = make_random_vectors(40, 6, 777);
|
|
let index = HnswIndex::build(&vectors, 6, 24);
|
|
|
|
let query = &vectors[10];
|
|
let original_results = index.search(query, 5, 24);
|
|
|
|
let bytes = index.to_hdf5_bytes().unwrap();
|
|
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
|
let loaded_results = loaded.search(query, 5, 24);
|
|
|
|
assert_eq!(original_results.len(), loaded_results.len());
|
|
for (orig, load) in original_results.iter().zip(loaded_results.iter()) {
|
|
assert_eq!(orig.0, load.0);
|
|
assert!((orig.1 - load.1).abs() < 1e-6);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cosine_roundtrip() {
|
|
let vectors = make_random_vectors(20, 3, 555);
|
|
let index = HnswIndex::build_with_metric(&vectors, 4, 16, DistanceMetric::Cosine);
|
|
|
|
let bytes = index.to_hdf5_bytes().unwrap();
|
|
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
|
assert_eq!(loaded.metric(), DistanceMetric::Cosine);
|
|
}
|
|
|
|
#[test]
|
|
fn index_metadata() {
|
|
let vectors = make_random_vectors(10, 5, 111);
|
|
let index = HnswIndex::build(&vectors, 3, 12);
|
|
assert_eq!(index.len(), 10);
|
|
assert_eq!(index.dimension(), 5);
|
|
assert!(index.num_layers() >= 1);
|
|
assert_eq!(index.metric(), DistanceMetric::L2);
|
|
}
|
|
|
|
#[test]
|
|
fn save_to_file_writer() {
|
|
let vectors = make_random_vectors(15, 3, 333);
|
|
let index = HnswIndex::build(&vectors, 4, 16);
|
|
|
|
let dir = std::env::temp_dir();
|
|
let path = dir.join("clawhdf5_ann_test_save.h5");
|
|
let mut writer = IoFileWriter::create(&path).unwrap();
|
|
index.save_to_hdf5(&mut writer).unwrap();
|
|
|
|
// Verify file exists and has HDF5 signature
|
|
let data = std::fs::read(&path).unwrap();
|
|
assert_eq!(&data[..8], b"\x89HDF\r\n\x1a\n");
|
|
|
|
// Load it back
|
|
let loaded = HnswIndex::load_from_hdf5(&data).unwrap();
|
|
assert_eq!(loaded.len(), 15);
|
|
|
|
std::fs::remove_file(&path).ok();
|
|
}
|
|
|
|
#[test]
|
|
fn search_accuracy_l2() {
|
|
// Build a small index and verify that brute-force search agrees
|
|
let vectors = make_random_vectors(100, 8, 42);
|
|
let index = HnswIndex::build(&vectors, 16, 64);
|
|
|
|
let query = &vectors[50];
|
|
let k = 10;
|
|
let results = index.search(query, k, 64);
|
|
|
|
// Brute-force nearest neighbors
|
|
let mut brute: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, compute_distance(query, v, DistanceMetric::L2)))
|
|
.collect();
|
|
brute.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
brute.truncate(k);
|
|
|
|
// With high ef and m, HNSW should find at least 80% of true neighbors
|
|
let hnsw_ids: HashSet<usize> = results.iter().map(|r| r.0).collect();
|
|
let brute_ids: HashSet<usize> = brute.iter().map(|r| r.0).collect();
|
|
let overlap = hnsw_ids.intersection(&brute_ids).count();
|
|
assert!(overlap >= k * 8 / 10, "HNSW recall too low: {overlap}/{k}");
|
|
}
|
|
|
|
#[test]
|
|
fn search_accuracy_cosine() {
|
|
let vectors = make_random_vectors(100, 8, 99);
|
|
let index = HnswIndex::build_with_metric(&vectors, 16, 64, DistanceMetric::Cosine);
|
|
|
|
let query = &vectors[25];
|
|
let k = 10;
|
|
let results = index.search(query, k, 64);
|
|
|
|
let mut brute: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, compute_distance(query, v, DistanceMetric::Cosine)))
|
|
.collect();
|
|
brute.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
brute.truncate(k);
|
|
|
|
let hnsw_ids: HashSet<usize> = results.iter().map(|r| r.0).collect();
|
|
let brute_ids: HashSet<usize> = brute.iter().map(|r| r.0).collect();
|
|
let overlap = hnsw_ids.intersection(&brute_ids).count();
|
|
assert!(
|
|
overlap >= k * 8 / 10,
|
|
"HNSW cosine recall too low: {overlap}/{k}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn distance_metric_str_roundtrip() {
|
|
assert_eq!(DistanceMetric::from_str("l2"), Some(DistanceMetric::L2));
|
|
assert_eq!(
|
|
DistanceMetric::from_str("cosine"),
|
|
Some(DistanceMetric::Cosine)
|
|
);
|
|
assert_eq!(DistanceMetric::from_str("unknown"), None);
|
|
assert_eq!(DistanceMetric::L2.as_str(), "l2");
|
|
assert_eq!(DistanceMetric::Cosine.as_str(), "cosine");
|
|
}
|
|
|
|
#[test]
|
|
fn cosine_zero_vector() {
|
|
let a = vec![0.0, 0.0];
|
|
let b = vec![1.0, 0.0];
|
|
let d = compute_distance(&a, &b, DistanceMetric::Cosine);
|
|
assert!((d - 1.0).abs() < 1e-6); // zero vector -> distance 1
|
|
}
|
|
|
|
#[test]
|
|
fn cosine_near_zero_vector() {
|
|
// Tiny-but-nonzero, identical-direction vectors: denom is well
|
|
// below f32::EPSILON but not exactly 0.0. Must still be treated
|
|
// as a degenerate/unreliable direction (distance 1, "maximally
|
|
// dissimilar"), not as an exact match (distance 0).
|
|
let a = vec![1e-4, 1e-4];
|
|
let b = vec![1e-4, 1e-4];
|
|
let d = compute_distance(&a, &b, DistanceMetric::Cosine);
|
|
assert!((d - 1.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn insert_into_empty_index() {
|
|
let mut index = HnswIndex::new(4, 16, DistanceMetric::L2);
|
|
assert!(index.is_empty());
|
|
let id = index.insert(vec![1.0, 0.0, 0.0]);
|
|
assert_eq!(id, 0);
|
|
assert_eq!(index.len(), 1);
|
|
let results = index.search(&[1.0, 0.0, 0.0], 1, 16);
|
|
assert_eq!(results, vec![(0, 0.0)]);
|
|
}
|
|
|
|
#[test]
|
|
fn incremental_insert_matches_batch_recall() {
|
|
// Build one index incrementally and one in batch from the same vectors,
|
|
// then confirm the incremental index has acceptable recall vs brute force.
|
|
let vectors = make_random_vectors(120, 8, 2024);
|
|
let mut incremental = HnswIndex::new(16, 64, DistanceMetric::L2);
|
|
for v in &vectors {
|
|
incremental.insert(v.clone());
|
|
}
|
|
assert_eq!(incremental.len(), vectors.len());
|
|
|
|
let query = &vectors[60];
|
|
let k = 10;
|
|
let results = incremental.search(query, k, 64);
|
|
|
|
let mut brute: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, compute_distance(query, v, DistanceMetric::L2)))
|
|
.collect();
|
|
brute.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
brute.truncate(k);
|
|
|
|
let hnsw_ids: HashSet<usize> = results.iter().map(|r| r.0).collect();
|
|
let brute_ids: HashSet<usize> = brute.iter().map(|r| r.0).collect();
|
|
let overlap = hnsw_ids.intersection(&brute_ids).count();
|
|
assert!(
|
|
overlap >= k * 8 / 10,
|
|
"incremental HNSW recall too low: {overlap}/{k}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn mark_deleted_excludes_from_search() {
|
|
let vectors = vec![
|
|
vec![1.0, 0.0],
|
|
vec![0.9, 0.1],
|
|
vec![0.0, 1.0],
|
|
vec![0.1, 0.9],
|
|
];
|
|
let mut index = HnswIndex::build(&vectors, 4, 16);
|
|
|
|
// Exact match on vector 0 before deletion.
|
|
let before = index.search(&[1.0, 0.0], 1, 16);
|
|
assert_eq!(before[0].0, 0);
|
|
|
|
assert!(index.mark_deleted(0));
|
|
assert!(index.is_deleted(0));
|
|
assert!(!index.mark_deleted(0)); // idempotent
|
|
assert_eq!(index.deleted_count(), 1);
|
|
assert_eq!(index.active_len(), 3);
|
|
|
|
// Vector 0 must no longer be returned; nearest is now vector 1.
|
|
let after = index.search(&[1.0, 0.0], 2, 16);
|
|
assert!(after.iter().all(|(id, _)| *id != 0));
|
|
assert_eq!(after[0].0, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn compact_drops_deleted_and_renumbers() {
|
|
let vectors = make_random_vectors(20, 4, 4242);
|
|
let mut index = HnswIndex::build(&vectors, 8, 32);
|
|
index.mark_deleted(3);
|
|
index.mark_deleted(7);
|
|
index.mark_deleted(11);
|
|
|
|
let mapping = index.compact();
|
|
assert_eq!(mapping.len(), 20);
|
|
assert_eq!(index.len(), 17);
|
|
assert_eq!(index.deleted_count(), 0);
|
|
// Deleted ids map to None; survivors map to a dense 0..17 range.
|
|
assert!(mapping[3].is_none() && mapping[7].is_none() && mapping[11].is_none());
|
|
let mut new_ids: Vec<usize> = mapping.iter().filter_map(|m| *m).collect();
|
|
new_ids.sort_unstable();
|
|
assert_eq!(new_ids, (0..17).collect::<Vec<_>>());
|
|
}
|
|
|
|
#[test]
|
|
fn compact_all_deleted_yields_empty_index() {
|
|
let vectors = make_random_vectors(5, 3, 1);
|
|
let mut index = HnswIndex::build(&vectors, 4, 16);
|
|
for i in 0..5 {
|
|
index.mark_deleted(i);
|
|
}
|
|
index.compact();
|
|
assert!(index.is_empty());
|
|
assert!(index.search(&[0.0, 0.0, 0.0], 3, 16).is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn versioned_roundtrip_preserves_deletions() {
|
|
let vectors = make_random_vectors(30, 4, 8080);
|
|
let mut index = HnswIndex::build(&vectors, 6, 24);
|
|
index.mark_deleted(5);
|
|
index.mark_deleted(12);
|
|
|
|
let bytes = index.to_hdf5_bytes().unwrap();
|
|
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
|
|
|
assert_eq!(loaded.len(), index.len());
|
|
assert!(loaded.is_deleted(5));
|
|
assert!(loaded.is_deleted(12));
|
|
assert_eq!(loaded.deleted_count(), 2);
|
|
// A deleted vector is not returned even after a round-trip.
|
|
let results = loaded.search(&vectors[5], 5, 24);
|
|
assert!(results.iter().all(|(id, _)| *id != 5));
|
|
}
|
|
|
|
#[test]
|
|
fn insert_then_save_load_search() {
|
|
let mut index = HnswIndex::new(8, 32, DistanceMetric::Cosine);
|
|
let vectors = make_random_vectors(25, 5, 31337);
|
|
for v in &vectors {
|
|
index.insert(v.clone());
|
|
}
|
|
let bytes = index.to_hdf5_bytes().unwrap();
|
|
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
|
assert_eq!(loaded.len(), 25);
|
|
assert_eq!(loaded.metric(), DistanceMetric::Cosine);
|
|
let results = loaded.search(&vectors[0], 3, 32);
|
|
assert_eq!(results.len(), 3);
|
|
assert_eq!(results[0].0, 0);
|
|
}
|
|
}
|