feat(ann): optional int8 storage for the index's vector copy
The HNSW index keeps its own copy of every vector, which at 100K x 384 f32 is ~146 MiB — the largest single item in the 2.43x footprint now that the agent stores embeddings once. `Storage::Int8` cuts that copy to a quarter by scaling each row to i8. The scale is per row, not global. Unit-length rows in d dimensions have components around 1/sqrt(d), so a fixed [-1, 1] scale spends fewer than 12 of the 255 levels on a 128-dimensional vector; measured against an exact ranking that gives 0.35 top-10 overlap. Scaling each row by its own largest component uses the full range and brings it to 0.99. Quantised distances still cost recall on their own, and `ef` does not buy it back because the loss is in the distances rather than the graph: at N=100K recall@10 tops out at 0.967 against f32's 0.9995. Re-scoring a wider candidate pool against the exact vectors removes the gap (0.9940 vs 0.9945 at ef=64) for ~13% of query throughput and ~16% of build time. That is the intended use, so it is what the test asserts — against ground truth, not against the f32 index, whose own mistakes a re-scored search is entitled to get right. Default is unchanged: `Storage::Float32`, chosen by every existing constructor. Serialized indexes carry f32 vectors and no storage tag, so a quantised index is rebuilt rather than loaded; `compact()` keeps the storage it was given. The harness grows `--int8` and `--rerank` axes, and reports the storage in each table header. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
osobh
+431
-62
@@ -154,6 +154,237 @@ impl Ord for FarCandidate {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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: `dim` terms of at most
|
||||
/// `127 * 127`, so `i32` suffices for any realistic dimension.
|
||||
fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||
// Four independent accumulators over 32-lane blocks: the widening product
|
||||
// has to sit in a fixed-length chunk for the vectoriser to see it, and the
|
||||
// separate accumulators keep it off one dependency chain.
|
||||
const LANE: usize = 8;
|
||||
let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>();
|
||||
let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>();
|
||||
let mut acc = [0i32; 4];
|
||||
for (x, y) in a_blocks.iter().zip(b_blocks) {
|
||||
for (lane, slot) in acc.iter_mut().enumerate() {
|
||||
let mut sum = 0i32;
|
||||
for k in 0..LANE {
|
||||
sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]);
|
||||
}
|
||||
*slot += sum;
|
||||
}
|
||||
}
|
||||
let tail: i32 = a_tail
|
||||
.iter()
|
||||
.zip(b_tail)
|
||||
.map(|(&x, &y)| i32::from(x) * i32::from(y))
|
||||
.sum();
|
||||
acc[0] + acc[1] + acc[2] + acc[3] + tail
|
||||
}
|
||||
|
||||
/// Magic for [`HnswIndex::graph_to_bytes`].
|
||||
const GRAPH_MAGIC: &[u8; 4] = b"CHG1";
|
||||
|
||||
@@ -174,8 +405,8 @@ pub const HNSW_FORMAT_VERSION: i64 = 2;
|
||||
/// HDF5 format.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HnswIndex {
|
||||
/// All vectors in the index.
|
||||
vectors: Vec<Vec<f32>>,
|
||||
/// 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
|
||||
@@ -214,6 +445,20 @@ impl HnswIndex {
|
||||
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");
|
||||
@@ -224,8 +469,11 @@ impl HnswIndex {
|
||||
|
||||
let m_max0 = m * 2;
|
||||
let n = vectors.len();
|
||||
let prepared: Vec<Vec<f32>> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect();
|
||||
let vectors: &[Vec<f32>] = &prepared;
|
||||
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);
|
||||
@@ -322,9 +570,20 @@ impl HnswIndex {
|
||||
/// 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 {
|
||||
vectors: Vec::new(),
|
||||
// The dimension is set by the first insert.
|
||||
vectors: Vectors::new(0, storage, metric),
|
||||
graph: Vec::new(),
|
||||
deleted: Vec::new(),
|
||||
entry_point: 0,
|
||||
@@ -351,7 +610,8 @@ impl HnswIndex {
|
||||
// Seed an empty index.
|
||||
if id == 0 {
|
||||
let node_level = assign_level(0, self.m);
|
||||
self.vectors.push(vector);
|
||||
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();
|
||||
@@ -361,12 +621,12 @@ impl HnswIndex {
|
||||
|
||||
assert_eq!(
|
||||
vector.len(),
|
||||
self.vectors[0].len(),
|
||||
self.vectors.dim(),
|
||||
"insert dimension mismatch"
|
||||
);
|
||||
|
||||
let node_level = assign_level(id, self.m);
|
||||
self.vectors.push(vector);
|
||||
self.vectors.push(&vector);
|
||||
self.deleted.push(false);
|
||||
self.node_levels.push(node_level);
|
||||
|
||||
@@ -387,7 +647,7 @@ impl HnswIndex {
|
||||
ep = greedy_closest(
|
||||
&self.vectors,
|
||||
&self.graph[layer],
|
||||
&self.vectors[id],
|
||||
&Target::Node(id),
|
||||
ep,
|
||||
self.metric,
|
||||
);
|
||||
@@ -400,7 +660,7 @@ impl HnswIndex {
|
||||
let neighbors = search_layer(
|
||||
&self.vectors,
|
||||
&self.graph[layer],
|
||||
&self.vectors[id],
|
||||
&Target::Node(id),
|
||||
ep,
|
||||
self.ef_construction,
|
||||
self.metric,
|
||||
@@ -466,16 +726,25 @@ impl HnswIndex {
|
||||
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, v) in self.vectors.iter().enumerate() {
|
||||
for (old, slot) in mapping.iter_mut().enumerate() {
|
||||
if !self.deleted[old] {
|
||||
mapping[old] = Some(surviving.len());
|
||||
surviving.push(v.clone());
|
||||
*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(self.m, self.ef_construction, self.metric)
|
||||
Self::new_with(self.m, self.ef_construction, self.metric, storage)
|
||||
} else {
|
||||
Self::build_with_metric(&surviving, self.m, self.ef_construction, self.metric)
|
||||
Self::build_with(
|
||||
&surviving,
|
||||
self.m,
|
||||
self.ef_construction,
|
||||
self.metric,
|
||||
storage,
|
||||
)
|
||||
};
|
||||
mapping
|
||||
}
|
||||
@@ -490,24 +759,22 @@ impl HnswIndex {
|
||||
/// # 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.is_empty() {
|
||||
if self.vectors.len() == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
assert_eq!(
|
||||
query.len(),
|
||||
self.vectors[0].len(),
|
||||
"query dimension mismatch"
|
||||
);
|
||||
assert_eq!(query.len(), self.vectors.dim(), "query dimension mismatch");
|
||||
let ef = ef.max(k);
|
||||
let prepared_query = prepare(query.to_vec(), self.metric);
|
||||
let query = prepared_query.as_slice();
|
||||
// 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], query, ep, self.metric);
|
||||
ep = greedy_closest(&self.vectors, &self.graph[layer], &target, ep, self.metric);
|
||||
}
|
||||
|
||||
// Search layer 0 for the ef nearest *live* nodes. Deleted nodes are
|
||||
@@ -516,7 +783,7 @@ impl HnswIndex {
|
||||
let candidates = search_layer(
|
||||
&self.vectors,
|
||||
&self.graph[0],
|
||||
query,
|
||||
&target,
|
||||
ep,
|
||||
ef,
|
||||
self.metric,
|
||||
@@ -543,14 +810,13 @@ impl HnswIndex {
|
||||
pub fn to_hdf5_bytes(&self) -> Result<Vec<u8>, FormatError> {
|
||||
let mut fw = FmtWriter::new();
|
||||
let n = self.vectors.len();
|
||||
let dim = if n > 0 { self.vectors[0].len() } else { 0 };
|
||||
let dim = self.vectors.dim();
|
||||
|
||||
// Flatten vectors into a 1D array for storage
|
||||
let flat_vectors: Vec<f32> = self
|
||||
.vectors
|
||||
.iter()
|
||||
.flat_map(|v| v.iter().copied())
|
||||
.collect();
|
||||
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");
|
||||
|
||||
@@ -709,7 +975,9 @@ impl HnswIndex {
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
vectors,
|
||||
// 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,
|
||||
@@ -773,6 +1041,16 @@ impl HnswIndex {
|
||||
/// `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()
|
||||
@@ -863,7 +1141,14 @@ impl HnswIndex {
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(),
|
||||
vectors: Vectors::from_rows(
|
||||
&vectors
|
||||
.into_iter()
|
||||
.map(|v| prepare(v, metric))
|
||||
.collect::<Vec<_>>(),
|
||||
storage,
|
||||
metric,
|
||||
),
|
||||
graph,
|
||||
deleted,
|
||||
entry_point,
|
||||
@@ -882,16 +1167,17 @@ impl HnswIndex {
|
||||
|
||||
/// Returns true if the index is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.vectors.is_empty()
|
||||
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 {
|
||||
if self.vectors.is_empty() {
|
||||
0
|
||||
} else {
|
||||
self.vectors[0].len()
|
||||
}
|
||||
self.vectors.dim()
|
||||
}
|
||||
|
||||
/// Returns the number of layers in the graph.
|
||||
@@ -916,17 +1202,17 @@ impl HnswIndex {
|
||||
|
||||
/// Greedy search: find the single closest node to `query` starting from `ep`.
|
||||
fn greedy_closest(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
target: &Target<'_>,
|
||||
mut ep: usize,
|
||||
metric: DistanceMetric,
|
||||
) -> usize {
|
||||
let mut best_dist = compute_distance(query, &vectors[ep], metric);
|
||||
let mut best_dist = vectors.dist_to(target, ep, metric);
|
||||
loop {
|
||||
let mut changed = false;
|
||||
for &neighbor in &layer[ep] {
|
||||
let d = compute_distance(query, &vectors[neighbor], metric);
|
||||
let d = vectors.dist_to(target, neighbor, metric);
|
||||
if d < best_dist {
|
||||
best_dist = d;
|
||||
ep = neighbor;
|
||||
@@ -950,15 +1236,15 @@ fn greedy_closest(
|
||||
/// 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: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
target: &Target<'_>,
|
||||
ep: usize,
|
||||
ef: usize,
|
||||
metric: DistanceMetric,
|
||||
skip: Option<&[bool]>,
|
||||
) -> Vec<Candidate> {
|
||||
let ep_dist = compute_distance(query, &vectors[ep], metric);
|
||||
let ep_dist = vectors.dist_to(target, ep, metric);
|
||||
|
||||
// Min-heap of candidates to explore
|
||||
let mut candidates = BinaryHeap::new();
|
||||
@@ -980,7 +1266,7 @@ fn search_layer(
|
||||
visited.begin(vectors.len());
|
||||
visited.insert(ep);
|
||||
search_layer_visit(
|
||||
vectors, layer, query, ef, metric, skip, visited, candidates, results,
|
||||
vectors, layer, target, ef, metric, skip, visited, candidates, results,
|
||||
)
|
||||
})
|
||||
}
|
||||
@@ -1023,9 +1309,9 @@ thread_local! {
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn search_layer_visit(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
target: &Target<'_>,
|
||||
ef: usize,
|
||||
metric: DistanceMetric,
|
||||
skip: Option<&[bool]>,
|
||||
@@ -1044,7 +1330,7 @@ fn search_layer_visit(
|
||||
continue;
|
||||
}
|
||||
|
||||
let d = compute_distance(query, &vectors[neighbor], metric);
|
||||
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 {
|
||||
@@ -1095,7 +1381,7 @@ fn search_layer_visit(
|
||||
/// remaining slots are then filled with the closest rejected candidates, so a
|
||||
/// node is never left under-connected.
|
||||
fn select_neighbors(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
candidates: &[(usize, f32)],
|
||||
max_conn: usize,
|
||||
metric: DistanceMetric,
|
||||
@@ -1111,7 +1397,7 @@ fn select_neighbors(
|
||||
}
|
||||
let diverse = selected
|
||||
.iter()
|
||||
.all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node);
|
||||
.all(|&s| vectors.dist(id, s, metric) > dist_to_node);
|
||||
if diverse {
|
||||
selected.push(id);
|
||||
} else {
|
||||
@@ -1136,7 +1422,7 @@ fn batch_len(linked: usize) -> usize {
|
||||
/// layers, found by searching the graph as it currently stands.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn plan_batch(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
graph: &[Vec<Vec<usize>>],
|
||||
node_levels: &[usize],
|
||||
batch: std::ops::Range<usize>,
|
||||
@@ -1150,7 +1436,7 @@ fn plan_batch(
|
||||
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], &vectors[i], ep, metric);
|
||||
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);
|
||||
@@ -1159,7 +1445,7 @@ fn plan_batch(
|
||||
let neighbors = search_layer(
|
||||
vectors,
|
||||
&graph[layer],
|
||||
&vectors[i],
|
||||
&Target::Node(i),
|
||||
ep,
|
||||
ef_construction,
|
||||
metric,
|
||||
@@ -1186,7 +1472,7 @@ fn plan_batch(
|
||||
/// 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: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
graph: &mut [Vec<Vec<usize>>],
|
||||
overflowed: Vec<(usize, usize)>,
|
||||
(m, m_max0): (usize, usize),
|
||||
@@ -1226,7 +1512,7 @@ const PARALLEL_MIN: usize = 8;
|
||||
/// 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: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
layer: &mut [Vec<usize>],
|
||||
new_id: usize,
|
||||
selected: &[usize],
|
||||
@@ -1248,7 +1534,7 @@ fn link_back(
|
||||
|
||||
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
|
||||
fn prune_connections(
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &Vectors,
|
||||
neighbors: &mut Vec<usize>,
|
||||
node: usize,
|
||||
max_conn: usize,
|
||||
@@ -1259,7 +1545,7 @@ fn prune_connections(
|
||||
}
|
||||
let mut scored: Vec<(usize, f32)> = neighbors
|
||||
.iter()
|
||||
.map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric)))
|
||||
.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);
|
||||
@@ -1519,6 +1805,88 @@ mod tests {
|
||||
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);
|
||||
@@ -1650,6 +2018,7 @@ mod tests {
|
||||
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| {
|
||||
(
|
||||
@@ -1659,12 +2028,12 @@ mod tests {
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(
|
||||
select_neighbors(&vectors, &scored, 2, DistanceMetric::L2),
|
||||
select_neighbors(&store, &scored, 2, DistanceMetric::L2),
|
||||
[1, 4]
|
||||
);
|
||||
// Spare capacity is filled with the closest rejected candidates.
|
||||
assert_eq!(
|
||||
select_neighbors(&vectors, &scored, 3, DistanceMetric::L2),
|
||||
select_neighbors(&store, &scored, 3, DistanceMetric::L2),
|
||||
[1, 4, 2]
|
||||
);
|
||||
}
|
||||
@@ -1790,7 +2159,7 @@ mod tests {
|
||||
|
||||
// Verify vectors match
|
||||
for i in 0..loaded.len() {
|
||||
assert_eq!(loaded.vectors[i], index.vectors[i]);
|
||||
assert_eq!(loaded.vectors.row(i), index.vectors.row(i));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -5,4 +5,4 @@
|
||||
|
||||
mod hnsw;
|
||||
|
||||
pub use hnsw::{DistanceMetric, HnswIndex};
|
||||
pub use hnsw::{DistanceMetric, HnswIndex, Storage};
|
||||
|
||||
Reference in New Issue
Block a user