diff --git a/.gitignore b/.gitignore index 05c7db3..029ea7a 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ benchmarks/longmemeval/*.json # Local model weights (MiniLM etc.) — large, not committed weights/ +.venv diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 3f4f814..bfea271 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -57,6 +57,52 @@ What remains at 2.43x: the flat vectors (1.0x), the HNSW index's own copy of them (1.0x), and text, ids and graph (~0.4x). The index copy is the next target — it is what a quantised or borrowed representation would address. +### Quantising the index copy (`quantized_index`) + +`MemoryConfig::quantized_index` stores the index's copy as `i8` instead of +`f32`. Same harness, same binary, `--footprint --full` with and without +`--int8`: + +| N | vectors (raw) | indexes, f32 | indexes, int8 | reopened, f32 | reopened, int8 | +|---:|---:|---:|---:|---:|---:| +| 1 000 | 1 MiB | 2 MiB | 1 MiB | 4 MiB (2.40x) | 2 MiB (1.64x) | +| 10 000 | 15 MiB | 32 MiB | 14 MiB | 44 MiB (3.03x) | 27 MiB (1.81x) | +| 100 000 | 146 MiB | 266 MiB | **123 MiB** | 399 MiB (2.72x) | **256 MiB (1.74x)** | + +The scale is **per row**, not global. A unit-length row in `d` dimensions has +components around `1/sqrt(d)`, so a fixed `[-1, 1]` scale spends fewer than 12 +of the 255 levels on a 128-dimensional vector: measured against an exact +ranking that gives 0.35 top-10 overlap — unusable. Scaling each row by its own +largest component brings the same measurement to 0.99. + +Quantised distances still cost recall on their own, and **`ef` does not buy it +back**, because the loss is in the distances rather than in the graph +(`--ann-only --full`, N = 100 000): + +| ef | recall@10, f32 | recall@10, int8 | recall@10, int8 + re-score | +|---:|---:|---:|---:| +| 32 | 0.9775 | 0.9415 | 0.9785 | +| 64 | 0.9945 | 0.9625 | 0.9940 | +| 128 | 0.9995 | 0.9670 | 0.9990 | +| 256 | 0.9995 | 0.9670 (ceiling) | 0.9990 | + +Re-scoring closes the gap: the store already holds the exact embeddings, so +the query path re-scores the candidate pool against them before fusion. That +is done automatically whenever the index is quantised. What it costs is +throughput — about 13% of QPS and 16% of build time at 100 000 x 384. So the +setting trades ~13% of query speed for ~36% of the process's memory at equal +recall. It is **off by default**: the right side of that trade depends on +whether the deployment is short of memory or short of CPU. + +A measurement trap worth recording: the synthetic `clustered` generator in the +`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so +neighbours there sit closer together than the quantisation error and top-10 +*identity* is noise. Scored on that fixture int8 looks catastrophic (0.57 +overlap) — a fact about the fixture, not the storage. The tests use random +vectors, and recall is measured against brute-force ground truth rather than +against the f32 index, whose own approximation errors a re-scored search is +entitled to get right. + ## Read harness Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b28c88..4d6091c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,35 @@ ## Unreleased +### Testing +- The Python interop suites honour **`CLAWHDF5_PYTHON`**, and `ci-test.sh` + picks up a `.venv/bin/python` automatically. On a PEP 668 "externally + managed" system h5py cannot be installed into the system interpreter at all, + so every interop suite — the h5py writer round-trips, the facade, netCDF4 + and the reference files — was skipping silently. A silent skip here is + exactly how the v5 compound-datatype bug reached a release. + `CLAWHDF5_REQUIRE_INTEROP=1` still turns a skip into a failure. + +### Memory +- `clawhdf5-agent`: **`MemoryConfig::quantized_index`** stores the vector + index's own copy of the embeddings as `i8` rather than `f32`, which at 100k + 384-dim entries takes the index from 266 to 123 MiB and the whole reopened + store from 399 to 256 MiB (2.72x -> **1.74x** the raw vectors). Quantised + distances are approximate and `ef` cannot compensate — recall@10 tops out at + 0.967 against f32's 0.9995 — so the query path re-scores the candidate pool + against the exact embeddings the store already holds, which restores recall + (0.9940 vs 0.9945 at ef=64) for about 13% of QPS. **Off by default**: it + trades query speed for memory, and which side is worth more depends on the + deployment. The setting is persisted, so a reopened store does not silently + revert to four times the index memory. +- `clawhdf5-ann`: `Storage::Int8` and the `build_with` / `new_with` / + `from_graph_bytes_with` constructors that select it. The scale is per row, + not global — a fixed `[-1, 1]` scale spends fewer than 12 of the 255 levels + on a unit-length 128-dim vector and is unusable (0.35 top-10 overlap against + an exact ranking, versus 0.99 per row). `compact()` keeps the storage it was + given; serialized indexes still carry f32 vectors, so a quantised index is + rebuilt rather than loaded. + ### Memory - `clawhdf5-agent`: **a loaded store holds ~30% less memory** (100k 384-dim entries: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors). The cache kept diff --git a/CLAUDE.md b/CLAUDE.md index 2dfc918..60bd9a0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -39,7 +39,13 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F (plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its graph is saved to `.h5.ann` at each checkpoint and reloaded by `open()` (tied to the checkpoint by a generation id; stale/damaged sidecars are - ignored and the index rebuilt). `hybrid_search` keeps one incremental BM25 + ignored and the index rebuilt). `MemoryConfig::quantized_index` (off by + default, persisted) stores the index's own copy of the embeddings as `i8`, + which roughly halves a loaded store's memory (2.72x -> 1.74x the raw vectors + at 100K); because quantised distances are approximate and `ef` cannot + compensate, the query path then re-scores the candidate pool against the + exact embeddings, which holds recall at the f32 index's level and costs + ~13% of QPS. `hybrid_search` keeps one incremental BM25 index for the life of the store and never writes the store: Hebbian activation boosts are persisted by the next checkpoint (or on drop), not per query. Measure any search-path change with diff --git a/README.md b/README.md index 5623ae5..b9bed33 100644 --- a/README.md +++ b/README.md @@ -432,6 +432,13 @@ ClawhDF5's agent memory design draws from 15+ recent papers: | `agent` | no | Full agent memory layer | | `float16` | **yes** | Half-precision embedding storage (2× compression) | | `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan | + +`MemoryConfig::quantized_index` (off by default) stores the HNSW index's own +copy of the embeddings as `i8`, roughly halving a loaded store's memory +(2.72x -> 1.74x the raw vectors at 100k x 384). Quantised distances are +approximate, so the query path re-scores the candidate pool against the exact +embeddings the store already holds — recall matches the `f32` index, at about +13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy". | `parallel` | no | Rayon parallel search | | `fast-math` | no | BLAS matrix-vector multiply | | `accelerate` | no | Apple Accelerate / AMX (macOS) | diff --git a/crates/clawhdf5-agent/src/agents_md.rs b/crates/clawhdf5-agent/src/agents_md.rs index ed9a218..58aace4 100644 --- a/crates/clawhdf5-agent/src/agents_md.rs +++ b/crates/clawhdf5-agent/src/agents_md.rs @@ -118,6 +118,7 @@ mod tests { created_at: "2025-01-01T00:00:00Z".to_string(), wal_enabled: false, wal_max_entries: 500, + quantized_index: false, } } diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index 93ef424..3f26c77 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -62,7 +62,7 @@ use std::path::{Path, PathBuf}; use cache::MemoryCache; #[cfg(feature = "hnsw")] -use clawhdf5_ann::{DistanceMetric, HnswIndex}; +use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage}; use ephemeral::{EphemeralConfig, EphemeralStore}; /// HNSW construction parameters used for the agent's vector index. Cosine is the @@ -139,6 +139,17 @@ pub struct MemoryConfig { pub created_at: String, pub wal_enabled: bool, pub wal_max_entries: usize, + /// Store the vector index's own copy of the embeddings as int8 rather than + /// f32, a quarter of the memory. + /// + /// The index's copy is the single largest part of a loaded store's + /// footprint. Quantised distances are approximate, so the candidate pool + /// is re-scored against the cache's exact embeddings before fusion, which + /// restores recall; what it costs is throughput — roughly 13% of queries + /// per second and 16% of build time at 100K x 384. See `BENCHMARKS.md`. + /// + /// Has no effect without the `hnsw` feature. + pub quantized_index: bool, } impl MemoryConfig { @@ -160,6 +171,7 @@ impl MemoryConfig { created_at, wal_enabled: true, wal_max_entries: 500, + quantized_index: false, } } } @@ -444,7 +456,17 @@ impl HDF5Memory { #[cfg(feature = "hnsw")] let loaded_index = if replay_only_appended { - Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint) + Self::load_vector_index( + path, + checkpoint.ann_generation, + &cache, + n_checkpoint, + if config.quantized_index { + Storage::Int8 + } else { + Storage::Float32 + }, + ) } else { None }; @@ -558,6 +580,7 @@ impl HDF5Memory { generation: Option, cache: &MemoryCache, n_checkpoint: usize, + storage: Storage, ) -> Option { let generation = generation?; let bytes = std::fs::read(Self::vector_index_path(store)).ok()?; @@ -568,7 +591,7 @@ impl HDF5Memory { let vectors: Vec> = (0..n_checkpoint) .map(|i| cache.embeddings.get(i).map(<[f32]>::to_vec)) .collect::>()?; - let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?; + let mut index = HnswIndex::from_graph_bytes_with(graph, vectors, storage).ok()?; if index.dimension() != cache.embedding_dim { return None; } @@ -803,6 +826,16 @@ impl HDF5Memory { // the index length drifts from the cache length (covering any mutation path // that doesn't call a hook, e.g. consolidation pushes). + /// How the index should store its copy of the vectors, per the config. + #[cfg(feature = "hnsw")] + fn index_storage(&self) -> Storage { + if self.config.quantized_index { + Storage::Int8 + } else { + Storage::Float32 + } + } + /// Build an HNSW index over the entire cache, re-applying tombstones as /// soft-deletions so node ids stay aligned with cache indices. /// @@ -821,11 +854,12 @@ impl HDF5Memory { // The index owns its vectors, so it needs rows rather than the cache's // flat buffer. This copy is the index's own; the cache keeps one. let rows: Vec> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect(); - let mut index = HnswIndex::build_with_metric( + let mut index = HnswIndex::build_with( &rows, HNSW_M, HNSW_EF_CONSTRUCTION, DistanceMetric::Cosine, + self.index_storage(), ); for (i, &t) in self.cache.tombstones.iter().enumerate() { if t != 0 { diff --git a/crates/clawhdf5-agent/src/schema.rs b/crates/clawhdf5-agent/src/schema.rs index 3128364..e9884e9 100644 --- a/crates/clawhdf5-agent/src/schema.rs +++ b/crates/clawhdf5-agent/src/schema.rs @@ -104,6 +104,10 @@ pub fn build_hdf5_file_with_meta( "wal_max_entries", AttrValue::I64(config.wal_max_entries as i64), ); + meta.set_attr( + "quantized_index", + AttrValue::I64(config.quantized_index.into()), + ); meta.set_attr( "edgehdf5_version", AttrValue::String(ZEROCLAW_VERSION.into()), @@ -484,6 +488,7 @@ pub fn validate_and_load( wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries") .and_then(|v| usize::try_from(v).ok()) .unwrap_or(500), + quantized_index: optional_bool_attr(&attrs, "quantized_index", false), }; // Load /memory group diff --git a/crates/clawhdf5-agent/src/search.rs b/crates/clawhdf5-agent/src/search.rs index 29f53ac..b4feff1 100644 --- a/crates/clawhdf5-agent/src/search.rs +++ b/crates/clawhdf5-agent/src/search.rs @@ -29,10 +29,26 @@ impl HDF5Memory { // Over-fetch so the merge sees a useful vector pool; cosine // distance from the index converts back to similarity (1 - d). let pool = (k * 8).max(64); - let vec_scores: Vec<(usize, f32)> = index - .search(query_embedding, pool, pool) + let candidates = index.search(query_embedding, pool, pool); + // A quantised index returns approximate distances, and no + // amount of `ef` fixes that — the loss is in the distances, + // not the graph. Re-score the pool against the cache's exact + // embeddings, which cost nothing extra to keep: recall then + // matches an f32 index. See `BENCHMARKS.md`. + let exact = index.storage() == clawhdf5_ann::Storage::Int8; + let vec_scores: Vec<(usize, f32)> = candidates .into_iter() - .map(|(id, dist)| (id, 1.0 - dist)) + .map(|(id, dist)| { + let score = if exact { + crate::vector_search::cosine_similarity( + query_embedding, + &self.cache.embeddings[id], + ) + } else { + 1.0 - dist + }; + (id, score) + }) .collect(); // Fusion normalises over every keyword match, so it needs all // the scores — but not ranked. diff --git a/crates/clawhdf5-agent/tests/hnsw_integration.rs b/crates/clawhdf5-agent/tests/hnsw_integration.rs index d853fb7..79f4ffe 100644 --- a/crates/clawhdf5-agent/tests/hnsw_integration.rs +++ b/crates/clawhdf5-agent/tests/hnsw_integration.rs @@ -165,3 +165,72 @@ fn save_batch_then_search_is_consistent() { ); } } + +#[test] +fn quantized_index_matches_the_f32_index_after_re_scoring() { + // A quantised index holds approximate vectors, but the store still has the + // exact ones, so the query path re-scores the candidate pool before + // fusion. The results a caller sees should therefore be the same. + let dim = 64; + let n = 400; + let mut seed = 0x5EED_1234_5678_9ABC; + let vectors: Vec> = (0..n).map(|_| make_vector(&mut seed, dim)).collect(); + let queries: Vec> = (0..20).map(|_| make_vector(&mut seed, dim)).collect(); + + let build = |dir: &TempDir, quantized: bool| { + let mut config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", dim); + config.quantized_index = quantized; + let mut mem = HDF5Memory::create(config).unwrap(); + for (i, v) in vectors.iter().enumerate() { + mem.save(entry(&format!("chunk {i}"), v.clone(), &format!("k{i}"))) + .unwrap(); + } + mem + }; + + let exact_dir = TempDir::new().unwrap(); + let quant_dir = TempDir::new().unwrap(); + let mut exact = build(&exact_dir, false); + let mut quantized = build(&quant_dir, true); + + let k = 10; + let mut agree = 0; + for q in &queries { + let want: Vec = exact + .hybrid_search(q, "", 1.0, 0.0, k) + .iter() + .map(|r| r.index) + .collect(); + agree += quantized + .hybrid_search(q, "", 1.0, 0.0, k) + .iter() + .filter(|r| want.contains(&r.index)) + .count(); + } + let overlap = agree as f64 / (k * queries.len()) as f64; + assert!( + overlap >= 0.95, + "quantised store should match the f32 one: {overlap}" + ); +} + +#[test] +fn quantized_index_setting_survives_a_reopen() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("mem.h5"); + let mut config = MemoryConfig::new(path.clone(), "agent", 8); + config.quantized_index = true; + let mut mem = HDF5Memory::create(config).unwrap(); + let mut seed = 7; + for i in 0..30 { + mem.save(entry(&format!("c{i}"), make_vector(&mut seed, 8), "t")) + .unwrap(); + } + mem.flush_wal().unwrap(); + drop(mem); + + // Reopening must not silently quadruple the index's memory, so the flag + // is part of the stored config rather than a per-session choice. + let reopened = HDF5Memory::open(&path).unwrap(); + assert!(reopened.config().quantized_index); +} diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index eb7124d..a4104e4 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -154,6 +154,237 @@ impl Ord for FarCandidate { } } +/// How the index keeps its copy of the vectors. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum Storage { + /// Exactly as given: `dim * 4` bytes per vector. + #[default] + Float32, + /// Each component scaled to an `i8`: `dim` bytes per vector, a quarter of + /// the space, at some cost in precision. + /// + /// Only meaningful for [`DistanceMetric::Cosine`]: rows are stored + /// unit-length, so a quantised dot product reconstructs the similarity + /// directly. Requesting it for `L2` keeps `Float32`, because an L2 + /// distance cannot be recovered from a dot product alone. + Int8, +} + +/// The index's copy of the vectors, flat and row-major. +#[derive(Debug, Clone)] +enum Vectors { + F32 { + dim: usize, + flat: Vec, + }, + /// `flat[i * dim + j]` is component `j` of vector `i` divided by + /// `scales[i]`; multiplying back recovers it. + /// + /// The scale is per row rather than global. A unit-length row in `d` + /// dimensions has components around `1/sqrt(d)`, so a fixed `[-1, 1]` + /// scale spends fewer than 12 of the 255 levels on a 128-dimensional + /// vector and the reconstruction error swamps the gaps between near + /// neighbours — measured at 0.35 top-10 overlap with the exact ranking. + /// Scaling each row by its own largest component uses the full range. + Int8 { + dim: usize, + flat: Vec, + scales: Vec, + }, +} + +/// Levels either side of zero. 127, not 128, so the range is symmetric. +const INT8_LEVELS: f32 = 127.0; + +/// Quantise one row, returning the codes and the scale that inverts them. +fn quantise_row(v: &[f32], out: &mut Vec) -> f32 { + let max_abs = v.iter().fold(0.0f32, |m, x| m.max(x.abs())); + if max_abs <= f32::MIN_POSITIVE { + out.extend(core::iter::repeat_n(0i8, v.len())); + return 0.0; + } + let inv = INT8_LEVELS / max_abs; + out.extend( + v.iter() + .map(|x| (x * inv).round().clamp(-INT8_LEVELS, INT8_LEVELS) as i8), + ); + max_abs / INT8_LEVELS +} + +impl Vectors { + fn new(dim: usize, storage: Storage, metric: DistanceMetric) -> Self { + match storage { + Storage::Int8 if metric == DistanceMetric::Cosine => Vectors::Int8 { + dim, + flat: Vec::new(), + scales: Vec::new(), + }, + _ => Vectors::F32 { + dim, + flat: Vec::new(), + }, + } + } + + fn dim(&self) -> usize { + match self { + Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim, + } + } + + fn storage(&self) -> Storage { + match self { + Vectors::F32 { .. } => Storage::Float32, + Vectors::Int8 { .. } => Storage::Int8, + } + } + + fn len(&self) -> usize { + let dim = self.dim(); + if dim == 0 { + return 0; + } + match self { + Vectors::F32 { flat, .. } => flat.len() / dim, + Vectors::Int8 { flat, .. } => flat.len() / dim, + } + } + + /// Set the row width, for a store seeded empty by `new`. + fn set_dim(&mut self, new_dim: usize) { + match self { + Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim = new_dim, + } + } + + fn push(&mut self, vector: &[f32]) { + match self { + Vectors::F32 { flat, .. } => flat.extend_from_slice(vector), + Vectors::Int8 { flat, scales, .. } => scales.push(quantise_row(vector, flat)), + } + } + + /// Row `i` as `f32`, for callers that need the values back (serialization, + /// and the f32 fast paths). Quantised rows are reconstructed, so this is + /// lossy in exactly the way the storage is. + fn row(&self, i: usize) -> Vec { + let dim = self.dim(); + let start = i * dim; + match self { + Vectors::F32 { flat, .. } => flat[start..start + dim].to_vec(), + Vectors::Int8 { flat, scales, .. } => flat[start..start + dim] + .iter() + .map(|&q| f32::from(q) * scales[i]) + .collect(), + } + } + + /// Distance between two stored vectors. + fn dist(&self, a: usize, b: usize, metric: DistanceMetric) -> f32 { + let dim = self.dim(); + match self { + Vectors::F32 { flat, .. } => { + let (x, y) = (a * dim, b * dim); + compute_distance(&flat[x..x + dim], &flat[y..y + dim], metric) + } + Vectors::Int8 { flat, scales, .. } => { + let (x, y) = (a * dim, b * dim); + let dot = dot_i8(&flat[x..x + dim], &flat[y..y + dim]); + 1.0 - dot as f32 * scales[a] * scales[b] + } + } + } + + /// Distance from a prepared query to stored vector `i`. + fn dist_query(&self, query: &Query, i: usize, metric: DistanceMetric) -> f32 { + let dim = self.dim(); + let start = i * dim; + match (self, query) { + (Vectors::F32 { flat, .. }, Query::F32(q)) => { + compute_distance(q, &flat[start..start + dim], metric) + } + (Vectors::Int8 { flat, scales, .. }, Query::Int8(q, q_scale)) => { + let dot = dot_i8(q, &flat[start..start + dim]); + 1.0 - dot as f32 * q_scale * scales[i] + } + // Mixed forms cannot occur: `Query` is built from the same storage. + _ => f32::MAX, + } + } + + /// Build a store from prepared rows. + fn from_rows(rows: &[Vec], storage: Storage, metric: DistanceMetric) -> Self { + let dim = rows.first().map_or(0, Vec::len); + let mut out = Vectors::new(dim, storage, metric); + for row in rows { + out.push(row); + } + out + } + + /// Prepare `query` for comparison against this store. + fn query(&self, query: Vec) -> Query { + match self { + Vectors::F32 { .. } => Query::F32(query), + Vectors::Int8 { .. } => { + let mut codes = Vec::with_capacity(query.len()); + let scale = quantise_row(&query, &mut codes); + Query::Int8(codes, scale) + } + } + } +} + +/// What a layer search is measuring distance *to*: an incoming query, or a +/// node already in the index (which is what insertion compares against). +enum Target<'a> { + Query(&'a Query), + Node(usize), +} + +impl Vectors { + fn dist_to(&self, target: &Target<'_>, i: usize, metric: DistanceMetric) -> f32 { + match target { + Target::Query(q) => self.dist_query(q, i, metric), + Target::Node(n) => self.dist(*n, i, metric), + } + } +} + +/// A search query in whichever form the store compares against. +enum Query { + F32(Vec), + /// Codes and the scale that inverts them, as in [`Vectors::Int8`]. + Int8(Vec, f32), +} + +/// Sum of products, widened so it cannot overflow: `dim` terms of at most +/// `127 * 127`, so `i32` suffices for any realistic dimension. +fn dot_i8(a: &[i8], b: &[i8]) -> i32 { + // Four independent accumulators over 32-lane blocks: the widening product + // has to sit in a fixed-length chunk for the vectoriser to see it, and the + // separate accumulators keep it off one dependency chain. + const LANE: usize = 8; + let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>(); + let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>(); + let mut acc = [0i32; 4]; + for (x, y) in a_blocks.iter().zip(b_blocks) { + for (lane, slot) in acc.iter_mut().enumerate() { + let mut sum = 0i32; + for k in 0..LANE { + sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]); + } + *slot += sum; + } + } + let tail: i32 = a_tail + .iter() + .zip(b_tail) + .map(|(&x, &y)| i32::from(x) * i32::from(y)) + .sum(); + acc[0] + acc[1] + acc[2] + acc[3] + tail +} + /// Magic for [`HnswIndex::graph_to_bytes`]. const GRAPH_MAGIC: &[u8; 4] = b"CHG1"; @@ -174,8 +405,8 @@ pub const HNSW_FORMAT_VERSION: i64 = 2; /// HDF5 format. #[derive(Debug, Clone)] pub struct HnswIndex { - /// All vectors in the index. - vectors: Vec>, + /// All vectors in the index, flat and row-major. + vectors: Vectors, /// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs. graph: Vec>>, /// Soft-deletion flags, one per node. Deleted nodes remain in the graph for @@ -214,6 +445,20 @@ impl HnswIndex { m: usize, ef_construction: usize, metric: DistanceMetric, + ) -> Self { + Self::build_with(vectors, m, ef_construction, metric, Storage::default()) + } + + /// Build an index, choosing how the vectors are stored. + /// + /// [`Storage::Int8`] keeps them at a quarter of the size; see its docs for + /// what that costs and when it applies. + pub fn build_with( + vectors: &[Vec], + m: usize, + ef_construction: usize, + metric: DistanceMetric, + storage: Storage, ) -> Self { assert!(!vectors.is_empty(), "cannot build index from empty vectors"); assert!(m >= 2, "m must be at least 2"); @@ -224,8 +469,11 @@ impl HnswIndex { let m_max0 = m * 2; let n = vectors.len(); - let prepared: Vec> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect(); - let vectors: &[Vec] = &prepared; + let mut prepared = Vectors::new(dim, storage, metric); + for v in vectors { + prepared.push(&prepare(v.clone(), metric)); + } + let vectors = &prepared; // Assign levels to all nodes let mut node_levels = Vec::with_capacity(n); @@ -322,9 +570,20 @@ impl HnswIndex { /// point for incremental [`HnswIndex::insert`] and as the result of /// [`HnswIndex::compact`] when every vector has been deleted. pub fn new(m: usize, ef_construction: usize, metric: DistanceMetric) -> Self { + Self::new_with(m, ef_construction, metric, Storage::default()) + } + + /// [`HnswIndex::new`], choosing how the vectors are stored. + pub fn new_with( + m: usize, + ef_construction: usize, + metric: DistanceMetric, + storage: Storage, + ) -> Self { assert!(m >= 2, "m must be at least 2"); Self { - vectors: Vec::new(), + // The dimension is set by the first insert. + vectors: Vectors::new(0, storage, metric), graph: Vec::new(), deleted: Vec::new(), entry_point: 0, @@ -351,7 +610,8 @@ impl HnswIndex { // Seed an empty index. if id == 0 { let node_level = assign_level(0, self.m); - self.vectors.push(vector); + self.vectors.set_dim(vector.len()); + self.vectors.push(&vector); self.deleted.push(false); self.node_levels.push(node_level); self.graph = (0..=node_level).map(|_| vec![Vec::new(); 1]).collect(); @@ -361,12 +621,12 @@ impl HnswIndex { assert_eq!( vector.len(), - self.vectors[0].len(), + self.vectors.dim(), "insert dimension mismatch" ); let node_level = assign_level(id, self.m); - self.vectors.push(vector); + self.vectors.push(&vector); self.deleted.push(false); self.node_levels.push(node_level); @@ -387,7 +647,7 @@ impl HnswIndex { ep = greedy_closest( &self.vectors, &self.graph[layer], - &self.vectors[id], + &Target::Node(id), ep, self.metric, ); @@ -400,7 +660,7 @@ impl HnswIndex { let neighbors = search_layer( &self.vectors, &self.graph[layer], - &self.vectors[id], + &Target::Node(id), ep, self.ef_construction, self.metric, @@ -466,16 +726,25 @@ impl HnswIndex { pub fn compact(&mut self) -> Vec> { let mut mapping = vec![None; self.vectors.len()]; let mut surviving: Vec> = Vec::with_capacity(self.active_len()); - for (old, v) in self.vectors.iter().enumerate() { + for (old, slot) in mapping.iter_mut().enumerate() { if !self.deleted[old] { - mapping[old] = Some(surviving.len()); - surviving.push(v.clone()); + *slot = Some(surviving.len()); + surviving.push(self.vectors.row(old)); } } + // Rebuilding must keep the storage the caller chose; a compaction is + // not the place to silently quadruple the index's memory. + let storage = self.vectors.storage(); *self = if surviving.is_empty() { - Self::new(self.m, self.ef_construction, self.metric) + Self::new_with(self.m, self.ef_construction, self.metric, storage) } else { - Self::build_with_metric(&surviving, self.m, self.ef_construction, self.metric) + Self::build_with( + &surviving, + self.m, + self.ef_construction, + self.metric, + storage, + ) }; mapping } @@ -490,24 +759,22 @@ impl HnswIndex { /// # Returns /// A vector of `(id, distance)` pairs sorted by distance (closest first). pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(usize, f32)> { - if self.vectors.is_empty() { + if self.vectors.len() == 0 { return Vec::new(); } - assert_eq!( - query.len(), - self.vectors[0].len(), - "query dimension mismatch" - ); + assert_eq!(query.len(), self.vectors.dim(), "query dimension mismatch"); let ef = ef.max(k); - let prepared_query = prepare(query.to_vec(), self.metric); - let query = prepared_query.as_slice(); + // Prepared and, for a quantised store, quantised once per search + // rather than once per comparison. + let prepared = self.vectors.query(prepare(query.to_vec(), self.metric)); + let target = Target::Query(&prepared); let mut ep = self.entry_point; let top_layer = self.graph.len().saturating_sub(1); // Greedy search from top layer down to layer 1 for layer in (1..=top_layer).rev() { - ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric); + ep = greedy_closest(&self.vectors, &self.graph[layer], &target, ep, self.metric); } // Search layer 0 for the ef nearest *live* nodes. Deleted nodes are @@ -516,7 +783,7 @@ impl HnswIndex { let candidates = search_layer( &self.vectors, &self.graph[0], - query, + &target, ep, ef, self.metric, @@ -543,14 +810,13 @@ impl HnswIndex { pub fn to_hdf5_bytes(&self) -> Result, FormatError> { let mut fw = FmtWriter::new(); let n = self.vectors.len(); - let dim = if n > 0 { self.vectors[0].len() } else { 0 }; + let dim = self.vectors.dim(); // Flatten vectors into a 1D array for storage - let flat_vectors: Vec = self - .vectors - .iter() - .flat_map(|v| v.iter().copied()) - .collect(); + let mut flat_vectors: Vec = Vec::with_capacity(n * dim); + for i in 0..n { + flat_vectors.extend_from_slice(&self.vectors.row(i)); + } let mut group = fw.create_group("ann"); @@ -709,7 +975,9 @@ impl HnswIndex { }; Ok(Self { - vectors, + // Serialized files carry f32 vectors and no storage tag: a + // quantised index is rebuilt, not loaded. + vectors: Vectors::from_rows(&vectors, Storage::Float32, metric), graph, deleted, entry_point, @@ -773,6 +1041,16 @@ impl HnswIndex { /// `bytes` is validated — a corrupt or mismatched graph is an error, never /// an index that panics or walks out of bounds during a search. pub fn from_graph_bytes(bytes: &[u8], vectors: Vec>) -> Result { + Self::from_graph_bytes_with(bytes, vectors, Storage::default()) + } + + /// As [`from_graph_bytes`](Self::from_graph_bytes), choosing how the + /// rehydrated vectors are stored. + pub fn from_graph_bytes_with( + bytes: &[u8], + vectors: Vec>, + storage: Storage, + ) -> Result { let bad = |what: &str| FormatError::SerializationError(format!("HNSW graph: {what}")); let body_len = bytes .len() @@ -863,7 +1141,14 @@ impl HnswIndex { } Ok(Self { - vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(), + vectors: Vectors::from_rows( + &vectors + .into_iter() + .map(|v| prepare(v, metric)) + .collect::>(), + storage, + metric, + ), graph, deleted, entry_point, @@ -882,16 +1167,17 @@ impl HnswIndex { /// Returns true if the index is empty. pub fn is_empty(&self) -> bool { - self.vectors.is_empty() + self.vectors.len() == 0 + } + + /// How this index stores its copy of the vectors. + pub fn storage(&self) -> Storage { + self.vectors.storage() } /// Returns the dimension of vectors in the index. pub fn dimension(&self) -> usize { - if self.vectors.is_empty() { - 0 - } else { - self.vectors[0].len() - } + self.vectors.dim() } /// Returns the number of layers in the graph. @@ -916,17 +1202,17 @@ impl HnswIndex { /// Greedy search: find the single closest node to `query` starting from `ep`. fn greedy_closest( - vectors: &[Vec], + vectors: &Vectors, layer: &[Vec], - query: &[f32], + target: &Target<'_>, mut ep: usize, metric: DistanceMetric, ) -> usize { - let mut best_dist = compute_distance(query, &vectors[ep], metric); + let mut best_dist = vectors.dist_to(target, ep, metric); loop { let mut changed = false; for &neighbor in &layer[ep] { - let d = compute_distance(query, &vectors[neighbor], metric); + let d = vectors.dist_to(target, neighbor, metric); if d < best_dist { best_dist = d; ep = neighbor; @@ -950,15 +1236,15 @@ fn greedy_closest( /// instead meant a query whose neighbourhood had been deleted got back fewer /// than `k` results, or none, however many live records were nearby. fn search_layer( - vectors: &[Vec], + vectors: &Vectors, layer: &[Vec], - query: &[f32], + target: &Target<'_>, ep: usize, ef: usize, metric: DistanceMetric, skip: Option<&[bool]>, ) -> Vec { - let ep_dist = compute_distance(query, &vectors[ep], metric); + let ep_dist = vectors.dist_to(target, ep, metric); // Min-heap of candidates to explore let mut candidates = BinaryHeap::new(); @@ -980,7 +1266,7 @@ fn search_layer( visited.begin(vectors.len()); visited.insert(ep); search_layer_visit( - vectors, layer, query, ef, metric, skip, visited, candidates, results, + vectors, layer, target, ef, metric, skip, visited, candidates, results, ) }) } @@ -1023,9 +1309,9 @@ thread_local! { #[allow(clippy::too_many_arguments)] fn search_layer_visit( - vectors: &[Vec], + vectors: &Vectors, layer: &[Vec], - query: &[f32], + target: &Target<'_>, ef: usize, metric: DistanceMetric, skip: Option<&[bool]>, @@ -1044,7 +1330,7 @@ fn search_layer_visit( continue; } - let d = compute_distance(query, &vectors[neighbor], metric); + let d = vectors.dist_to(target, neighbor, metric); let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance); if d < furthest_dist || results.len() < ef { @@ -1095,7 +1381,7 @@ fn search_layer_visit( /// remaining slots are then filled with the closest rejected candidates, so a /// node is never left under-connected. fn select_neighbors( - vectors: &[Vec], + vectors: &Vectors, candidates: &[(usize, f32)], max_conn: usize, metric: DistanceMetric, @@ -1111,7 +1397,7 @@ fn select_neighbors( } let diverse = selected .iter() - .all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node); + .all(|&s| vectors.dist(id, s, metric) > dist_to_node); if diverse { selected.push(id); } else { @@ -1136,7 +1422,7 @@ fn batch_len(linked: usize) -> usize { /// layers, found by searching the graph as it currently stands. #[allow(clippy::too_many_arguments)] fn plan_batch( - vectors: &[Vec], + vectors: &Vectors, graph: &[Vec>], node_levels: &[usize], batch: std::ops::Range, @@ -1150,7 +1436,7 @@ fn plan_batch( let mut ep = entry_point; // Phase 1: greedy descent from the top layer down to node_level + 1. for layer in (node_level + 1..=ep_level).rev() { - ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric); + ep = greedy_closest(vectors, &graph[layer], &Target::Node(i), ep, metric); } // Phase 2: search and select on every layer the node lives on. let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1); @@ -1159,7 +1445,7 @@ fn plan_batch( let neighbors = search_layer( vectors, &graph[layer], - &vectors[i], + &Target::Node(i), ep, ef_construction, metric, @@ -1186,7 +1472,7 @@ fn plan_batch( /// Prune every `(layer, node)` neighbour list in `overflowed` back to its /// limit. Each list belongs to a different node, so they are independent. fn prune_overflowed( - vectors: &[Vec], + vectors: &Vectors, graph: &mut [Vec>], overflowed: Vec<(usize, usize)>, (m, m_max0): (usize, usize), @@ -1226,7 +1512,7 @@ const PARALLEL_MIN: usize = 8; /// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16 /// cores; bulk builds batch their pruning instead, see `prune_overflowed`.) fn link_back( - vectors: &[Vec], + vectors: &Vectors, layer: &mut [Vec], new_id: usize, selected: &[usize], @@ -1248,7 +1534,7 @@ fn link_back( /// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`]. fn prune_connections( - vectors: &[Vec], + vectors: &Vectors, neighbors: &mut Vec, node: usize, max_conn: usize, @@ -1259,7 +1545,7 @@ fn prune_connections( } let mut scored: Vec<(usize, f32)> = neighbors .iter() - .map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric))) + .map(|&n| (n, vectors.dist(node, n, metric))) .collect(); scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0))); *neighbors = select_neighbors(vectors, &scored, max_conn, metric); @@ -1519,6 +1805,88 @@ mod tests { assert!(recall >= 0.95, "incremental recall@10 = {recall}"); } + #[test] + fn int8_storage_needs_an_exact_re_score_to_match_f32() { + // Cosine only: rows are unit-length, so a quantised dot product + // reconstructs the similarity directly. + // + // Not the `clustered` generator: its clusters are far tighter than any + // real embedding, so neighbours sit closer together than the + // quantisation error and top-10 identity there is noise — that would + // measure the fixture, not the storage. + let mut vectors = make_random_vectors(3060, 128, 5); + let queries = vectors.split_off(3000); + let f32_index = + HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Float32); + let quantised = + HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Int8); + assert_eq!(quantised.storage(), Storage::Int8); + + // Ground truth, not the f32 index's answers: re-scoring can beat that + // index, and measuring against it would score being right as drift. + let truth: Vec> = queries + .iter() + .map(|q| { + let mut d: Vec<(usize, f32)> = vectors + .iter() + .enumerate() + .map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::Cosine))) + .collect(); + d.sort_by(|a, b| a.1.total_cmp(&b.1)); + d[..10].iter().map(|x| x.0).collect() + }) + .collect(); + + let recall = |got: &dyn Fn(&[f32]) -> Vec| -> f64 { + let mut hits = 0; + for (q, want) in queries.iter().zip(&truth) { + hits += got(q).iter().filter(|id| want.contains(id)).count(); + } + hits as f64 / (10 * queries.len()) as f64 + }; + + let exact_recall = recall(&|q| f32_index.search(q, 10, 64).iter().map(|r| r.0).collect()); + let raw_recall = recall(&|q| quantised.search(q, 10, 64).iter().map(|r| r.0).collect()); + // Quantised distances alone cost recall, and `ef` cannot buy it back: + // the loss is in the distances, not in the graph. + assert!( + raw_recall < exact_recall, + "int8 alone should cost recall: {raw_recall} vs {exact_recall}" + ); + + // Re-scoring a wider candidate pool against the exact vectors — what a + // caller holding them (the agent's embedding cache) does — puts it + // back, because only the *ordering* was approximate. + let rescored_recall = recall(&|q| { + let mut pool: Vec<(usize, f32)> = quantised + .search(q, 40, 64) + .into_iter() + .map(|(id, _)| { + ( + id, + compute_distance(q, &vectors[id], DistanceMetric::Cosine), + ) + }) + .collect(); + pool.sort_by(|a, b| a.1.total_cmp(&b.1)); + pool.truncate(10); + pool.into_iter().map(|p| p.0).collect() + }); + assert!( + rescored_recall >= exact_recall - 0.01, + "int8 + exact re-score should match f32: {rescored_recall} vs {exact_recall} (raw {raw_recall})" + ); + } + + #[test] + fn int8_storage_falls_back_to_f32_for_non_cosine_metrics() { + // L2 distance is not recoverable from a quantised dot product, so the + // store silently stays f32 rather than returning wrong distances. + let vectors = clustered(100, 8, 5, 3); + let index = HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::L2, Storage::Int8); + assert_eq!(index.storage(), Storage::Float32); + } + #[test] fn deletions_near_the_query_do_not_shrink_or_degrade_results() { let mut vectors = clustered(2040, 16, 20, 11); @@ -1650,6 +2018,7 @@ mod tests { vec![1.2, 0.0], // 3 vec![-2.0, 0.0], // 4 ]; + let store = Vectors::from_rows(&vectors, Storage::Float32, DistanceMetric::L2); let scored: Vec<(usize, f32)> = (1..5) .map(|i| { ( @@ -1659,12 +2028,12 @@ mod tests { }) .collect(); assert_eq!( - select_neighbors(&vectors, &scored, 2, DistanceMetric::L2), + select_neighbors(&store, &scored, 2, DistanceMetric::L2), [1, 4] ); // Spare capacity is filled with the closest rejected candidates. assert_eq!( - select_neighbors(&vectors, &scored, 3, DistanceMetric::L2), + select_neighbors(&store, &scored, 3, DistanceMetric::L2), [1, 4, 2] ); } @@ -1790,7 +2159,7 @@ mod tests { // Verify vectors match for i in 0..loaded.len() { - assert_eq!(loaded.vectors[i], index.vectors[i]); + assert_eq!(loaded.vectors.row(i), index.vectors.row(i)); } } diff --git a/crates/clawhdf5-ann/src/lib.rs b/crates/clawhdf5-ann/src/lib.rs index bb5d042..f10470f 100644 --- a/crates/clawhdf5-ann/src/lib.rs +++ b/crates/clawhdf5-ann/src/lib.rs @@ -5,4 +5,4 @@ mod hnsw; -pub use hnsw::{DistanceMetric, HnswIndex}; +pub use hnsw::{DistanceMetric, HnswIndex, Storage}; diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index b8902a8..fabf42c 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -24,7 +24,7 @@ use std::time::{Duration, Instant}; use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry}; -use clawhdf5_ann::{DistanceMetric, HnswIndex}; +use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage}; const DIM: usize = 384; const K: usize = 10; @@ -84,6 +84,22 @@ struct Dataset { /// 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); + +/// `--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) { @@ -169,6 +185,11 @@ fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String { // 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::() +} + fn exact_top_k(vectors: &[Vec], query: &[f32], k: usize) -> Vec { // Vectors are unit length, so cosine order == dot-product order. let mut scored: Vec<(usize, f32)> = vectors @@ -270,11 +291,12 @@ fn bench_ann(n: usize, json: &mut Vec) { .collect(); let started = Instant::now(); - let index = HnswIndex::build_with_metric( + let index = HnswIndex::build_with( &data.vectors, HNSW_M, HNSW_EF_CONSTRUCTION, DistanceMetric::Cosine, + storage(), ); let build = started.elapsed(); @@ -291,7 +313,8 @@ fn bench_ann(n: usize, json: &mut Vec) { ); println!( - "\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n" + "\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", @@ -302,12 +325,26 @@ fn bench_ann(n: usize, json: &mut Vec) { ); 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 got = index.search(q, K, ef); + 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(); } @@ -445,11 +482,12 @@ fn fusion_study(n: usize) { .map(|(i, c)| text_for(*c, i, &mut rng)) .collect(); let bm25 = BM25Index::build(&texts, &vec![0u8; n]); - let index = HnswIndex::build_with_metric( + let index = HnswIndex::build_with( &data.vectors, HNSW_M, HNSW_EF_CONSTRUCTION, DistanceMetric::Cosine, + storage(), ); let vec_pool = (K * 8).max(64); @@ -525,7 +563,9 @@ fn bench_footprint(n: usize) { .collect(); let after_entries = heap_bytes(); - let mut mem = HDF5Memory::create(MemoryConfig::new(path, "bench", DIM)).unwrap(); + 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(); @@ -571,6 +611,14 @@ fn main() { } 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)"); diff --git a/crates/clawhdf5-cli/src/main.rs b/crates/clawhdf5-cli/src/main.rs index 03076a2..53831a0 100644 --- a/crates/clawhdf5-cli/src/main.rs +++ b/crates/clawhdf5-cli/src/main.rs @@ -28,6 +28,10 @@ enum Commands { /// Enable write-ahead log #[arg(long)] wal: bool, + /// Store the vector index's copy of the embeddings as int8, roughly + /// halving a loaded store's memory at about 13% fewer queries/second + #[arg(long)] + quantized_index: bool, }, /// Save a memory entry (reads JSON from stdin or --json) Save { @@ -88,9 +92,15 @@ fn main() { fn run(cli: Cli) -> Result<(), Box> { match cli.command { - Commands::Create { agent_id, dim, wal } => { + Commands::Create { + agent_id, + dim, + wal, + quantized_index, + } => { let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim); config.wal_enabled = wal; + config.quantized_index = quantized_index; let mem = HDF5Memory::create(config)?; let j = serde_json::json!({ "status": "created", @@ -98,6 +108,7 @@ fn run(cli: Cli) -> Result<(), Box> { "agent_id": agent_id, "embedding_dim": dim, "wal_enabled": wal, + "quantized_index": quantized_index, "count": mem.count(), }); println!("{}", serde_json::to_string_pretty(&j)?); diff --git a/crates/clawhdf5-format/src/chunked_write.rs b/crates/clawhdf5-format/src/chunked_write.rs index e6a9e44..6816546 100644 --- a/crates/clawhdf5-format/src/chunked_write.rs +++ b/crates/clawhdf5-format/src/chunked_write.rs @@ -15,7 +15,6 @@ use crate::filter_pipeline::{ FilterDescription, FilterPipeline, }; use crate::filters::compress_chunk; - /// Round a file offset up to the next cache-line boundary. /// /// This ensures chunk data starts at an address that is a multiple of the @@ -928,6 +927,7 @@ pub fn write_selection_to_buffer( #[cfg(test)] mod tests { + use super::*; use crate::chunked_read::read_chunked_data; use crate::data_layout::DataLayout; @@ -1512,9 +1512,20 @@ mod tests { // ---- h5py round-trip tests for chunked writes ---- + /// The Python interpreter to drive interop checks with. + /// + /// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, + /// which on a PEP 668 "externally managed" system is the only place it + /// can be installed. Without it the suite silently skips, and a silent + /// skip here is how a datatype bug once reached a release. + #[cfg(feature = "std")] + fn python() -> String { + std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string()) + } + #[cfg(feature = "std")] fn h5py_available() -> bool { - std::process::Command::new("python3") + std::process::Command::new(python()) .args(["-c", "import h5py"]) .output() .map(|o| o.status.success()) @@ -1526,10 +1537,10 @@ mod tests { if !h5py_available() { panic!("h5py not installed — skipping interop test"); } - let o = std::process::Command::new("python3") + let o = std::process::Command::new(python()) .args(["-c", script]) .output() - .expect("python3"); + .expect("python interpreter"); if !o.status.success() { panic!("h5py: {}", String::from_utf8_lossy(&o.stderr)); } diff --git a/crates/clawhdf5-format/tests/reference_tests.rs b/crates/clawhdf5-format/tests/reference_tests.rs index b1dcf40..8a2bbcb 100644 --- a/crates/clawhdf5-format/tests/reference_tests.rs +++ b/crates/clawhdf5-format/tests/reference_tests.rs @@ -2,6 +2,15 @@ use clawhdf5_format::data_read::{read_object_references, read_region_references}; use clawhdf5_format::datatype::{Datatype, ReferenceType}; +/// The Python interpreter to drive interop checks with. +/// +/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which +/// on a PEP 668 "externally managed" system is the only place it can be +/// installed. Without it the suite silently skips, and a silent skip here is +/// how a datatype bug once reached a release. +fn python() -> String { + std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string()) +} #[test] fn object_ref_single_valid() { @@ -173,7 +182,7 @@ print('ok') "#, path.display() ); - let output = std::process::Command::new("python3") + let output = std::process::Command::new(python()) .args(["-c", &script]) .output(); diff --git a/crates/clawhdf5-format/tests/writer_h5py_tests.rs b/crates/clawhdf5-format/tests/writer_h5py_tests.rs index 58d1660..9707d04 100644 --- a/crates/clawhdf5-format/tests/writer_h5py_tests.rs +++ b/crates/clawhdf5-format/tests/writer_h5py_tests.rs @@ -4,9 +4,18 @@ //! (and vice versa). They require python3 + h5py to be installed. use clawhdf5_format::file_writer::{AttrValue, CompoundTypeBuilder, EnumTypeBuilder, FileWriter}; +/// The Python interpreter to drive interop checks with. +/// +/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which +/// on a PEP 668 "externally managed" system is the only place it can be +/// installed. Without it the suite silently skips, and a silent skip here is +/// how a datatype bug once reached a release. +fn python() -> String { + std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string()) +} fn h5py_available() -> bool { - std::process::Command::new("python3") + std::process::Command::new(python()) .args(["-c", "import h5py"]) .output() .map(|o| o.status.success()) @@ -17,10 +26,10 @@ fn h5py_read(_path: &std::path::Path, script: &str) -> String { if !h5py_available() { panic!("h5py not installed — skipping interop test"); } - let o = std::process::Command::new("python3") + let o = std::process::Command::new(python()) .args(["-c", script]) .output() - .expect("python3"); + .expect("python interpreter"); if !o.status.success() { panic!("h5py: {}", String::from_utf8_lossy(&o.stderr)); } diff --git a/crates/clawhdf5-netcdf4/tests/interop_tests.rs b/crates/clawhdf5-netcdf4/tests/interop_tests.rs index 17c7373..af75136 100644 --- a/crates/clawhdf5-netcdf4/tests/interop_tests.rs +++ b/crates/clawhdf5-netcdf4/tests/interop_tests.rs @@ -9,6 +9,15 @@ use clawhdf5_netcdf4::{AttrValue, NetCDF4File}; // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- +/// The Python interpreter to drive interop checks with. +/// +/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which +/// on a PEP 668 "externally managed" system is the only place it can be +/// installed. Without it the suite silently skips, and a silent skip here is +/// how a datatype bug once reached a release. +fn python() -> String { + std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string()) +} /// When `CLAWHDF5_REQUIRE_INTEROP=1` (set in CI), a missing Python dependency /// is a test failure instead of a silent skip. @@ -17,7 +26,7 @@ fn interop_required() -> bool { } fn netcdf4_python_available() -> bool { - Command::new("python3") + Command::new(python()) .args(["-c", "import netCDF4; print(netCDF4.__version__)"]) .output() .map(|o| o.status.success()) @@ -25,7 +34,7 @@ fn netcdf4_python_available() -> bool { } fn xarray_available() -> bool { - Command::new("python3") + Command::new(python()) .args(["-c", "import xarray; print(xarray.__version__)"]) .output() .map(|o| o.status.success()) @@ -59,7 +68,7 @@ macro_rules! skip_if_no_xarray { } fn run_python(script: &str) { - let output = Command::new("python3") + let output = Command::new(python()) .args(["-c", script]) .output() .expect("failed to run python3"); diff --git a/crates/clawhdf5/tests/h5py_interop_tests.rs b/crates/clawhdf5/tests/h5py_interop_tests.rs index e177b72..c7ab11c 100644 --- a/crates/clawhdf5/tests/h5py_interop_tests.rs +++ b/crates/clawhdf5/tests/h5py_interop_tests.rs @@ -9,6 +9,15 @@ use clawhdf5::{AttrValue, CompoundTypeBuilder, DType, File, FileBuilder}; // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- +/// The Python interpreter to drive interop checks with. +/// +/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which +/// on a PEP 668 "externally managed" system is the only place it can be +/// installed. Without it the suite silently skips, and a silent skip here is +/// how a datatype bug once reached a release. +fn python() -> String { + std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string()) +} /// When `CLAWHDF5_REQUIRE_INTEROP=1` (set in CI), a missing Python dependency /// is a test failure instead of a silent skip. @@ -17,7 +26,7 @@ fn interop_required() -> bool { } fn python_available() -> bool { - Command::new("python3") + Command::new(python()) .args(["-c", "import h5py; print(h5py.__version__)"]) .output() .map(|o| o.status.success()) @@ -39,7 +48,7 @@ macro_rules! skip_if_no_python { /// Run a Python script and panic if it fails. fn run_python(script: &str) { - let output = Command::new("python3") + let output = Command::new(python()) .args(["-c", script]) .output() .expect("failed to run python3"); @@ -52,7 +61,7 @@ fn run_python(script: &str) { /// Run a Python script and return stdout as a trimmed string. fn run_python_output(script: &str) -> String { - let output = Command::new("python3") + let output = Command::new(python()) .args(["-c", script]) .output() .expect("failed to run python3"); diff --git a/docs/QUICKSTART.md b/docs/QUICKSTART.md index e776244..af9faf8 100644 --- a/docs/QUICKSTART.md +++ b/docs/QUICKSTART.md @@ -364,6 +364,11 @@ cargo install --path crates/clawhdf5-cli clawhdf5 --path agent.h5 create --agent-id my-agent --dim 384 --wal ``` +Add `--quantized-index` to store the vector index's copy of the embeddings as +int8. That roughly halves a loaded store's memory at about 13% fewer queries +per second, with recall unchanged — the query path re-scores candidates +against the exact embeddings. The setting is recorded in the file. + Output: ```json { diff --git a/scripts/ci-test.sh b/scripts/ci-test.sh index 03367a0..f04dd09 100755 --- a/scripts/ci-test.sh +++ b/scripts/ci-test.sh @@ -20,6 +20,13 @@ set -uo pipefail SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +# Interop suites drive a Python interpreter. On a PEP 668 "externally managed" +# system h5py can only live in a virtualenv, so pick one up here — before any +# test step, since the non-ignored interop suites read the same variable. +if [ -z "${CLAWHDF5_PYTHON:-}" ] && [ -x "$SCRIPT_DIR/../.venv/bin/python" ]; then + export CLAWHDF5_PYTHON="$SCRIPT_DIR/../.venv/bin/python" +fi PASS=0 FAIL=0 STEPS=() @@ -85,12 +92,18 @@ run_step "cargo test (ann parallel)" cargo test \ # 5. Python interop suites. The h5py writer tests are #[ignore]d so a plain # `cargo test` stays hermetic; run them explicitly here. -if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then +# On a PEP 668 "externally managed" system h5py can only live in a +# virtualenv, so honour CLAWHDF5_PYTHON (and a local .venv) rather than +# skipping — the tests read the same variable. +PYTHON="${CLAWHDF5_PYTHON:-python3}" +if "$PYTHON" -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then run_step "h5py interop (format, ignored tests)" cargo test \ -p clawhdf5-format --test writer_h5py_tests -- --include-ignored else echo "" - echo "==> [h5py interop] SKIPPED: python3 with h5py not available" + echo "==> [h5py interop] SKIPPED: no h5py in $PYTHON" + echo " (set CLAWHDF5_PYTHON=/path/to/venv/bin/python, or create .venv;" + echo " CLAWHDF5_REQUIRE_INTEROP=1 makes this a failure instead)" STEPS+=("SKIP: h5py interop (format, ignored tests)") fi