perf(agent): cache the knowledge graph's adjacency index
bfs_neighbors and spreading_activation built an adjacency index over the
whole graph on every call (1efd82c), so a 2-hop BFS over 1K entities
paid to index every entity and relation first: 155 us, 6.5x the 24 us
the README quoted. Found by the dated benchmark re-run.
The index is now cached on KnowledgeCache and checked against a
fingerprint of the graph on each use — one pass over entity ids and
relation endpoints, no allocation — so any change, including direct
edits of the public entities/relations Vecs (schema.rs's load path
pushes to them), still triggers a rebuild. A test edits the graph
directly in every way (push, in-place rewire, pop + push at equal
length) between traversals.
tank, 2026-09-24: BFS 1K entities 155.1 -> 23.1 us, 100 entities
17.5 -> 5.23 us, spreading activation 100 22.8 -> 10.1 us.
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
@@ -163,12 +163,13 @@ fn levenshtein(a: &str, b: &str) -> usize {
|
||||
/// entities-slice-index map, and an entity-id -> relation-indices map (edges
|
||||
/// touching that entity as either source or target).
|
||||
///
|
||||
/// Built fresh per traversal call rather than cached on `KnowledgeCache`:
|
||||
/// entities/relations are plain `pub` `Vec`s that get pushed to directly
|
||||
/// (e.g. `schema.rs`'s load path bypasses `add_entity`/`add_relation`), so a
|
||||
/// persistent index would need extra bookkeeping to avoid drifting stale. A
|
||||
/// one-off O(V+E) build per call is still a large win over the O(V·E) (BFS)
|
||||
/// / O(steps·active·E) (spreading activation) scans it replaces.
|
||||
/// Cached on `KnowledgeCache` and checked against a fingerprint of the graph
|
||||
/// on every use ([`graph_fingerprint`]). entities/relations are plain `pub`
|
||||
/// `Vec`s that get changed directly (e.g. `schema.rs`'s load path bypasses
|
||||
/// `add_entity`/`add_relation`), so the cache cannot rely on being told about
|
||||
/// changes; the fingerprint notices any of them. Rebuilding it on every
|
||||
/// traversal instead made a 2-hop BFS over 1K entities 6.5x slower than the
|
||||
/// scan it replaced (24 -> 155 µs; `BENCHMARKS.md`, "Knowledge Graph").
|
||||
struct AdjacencyIndex {
|
||||
entity_index: HashMap<u64, usize>,
|
||||
by_entity: HashMap<u64, Vec<usize>>,
|
||||
@@ -204,6 +205,45 @@ impl AdjacencyIndex {
|
||||
}
|
||||
}
|
||||
|
||||
/// A hash of everything [`AdjacencyIndex`] depends on — each entity's id and
|
||||
/// position, each relation's endpoints and position. One linear pass, no
|
||||
/// allocation: far cheaper than building the index, which hashes the same
|
||||
/// values into two maps.
|
||||
fn graph_fingerprint(entities: &[Entity], relations: &[Relation]) -> u64 {
|
||||
// splitmix64-style mixing; order matters, so positions are covered.
|
||||
fn mix(h: u64, v: u64) -> u64 {
|
||||
let mut z = (h ^ v).wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||
z ^ (z >> 31)
|
||||
}
|
||||
let mut h = mix(entities.len() as u64, relations.len() as u64);
|
||||
for e in entities {
|
||||
h = mix(h, e.id);
|
||||
}
|
||||
for r in relations {
|
||||
h = mix(mix(h, r.src), r.tgt);
|
||||
}
|
||||
h
|
||||
}
|
||||
|
||||
/// The cached [`AdjacencyIndex`] and the fingerprint it was built for.
|
||||
/// Cloning a `KnowledgeCache` starts the clone with an empty cache.
|
||||
#[derive(Default)]
|
||||
struct AdjacencyCache(std::sync::Mutex<Option<(u64, std::sync::Arc<AdjacencyIndex>)>>);
|
||||
|
||||
impl Clone for AdjacencyCache {
|
||||
fn clone(&self) -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AdjacencyCache {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("AdjacencyCache")
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// KnowledgeCache
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -216,6 +256,7 @@ pub struct KnowledgeCache {
|
||||
pub alias_strings: Vec<String>,
|
||||
pub alias_entity_ids: Vec<i64>,
|
||||
next_entity_id: u64,
|
||||
adjacency: AdjacencyCache,
|
||||
}
|
||||
|
||||
impl KnowledgeCache {
|
||||
@@ -226,6 +267,7 @@ impl KnowledgeCache {
|
||||
alias_strings: Vec::new(),
|
||||
alias_entity_ids: Vec::new(),
|
||||
next_entity_id: 0,
|
||||
adjacency: AdjacencyCache::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -236,9 +278,29 @@ impl KnowledgeCache {
|
||||
alias_strings: Vec::new(),
|
||||
alias_entity_ids: Vec::new(),
|
||||
next_entity_id: next_id,
|
||||
adjacency: AdjacencyCache::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The adjacency index for the graph as it is now: the cached one if the
|
||||
/// graph's fingerprint still matches, otherwise rebuilt and cached.
|
||||
fn adjacency_index(&self) -> std::sync::Arc<AdjacencyIndex> {
|
||||
let fp = graph_fingerprint(&self.entities, &self.relations);
|
||||
let mut slot = self
|
||||
.adjacency
|
||||
.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some((cached_fp, idx)) = slot.as_ref()
|
||||
&& *cached_fp == fp
|
||||
{
|
||||
return idx.clone();
|
||||
}
|
||||
let idx = std::sync::Arc::new(AdjacencyIndex::build(&self.entities, &self.relations));
|
||||
*slot = Some((fp, idx.clone()));
|
||||
idx
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Entity management
|
||||
// -----------------------------------------------------------------------
|
||||
@@ -397,7 +459,7 @@ impl KnowledgeCache {
|
||||
/// together with their discovered depth. The seed entity itself is NOT
|
||||
/// included. Traversal follows both outgoing and incoming relation edges.
|
||||
pub fn bfs_neighbors(&self, entity_id: u64, max_depth: usize) -> Vec<(Entity, usize)> {
|
||||
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
||||
let idx = self.adjacency_index();
|
||||
let mut visited: HashSet<u64> = HashSet::new();
|
||||
let mut queue: VecDeque<(u64, usize)> = VecDeque::new();
|
||||
let mut results: Vec<(Entity, usize)> = Vec::new();
|
||||
@@ -502,7 +564,7 @@ impl KnowledgeCache {
|
||||
min_activation: f32,
|
||||
max_steps: usize,
|
||||
) -> Vec<(u64, f32)> {
|
||||
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
||||
let idx = self.adjacency_index();
|
||||
let mut activation: HashMap<u64, f32> = HashMap::new();
|
||||
|
||||
// Initialise seeds with activation 1.0.
|
||||
@@ -631,6 +693,51 @@ impl Default for KnowledgeCache {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cached_adjacency_sees_direct_changes_to_the_graph() {
|
||||
// The index is cached across traversals, but entities/relations are
|
||||
// pub Vecs anyone can edit; every kind of edit must be seen.
|
||||
let mut kg = KnowledgeCache::new();
|
||||
let a = kg.add_entity("a", "t", -1);
|
||||
let b = kg.add_entity("b", "t", -1);
|
||||
let c = kg.add_entity("c", "t", -1);
|
||||
kg.add_relation(a, b, "r", 1.0);
|
||||
let ids = |kg: &KnowledgeCache| -> Vec<u64> {
|
||||
let mut v: Vec<u64> = kg.bfs_neighbors(a, 3).iter().map(|(e, _)| e.id).collect();
|
||||
v.sort();
|
||||
v
|
||||
};
|
||||
assert_eq!(ids(&kg), vec![b]);
|
||||
assert_eq!(ids(&kg), vec![b], "cached index reused");
|
||||
|
||||
// Pushed directly, bypassing add_relation.
|
||||
kg.relations.push(Relation {
|
||||
src: b,
|
||||
tgt: c,
|
||||
..Relation::default()
|
||||
});
|
||||
assert_eq!(ids(&kg), vec![b, c]);
|
||||
|
||||
// Rewired in place: same lengths, different edge.
|
||||
kg.relations[1].tgt = a;
|
||||
assert_eq!(ids(&kg), vec![b]);
|
||||
|
||||
// Removed and replaced: same lengths again.
|
||||
kg.relations.pop();
|
||||
kg.relations.push(Relation {
|
||||
src: a,
|
||||
tgt: c,
|
||||
..Relation::default()
|
||||
});
|
||||
assert_eq!(ids(&kg), vec![b, c]);
|
||||
let act: Vec<u64> = kg
|
||||
.spreading_activation(&[a], 0.5, 0.0, 2)
|
||||
.iter()
|
||||
.map(|(id, _)| *id)
|
||||
.collect();
|
||||
assert!(act.contains(&c));
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Original tests — must remain passing
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user