osobhandClaude Opus 5.5
db9af7972c
Makes the README's "cryptographically verifiable memory" true. With HDF5Memory::set_signing_key(key), every checkpoint stores a signed manifest of the store: a SHA-256 per memory record (text, embedding as stored, channel, timestamp, session, tags, deleted flag, activation) in a Merkle tree, plus hashes of the settings (and WAL mark), sessions and knowledge graph. The signature, public key and manifest hashes go in /meta; the per-record hashes in /integrity/record_hashes, so HDF5Memory::verify(path, &public_key) can say which records changed, not just that something did. A forged manifest fails the signature. Decisions, as agreed: - the key is set on the open store and never persisted; - a signed store refuses to checkpoint without its key (MemoryError::SigningKeyRequired); remove_signature() is the deliberate way back to unsigned; - checkpoints only: saves still in the WAL are not covered, and verify reports how many there are. The hashes cover exactly what the file persists, in the form the loader returns it (strings lose trailing NULs; an empty WAL mark is not written), so untouched stores verify across any number of reopen and checkpoint cycles. MemoryError becomes #[non_exhaustive] (it already gains variants in this unreleased version). CLI: keygen (owner-only key file), --signing-key / CLAWHDF5_SIGNING_KEY on writing commands (create signs immediately), verify --public-key (JSON; exit 2 if not valid), `signed` in create/stats output. Tests: reopen/checkpoint cycles with awkward strings (f16 and f32), refusal without the key, wrong and rotated keys, eight kinds of edit each detected and located, a forged manifest, unsigned stores, NULs in text, and an edit made in place with h5py that verify pinpoints. Cost on tank (search_harness --signing-study --full, 3 runs): ~20% of a checkpoint (+9 ms at 10K, +89-112 ms at 100K), verify 18.6 ms / 247 ms, 32 bytes per record in the file. New deps ed25519-dalek, sha2, rand_core: pure Rust, the no-C check passes. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
1151 lines
41 KiB
Rust
1151 lines
41 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
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full
|
|
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --signing-study --full
|
|
//! ```
|
|
|
|
use std::time::{Duration, Instant};
|
|
|
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
|
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
|
|
|
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);
|
|
|
|
/// `--int8`: build the HNSW index over int8-quantised vectors (a quarter of
|
|
/// the memory) instead of f32, to price the recall it costs.
|
|
static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
|
|
|
/// `--f16-first`: in `--float16-study`, run the float16 store first.
|
|
static F16_FIRST: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
|
|
|
/// `--rerank`: re-score the candidate pool against the exact vectors before
|
|
/// taking the top K.
|
|
static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
|
|
|
fn storage() -> Storage {
|
|
if INT8.load(std::sync::atomic::Ordering::Relaxed) {
|
|
Storage::Int8
|
|
} else {
|
|
Storage::Float32
|
|
}
|
|
}
|
|
|
|
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
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Exact cosine distance between unit-length vectors.
|
|
fn exact_dist(a: &[f32], b: &[f32]) -> f32 {
|
|
1.0 - a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>()
|
|
}
|
|
|
|
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(),
|
|
}
|
|
}
|
|
|
|
/// Counts live heap bytes, so a structure's cost can be measured by
|
|
/// difference.
|
|
///
|
|
/// RSS cannot do this from inside one process: freeing a large structure
|
|
/// returns its pages to the allocator's pool rather than to the OS, so
|
|
/// allocating the next one shows no change. Measured that way, a store that
|
|
/// holds the corpus twice and one that holds it once look identical.
|
|
struct CountingAllocator;
|
|
|
|
static LIVE_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);
|
|
|
|
/// High-water mark of [`LIVE_BYTES`] since it was last reset.
|
|
///
|
|
/// Live bytes at a checkpoint cannot see a buffer that was allocated and
|
|
/// freed in between, and that is exactly the shape of a transient copy —
|
|
/// which still has to fit in memory while it exists.
|
|
static PEAK_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);
|
|
|
|
fn note_peak(live: i64) {
|
|
PEAK_BYTES.fetch_max(live, std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
|
|
// SAFETY: every method forwards to the system allocator with the same layout
|
|
// it was given, and only adds bookkeeping around it.
|
|
unsafe impl std::alloc::GlobalAlloc for CountingAllocator {
|
|
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
|
|
let ptr = unsafe { std::alloc::System.alloc(layout) };
|
|
if !ptr.is_null() {
|
|
let live = LIVE_BYTES
|
|
.fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed)
|
|
+ layout.size() as i64;
|
|
note_peak(live);
|
|
}
|
|
ptr
|
|
}
|
|
|
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
|
|
LIVE_BYTES.fetch_sub(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
|
unsafe { std::alloc::System.dealloc(ptr, layout) }
|
|
}
|
|
|
|
unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
|
|
let new_ptr = unsafe { std::alloc::System.realloc(ptr, layout, new_size) };
|
|
if !new_ptr.is_null() {
|
|
let delta = new_size as i64 - layout.size() as i64;
|
|
let live = LIVE_BYTES.fetch_add(delta, std::sync::atomic::Ordering::Relaxed) + delta;
|
|
note_peak(live);
|
|
}
|
|
new_ptr
|
|
}
|
|
}
|
|
|
|
#[global_allocator]
|
|
static ALLOCATOR: CountingAllocator = CountingAllocator;
|
|
|
|
/// Live heap bytes right now.
|
|
fn heap_bytes() -> u64 {
|
|
LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
|
|
}
|
|
|
|
/// Start watching for a new high-water mark from the current live total.
|
|
fn reset_peak() {
|
|
PEAK_BYTES.store(
|
|
LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed),
|
|
std::sync::atomic::Ordering::Relaxed,
|
|
);
|
|
}
|
|
|
|
/// The highest live total seen since [`reset_peak`].
|
|
fn peak_bytes() -> u64 {
|
|
PEAK_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
|
|
}
|
|
|
|
fn mib(bytes: u64) -> f64 {
|
|
bytes as f64 / (1 << 20) as 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(
|
|
&data.vectors,
|
|
HNSW_M,
|
|
HNSW_EF_CONSTRUCTION,
|
|
DistanceMetric::Cosine,
|
|
storage(),
|
|
);
|
|
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}, storage = {:?}\n",
|
|
index.storage()
|
|
);
|
|
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!("|---:|---:|---:|---:|---:|");
|
|
// With a quantised index the distances it returns are approximate, so
|
|
// the candidates are re-scored against the exact vectors the caller
|
|
// already holds (in the agent, the embedding cache) before taking the
|
|
// top K. `--rerank` prices that: it costs one exact distance per
|
|
// candidate and is what decides whether int8 is usable.
|
|
let rerank = RERANK.load(std::sync::atomic::Ordering::Relaxed);
|
|
let pool = if rerank { K * 4 } else { K };
|
|
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 mut got = index.search(q, pool, ef.max(pool));
|
|
if rerank {
|
|
for cand in &mut got {
|
|
cand.1 = exact_dist(&data.vectors[cand.0], q);
|
|
}
|
|
got.select_nth_unstable_by(K - 1, |a, b| a.1.total_cmp(&b.1));
|
|
got.truncate(K);
|
|
}
|
|
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,
|
|
}));
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Signing study: what does an Ed25519-signed checkpoint cost?
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// `--signing-study`: checkpoint time unsigned vs signed, `verify` time, and
|
|
/// the file-size cost of the stored per-record hashes. Default store
|
|
/// settings (float16, int8 index). Medians of five checkpoints / three
|
|
/// verifies.
|
|
fn signing_study(n: usize) {
|
|
use clawhdf5_agent::signing::SigningKey;
|
|
let data = make_dataset(n, 0x516 ^ n as u64);
|
|
let mut rng = Rng(9);
|
|
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 dir = tempfile::tempdir().unwrap();
|
|
let path = dir.path().join("sign.h5");
|
|
let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
|
|
mem.save_batch(entries).unwrap();
|
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));
|
|
|
|
let median = |mut v: Vec<Duration>| {
|
|
v.sort();
|
|
v[v.len() / 2]
|
|
};
|
|
let checkpoint = |mem: &mut HDF5Memory| {
|
|
median(
|
|
(0..5)
|
|
.map(|_| {
|
|
let t = Instant::now();
|
|
mem.flush_wal().unwrap();
|
|
t.elapsed()
|
|
})
|
|
.collect(),
|
|
)
|
|
};
|
|
let unsigned = checkpoint(&mut mem);
|
|
let unsigned_bytes = std::fs::metadata(&path).unwrap().len();
|
|
let key = SigningKey::from_bytes(&[7; 32]);
|
|
mem.set_signing_key(key.clone());
|
|
let signed = checkpoint(&mut mem);
|
|
let signed_bytes = std::fs::metadata(&path).unwrap().len();
|
|
drop(mem);
|
|
let vk = key.verifying_key();
|
|
let verify = median(
|
|
(0..3)
|
|
.map(|_| {
|
|
let t = Instant::now();
|
|
let r = HDF5Memory::verify(&path, &vk).unwrap();
|
|
let d = t.elapsed();
|
|
assert!(r.is_valid());
|
|
d
|
|
})
|
|
.collect(),
|
|
);
|
|
println!(
|
|
"| {n} | {:.1} | {:.1} | {:+.1} | {:.1} | {:+.2} |",
|
|
millis(unsigned),
|
|
millis(signed),
|
|
millis(signed) - millis(unsigned),
|
|
millis(verify),
|
|
(signed_bytes as f64 - unsigned_bytes as f64) / (1024.0 * 1024.0),
|
|
);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Search options study: source filters, re-ranking, confidence rejection
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// `--options-study`: what `HDF5Memory::search`'s options cost and whether a
|
|
/// filtered search finds the right records. Filters keep 50%, 10% or 1% of
|
|
/// the store at random, or two whole clusters away from the query (the case
|
|
/// the index cannot serve, which falls back to an exact scan). Recall is
|
|
/// vector-only against an exact scan of the allowed records; latency is full
|
|
/// hybrid search. Hebbian boosting is off.
|
|
fn options_study(n: usize) {
|
|
use clawhdf5_agent::SearchOptions;
|
|
use clawhdf5_agent::confidence::ConfidenceConfig;
|
|
use clawhdf5_agent::hybrid::Fusion;
|
|
use clawhdf5_agent::reranker::ReRankConfig;
|
|
|
|
let data = make_dataset(n, 0x0B7 ^ n as u64);
|
|
let n_clusters = data.cluster_of.iter().max().map_or(1, |m| m + 1);
|
|
let mut rng = Rng(5);
|
|
let bucket_of: Vec<usize> = (0..n).map(|_| rng.below(100)).collect();
|
|
let bucket = &bucket_of;
|
|
let query_texts: Vec<String> = data
|
|
.query_cluster
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
|
.collect();
|
|
let exact_top = |q: &[f32], allowed: &dyn Fn(usize) -> bool| -> Vec<usize> {
|
|
let mut s: Vec<(usize, f32)> = (0..n)
|
|
.filter(|&i| allowed(i))
|
|
.map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
|
|
.collect();
|
|
s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
|
s.into_iter().take(K).map(|(i, _)| i).collect()
|
|
};
|
|
|
|
// Two stores: channel = random bucket, and channel = cluster.
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let mut stores = Vec::new();
|
|
for by_cluster in [false, true] {
|
|
let mut rng = Rng(3);
|
|
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: if by_cluster {
|
|
format!("c{}", data.cluster_of[i])
|
|
} else {
|
|
format!("b{}", bucket[i])
|
|
},
|
|
timestamp: i as f64,
|
|
session_id: format!("s{}", i % 50),
|
|
tags: format!("t{i}"),
|
|
})
|
|
.collect();
|
|
let mut config = MemoryConfig::new(
|
|
dir.path().join(format!("opt_{by_cluster}.h5")),
|
|
"bench",
|
|
DIM,
|
|
);
|
|
config.hebbian_boost = 0.0;
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
mem.save_batch(entries).unwrap();
|
|
std::hint::black_box(mem.search(&data.queries[0], "", &SearchOptions::new(K)));
|
|
stores.push(mem);
|
|
}
|
|
|
|
let vector_only = SearchOptions::new(K).with_fusion(Fusion::Weighted {
|
|
vector: 1.0,
|
|
keyword: 0.0,
|
|
});
|
|
// (label, store, channels for query i, allowed(i, record))
|
|
type Case<'a> = (
|
|
String,
|
|
usize,
|
|
Box<dyn Fn(usize) -> Option<Vec<String>> + 'a>,
|
|
Box<dyn Fn(usize, usize) -> bool + 'a>,
|
|
);
|
|
let mut cases: Vec<Case> = vec![(
|
|
"no filter".into(),
|
|
0,
|
|
Box::new(|_| None),
|
|
Box::new(|_, _| true),
|
|
)];
|
|
for pct in [50usize, 10, 1] {
|
|
cases.push((
|
|
format!("random {pct}%"),
|
|
0,
|
|
Box::new(move |_| Some((0..pct).map(|b| format!("b{b}")).collect())),
|
|
Box::new(move |_, i| bucket[i] < pct),
|
|
));
|
|
}
|
|
let d = &data;
|
|
let away = move |qi: usize| {
|
|
let qc = d.query_cluster[qi];
|
|
[
|
|
(qc + n_clusters / 3) % n_clusters,
|
|
(qc + 2 * n_clusters / 3) % n_clusters,
|
|
]
|
|
};
|
|
cases.push((
|
|
"2 clusters away from the query".into(),
|
|
1,
|
|
Box::new(move |qi| Some(away(qi).iter().map(|c| format!("c{c}")).collect())),
|
|
Box::new(move |qi, i| away(qi).contains(&d.cluster_of[i])),
|
|
));
|
|
|
|
for (label, store, channels, allowed) in &cases {
|
|
let mem = &mut stores[*store];
|
|
let mut hits = 0;
|
|
let mut kept = 0;
|
|
for (qi, q) in data.queries.iter().enumerate() {
|
|
let mut opts = vector_only.clone();
|
|
opts.source_channels = channels(qi);
|
|
let got = mem.search(q, "", &opts);
|
|
let want = exact_top(q, &|i| allowed(qi, i));
|
|
kept += want.len();
|
|
hits += got.iter().filter(|r| want.contains(&r.index)).count();
|
|
}
|
|
let latency = summarize(
|
|
(0..N_QUERIES)
|
|
.map(|qi| {
|
|
let mut opts = SearchOptions::new(K);
|
|
opts.source_channels = channels(qi);
|
|
let t = Instant::now();
|
|
std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
|
|
t.elapsed()
|
|
})
|
|
.collect(),
|
|
);
|
|
println!(
|
|
"| {n} | {label} | {:.4} | {:.3} | {:.3} |",
|
|
hits as f64 / kept.max(1) as f64,
|
|
millis(latency.p50),
|
|
millis(latency.p99),
|
|
);
|
|
}
|
|
|
|
let mem = &mut stores[0];
|
|
for (label, opts) in [
|
|
(
|
|
"re-rank",
|
|
SearchOptions::new(K).with_rerank(ReRankConfig::default()),
|
|
),
|
|
(
|
|
"re-rank + confidence",
|
|
SearchOptions::new(K)
|
|
.with_rerank(ReRankConfig::default())
|
|
.with_confidence(ConfidenceConfig::default()),
|
|
),
|
|
] {
|
|
let latency = summarize(
|
|
(0..N_QUERIES)
|
|
.map(|qi| {
|
|
let t = Instant::now();
|
|
std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
|
|
t.elapsed()
|
|
})
|
|
.collect(),
|
|
);
|
|
println!(
|
|
"| {n} | {label} | — | {:.3} | {:.3} |",
|
|
millis(latency.p50),
|
|
millis(latency.p99)
|
|
);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// float16 study: what does half-precision embedding storage cost?
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// `--float16-study`: the same data in an `f32` store and a `float16` store.
|
|
/// Reports file size, checkpoint and open time, vector-search recall@10
|
|
/// against an exact scan of the *original* f32 vectors, how often the two
|
|
/// stores return the same top 10, and `hybrid_search` latency. Hebbian
|
|
/// boosting is off, so every query sees the same store.
|
|
fn float16_study(n: usize) {
|
|
let data = make_dataset(n, 0xF16 ^ n as u64);
|
|
let mut rng = Rng(11);
|
|
let query_texts: Vec<String> = data
|
|
.query_cluster
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, c)| text_for(*c, i, &mut rng))
|
|
.collect();
|
|
|
|
// Exact top K by cosine (the vectors are unit length) on the f32 inputs.
|
|
let exact: Vec<Vec<usize>> = data
|
|
.queries
|
|
.iter()
|
|
.map(|q| {
|
|
let mut scored: Vec<(usize, f32)> = data
|
|
.vectors
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, v.iter().zip(q).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.into_iter().take(K).map(|(i, _)| i).collect()
|
|
})
|
|
.collect();
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let mut per_variant: Vec<(bool, Vec<Vec<usize>>)> = Vec::new();
|
|
// `--f16-first` swaps the order, to check the numbers do not depend on
|
|
// which store runs first (page cache, allocator, CPU frequency).
|
|
let order = if F16_FIRST.load(std::sync::atomic::Ordering::Relaxed) {
|
|
[true, false]
|
|
} else {
|
|
[false, true]
|
|
};
|
|
for float16 in order {
|
|
let path = dir.path().join(format!("f16study_{float16}.h5"));
|
|
let mut rng = Rng(3);
|
|
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 mut config = MemoryConfig::new(path.clone(), "bench", DIM);
|
|
config.float16 = float16;
|
|
config.hebbian_boost = 0.0;
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
mem.save_batch(entries).unwrap();
|
|
// Build the indexes, then time a checkpoint that writes everything.
|
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));
|
|
let t = Instant::now();
|
|
mem.flush_wal().unwrap();
|
|
let checkpoint = t.elapsed();
|
|
drop(mem);
|
|
let file_bytes = std::fs::metadata(&path).unwrap().len();
|
|
|
|
// Median of three opens.
|
|
let mut opens: Vec<Duration> = (0..3)
|
|
.map(|_| {
|
|
let t = Instant::now();
|
|
let m = HDF5Memory::open(&path).unwrap();
|
|
let d = t.elapsed();
|
|
drop(m);
|
|
d
|
|
})
|
|
.collect();
|
|
opens.sort();
|
|
let mut mem = HDF5Memory::open(&path).unwrap();
|
|
|
|
// Vector-only search: empty text, all weight on the vector stage.
|
|
let results: Vec<Vec<usize>> = data
|
|
.queries
|
|
.iter()
|
|
.map(|q| {
|
|
mem.hybrid_search(q, "", 1.0, 0.0, K)
|
|
.iter()
|
|
.map(|r| r.index)
|
|
.collect()
|
|
})
|
|
.collect();
|
|
let hits: usize = results
|
|
.iter()
|
|
.zip(&exact)
|
|
.map(|(got, want)| got.iter().filter(|i| want.contains(i)).count())
|
|
.sum();
|
|
let recall = hits as f64 / (K * data.queries.len()) as f64;
|
|
|
|
let latency = summarize(
|
|
(0..N_QUERIES)
|
|
.map(|i| {
|
|
let t = Instant::now();
|
|
std::hint::black_box(mem.hybrid_search(
|
|
&data.queries[i],
|
|
&query_texts[i],
|
|
0.4,
|
|
0.6,
|
|
K,
|
|
));
|
|
t.elapsed()
|
|
})
|
|
.collect(),
|
|
);
|
|
let overlap = match per_variant.first() {
|
|
Some((_, other)) => {
|
|
let same: usize = results
|
|
.iter()
|
|
.zip(other)
|
|
.map(|(a, b)| a.iter().filter(|i| b.contains(i)).count())
|
|
.sum();
|
|
format!("{:.4}", same as f64 / (K * data.queries.len()) as f64)
|
|
}
|
|
None => "—".into(),
|
|
};
|
|
println!(
|
|
"| {n} | {} | {:.1} | {:.0} | {:.1} | {recall:.4} | {overlap} | {:.3} |",
|
|
if float16 { "float16" } else { "f32" },
|
|
mib(file_bytes),
|
|
millis(checkpoint),
|
|
millis(opens[1]),
|
|
millis(latency.p50),
|
|
);
|
|
per_variant.push((float16, results));
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 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(
|
|
&data.vectors,
|
|
HNSW_M,
|
|
HNSW_EF_CONSTRUCTION,
|
|
DistanceMetric::Cosine,
|
|
storage(),
|
|
);
|
|
|
|
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
|
|
);
|
|
}
|
|
}
|
|
|
|
/// What an in-memory store costs, stage by stage. The vectors are the floor:
|
|
/// everything above it is bookkeeping that could in principle be shared.
|
|
fn bench_footprint(n: usize) {
|
|
let data = make_dataset(n, 0xF007 ^ n as u64);
|
|
let mut rng = Rng(11);
|
|
let dir = tempfile::TempDir::new().unwrap();
|
|
let path = dir.path().join("footprint.h5");
|
|
|
|
let base = heap_bytes();
|
|
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 after_entries = heap_bytes();
|
|
|
|
let mut config = MemoryConfig::new(path, "bench", DIM);
|
|
config.quantized_index = INT8.load(std::sync::atomic::Ordering::Relaxed);
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
mem.save_batch(entries).unwrap();
|
|
let after_store = heap_bytes();
|
|
|
|
// First query builds the vector and keyword indexes.
|
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "record", 0.7, 0.3, K));
|
|
let after_indexes = heap_bytes();
|
|
|
|
// Reopening is the figure that matters for a long-lived process, and the
|
|
// only one RSS reports honestly: memory freed when the ingest buffers went
|
|
// away stays in the allocator's pool, so the stage deltas above understate
|
|
// what was given back.
|
|
let path = mem.config().path.clone();
|
|
drop(mem);
|
|
let before_open = heap_bytes();
|
|
reset_peak();
|
|
let reopened = HDF5Memory::open(&path).unwrap();
|
|
let after_open = heap_bytes();
|
|
let loaded = after_open.saturating_sub(before_open);
|
|
// Peak over the open, not just what it leaves behind: a buffer allocated
|
|
// and freed during the parse never shows up in the live total.
|
|
let peak = peak_bytes().saturating_sub(before_open);
|
|
drop(reopened);
|
|
|
|
let raw = (n * DIM * 4) as u64;
|
|
println!(
|
|
"| {n} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.2}x |",
|
|
mib(raw),
|
|
mib(after_entries.saturating_sub(base)),
|
|
mib(after_store.saturating_sub(after_entries)),
|
|
mib(after_indexes.saturating_sub(after_store)),
|
|
mib(loaded),
|
|
mib(peak),
|
|
loaded as f64 / raw 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 == "--signing-study") {
|
|
println!("## Signed checkpoints ({DIM}-dim, float16, int8 index)\n");
|
|
println!(
|
|
"| N | checkpoint ms, unsigned | checkpoint ms, signed | signing adds ms | verify ms | file MiB added |"
|
|
);
|
|
println!("|---:|---:|---:|---:|---:|---:|");
|
|
for &n in if full {
|
|
&[1_000, 10_000, 100_000][..]
|
|
} else {
|
|
&[1_000, 10_000][..]
|
|
} {
|
|
signing_study(n);
|
|
}
|
|
return;
|
|
}
|
|
if args.iter().any(|a| a == "--options-study") {
|
|
println!("## Search options ({DIM}-dim, k = {K}, Hebbian boost off)\n");
|
|
println!("| N | options | filtered recall@10 | p50 ms | p99 ms |");
|
|
println!("|---:|---|---:|---:|---:|");
|
|
for &n in if full {
|
|
&[10_000, 100_000][..]
|
|
} else {
|
|
&[10_000][..]
|
|
} {
|
|
options_study(n);
|
|
}
|
|
return;
|
|
}
|
|
if args.iter().any(|a| a == "--f16-first") {
|
|
F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
if args.iter().any(|a| a == "--float16-study") {
|
|
println!("## float16 embedding storage ({DIM}-dim, int8 index, Hebbian boost off)\n");
|
|
println!(
|
|
"| N | embeddings | file MiB | checkpoint ms | open ms | recall@10 | top-10 overlap with the other | hybrid p50 ms |"
|
|
);
|
|
println!("|---:|---|---:|---:|---:|---:|---:|---:|");
|
|
for &n in if full {
|
|
&[1_000, 10_000, 100_000][..]
|
|
} else {
|
|
&[1_000, 10_000][..]
|
|
} {
|
|
float16_study(n);
|
|
}
|
|
return;
|
|
}
|
|
if args.iter().any(|a| a == "--int8") {
|
|
INT8.store(true, std::sync::atomic::Ordering::Relaxed);
|
|
println!("(int8-quantised index vectors)");
|
|
}
|
|
if args.iter().any(|a| a == "--rerank") {
|
|
RERANK.store(true, std::sync::atomic::Ordering::Relaxed);
|
|
println!("(candidates re-scored against exact vectors)");
|
|
}
|
|
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");
|
|
|
|
if args.iter().any(|a| a == "--footprint") {
|
|
println!("\n### Resident memory, {DIM}-dim f32\n");
|
|
println!(
|
|
"| N | vectors (raw) | entries MiB | store MiB | indexes MiB | reopened MiB | peak during open MiB | reopened / raw |"
|
|
);
|
|
println!("|---:|---:|---:|---:|---:|---:|---:|---:|");
|
|
for &n in sizes {
|
|
bench_footprint(n);
|
|
}
|
|
return;
|
|
}
|
|
// `--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}");
|
|
}
|
|
}
|