506 lines
16 KiB
Rust
506 lines
16 KiB
Rust
//! HNSW Benchmark Suite for 7sense
|
|
//!
|
|
//! Performance targets from ADR-004:
|
|
//! - HNSW Search: 150x speedup vs brute force
|
|
//! - Query Latency p99: < 50ms
|
|
//! - Recall@10: >= 0.95
|
|
//! - Recall@100: >= 0.98
|
|
//! - Insert Throughput: >= 10,000 vectors/s
|
|
//! - Build Time: < 30 min for 1M vectors
|
|
|
|
use criterion::{
|
|
black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
|
|
};
|
|
use std::time::{Duration, Instant};
|
|
|
|
mod utils;
|
|
use utils::*;
|
|
|
|
/// Index sizes to benchmark
|
|
const SMALL_INDEX: usize = 10_000;
|
|
const MEDIUM_INDEX: usize = 100_000;
|
|
const LARGE_INDEX: usize = 500_000;
|
|
|
|
/// K values for search benchmarks
|
|
const K_VALUES: &[usize] = &[10, 50, 100];
|
|
|
|
// ============================================================================
|
|
// HNSW Search Benchmarks
|
|
// ============================================================================
|
|
|
|
/// Benchmark HNSW search performance with different index sizes and k values
|
|
fn benchmark_hnsw_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_search");
|
|
group.sample_size(50);
|
|
group.measurement_time(Duration::from_secs(10));
|
|
|
|
// Generate query vectors once
|
|
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
|
|
|
for &size in &[SMALL_INDEX, MEDIUM_INDEX] {
|
|
// Build index
|
|
println!("Building index with {} vectors...", size);
|
|
let index = setup_test_index(size);
|
|
|
|
for &k in K_VALUES {
|
|
group.throughput(Throughput::Elements(queries.len() as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("size_{}_k_{}", size, k), k),
|
|
&k,
|
|
|b, &k| {
|
|
b.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search(query, k));
|
|
}
|
|
});
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark HNSW search with different ef_search values
|
|
fn benchmark_hnsw_search_ef(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_search_ef");
|
|
group.sample_size(30);
|
|
group.measurement_time(Duration::from_secs(8));
|
|
|
|
let size = MEDIUM_INDEX;
|
|
let mut index = setup_test_index(size);
|
|
let queries = generate_random_vectors(50, PERCH_EMBEDDING_DIM);
|
|
let k = 10;
|
|
|
|
for ef in [64, 128, 256, 512] {
|
|
index.set_ef_search(ef);
|
|
|
|
group.throughput(Throughput::Elements(queries.len() as u64));
|
|
group.bench_with_input(BenchmarkId::new("ef", ef), &ef, |b, _| {
|
|
b.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search(query, k));
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark HNSW vs brute force to calculate speedup ratio
|
|
fn benchmark_hnsw_vs_brute_force(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_vs_brute_force");
|
|
group.sample_size(20);
|
|
group.measurement_time(Duration::from_secs(15));
|
|
|
|
// Use smaller index for brute force comparison
|
|
let size = 10_000;
|
|
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
|
let mut index = SimpleHnswIndex::new_default();
|
|
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
|
|
let queries = generate_random_vectors(20, PERCH_EMBEDDING_DIM);
|
|
let k = 10;
|
|
|
|
// Benchmark brute force
|
|
group.bench_function("brute_force", |b| {
|
|
b.iter(|| {
|
|
for query in &queries {
|
|
black_box(brute_force_knn(query, &vectors, k));
|
|
}
|
|
});
|
|
});
|
|
|
|
// Benchmark HNSW
|
|
group.bench_function("hnsw", |b| {
|
|
b.iter(|| {
|
|
for query in &queries {
|
|
black_box(index.search(query, k));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// HNSW Insert Benchmarks
|
|
// ============================================================================
|
|
|
|
/// Benchmark single vector insertion
|
|
fn benchmark_hnsw_insert_single(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_insert_single");
|
|
group.sample_size(50);
|
|
group.measurement_time(Duration::from_secs(10));
|
|
|
|
// Benchmark insertion into indices of different sizes
|
|
for &initial_size in &[1000, 10_000, 50_000] {
|
|
let vectors_to_insert = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("initial_size", initial_size),
|
|
&initial_size,
|
|
|b, &size| {
|
|
b.iter_batched(
|
|
|| {
|
|
// Setup: create index with initial vectors
|
|
setup_test_index(size)
|
|
},
|
|
|mut index| {
|
|
// Insert new vectors
|
|
for vec in &vectors_to_insert {
|
|
black_box(index.add(vec.clone()));
|
|
}
|
|
},
|
|
criterion::BatchSize::SmallInput,
|
|
);
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark batch vector insertion
|
|
fn benchmark_hnsw_insert_batch(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_insert_batch");
|
|
group.sample_size(20);
|
|
group.measurement_time(Duration::from_secs(15));
|
|
|
|
for &batch_size in &[100, 1000, 5000] {
|
|
let vectors = generate_random_vectors(batch_size, PERCH_EMBEDDING_DIM);
|
|
|
|
group.throughput(Throughput::Elements(batch_size as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("batch_size", batch_size),
|
|
&batch_size,
|
|
|b, _| {
|
|
b.iter_batched(
|
|
|| {
|
|
// Setup: create empty index
|
|
SimpleHnswIndex::new_default()
|
|
},
|
|
|mut index| {
|
|
// Insert batch
|
|
black_box(index.batch_add(vectors.clone()));
|
|
},
|
|
criterion::BatchSize::SmallInput,
|
|
);
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// HNSW Build Benchmarks
|
|
// ============================================================================
|
|
|
|
/// Benchmark index construction time
|
|
fn benchmark_hnsw_build(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_build");
|
|
group.sample_size(10);
|
|
group.measurement_time(Duration::from_secs(30));
|
|
|
|
for &size in &[1000, 5000, 10_000] {
|
|
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
|
|
|
group.throughput(Throughput::Elements(size as u64));
|
|
group.bench_with_input(BenchmarkId::new("vectors", size), &size, |b, _| {
|
|
b.iter(|| {
|
|
let mut index = SimpleHnswIndex::new_default();
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
black_box(index)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark index construction with different M parameters
|
|
fn benchmark_hnsw_build_m_param(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_build_m_param");
|
|
group.sample_size(10);
|
|
group.measurement_time(Duration::from_secs(20));
|
|
|
|
let size = 5000;
|
|
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
|
|
|
for m in [16, 24, 32, 48] {
|
|
group.bench_with_input(BenchmarkId::new("M", m), &m, |b, &m| {
|
|
b.iter(|| {
|
|
let mut index =
|
|
SimpleHnswIndex::new(PERCH_EMBEDDING_DIM, m, DEFAULT_EF_CONSTRUCTION, DEFAULT_EF_SEARCH);
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
black_box(index)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Recall Measurement
|
|
// ============================================================================
|
|
|
|
/// Measure and report recall metrics (not a benchmark, but a validation)
|
|
fn measure_recall(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("recall_measurement");
|
|
group.sample_size(10);
|
|
|
|
let size = 10_000;
|
|
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
|
let mut index = SimpleHnswIndex::new_default();
|
|
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
|
|
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
|
|
|
// This benchmark measures time to compute recall (including brute force)
|
|
group.bench_function("recall_computation", |b| {
|
|
b.iter(|| {
|
|
let mut total_recall_10 = 0.0;
|
|
let mut total_recall_100 = 0.0;
|
|
|
|
for query in &queries {
|
|
let hnsw_results = index.search(query, 100);
|
|
let ground_truth = brute_force_knn(query, &vectors, 100);
|
|
|
|
total_recall_10 += measure_recall_at_k(&hnsw_results, &ground_truth, 10);
|
|
total_recall_100 += measure_recall_at_k(&hnsw_results, &ground_truth, 100);
|
|
}
|
|
|
|
let avg_recall_10 = total_recall_10 / queries.len() as f32;
|
|
let avg_recall_100 = total_recall_100 / queries.len() as f32;
|
|
|
|
black_box((avg_recall_10, avg_recall_100))
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Speedup Ratio Calculation
|
|
// ============================================================================
|
|
|
|
/// Calculate and report the speedup ratio of HNSW vs brute force
|
|
/// This is run as a single iteration with detailed output
|
|
fn calculate_speedup_ratio() {
|
|
println!("\n=== HNSW vs Brute Force Speedup Analysis ===\n");
|
|
|
|
for &size in &[1_000, 5_000, 10_000, 50_000] {
|
|
println!("Index size: {} vectors", size);
|
|
println!("Dimension: {}", PERCH_EMBEDDING_DIM);
|
|
|
|
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
|
let mut index = SimpleHnswIndex::new_default();
|
|
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
|
|
let queries = generate_random_vectors(100, PERCH_EMBEDDING_DIM);
|
|
let k = 10;
|
|
|
|
// Time brute force
|
|
let bf_start = Instant::now();
|
|
for query in &queries {
|
|
let _ = brute_force_knn(query, &vectors, k);
|
|
}
|
|
let bf_time = bf_start.elapsed();
|
|
|
|
// Time HNSW
|
|
let hnsw_start = Instant::now();
|
|
for query in &queries {
|
|
let _ = index.search(query, k);
|
|
}
|
|
let hnsw_time = hnsw_start.elapsed();
|
|
|
|
let speedup = bf_time.as_secs_f64() / hnsw_time.as_secs_f64();
|
|
|
|
// Calculate recall
|
|
let mut total_recall = 0.0;
|
|
for query in &queries {
|
|
let hnsw_results = index.search(query, k);
|
|
let ground_truth = brute_force_knn(query, &vectors, k);
|
|
total_recall += measure_recall_at_k(&hnsw_results, &ground_truth, k);
|
|
}
|
|
let avg_recall = total_recall / queries.len() as f32;
|
|
|
|
println!(" Brute Force: {:?} ({} queries)", bf_time, queries.len());
|
|
println!(" HNSW: {:?} ({} queries)", hnsw_time, queries.len());
|
|
println!(" Speedup: {:.1}x", speedup);
|
|
println!(" Recall@{}: {:.3}", k, avg_recall);
|
|
println!(
|
|
" Target: {}x speedup ({})",
|
|
targets::HNSW_SPEEDUP_VS_BRUTE_FORCE,
|
|
if speedup >= targets::HNSW_SPEEDUP_VS_BRUTE_FORCE {
|
|
"PASS"
|
|
} else {
|
|
"FAIL"
|
|
}
|
|
);
|
|
println!();
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Latency Distribution Analysis
|
|
// ============================================================================
|
|
|
|
/// Analyze query latency distribution
|
|
fn analyze_latency_distribution() {
|
|
println!("\n=== Query Latency Distribution Analysis ===\n");
|
|
|
|
let size = MEDIUM_INDEX;
|
|
println!("Building index with {} vectors...", size);
|
|
let index = setup_test_index(size);
|
|
|
|
let queries = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
|
let k = 10;
|
|
|
|
let mut latencies = Vec::with_capacity(queries.len());
|
|
|
|
for query in &queries {
|
|
let start = Instant::now();
|
|
let _ = index.search(query, k);
|
|
latencies.push(start.elapsed());
|
|
}
|
|
|
|
let stats = PerformanceStats::from_latencies(latencies);
|
|
|
|
println!("Query latency statistics (k={}, {} queries):", k, queries.len());
|
|
println!("{}", stats.report());
|
|
println!();
|
|
println!("Performance targets:");
|
|
println!(
|
|
" p50 target: {}ms ({})",
|
|
targets::QUERY_LATENCY_P50_MS,
|
|
if stats.p50 <= Duration::from_millis(targets::QUERY_LATENCY_P50_MS) {
|
|
"PASS"
|
|
} else {
|
|
"FAIL"
|
|
}
|
|
);
|
|
println!(
|
|
" p99 target: {}ms ({})",
|
|
targets::QUERY_LATENCY_P99_MS,
|
|
if stats.p99 <= Duration::from_millis(targets::QUERY_LATENCY_P99_MS) {
|
|
"PASS"
|
|
} else {
|
|
"FAIL"
|
|
}
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Criterion Groups
|
|
// ============================================================================
|
|
|
|
criterion_group!(
|
|
name = search_benches;
|
|
config = Criterion::default().with_output_color(true);
|
|
targets = benchmark_hnsw_search, benchmark_hnsw_search_ef, benchmark_hnsw_vs_brute_force
|
|
);
|
|
|
|
criterion_group!(
|
|
name = insert_benches;
|
|
config = Criterion::default().with_output_color(true);
|
|
targets = benchmark_hnsw_insert_single, benchmark_hnsw_insert_batch
|
|
);
|
|
|
|
criterion_group!(
|
|
name = build_benches;
|
|
config = Criterion::default().with_output_color(true);
|
|
targets = benchmark_hnsw_build, benchmark_hnsw_build_m_param
|
|
);
|
|
|
|
criterion_group!(
|
|
name = recall_benches;
|
|
config = Criterion::default().with_output_color(true);
|
|
targets = measure_recall
|
|
);
|
|
|
|
criterion_main!(search_benches, insert_benches, build_benches, recall_benches);
|
|
|
|
// ============================================================================
|
|
// Additional Analysis Functions (run separately)
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod analysis {
|
|
use super::*;
|
|
|
|
#[test]
|
|
#[ignore] // Run with: cargo test --release -- --ignored --nocapture
|
|
fn run_speedup_analysis() {
|
|
calculate_speedup_ratio();
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn run_latency_analysis() {
|
|
analyze_latency_distribution();
|
|
}
|
|
|
|
#[test]
|
|
fn test_target_recall_at_10() {
|
|
let size = 5_000;
|
|
let vectors = generate_random_vectors(size, PERCH_EMBEDDING_DIM);
|
|
let mut index = SimpleHnswIndex::new_default();
|
|
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
|
|
let queries = generate_random_vectors(50, PERCH_EMBEDDING_DIM);
|
|
|
|
let mut total_recall = 0.0;
|
|
for query in &queries {
|
|
let hnsw_results = index.search(query, 10);
|
|
let ground_truth = brute_force_knn(query, &vectors, 10);
|
|
total_recall += measure_recall_at_k(&hnsw_results, &ground_truth, 10);
|
|
}
|
|
|
|
let avg_recall = total_recall / queries.len() as f32;
|
|
println!("Average Recall@10: {:.3}", avg_recall);
|
|
assert!(
|
|
avg_recall as f64 >= targets::RECALL_AT_10,
|
|
"Recall@10 {} below target {}",
|
|
avg_recall,
|
|
targets::RECALL_AT_10
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_insert_throughput() {
|
|
let vectors = generate_random_vectors(1000, PERCH_EMBEDDING_DIM);
|
|
let mut index = SimpleHnswIndex::new_default();
|
|
|
|
let start = Instant::now();
|
|
for vec in &vectors {
|
|
index.add(vec.clone());
|
|
}
|
|
let elapsed = start.elapsed();
|
|
|
|
let throughput = vectors.len() as f64 / elapsed.as_secs_f64();
|
|
println!("Insert throughput: {:.0} vectors/sec", throughput);
|
|
|
|
// Note: This is a simplified index, real HNSW should achieve higher throughput
|
|
}
|
|
}
|