use crate::{Vector, VectorIndex};
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkConfig {
pub warmup_runs: usize,
pub benchmark_runs: usize,
pub max_duration: Duration,
pub profile_memory: bool,
pub detailed_timing: bool,
pub quality_metrics: bool,
pub random_seed: Option<u64>,
pub output_format: BenchmarkOutputFormat,
}
impl Default for BenchmarkConfig {
fn default() -> Self {
Self {
warmup_runs: 3,
benchmark_runs: 10,
max_duration: Duration::from_secs(300), profile_memory: true,
detailed_timing: true,
quality_metrics: true,
random_seed: Some(42),
output_format: BenchmarkOutputFormat::Json,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BenchmarkOutputFormat {
Json,
Csv,
Table,
AnnBenchmarks,
}
pub struct BenchmarkSuite {
config: BenchmarkConfig,
datasets: Vec<BenchmarkDataset>,
algorithms: Vec<Box<dyn VectorIndex>>,
results: Vec<BenchmarkResult>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkDataset {
pub name: String,
pub description: String,
pub train_vectors: Vec<Vector>,
pub query_vectors: Vec<Vector>,
pub ground_truth: Option<Vec<Vec<usize>>>, pub metadata: HashMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkTestCase {
pub name: String,
pub description: String,
pub dataset: String,
pub algorithm: String,
pub parameters: HashMap<String, serde_json::Value>,
pub query_count: usize,
pub k: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResult {
pub test_case: BenchmarkTestCase,
pub performance: PerformanceMetrics,
pub quality: Option<QualityMetrics>,
pub memory: Option<MemoryMetrics>,
pub scalability: Option<ScalabilityMetrics>,
pub system_info: SystemInfo,
pub timestamp: std::time::SystemTime,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PerformanceMetrics {
pub avg_query_time: Duration,
pub median_query_time: Duration,
pub p95_query_time: Duration,
pub p99_query_time: Duration,
pub min_query_time: Duration,
pub max_query_time: Duration,
pub std_dev_query_time: Duration,
pub queries_per_second: f64,
pub index_build_time: Option<Duration>,
pub index_update_time: Option<Duration>,
pub throughput: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QualityMetrics {
pub recall_at_k: f64,
pub precision_at_k: f64,
pub mean_average_precision: f64,
pub ndcg_at_k: f64,
pub distance_ratio: Option<f64>,
pub relative_error: Option<f64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryMetrics {
pub peak_memory_bytes: usize,
pub avg_memory_bytes: usize,
pub memory_per_vector: f64,
pub index_overhead_bytes: usize,
pub memory_efficiency: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScalabilityMetrics {
pub performance_scaling: Vec<(usize, Duration)>, pub memory_scaling: Vec<(usize, usize)>, pub build_time_scaling: Vec<(usize, Duration)>, pub concurrency_scaling: Vec<(usize, f64)>, }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemInfo {
pub cpu_info: String,
pub total_ram_gb: f64,
pub available_ram_gb: f64,
pub os: String,
pub rust_version: String,
pub simd_features: Vec<String>,
pub gpu_info: Option<String>,
}
pub struct PerformanceProfiler {
start_time: Instant,
checkpoints: Vec<(String, Instant)>,
memory_samples: Vec<(Instant, usize)>,
}
impl PerformanceProfiler {
pub fn new() -> Self {
Self {
start_time: Instant::now(),
checkpoints: Vec::new(),
memory_samples: Vec::new(),
}
}
pub fn checkpoint(&mut self, name: &str) {
self.checkpoints.push((name.to_string(), Instant::now()));
}
pub fn sample_memory(&mut self) {
let memory_usage = self.get_current_memory_usage();
self.memory_samples.push((Instant::now(), memory_usage));
}
pub fn get_timing_breakdown(&self) -> Vec<(String, Duration)> {
let mut breakdown = Vec::new();
let mut last_time = self.start_time;
for (name, time) in &self.checkpoints {
breakdown.push((name.clone(), time.duration_since(last_time)));
last_time = *time;
}
breakdown
}
pub fn get_memory_profile(&self) -> Vec<(Duration, usize)> {
self.memory_samples
.iter()
.map(|(time, memory)| (time.duration_since(self.start_time), *memory))
.collect()
}
fn get_current_memory_usage(&self) -> usize {
#[cfg(target_os = "linux")]
{
if let Ok(content) = std::fs::read_to_string("/proc/self/status") {
for line in content.lines() {
if line.starts_with("VmRSS:") {
if let Some(kb_str) = line.split_whitespace().nth(1) {
if let Ok(kb) = kb_str.parse::<usize>() {
return kb * 1024; }
}
}
}
}
}
#[cfg(target_os = "macos")]
{
use std::process::Command;
if let Ok(output) = Command::new("ps")
.args(["-o", "rss=", "-p"])
.arg(std::process::id().to_string())
.output()
{
if let Ok(rss_str) = String::from_utf8(output.stdout) {
if let Ok(rss_kb) = rss_str.trim().parse::<usize>() {
return rss_kb * 1024; }
}
}
}
0
}
}
impl Default for PerformanceProfiler {
fn default() -> Self {
Self::new()
}
}
impl BenchmarkSuite {
pub fn new(config: BenchmarkConfig) -> Self {
Self {
config,
datasets: Vec::new(),
algorithms: Vec::new(),
results: Vec::new(),
}
}
pub fn add_dataset(&mut self, dataset: BenchmarkDataset) {
self.datasets.push(dataset);
}
pub fn add_algorithm(&mut self, algorithm: Box<dyn VectorIndex>) {
self.algorithms.push(algorithm);
}
pub fn generate_synthetic_datasets(&mut self) -> Result<()> {
self.generate_random_dataset("random_1000", 1000, 128, 100)?;
self.generate_random_dataset("random_10000", 10000, 256, 1000)?;
self.generate_clustered_dataset("clustered_5000", 5000, 384, 500, 10)?;
self.generate_uniform_dataset("uniform_2000", 2000, 512, 200)?;
Ok(())
}
fn generate_random_dataset(
&mut self,
name: &str,
size: usize,
dimensions: usize,
query_count: usize,
) -> Result<()> {
let mut train_vectors = Vec::new();
let mut query_vectors = Vec::new();
for i in 0..size {
let vector = crate::utils::random_vector(dimensions, Some(i as u64));
train_vectors.push(vector);
}
for i in 0..query_count {
let vector = crate::utils::random_vector(dimensions, Some((size + i) as u64));
query_vectors.push(vector);
}
let dataset = BenchmarkDataset {
name: name.to_string(),
description: format!("Random dataset with {size} vectors of {dimensions} dimensions"),
train_vectors,
query_vectors,
ground_truth: None,
metadata: {
let mut meta = HashMap::new();
meta.insert("type".to_string(), "random".to_string());
meta.insert("size".to_string(), size.to_string());
meta.insert("dimensions".to_string(), dimensions.to_string());
meta
},
};
self.add_dataset(dataset);
Ok(())
}
fn generate_clustered_dataset(
&mut self,
name: &str,
size: usize,
dimensions: usize,
query_count: usize,
num_clusters: usize,
) -> Result<()> {
let mut train_vectors = Vec::new();
let mut query_vectors = Vec::new();
let mut cluster_centers = Vec::new();
for i in 0..num_clusters {
let center = crate::utils::random_vector(dimensions, Some(i as u64));
cluster_centers.push(center);
}
for i in 0..size {
let cluster_idx = i % num_clusters;
let center = &cluster_centers[cluster_idx];
let center_f32 = center.as_f32();
let mut noisy_vector = center_f32.clone();
let noise_scale = 0.1;
for val in &mut noisy_vector {
let noise = self.gaussian_random(0.0, noise_scale, i as u64);
*val += noise;
}
train_vectors.push(Vector::new(noisy_vector));
}
for i in 0..query_count {
let cluster_idx = i % num_clusters;
let center = &cluster_centers[cluster_idx];
let center_f32 = center.as_f32();
let mut noisy_vector = center_f32.clone();
let noise_scale = 0.05;
for val in &mut noisy_vector {
let noise = self.gaussian_random(0.0, noise_scale, (size + i) as u64);
*val += noise;
}
query_vectors.push(Vector::new(noisy_vector));
}
let dataset = BenchmarkDataset {
name: name.to_string(),
description: format!(
"Clustered dataset with {size} vectors in {num_clusters} clusters of {dimensions} dimensions"
),
train_vectors,
query_vectors,
ground_truth: None,
metadata: {
let mut meta = HashMap::new();
meta.insert("type".to_string(), "clustered".to_string());
meta.insert("size".to_string(), size.to_string());
meta.insert("dimensions".to_string(), dimensions.to_string());
meta.insert("clusters".to_string(), num_clusters.to_string());
meta
},
};
self.add_dataset(dataset);
Ok(())
}
fn generate_uniform_dataset(
&mut self,
name: &str,
size: usize,
dimensions: usize,
query_count: usize,
) -> Result<()> {
let mut train_vectors = Vec::new();
let mut query_vectors = Vec::new();
for i in 0..size {
let mut values = Vec::with_capacity(dimensions);
let mut state = i as u64;
for _ in 0..dimensions {
state = state.wrapping_mul(1103515245).wrapping_add(12345);
let normalized = (state as f32) / (u64::MAX as f32);
values.push(normalized); }
train_vectors.push(Vector::new(values));
}
for i in 0..query_count {
let mut values = Vec::with_capacity(dimensions);
let mut state = (size + i) as u64;
for _ in 0..dimensions {
state = state.wrapping_mul(1103515245).wrapping_add(12345);
let normalized = (state as f32) / (u64::MAX as f32);
values.push(normalized);
}
query_vectors.push(Vector::new(values));
}
let dataset = BenchmarkDataset {
name: name.to_string(),
description: format!("Uniform dataset with {size} vectors of {dimensions} dimensions"),
train_vectors,
query_vectors,
ground_truth: None,
metadata: {
let mut meta = HashMap::new();
meta.insert("type".to_string(), "uniform".to_string());
meta.insert("size".to_string(), size.to_string());
meta.insert("dimensions".to_string(), dimensions.to_string());
meta
},
};
self.add_dataset(dataset);
Ok(())
}
fn gaussian_random(&self, mean: f32, std_dev: f32, seed: u64) -> f32 {
let mut state = seed;
state = state.wrapping_mul(1103515245).wrapping_add(12345);
let u1 = (state as f32) / (u64::MAX as f32);
state = state.wrapping_mul(1103515245).wrapping_add(12345);
let u2 = (state as f32) / (u64::MAX as f32);
let z0 = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
mean + std_dev * z0
}
pub fn run_all_benchmarks(&mut self) -> Result<Vec<BenchmarkResult>> {
let mut all_results = Vec::new();
for dataset in &self.datasets {
for (alg_idx, algorithm) in self.algorithms.iter().enumerate() {
let test_case = BenchmarkTestCase {
name: format!("{}_{alg_idx}", dataset.name),
description: format!("Benchmark {alg_idx} on {}", dataset.name),
dataset: dataset.name.clone(),
algorithm: format!("algorithm_{alg_idx}"),
parameters: HashMap::new(),
query_count: dataset.query_vectors.len(),
k: 10,
};
let result = self.run_single_benchmark(&test_case, dataset, algorithm.as_ref())?;
all_results.push(result);
}
}
self.results.extend(all_results.clone());
Ok(all_results)
}
fn run_single_benchmark(
&self,
test_case: &BenchmarkTestCase,
dataset: &BenchmarkDataset,
algorithm: &dyn VectorIndex,
) -> Result<BenchmarkResult> {
let mut profiler = PerformanceProfiler::new();
tracing::info!("Running benchmark: {}", test_case.name);
profiler.checkpoint("benchmark_start");
profiler.checkpoint("index_build_start");
let mut index = self.create_index_copy(algorithm)?;
for (i, vector) in dataset.train_vectors.iter().enumerate() {
index.insert(format!("vec_{i}"), vector.clone())?;
}
profiler.checkpoint("index_build_end");
profiler.checkpoint("warmup_start");
for _ in 0..self.config.warmup_runs {
for query in dataset.query_vectors.iter().take(10) {
let _ = index.search_knn(query, test_case.k)?;
}
}
profiler.checkpoint("warmup_end");
profiler.checkpoint("benchmark_queries_start");
let mut query_times = Vec::new();
for query in &dataset.query_vectors {
profiler.sample_memory();
let start = Instant::now();
let _results = index.search_knn(query, test_case.k)?;
let query_time = start.elapsed();
query_times.push(query_time);
}
profiler.checkpoint("benchmark_queries_end");
let performance = self.calculate_performance_metrics(&query_times)?;
let quality = if let Some(ground_truth) = &dataset.ground_truth {
Some(self.calculate_quality_metrics(
index.as_ref(),
&dataset.query_vectors,
ground_truth,
test_case.k,
)?)
} else {
None
};
let memory = if self.config.profile_memory {
Some(self.calculate_memory_metrics(&profiler, dataset.train_vectors.len())?)
} else {
None
};
let system_info = self.get_system_info();
Ok(BenchmarkResult {
test_case: test_case.clone(),
performance,
quality,
memory,
scalability: None, system_info,
timestamp: std::time::SystemTime::now(),
})
}
fn create_index_copy(&self, _algorithm: &dyn VectorIndex) -> Result<Box<dyn VectorIndex>> {
Ok(Box::new(crate::MemoryVectorIndex::new()))
}
fn calculate_performance_metrics(
&self,
query_times: &[Duration],
) -> Result<PerformanceMetrics> {
if query_times.is_empty() {
return Err(anyhow!("No query times to analyze"));
}
let mut sorted_times = query_times.to_vec();
sorted_times.sort();
let avg_query_time = Duration::from_nanos(
(query_times.iter().map(|d| d.as_nanos()).sum::<u128>() / query_times.len() as u128)
.try_into()
.expect("average query time should fit in u64"),
);
let median_query_time = sorted_times[sorted_times.len() / 2];
let p95_idx = (sorted_times.len() as f64 * 0.95) as usize;
let p99_idx = (sorted_times.len() as f64 * 0.99) as usize;
let p95_query_time = sorted_times[p95_idx.min(sorted_times.len() - 1)];
let p99_query_time = sorted_times[p99_idx.min(sorted_times.len() - 1)];
let min_query_time = sorted_times[0];
let max_query_time = sorted_times[sorted_times.len() - 1];
let mean_nanos = avg_query_time.as_nanos() as f64;
let variance = query_times
.iter()
.map(|d| {
let diff = d.as_nanos() as f64 - mean_nanos;
diff * diff
})
.sum::<f64>()
/ query_times.len() as f64;
let std_dev_query_time = Duration::from_nanos(variance.sqrt() as u64);
let queries_per_second = 1.0 / avg_query_time.as_secs_f64();
let throughput =
query_times.len() as f64 / query_times.iter().map(|d| d.as_secs_f64()).sum::<f64>();
Ok(PerformanceMetrics {
avg_query_time,
median_query_time,
p95_query_time,
p99_query_time,
min_query_time,
max_query_time,
std_dev_query_time,
queries_per_second,
index_build_time: None, index_update_time: None,
throughput,
})
}
fn calculate_quality_metrics(
&self,
index: &dyn VectorIndex,
queries: &[Vector],
ground_truth: &[Vec<usize>],
k: usize,
) -> Result<QualityMetrics> {
let mut total_recall = 0.0;
let mut total_precision = 0.0;
let mut total_queries = 0;
for (query_idx, query) in queries.iter().enumerate() {
if query_idx >= ground_truth.len() {
break;
}
let results = index.search_knn(query, k)?;
let returned_indices: Vec<usize> = results
.iter()
.filter_map(|(uri, _)| {
uri.strip_prefix("vec_")
.and_then(|s| s.parse::<usize>().ok())
})
.collect();
let true_neighbors = &ground_truth[query_idx];
let true_neighbors_k: std::collections::HashSet<usize> =
true_neighbors.iter().take(k).copied().collect();
let found_true = returned_indices
.iter()
.filter(|&idx| true_neighbors_k.contains(idx))
.count();
let recall = found_true as f64 / k.min(true_neighbors.len()) as f64;
let precision = found_true as f64 / returned_indices.len() as f64;
total_recall += recall;
total_precision += precision;
total_queries += 1;
}
let avg_recall = total_recall / total_queries as f64;
let avg_precision = total_precision / total_queries as f64;
Ok(QualityMetrics {
recall_at_k: avg_recall,
precision_at_k: avg_precision,
mean_average_precision: avg_precision, ndcg_at_k: avg_recall, distance_ratio: None,
relative_error: None,
})
}
fn calculate_memory_metrics(
&self,
profiler: &PerformanceProfiler,
vector_count: usize,
) -> Result<MemoryMetrics> {
let memory_profile = profiler.get_memory_profile();
if memory_profile.is_empty() {
return Ok(MemoryMetrics {
peak_memory_bytes: 0,
avg_memory_bytes: 0,
memory_per_vector: 0.0,
index_overhead_bytes: 0,
memory_efficiency: 0.0,
});
}
let peak_memory_bytes = memory_profile
.iter()
.map(|(_, mem)| *mem)
.max()
.unwrap_or(0);
let avg_memory_bytes =
memory_profile.iter().map(|(_, mem)| *mem).sum::<usize>() / memory_profile.len();
let memory_per_vector = avg_memory_bytes as f64 / vector_count as f64;
Ok(MemoryMetrics {
peak_memory_bytes,
avg_memory_bytes,
memory_per_vector,
index_overhead_bytes: 0, memory_efficiency: 1.0, })
}
fn get_system_info(&self) -> SystemInfo {
SystemInfo {
cpu_info: self.get_cpu_info(),
total_ram_gb: self.get_total_ram_gb(),
available_ram_gb: self.get_available_ram_gb(),
os: std::env::consts::OS.to_string(),
rust_version: self.get_rust_version(),
simd_features: self.get_simd_features(),
gpu_info: self.get_gpu_info(),
}
}
fn get_cpu_info(&self) -> String {
#[cfg(target_os = "linux")]
{
if let Ok(content) = std::fs::read_to_string("/proc/cpuinfo") {
for line in content.lines() {
if line.starts_with("model name") {
if let Some(name) = line.split(':').nth(1) {
return name.trim().to_string();
}
}
}
}
}
#[cfg(target_os = "macos")]
{
use std::process::Command;
if let Ok(output) = Command::new("sysctl")
.args(["-n", "machdep.cpu.brand_string"])
.output()
{
if let Ok(cpu_name) = String::from_utf8(output.stdout) {
return cpu_name.trim().to_string();
}
}
}
"Unknown CPU".to_string()
}
fn get_total_ram_gb(&self) -> f64 {
#[cfg(target_os = "linux")]
{
if let Ok(content) = std::fs::read_to_string("/proc/meminfo") {
for line in content.lines() {
if line.starts_with("MemTotal:") {
if let Some(kb_str) = line.split_whitespace().nth(1) {
if let Ok(kb) = kb_str.parse::<u64>() {
return kb as f64 / 1024.0 / 1024.0; }
}
}
}
}
}
#[cfg(target_os = "macos")]
{
use std::process::Command;
if let Ok(output) = Command::new("sysctl").args(["-n", "hw.memsize"]).output() {
if let Ok(mem_str) = String::from_utf8(output.stdout) {
if let Ok(bytes) = mem_str.trim().parse::<u64>() {
return bytes as f64 / 1024.0 / 1024.0 / 1024.0; }
}
}
}
8.0 }
fn get_available_ram_gb(&self) -> f64 {
self.get_total_ram_gb() * 0.8
}
fn get_rust_version(&self) -> String {
std::env::var("RUSTC_VERSION").unwrap_or_else(|_| "unknown".to_string())
}
fn get_simd_features(&self) -> Vec<String> {
let mut features = Vec::new();
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("sse") {
features.push("SSE".to_string());
}
if is_x86_feature_detected!("sse2") {
features.push("SSE2".to_string());
}
if is_x86_feature_detected!("sse3") {
features.push("SSE3".to_string());
}
if is_x86_feature_detected!("sse4.1") {
features.push("SSE4.1".to_string());
}
if is_x86_feature_detected!("sse4.2") {
features.push("SSE4.2".to_string());
}
if is_x86_feature_detected!("avx") {
features.push("AVX".to_string());
}
if is_x86_feature_detected!("avx2") {
features.push("AVX2".to_string());
}
}
#[cfg(target_arch = "aarch64")]
{
features.push("NEON".to_string());
}
features
}
fn get_gpu_info(&self) -> Option<String> {
None
}
pub fn export_results(&self, format: BenchmarkOutputFormat) -> Result<String> {
match format {
BenchmarkOutputFormat::Json => serde_json::to_string_pretty(&self.results)
.map_err(|e| anyhow!("Failed to serialize to JSON: {}", e)),
BenchmarkOutputFormat::Csv => self.export_csv(),
BenchmarkOutputFormat::Table => self.export_table(),
BenchmarkOutputFormat::AnnBenchmarks => self.export_ann_benchmarks(),
}
}
fn export_csv(&self) -> Result<String> {
let mut csv = String::new();
csv.push_str(
"dataset,algorithm,avg_query_time_ms,queries_per_second,recall_at_k,memory_mb\n",
);
for result in &self.results {
csv.push_str(&format!(
"{},{},{:.3},{:.2},{:.3},{:.2}\n",
result.test_case.dataset,
result.test_case.algorithm,
result.performance.avg_query_time.as_millis(),
result.performance.queries_per_second,
result
.quality
.as_ref()
.map(|q| q.recall_at_k)
.unwrap_or(0.0),
result
.memory
.as_ref()
.map(|m| m.avg_memory_bytes as f64 / 1024.0 / 1024.0)
.unwrap_or(0.0),
));
}
Ok(csv)
}
fn export_table(&self) -> Result<String> {
let mut table = String::new();
table.push_str(&format!(
"{:<20} {:<15} {:<15} {:<15} {:<10}\n",
"Dataset", "Algorithm", "Avg Time (ms)", "QPS", "Recall@K"
));
table.push_str(&"-".repeat(80));
table.push('\n');
for result in &self.results {
table.push_str(&format!(
"{:<20} {:<15} {:<15.3} {:<15.2} {:<10.3}\n",
result.test_case.dataset,
result.test_case.algorithm,
result.performance.avg_query_time.as_millis(),
result.performance.queries_per_second,
result
.quality
.as_ref()
.map(|q| q.recall_at_k)
.unwrap_or(0.0),
));
}
Ok(table)
}
fn export_ann_benchmarks(&self) -> Result<String> {
let mut results = serde_json::Map::new();
for result in &self.results {
let mut entry = serde_json::Map::new();
entry.insert(
"k".to_string(),
serde_json::Value::Number(result.test_case.k.into()),
);
entry.insert(
"recall".to_string(),
serde_json::Value::Number(
serde_json::Number::from_f64(
result
.quality
.as_ref()
.map(|q| q.recall_at_k)
.unwrap_or(0.0),
)
.unwrap_or(serde_json::Number::from(0)),
),
);
entry.insert(
"qps".to_string(),
serde_json::Value::Number(
serde_json::Number::from_f64(result.performance.queries_per_second)
.unwrap_or(serde_json::Number::from(0)),
),
);
let key = format!(
"{}_{}",
result.test_case.dataset, result.test_case.algorithm
);
results.insert(key, serde_json::Value::Object(entry));
}
serde_json::to_string_pretty(&results)
.map_err(|e| anyhow!("Failed to serialize ANN-Benchmarks format: {}", e))
}
pub fn run_scalability_benchmark(&mut self, dataset_name: &str) -> Result<ScalabilityMetrics> {
let dataset = self
.datasets
.iter()
.find(|d| d.name == dataset_name)
.ok_or_else(|| anyhow!("Dataset not found: {}", dataset_name))?;
let mut performance_scaling = Vec::new();
let mut memory_scaling = Vec::new();
let mut build_time_scaling = Vec::new();
let test_sizes = [100, 500, 1000, 2000, 5000];
for &size in &test_sizes {
if size > dataset.train_vectors.len() {
continue;
}
let subset_vectors = &dataset.train_vectors[..size];
let test_queries = &dataset.query_vectors[..10.min(dataset.query_vectors.len())];
let build_start = Instant::now();
let mut index = Box::new(crate::MemoryVectorIndex::new());
for (i, vector) in subset_vectors.iter().enumerate() {
index.insert(format!("vec_{i}"), vector.clone())?;
}
let build_time = build_start.elapsed();
let mut query_times = Vec::new();
let memory_start = self.get_current_memory_usage();
for query in test_queries {
let start = Instant::now();
let _ = index.search_knn(query, 10)?;
query_times.push(start.elapsed());
}
let memory_end = self.get_current_memory_usage();
let avg_query_time = Duration::from_nanos(
(query_times.iter().map(|d| d.as_nanos()).sum::<u128>()
/ query_times.len() as u128)
.try_into()
.expect("average query time should fit in u64"),
);
performance_scaling.push((size, avg_query_time));
memory_scaling.push((size, memory_end.saturating_sub(memory_start)));
build_time_scaling.push((size, build_time));
}
let mut concurrency_scaling = Vec::new();
let thread_counts = [1, 2, 4, 8];
for &thread_count in &thread_counts {
let qps = self.measure_concurrent_performance(
&dataset.query_vectors[..100.min(dataset.query_vectors.len())],
thread_count,
)?;
concurrency_scaling.push((thread_count, qps));
}
Ok(ScalabilityMetrics {
performance_scaling,
memory_scaling,
build_time_scaling,
concurrency_scaling,
})
}
fn get_current_memory_usage(&self) -> usize {
0 }
fn measure_concurrent_performance(
&self,
queries: &[Vector],
thread_count: usize,
) -> Result<f64> {
use std::sync::Arc;
use std::thread;
let index = Arc::new(crate::MemoryVectorIndex::new());
let queries = Arc::new(queries.to_vec());
let start_time = Instant::now();
let mut handles = Vec::new();
for thread_id in 0..thread_count {
let index = Arc::clone(&index);
let queries = Arc::clone(&queries);
let handle = thread::spawn(move || {
let queries_per_thread = queries.len() / thread_count;
let start_idx = thread_id * queries_per_thread;
let end_idx = if thread_id == thread_count - 1 {
queries.len()
} else {
start_idx + queries_per_thread
};
let mut query_count = 0;
for query in &queries[start_idx..end_idx] {
if index.search_knn(query, 10).is_ok() {
query_count += 1;
}
}
query_count
});
handles.push(handle);
}
let total_queries: usize = handles.into_iter().map(|h| h.join().unwrap_or(0)).sum();
let elapsed = start_time.elapsed();
Ok(total_queries as f64 / elapsed.as_secs_f64())
}
}
pub struct BenchmarkRunner;
impl BenchmarkRunner {
pub fn run_standard_benchmarks() -> Result<Vec<BenchmarkResult>> {
let config = BenchmarkConfig::default();
let mut suite = BenchmarkSuite::new(config);
suite.generate_synthetic_datasets()?;
suite.add_algorithm(Box::new(crate::MemoryVectorIndex::new()));
suite.run_all_benchmarks()
}
pub fn run_quick_benchmarks() -> Result<Vec<BenchmarkResult>> {
let config = BenchmarkConfig {
warmup_runs: 1,
benchmark_runs: 3,
max_duration: Duration::from_secs(30),
..BenchmarkConfig::default()
};
let mut suite = BenchmarkSuite::new(config);
suite.generate_random_dataset("quick_test", 100, 64, 10)?;
suite.add_algorithm(Box::new(crate::MemoryVectorIndex::new()));
suite.run_all_benchmarks()
}
pub fn run_comprehensive_benchmarks() -> Result<String> {
let results = Self::run_standard_benchmarks()?;
let config = BenchmarkConfig::default();
let suite = BenchmarkSuite {
config,
datasets: Vec::new(),
algorithms: Vec::new(),
results,
};
suite.export_results(BenchmarkOutputFormat::Table)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_benchmark_suite_creation() {
let config = BenchmarkConfig::default();
let suite = BenchmarkSuite::new(config);
assert_eq!(suite.datasets.len(), 0);
assert_eq!(suite.algorithms.len(), 0);
}
#[test]
fn test_synthetic_dataset_generation() -> Result<()> {
let config = BenchmarkConfig::default();
let mut suite = BenchmarkSuite::new(config);
suite.generate_synthetic_datasets()?;
assert!(!suite.datasets.is_empty());
for dataset in &suite.datasets {
assert!(!dataset.train_vectors.is_empty());
assert!(!dataset.query_vectors.is_empty());
}
Ok(())
}
#[test]
fn test_performance_metrics_calculation() -> Result<()> {
let config = BenchmarkConfig::default();
let suite = BenchmarkSuite::new(config);
let query_times = vec![
Duration::from_millis(10),
Duration::from_millis(15),
Duration::from_millis(12),
Duration::from_millis(20),
Duration::from_millis(8),
];
let metrics = suite.calculate_performance_metrics(&query_times)?;
assert!(metrics.avg_query_time.as_millis() > 0);
assert!(metrics.queries_per_second > 0.0);
assert!(metrics.min_query_time <= metrics.median_query_time);
assert!(metrics.median_query_time <= metrics.max_query_time);
Ok(())
}
#[test]
fn test_quick_benchmarks() -> Result<()> {
let results = BenchmarkRunner::run_quick_benchmarks();
assert!(results.is_ok());
let results = results?;
assert!(!results.is_empty());
for result in results {
assert!(result.performance.avg_query_time.as_nanos() > 0);
assert!(result.performance.queries_per_second > 0.0);
}
Ok(())
}
#[test]
fn test_export_formats() {
let config = BenchmarkConfig::default();
let mut suite = BenchmarkSuite::new(config);
let test_case = BenchmarkTestCase {
name: "test".to_string(),
description: "test case".to_string(),
dataset: "test_data".to_string(),
algorithm: "test_alg".to_string(),
parameters: HashMap::new(),
query_count: 10,
k: 5,
};
let result = BenchmarkResult {
test_case,
performance: PerformanceMetrics {
avg_query_time: Duration::from_millis(10),
median_query_time: Duration::from_millis(10),
p95_query_time: Duration::from_millis(15),
p99_query_time: Duration::from_millis(20),
min_query_time: Duration::from_millis(5),
max_query_time: Duration::from_millis(25),
std_dev_query_time: Duration::from_millis(3),
queries_per_second: 100.0,
index_build_time: None,
index_update_time: None,
throughput: 100.0,
},
quality: None,
memory: None,
scalability: None,
system_info: SystemInfo {
cpu_info: "Test CPU".to_string(),
total_ram_gb: 16.0,
available_ram_gb: 12.0,
os: "test".to_string(),
rust_version: "1.70.0".to_string(),
simd_features: vec!["AVX2".to_string()],
gpu_info: None,
},
timestamp: std::time::SystemTime::now(),
};
suite.results.push(result);
let json_output = suite.export_results(BenchmarkOutputFormat::Json);
assert!(json_output.is_ok());
let csv_output = suite.export_results(BenchmarkOutputFormat::Csv);
assert!(csv_output.is_ok());
let table_output = suite.export_results(BenchmarkOutputFormat::Table);
assert!(table_output.is_ok());
}
#[test]
fn test_profiler() {
let mut profiler = PerformanceProfiler::new();
profiler.checkpoint("start");
std::thread::sleep(Duration::from_millis(50));
profiler.checkpoint("middle");
std::thread::sleep(Duration::from_millis(100));
profiler.checkpoint("end");
let breakdown = profiler.get_timing_breakdown();
assert_eq!(breakdown.len(), 3);
for (name, duration) in breakdown {
assert!(!name.is_empty());
assert!(duration.as_nanos() > 0);
}
}
}