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:
co-authored by
Claude Fable 5.1
parent
2bfbb7fb4b
commit
0ee698accd
@@ -159,6 +159,23 @@ index removes that.
|
|||||||
| 10000 | 104 | 31.1 | 10.9 | 2627.1 | 2.09 | 2.11 | 479.5 |
|
| 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 |
|
| 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
|
||||||
|
`<store>.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
|
## Vector Search Latency
|
||||||
|
|
||||||
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
|
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
|
||||||
|
|||||||
@@ -23,6 +23,17 @@
|
|||||||
immediately; a crash in between forgets only the boosts since the last
|
immediately; a crash in between forgets only the boosts since the last
|
||||||
checkpoint. Activation weights are now capped (16.0) - they grew without
|
checkpoint. Activation weights are now capped (16.0) - they grew without
|
||||||
bound.
|
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 `<store>.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),
|
- `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
|
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.
|
bound and then ignored it is gone. IDF is computed per query.
|
||||||
|
|||||||
@@ -390,8 +390,13 @@ impl HDF5Memory {
|
|||||||
} else {
|
} else {
|
||||||
Some(store_lock::StoreLock::acquire(path)?)
|
Some(store_lock::StoreLock::acquire(path)?)
|
||||||
};
|
};
|
||||||
let ((config, mut cache, sessions, knowledge), wal_applied) =
|
let ((config, mut cache, sessions, knowledge), checkpoint) =
|
||||||
storage::read_from_disk_with_mark(path)?;
|
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
|
// Replay WAL if present
|
||||||
let wal_path = path.with_extension("h5.wal");
|
let wal_path = path.with_extension("h5.wal");
|
||||||
@@ -408,6 +413,9 @@ impl HDF5Memory {
|
|||||||
&& let Ok(entries) =
|
&& let Ok(entries) =
|
||||||
wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)
|
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);
|
wal::replay_into_cache(&entries, &mut cache);
|
||||||
}
|
}
|
||||||
None
|
None
|
||||||
@@ -419,6 +427,9 @@ impl HDF5Memory {
|
|||||||
// in case the process died between writing the .h5 and
|
// in case the process died between writing the .h5 and
|
||||||
// truncating the WAL.
|
// truncating the WAL.
|
||||||
let entries = wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)?;
|
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);
|
wal::replay_into_cache(&entries, &mut cache);
|
||||||
Some(wal::WalFile::open(&wal_path)?)
|
Some(wal::WalFile::open(&wal_path)?)
|
||||||
} else if config.wal_enabled {
|
} else if config.wal_enabled {
|
||||||
@@ -427,6 +438,25 @@ impl HDF5Memory {
|
|||||||
None
|
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 {
|
Ok(Self {
|
||||||
config,
|
config,
|
||||||
cache,
|
cache,
|
||||||
@@ -435,14 +465,14 @@ impl HDF5Memory {
|
|||||||
wal,
|
wal,
|
||||||
strategy: None,
|
strategy: None,
|
||||||
ephemeral: None,
|
ephemeral: None,
|
||||||
// Existing data is loaded from disk + WAL replay; mark the index
|
// Reuse the vector index saved with the checkpoint when there is
|
||||||
// dirty so it is (re)built from the cache on the first search.
|
// one; otherwise mark it dirty so the first search builds it.
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
hnsw: None,
|
hnsw_dirty: loaded_index.is_none(),
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
hnsw_dirty: true,
|
hnsw_synced_len: synced_len,
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
hnsw_synced_len: 0,
|
hnsw: loaded_index,
|
||||||
// No on-disk provenance ledger exists yet (see CLAUDE.md), so
|
// No on-disk provenance ledger exists yet (see CLAUDE.md), so
|
||||||
// there's no historical hash to verify loaded records against —
|
// there's no historical hash to verify loaded records against —
|
||||||
// the store starts empty and is populated as records are
|
// 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.
|
/// 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
|
/// 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
|
// Record which WAL prefix this checkpoint contains, so a crash before
|
||||||
// the truncate below can't replay those entries a second time.
|
// the truncate below can't replay those entries a second time.
|
||||||
let wal_applied = self.wal.as_ref().map(|w| w.mark());
|
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.path,
|
||||||
&self.config,
|
&self.config,
|
||||||
&self.cache,
|
&self.cache,
|
||||||
&self.sessions,
|
&self.sessions,
|
||||||
&self.knowledge,
|
&self.knowledge,
|
||||||
|
&schema::CheckpointMeta {
|
||||||
wal_applied,
|
wal_applied,
|
||||||
|
ann_generation,
|
||||||
|
},
|
||||||
)?;
|
)?;
|
||||||
if let Some(ref mut w) = self.wal {
|
if let Some(ref mut w) = self.wal {
|
||||||
w.truncate()?;
|
w.truncate()?;
|
||||||
@@ -686,6 +815,30 @@ impl HDF5Memory {
|
|||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
fn ensure_hnsw_fresh(&mut self) {
|
fn ensure_hnsw_fresh(&mut self) {
|
||||||
let n = self.cache.embeddings.len();
|
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 {
|
if self.hnsw_dirty || self.hnsw_synced_len != n {
|
||||||
self.hnsw = self.build_hnsw_from_cache();
|
self.hnsw = self.build_hnsw_from_cache();
|
||||||
self.hnsw_synced_len = n;
|
self.hnsw_synced_len = n;
|
||||||
@@ -928,8 +1081,8 @@ impl AgentMemory for HDF5Memory {
|
|||||||
);
|
);
|
||||||
indices.push(idx);
|
indices.push(idx);
|
||||||
}
|
}
|
||||||
// Batch inserts rebuild the index once rather than node-by-node.
|
// The vector and keyword indexes pick the new records up
|
||||||
self.hnsw_mark_dirty();
|
// incrementally the next time they are needed.
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
Ok(indices)
|
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() {
|
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"))?;
|
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)
|
Ok(snapshot)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1674,6 +1834,126 @@ mod tests {
|
|||||||
assert_eq!(restored.cache.chunks, ["checkpointed", "wal-only"]);
|
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]
|
#[test]
|
||||||
fn keyword_index_stays_in_sync_through_every_mutation() {
|
fn keyword_index_stays_in_sync_through_every_mutation() {
|
||||||
let dir = TempDir::new().unwrap();
|
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.
|
/// the checkpoint was taken with an empty WAL.
|
||||||
const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len";
|
const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len";
|
||||||
const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc";
|
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.
|
/// Build a complete HDF5 file from the in-memory state.
|
||||||
pub fn build_hdf5_file(
|
pub fn build_hdf5_file(
|
||||||
@@ -43,6 +44,34 @@ pub fn build_hdf5_file_with_mark(
|
|||||||
knowledge: &KnowledgeCache,
|
knowledge: &KnowledgeCache,
|
||||||
wal_applied: Option<WalMark>,
|
wal_applied: Option<WalMark>,
|
||||||
) -> Result<Vec<u8>, MemoryError> {
|
) -> 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();
|
let mut builder = clawhdf5::FileBuilder::new();
|
||||||
|
|
||||||
// /meta group with schema attributes
|
// /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_LEN_ATTR, AttrValue::I64(mark.len as i64));
|
||||||
meta.set_attr(WAL_APPLIED_CRC_ATTR, AttrValue::I64(i64::from(mark.crc)));
|
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
|
// Need at least one dataset in the group for it to be a proper group
|
||||||
meta.create_dataset("_marker").with_u8_data(&[1]).compact();
|
meta.create_dataset("_marker").with_u8_data(&[1]).compact();
|
||||||
let finished_meta = meta.finish();
|
let finished_meta = meta.finish();
|
||||||
@@ -386,6 +420,22 @@ pub fn read_wal_mark(file: &clawhdf5::File) -> Option<WalMark> {
|
|||||||
Some(WalMark { len, crc })
|
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(
|
pub fn validate_and_load(
|
||||||
file: &clawhdf5::File,
|
file: &clawhdf5::File,
|
||||||
) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> {
|
) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> {
|
||||||
|
|||||||
@@ -34,7 +34,23 @@ pub fn write_to_disk_with_mark(
|
|||||||
knowledge: &KnowledgeCache,
|
knowledge: &KnowledgeCache,
|
||||||
wal_applied: Option<WalMark>,
|
wal_applied: Option<WalMark>,
|
||||||
) -> Result<(), MemoryError> {
|
) -> 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() {
|
if bytes.is_empty() {
|
||||||
return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into()));
|
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.
|
/// 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;
|
use std::io::Write;
|
||||||
let mut f = std::fs::File::create(path).map_err(MemoryError::Io)?;
|
let mut f = std::fs::File::create(path).map_err(MemoryError::Io)?;
|
||||||
f.write_all(bytes).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
|
/// This is per-checkpoint/snapshot cost only (each is already a full file
|
||||||
/// write). Individual WAL appends are deliberately not synced — see the
|
/// write). Individual WAL appends are deliberately not synced — see the
|
||||||
/// durability notes in the crate docs.
|
/// 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)?;
|
std::fs::rename(from, to).map_err(MemoryError::Io)?;
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
if let Some(dir) = to.parent() {
|
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))
|
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.
|
/// Copy an HDF5 file atomically to a destination.
|
||||||
pub fn snapshot_file(src: &Path, dest: &Path) -> Result<std::path::PathBuf, MemoryError> {
|
pub fn snapshot_file(src: &Path, dest: &Path) -> Result<std::path::PathBuf, MemoryError> {
|
||||||
let dest_file = if dest.is_dir() {
|
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.
|
/// On-disk format version for the serialized HNSW index.
|
||||||
///
|
///
|
||||||
/// - Version 1: original layout (`vectors`, `graph_layer_*`, `config`), no
|
/// - 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.
|
/// Returns the number of vectors in the index.
|
||||||
pub fn len(&self) -> usize {
|
pub fn len(&self) -> usize {
|
||||||
self.vectors.len()
|
self.vectors.len()
|
||||||
@@ -1147,6 +1304,73 @@ mod tests {
|
|||||||
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
|
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]
|
#[test]
|
||||||
fn select_neighbors_prefers_diverse_directions_and_fills_up() {
|
fn select_neighbors_prefers_diverse_directions_and_fills_up() {
|
||||||
// Node at the origin. Three candidates bunched together on the right,
|
// 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
|
//! * **ANN** — index build time, and for each `ef`: recall@10 against an exact
|
||||||
//! brute-force scan, queries/second, and p50/p99 latency.
|
//! brute-force scan, queries/second, and p50/p99 latency.
|
||||||
//! * **End to end** — `HDF5Memory`: ingest time, checkpoint time, `open()`
|
//! * **End to end** — `HDF5Memory`: ingest time, checkpoint time, `open()`
|
||||||
//! time, the first query after open (which pays for any index rebuild), and
|
//! time, the one-off cold index build (first query ever), the first query
|
||||||
//! steady-state `hybrid_search` p50/p99 at each store size.
|
//! after a reopen, and steady-state `hybrid_search` p50/p99 at each size.
|
||||||
//!
|
//!
|
||||||
//! Data is *clustered* (points = cluster centre + noise, unit-normalised), not
|
//! Data is *clustered* (points = cluster centre + noise, unit-normalised), not
|
||||||
//! uniform: uniform random high-dimensional vectors are nearly equidistant,
|
//! 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();
|
let t = Instant::now();
|
||||||
mem.save_batch(entries).unwrap();
|
mem.save_batch(entries).unwrap();
|
||||||
let ingest = t.elapsed();
|
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();
|
let t = Instant::now();
|
||||||
mem.flush_wal().unwrap();
|
mem.flush_wal().unwrap();
|
||||||
let checkpoint = t.elapsed();
|
let checkpoint = t.elapsed();
|
||||||
@@ -343,8 +350,9 @@ fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
|
|||||||
);
|
);
|
||||||
|
|
||||||
println!(
|
println!(
|
||||||
"| {n} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |",
|
"| {n} | {:.0} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |",
|
||||||
millis(ingest),
|
millis(ingest),
|
||||||
|
millis(cold_build),
|
||||||
millis(checkpoint),
|
millis(checkpoint),
|
||||||
millis(open),
|
millis(open),
|
||||||
millis(first_query),
|
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!({
|
json.push(serde_json::json!({
|
||||||
"bench": "hybrid_search", "n": n,
|
"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),
|
"open_ms": millis(open), "first_query_ms": millis(first_query),
|
||||||
"p50_ms": millis(steady.p50), "p99_ms": millis(steady.p99), "qps": steady.qps,
|
"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### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n");
|
||||||
println!(
|
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 {
|
for &n in sizes {
|
||||||
bench_end_to_end(n, &mut json);
|
bench_end_to_end(n, &mut json);
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user