osobhandClaude Fable 5.1
f507803ec1
`parallel` joins the agent's default features, so the HNSW bulk build uses the thread pool: cold index build at 10K records 1152 -> ~380 ms in a same-moment A/B (the graph is identical either way). Nothing else on the measured paths changes — ingest, checkpoint, open and steady-state query times are the same with the feature on or off. Adds rayon to the default dependency set; opt out with `--no-default-features --features float16,hnsw`. Harness: `--e2e-only` runs the end-to-end section without the index benchmarks. Note for anyone comparing numbers: this machine's absolute timings drifted ~1.5x over a long session, so only same-moment A/B runs are comparable. Co-Authored-By: Claude Fable 5.1 <[email protected]>
513 lines
18 KiB
Rust
513 lines
18 KiB
Rust
//! Search measurement harness: recall vs. speed for the HNSW index, and
|
|
//! end-to-end `hybrid_search` latency as the store grows.
|
|
//!
|
|
//! Every search-path change should be justified by a before/after run of this
|
|
//! binary. It reports, for deterministic synthetic data:
|
|
//!
|
|
//! * **ANN** — index build time, and for each `ef`: recall@10 against an exact
|
|
//! brute-force scan, queries/second, and p50/p99 latency.
|
|
//! * **End to end** — `HDF5Memory`: ingest time, checkpoint time, `open()`
|
|
//! time, the one-off cold index build (first query ever), the first query
|
|
//! after a reopen, and steady-state `hybrid_search` p50/p99 at each size.
|
|
//!
|
|
//! Data is *clustered* (points = cluster centre + noise, unit-normalised), not
|
|
//! uniform: uniform random high-dimensional vectors are nearly equidistant,
|
|
//! which makes recall numbers meaningless and is nothing like embeddings.
|
|
//!
|
|
//! ```text
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness # 1K, 10K
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --full # + 100K
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform
|
|
//! ```
|
|
|
|
use std::time::{Duration, Instant};
|
|
|
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
|
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
|
|
|
const DIM: usize = 384;
|
|
const K: usize = 10;
|
|
const N_QUERIES: usize = 200;
|
|
const HNSW_M: usize = 16;
|
|
const HNSW_EF_CONSTRUCTION: usize = 64;
|
|
const EF_VALUES: [usize; 5] = [16, 32, 64, 128, 256];
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Deterministic data
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct Rng(u64);
|
|
|
|
impl Rng {
|
|
fn next_u64(&mut self) -> u64 {
|
|
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
|
let mut z = self.0;
|
|
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
|
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
|
z ^ (z >> 31)
|
|
}
|
|
|
|
/// Uniform in [0, 1).
|
|
fn unit(&mut self) -> f32 {
|
|
(self.next_u64() >> 40) as f32 / (1u64 << 24) as f32
|
|
}
|
|
|
|
/// Approximately standard normal (sum of uniforms).
|
|
fn gauss(&mut self) -> f32 {
|
|
let sum: f32 = (0..6).map(|_| self.unit()).sum();
|
|
(sum - 3.0) * std::f32::consts::SQRT_2
|
|
}
|
|
|
|
fn below(&mut self, n: usize) -> usize {
|
|
(self.next_u64() % n as u64) as usize
|
|
}
|
|
}
|
|
|
|
fn normalize(v: &mut [f32]) {
|
|
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
if norm > 0.0 {
|
|
v.iter_mut().for_each(|x| *x /= norm);
|
|
}
|
|
}
|
|
|
|
struct Dataset {
|
|
vectors: Vec<Vec<f32>>,
|
|
queries: Vec<Vec<f32>>,
|
|
/// Cluster id of each vector (used to give records topical text).
|
|
cluster_of: Vec<usize>,
|
|
query_cluster: Vec<usize>,
|
|
}
|
|
|
|
/// `--uniform`: isotropic random unit vectors instead of clusters. Not a
|
|
/// realistic workload, but a useful second distribution — a recall problem
|
|
/// that appears only on clustered data points at graph connectivity.
|
|
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
|
|
|
fn make_dataset(n: usize, seed: u64) -> Dataset {
|
|
let mut rng = Rng(seed);
|
|
if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
|
|
let random_unit = |rng: &mut Rng| {
|
|
let mut v: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
|
|
normalize(&mut v);
|
|
v
|
|
};
|
|
return Dataset {
|
|
vectors: (0..n).map(|_| random_unit(&mut rng)).collect(),
|
|
queries: (0..N_QUERIES).map(|_| random_unit(&mut rng)).collect(),
|
|
cluster_of: vec![0; n],
|
|
query_cluster: vec![0; N_QUERIES],
|
|
};
|
|
}
|
|
let n_clusters = (n / 100).clamp(8, 512);
|
|
let centres: Vec<Vec<f32>> = (0..n_clusters)
|
|
.map(|_| {
|
|
let mut c: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
|
|
normalize(&mut c);
|
|
c
|
|
})
|
|
.collect();
|
|
let point = |rng: &mut Rng, cluster: usize| {
|
|
// Noise comparable to the centre's per-dimension magnitude, so
|
|
// clusters overlap and the nearest neighbours are non-trivial.
|
|
let scale = 0.6 / (DIM as f32).sqrt();
|
|
let mut v: Vec<f32> = centres[cluster]
|
|
.iter()
|
|
.map(|c| c + rng.gauss() * scale)
|
|
.collect();
|
|
normalize(&mut v);
|
|
v
|
|
};
|
|
let mut vectors = Vec::with_capacity(n);
|
|
let mut cluster_of = Vec::with_capacity(n);
|
|
for _ in 0..n {
|
|
let c = rng.below(n_clusters);
|
|
vectors.push(point(&mut rng, c));
|
|
cluster_of.push(c);
|
|
}
|
|
let mut queries = Vec::with_capacity(N_QUERIES);
|
|
let mut query_cluster = Vec::with_capacity(N_QUERIES);
|
|
for _ in 0..N_QUERIES {
|
|
let c = rng.below(n_clusters);
|
|
queries.push(point(&mut rng, c));
|
|
query_cluster.push(c);
|
|
}
|
|
Dataset {
|
|
vectors,
|
|
queries,
|
|
cluster_of,
|
|
query_cluster,
|
|
}
|
|
}
|
|
|
|
const WORDS: &[&str] = &[
|
|
"deploy", "latency", "cache", "schema", "index", "vector", "memory", "agent", "kernel",
|
|
"buffer", "socket", "thread", "tensor", "gradient", "ledger", "invoice", "meeting", "roadmap",
|
|
"customer", "contract", "sensor", "orbit", "protein", "genome", "harbor", "bridge", "engine",
|
|
"battery", "harvest", "weather", "museum", "recipe",
|
|
];
|
|
|
|
/// Text whose vocabulary is biased by cluster, so keyword and vector signals
|
|
/// agree the way they do for real embedded text.
|
|
fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
|
|
let topic = [
|
|
WORDS[cluster % WORDS.len()],
|
|
WORDS[(cluster / 7 + 3) % WORDS.len()],
|
|
];
|
|
let mut words = Vec::with_capacity(14);
|
|
for j in 0..14 {
|
|
if j % 3 == 0 {
|
|
words.push(topic[j / 3 % 2]);
|
|
} else {
|
|
words.push(WORDS[rng.below(WORDS.len())]);
|
|
}
|
|
}
|
|
format!("record {i}: {}", words.join(" "))
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Measurement helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
|
|
// Vectors are unit length, so cosine order == dot-product order.
|
|
let mut scored: Vec<(usize, f32)> = vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, v.iter().zip(query).map(|(a, b)| a * b).sum()))
|
|
.collect();
|
|
scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
|
scored.truncate(k);
|
|
scored.into_iter().map(|(i, _)| i).collect()
|
|
}
|
|
|
|
struct Latency {
|
|
p50: Duration,
|
|
p99: Duration,
|
|
qps: f64,
|
|
}
|
|
|
|
fn summarize(mut samples: Vec<Duration>) -> Latency {
|
|
samples.sort();
|
|
let total: Duration = samples.iter().sum();
|
|
let at = |q: f64| samples[((samples.len() - 1) as f64 * q).round() as usize];
|
|
Latency {
|
|
p50: at(0.50),
|
|
p99: at(0.99),
|
|
qps: samples.len() as f64 / total.as_secs_f64(),
|
|
}
|
|
}
|
|
|
|
fn micros(d: Duration) -> f64 {
|
|
d.as_secs_f64() * 1e6
|
|
}
|
|
|
|
fn millis(d: Duration) -> f64 {
|
|
d.as_secs_f64() * 1e3
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ANN: recall vs speed
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
|
|
let data = make_dataset(n, 0xA11CE ^ n as u64);
|
|
let truth: Vec<Vec<usize>> = data
|
|
.queries
|
|
.iter()
|
|
.map(|q| exact_top_k(&data.vectors, q, K))
|
|
.collect();
|
|
|
|
let started = Instant::now();
|
|
let index = HnswIndex::build_with_metric(
|
|
&data.vectors,
|
|
HNSW_M,
|
|
HNSW_EF_CONSTRUCTION,
|
|
DistanceMetric::Cosine,
|
|
);
|
|
let build = started.elapsed();
|
|
|
|
// Exact scan baseline, for scale.
|
|
let exact = summarize(
|
|
data.queries
|
|
.iter()
|
|
.map(|q| {
|
|
let t = Instant::now();
|
|
std::hint::black_box(exact_top_k(&data.vectors, q, K));
|
|
t.elapsed()
|
|
})
|
|
.collect(),
|
|
);
|
|
|
|
println!(
|
|
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n"
|
|
);
|
|
println!(
|
|
"build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
|
|
millis(build),
|
|
n as f64 / build.as_secs_f64(),
|
|
exact.qps,
|
|
micros(exact.p50)
|
|
);
|
|
println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
|
|
println!("|---:|---:|---:|---:|---:|");
|
|
for ef in EF_VALUES {
|
|
let mut hits = 0usize;
|
|
let mut samples = Vec::with_capacity(data.queries.len());
|
|
for (q, want) in data.queries.iter().zip(&truth) {
|
|
let t = Instant::now();
|
|
let got = index.search(q, K, ef);
|
|
samples.push(t.elapsed());
|
|
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
|
|
}
|
|
let recall = hits as f64 / (K * data.queries.len()) as f64;
|
|
let lat = summarize(samples);
|
|
println!(
|
|
"| {ef} | {recall:.4} | {:.0} | {:.0} | {:.0} |",
|
|
lat.qps,
|
|
micros(lat.p50),
|
|
micros(lat.p99)
|
|
);
|
|
json.push(serde_json::json!({
|
|
"bench": "hnsw", "n": n, "ef": ef, "recall_at_10": recall,
|
|
"qps": lat.qps, "p50_us": micros(lat.p50), "p99_us": micros(lat.p99),
|
|
"build_ms": millis(build),
|
|
}));
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// End to end: HDF5Memory::hybrid_search
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
|
|
let data = make_dataset(n, 0xE2E ^ n as u64);
|
|
let dir = tempfile::TempDir::new().unwrap();
|
|
let path = dir.path().join("store.h5");
|
|
let mut rng = Rng(7);
|
|
|
|
let entries: Vec<MemoryEntry> = data
|
|
.vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| MemoryEntry {
|
|
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
|
embedding: v.clone(),
|
|
source_channel: "bench".into(),
|
|
timestamp: i as f64,
|
|
session_id: format!("s{}", i % 50),
|
|
tags: format!("t{i}"),
|
|
})
|
|
.collect();
|
|
let query_texts: Vec<String> = data
|
|
.query_cluster
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
|
.collect();
|
|
|
|
let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
|
|
let t = Instant::now();
|
|
mem.save_batch(entries).unwrap();
|
|
let ingest = t.elapsed();
|
|
// The very first query builds the vector and keyword indexes from
|
|
// scratch. It happens once per store, not once per session: the checkpoint
|
|
// below saves the vector index, so a later `open()` reloads it.
|
|
let t = Instant::now();
|
|
std::hint::black_box(mem.hybrid_search(&data.queries[1], &query_texts[1], 0.7, 0.3, K));
|
|
let cold_build = t.elapsed();
|
|
|
|
let t = Instant::now();
|
|
mem.flush_wal().unwrap();
|
|
let checkpoint = t.elapsed();
|
|
drop(mem);
|
|
|
|
let t = Instant::now();
|
|
let mut mem = HDF5Memory::open(&path).unwrap();
|
|
let open = t.elapsed();
|
|
|
|
// The first query after open pays for whatever is rebuilt lazily.
|
|
let t = Instant::now();
|
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], &query_texts[0], 0.7, 0.3, K));
|
|
let first_query = t.elapsed();
|
|
|
|
// Fewer steady-state samples at large N: each query is currently O(N).
|
|
let samples_wanted = if n >= 100_000 { 20 } else { N_QUERIES.min(100) };
|
|
let steady = summarize(
|
|
(0..samples_wanted)
|
|
.map(|i| {
|
|
let t = Instant::now();
|
|
std::hint::black_box(mem.hybrid_search(
|
|
&data.queries[i % N_QUERIES],
|
|
&query_texts[i % N_QUERIES],
|
|
0.7,
|
|
0.3,
|
|
K,
|
|
));
|
|
t.elapsed()
|
|
})
|
|
.collect(),
|
|
);
|
|
|
|
println!(
|
|
"| {n} | {:.0} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |",
|
|
millis(ingest),
|
|
millis(cold_build),
|
|
millis(checkpoint),
|
|
millis(open),
|
|
millis(first_query),
|
|
millis(steady.p50),
|
|
millis(steady.p99),
|
|
steady.qps
|
|
);
|
|
json.push(serde_json::json!({
|
|
"bench": "hybrid_search", "n": n,
|
|
"ingest_ms": millis(ingest), "cold_index_build_ms": millis(cold_build),
|
|
"checkpoint_ms": millis(checkpoint),
|
|
"open_ms": millis(open), "first_query_ms": millis(first_query),
|
|
"p50_ms": millis(steady.p50), "p99_ms": millis(steady.p99), "qps": steady.qps,
|
|
}));
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Fusion study: does capping the keyword candidate pool change the ranking?
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// `hybrid_search` min-max normalises each signal over the candidates it is
|
|
/// given. The vector stage supplies a pool of `max(8k, 64)`; the keyword stage
|
|
/// supplies *every* matching record, which is what now dominates query time.
|
|
/// This compares the current fusion with one whose keyword stage is capped to
|
|
/// a pool, reporting how often the final top-k agree and what each costs.
|
|
fn fusion_study(n: usize) {
|
|
use clawhdf5_agent::bm25::BM25Index;
|
|
use clawhdf5_agent::hybrid::merge_vector_keyword;
|
|
|
|
let data = make_dataset(n, 0xE2E ^ n as u64);
|
|
let mut rng = Rng(7);
|
|
let texts: Vec<String> = (0..n)
|
|
.map(|i| text_for(data.cluster_of[i], i, &mut rng))
|
|
.collect();
|
|
let query_texts: Vec<String> = data
|
|
.query_cluster
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
|
.collect();
|
|
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
|
|
let index = HnswIndex::build_with_metric(
|
|
&data.vectors,
|
|
HNSW_M,
|
|
HNSW_EF_CONSTRUCTION,
|
|
DistanceMetric::Cosine,
|
|
);
|
|
|
|
let vec_pool = (K * 8).max(64);
|
|
println!("\n### Fusion study, N = {n} (k = {K}, weights 0.7 / 0.3, vector pool {vec_pool})\n");
|
|
println!(
|
|
"| keyword pool | top-{K} overlap vs full | identical top-{K} | same #1 | keyword+merge µs |"
|
|
);
|
|
println!("|---:|---:|---:|---:|---:|");
|
|
|
|
let fuse = |q: usize, kw_pool: usize| -> (Vec<usize>, Duration) {
|
|
let vec_scores: Vec<(usize, f32)> = index
|
|
.search(&data.queries[q], vec_pool, vec_pool)
|
|
.into_iter()
|
|
.map(|(id, d)| (id, 1.0 - d))
|
|
.collect();
|
|
let t = Instant::now();
|
|
let kw = bm25.search(&query_texts[q], kw_pool);
|
|
let merged = merge_vector_keyword(vec_scores, kw, 0.7, 0.3, K);
|
|
let took = t.elapsed();
|
|
(merged.into_iter().map(|(id, _)| id).collect(), took)
|
|
};
|
|
|
|
let full: Vec<(Vec<usize>, Duration)> = (0..N_QUERIES).map(|q| fuse(q, n)).collect();
|
|
let full_time: Duration = full.iter().map(|f| f.1).sum();
|
|
println!(
|
|
"| all ({n}) | 1.0000 | 100.0% | 100.0% | {:.0} |",
|
|
micros(full_time) / N_QUERIES as f64
|
|
);
|
|
for pool in [vec_pool, vec_pool * 4, 1000] {
|
|
if pool >= n {
|
|
continue;
|
|
}
|
|
let (mut overlap, mut identical, mut same_first) = (0usize, 0usize, 0usize);
|
|
let mut time = Duration::ZERO;
|
|
for (q, (want, _)) in full.iter().enumerate() {
|
|
let (got, took) = fuse(q, pool);
|
|
time += took;
|
|
overlap += got.iter().filter(|id| want.contains(id)).count();
|
|
identical += usize::from(&got == want);
|
|
same_first += usize::from(got.first() == want.first());
|
|
}
|
|
println!(
|
|
"| {pool} | {:.4} | {:.1}% | {:.1}% | {:.0} |",
|
|
overlap as f64 / (K * N_QUERIES) as f64,
|
|
100.0 * identical as f64 / N_QUERIES as f64,
|
|
100.0 * same_first as f64 / N_QUERIES as f64,
|
|
micros(time) / N_QUERIES as f64
|
|
);
|
|
}
|
|
}
|
|
|
|
fn main() {
|
|
let args: Vec<String> = std::env::args().skip(1).collect();
|
|
let full = args.iter().any(|a| a == "--full");
|
|
let ann_only = args.iter().any(|a| a == "--ann-only");
|
|
if args.iter().any(|a| a == "--fusion-study") {
|
|
for &n in if full {
|
|
&[10_000, 100_000][..]
|
|
} else {
|
|
&[10_000][..]
|
|
} {
|
|
fusion_study(n);
|
|
}
|
|
return;
|
|
}
|
|
if args.iter().any(|a| a == "--uniform") {
|
|
UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
|
|
println!("(uniform random data)");
|
|
}
|
|
let json_path = args
|
|
.iter()
|
|
.position(|a| a == "--json")
|
|
.and_then(|i| args.get(i + 1))
|
|
.cloned();
|
|
let sizes: &[usize] = if full {
|
|
&[1_000, 10_000, 100_000]
|
|
} else {
|
|
&[1_000, 10_000]
|
|
};
|
|
|
|
if cfg!(debug_assertions) {
|
|
eprintln!("warning: debug build — numbers are meaningless. Use --release.");
|
|
}
|
|
|
|
let mut json = Vec::new();
|
|
println!("## Search harness");
|
|
// `--e2e-only` skips the index benchmarks, so the end-to-end section runs
|
|
// in a process that has not already spun up a thread pool.
|
|
if !args.iter().any(|a| a == "--e2e-only") {
|
|
for &n in sizes {
|
|
bench_ann(n, &mut json);
|
|
}
|
|
}
|
|
|
|
if ann_only {
|
|
return;
|
|
}
|
|
println!("\n### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n");
|
|
println!(
|
|
"| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |"
|
|
);
|
|
println!("|---:|---:|---:|---:|---:|---:|---:|---:|---:|");
|
|
for &n in sizes {
|
|
bench_end_to_end(n, &mut json);
|
|
}
|
|
|
|
if let Some(path) = json_path {
|
|
std::fs::write(&path, serde_json::to_string_pretty(&json).unwrap()).unwrap();
|
|
eprintln!("wrote {path}");
|
|
}
|
|
}
|