use std::time::Instant;
use iqdb_index::IndexCore;
use iqdb_types::SearchParams;
use crate::error::{EvalError, Result};
use crate::report::LatencyReport;
#[derive(Debug, Clone, Copy, Default)]
pub struct LatencyConfig {
pub warmup: usize,
}
pub fn latency<I: IndexCore>(
index: &I,
queries: &[Vec<f32>],
params: &SearchParams,
config: &LatencyConfig,
) -> Result<LatencyReport> {
if queries.is_empty() {
return Err(EvalError::EmptyInput { kind: "queries" });
}
let dim = index.dim();
for query in queries {
if query.len() != dim {
return Err(EvalError::DimensionMismatch {
expected: dim,
found: query.len(),
});
}
}
let span = tracing::info_span!(
"eval.latency",
n_queries = queries.len(),
warmup = config.warmup,
);
let _enter = span.enter();
for i in 0..config.warmup {
let q = &queries[i % queries.len()];
let _ = index.search(q, params)?;
}
let mut samples_us: Vec<f64> = Vec::with_capacity(queries.len());
for query in queries {
let t0 = Instant::now();
let _hits = index.search(query, params)?;
let elapsed_us = t0.elapsed().as_nanos() as f64 / 1_000.0;
samples_us.push(elapsed_us);
}
Ok(summarize(&mut samples_us))
}
fn summarize(samples: &mut [f64]) -> LatencyReport {
samples.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let n = samples.len();
let sum: f64 = samples.iter().sum();
let mean_us = sum / n as f64;
let min_us = samples[0];
let max_us = samples[n - 1];
let p50_us = samples[percentile_index(0.50, n)];
let p95_us = samples[percentile_index(0.95, n)];
let p99_us = samples[percentile_index(0.99, n)];
let total_secs = sum / 1_000_000.0;
let qps = if total_secs > 0.0 {
n as f64 / total_secs
} else {
f64::INFINITY
};
LatencyReport {
query_count: n,
mean_us,
min_us,
max_us,
p50_us,
p95_us,
p99_us,
qps,
}
}
fn percentile_index(p: f64, n: usize) -> usize {
debug_assert!(n >= 1, "percentile_index requires n >= 1");
let raw = (p * n as f64).ceil() as isize - 1;
if raw < 0 {
0
} else if (raw as usize) >= n {
n - 1
} else {
raw as usize
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
#[test]
fn percentile_index_typical_n100() {
assert_eq!(percentile_index(0.50, 100), 49);
assert_eq!(percentile_index(0.95, 100), 94);
assert_eq!(percentile_index(0.99, 100), 98);
}
#[test]
fn percentile_index_n1_returns_0() {
assert_eq!(percentile_index(0.50, 1), 0);
assert_eq!(percentile_index(0.99, 1), 0);
}
#[test]
fn percentile_index_n10_p99_clamps_to_last() {
assert_eq!(percentile_index(0.99, 10), 9);
}
#[test]
fn summarize_orders_percentiles() {
let mut samples = vec![5.0, 1.0, 4.0, 2.0, 3.0];
let r = summarize(&mut samples);
assert_eq!(r.query_count, 5);
assert!(r.min_us <= r.p50_us);
assert!(r.p50_us <= r.p95_us);
assert!(r.p95_us <= r.p99_us);
assert!(r.p99_us <= r.max_us);
assert_eq!(r.min_us, 1.0);
assert_eq!(r.max_us, 5.0);
}
#[test]
fn summarize_single_sample_collapses_all_fields() {
let mut samples = vec![42.0];
let r = summarize(&mut samples);
assert_eq!(r.query_count, 1);
assert_eq!(r.min_us, 42.0);
assert_eq!(r.max_us, 42.0);
assert_eq!(r.p50_us, 42.0);
assert_eq!(r.p95_us, 42.0);
assert_eq!(r.p99_us, 42.0);
assert!(r.qps > 0.0);
}
#[test]
fn summarize_zero_total_time_yields_infinite_qps() {
let mut samples = vec![0.0, 0.0, 0.0];
let r = summarize(&mut samples);
assert!(r.qps.is_infinite());
}
mod with_index {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use crate::{LatencyConfig, build_index_from_base, latency};
use iqdb_flat::{FlatConfig, FlatIndex};
use iqdb_types::{DistanceMetric, SearchParams};
const M: DistanceMetric = DistanceMetric::Euclidean;
fn index() -> FlatIndex {
let base: Vec<Vec<f32>> = vec![vec![0.0], vec![1.0], vec![2.0]];
build_index_from_base(FlatConfig, 1, M, &base).unwrap()
}
#[test]
fn single_query_reports_one_sample() {
let idx = index();
let r = latency(
&idx,
&[vec![0.0]],
&SearchParams::new(1, M),
&LatencyConfig::default(),
)
.unwrap();
assert_eq!(r.query_count, 1);
assert_eq!(r.min_us, r.max_us);
assert_eq!(r.p50_us, r.p99_us);
}
#[test]
fn warmup_larger_than_query_set_cycles_and_excludes_itself() {
let idx = index();
let queries = vec![vec![0.0], vec![2.0]];
let cfg = LatencyConfig { warmup: 10 };
let r = latency(&idx, &queries, &SearchParams::new(1, M), &cfg).unwrap();
assert_eq!(r.query_count, 2);
}
}
}