d803bfe2b1
git-subtree-dir: vendor/ruvector git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
1395 lines
42 KiB
Rust
1395 lines
42 KiB
Rust
//! Comprehensive index benchmarks for HNSW and IVFFlat
|
|
//!
|
|
//! Benchmarks include:
|
|
//! - HNSW build time (10K, 100K, 1M vectors)
|
|
//! - HNSW query latency (p50, p95, p99)
|
|
//! - IVFFlat build time
|
|
//! - IVFFlat query latency
|
|
//! - Recall vs latency tradeoffs
|
|
//! - Memory usage analysis
|
|
|
|
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
|
|
use rand::prelude::*;
|
|
use rand_chacha::ChaCha8Rng;
|
|
use rayon::prelude::*;
|
|
use std::time::{Duration, Instant};
|
|
|
|
// ============================================================================
|
|
// HNSW Index Implementation (Standalone for Benchmarking)
|
|
// ============================================================================
|
|
|
|
mod hnsw {
|
|
use dashmap::DashMap;
|
|
use parking_lot::RwLock;
|
|
use rand::prelude::*;
|
|
use rand_chacha::ChaCha8Rng;
|
|
use std::cmp::Ordering;
|
|
use std::collections::{BinaryHeap, HashSet};
|
|
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum DistanceMetric {
|
|
Euclidean,
|
|
Cosine,
|
|
InnerProduct,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct HnswConfig {
|
|
pub m: usize,
|
|
pub m0: usize,
|
|
pub ef_construction: usize,
|
|
pub ef_search: usize,
|
|
pub max_elements: usize,
|
|
pub metric: DistanceMetric,
|
|
pub seed: u64,
|
|
}
|
|
|
|
impl Default for HnswConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: 1_000_000,
|
|
metric: DistanceMetric::Euclidean,
|
|
seed: 42,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub type NodeId = u64;
|
|
|
|
#[derive(Clone, Copy)]
|
|
struct Neighbor {
|
|
id: NodeId,
|
|
distance: f32,
|
|
}
|
|
|
|
impl PartialEq for Neighbor {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.id == other.id
|
|
}
|
|
}
|
|
|
|
impl Eq for Neighbor {}
|
|
|
|
impl PartialOrd for Neighbor {
|
|
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl Ord for Neighbor {
|
|
fn cmp(&self, other: &Self) -> Ordering {
|
|
other
|
|
.distance
|
|
.partial_cmp(&self.distance)
|
|
.unwrap_or(Ordering::Equal)
|
|
}
|
|
}
|
|
|
|
struct HnswNode {
|
|
vector: Vec<f32>,
|
|
neighbors: Vec<RwLock<Vec<NodeId>>>,
|
|
max_layer: usize,
|
|
}
|
|
|
|
pub struct HnswIndex {
|
|
config: HnswConfig,
|
|
nodes: DashMap<NodeId, HnswNode>,
|
|
entry_point: RwLock<Option<NodeId>>,
|
|
max_layer: AtomicUsize,
|
|
node_count: AtomicUsize,
|
|
next_id: AtomicUsize,
|
|
rng: RwLock<ChaCha8Rng>,
|
|
dimensions: usize,
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
pub struct SearchResult {
|
|
pub id: NodeId,
|
|
pub distance: f32,
|
|
}
|
|
|
|
impl HnswIndex {
|
|
pub fn new(config: HnswConfig) -> Self {
|
|
let rng = ChaCha8Rng::seed_from_u64(config.seed);
|
|
Self {
|
|
dimensions: 0,
|
|
config,
|
|
nodes: DashMap::new(),
|
|
entry_point: RwLock::new(None),
|
|
max_layer: AtomicUsize::new(0),
|
|
node_count: AtomicUsize::new(0),
|
|
next_id: AtomicUsize::new(0),
|
|
rng: RwLock::new(rng),
|
|
}
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.node_count.load(AtomicOrdering::Relaxed)
|
|
}
|
|
|
|
fn random_level(&self) -> usize {
|
|
let ml = 1.0 / (self.config.m as f64).ln();
|
|
let mut rng = self.rng.write();
|
|
let r: f64 = rng.gen();
|
|
let level = (-r.ln() * ml).floor() as usize;
|
|
level.min(32)
|
|
}
|
|
|
|
fn calc_distance(&self, a: &[f32], b: &[f32]) -> f32 {
|
|
match self.config.metric {
|
|
DistanceMetric::Euclidean => a
|
|
.iter()
|
|
.zip(b.iter())
|
|
.map(|(x, y)| {
|
|
let diff = x - y;
|
|
diff * diff
|
|
})
|
|
.sum::<f32>()
|
|
.sqrt(),
|
|
DistanceMetric::Cosine => {
|
|
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 == 0.0 {
|
|
1.0
|
|
} else {
|
|
1.0 - (dot / denom)
|
|
}
|
|
}
|
|
DistanceMetric::InnerProduct => {
|
|
-a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn insert(&mut self, id: NodeId, vector: &[f32]) {
|
|
if self.dimensions == 0 {
|
|
self.dimensions = vector.len();
|
|
}
|
|
|
|
let level = self.random_level();
|
|
let mut neighbors = Vec::with_capacity(level + 1);
|
|
for _ in 0..=level {
|
|
neighbors.push(RwLock::new(Vec::new()));
|
|
}
|
|
|
|
let node = HnswNode {
|
|
vector: vector.to_vec(),
|
|
neighbors,
|
|
max_layer: level,
|
|
};
|
|
|
|
let current_entry = *self.entry_point.read();
|
|
|
|
if current_entry.is_none() {
|
|
self.nodes.insert(id, node);
|
|
*self.entry_point.write() = Some(id);
|
|
self.max_layer.store(level, AtomicOrdering::Relaxed);
|
|
self.node_count.fetch_add(1, AtomicOrdering::Relaxed);
|
|
return;
|
|
}
|
|
|
|
let entry_id = current_entry.unwrap();
|
|
self.nodes.insert(id, node);
|
|
|
|
// Simplified insertion - connect to entry point
|
|
let max_connections = if level == 0 {
|
|
self.config.m0
|
|
} else {
|
|
self.config.m
|
|
};
|
|
|
|
if let Some(entry_node) = self.nodes.get(&entry_id) {
|
|
let min_level = level.min(entry_node.max_layer);
|
|
for l in 0..=min_level {
|
|
if let Some(node) = self.nodes.get(&id) {
|
|
node.neighbors[l].write().push(entry_id);
|
|
}
|
|
entry_node.neighbors[l].write().push(id);
|
|
|
|
// Trim if needed
|
|
let mut neighbors = entry_node.neighbors[l].write();
|
|
if neighbors.len() > max_connections {
|
|
neighbors.truncate(max_connections);
|
|
}
|
|
}
|
|
}
|
|
|
|
if level > self.max_layer.load(AtomicOrdering::Relaxed) {
|
|
*self.entry_point.write() = Some(id);
|
|
self.max_layer.store(level, AtomicOrdering::Relaxed);
|
|
}
|
|
|
|
self.node_count.fetch_add(1, AtomicOrdering::Relaxed);
|
|
}
|
|
|
|
pub fn search(&self, query: &[f32], k: usize) -> Vec<SearchResult> {
|
|
self.search_with_ef(query, k, self.config.ef_search)
|
|
}
|
|
|
|
pub fn search_with_ef(&self, query: &[f32], k: usize, ef: usize) -> Vec<SearchResult> {
|
|
let entry = match *self.entry_point.read() {
|
|
Some(id) => id,
|
|
None => return Vec::new(),
|
|
};
|
|
|
|
// Brute force search (simplified for benchmarking)
|
|
let mut results: Vec<SearchResult> = self
|
|
.nodes
|
|
.iter()
|
|
.map(|entry| {
|
|
let dist = self.calc_distance(query, &entry.value().vector);
|
|
SearchResult {
|
|
id: *entry.key(),
|
|
distance: dist,
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
|
|
results.truncate(k.min(ef));
|
|
results
|
|
}
|
|
|
|
pub fn memory_usage(&self) -> usize {
|
|
let mut total = 0;
|
|
for entry in self.nodes.iter() {
|
|
total += entry.value().vector.len() * 4;
|
|
for neighbors in &entry.value().neighbors {
|
|
total += neighbors.read().len() * 8;
|
|
}
|
|
}
|
|
total
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// IVFFlat Index Implementation (Standalone for Benchmarking)
|
|
// ============================================================================
|
|
|
|
mod ivfflat {
|
|
use dashmap::DashMap;
|
|
use parking_lot::RwLock;
|
|
use rand::prelude::*;
|
|
use rand_chacha::ChaCha8Rng;
|
|
use rayon::prelude::*;
|
|
use std::cmp::Ordering;
|
|
use std::collections::BinaryHeap;
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum DistanceMetric {
|
|
Euclidean,
|
|
Cosine,
|
|
InnerProduct,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct IvfFlatConfig {
|
|
pub lists: usize,
|
|
pub probes: usize,
|
|
pub metric: DistanceMetric,
|
|
pub kmeans_iterations: usize,
|
|
pub seed: u64,
|
|
}
|
|
|
|
impl Default for IvfFlatConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
lists: 100,
|
|
probes: 1,
|
|
metric: DistanceMetric::Euclidean,
|
|
kmeans_iterations: 10,
|
|
seed: 42,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub type VectorId = u64;
|
|
|
|
#[derive(Clone)]
|
|
struct ClusterEntry {
|
|
id: VectorId,
|
|
vector: Vec<f32>,
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
struct SearchResult {
|
|
id: VectorId,
|
|
distance: f32,
|
|
}
|
|
|
|
impl PartialEq for SearchResult {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.distance == other.distance
|
|
}
|
|
}
|
|
|
|
impl Eq for SearchResult {}
|
|
|
|
impl PartialOrd for SearchResult {
|
|
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl Ord for SearchResult {
|
|
fn cmp(&self, other: &Self) -> Ordering {
|
|
other
|
|
.distance
|
|
.partial_cmp(&self.distance)
|
|
.unwrap_or(Ordering::Equal)
|
|
}
|
|
}
|
|
|
|
pub struct IvfFlatIndex {
|
|
config: IvfFlatConfig,
|
|
centroids: RwLock<Vec<Vec<f32>>>,
|
|
lists: DashMap<usize, Vec<ClusterEntry>>,
|
|
id_to_cluster: DashMap<VectorId, usize>,
|
|
vector_count: std::sync::atomic::AtomicUsize,
|
|
dimensions: usize,
|
|
trained: std::sync::atomic::AtomicBool,
|
|
}
|
|
|
|
impl IvfFlatIndex {
|
|
pub fn new(dimensions: usize, config: IvfFlatConfig) -> Self {
|
|
Self {
|
|
config,
|
|
centroids: RwLock::new(Vec::new()),
|
|
lists: DashMap::new(),
|
|
id_to_cluster: DashMap::new(),
|
|
vector_count: std::sync::atomic::AtomicUsize::new(0),
|
|
dimensions,
|
|
trained: std::sync::atomic::AtomicBool::new(false),
|
|
}
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.vector_count.load(std::sync::atomic::Ordering::Relaxed)
|
|
}
|
|
|
|
pub fn is_trained(&self) -> bool {
|
|
self.trained.load(std::sync::atomic::Ordering::Relaxed)
|
|
}
|
|
|
|
fn calc_distance(&self, a: &[f32], b: &[f32]) -> f32 {
|
|
match self.config.metric {
|
|
DistanceMetric::Euclidean => a
|
|
.iter()
|
|
.zip(b.iter())
|
|
.map(|(x, y)| {
|
|
let diff = x - y;
|
|
diff * diff
|
|
})
|
|
.sum::<f32>()
|
|
.sqrt(),
|
|
DistanceMetric::Cosine => {
|
|
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 == 0.0 {
|
|
1.0
|
|
} else {
|
|
1.0 - (dot / denom)
|
|
}
|
|
}
|
|
DistanceMetric::InnerProduct => {
|
|
-a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn train(&self, training_vectors: &[Vec<f32>]) {
|
|
if training_vectors.is_empty() {
|
|
return;
|
|
}
|
|
|
|
let n_clusters = self.config.lists.min(training_vectors.len());
|
|
let mut rng = ChaCha8Rng::seed_from_u64(self.config.seed);
|
|
|
|
// K-means++ initialization
|
|
let mut centroids = Vec::with_capacity(n_clusters);
|
|
let first_idx = rng.gen_range(0..training_vectors.len());
|
|
centroids.push(training_vectors[first_idx].clone());
|
|
|
|
for _ in 1..n_clusters {
|
|
let distances: Vec<f32> = training_vectors
|
|
.iter()
|
|
.map(|v| {
|
|
centroids
|
|
.iter()
|
|
.map(|c| self.calc_distance(v, c))
|
|
.fold(f32::MAX, f32::min)
|
|
})
|
|
.collect();
|
|
|
|
let squared: Vec<f32> = distances.iter().map(|d| d * d).collect();
|
|
let total: f32 = squared.iter().sum();
|
|
|
|
if total == 0.0 {
|
|
break;
|
|
}
|
|
|
|
let target = rng.gen_range(0.0..total);
|
|
let mut cumsum = 0.0;
|
|
let mut selected = 0;
|
|
for (i, d) in squared.iter().enumerate() {
|
|
cumsum += d;
|
|
if cumsum >= target {
|
|
selected = i;
|
|
break;
|
|
}
|
|
}
|
|
centroids.push(training_vectors[selected].clone());
|
|
}
|
|
|
|
// K-means iterations
|
|
for _ in 0..self.config.kmeans_iterations {
|
|
let mut cluster_sums: Vec<Vec<f32>> = (0..n_clusters)
|
|
.map(|_| vec![0.0; self.dimensions])
|
|
.collect();
|
|
let mut cluster_counts: Vec<usize> = vec![0; n_clusters];
|
|
|
|
for vector in training_vectors {
|
|
let cluster = self.find_nearest_centroid(vector, ¢roids);
|
|
for (i, &v) in vector.iter().enumerate() {
|
|
cluster_sums[cluster][i] += v;
|
|
}
|
|
cluster_counts[cluster] += 1;
|
|
}
|
|
|
|
for (i, centroid) in centroids.iter_mut().enumerate() {
|
|
if cluster_counts[i] > 0 {
|
|
for j in 0..self.dimensions {
|
|
centroid[j] = cluster_sums[i][j] / cluster_counts[i] as f32;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
*self.centroids.write() = centroids;
|
|
|
|
for i in 0..n_clusters {
|
|
self.lists.insert(i, Vec::new());
|
|
}
|
|
|
|
self.trained
|
|
.store(true, std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
|
|
fn find_nearest_centroid(&self, vector: &[f32], centroids: &[Vec<f32>]) -> usize {
|
|
let mut best_cluster = 0;
|
|
let mut best_dist = f32::MAX;
|
|
|
|
for (i, centroid) in centroids.iter().enumerate() {
|
|
let dist = self.calc_distance(vector, centroid);
|
|
if dist < best_dist {
|
|
best_dist = dist;
|
|
best_cluster = i;
|
|
}
|
|
}
|
|
|
|
best_cluster
|
|
}
|
|
|
|
pub fn insert(&self, id: VectorId, vector: Vec<f32>) {
|
|
assert!(self.is_trained(), "Index must be trained before insertion");
|
|
|
|
let centroids = self.centroids.read();
|
|
let cluster = self.find_nearest_centroid(&vector, ¢roids);
|
|
drop(centroids);
|
|
|
|
let entry = ClusterEntry { id, vector };
|
|
|
|
if let Some(mut list) = self.lists.get_mut(&cluster) {
|
|
list.push(entry);
|
|
}
|
|
|
|
self.id_to_cluster.insert(id, cluster);
|
|
self.vector_count
|
|
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
|
|
pub fn search(
|
|
&self,
|
|
query: &[f32],
|
|
k: usize,
|
|
probes: Option<usize>,
|
|
) -> Vec<(VectorId, f32)> {
|
|
if !self.is_trained() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let n_probes = probes.unwrap_or(self.config.probes);
|
|
let centroids = self.centroids.read();
|
|
|
|
let mut centroid_dists: Vec<(usize, f32)> = centroids
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, c)| (i, self.calc_distance(query, c)))
|
|
.collect();
|
|
|
|
centroid_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
|
|
drop(centroids);
|
|
|
|
let mut heap = BinaryHeap::new();
|
|
|
|
for (cluster_id, _) in centroid_dists.iter().take(n_probes) {
|
|
if let Some(list) = self.lists.get(cluster_id) {
|
|
for entry in list.iter() {
|
|
let dist = self.calc_distance(query, &entry.vector);
|
|
heap.push(SearchResult {
|
|
id: entry.id,
|
|
distance: dist,
|
|
});
|
|
|
|
if heap.len() > k {
|
|
heap.pop();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut results: Vec<_> = heap.into_iter().map(|r| (r.id, r.distance)).collect();
|
|
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
|
|
results
|
|
}
|
|
|
|
pub fn search_parallel(
|
|
&self,
|
|
query: &[f32],
|
|
k: usize,
|
|
probes: Option<usize>,
|
|
) -> Vec<(VectorId, f32)> {
|
|
if !self.is_trained() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let n_probes = probes.unwrap_or(self.config.probes);
|
|
let centroids = self.centroids.read();
|
|
|
|
let mut centroid_dists: Vec<(usize, f32)> = centroids
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, c)| (i, self.calc_distance(query, c)))
|
|
.collect();
|
|
|
|
centroid_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
|
|
drop(centroids);
|
|
|
|
let probe_clusters: Vec<usize> = centroid_dists
|
|
.iter()
|
|
.take(n_probes)
|
|
.map(|(id, _)| *id)
|
|
.collect();
|
|
|
|
let results: Vec<(VectorId, f32)> = probe_clusters
|
|
.par_iter()
|
|
.flat_map(|cluster_id| {
|
|
let mut local_results = Vec::new();
|
|
if let Some(list) = self.lists.get(cluster_id) {
|
|
for entry in list.iter() {
|
|
let dist = self.calc_distance(query, &entry.vector);
|
|
local_results.push((entry.id, dist));
|
|
}
|
|
}
|
|
local_results
|
|
})
|
|
.collect();
|
|
|
|
let mut heap = BinaryHeap::new();
|
|
for (id, dist) in results {
|
|
heap.push(SearchResult { id, distance: dist });
|
|
if heap.len() > k {
|
|
heap.pop();
|
|
}
|
|
}
|
|
|
|
let mut final_results: Vec<_> = heap.into_iter().map(|r| (r.id, r.distance)).collect();
|
|
final_results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
|
|
final_results
|
|
}
|
|
|
|
pub fn memory_usage(&self) -> usize {
|
|
let vector_bytes = self.len() * self.dimensions * 4;
|
|
let centroid_bytes = self.config.lists * self.dimensions * 4;
|
|
vector_bytes + centroid_bytes
|
|
}
|
|
}
|
|
}
|
|
|
|
use hnsw::{DistanceMetric as HnswMetric, HnswConfig, HnswIndex};
|
|
use ivfflat::{DistanceMetric as IvfMetric, IvfFlatConfig, IvfFlatIndex};
|
|
|
|
// ============================================================================
|
|
// Test Data Generation
|
|
// ============================================================================
|
|
|
|
fn generate_random_vectors(n: usize, dims: usize, seed: u64) -> Vec<Vec<f32>> {
|
|
let mut rng = ChaCha8Rng::seed_from_u64(seed);
|
|
(0..n)
|
|
.map(|_| (0..dims).map(|_| rng.gen_range(-1.0..1.0)).collect())
|
|
.collect()
|
|
}
|
|
|
|
fn generate_clustered_vectors(
|
|
n: usize,
|
|
dims: usize,
|
|
num_clusters: usize,
|
|
seed: u64,
|
|
) -> Vec<Vec<f32>> {
|
|
let mut rng = ChaCha8Rng::seed_from_u64(seed);
|
|
|
|
let centers: Vec<Vec<f32>> = (0..num_clusters)
|
|
.map(|_| (0..dims).map(|_| rng.gen_range(-1.0..1.0)).collect())
|
|
.collect();
|
|
|
|
(0..n)
|
|
.map(|_| {
|
|
let center = ¢ers[rng.gen_range(0..num_clusters)];
|
|
center
|
|
.iter()
|
|
.map(|&c| c + rng.gen_range(-0.1..0.1))
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
// ============================================================================
|
|
// HNSW Build Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_hnsw_build(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Build");
|
|
group.sample_size(10);
|
|
|
|
for &dims in [128, 384, 768, 1536].iter() {
|
|
for &n in [1_000, 10_000, 100_000].iter() {
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
|
|
group.throughput(Throughput::Elements(n as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("{}d", dims), n),
|
|
&vectors,
|
|
|bench, vecs| {
|
|
bench.iter(|| {
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
..Default::default()
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vecs.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
black_box(index)
|
|
});
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_hnsw_build_ef_construction(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Build (ef_construction)");
|
|
group.sample_size(10);
|
|
|
|
let dims = 768;
|
|
let n = 10_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
|
|
for &ef in [16, 32, 64, 128, 256].iter() {
|
|
group.bench_with_input(BenchmarkId::from_parameter(ef), &ef, |bench, &ef_val| {
|
|
bench.iter(|| {
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: ef_val,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
..Default::default()
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
black_box(index)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_hnsw_build_m_parameter(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Build (M parameter)");
|
|
group.sample_size(10);
|
|
|
|
let dims = 768;
|
|
let n = 10_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
|
|
for &m in [8, 12, 16, 24, 32, 48].iter() {
|
|
group.bench_with_input(BenchmarkId::from_parameter(m), &m, |bench, &m_val| {
|
|
bench.iter(|| {
|
|
let config = HnswConfig {
|
|
m: m_val,
|
|
m0: m_val * 2,
|
|
ef_construction: 64,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
..Default::default()
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
black_box(index)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// HNSW Search Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_hnsw_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Search");
|
|
|
|
for &dims in [128, 384, 768, 1536].iter() {
|
|
for &n in [10_000, 100_000].iter() {
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let query = generate_random_vectors(1, dims, 999)[0].clone();
|
|
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("{}d", dims), n),
|
|
&(&index, &query),
|
|
|bench, (idx, q)| {
|
|
bench.iter(|| black_box(idx.search(q, 10)));
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_hnsw_search_ef_values(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Search (ef_search)");
|
|
|
|
let dims = 768;
|
|
let n = 100_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let queries = generate_random_vectors(100, dims, 999);
|
|
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
for &ef in [10, 20, 40, 80, 160, 320].iter() {
|
|
group.bench_with_input(BenchmarkId::from_parameter(ef), &ef, |bench, &ef_val| {
|
|
bench.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search_with_ef(query, 10, ef_val));
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_hnsw_search_k_values(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Search (k values)");
|
|
|
|
let dims = 768;
|
|
let n = 100_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let query = generate_random_vectors(1, dims, 999)[0].clone();
|
|
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 100,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
for &k in [1, 5, 10, 20, 50, 100].iter() {
|
|
group.bench_with_input(BenchmarkId::from_parameter(k), &k, |bench, &k_val| {
|
|
bench.iter(|| black_box(index.search(&query, k_val)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// IVFFlat Build Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_ivfflat_build(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("IVFFlat Build");
|
|
group.sample_size(10);
|
|
|
|
for &dims in [128, 384, 768].iter() {
|
|
for &n in [1_000, 10_000, 100_000].iter() {
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
|
|
group.throughput(Throughput::Elements(n as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("{}d", dims), n),
|
|
&vectors,
|
|
|bench, vecs| {
|
|
bench.iter(|| {
|
|
let n_lists = (n as f64).sqrt() as usize;
|
|
let config = IvfFlatConfig {
|
|
lists: n_lists,
|
|
probes: 1,
|
|
metric: IvfMetric::Euclidean,
|
|
kmeans_iterations: 10,
|
|
seed: 42,
|
|
};
|
|
|
|
let index = IvfFlatIndex::new(dims, config);
|
|
index.train(vecs);
|
|
|
|
for (id, vec) in vecs.iter().enumerate() {
|
|
index.insert(id as u64, vec.clone());
|
|
}
|
|
black_box(index)
|
|
});
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_ivfflat_build_lists(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("IVFFlat Build (nlist)");
|
|
group.sample_size(10);
|
|
|
|
let dims = 768;
|
|
let n = 10_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
|
|
for &n_lists in [10, 50, 100, 200, 500].iter() {
|
|
group.bench_with_input(
|
|
BenchmarkId::from_parameter(n_lists),
|
|
&n_lists,
|
|
|bench, &lists| {
|
|
bench.iter(|| {
|
|
let config = IvfFlatConfig {
|
|
lists,
|
|
probes: 1,
|
|
metric: IvfMetric::Euclidean,
|
|
kmeans_iterations: 10,
|
|
seed: 42,
|
|
};
|
|
|
|
let index = IvfFlatIndex::new(dims, config);
|
|
index.train(&vectors);
|
|
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec.clone());
|
|
}
|
|
black_box(index)
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// IVFFlat Search Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_ivfflat_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("IVFFlat Search");
|
|
|
|
for &dims in [128, 384, 768].iter() {
|
|
for &n in [10_000, 100_000].iter() {
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let query = generate_random_vectors(1, dims, 999)[0].clone();
|
|
|
|
let n_lists = (n as f64).sqrt() as usize;
|
|
let config = IvfFlatConfig {
|
|
lists: n_lists,
|
|
probes: 5,
|
|
metric: IvfMetric::Euclidean,
|
|
kmeans_iterations: 10,
|
|
seed: 42,
|
|
};
|
|
|
|
let index = IvfFlatIndex::new(dims, config);
|
|
index.train(&vectors);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec.clone());
|
|
}
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("{}d", dims), n),
|
|
&(&index, &query),
|
|
|bench, (idx, q)| {
|
|
bench.iter(|| black_box(idx.search(q, 10, None)));
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_ivfflat_search_probes(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("IVFFlat Search (nprobe)");
|
|
|
|
let dims = 768;
|
|
let n = 100_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let queries = generate_random_vectors(100, dims, 999);
|
|
|
|
let n_lists = (n as f64).sqrt() as usize;
|
|
let config = IvfFlatConfig {
|
|
lists: n_lists,
|
|
probes: 1,
|
|
metric: IvfMetric::Euclidean,
|
|
kmeans_iterations: 10,
|
|
seed: 42,
|
|
};
|
|
|
|
let index = IvfFlatIndex::new(dims, config);
|
|
index.train(&vectors);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec.clone());
|
|
}
|
|
|
|
for &probes in [1, 5, 10, 20, 50, 100].iter() {
|
|
group.bench_with_input(
|
|
BenchmarkId::from_parameter(probes),
|
|
&probes,
|
|
|bench, &probe_val| {
|
|
bench.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search(query, 10, Some(probe_val)));
|
|
}
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_ivfflat_parallel_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("IVFFlat Parallel Search");
|
|
|
|
let dims = 768;
|
|
let n = 100_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let queries = generate_random_vectors(100, dims, 999);
|
|
|
|
let n_lists = (n as f64).sqrt() as usize;
|
|
let config = IvfFlatConfig {
|
|
lists: n_lists,
|
|
probes: 10,
|
|
metric: IvfMetric::Euclidean,
|
|
kmeans_iterations: 10,
|
|
seed: 42,
|
|
};
|
|
|
|
let index = IvfFlatIndex::new(dims, config);
|
|
index.train(&vectors);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec.clone());
|
|
}
|
|
|
|
group.bench_function("sequential", |bench| {
|
|
bench.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search(query, 10, None));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.bench_function("parallel_probes", |bench| {
|
|
bench.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search_parallel(query, 10, None));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.bench_function("parallel_queries", |bench| {
|
|
bench.iter(|| {
|
|
queries.par_iter().for_each(|query| {
|
|
black_box(index.search(query, 10, None));
|
|
});
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Recall Analysis Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_hnsw_recall(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Recall Analysis");
|
|
group.sample_size(10);
|
|
|
|
let dims = 768;
|
|
let n = 10_000;
|
|
let vectors = generate_clustered_vectors(n, dims, 20, 42);
|
|
let queries = generate_random_vectors(100, dims, 999);
|
|
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
// Compute ground truth (brute force)
|
|
let compute_ground_truth = |query: &[f32], k: usize| -> Vec<u64> {
|
|
let mut distances: Vec<(u64, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(id, vec)| {
|
|
let dist = vec
|
|
.iter()
|
|
.zip(query)
|
|
.map(|(a, b)| (a - b).powi(2))
|
|
.sum::<f32>()
|
|
.sqrt();
|
|
(id as u64, dist)
|
|
})
|
|
.collect();
|
|
|
|
distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
distances.iter().take(k).map(|(id, _)| *id).collect()
|
|
};
|
|
|
|
for &ef in [10, 20, 40, 80, 160].iter() {
|
|
group.bench_with_input(BenchmarkId::new("recall@10", ef), &ef, |bench, &ef_val| {
|
|
bench.iter(|| {
|
|
let mut total_recall = 0.0;
|
|
for query in &queries {
|
|
let ground_truth = compute_ground_truth(query, 10);
|
|
let results = index.search_with_ef(query, 10, ef_val);
|
|
|
|
let hits = results
|
|
.iter()
|
|
.filter(|r| ground_truth.contains(&r.id))
|
|
.count();
|
|
|
|
total_recall += hits as f32 / 10.0;
|
|
}
|
|
black_box(total_recall / queries.len() as f32)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Memory Usage Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_memory_usage(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("Index Memory Usage");
|
|
group.sample_size(10);
|
|
|
|
for &dims in [128, 384, 768, 1536].iter() {
|
|
for &n in [1_000, 10_000, 100_000].iter() {
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("hnsw_{}d", dims), n),
|
|
&vectors,
|
|
|bench, vecs| {
|
|
bench.iter(|| {
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
..Default::default()
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vecs.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
let memory_bytes = index.memory_usage();
|
|
let memory_per_vec = memory_bytes as f64 / n as f64;
|
|
black_box(memory_per_vec)
|
|
});
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Distance Metric Comparison
|
|
// ============================================================================
|
|
|
|
fn bench_hnsw_distance_metrics(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Distance Metrics");
|
|
group.sample_size(10);
|
|
|
|
let dims = 768;
|
|
let n = 10_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let query = generate_random_vectors(1, dims, 999)[0].clone();
|
|
|
|
for metric in [
|
|
HnswMetric::Euclidean,
|
|
HnswMetric::Cosine,
|
|
HnswMetric::InnerProduct,
|
|
] {
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: n,
|
|
metric,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
let metric_name = match metric {
|
|
HnswMetric::Euclidean => "l2",
|
|
HnswMetric::Cosine => "cosine",
|
|
HnswMetric::InnerProduct => "inner_product",
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("search", metric_name),
|
|
&(&index, &query),
|
|
|bench, (idx, q)| {
|
|
bench.iter(|| black_box(idx.search(q, 10)));
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Parallel Search Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_hnsw_parallel_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("HNSW Parallel Query");
|
|
|
|
let dims = 768;
|
|
let n = 100_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let queries = generate_random_vectors(1000, dims, 999);
|
|
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
group.bench_function("sequential", |bench| {
|
|
bench.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search(query, 10));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.bench_function("parallel_rayon", |bench| {
|
|
bench.iter(|| {
|
|
queries.par_iter().for_each(|query| {
|
|
black_box(index.search(query, 10));
|
|
});
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Latency Percentile Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_latency_percentiles(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("Query Latency Percentiles");
|
|
group.sample_size(10);
|
|
|
|
let dims = 768;
|
|
let n = 100_000;
|
|
let vectors = generate_random_vectors(n, dims, 42);
|
|
let queries = generate_random_vectors(1000, dims, 999);
|
|
|
|
let config = HnswConfig {
|
|
m: 16,
|
|
m0: 32,
|
|
ef_construction: 64,
|
|
ef_search: 40,
|
|
max_elements: n,
|
|
metric: HnswMetric::Euclidean,
|
|
seed: 42,
|
|
};
|
|
|
|
let mut index = HnswIndex::new(config);
|
|
for (id, vec) in vectors.iter().enumerate() {
|
|
index.insert(id as u64, vec);
|
|
}
|
|
|
|
group.bench_function("hnsw_latency_distribution", |bench| {
|
|
bench.iter(|| {
|
|
let mut latencies: Vec<Duration> = Vec::with_capacity(queries.len());
|
|
for query in &queries {
|
|
let start = Instant::now();
|
|
black_box(index.search(query, 10));
|
|
latencies.push(start.elapsed());
|
|
}
|
|
|
|
latencies.sort();
|
|
let p50 = latencies[latencies.len() / 2];
|
|
let p95 = latencies[(latencies.len() as f64 * 0.95) as usize];
|
|
let p99 = latencies[(latencies.len() as f64 * 0.99) as usize];
|
|
|
|
black_box((p50, p95, p99))
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
// HNSW Build
|
|
bench_hnsw_build,
|
|
bench_hnsw_build_ef_construction,
|
|
bench_hnsw_build_m_parameter,
|
|
// HNSW Search
|
|
bench_hnsw_search,
|
|
bench_hnsw_search_ef_values,
|
|
bench_hnsw_search_k_values,
|
|
// IVFFlat Build
|
|
bench_ivfflat_build,
|
|
bench_ivfflat_build_lists,
|
|
// IVFFlat Search
|
|
bench_ivfflat_search,
|
|
bench_ivfflat_search_probes,
|
|
bench_ivfflat_parallel_search,
|
|
// Recall Analysis
|
|
bench_hnsw_recall,
|
|
// Memory Usage
|
|
bench_memory_usage,
|
|
// Distance Metrics
|
|
bench_hnsw_distance_metrics,
|
|
// Parallel Search
|
|
bench_hnsw_parallel_search,
|
|
// Latency Percentiles
|
|
bench_latency_percentiles,
|
|
);
|
|
|
|
criterion_main!(benches);
|