1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
//! Cross-graph embedding carry — `DirGraph::copy_embeddings_from`.
//!
//! Extracted from `dir_graph.rs` to keep that file under the god-file LoC
//! ceiling. The method lives on `DirGraph` (an `impl` block here is fine —
//! same crate/module tree); it's the core behind the Python
//! `KnowledgeGraph.copy_embeddings_from`.
use std::collections::HashMap;
use petgraph::graph::NodeIndex;
use crate::graph::dir_graph::DirGraph;
use crate::graph::schema::EmbeddingStore;
use crate::graph::storage::GraphRead;
impl DirGraph {
/// Copy every embedding store from `src` into this graph, matching
/// vectors by **node id** (not internal index — those differ across a
/// rebuild). For the dominant embedding workflow: rebuild a fresh graph
/// from a source of truth on each load, then `new.copy_embeddings_from(old)`
/// to carry the vectors forward in one call — instead of the manual
/// `embeddings()` snapshot → `add_embeddings()` restore dance. Carries the
/// store's `dimension`, `metric`, `model_id`, and per-node `text_hashes`
/// (so a subsequent `embed_texts(mode='changed')` re-embeds only what
/// actually changed). A vector whose source id has no matching node here is
/// skipped. Returns `(stores_copied, vectors_copied, vectors_skipped)`.
pub fn copy_embeddings_from(&mut self, src: &DirGraph) -> (usize, usize, usize) {
let mut stores_copied = 0usize;
let mut vectors_copied = 0usize;
let mut vectors_skipped = 0usize;
// Snapshot the store keys + node types first so we can build each
// type's id index (a `&mut self` op) before the immutable id lookups.
let store_keys: Vec<(String, String)> = src.embeddings.keys().cloned().collect();
for (node_type, prop) in store_keys {
let Some(src_store) = src.embeddings.get(&(node_type.clone(), prop.clone())) else {
continue;
};
// Build this type's id index once so lookups are O(1).
self.build_id_index(&node_type);
let mut dst_store = EmbeddingStore::new(src_store.dimension);
dst_store.metric = src_store.metric.clone();
dst_store.model_id = src_store.model_id.clone();
for &src_idx in src_store.node_to_slot.keys() {
let Some(src_node) = src.graph.node_weight(NodeIndex::new(src_idx)) else {
vectors_skipped += 1;
continue;
};
let id = src_node.id().into_owned();
let Some(embedding) = src_store.get_embedding(src_idx) else {
vectors_skipped += 1;
continue;
};
match self.lookup_by_id_readonly(&node_type, &id) {
Some(dst_idx) => {
dst_store.set_embedding(dst_idx.index(), embedding);
if let Some(&h) = src_store.text_hashes.get(&src_idx) {
dst_store.set_text_hash(dst_idx.index(), h);
}
vectors_copied += 1;
}
None => vectors_skipped += 1,
}
}
self.embeddings.insert((node_type, prop), dst_store);
stores_copied += 1;
}
(stores_copied, vectors_copied, vectors_skipped)
}
/// Remap every embedding store's internal node indices through `old_to_new`
/// after a `vacuum()` rebuilds the graph with contiguous indices. Drops
/// vectors whose node was deleted (absent from the map), compacts the data
/// buffer to the surviving slots, and resyncs the cached-norm column.
/// Extracted from `vacuum()` to keep `dir_graph.rs` under the god-file
/// ceiling.
pub(crate) fn remap_embedding_slots(&mut self, old_to_new: &HashMap<NodeIndex, NodeIndex>) {
for store in self.embeddings.values_mut() {
let mut new_node_to_slot = HashMap::with_capacity(store.node_to_slot.len());
let mut new_slot_to_node = Vec::with_capacity(store.slot_to_node.len());
let mut new_data = Vec::with_capacity(store.data.len());
for (&old_node_raw, &slot) in &store.node_to_slot {
let old_idx = NodeIndex::new(old_node_raw);
if let Some(&new_idx) = old_to_new.get(&old_idx) {
let new_slot = new_slot_to_node.len();
new_node_to_slot.insert(new_idx.index(), new_slot);
new_slot_to_node.push(new_idx.index());
let start = slot * store.dimension;
let end = start + store.dimension;
new_data.extend_from_slice(&store.data[start..end]);
}
// Deleted nodes (not in old_to_new) are dropped.
}
store.node_to_slot = new_node_to_slot;
store.slot_to_node = new_slot_to_node;
store.data = new_data;
// Slots were remapped wholesale; resync the cached-norm column and
// drop any HNSW index (its slot ids are now stale — rebuild on demand).
store.rebuild_norms();
store.invalidate_index();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::datatypes::Value;
use crate::graph::schema::NodeData;
use crate::graph::storage::GraphWrite;
use std::collections::HashMap;
fn graph_with_docs(ids: &[i64]) -> DirGraph {
let mut g = DirGraph::new();
for &id in ids {
let nd = NodeData::new(
Value::Int64(id),
Value::String(format!("d{id}")),
"Doc".to_string(),
HashMap::new(),
&mut g.interner,
);
let idx = GraphWrite::add_node(&mut g.graph, nd);
g.type_indices.entry_or_default("Doc".to_string()).push(idx);
}
g.build_id_index("Doc");
g
}
/// Vectors carry by node id (not internal index), with dimension, model_id
/// and text-hashes preserved; ids absent in the destination are skipped.
#[test]
fn copies_vectors_by_id_with_provenance() {
let mut src = graph_with_docs(&[1, 2, 3]);
let mut store = EmbeddingStore::new(2);
store.model_id = Some("m/1".to_string());
for &id in &[1i64, 2, 3] {
let idx = src.lookup_by_id_readonly("Doc", &Value::Int64(id)).unwrap();
store.set_embedding(idx.index(), &[id as f32, 0.0]);
store.set_text_hash(idx.index(), EmbeddingStore::text_hash(&format!("t{id}")));
}
src.embeddings
.insert(("Doc".to_string(), "summary_emb".to_string()), store);
// Destination is a fresh rebuild missing id 3.
let mut dst = graph_with_docs(&[1, 2]);
let (stores, vectors, skipped) = dst.copy_embeddings_from(&src);
assert_eq!((stores, vectors, skipped), (1, 2, 1));
let dst_store = dst
.embeddings
.get(&("Doc".to_string(), "summary_emb".to_string()))
.unwrap();
assert_eq!(dst_store.dimension, 2);
assert_eq!(dst_store.model_id.as_deref(), Some("m/1"));
assert_eq!(dst_store.len(), 2);
assert_eq!(dst_store.text_hashes.len(), 2);
// The carried vector landed on the dst node with the matching id.
let dst_idx = dst.lookup_by_id_readonly("Doc", &Value::Int64(2)).unwrap();
assert_eq!(
dst_store.get_embedding(dst_idx.index()),
Some(&[2.0f32, 0.0][..])
);
}
}