diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 2701517..0a80afe 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -159,6 +159,23 @@ index removes that. | 10000 | 104 | 31.1 | 10.9 | 2627.1 | 2.09 | 2.11 | 479.5 | | 100000 | 1436 | 684.7 | 278.0 | 36308.1 | 22.90 | 25.46 | 43.5 | +### After: vector index persisted with the checkpoint + +The HNSW graph (not the vectors, which the store already holds) is saved to +`.h5.ann` at each checkpoint and reloaded by `open()`, tied to that +checkpoint by a generation id. The index is now built once per store (the *cold +index build* column — the first query ever), not once per session. First query +after `open()`: **196 → 1.7 ms** (1K), **2627 → 15 ms** (10K), +**36308 → 159 ms** (100K); what remains is the one-off keyword index build. +Batch saves no longer force a full rebuild either: appended records join the +index incrementally. + +| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS | +|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 1000 | 14 | 220 | 6.1 | 1.2 | 1.7 | 0.24 | 0.27 | 4049.9 | +| 10000 | 120 | 2916 | 33.1 | 14.0 | 15.4 | 2.15 | 3.30 | 421.2 | +| 100000 | 1591 | 40515 | 747.3 | 324.7 | 158.9 | 23.07 | 30.42 | 41.3 | + ## Vector Search Latency Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size). diff --git a/CHANGELOG.md b/CHANGELOG.md index e5f16ec..28a74fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,17 @@ immediately; a crash in between forgets only the boosts since the last checkpoint. Activation weights are now capped (16.0) - they grew without bound. +- `clawhdf5-agent`: **the vector index is persisted**, so `open()` no longer + rebuilds it on the first search (first query after open: 2627 -> 15 ms at 10K + records, 36 s -> 159 ms at 100K). The HNSW graph — not the vectors, which the + store already holds — is written to `.h5.ann` at each checkpoint and + tied to it by a generation id in `/meta`; a missing, stale, damaged or + structurally invalid sidecar is ignored and the index rebuilt. Records + replayed from the WAL join the loaded index incrementally; a replayed update + or delete invalidates it. `snapshot()` copies it. Batch saves no longer force + a full index rebuild. +- `clawhdf5-ann`: `HnswIndex::graph_to_bytes` / `from_graph_bytes` — graph-only + serialization (checksummed, every neighbour id and level validated on load). - `clawhdf5-agent`: BM25 results are deterministic (ties break by record id), top-k uses a bounded heap, and the "WAND early termination" that computed a bound and then ignored it is gone. IDF is computed per query. diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index ad624f0..746152e 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -390,8 +390,13 @@ impl HDF5Memory { } else { Some(store_lock::StoreLock::acquire(path)?) }; - let ((config, mut cache, sessions, knowledge), wal_applied) = - storage::read_from_disk_with_mark(path)?; + let ((config, mut cache, sessions, knowledge), checkpoint) = + storage::read_from_disk_with_meta(path)?; + let wal_applied = checkpoint.wal_applied; + let n_checkpoint = cache.len(); + // Set if WAL replay did anything other than append records; the saved + // vector index then no longer describes the first `n_checkpoint` ones. + let mut replay_only_appended = true; // Replay WAL if present let wal_path = path.with_extension("h5.wal"); @@ -408,6 +413,9 @@ impl HDF5Memory { && let Ok(entries) = wal::WalFile::read_entries_for_migration(&wal_path, wal_applied) { + replay_only_appended &= entries + .iter() + .all(|e| e.entry_type == wal::WalEntryType::Save); wal::replay_into_cache(&entries, &mut cache); } None @@ -419,6 +427,9 @@ impl HDF5Memory { // in case the process died between writing the .h5 and // truncating the WAL. let entries = wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)?; + replay_only_appended &= entries + .iter() + .all(|e| e.entry_type == wal::WalEntryType::Save); wal::replay_into_cache(&entries, &mut cache); Some(wal::WalFile::open(&wal_path)?) } else if config.wal_enabled { @@ -427,6 +438,25 @@ impl HDF5Memory { None }; + #[cfg(feature = "hnsw")] + let loaded_index = if replay_only_appended { + Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint) + } else { + None + }; + #[cfg(not(feature = "hnsw"))] + let _ = ( + n_checkpoint, + replay_only_appended, + checkpoint.ann_generation, + ); + #[cfg(feature = "hnsw")] + let synced_len = if loaded_index.is_some() { + cache.len() + } else { + 0 + }; + Ok(Self { config, cache, @@ -435,14 +465,14 @@ impl HDF5Memory { wal, strategy: None, ephemeral: None, - // Existing data is loaded from disk + WAL replay; mark the index - // dirty so it is (re)built from the cache on the first search. + // Reuse the vector index saved with the checkpoint when there is + // one; otherwise mark it dirty so the first search builds it. #[cfg(feature = "hnsw")] - hnsw: None, + hnsw_dirty: loaded_index.is_none(), #[cfg(feature = "hnsw")] - hnsw_dirty: true, + hnsw_synced_len: synced_len, #[cfg(feature = "hnsw")] - hnsw_synced_len: 0, + hnsw: loaded_index, // No on-disk provenance ledger exists yet (see CLAUDE.md), so // there's no historical hash to verify loaded records against — // the store starts empty and is populated as records are @@ -458,6 +488,99 @@ impl HDF5Memory { }) } + /// Where the vector index graph is kept between sessions. + #[cfg_attr(not(feature = "hnsw"), allow(dead_code))] + fn vector_index_path(store: &Path) -> PathBuf { + store.with_extension("h5.ann") + } + + /// Save the vector index graph next to the store, returning the generation + /// id the checkpoint must record for it. Only an index that exactly mirrors + /// the cache is saved; otherwise any stale sidecar is removed and `None` + /// returned, and the next session rebuilds. Failures are not errors — the + /// sidecar is a cache of derived data. + #[cfg(feature = "hnsw")] + fn persist_vector_index(&self) -> Option { + let path = Self::vector_index_path(&self.config.path); + let index = match self.hnsw.as_ref() { + Some(index) + if !self.hnsw_dirty + && self.hnsw_synced_len == self.cache.embeddings.len() + && index.len() == self.cache.embeddings.len() => + { + index + } + _ => { + let _ = std::fs::remove_file(&path); + return None; + } + }; + static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_nanos() as u64); + let generation = nanos + ^ (u64::from(std::process::id()) << 32) + ^ COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + let mut bytes = generation.to_le_bytes().to_vec(); + bytes.extend_from_slice(&index.graph_to_bytes()); + let tmp = path.with_extension("ann.tmp"); + let written = + storage::write_synced(&tmp, &bytes).and_then(|()| storage::rename_synced(&tmp, &path)); + match written { + Ok(()) => Some(generation), + Err(_) => { + let _ = std::fs::remove_file(&tmp); + let _ = std::fs::remove_file(&path); + None + } + } + } + + #[cfg(not(feature = "hnsw"))] + fn persist_vector_index(&self) -> Option { + None + } + + /// Load the vector index saved with the checkpoint identified by + /// `generation`, covering the first `n_checkpoint` records of `cache`. + /// Anything unexpected — no sidecar, another generation, a damaged or + /// mismatched graph — yields `None` and the index is rebuilt on demand. + #[cfg(feature = "hnsw")] + fn load_vector_index( + store: &Path, + generation: Option, + cache: &MemoryCache, + n_checkpoint: usize, + ) -> Option { + let generation = generation?; + let bytes = std::fs::read(Self::vector_index_path(store)).ok()?; + let (stamp, graph) = bytes.split_at_checked(8)?; + if u64::from_le_bytes(stamp.try_into().ok()?) != generation { + return None; + } + let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec(); + let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?; + if index.dimension() != cache.embedding_dim { + return None; + } + // Records appended since (replayed from the WAL) join incrementally. + for id in n_checkpoint..cache.embeddings.len() { + if cache.embeddings[id].len() != index.dimension() + || index.insert(cache.embeddings[id].clone()) != id + { + return None; + } + } + for (id, &t) in cache.tombstones.iter().enumerate() { + if t != 0 { + index.mark_deleted(id); + } + } + Some(index) + } + /// Bring the keyword index in line with the cache and return it. /// /// Appends need no hook: records the index hasn't seen yet (whatever path @@ -517,13 +640,19 @@ impl HDF5Memory { // Record which WAL prefix this checkpoint contains, so a crash before // the truncate below can't replay those entries a second time. let wal_applied = self.wal.as_ref().map(|w| w.mark()); - storage::write_to_disk_with_mark( + // Written before the .h5 so a crash in between leaves a sidecar whose + // generation matches no checkpoint (ignored), never the reverse. + let ann_generation = self.persist_vector_index(); + storage::write_to_disk_with_meta( &self.config.path, &self.config, &self.cache, &self.sessions, &self.knowledge, - wal_applied, + &schema::CheckpointMeta { + wal_applied, + ann_generation, + }, )?; if let Some(ref mut w) = self.wal { w.truncate()?; @@ -686,6 +815,30 @@ impl HDF5Memory { #[cfg(feature = "hnsw")] fn ensure_hnsw_fresh(&mut self) { let n = self.cache.embeddings.len(); + // Records appended since the index was last in sync (a batch save, or + // any path that pushes to the cache without a hook) are inserted + // incrementally rather than triggering a rebuild of the whole graph. + if !self.hnsw_dirty + && self.hnsw_synced_len < n + && let Some(index) = self.hnsw.as_mut() + && index.len() == self.hnsw_synced_len + { + let dim = index.dimension(); + let appended = (self.hnsw_synced_len..n).all(|id| { + self.cache.embeddings[id].len() == dim + && index.insert(self.cache.embeddings[id].clone()) == id + }); + if appended { + for id in self.hnsw_synced_len..n { + if self.cache.tombstones[id] != 0 { + index.mark_deleted(id); + } + } + self.hnsw_synced_len = n; + } else { + self.hnsw_dirty = true; + } + } if self.hnsw_dirty || self.hnsw_synced_len != n { self.hnsw = self.build_hnsw_from_cache(); self.hnsw_synced_len = n; @@ -928,8 +1081,8 @@ impl AgentMemory for HDF5Memory { ); indices.push(idx); } - // Batch inserts rebuild the index once rather than node-by-node. - self.hnsw_mark_dirty(); + // The vector and keyword indexes pick the new records up + // incrementally the next time they are needed. self.flush()?; Ok(indices) } @@ -987,6 +1140,13 @@ impl AgentMemory for HDF5Memory { if self.wal.as_ref().is_some_and(|w| !w.is_empty()) && wal_path.exists() { storage::snapshot_file(&wal_path, &snapshot.with_extension("h5.wal"))?; } + // The saved vector index belongs to the checkpoint just copied (its + // generation id is in that .h5), so it is valid for the snapshot too. + // Best effort: without it the snapshot simply rebuilds on first search. + let ann_path = Self::vector_index_path(&self.config.path); + if ann_path.exists() { + let _ = storage::snapshot_file(&ann_path, &Self::vector_index_path(&snapshot)); + } Ok(snapshot) } @@ -1674,6 +1834,126 @@ mod tests { assert_eq!(restored.cache.chunks, ["checkpointed", "wal-only"]); } + /// A store with `n` records spread over a few directions, WAL on. + #[cfg(feature = "hnsw")] + fn indexed_store(dir: &TempDir, n: usize) -> (HDF5Memory, PathBuf) { + let mut config = make_config(dir); + config.wal_enabled = true; + config.wal_max_entries = 10_000; + config.compact_threshold = 0.0; + let path = config.path.clone(); + let mut mem = HDF5Memory::create(config).unwrap(); + for i in 0..n { + let a = i as f32 * 0.37; + mem.save(make_entry( + &format!("rec{i}"), + &[a.cos(), a.sin(), (a * 0.5).cos(), 0.1], + )) + .unwrap(); + } + (mem, path) + } + + #[cfg(feature = "hnsw")] + fn top_ids(mem: &mut HDF5Memory, q: &[f32]) -> Vec { + mem.hybrid_search(q, "", 1.0, 0.0, 5) + .into_iter() + .map(|r| r.index) + .collect() + } + + #[cfg(feature = "hnsw")] + #[test] + fn vector_index_is_reloaded_not_rebuilt() { + let dir = TempDir::new().unwrap(); + let (mut mem, path) = indexed_store(&dir, 60); + let q = [0.3f32.cos(), 0.3f32.sin(), 0.9, 0.1]; + let expected = top_ids(&mut mem, &q); // builds the index + mem.flush_wal().unwrap(); // checkpoint + sidecar + drop(mem); + assert!(HDF5Memory::vector_index_path(&path).exists()); + + let mut reopened = HDF5Memory::open(&path).unwrap(); + assert!(!reopened.hnsw_dirty, "index should come from the sidecar"); + assert_eq!(reopened.hnsw.as_ref().unwrap().len(), 60); + assert_eq!(top_ids(&mut reopened, &q), expected); + } + + #[cfg(feature = "hnsw")] + #[test] + fn records_appended_after_the_checkpoint_join_the_loaded_index() { + let dir = TempDir::new().unwrap(); + let (mut mem, path) = indexed_store(&dir, 40); + top_ids(&mut mem, &[1.0, 0.0, 0.0, 0.0]); + mem.flush_wal().unwrap(); + // Only in the WAL when the process "dies". + mem.save(make_entry("late", &[0.0, 0.0, 0.0, 1.0])).unwrap(); + drop(mem); + + let mut reopened = HDF5Memory::open(&path).unwrap(); + assert!(!reopened.hnsw_dirty); + assert_eq!(reopened.hnsw.as_ref().unwrap().len(), 41); + assert_eq!(top_ids(&mut reopened, &[0.0, 0.0, 0.0, 1.0])[0], 40); + } + + #[cfg(feature = "hnsw")] + #[test] + fn replayed_update_invalidates_the_saved_index() { + let dir = TempDir::new().unwrap(); + let (mut mem, path) = indexed_store(&dir, 40); + top_ids(&mut mem, &[1.0, 0.0, 0.0, 0.0]); + mem.flush_wal().unwrap(); + // An in-place update after the checkpoint changes record 0's vector; + // the saved graph was built over the old one. + let mut moved = make_entry("rec0 moved", &[0.0, 0.0, 0.0, 1.0]); + moved.tags = mem.cache.tags[0].clone(); + mem.save_or_update(moved).unwrap(); + let expected = top_ids(&mut mem, &[0.0, 0.0, 0.0, 1.0]); + std::mem::forget(mem); // die without the drop-time checkpoint + + let mut reopened = HDF5Memory::open_read_only(&path).unwrap(); + assert!(reopened.hnsw_dirty, "saved index must not be reused"); + assert_eq!(top_ids(&mut reopened, &[0.0, 0.0, 0.0, 1.0]), expected); + } + + #[cfg(feature = "hnsw")] + #[test] + fn stale_or_damaged_index_sidecar_is_ignored() { + let dir = TempDir::new().unwrap(); + let (mut mem, path) = indexed_store(&dir, 40); + let q = [1.0, 0.0, 0.0, 0.0]; + top_ids(&mut mem, &q); // builds the index + mem.flush_wal().unwrap(); + let ann = HDF5Memory::vector_index_path(&path); + let first_sidecar = std::fs::read(&ann).unwrap(); + // A second checkpoint gets a new generation. + mem.save(make_entry("more", &[0.5, 0.5, 0.0, 0.0])).unwrap(); + top_ids(&mut mem, &q); + mem.flush_wal().unwrap(); + let expected_after = top_ids(&mut mem, &q); + drop(mem); + + // Sidecar from the earlier checkpoint: wrong generation. + std::fs::write(&ann, &first_sidecar).unwrap(); + let mut reopened = HDF5Memory::open_read_only(&path).unwrap(); + assert!(reopened.hnsw_dirty); + assert_eq!(top_ids(&mut reopened, &q), expected_after); + drop(reopened); + + // Right generation, damaged graph. + let mut mem = HDF5Memory::open(&path).unwrap(); + top_ids(&mut mem, &q); + mem.flush_wal().unwrap(); + drop(mem); + let mut bytes = std::fs::read(&ann).unwrap(); + let mid = bytes.len() / 2; + bytes[mid] ^= 0xFF; + std::fs::write(&ann, &bytes).unwrap(); + let mut reopened = HDF5Memory::open_read_only(&path).unwrap(); + assert!(reopened.hnsw_dirty); + assert_eq!(top_ids(&mut reopened, &q), expected_after); + } + #[test] fn keyword_index_stays_in_sync_through_every_mutation() { let dir = TempDir::new().unwrap(); diff --git a/crates/clawhdf5-agent/src/schema.rs b/crates/clawhdf5-agent/src/schema.rs index 5376d13..c7cb92a 100644 --- a/crates/clawhdf5-agent/src/schema.rs +++ b/crates/clawhdf5-agent/src/schema.rs @@ -22,6 +22,7 @@ pub const ZEROCLAW_VERSION: &str = "0.8.0"; /// the checkpoint was taken with an empty WAL. const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len"; const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc"; +const ANN_GENERATION_ATTR: &str = "ann_generation"; /// Build a complete HDF5 file from the in-memory state. pub fn build_hdf5_file( @@ -43,6 +44,34 @@ pub fn build_hdf5_file_with_mark( knowledge: &KnowledgeCache, wal_applied: Option, ) -> Result, MemoryError> { + let meta = CheckpointMeta { + wal_applied, + ann_generation: None, + }; + build_hdf5_file_with_meta(config, cache, sessions, knowledge, &meta) +} + +/// Bookkeeping a checkpoint records in `/meta` beside the store's contents. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CheckpointMeta { + /// The WAL prefix this checkpoint already contains; see [`WalMark`]. + pub wal_applied: Option, + /// Identifies the vector-index sidecar (`.h5.ann`) written with this + /// checkpoint. A sidecar is loaded only if it carries the same value, so + /// one left over from another checkpoint can never be attached to records + /// it wasn't built from. + pub ann_generation: Option, +} + +/// [`build_hdf5_file`] with checkpoint bookkeeping. +pub fn build_hdf5_file_with_meta( + config: &MemoryConfig, + cache: &MemoryCache, + sessions: &SessionCache, + knowledge: &KnowledgeCache, + checkpoint: &CheckpointMeta, +) -> Result, MemoryError> { + let wal_applied = checkpoint.wal_applied; let mut builder = clawhdf5::FileBuilder::new(); // /meta group with schema attributes @@ -83,6 +112,11 @@ pub fn build_hdf5_file_with_mark( meta.set_attr(WAL_APPLIED_LEN_ATTR, AttrValue::I64(mark.len as i64)); meta.set_attr(WAL_APPLIED_CRC_ATTR, AttrValue::I64(i64::from(mark.crc))); } + if let Some(generation) = checkpoint.ann_generation { + // Stored as the i64 with the same bits; attributes have no u64 scalar + // round trip through every reader. + meta.set_attr(ANN_GENERATION_ATTR, AttrValue::I64(generation as i64)); + } // Need at least one dataset in the group for it to be a proper group meta.create_dataset("_marker").with_u8_data(&[1]).compact(); let finished_meta = meta.finish(); @@ -386,6 +420,22 @@ pub fn read_wal_mark(file: &clawhdf5::File) -> Option { Some(WalMark { len, crc }) } +/// Read the checkpoint bookkeeping from `/meta`. +pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta { + let ann_generation = file + .group("meta") + .ok() + .and_then(|g| g.attrs().ok()) + .and_then(|attrs| match attrs.get(ANN_GENERATION_ATTR) { + Some(AttrValue::I64(v)) => Some(*v as u64), + _ => None, + }); + CheckpointMeta { + wal_applied: read_wal_mark(file), + ann_generation, + } +} + pub fn validate_and_load( file: &clawhdf5::File, ) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> { diff --git a/crates/clawhdf5-agent/src/storage.rs b/crates/clawhdf5-agent/src/storage.rs index c9000a1..533bad3 100644 --- a/crates/clawhdf5-agent/src/storage.rs +++ b/crates/clawhdf5-agent/src/storage.rs @@ -34,7 +34,23 @@ pub fn write_to_disk_with_mark( knowledge: &KnowledgeCache, wal_applied: Option, ) -> Result<(), MemoryError> { - let bytes = schema::build_hdf5_file_with_mark(config, cache, sessions, knowledge, wal_applied)?; + let meta = schema::CheckpointMeta { + wal_applied, + ann_generation: None, + }; + write_to_disk_with_meta(path, config, cache, sessions, knowledge, &meta) +} + +/// [`write_to_disk`] with full checkpoint bookkeeping. +pub fn write_to_disk_with_meta( + path: &Path, + config: &MemoryConfig, + cache: &MemoryCache, + sessions: &SessionCache, + knowledge: &KnowledgeCache, + checkpoint: &schema::CheckpointMeta, +) -> Result<(), MemoryError> { + let bytes = schema::build_hdf5_file_with_meta(config, cache, sessions, knowledge, checkpoint)?; if bytes.is_empty() { return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into())); @@ -47,7 +63,7 @@ pub fn write_to_disk_with_mark( } /// Write `bytes` to `path` and flush them to stable storage. -fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> { +pub(crate) fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> { use std::io::Write; let mut f = std::fs::File::create(path).map_err(MemoryError::Io)?; f.write_all(bytes).map_err(MemoryError::Io)?; @@ -62,7 +78,7 @@ fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> { /// This is per-checkpoint/snapshot cost only (each is already a full file /// write). Individual WAL appends are deliberately not synced — see the /// durability notes in the crate docs. -fn rename_synced(from: &Path, to: &Path) -> Result<(), MemoryError> { +pub(crate) fn rename_synced(from: &Path, to: &Path) -> Result<(), MemoryError> { std::fs::rename(from, to).map_err(MemoryError::Io)?; #[cfg(unix)] if let Some(dir) = to.parent() { @@ -113,6 +129,20 @@ pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option Result<(StoreState, schema::CheckpointMeta), MemoryError> { + let mmap = clawhdf5_io::MmapReader::open(path).map_err(MemoryError::Io)?; + mmap.advise_willneed(0, mmap.len()); + let file = clawhdf5::File::from_bytes(mmap.as_bytes().to_vec()) + .map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?; + let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?; + config.path = path.to_path_buf(); + let meta = schema::read_checkpoint_meta(&file); + Ok(((config, cache, sessions, knowledge), meta)) +} + /// Copy an HDF5 file atomically to a destination. pub fn snapshot_file(src: &Path, dest: &Path) -> Result { let dest_file = if dest.is_dir() { diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index 5755f2b..2644add 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -134,6 +134,9 @@ impl Ord for FarCandidate { } } +/// Magic for [`HnswIndex::graph_to_bytes`]. +const GRAPH_MAGIC: &[u8; 4] = b"CHG1"; + /// On-disk format version for the serialized HNSW index. /// /// - Version 1: original layout (`vectors`, `graph_layer_*`, `config`), no @@ -690,6 +693,160 @@ impl HnswIndex { }) } + /// Serialize the **graph only** — levels, tombstones and adjacency, not the + /// vectors — for a caller that already stores the vectors elsewhere (the + /// agent's record cache). [`HnswIndex::to_hdf5_bytes`] writes a complete, + /// self-contained index including a full copy of every vector, which would + /// double such a store's size. Reattach with + /// [`HnswIndex::from_graph_bytes`]. + /// + /// Layout (little endian): magic `CHG1`, then u32 fields `n`, `m`, + /// `m_max0`, `ef_construction`, `entry_point`, `num_layers`, `metric`; + /// `n` level bytes; `n` tombstone bytes; per layer, per node that exists on + /// that layer: u32 neighbour count + u32 ids; trailing CRC32 of all of it. + pub fn graph_to_bytes(&self) -> Vec { + let n = self.vectors.len(); + let mut out = Vec::with_capacity(32 + n * 2 + n * self.m_max0 * 4); + out.extend_from_slice(GRAPH_MAGIC); + for field in [ + n, + self.m, + self.m_max0, + self.ef_construction, + self.entry_point, + self.graph.len(), + match self.metric { + DistanceMetric::L2 => 0, + DistanceMetric::Cosine => 1, + }, + ] { + out.extend_from_slice(&(field as u32).to_le_bytes()); + } + out.extend(self.node_levels.iter().map(|&l| l.min(255) as u8)); + out.extend(self.deleted.iter().map(|&d| u8::from(d))); + for (layer, adjacency) in self.graph.iter().enumerate() { + for (node, neighbors) in adjacency.iter().enumerate() { + if self.node_levels[node] < layer { + continue; // node does not exist on this layer + } + out.extend_from_slice(&(neighbors.len() as u32).to_le_bytes()); + for &id in neighbors { + out.extend_from_slice(&(id as u32).to_le_bytes()); + } + } + } + let crc = clawhdf5_format::checksum::crc32(&out); + out.extend_from_slice(&crc.to_le_bytes()); + out + } + + /// Rebuild an index from [`HnswIndex::graph_to_bytes`] output and the + /// vectors it was built over (same order). Every structural claim in + /// `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 { + let bad = |what: &str| FormatError::SerializationError(format!("HNSW graph: {what}")); + let body_len = bytes + .len() + .checked_sub(4) + .filter(|&l| l >= GRAPH_MAGIC.len() + 7 * 4) + .ok_or_else(|| bad("truncated"))?; + let (body, crc_bytes) = bytes.split_at(body_len); + if &body[..4] != GRAPH_MAGIC { + return Err(bad("bad magic")); + } + let stored_crc = + u32::from_le_bytes([crc_bytes[0], crc_bytes[1], crc_bytes[2], crc_bytes[3]]); + if clawhdf5_format::checksum::crc32(body) != stored_crc { + return Err(bad("checksum mismatch")); + } + + let mut pos = 4; + let next_u32 = |pos: &mut usize| -> Result { + let b = body.get(*pos..*pos + 4).ok_or_else(|| bad("truncated"))?; + *pos += 4; + Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize) + }; + let n = next_u32(&mut pos)?; + let m = next_u32(&mut pos)?; + let m_max0 = next_u32(&mut pos)?; + let ef_construction = next_u32(&mut pos)?; + let entry_point = next_u32(&mut pos)?; + let num_layers = next_u32(&mut pos)?; + let metric = match next_u32(&mut pos)? { + 0 => DistanceMetric::L2, + 1 => DistanceMetric::Cosine, + _ => return Err(bad("unknown metric")), + }; + if n != vectors.len() { + return Err(bad("vector count does not match the graph")); + } + if n == 0 || entry_point >= n || m < 2 || num_layers == 0 || num_layers > 256 { + return Err(bad("invalid header")); + } + let dim = vectors[0].len(); + if vectors.iter().any(|v| v.len() != dim) { + return Err(bad("vectors have mixed dimensions")); + } + + let levels = body.get(pos..pos + n).ok_or_else(|| bad("truncated"))?; + pos += n; + let node_levels: Vec = levels.iter().map(|&l| l as usize).collect(); + if node_levels.iter().any(|&l| l >= num_layers) + || node_levels[entry_point] + 1 != num_layers + { + return Err(bad("levels inconsistent with layer count")); + } + let deleted: Vec = body + .get(pos..pos + n) + .ok_or_else(|| bad("truncated"))? + .iter() + .map(|&d| d != 0) + .collect(); + pos += n; + + let mut graph: Vec>> = Vec::with_capacity(num_layers); + for layer in 0..num_layers { + let max_conn = if layer == 0 { m_max0 } else { m }; + let mut adjacency = vec![Vec::new(); n]; + for (node, slot) in adjacency.iter_mut().enumerate() { + if node_levels[node] < layer { + continue; + } + let count = next_u32(&mut pos)?; + if count > max_conn { + return Err(bad("neighbour list exceeds the connection limit")); + } + let mut neighbors = Vec::with_capacity(count); + for _ in 0..count { + let id = next_u32(&mut pos)?; + // A neighbour must exist, and exist on this layer. + if id >= n || node_levels[id] < layer { + return Err(bad("neighbour id out of range for its layer")); + } + neighbors.push(id); + } + *slot = neighbors; + } + graph.push(adjacency); + } + if pos != body.len() { + return Err(bad("trailing bytes")); + } + + Ok(Self { + vectors, + graph, + deleted, + entry_point, + m, + m_max0, + ef_construction, + node_levels, + metric, + }) + } + /// Returns the number of vectors in the index. pub fn len(&self) -> usize { self.vectors.len() @@ -1147,6 +1304,73 @@ mod tests { assert!(recall >= 0.95, "incremental recall@10 = {recall}"); } + #[test] + fn graph_bytes_round_trip_gives_identical_searches() { + let mut vectors = clustered(1260, 16, 12, 9); + let queries = vectors.split_off(1200); + let mut index = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2); + index.mark_deleted(3); + index.mark_deleted(700); + + let bytes = index.graph_to_bytes(); + // The graph is a small fraction of the vectors it indexes... not + // necessarily at dim 16, but it must not embed them. + assert!(bytes.len() < 1200 * (16 * 2 + 2) * 4); + let restored = HnswIndex::from_graph_bytes(&bytes, vectors.clone()).unwrap(); + assert_eq!(restored.deleted_count(), 2); + for q in &queries { + assert_eq!(restored.search(q, 10, 50), index.search(q, 10, 50)); + } + // A restored index keeps working incrementally. + let mut restored = restored; + let id = restored.insert(queries[0].clone()); + assert_eq!(restored.search(&queries[0], 1, 50)[0].0, id); + } + + #[test] + fn damaged_or_mismatched_graph_bytes_are_errors() { + let vectors = clustered(300, 8, 6, 4); + let index = HnswIndex::build_with_metric(&vectors, 6, 30, DistanceMetric::Cosine); + let bytes = index.graph_to_bytes(); + + // Wrong vector set. + assert!(HnswIndex::from_graph_bytes(&bytes, vectors[..299].to_vec()).is_err()); + // Every truncation. + for len in 0..bytes.len() { + assert!( + HnswIndex::from_graph_bytes(&bytes[..len], vectors.clone()).is_err(), + "truncated to {len}" + ); + } + // A flipped bit anywhere. + for i in (0..bytes.len()).step_by(7) { + let mut damaged = bytes.clone(); + damaged[i] ^= 0x10; + assert!( + HnswIndex::from_graph_bytes(&damaged, vectors.clone()).is_err(), + "bit flip at {i}" + ); + } + } + + #[test] + fn structurally_invalid_graph_with_a_valid_checksum_is_rejected() { + // The CRC only proves the bytes are what was written; a hostile or + // buggy writer can checksum nonsense. Out-of-range neighbour ids must + // still be caught, or search would index out of bounds. + let vectors = clustered(50, 4, 3, 5); + let index = HnswIndex::build_with_metric(&vectors, 4, 20, DistanceMetric::L2); + let mut bytes = index.graph_to_bytes(); + let body_len = bytes.len() - 4; + // First neighbour id of node 0 on layer 0 sits right after the header, + // levels, tombstones and node 0's count. + let at = 4 + 7 * 4 + 50 + 50 + 4; + bytes[at..at + 4].copy_from_slice(&9999u32.to_le_bytes()); + let crc = clawhdf5_format::checksum::crc32(&bytes[..body_len]); + bytes[body_len..].copy_from_slice(&crc.to_le_bytes()); + assert!(HnswIndex::from_graph_bytes(&bytes, vectors).is_err()); + } + #[test] fn select_neighbors_prefers_diverse_directions_and_fills_up() { // Node at the origin. Three candidates bunched together on the right, diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index d6e5d19..4e50c75 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -7,8 +7,8 @@ //! * **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 first query after open (which pays for any index rebuild), and -//! steady-state `hybrid_search` p50/p99 at each store size. +//! 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, @@ -310,6 +310,13 @@ fn bench_end_to_end(n: usize, json: &mut Vec) { 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(); @@ -343,8 +350,9 @@ fn bench_end_to_end(n: usize, json: &mut Vec) { ); println!( - "| {n} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |", + "| {n} | {:.0} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |", millis(ingest), + millis(cold_build), millis(checkpoint), millis(open), millis(first_query), @@ -354,7 +362,8 @@ fn bench_end_to_end(n: usize, json: &mut Vec) { ); json.push(serde_json::json!({ "bench": "hybrid_search", "n": n, - "ingest_ms": millis(ingest), "checkpoint_ms": millis(checkpoint), + "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, })); @@ -394,9 +403,9 @@ fn main() { } println!("\n### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n"); println!( - "| N | ingest ms | checkpoint ms | open ms | first query ms | p50 ms | p99 ms | QPS |" + "| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |" ); - println!("|---:|---:|---:|---:|---:|---:|---:|---:|"); + println!("|---:|---:|---:|---:|---:|---:|---:|---:|---:|"); for &n in sizes { bench_end_to_end(n, &mut json); }