dearsky/wifi-densepose
1
0
Fork 0
Code Issues 20 Pull Requests 5 Actions Packages Projects Releases 1 Wiki Activity

Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

Browse Source
This commit is contained in:
ruv
2026-02-28 14:39:40 -05:00
parent 7885bf6278 d803bfe2b1
commit cd5943df23
7854 changed files with 3522914 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

271
vendor/ruvector/crates/rvf/rvf-index/src/builder.rs vendored Normal file
View File

@@ -0,0 +1,271 @@
//! Index construction: building Layer A, B, C from vectors and an HNSW graph.
extern crate alloc;
use alloc::collections::BTreeMap;
use alloc::collections::BTreeSet;
use alloc::vec::Vec;
use crate::hnsw::{HnswConfig, HnswGraph, HnswLayer};
use crate::layers::{LayerA, LayerB, LayerC, PartitionEntry};
use crate::traits::VectorStore;
/// Build the full HNSW graph from a set of vectors.
///
/// `rng_values`: one random value per vector for level selection.
/// These should be uniform in (0, 1).
pub fn build_full_index(
vectors: &dyn VectorStore,
num_vectors: usize,
config: &HnswConfig,
rng_values: &[f64],
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> HnswGraph {
assert!(
rng_values.len() >= num_vectors,
"Need at least one rng value per vector"
);
let mut graph = HnswGraph::new(config);
for (i, &rng_val) in rng_values.iter().enumerate().take(num_vectors) {
graph.insert(i as u64, rng_val, vectors, distance_fn);
}
graph
}
/// Build Layer A from an existing HNSW graph.
///
/// Extracts entry points, top-layer adjacency, centroids, and a partition map.
///
/// `centroids`: precomputed cluster centroids.
/// `assignments`: for each vector ID, the centroid index it's assigned to.
pub fn build_layer_a(
graph: &HnswGraph,
centroids: &[Vec<f32>],
assignments: &[u32],
_num_vectors: u64,
) -> LayerA {
let entry_points = match graph.entry_point {
Some(ep) => vec![(ep, graph.max_layer as u32)],
None => vec![],
};
// Extract top layers. "Top" = layers above the threshold.
// For progressive indexing, we take layers >= max_layer - 1 (at least
// the top 2 layers). Adjust based on graph size.
let threshold = graph.max_layer.saturating_sub(1);
let top_layers: Vec<HnswLayer> = graph.layers[threshold..].to_vec();
// Build partition map from assignments.
let mut partitions: BTreeMap<u32, (u64, u64)> = BTreeMap::new();
for (vid, &centroid_id) in assignments.iter().enumerate() {
let entry = partitions
.entry(centroid_id)
.or_insert((vid as u64, vid as u64));
entry.0 = entry.0.min(vid as u64);
entry.1 = entry.1.max(vid as u64 + 1);
}
let partition_map: Vec<PartitionEntry> = partitions
.into_iter()
.map(|(centroid_id, (start, end))| PartitionEntry {
centroid_id,
vector_id_start: start,
vector_id_end: end,
segment_ref: 0,
block_ref: 0,
})
.collect();
LayerA {
entry_points,
top_layers,
top_layer_start: threshold,
centroids: centroids.to_vec(),
partition_map,
}
}
/// Build Layer B from an existing HNSW graph, keeping only hot nodes.
///
/// `hot_node_ids`: the set of node IDs in the hot working set.
pub fn build_layer_b(graph: &HnswGraph, hot_node_ids: &BTreeSet<u64>) -> LayerB {
let mut partial_adjacency = BTreeMap::new();
// For each hot node, include its layer 0 neighbors.
if let Some(layer0) = graph.layers.first() {
for &nid in hot_node_ids {
if let Some(neighbors) = layer0.adjacency.get(&nid) {
partial_adjacency.insert(nid, neighbors.clone());
}
}
}
// Compute covered ranges from the hot node set.
let covered_ranges = compute_ranges(hot_node_ids);
LayerB {
partial_adjacency,
covered_ranges,
}
}
/// Build Layer C from the full HNSW graph (just wraps all adjacency).
pub fn build_layer_c(graph: &HnswGraph) -> LayerC {
LayerC {
full_adjacency: graph.layers.clone(),
}
}
/// Incrementally add a vector to an existing HNSW graph.
pub fn incremental_insert(
graph: &mut HnswGraph,
id: u64,
rng_val: f64,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) {
graph.insert(id, rng_val, vectors, distance_fn);
}
/// Compute contiguous ranges from a sorted set of IDs.
fn compute_ranges(ids: &BTreeSet<u64>) -> Vec<(u64, u64)> {
if ids.is_empty() {
return Vec::new();
}
let mut ranges = Vec::new();
let mut iter = ids.iter();
let &first = iter.next().unwrap();
let mut start = first;
let mut end = first + 1;
for &id in iter {
if id == end {
end = id + 1;
} else {
ranges.push((start, end));
start = id;
end = id + 1;
}
}
ranges.push((start, end));
ranges
}
#[cfg(test)]
mod tests {
use super::*;
use crate::distance::l2_distance;
use crate::traits::InMemoryVectorStore;
#[test]
fn build_full_index_basic() {
let n = 50;
let dim = 4;
let vecs: Vec<Vec<f32>> = (0..n)
.map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect())
.collect();
let store = InMemoryVectorStore::new(vecs);
let config = HnswConfig {
m: 8,
m0: 16,
ef_construction: 50,
};
let rng_vals: Vec<f64> = (0..n).map(|i| ((i * 7 + 3) % 100) as f64 / 100.0).collect();
let graph = build_full_index(&store, n, &config, &rng_vals, &l2_distance);
assert_eq!(graph.node_count(), n);
assert!(graph.entry_point.is_some());
}
#[test]
fn build_layer_a_from_graph() {
let n = 100;
let dim = 4;
let vecs: Vec<Vec<f32>> = (0..n)
.map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect())
.collect();
let store = InMemoryVectorStore::new(vecs.clone());
let config = HnswConfig::default();
let rng_vals: Vec<f64> = (0..n).map(|i| ((i * 7 + 3) % 100) as f64 / 100.0).collect();
let graph = build_full_index(&store, n, &config, &rng_vals, &l2_distance);
let centroids = vec![vecs[25].clone(), vecs[75].clone()];
let assignments: Vec<u32> = (0..n).map(|i| if i < 50 { 0 } else { 1 }).collect();
let layer_a = build_layer_a(&graph, &centroids, &assignments, n as u64);
assert!(!layer_a.entry_points.is_empty());
assert_eq!(layer_a.centroids.len(), 2);
assert!(!layer_a.partition_map.is_empty());
}
#[test]
fn build_layer_b_from_graph() {
let n = 50;
let dim = 4;
let vecs: Vec<Vec<f32>> = (0..n)
.map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect())
.collect();
let store = InMemoryVectorStore::new(vecs);
let config = HnswConfig {
m: 8,
m0: 16,
ef_construction: 50,
};
let rng_vals: Vec<f64> = (0..n).map(|i| ((i * 7 + 3) % 100) as f64 / 100.0).collect();
let graph = build_full_index(&store, n, &config, &rng_vals, &l2_distance);
// Mark first 25 nodes as hot.
let hot: BTreeSet<u64> = (0..25).collect();
let layer_b = build_layer_b(&graph, &hot);
assert!(!layer_b.partial_adjacency.is_empty());
assert!(layer_b.has_node(0));
assert!(!layer_b.has_node(49));
}
#[test]
fn compute_ranges_basic() {
let ids: BTreeSet<u64> = [1, 2, 3, 5, 6, 10].into_iter().collect();
let ranges = compute_ranges(&ids);
assert_eq!(ranges, vec![(1, 4), (5, 7), (10, 11)]);
}
#[test]
fn compute_ranges_empty() {
let ids: BTreeSet<u64> = BTreeSet::new();
assert!(compute_ranges(&ids).is_empty());
}
#[test]
fn incremental_insert_works() {
let n = 20;
let dim = 4;
let mut vecs: Vec<Vec<f32>> = (0..n)
.map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect())
.collect();
let store = InMemoryVectorStore::new(vecs.clone());
let config = HnswConfig {
m: 8,
m0: 16,
ef_construction: 50,
};
let rng_vals: Vec<f64> = (0..n).map(|i| ((i * 7 + 3) % 100) as f64 / 100.0).collect();
let mut graph = build_full_index(&store, n, &config, &rng_vals, &l2_distance);
assert_eq!(graph.node_count(), n);
// Add one more vector.
vecs.push((0..dim).map(|d| (n * dim + d) as f32).collect());
let store2 = InMemoryVectorStore::new(vecs);
incremental_insert(&mut graph, n as u64, 0.5, &store2, &l2_distance);
assert_eq!(graph.node_count(), n + 1);
}
}

503
vendor/ruvector/crates/rvf/rvf-index/src/codec.rs vendored Normal file
View File

@@ -0,0 +1,503 @@
//! INDEX_SEG encode/decode: varint delta encoding with restart points.
//!
//! Implements the binary layout from the RVF wire spec for INDEX_SEG payloads.
extern crate alloc;
use alloc::vec::Vec;
/// Default restart interval for varint delta encoding.
pub const DEFAULT_RESTART_INTERVAL: u32 = 64;
/// Index segment header (64-byte aligned).
#[derive(Clone, Debug, PartialEq)]
pub struct IndexSegHeader {
/// 0 = HNSW, 1 = IVF, 2 = flat.
pub index_type: u8,
/// Layer level: 0 = A, 1 = B, 2 = C.
pub layer_level: u8,
/// HNSW max neighbors per layer.
pub m: u16,
/// ef_construction parameter.
pub ef_construction: u32,
/// Number of nodes in this segment.
pub node_count: u64,
}
/// Encoded adjacency data for a single node.
#[derive(Clone, Debug, PartialEq)]
pub struct NodeAdjacency {
/// The node ID.
pub node_id: u64,
/// Neighbor IDs per HNSW layer (index 0 = layer 0).
pub layers: Vec<Vec<u64>>,
}
/// Full decoded index segment data.
#[derive(Clone, Debug, PartialEq)]
pub struct IndexSegData {
pub header: IndexSegHeader,
pub restart_interval: u32,
pub nodes: Vec<NodeAdjacency>,
}
// ── Varint Encoding (LEB128) ─────────────────────────────────────
/// Encode a u64 as LEB128 varint.
pub fn encode_varint(mut value: u64, buf: &mut Vec<u8>) {
loop {
let mut byte = (value & 0x7F) as u8;
value >>= 7;
if value != 0 {
byte |= 0x80;
}
buf.push(byte);
if value == 0 {
break;
}
}
}
/// Decode a LEB128 varint from a byte slice. Returns `(value, bytes_consumed)`.
pub fn decode_varint(data: &[u8]) -> Option<(u64, usize)> {
let mut value: u64 = 0;
let mut shift: u32 = 0;
for (i, &byte) in data.iter().enumerate() {
if shift >= 64 {
return None; // Overflow.
}
value |= ((byte & 0x7F) as u64) << shift;
shift += 7;
if byte & 0x80 == 0 {
return Some((value, i + 1));
}
}
None // Incomplete.
}
// ── Delta Encoding ───────────────────────────────────────────────
/// Delta-encode a sorted sequence of u64 values.
pub fn delta_encode(sorted_ids: &[u64]) -> Vec<u64> {
if sorted_ids.is_empty() {
return Vec::new();
}
let mut deltas = Vec::with_capacity(sorted_ids.len());
deltas.push(sorted_ids[0]);
for i in 1..sorted_ids.len() {
deltas.push(sorted_ids[i] - sorted_ids[i - 1]);
}
deltas
}
/// Decode delta-encoded values back to absolute IDs.
pub fn delta_decode(deltas: &[u64]) -> Vec<u64> {
if deltas.is_empty() {
return Vec::new();
}
let mut ids = Vec::with_capacity(deltas.len());
ids.push(deltas[0]);
for i in 1..deltas.len() {
ids.push(ids[i - 1] + deltas[i]);
}
ids
}
// ── INDEX_SEG Encode ─────────────────────────────────────────────
/// Encode an INDEX_SEG payload.
///
/// Layout:
/// 1. Index header (padded to 64 bytes)
/// 2. Restart point index (padded to 64 bytes)
/// 3. Adjacency data with delta-encoded neighbor lists
pub fn encode_index_seg(data: &IndexSegData) -> Vec<u8> {
let mut buf = Vec::new();
// 1. Header (pad to 64 bytes).
buf.push(data.header.index_type);
buf.push(data.header.layer_level);
buf.extend_from_slice(&data.header.m.to_le_bytes());
buf.extend_from_slice(&data.header.ef_construction.to_le_bytes());
buf.extend_from_slice(&data.header.node_count.to_le_bytes());
pad_to_alignment(&mut buf, 64);
// 2. Encode adjacency data with restart points.
let restart_interval = data.restart_interval;
let mut adj_buf = Vec::new();
let mut restart_offsets: Vec<u32> = Vec::new();
for (idx, node) in data.nodes.iter().enumerate() {
if (idx as u32).is_multiple_of(restart_interval) {
restart_offsets.push(adj_buf.len() as u32);
}
// Encode layer count.
encode_varint(node.layers.len() as u64, &mut adj_buf);
// Encode each layer's neighbors.
for neighbors in &node.layers {
encode_varint(neighbors.len() as u64, &mut adj_buf);
// Delta-encode sorted neighbor IDs.
let mut sorted = neighbors.clone();
sorted.sort();
let is_restart = (idx as u32).is_multiple_of(restart_interval);
if is_restart {
// At restart points, encode absolute IDs.
for &nid in &sorted {
encode_varint(nid, &mut adj_buf);
}
} else {
// Delta encode.
let deltas = delta_encode(&sorted);
for &d in &deltas {
encode_varint(d, &mut adj_buf);
}
}
}
}
// Write restart point index.
buf.extend_from_slice(&restart_interval.to_le_bytes());
let restart_count = restart_offsets.len() as u32;
buf.extend_from_slice(&restart_count.to_le_bytes());
for offset in &restart_offsets {
buf.extend_from_slice(&offset.to_le_bytes());
}
pad_to_alignment(&mut buf, 64);
// Write adjacency data.
buf.extend_from_slice(&adj_buf);
pad_to_alignment(&mut buf, 64);
buf
}
/// Decode an INDEX_SEG payload.
pub fn decode_index_seg(data: &[u8]) -> Result<IndexSegData, CodecError> {
if data.len() < 64 {
return Err(CodecError::TooShort);
}
// 1. Parse header.
let index_type = data[0];
let layer_level = data[1];
let m = u16::from_le_bytes([data[2], data[3]]);
let ef_construction = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
let node_count = u64::from_le_bytes([
data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
]);
let header = IndexSegHeader {
index_type,
layer_level,
m,
ef_construction,
node_count,
};
// Skip header padding.
let mut pos = 64;
// 2. Parse restart point index.
if pos + 8 > data.len() {
return Err(CodecError::TooShort);
}
let restart_interval =
u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
pos += 4;
let restart_count =
u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
pos += 4;
let mut restart_offsets = Vec::with_capacity(restart_count as usize);
for _ in 0..restart_count {
if pos + 4 > data.len() {
return Err(CodecError::TooShort);
}
let offset = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
restart_offsets.push(offset);
pos += 4;
}
// Skip padding to 64-byte alignment.
pos = align_up(pos, 64);
// 3. Parse adjacency data.
let adj_start = pos;
let adj_data = &data[adj_start..];
let mut nodes = Vec::new();
let mut adj_pos = 0;
for node_idx in 0..node_count as usize {
let is_restart = (node_idx as u32).is_multiple_of(restart_interval);
// Decode layer count.
let (layer_count, consumed) =
decode_varint(&adj_data[adj_pos..]).ok_or(CodecError::InvalidVarint)?;
adj_pos += consumed;
let mut layers = Vec::with_capacity(layer_count as usize);
for _ in 0..layer_count {
let (neighbor_count, consumed) =
decode_varint(&adj_data[adj_pos..]).ok_or(CodecError::InvalidVarint)?;
adj_pos += consumed;
let mut neighbor_ids = Vec::with_capacity(neighbor_count as usize);
if is_restart {
// Absolute IDs at restart points.
for _ in 0..neighbor_count {
let (nid, consumed) =
decode_varint(&adj_data[adj_pos..]).ok_or(CodecError::InvalidVarint)?;
adj_pos += consumed;
neighbor_ids.push(nid);
}
} else {
// Delta-encoded IDs.
let mut deltas = Vec::with_capacity(neighbor_count as usize);
for _ in 0..neighbor_count {
let (d, consumed) =
decode_varint(&adj_data[adj_pos..]).ok_or(CodecError::InvalidVarint)?;
adj_pos += consumed;
deltas.push(d);
}
neighbor_ids = delta_decode(&deltas);
}
layers.push(neighbor_ids);
}
nodes.push(NodeAdjacency {
node_id: node_idx as u64,
layers,
});
}
Ok(IndexSegData {
header,
restart_interval,
nodes,
})
}
/// Errors that can occur during INDEX_SEG codec operations.
#[derive(Clone, Debug, PartialEq)]
pub enum CodecError {
/// Input data is shorter than expected.
TooShort,
/// Invalid varint encountered.
InvalidVarint,
/// Unknown index type.
UnknownIndexType(u8),
}
impl core::fmt::Display for CodecError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::TooShort => write!(f, "input data too short"),
Self::InvalidVarint => write!(f, "invalid varint encoding"),
Self::UnknownIndexType(t) => write!(f, "unknown index type: {}", t),
}
}
}
// ── Helpers ──────────────────────────────────────────────────────
/// Pad `buf` with zeros to the next multiple of `alignment`.
fn pad_to_alignment(buf: &mut Vec<u8>, alignment: usize) {
let rem = buf.len() % alignment;
if rem != 0 {
buf.resize(buf.len() + (alignment - rem), 0);
}
}
/// Round `offset` up to the next multiple of `alignment`.
fn align_up(offset: usize, alignment: usize) -> usize {
let rem = offset % alignment;
if rem == 0 {
offset
} else {
offset + (alignment - rem)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn varint_round_trip() {
let values = [0, 1, 127, 128, 16383, 16384, 2097151, u64::MAX];
for &val in &values {
let mut buf = Vec::new();
encode_varint(val, &mut buf);
let (decoded, consumed) = decode_varint(&buf).unwrap();
assert_eq!(decoded, val);
assert_eq!(consumed, buf.len());
}
}
#[test]
fn varint_encoding_sizes() {
let mut buf = Vec::new();
encode_varint(0, &mut buf);
assert_eq!(buf.len(), 1);
buf.clear();
encode_varint(127, &mut buf);
assert_eq!(buf.len(), 1);
buf.clear();
encode_varint(128, &mut buf);
assert_eq!(buf.len(), 2);
buf.clear();
encode_varint(16383, &mut buf);
assert_eq!(buf.len(), 2);
buf.clear();
encode_varint(16384, &mut buf);
assert_eq!(buf.len(), 3);
}
#[test]
fn delta_encode_decode_round_trip() {
let ids = vec![100, 105, 108, 120, 200];
let deltas = delta_encode(&ids);
assert_eq!(deltas, vec![100, 5, 3, 12, 80]);
let decoded = delta_decode(&deltas);
assert_eq!(decoded, ids);
}
#[test]
fn delta_encode_empty() {
assert!(delta_encode(&[]).is_empty());
assert!(delta_decode(&[]).is_empty());
}
#[test]
fn index_seg_round_trip() {
let data = IndexSegData {
header: IndexSegHeader {
index_type: 0, // HNSW
layer_level: 2, // Layer C
m: 16,
ef_construction: 200,
node_count: 5,
},
restart_interval: 3,
nodes: vec![
NodeAdjacency {
node_id: 0,
layers: vec![vec![1, 2, 3], vec![1]],
},
NodeAdjacency {
node_id: 1,
layers: vec![vec![0, 2, 4]],
},
NodeAdjacency {
node_id: 2,
layers: vec![vec![0, 1, 3, 4]],
},
NodeAdjacency {
node_id: 3,
layers: vec![vec![0, 2, 4], vec![4]],
},
NodeAdjacency {
node_id: 4,
layers: vec![vec![1, 2, 3]],
},
],
};
let encoded = encode_index_seg(&data);
let decoded = decode_index_seg(&encoded).unwrap();
assert_eq!(decoded.header, data.header);
assert_eq!(decoded.restart_interval, data.restart_interval);
assert_eq!(decoded.nodes.len(), data.nodes.len());
// Verify each node's adjacency. Note: neighbors are sorted during encoding.
for (orig, dec) in data.nodes.iter().zip(decoded.nodes.iter()) {
assert_eq!(dec.node_id, orig.node_id);
assert_eq!(dec.layers.len(), orig.layers.len());
for (ol, dl) in orig.layers.iter().zip(dec.layers.iter()) {
let mut sorted_orig = ol.clone();
sorted_orig.sort();
assert_eq!(*dl, sorted_orig);
}
}
}
#[test]
fn index_seg_larger_with_restart() {
// Test with enough nodes to exercise multiple restart groups.
let num_nodes = 200;
let restart_interval = 64;
let nodes: Vec<NodeAdjacency> = (0..num_nodes)
.map(|i| {
let neighbors: Vec<u64> =
(0..8).map(|j| ((i + j + 1) % num_nodes) as u64).collect();
NodeAdjacency {
node_id: i as u64,
layers: vec![neighbors],
}
})
.collect();
let data = IndexSegData {
header: IndexSegHeader {
index_type: 0,
layer_level: 2,
m: 16,
ef_construction: 200,
node_count: num_nodes as u64,
},
restart_interval,
nodes,
};
let encoded = encode_index_seg(&data);
let decoded = decode_index_seg(&encoded).unwrap();
assert_eq!(decoded.header, data.header);
assert_eq!(decoded.nodes.len(), data.nodes.len());
for (orig, dec) in data.nodes.iter().zip(decoded.nodes.iter()) {
assert_eq!(dec.layers.len(), orig.layers.len());
for (ol, dl) in orig.layers.iter().zip(dec.layers.iter()) {
let mut sorted_orig = ol.clone();
sorted_orig.sort();
assert_eq!(*dl, sorted_orig);
}
}
}
#[test]
fn delta_encoding_sorted_u64_sequences() {
// Verify exact round-trip for various sorted u64 sequences.
let sequences: Vec<Vec<u64>> = vec![
vec![0, 1, 2, 3, 4],
vec![1000, 2000, 3000, 4000],
vec![0, 100, 200, 300, 400, 500],
vec![
u64::MAX - 4,
u64::MAX - 3,
u64::MAX - 2,
u64::MAX - 1,
u64::MAX,
],
];
for seq in sequences {
let deltas = delta_encode(&seq);
let decoded = delta_decode(&deltas);
assert_eq!(decoded, seq, "Failed for sequence: {:?}", seq);
}
}
}

516
vendor/ruvector/crates/rvf/rvf-index/src/distance.rs vendored Normal file
View File

@@ -0,0 +1,516 @@
//! Distance functions for vector similarity search.
//!
//! Provides L2 (Euclidean), cosine, and inner product distance metrics.
//! Includes platform-specific SIMD implementations (AVX2+FMA on x86_64,
//! NEON on aarch64) with automatic runtime dispatch.
// ── Scalar implementations ─────────────────────────────────────────
/// Scalar squared L2 (Euclidean) distance between two vectors.
///
/// Returns the sum of squared differences. Does NOT take the square root
/// because the ordering is preserved and sqrt is monotonic.
#[inline]
fn l2_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
debug_assert_eq!(a.len(), b.len());
a.iter()
.zip(b.iter())
.map(|(x, y)| {
let d = x - y;
d * d
})
.sum()
}
/// Scalar cosine distance: `1 - cosine_similarity`.
///
/// Returns a value in `[0, 2]` where 0 means identical direction.
/// If either vector has zero norm, returns `1.0`.
#[inline]
fn cosine_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
debug_assert_eq!(a.len(), b.len());
let mut dot = 0.0f32;
let mut norm_a = 0.0f32;
let mut norm_b = 0.0f32;
for (x, y) in a.iter().zip(b.iter()) {
dot += x * y;
norm_a += x * x;
norm_b += y * y;
}
let denom = (norm_a * norm_b).sqrt();
if denom < f32::EPSILON {
return 1.0;
}
1.0 - dot / denom
}
/// Scalar inner (dot) product distance: `-dot(a, b)`.
///
/// Negated so that higher similarity yields a lower distance value,
/// which is consistent with the min-heap search ordering.
#[inline]
fn dot_product_scalar(a: &[f32], b: &[f32]) -> f32 {
debug_assert_eq!(a.len(), b.len());
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
-dot
}
// ── x86_64 AVX2+FMA implementations ────────────────────────────────
#[cfg(target_arch = "x86_64")]
mod avx2 {
#[target_feature(enable = "avx2", enable = "fma")]
pub(super) unsafe fn l2_distance_avx2(a: &[f32], b: &[f32]) -> f32 {
use core::arch::x86_64::*;
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let chunks = n / 8;
let remainder = n % 8;
let mut sum = _mm256_setzero_ps();
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
for i in 0..chunks {
let offset = i * 8;
let va = _mm256_loadu_ps(a_ptr.add(offset));
let vb = _mm256_loadu_ps(b_ptr.add(offset));
let diff = _mm256_sub_ps(va, vb);
sum = _mm256_fmadd_ps(diff, diff, sum);
}
// Horizontal sum of the 8 lanes.
// sum = [s0, s1, s2, s3, s4, s5, s6, s7]
let hi128 = _mm256_extractf128_ps(sum, 1);
let lo128 = _mm256_castps256_ps128(sum);
let sum128 = _mm_add_ps(lo128, hi128);
let shuf = _mm_movehdup_ps(sum128);
let sums = _mm_add_ps(sum128, shuf);
let shuf2 = _mm_movehl_ps(sums, sums);
let result = _mm_add_ss(sums, shuf2);
let mut total = _mm_cvtss_f32(result);
// Handle remainder with scalar.
let base = chunks * 8;
for i in 0..remainder {
let d = a[base + i] - b[base + i];
total += d * d;
}
total
}
#[target_feature(enable = "avx2", enable = "fma")]
pub(super) unsafe fn cosine_distance_avx2(a: &[f32], b: &[f32]) -> f32 {
use core::arch::x86_64::*;
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let chunks = n / 8;
let remainder = n % 8;
let mut dot_acc = _mm256_setzero_ps();
let mut norm_a_acc = _mm256_setzero_ps();
let mut norm_b_acc = _mm256_setzero_ps();
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
for i in 0..chunks {
let offset = i * 8;
let va = _mm256_loadu_ps(a_ptr.add(offset));
let vb = _mm256_loadu_ps(b_ptr.add(offset));
dot_acc = _mm256_fmadd_ps(va, vb, dot_acc);
norm_a_acc = _mm256_fmadd_ps(va, va, norm_a_acc);
norm_b_acc = _mm256_fmadd_ps(vb, vb, norm_b_acc);
}
// Horizontal sums.
let hsum = |v: __m256| -> f32 {
let hi128 = _mm256_extractf128_ps(v, 1);
let lo128 = _mm256_castps256_ps128(v);
let sum128 = _mm_add_ps(lo128, hi128);
let shuf = _mm_movehdup_ps(sum128);
let sums = _mm_add_ps(sum128, shuf);
let shuf2 = _mm_movehl_ps(sums, sums);
let result = _mm_add_ss(sums, shuf2);
_mm_cvtss_f32(result)
};
let mut dot = hsum(dot_acc);
let mut norm_a = hsum(norm_a_acc);
let mut norm_b = hsum(norm_b_acc);
// Remainder.
let base = chunks * 8;
for i in 0..remainder {
let x = a[base + i];
let y = b[base + i];
dot += x * y;
norm_a += x * x;
norm_b += y * y;
}
let denom = (norm_a * norm_b).sqrt();
if denom < f32::EPSILON {
return 1.0;
}
1.0 - dot / denom
}
#[target_feature(enable = "avx2", enable = "fma")]
pub(super) unsafe fn dot_product_avx2(a: &[f32], b: &[f32]) -> f32 {
use core::arch::x86_64::*;
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let chunks = n / 8;
let remainder = n % 8;
let mut dot_acc = _mm256_setzero_ps();
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
for i in 0..chunks {
let offset = i * 8;
let va = _mm256_loadu_ps(a_ptr.add(offset));
let vb = _mm256_loadu_ps(b_ptr.add(offset));
dot_acc = _mm256_fmadd_ps(va, vb, dot_acc);
}
let hi128 = _mm256_extractf128_ps(dot_acc, 1);
let lo128 = _mm256_castps256_ps128(dot_acc);
let sum128 = _mm_add_ps(lo128, hi128);
let shuf = _mm_movehdup_ps(sum128);
let sums = _mm_add_ps(sum128, shuf);
let shuf2 = _mm_movehl_ps(sums, sums);
let result = _mm_add_ss(sums, shuf2);
let mut dot = _mm_cvtss_f32(result);
let base = chunks * 8;
for i in 0..remainder {
dot += a[base + i] * b[base + i];
}
-dot
}
}
// ── aarch64 NEON implementations ────────────────────────────────────
#[cfg(target_arch = "aarch64")]
mod neon {
#[target_feature(enable = "neon")]
pub(super) unsafe fn l2_distance_neon(a: &[f32], b: &[f32]) -> f32 {
use core::arch::aarch64::*;
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let chunks = n / 4;
let remainder = n % 4;
let mut sum = vdupq_n_f32(0.0);
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
for i in 0..chunks {
let offset = i * 4;
let va = vld1q_f32(a_ptr.add(offset));
let vb = vld1q_f32(b_ptr.add(offset));
let diff = vsubq_f32(va, vb);
sum = vfmaq_f32(sum, diff, diff);
}
let mut total = vaddvq_f32(sum);
let base = chunks * 4;
for i in 0..remainder {
let d = a[base + i] - b[base + i];
total += d * d;
}
total
}
#[target_feature(enable = "neon")]
pub(super) unsafe fn cosine_distance_neon(a: &[f32], b: &[f32]) -> f32 {
use core::arch::aarch64::*;
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let chunks = n / 4;
let remainder = n % 4;
let mut dot_acc = vdupq_n_f32(0.0);
let mut norm_a_acc = vdupq_n_f32(0.0);
let mut norm_b_acc = vdupq_n_f32(0.0);
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
for i in 0..chunks {
let offset = i * 4;
let va = vld1q_f32(a_ptr.add(offset));
let vb = vld1q_f32(b_ptr.add(offset));
dot_acc = vfmaq_f32(dot_acc, va, vb);
norm_a_acc = vfmaq_f32(norm_a_acc, va, va);
norm_b_acc = vfmaq_f32(norm_b_acc, vb, vb);
}
let mut dot = vaddvq_f32(dot_acc);
let mut norm_a = vaddvq_f32(norm_a_acc);
let mut norm_b = vaddvq_f32(norm_b_acc);
let base = chunks * 4;
for i in 0..remainder {
let x = a[base + i];
let y = b[base + i];
dot += x * y;
norm_a += x * x;
norm_b += y * y;
}
let denom = (norm_a * norm_b).sqrt();
if denom < f32::EPSILON {
return 1.0;
}
1.0 - dot / denom
}
#[target_feature(enable = "neon")]
pub(super) unsafe fn dot_product_neon(a: &[f32], b: &[f32]) -> f32 {
use core::arch::aarch64::*;
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let chunks = n / 4;
let remainder = n % 4;
let mut dot_acc = vdupq_n_f32(0.0);
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
for i in 0..chunks {
let offset = i * 4;
let va = vld1q_f32(a_ptr.add(offset));
let vb = vld1q_f32(b_ptr.add(offset));
dot_acc = vfmaq_f32(dot_acc, va, vb);
}
let mut dot = vaddvq_f32(dot_acc);
let base = chunks * 4;
for i in 0..remainder {
dot += a[base + i] * b[base + i];
}
-dot
}
}
// ── Runtime dispatch ────────────────────────────────────────────────
/// Squared L2 (Euclidean) distance between two vectors.
///
/// Returns the sum of squared differences. Does NOT take the square root
/// because the ordering is preserved and sqrt is monotonic.
///
/// Automatically selects the best SIMD implementation at runtime:
/// - x86_64: AVX2+FMA (processes 8 floats per cycle)
/// - aarch64: NEON (processes 4 floats per cycle)
/// - Fallback: scalar loop
#[inline]
pub fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
return unsafe { avx2::l2_distance_avx2(a, b) };
}
}
#[cfg(target_arch = "aarch64")]
{
if std::arch::is_aarch64_feature_detected!("neon") {
return unsafe { neon::l2_distance_neon(a, b) };
}
}
l2_distance_scalar(a, b)
}
/// Cosine distance: `1 - cosine_similarity`.
///
/// Returns a value in `[0, 2]` where 0 means identical direction.
/// If either vector has zero norm, returns `1.0`.
///
/// Automatically selects the best SIMD implementation at runtime.
#[inline]
pub fn cosine_distance(a: &[f32], b: &[f32]) -> f32 {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
return unsafe { avx2::cosine_distance_avx2(a, b) };
}
}
#[cfg(target_arch = "aarch64")]
{
if std::arch::is_aarch64_feature_detected!("neon") {
return unsafe { neon::cosine_distance_neon(a, b) };
}
}
cosine_distance_scalar(a, b)
}
/// Inner (dot) product distance: `-dot(a, b)`.
///
/// Negated so that higher similarity yields a lower distance value,
/// which is consistent with the min-heap search ordering.
///
/// Automatically selects the best SIMD implementation at runtime.
#[inline]
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
return unsafe { avx2::dot_product_avx2(a, b) };
}
}
#[cfg(target_arch = "aarch64")]
{
if std::arch::is_aarch64_feature_detected!("neon") {
return unsafe { neon::dot_product_neon(a, b) };
}
}
dot_product_scalar(a, b)
}
// ── SIMD feature-gated wrappers (backward compatibility) ────────────
/// SIMD-accelerated squared L2 distance (same as `l2_distance` with runtime dispatch).
#[cfg(feature = "simd")]
#[inline]
pub fn l2_distance_simd(a: &[f32], b: &[f32]) -> f32 {
l2_distance(a, b)
}
/// SIMD-accelerated cosine distance (same as `cosine_distance` with runtime dispatch).
#[cfg(feature = "simd")]
#[inline]
pub fn cosine_distance_simd(a: &[f32], b: &[f32]) -> f32 {
cosine_distance(a, b)
}
/// SIMD-accelerated negative dot product distance (same as `dot_product` with runtime dispatch).
#[cfg(feature = "simd")]
#[inline]
pub fn dot_product_simd(a: &[f32], b: &[f32]) -> f32 {
dot_product(a, b)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn l2_identical_is_zero() {
let v = vec![1.0, 2.0, 3.0];
assert!((l2_distance(&v, &v) - 0.0).abs() < f32::EPSILON);
}
#[test]
fn l2_known_value() {
let a = vec![0.0, 0.0];
let b = vec![3.0, 4.0];
assert!((l2_distance(&a, &b) - 25.0).abs() < f32::EPSILON);
}
#[test]
fn l2_large_vector() {
// Test with a vector large enough to exercise SIMD paths (>8 elements).
let n = 256;
let a: Vec<f32> = (0..n).map(|i| i as f32 * 0.1).collect();
let b: Vec<f32> = (0..n).map(|i| i as f32 * 0.1 + 0.5).collect();
let dist = l2_distance(&a, &b);
let expected = l2_distance_scalar(&a, &b);
assert!(
(dist - expected).abs() < 1e-3,
"SIMD L2 mismatch: got {dist}, expected {expected}"
);
}
#[test]
fn l2_odd_length() {
// Non-multiple-of-8 length to test remainder handling.
let a: Vec<f32> = (0..13).map(|i| i as f32).collect();
let b: Vec<f32> = (0..13).map(|i| (i as f32) + 1.0).collect();
let dist = l2_distance(&a, &b);
// Each diff is 1.0, so sum = 13.0.
assert!((dist - 13.0).abs() < 1e-4);
}
#[test]
fn cosine_identical_is_zero() {
let v = vec![1.0, 2.0, 3.0];
assert!(cosine_distance(&v, &v) < 1e-6);
}
#[test]
fn cosine_orthogonal_is_one() {
let a = vec![1.0, 0.0];
let b = vec![0.0, 1.0];
assert!((cosine_distance(&a, &b) - 1.0).abs() < 1e-6);
}
#[test]
fn cosine_zero_vector() {
let a = vec![0.0, 0.0];
let b = vec![1.0, 2.0];
assert!((cosine_distance(&a, &b) - 1.0).abs() < f32::EPSILON);
}
#[test]
fn cosine_large_vector() {
let n = 256;
let a: Vec<f32> = (0..n).map(|i| (i as f32 + 1.0).sin()).collect();
let b: Vec<f32> = (0..n).map(|i| (i as f32 + 2.0).cos()).collect();
let dist = cosine_distance(&a, &b);
let expected = cosine_distance_scalar(&a, &b);
assert!(
(dist - expected).abs() < 1e-4,
"SIMD cosine mismatch: got {dist}, expected {expected}"
);
}
#[test]
fn dot_product_known_value() {
let a = vec![1.0, 2.0, 3.0];
let b = vec![4.0, 5.0, 6.0];
// dot = 4 + 10 + 18 = 32, negated = -32
assert!((dot_product(&a, &b) - (-32.0)).abs() < f32::EPSILON);
}
#[test]
fn dot_product_large_vector() {
let n = 256;
let a: Vec<f32> = (0..n).map(|i| i as f32 * 0.01).collect();
let b: Vec<f32> = (0..n).map(|i| (n - i) as f32 * 0.01).collect();
let dist = dot_product(&a, &b);
let expected = dot_product_scalar(&a, &b);
assert!(
(dist - expected).abs() < 1e-2,
"SIMD dot mismatch: got {dist}, expected {expected}"
);
}
#[test]
fn scalar_matches_dispatch() {
// Ensure the dispatched version matches scalar on various sizes.
for n in [1, 2, 3, 7, 8, 9, 15, 16, 17, 31, 32, 100] {
let a: Vec<f32> = (0..n).map(|i| (i as f32 * 1.7).sin()).collect();
let b: Vec<f32> = (0..n).map(|i| (i as f32 * 2.3).cos()).collect();
let l2 = l2_distance(&a, &b);
let l2s = l2_distance_scalar(&a, &b);
assert!((l2 - l2s).abs() < 1e-3, "L2 mismatch for n={n}");
let cos = cosine_distance(&a, &b);
let coss = cosine_distance_scalar(&a, &b);
assert!((cos - coss).abs() < 1e-4, "Cosine mismatch for n={n}");
let dp = dot_product(&a, &b);
let dps = dot_product_scalar(&a, &b);
assert!((dp - dps).abs() < 1e-3, "Dot mismatch for n={n}");
}
}
}

540
vendor/ruvector/crates/rvf/rvf-index/src/hnsw.rs vendored Normal file
View File

@@ -0,0 +1,540 @@
//! Core HNSW (Hierarchical Navigable Small World) graph implementation.
//!
//! Implements the algorithm from Malkov & Yashunin (2018) with:
//! - Configurable M (max neighbors per layer) and ef_construction
//! - Layer selection via P = 1/ln(M), level = floor(-ln(random) * P)
//! - Greedy search at upper layers, beam search at layer 0
extern crate alloc;
use alloc::collections::BTreeMap;
use alloc::vec;
use alloc::vec::Vec;
use crate::traits::VectorStore;
/// Configuration for HNSW graph construction.
#[derive(Clone, Debug)]
pub struct HnswConfig {
/// Maximum number of neighbors per node per layer (layer > 0).
pub m: usize,
/// Maximum number of neighbors at layer 0 (typically 2*M).
pub m0: usize,
/// Size of the dynamic candidate list during construction.
pub ef_construction: usize,
}
impl Default for HnswConfig {
fn default() -> Self {
Self {
m: 16,
m0: 32,
ef_construction: 200,
}
}
}
/// A single layer of the HNSW graph, mapping node IDs to their neighbor lists.
#[derive(Clone, Debug, Default)]
pub struct HnswLayer {
/// Node ID -> sorted list of neighbor IDs.
pub adjacency: BTreeMap<u64, Vec<u64>>,
}
impl HnswLayer {
/// Returns true if this layer contains the given node.
#[inline]
pub fn contains(&self, id: u64) -> bool {
self.adjacency.contains_key(&id)
}
/// Returns the neighbors of a node, or an empty slice if not present.
#[inline]
pub fn neighbors(&self, id: u64) -> &[u64] {
self.adjacency.get(&id).map_or(&[], |v| v.as_slice())
}
/// Number of nodes in this layer.
#[inline]
pub fn len(&self) -> usize {
self.adjacency.len()
}
/// Returns true if the layer has no nodes.
#[inline]
pub fn is_empty(&self) -> bool {
self.adjacency.is_empty()
}
}
/// The full HNSW graph structure.
#[derive(Clone, Debug)]
pub struct HnswGraph {
/// Layers from bottom (0) to top.
pub layers: Vec<HnswLayer>,
/// The entry point node ID (node at the highest layer).
pub entry_point: Option<u64>,
/// The highest occupied layer index.
pub max_layer: usize,
/// Max neighbors per layer (> 0).
pub m: usize,
/// Max neighbors at layer 0.
pub m0: usize,
/// ef_construction parameter.
pub ef_construction: usize,
/// Level normalization factor: 1 / ln(M).
ml: f64,
}
impl HnswGraph {
/// Create a new empty HNSW graph with the given configuration.
pub fn new(config: &HnswConfig) -> Self {
Self {
layers: vec![HnswLayer::default()],
entry_point: None,
max_layer: 0,
m: config.m,
m0: config.m0,
ef_construction: config.ef_construction,
ml: 1.0 / (config.m as f64).ln(),
}
}
/// Select a random level for a new node.
/// Level = floor(-ln(uniform(0,1)) * ml).
fn random_level(&self, rng_val: f64) -> usize {
let r = if rng_val <= 0.0 { 1e-10 } else { rng_val };
(-r.ln() * self.ml).floor() as usize
}
/// Insert a new node into the HNSW graph.
///
/// `id`: the node ID to insert.
/// `rng_val`: a uniform random value in (0, 1) for level selection.
/// `vectors`: provides access to all vectors by ID.
/// `distance_fn`: distance function between two vectors.
pub fn insert(
&mut self,
id: u64,
rng_val: f64,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) {
let level = self.random_level(rng_val);
// Ensure we have enough layers.
while self.layers.len() <= level {
self.layers.push(HnswLayer::default());
}
// Add the node to each layer from 0 to `level`.
for l in 0..=level {
self.layers[l].adjacency.entry(id).or_default();
}
let query_vec = match vectors.get_vector(id) {
Some(v) => v,
None => return,
};
if self.entry_point.is_none() {
// First node.
self.entry_point = Some(id);
self.max_layer = level;
return;
}
let ep = self.entry_point.unwrap();
// Phase 1: greedy search from top layer down to level+1.
let mut current_ep = ep;
let top = self.max_layer;
if top > level {
for l in (level + 1..=top).rev() {
current_ep = self.greedy_closest(query_vec, current_ep, l, vectors, distance_fn);
}
}
// Phase 2: at each layer from min(level, max_layer) down to 0,
// do a beam search and connect neighbors.
let start_layer = level.min(top);
let mut entry_points = vec![current_ep];
for l in (0..=start_layer).rev() {
let max_neighbors = if l == 0 { self.m0 } else { self.m };
let candidates = self.search_layer(
query_vec,
&entry_points,
self.ef_construction,
l,
vectors,
distance_fn,
);
// Select the closest `max_neighbors` candidates.
let selected: Vec<(u64, f32)> =
candidates.iter().take(max_neighbors).cloned().collect();
// Connect the new node to selected neighbors.
let neighbor_ids: Vec<u64> = selected.iter().map(|&(nid, _)| nid).collect();
self.layers[l].adjacency.insert(id, neighbor_ids.clone());
// Bidirectional: add the new node as a neighbor of each selected node,
// then prune if over the limit.
for &nid in &neighbor_ids {
let nlist = self.layers[l].adjacency.entry(nid).or_default();
if !nlist.contains(&id) {
nlist.push(id);
}
if nlist.len() > max_neighbors {
// Prune: keep only the closest max_neighbors.
self.prune_neighbors(nid, l, max_neighbors, vectors, distance_fn);
}
}
// Use the selected candidates as entry points for the next layer down.
entry_points = selected.iter().map(|&(nid, _)| nid).collect();
}
// Update entry point if the new node is at a higher layer.
if level > self.max_layer {
self.entry_point = Some(id);
self.max_layer = level;
}
}
/// Greedy search: starting from `ep`, walk to the closest node at `layer`.
fn greedy_closest(
&self,
query: &[f32],
ep: u64,
layer: usize,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> u64 {
let mut current = ep;
let mut current_dist = match vectors.get_vector(ep) {
Some(v) => distance_fn(query, v),
None => return ep,
};
loop {
let mut changed = false;
let neighbors = self.layers[layer].neighbors(current);
for &nid in neighbors {
if let Some(nv) = vectors.get_vector(nid) {
let d = distance_fn(query, nv);
if d < current_dist {
current = nid;
current_dist = d;
changed = true;
}
}
}
if !changed {
break;
}
}
current
}
/// Beam search at a given layer. Returns candidates sorted by distance (ascending).
fn search_layer(
&self,
query: &[f32],
entry_points: &[u64],
ef: usize,
layer: usize,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
#[cfg(not(feature = "std"))]
use alloc::collections::BTreeSet as HashSet;
#[cfg(feature = "std")]
use std::collections::HashSet;
let mut visited = HashSet::new();
// candidates sorted by (distance, id) — acts as a min-heap.
let mut candidates: Vec<(u64, f32)> = Vec::new();
let mut results: Vec<(u64, f32)> = Vec::new();
for &ep in entry_points {
if visited.insert(ep) {
if let Some(v) = vectors.get_vector(ep) {
let d = distance_fn(query, v);
candidates.push((ep, d));
results.push((ep, d));
}
}
}
candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
let mut candidate_idx = 0;
while candidate_idx < candidates.len() {
let (cid, cdist) = candidates[candidate_idx];
candidate_idx += 1;
// If the closest candidate is farther than the worst result and
// we already have `ef` results, stop.
if results.len() >= ef {
let worst_dist = results.last().map_or(f32::MAX, |r| r.1);
if cdist > worst_dist {
break;
}
}
let neighbors = self.layers[layer].neighbors(cid);
for &nid in neighbors {
if !visited.insert(nid) {
continue;
}
if let Some(nv) = vectors.get_vector(nid) {
let d = distance_fn(query, nv);
let worst_dist = if results.len() >= ef {
results.last().map_or(f32::MAX, |r| r.1)
} else {
f32::MAX
};
if d < worst_dist || results.len() < ef {
// Insert into candidates (sorted).
let pos = candidates[candidate_idx..]
.binary_search_by(|probe| {
probe
.1
.partial_cmp(&d)
.unwrap_or(core::cmp::Ordering::Equal)
})
.unwrap_or_else(|e| e);
candidates.insert(candidate_idx + pos, (nid, d));
// Insert into results (sorted).
let rpos = results
.binary_search_by(|probe| {
probe
.1
.partial_cmp(&d)
.unwrap_or(core::cmp::Ordering::Equal)
})
.unwrap_or_else(|e| e);
results.insert(rpos, (nid, d));
if results.len() > ef {
results.pop();
}
}
}
}
}
results
}
/// Prune neighbors of a node to keep only the closest `max_neighbors`.
fn prune_neighbors(
&mut self,
node: u64,
layer: usize,
max_neighbors: usize,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) {
let node_vec = match vectors.get_vector(node) {
Some(v) => v,
None => return,
};
let neighbors = match self.layers[layer].adjacency.get(&node) {
Some(n) => n.clone(),
None => return,
};
let mut scored: Vec<(u64, f32)> = neighbors
.iter()
.filter_map(|&nid| {
vectors
.get_vector(nid)
.map(|nv| (nid, distance_fn(node_vec, nv)))
})
.collect();
scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
scored.truncate(max_neighbors);
let pruned: Vec<u64> = scored.into_iter().map(|(nid, _)| nid).collect();
self.layers[layer].adjacency.insert(node, pruned);
}
/// Search the HNSW graph for the `k` nearest neighbors of `query`.
///
/// `ef_search`: size of the dynamic candidate list during search.
/// Returns a list of `(node_id, distance)` sorted by distance (ascending).
pub fn search(
&self,
query: &[f32],
k: usize,
ef_search: usize,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
let ep = match self.entry_point {
Some(ep) => ep,
None => return Vec::new(),
};
let ef = ef_search.max(k);
// Phase 1: greedy search from top layer down to layer 1.
let mut current_ep = ep;
for l in (1..=self.max_layer).rev() {
current_ep = self.greedy_closest(query, current_ep, l, vectors, distance_fn);
}
// Phase 2: beam search at layer 0.
let mut results = self.search_layer(query, &[current_ep], ef, 0, vectors, distance_fn);
results.truncate(k);
results
}
/// Returns the total number of nodes across all layers.
pub fn node_count(&self) -> usize {
self.layers.first().map_or(0, |l| l.adjacency.len())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::distance::l2_distance;
use crate::traits::InMemoryVectorStore;
fn make_config() -> HnswConfig {
HnswConfig {
m: 8,
m0: 16,
ef_construction: 100,
}
}
#[test]
fn empty_graph_search_returns_empty() {
let config = make_config();
let graph = HnswGraph::new(&config);
let store = InMemoryVectorStore::new(vec![vec![0.0; 4]]);
let results = graph.search(&[0.0; 4], 5, 50, &store, &l2_distance);
assert!(results.is_empty());
}
#[test]
fn insert_single_node() {
let config = make_config();
let mut graph = HnswGraph::new(&config);
let store = InMemoryVectorStore::new(vec![vec![1.0, 2.0, 3.0]]);
graph.insert(0, 0.5, &store, &l2_distance);
assert_eq!(graph.entry_point, Some(0));
assert_eq!(graph.node_count(), 1);
}
#[test]
fn insert_and_search_small() {
let config = make_config();
let mut graph = HnswGraph::new(&config);
let vectors: Vec<Vec<f32>> = (0..20)
.map(|i| vec![i as f32, (i * 2) as f32, (i * 3) as f32])
.collect();
let store = InMemoryVectorStore::new(vectors);
// Insert all with deterministic pseudo-random values.
for i in 0..20u64 {
let rng = ((i * 7 + 3) % 100) as f64 / 100.0;
graph.insert(i, rng, &store, &l2_distance);
}
// Search for a query near node 10.
let query = [10.0, 20.0, 30.0];
let results = graph.search(&query, 3, 50, &store, &l2_distance);
assert!(!results.is_empty());
// Node 10 should be the closest (exact match).
assert_eq!(results[0].0, 10);
}
/// Build HNSW with 1000 random vectors, verify recall@10 >= 0.95.
#[test]
fn recall_at_10_1000_vectors() {
use alloc::collections::BTreeSet;
let n = 1000;
let dim = 32;
// Generate deterministic pseudo-random vectors using a simple LCG.
let mut seed: u64 = 42;
let vectors: Vec<Vec<f32>> = (0..n)
.map(|_| {
(0..dim)
.map(|_| {
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
(seed >> 33) as f32 / (1u64 << 31) as f32
})
.collect()
})
.collect();
let store = InMemoryVectorStore::new(vectors.clone());
// Build the graph.
let config = HnswConfig {
m: 16,
m0: 32,
ef_construction: 200,
};
let mut graph = HnswGraph::new(&config);
let mut rng_seed: u64 = 123;
for i in 0..n as u64 {
rng_seed = rng_seed.wrapping_mul(6364136223846793005).wrapping_add(1);
let rng_val = (rng_seed >> 33) as f64 / (1u64 << 31) as f64;
let rng_val = rng_val.clamp(0.001, 0.999);
graph.insert(i, rng_val, &store, &l2_distance);
}
// Compute brute-force ground truth and measure recall.
let num_queries = 50;
let k = 10;
let ef_search = 200;
let mut total_recall = 0.0;
let mut query_seed: u64 = 999;
for _ in 0..num_queries {
// Generate a random query.
let query: Vec<f32> = (0..dim)
.map(|_| {
query_seed = query_seed.wrapping_mul(6364136223846793005).wrapping_add(1);
(query_seed >> 33) as f32 / (1u64 << 31) as f32
})
.collect();
// Brute-force top-k.
let mut all_dists: Vec<(u64, f32)> = (0..n as u64)
.map(|i| (i, l2_distance(&query, &vectors[i as usize])))
.collect();
all_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
let gt_set: BTreeSet<u64> = all_dists.iter().take(k).map(|&(id, _)| id).collect();
// HNSW search.
let results = graph.search(&query, k, ef_search, &store, &l2_distance);
let result_set: BTreeSet<u64> = results.iter().map(|&(id, _)| id).collect();
let overlap = gt_set.intersection(&result_set).count();
total_recall += overlap as f64 / k as f64;
}
let avg_recall = total_recall / num_queries as f64;
assert!(
avg_recall >= 0.95,
"Recall@10 = {:.3}, expected >= 0.95",
avg_recall
);
}
}

248
vendor/ruvector/crates/rvf/rvf-index/src/layers.rs vendored Normal file
View File

@@ -0,0 +1,248 @@
//! Progressive layer model (Layer A / B / C) for RVF indexing.
//!
//! Each layer is independently useful and stores a different granularity
//! of the HNSW graph, enabling progressive availability.
extern crate alloc;
use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use crate::hnsw::HnswLayer;
/// Which index layer a piece of data belongs to.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum IndexLayer {
/// Entry points + coarse routing. Always present, loaded first (< 5ms).
A = 0,
/// Partial adjacency for the hot region. Loaded second (100ms-1s).
B = 1,
/// Full adjacency for every node. Loaded last (seconds to minutes).
C = 2,
}
impl TryFrom<u8> for IndexLayer {
type Error = u8;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
0 => Ok(Self::A),
1 => Ok(Self::B),
2 => Ok(Self::C),
other => Err(other),
}
}
}
/// Entry in the centroid-to-partition map.
#[derive(Clone, Debug)]
pub struct PartitionEntry {
/// Which centroid owns this partition.
pub centroid_id: u32,
/// First vector ID in this partition.
pub vector_id_start: u64,
/// Last vector ID in this partition (exclusive).
pub vector_id_end: u64,
/// Segment ID containing the vector data.
pub segment_ref: u64,
/// Block offset within the segment.
pub block_ref: u32,
}
/// Layer A: Entry Points + Coarse Routing.
///
/// Contains:
/// - HNSW entry points (node ID + layer)
/// - Top-layer adjacency lists (layers >= threshold)
/// - Cluster centroids for IVF-style partition routing
/// - Centroid-to-partition map
#[derive(Clone, Debug)]
pub struct LayerA {
/// Entry points: `(node_id, max_layer)`.
pub entry_points: Vec<(u64, u32)>,
/// Top-layer adjacency: HNSW layers at the highest levels.
/// Index 0 = the highest layer, etc.
pub top_layers: Vec<HnswLayer>,
/// The HNSW layer index where top_layers[0] starts.
pub top_layer_start: usize,
/// Cluster centroids for partition routing.
pub centroids: Vec<Vec<f32>>,
/// Map from centroid to vector ID ranges.
pub partition_map: Vec<PartitionEntry>,
}
/// Layer B: Partial Adjacency for the hot working set.
///
/// Contains neighbor lists for the most-accessed nodes (determined by
/// temperature sketch). Typically covers 10-20% of total nodes.
#[derive(Clone, Debug)]
pub struct LayerB {
/// Partial adjacency: node_id -> neighbor list.
/// Only nodes in the hot region are present.
pub partial_adjacency: BTreeMap<u64, Vec<u64>>,
/// Ranges of node IDs covered by this layer.
pub covered_ranges: Vec<(u64, u64)>,
}
impl LayerB {
/// Returns true if the given node has adjacency data in this layer.
#[inline]
pub fn has_node(&self, id: u64) -> bool {
self.partial_adjacency.contains_key(&id)
}
/// Returns neighbors for a node, or `None` if not in the hot region.
#[inline]
pub fn neighbors(&self, id: u64) -> Option<&[u64]> {
self.partial_adjacency.get(&id).map(|v| v.as_slice())
}
}
/// Layer C: Full Adjacency.
///
/// Complete neighbor lists for every node at every HNSW level.
/// This is the traditional full HNSW graph.
#[derive(Clone, Debug)]
pub struct LayerC {
/// Full adjacency at every HNSW layer. Index 0 = layer 0 (bottom).
pub full_adjacency: Vec<HnswLayer>,
}
/// Aggregated state of all loaded index layers.
#[derive(Clone, Debug)]
pub struct IndexState {
pub layer_a: Option<LayerA>,
pub layer_b: Option<LayerB>,
pub layer_c: Option<LayerC>,
/// Total number of nodes in the full graph (known from metadata).
pub total_nodes: u64,
}
/// Estimate recall@10 based on which layers are currently loaded.
///
/// These are approximate lower-bound estimates based on the spec:
/// - A only: 0.65-0.75
/// - A + B: 0.85-0.92
/// - A + B + C: 0.95-0.99
pub fn available_recall(state: &IndexState) -> f32 {
match (&state.layer_a, &state.layer_b, &state.layer_c) {
(None, _, _) => 0.0,
(Some(_), None, None) => 0.70,
(Some(_), Some(b), None) => {
// Recall scales with coverage of partial adjacency.
let covered_nodes: u64 = b
.covered_ranges
.iter()
.map(|(start, end)| end.saturating_sub(*start))
.sum();
let coverage = if state.total_nodes > 0 {
covered_nodes as f32 / state.total_nodes as f32
} else {
0.0
};
// Scale between 0.70 (no B coverage) and 0.92 (full B coverage).
0.70 + coverage * 0.22
}
(Some(_), _, Some(_)) => 0.97,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn index_layer_round_trip() {
assert_eq!(IndexLayer::try_from(0), Ok(IndexLayer::A));
assert_eq!(IndexLayer::try_from(1), Ok(IndexLayer::B));
assert_eq!(IndexLayer::try_from(2), Ok(IndexLayer::C));
assert_eq!(IndexLayer::try_from(3), Err(3));
}
#[test]
fn recall_no_layers() {
let state = IndexState {
layer_a: None,
layer_b: None,
layer_c: None,
total_nodes: 1000,
};
assert!((available_recall(&state) - 0.0).abs() < f32::EPSILON);
}
#[test]
fn recall_a_only() {
let state = IndexState {
layer_a: Some(LayerA {
entry_points: vec![(0, 5)],
top_layers: vec![],
top_layer_start: 5,
centroids: vec![],
partition_map: vec![],
}),
layer_b: None,
layer_c: None,
total_nodes: 1000,
};
assert!((available_recall(&state) - 0.70).abs() < 0.01);
}
#[test]
fn recall_a_plus_b() {
let state = IndexState {
layer_a: Some(LayerA {
entry_points: vec![(0, 5)],
top_layers: vec![],
top_layer_start: 5,
centroids: vec![],
partition_map: vec![],
}),
layer_b: Some(LayerB {
partial_adjacency: BTreeMap::new(),
covered_ranges: vec![(0, 500)],
}),
layer_c: None,
total_nodes: 1000,
};
let recall = available_recall(&state);
assert!(recall > 0.70);
assert!(recall < 0.93);
}
#[test]
fn recall_full() {
let state = IndexState {
layer_a: Some(LayerA {
entry_points: vec![(0, 5)],
top_layers: vec![],
top_layer_start: 5,
centroids: vec![],
partition_map: vec![],
}),
layer_b: Some(LayerB {
partial_adjacency: BTreeMap::new(),
covered_ranges: vec![(0, 1000)],
}),
layer_c: Some(LayerC {
full_adjacency: vec![],
}),
total_nodes: 1000,
};
assert!(available_recall(&state) >= 0.95);
}
#[test]
fn layer_b_has_node() {
let mut adj = BTreeMap::new();
adj.insert(42, vec![1, 2, 3]);
let b = LayerB {
partial_adjacency: adj,
covered_ranges: vec![(0, 100)],
};
assert!(b.has_node(42));
assert!(!b.has_node(99));
assert_eq!(b.neighbors(42), Some([1u64, 2, 3].as_slice()));
assert_eq!(b.neighbors(99), None);
}
}

30
vendor/ruvector/crates/rvf/rvf-index/src/lib.rs vendored Normal file
View File

@@ -0,0 +1,30 @@
//! Progressive HNSW indexing for the RuVector Format (RVF).
//!
//! This crate implements the three-layer progressive indexing model:
//!
//! - **Layer A**: Entry points + coarse routing (< 5ms load, ~0.70 recall)
//! - **Layer B**: Partial adjacency for hot region (100ms-1s load, ~0.85 recall)
//! - **Layer C**: Full HNSW adjacency (seconds load, >= 0.95 recall)
#![cfg_attr(not(feature = "std"), no_std)]
extern crate alloc;
pub mod builder;
pub mod codec;
pub mod distance;
pub mod hnsw;
pub mod layers;
pub mod progressive;
pub mod traits;
pub use builder::{build_full_index, build_layer_a, build_layer_b, build_layer_c};
pub use codec::{decode_index_seg, encode_index_seg, CodecError, IndexSegData, IndexSegHeader};
pub use distance::{cosine_distance, dot_product, l2_distance};
pub use hnsw::{HnswConfig, HnswGraph, HnswLayer};
pub use layers::{IndexLayer, IndexState, LayerA, LayerB, LayerC, PartitionEntry};
pub use progressive::ProgressiveIndex;
pub use traits::VectorStore;
#[cfg(feature = "std")]
pub use traits::InMemoryVectorStore;

519
vendor/ruvector/crates/rvf/rvf-index/src/progressive.rs vendored Normal file
View File

@@ -0,0 +1,519 @@
//! Progressive search combining Layers A, B, and C.
//!
//! Depending on which layers are loaded, the search adapts:
//! - Layer A only: centroid routing + top-layer HNSW + hot cache scan
//! - A + B: HNSW through hot region, fallback centroid scan for cold
//! - A + B + C: full HNSW at all layers
extern crate alloc;
use alloc::collections::BTreeSet;
use alloc::vec::Vec;
use crate::distance::l2_distance;
use crate::layers::{LayerA, LayerB, LayerC};
use crate::traits::VectorStore;
/// Progressive index that adapts search quality based on loaded layers.
#[derive(Clone, Debug)]
pub struct ProgressiveIndex {
pub layer_a: Option<LayerA>,
pub layer_b: Option<LayerB>,
pub layer_c: Option<LayerC>,
}
impl ProgressiveIndex {
/// Create a new empty progressive index.
pub fn new() -> Self {
Self {
layer_a: None,
layer_b: None,
layer_c: None,
}
}
/// Search using whatever layers are available.
///
/// Returns `(node_id, distance)` pairs sorted by distance ascending.
pub fn search(
&self,
query: &[f32],
k: usize,
ef_search: usize,
vectors: &dyn VectorStore,
) -> Vec<(u64, f32)> {
self.search_with_distance(query, k, ef_search, vectors, &l2_distance)
}
/// Search with a custom distance function.
pub fn search_with_distance(
&self,
query: &[f32],
k: usize,
ef_search: usize,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
match (&self.layer_a, &self.layer_b, &self.layer_c) {
(None, _, _) => Vec::new(),
(Some(a), None, None) => self.search_layer_a_only(query, k, a, vectors, distance_fn),
(Some(a), Some(b), None) => {
self.search_a_plus_b(query, k, ef_search, a, b, vectors, distance_fn)
}
(Some(_a), _, Some(c)) => {
self.search_full(query, k, ef_search, c, vectors, distance_fn)
}
}
}
/// Search using only Layer A: centroid routing + top-layer HNSW traversal.
fn search_layer_a_only(
&self,
query: &[f32],
k: usize,
layer_a: &LayerA,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
let mut candidates: Vec<(u64, f32)> = Vec::new();
// Step 1: find nearest centroids.
let n_probe = 10.min(layer_a.centroids.len());
let mut centroid_dists: Vec<(usize, f32)> = layer_a
.centroids
.iter()
.enumerate()
.map(|(i, c)| (i, distance_fn(query, c)))
.collect();
centroid_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
centroid_dists.truncate(n_probe);
// Step 2: HNSW search through top layers using Layer A entry points.
if let Some(&(ep, _)) = layer_a.entry_points.first() {
let mut current = ep;
// Greedy walk through top layers.
for tl in &layer_a.top_layers {
current = greedy_walk(query, current, tl, vectors, distance_fn);
}
if let Some(v) = vectors.get_vector(current) {
candidates.push((current, distance_fn(query, v)));
}
// Also check neighbors of the landing node at the lowest top layer.
if let Some(last_tl) = layer_a.top_layers.last() {
for &nid in last_tl.neighbors(current) {
if let Some(nv) = vectors.get_vector(nid) {
candidates.push((nid, distance_fn(query, nv)));
}
}
}
}
// Step 3: scan vectors in the nearest centroid partitions.
for &(ci, _) in &centroid_dists {
for part in &layer_a.partition_map {
if part.centroid_id == ci as u32 {
// Scan vectors in this partition.
for vid in part.vector_id_start..part.vector_id_end {
if let Some(v) = vectors.get_vector(vid) {
candidates.push((vid, distance_fn(query, v)));
}
}
}
}
}
// Deduplicate and return top-k.
dedup_top_k(&mut candidates, k)
}
/// Search using Layers A + B: HNSW through hot region, fallback for cold.
#[allow(clippy::too_many_arguments)]
fn search_a_plus_b(
&self,
query: &[f32],
k: usize,
ef_search: usize,
layer_a: &LayerA,
layer_b: &LayerB,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
let ef = ef_search.max(k);
let mut visited = BTreeSet::new();
let mut results: Vec<(u64, f32)> = Vec::new();
// Start with Layer A routing to find the best entry into hot region.
let entry = layer_a.entry_points.first().map(|&(ep, _)| ep).unwrap_or(0);
let mut current = entry;
for tl in &layer_a.top_layers {
current = greedy_walk(query, current, tl, vectors, distance_fn);
}
// Beam search through Layer B's partial adjacency.
let mut candidates: Vec<(u64, f32)> = Vec::new();
if let Some(v) = vectors.get_vector(current) {
let d = distance_fn(query, v);
candidates.push((current, d));
results.push((current, d));
visited.insert(current);
}
let mut idx = 0;
while idx < candidates.len() {
let (cid, cdist) = candidates[idx];
idx += 1;
if results.len() >= ef {
let worst = results.last().map_or(f32::MAX, |r| r.1);
if cdist > worst {
break;
}
}
// Get neighbors: prefer Layer B, fallback to Layer A's top layers.
let neighbor_ids: Vec<u64> = if let Some(neighbors) = layer_b.neighbors(cid) {
neighbors.to_vec()
} else {
// Fallback: check top layers for any adjacency.
let mut fallback = Vec::new();
for tl in &layer_a.top_layers {
fallback.extend_from_slice(tl.neighbors(cid));
}
fallback
};
for nid in neighbor_ids {
if !visited.insert(nid) {
continue;
}
if let Some(nv) = vectors.get_vector(nid) {
let d = distance_fn(query, nv);
let worst = if results.len() >= ef {
results.last().map_or(f32::MAX, |r| r.1)
} else {
f32::MAX
};
if d < worst || results.len() < ef {
let pos = candidates[idx..]
.binary_search_by(|p| {
p.1.partial_cmp(&d).unwrap_or(core::cmp::Ordering::Equal)
})
.unwrap_or_else(|e| e);
candidates.insert(idx + pos, (nid, d));
let rpos = results
.binary_search_by(|p| {
p.1.partial_cmp(&d).unwrap_or(core::cmp::Ordering::Equal)
})
.unwrap_or_else(|e| e);
results.insert(rpos, (nid, d));
if results.len() > ef {
results.pop();
}
}
}
}
}
results.truncate(k);
results
}
/// Search using full Layer C HNSW graph.
fn search_full(
&self,
query: &[f32],
k: usize,
ef_search: usize,
layer_c: &LayerC,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
let ef = ef_search.max(k);
let max_layer = if layer_c.full_adjacency.is_empty() {
return Vec::new();
} else {
layer_c.full_adjacency.len() - 1
};
// Find the entry point: any node at the highest layer.
let entry = match layer_c.full_adjacency[max_layer].adjacency.keys().next() {
Some(&ep) => ep,
None => return Vec::new(),
};
// Phase 1: greedy descent through upper layers.
let mut current = entry;
for l in (1..=max_layer).rev() {
current = greedy_walk(
query,
current,
&layer_c.full_adjacency[l],
vectors,
distance_fn,
);
}
// Phase 2: beam search at layer 0.
beam_search_layer(
query,
&[current],
ef,
k,
&layer_c.full_adjacency[0],
vectors,
distance_fn,
)
}
}
impl Default for ProgressiveIndex {
fn default() -> Self {
Self::new()
}
}
// ── Helpers ──────────────────────────────────────────────────────
/// Greedy walk to the closest node in a single HNSW layer.
fn greedy_walk(
query: &[f32],
start: u64,
layer: &crate::hnsw::HnswLayer,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> u64 {
let mut current = start;
let mut current_dist = match vectors.get_vector(start) {
Some(v) => distance_fn(query, v),
None => return start,
};
loop {
let mut improved = false;
for &nid in layer.neighbors(current) {
if let Some(nv) = vectors.get_vector(nid) {
let d = distance_fn(query, nv);
if d < current_dist {
current = nid;
current_dist = d;
improved = true;
}
}
}
if !improved {
break;
}
}
current
}
/// Beam search at a single HNSW layer. Returns top-k results sorted by distance.
fn beam_search_layer(
query: &[f32],
entry_points: &[u64],
ef: usize,
k: usize,
layer: &crate::hnsw::HnswLayer,
vectors: &dyn VectorStore,
distance_fn: &dyn Fn(&[f32], &[f32]) -> f32,
) -> Vec<(u64, f32)> {
let mut visited = BTreeSet::new();
let mut candidates: Vec<(u64, f32)> = Vec::new();
let mut results: Vec<(u64, f32)> = Vec::new();
for &ep in entry_points {
if visited.insert(ep) {
if let Some(v) = vectors.get_vector(ep) {
let d = distance_fn(query, v);
candidates.push((ep, d));
results.push((ep, d));
}
}
}
candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
let mut idx = 0;
while idx < candidates.len() {
let (cid, cdist) = candidates[idx];
idx += 1;
if results.len() >= ef {
let worst = results.last().map_or(f32::MAX, |r| r.1);
if cdist > worst {
break;
}
}
for &nid in layer.neighbors(cid) {
if !visited.insert(nid) {
continue;
}
if let Some(nv) = vectors.get_vector(nid) {
let d = distance_fn(query, nv);
let worst = if results.len() >= ef {
results.last().map_or(f32::MAX, |r| r.1)
} else {
f32::MAX
};
if d < worst || results.len() < ef {
let pos = candidates[idx..]
.binary_search_by(|p| {
p.1.partial_cmp(&d).unwrap_or(core::cmp::Ordering::Equal)
})
.unwrap_or_else(|e| e);
candidates.insert(idx + pos, (nid, d));
let rpos = results
.binary_search_by(|p| {
p.1.partial_cmp(&d).unwrap_or(core::cmp::Ordering::Equal)
})
.unwrap_or_else(|e| e);
results.insert(rpos, (nid, d));
if results.len() > ef {
results.pop();
}
}
}
}
}
results.truncate(k);
results
}
/// Deduplicate candidates by node ID and return top-k by distance.
fn dedup_top_k(candidates: &mut [(u64, f32)], k: usize) -> Vec<(u64, f32)> {
candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
let mut seen = BTreeSet::new();
let mut result = Vec::with_capacity(k);
for &(id, dist) in candidates.iter() {
if seen.insert(id) {
result.push((id, dist));
if result.len() == k {
break;
}
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use crate::hnsw::{HnswConfig, HnswGraph};
use crate::layers::{LayerA, LayerC, PartitionEntry};
use crate::traits::InMemoryVectorStore;
fn make_test_vectors(n: usize, dim: usize) -> Vec<Vec<f32>> {
(0..n)
.map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect())
.collect()
}
#[test]
fn progressive_empty_returns_empty() {
let idx = ProgressiveIndex::new();
let store = InMemoryVectorStore::new(vec![vec![0.0; 4]]);
let results = idx.search(&[0.0; 4], 5, 50, &store);
assert!(results.is_empty());
}
#[test]
fn progressive_layer_a_only() {
let vectors = make_test_vectors(100, 4);
let store = InMemoryVectorStore::new(vectors.clone());
// Build a centroid from first 50 vectors (partition 0) and last 50 (partition 1).
let centroid_0: Vec<f32> = (0..4)
.map(|d| (0..50).map(|i| vectors[i][d]).sum::<f32>() / 50.0)
.collect();
let centroid_1: Vec<f32> = (0..4)
.map(|d| (50..100).map(|i| vectors[i][d]).sum::<f32>() / 50.0)
.collect();
let idx = ProgressiveIndex {
layer_a: Some(LayerA {
entry_points: vec![(0, 0)],
top_layers: vec![],
top_layer_start: 0,
centroids: vec![centroid_0, centroid_1],
partition_map: vec![
PartitionEntry {
centroid_id: 0,
vector_id_start: 0,
vector_id_end: 50,
segment_ref: 0,
block_ref: 0,
},
PartitionEntry {
centroid_id: 1,
vector_id_start: 50,
vector_id_end: 100,
segment_ref: 0,
block_ref: 0,
},
],
}),
layer_b: None,
layer_c: None,
};
let query = vectors[25].clone();
let results = idx.search(&query, 5, 50, &store);
assert!(!results.is_empty());
// The exact match should be found.
assert_eq!(results[0].0, 25);
}
#[test]
fn progressive_full_layer_c() {
let n = 200;
let dim = 4;
let vectors = make_test_vectors(n, dim);
let store = InMemoryVectorStore::new(vectors.clone());
// Build a full HNSW graph, then extract it as Layer C.
let config = HnswConfig {
m: 8,
m0: 16,
ef_construction: 100,
};
let mut graph = HnswGraph::new(&config);
for i in 0..n as u64 {
let rng = ((i * 7 + 3) % 100) as f64 / 100.0;
graph.insert(i, rng, &store, &l2_distance);
}
let layer_c = LayerC {
full_adjacency: graph.layers.clone(),
};
let idx = ProgressiveIndex {
layer_a: Some(LayerA {
entry_points: vec![(graph.entry_point.unwrap(), graph.max_layer as u32)],
top_layers: vec![],
top_layer_start: 0,
centroids: vec![],
partition_map: vec![],
}),
layer_b: None,
layer_c: Some(layer_c),
};
// Query for a known vector.
let target = 100;
let query = vectors[target].clone();
let results = idx.search(&query, 10, 100, &store);
assert!(!results.is_empty());
assert_eq!(results[0].0, target as u64);
}
}

46
vendor/ruvector/crates/rvf/rvf-index/src/traits.rs vendored Normal file
View File

@@ -0,0 +1,46 @@
//! Vector storage trait for abstract access to vector data.
/// Provides access to vectors by ID without requiring a specific storage layout.
///
/// Implementations may back this with in-memory arrays, mmap'd VEC_SEGs,
/// or any other source of vector data.
pub trait VectorStore {
/// Return the vector for the given node ID, or `None` if not present.
fn get_vector(&self, id: u64) -> Option<&[f32]>;
/// The dimensionality of all vectors in this store.
fn dimension(&self) -> usize;
}
/// Simple in-memory vector store backed by a `Vec<Vec<f32>>`.
///
/// IDs are assumed to be contiguous starting from 0.
#[cfg(feature = "std")]
pub struct InMemoryVectorStore {
vectors: Vec<Vec<f32>>,
dim: usize,
}
#[cfg(feature = "std")]
impl InMemoryVectorStore {
/// Create a new store from a collection of vectors.
///
/// # Panics
///
/// Panics if `vectors` is empty.
pub fn new(vectors: Vec<Vec<f32>>) -> Self {
let dim = vectors.first().map_or(0, |v| v.len());
Self { vectors, dim }
}
}
#[cfg(feature = "std")]
impl VectorStore for InMemoryVectorStore {
fn get_vector(&self, id: u64) -> Option<&[f32]> {
self.vectors.get(id as usize).map(|v| v.as_slice())
}
fn dimension(&self) -> usize {
self.dim
}
}