perf(agent): persist the vector index graph; incremental catch-up
open() marked the HNSW index dirty, so the first search of every session
rebuilt it from scratch — 36 s at 100K records with the (better, slower)
heuristic build. First query after open is now 1.7 / 15 / 159 ms at
1K / 10K / 100K; what remains is the one-off keyword index build.
- clawhdf5-ann: HnswIndex::graph_to_bytes / from_graph_bytes serialize the
graph only (levels, tombstones, adjacency as u32, CRC32). The existing HDF5
serializer embeds a full copy of every vector, which would double a store
that already holds them. Loading validates everything — counts, levels vs
layer count, connection limits, every neighbour id and the layer it must
exist on — so a damaged graph, or a hostile one with a valid checksum, is an
error rather than an out-of-bounds walk during search.
- clawhdf5-agent: each checkpoint writes the graph to <store>.h5.ann (synced,
atomic, before the .h5) and records a fresh generation id in /meta. open()
loads the sidecar only if its generation matches that checkpoint; missing,
stale, damaged or mismatched sidecars are ignored and the index rebuilt.
Records appended through WAL replay join the loaded index incrementally; a
replayed Update or Tombstone invalidates it. snapshot() copies it. Only an
index that exactly mirrors the cache is saved; otherwise a stale sidecar is
removed.
- ensure_hnsw_fresh inserts records appended since the last sync instead of
rebuilding, so save_batch no longer marks the whole index dirty.
- CheckpointMeta { wal_applied, ann_generation } with *_with_meta build/write/
read functions; the *_with_mark ones delegate.
- Harness reports the one-off cold index build separately from the first query
after a reopen.
Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
2bfbb7fb4b
commit
0ee698accd
@@ -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<u64> {
|
||||
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<u64> {
|
||||
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<u64>,
|
||||
cache: &MemoryCache,
|
||||
n_checkpoint: usize,
|
||||
) -> Option<HnswIndex> {
|
||||
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<usize> {
|
||||
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();
|
||||
|
||||
@@ -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<WalMark>,
|
||||
) -> Result<Vec<u8>, 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<WalMark>,
|
||||
/// Identifies the vector-index sidecar (`<store>.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<u64>,
|
||||
}
|
||||
|
||||
/// [`build_hdf5_file`] with checkpoint bookkeeping.
|
||||
pub fn build_hdf5_file_with_meta(
|
||||
config: &MemoryConfig,
|
||||
cache: &MemoryCache,
|
||||
sessions: &SessionCache,
|
||||
knowledge: &KnowledgeCache,
|
||||
checkpoint: &CheckpointMeta,
|
||||
) -> Result<Vec<u8>, 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<WalMark> {
|
||||
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> {
|
||||
|
||||
@@ -34,7 +34,23 @@ pub fn write_to_disk_with_mark(
|
||||
knowledge: &KnowledgeCache,
|
||||
wal_applied: Option<WalMark>,
|
||||
) -> 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<WalMa
|
||||
Ok(((config, cache, sessions, knowledge), wal_applied))
|
||||
}
|
||||
|
||||
/// [`read_from_disk`], plus all checkpoint bookkeeping.
|
||||
pub fn read_from_disk_with_meta(
|
||||
path: &Path,
|
||||
) -> 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<std::path::PathBuf, MemoryError> {
|
||||
let dest_file = if dest.is_dir() {
|
||||
|
||||
@@ -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<u8> {
|
||||
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<Vec<f32>>) -> Result<Self, FormatError> {
|
||||
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<usize, FormatError> {
|
||||
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<usize> = 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<bool> = body
|
||||
.get(pos..pos + n)
|
||||
.ok_or_else(|| bad("truncated"))?
|
||||
.iter()
|
||||
.map(|&d| d != 0)
|
||||
.collect();
|
||||
pos += n;
|
||||
|
||||
let mut graph: Vec<Vec<Vec<usize>>> = 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,
|
||||
|
||||
@@ -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<serde_json::Value>) {
|
||||
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<serde_json::Value>) {
|
||||
);
|
||||
|
||||
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<serde_json::Value>) {
|
||||
);
|
||||
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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user