p_memory/
graph_search.rs

1//! 图谱搜索：基于 petgraph 的**内存态图**。
2//!
3//! 存储仍是 SQLite（权威）；本模块按需把 `relations` 投影表读进内存建图，
4//! 调 petgraph 现成算法，用完即弃。算法一律用 petgraph 实现，不自己造。
5
6use crate::graph::{Entity, GraphStore};
7use crate::storage::{self, filter_sql};
8use crate::types::{ReadFilter, RecordKind, WriteReceipt};
9use crate::{Error, Result};
10use petgraph::algo::{astar, connected_components, tarjan_scc};
11use petgraph::graph::{DiGraph, Graph, NodeIndex};
12use rusqlite::{params, params_from_iter, Connection};
13use std::collections::{BTreeMap, HashMap, HashSet};
14
15/// 别名等价关系的内置谓词：建图时并查集缩点，把互为别名的实体折成同一超节点。
16const SAME_AS: &str = "sys:same_as";
17
18/// 查询期并查集：只服务 `sys:same_as` 的缩点，用完即弃，不落盘。
19#[derive(Default)]
20struct UnionFind { parent: HashMap<i64, i64> }
21
22impl UnionFind {
23    fn find(&mut self, x: i64) -> i64 {
24        if !self.parent.contains_key(&x) { self.parent.insert(x, x); return x; }
25        let mut root = x;
26        while self.parent[&root] != root { root = self.parent[&root]; }
27        let mut cursor = x;
28        while self.parent[&cursor] != root { let next = self.parent[&cursor]; self.parent.insert(cursor, root); cursor = next; }
29        root
30    }
31    fn union(&mut self, a: i64, b: i64) {
32        let (ra, rb) = (self.find(a), self.find(b));
33        if ra != rb { self.parent.insert(ra, rb); }
34    }
35}
36
37/// 从 `relations` 构建的内存态图。节点权重 = 实体 record_id，边权重 = 谓词文本。
38///
39/// 图是**有向**的（petgraph `Graph` 默认 `Directed`，`neighbors` 只返回出边）：
40/// 反向查询（如由「父亲」反查「子女」）不会自动连通，靠 `predicate_rules` 在 `snapshot`
41/// 里补出对称/逆关系的虚拟边来打通，物理表不落双向边。
42///
43/// 节点 id 是**缩点后的代表 id**：互为 `sys:same_as` 的实体被并查集折叠成同一节点，
44/// `canonical` 保存「原始 id -> 代表 id」的映射，查询入口先做一次换算。
45pub struct GraphView {
46    graph: Graph<i64, String>,
47    index: HashMap<i64, NodeIndex>,
48    canonical: HashMap<i64, i64>,
49}
50
51/// 读一次快照：`filter` 范围内的实体 id，以及有向边 (subject, object, predicate)。
52///
53/// 边里已按 `predicate_rules` 补入内存虚拟边——对称谓词补反向同谓词边、有逆谓词的补对偶边。
54fn snapshot(conn: &Connection, filter: &ReadFilter) -> Result<(Vec<i64>, Vec<(i64, i64, String)>)> {
55    let (cond, values) = filter_sql(filter, &[RecordKind::Entity], false)?;
56    let mut stmt = conn.prepare(&format!(
57        "SELECT e.record_id FROM entities e JOIN records r ON r.id=e.record_id WHERE {cond}"
58    ))?;
59    let nodes = stmt
60        .query_map(params_from_iter(values), |row| row.get::<_, i64>(0))?
61        .collect::<rusqlite::Result<Vec<_>>>()?;
62
63    let (cond, values) = filter_sql(filter, &[RecordKind::Relation], false)?;
64    let mut stmt = conn.prepare(&format!(
65        "SELECT rel.subject_id, rel.object_id, s.text, pr.is_symmetric, inv.text \
66         FROM relations rel JOIN records r ON r.id=rel.record_id \
67         JOIN strings s ON s.id=rel.predicate_id \
68         LEFT JOIN predicate_rules pr ON pr.predicate_id=rel.predicate_id \
69         LEFT JOIN strings inv ON inv.id=pr.inverse_predicate_id WHERE {cond}"
70    ))?;
71    let mut edges: Vec<(i64, i64, String)> = Vec::new();
72    let mut rows = stmt.query(params_from_iter(values))?;
73    while let Some(row) = rows.next()? {
74        let subject: i64 = row.get(0)?;
75        let object: i64 = row.get(1)?;
76        let predicate: String = row.get(2)?;
77        let symmetric: Option<i64> = row.get(3)?;
78        let inverse: Option<String> = row.get(4)?;
79        edges.push((subject, object, predicate.clone()));
80        if symmetric == Some(1) {
81            edges.push((object, subject, predicate));
82        } else if let Some(inverse) = inverse {
83            edges.push((object, subject, inverse));
84        }
85    }
86    Ok((nodes, edges))
87}
88
89/// 按 `sys:same_as` 求每个 id 的代表 id（含 nodes 与所有边端点）。
90fn canonical_ids(nodes: &[i64], edges: &[(i64, i64, String)]) -> HashMap<i64, i64> {
91    let mut uf = UnionFind::default();
92    for (s, o, pred) in edges {
93        if pred == SAME_AS { uf.union(*s, *o); }
94    }
95    let mut canonical: HashMap<i64, i64> = HashMap::new();
96    for id in nodes.iter().copied().chain(edges.iter().flat_map(|(s, o, _)| [*s, *o])) {
97        let representative = uf.find(id);
98        canonical.insert(id, representative);
99    }
100    canonical
101}
102
103impl GraphStore {
104    /// 把 `filter` 范围（namespace/scope/tag）内的实体铺成节点、关系连成边，建一张有向图。
105    ///
106    /// 节点来自 entity 记录（含没有关系的孤立实体，「谁没连进主图」才看得出来），
107    /// 边来自 relation 记录。范围外的记录不进图——多命名空间不会串。
108    pub fn build_graph(&self, filter: &ReadFilter) -> Result<GraphView> {
109        let state = self.0.read()?;
110        let (nodes, edges) = snapshot(state.conn(), filter)?;
111        let canonical = canonical_ids(&nodes, &edges);
112        let mut graph = Graph::<i64, String>::new();
113        let mut index: HashMap<i64, NodeIndex> = HashMap::new();
114        for raw in nodes.iter().copied().chain(edges.iter().flat_map(|(s, o, _)| [*s, *o])) {
115            let representative = canonical[&raw];
116            if !index.contains_key(&representative) {
117                let i = graph.add_node(representative);
118                index.insert(representative, i);
119            }
120        }
121        for (s, o, predicate) in edges {
122            if predicate == SAME_AS { continue; }
123            graph.add_edge(index[&canonical[&s]], index[&canonical[&o]], predicate);
124        }
125        Ok(GraphView { graph, index, canonical })
126    }
127
128    /// 强连通分量（有向，`tarjan_scc`）：互相可达的实体环，如「互为对手/同伙」的闭环。
129    /// 只返回大小 > 1 的分量，孤点自环被滤掉。
130    pub fn strongly_connected(&self, filter: &ReadFilter) -> Result<Vec<Vec<i64>>> {
131        let (nodes, edges) = {
132            let state = self.0.read()?;
133            snapshot(state.conn(), filter)?
134        };
135        let canonical = canonical_ids(&nodes, &edges);
136        let mut graph = DiGraph::<i64, ()>::new();
137        let mut index: HashMap<i64, NodeIndex> = HashMap::new();
138        for raw in nodes.iter().copied().chain(edges.iter().flat_map(|(s, o, _)| [*s, *o])) {
139            let representative = canonical[&raw];
140            if !index.contains_key(&representative) {
141                let i = graph.add_node(representative);
142                index.insert(representative, i);
143            }
144        }
145        for (s, o, _pred) in edges {
146            let (si, oi) = (index[&canonical[&s]], index[&canonical[&o]]);
147            if si != oi { graph.add_edge(si, oi, ()); }
148        }
149        Ok(tarjan_scc(&graph)
150            .into_iter()
151            .filter(|c| c.len() > 1)
152            .map(|c| c.into_iter().map(|n| graph[n]).collect())
153            .collect())
154    }
155
156    /// 从 `root` 出发 `depth` 跳内的实体（带 name / aliases / attributes）。
157    pub fn ego(&self, root: i64, depth: usize, filter: &ReadFilter, limit: usize) -> Result<Vec<Entity>> {
158        let ids = self.build_graph(filter)?.ego_ids(root, depth, limit);
159        let state = self.0.read()?;
160        // 一次批量取回：逐条 `get` 会为每个实体各跑一遍过滤与装配。
161        let mut loaded: BTreeMap<i64, Entity> = storage::load_many(state.conn(), &ids, filter)?;
162        Ok(ids.iter().filter_map(|id| loaded.remove(id)).collect())
163    }
164
165    /// `from` → `to` 桥接路径上的实体（含两端）。不连通返回 None。
166    pub fn path(&self, from: i64, to: i64, filter: &ReadFilter) -> Result<Option<Vec<Entity>>> {
167        let Some(ids) = self.build_graph(filter)?.path_ids(from, to) else {
168            return Ok(None);
169        };
170        let state = self.0.read()?;
171        let mut loaded: BTreeMap<i64, Entity> = storage::load_many(state.conn(), &ids, filter)?;
172        Ok(Some(ids.iter().filter_map(|id| loaded.remove(id)).collect()))
173    }
174
175    /// 登记谓词元规则：`symmetric` 声明对称谓词（反向即自身），`inverse` 声明逆谓词
176    /// （反向补一条对偶边，如 `父亲` 的逆是 `子女`）。二者互斥，规则只影响建图时的内存补边。
177    pub fn set_predicate_rule(&self, predicate: &str, inverse: Option<&str>, symmetric: bool) -> Result<WriteReceipt<()>> {
178        storage::validate_identity("predicate", predicate)?;
179        if symmetric && inverse.is_some() {
180            return Err(Error::Validation("a symmetric predicate must not declare an inverse".into()));
181        }
182        if let Some(inverse) = inverse { storage::validate_identity("inverse predicate", inverse)?; }
183        self.0.mutate_meta(|tx| {
184            let predicate_id = storage::term_id(tx, predicate)?;
185            let inverse_id = match inverse { Some(text) => Some(storage::term_id(tx, text)?), None => None };
186            tx.execute("INSERT INTO predicate_rules(predicate_id,inverse_predicate_id,is_symmetric) VALUES (?1,?2,?3)
187                ON CONFLICT(predicate_id) DO UPDATE SET inverse_predicate_id=excluded.inverse_predicate_id,is_symmetric=excluded.is_symmetric",
188                params![predicate_id, inverse_id, i64::from(symmetric)])?;
189            Ok(())
190        })
191    }
192}
193
194impl GraphView {
195    pub fn node_count(&self) -> usize {
196        self.graph.node_count()
197    }
198
199    pub fn edge_count(&self) -> usize {
200        self.graph.edge_count()
201    }
202
203    /// 从 `root` 出发 `depth` 跳内的实体 id（不含 root）。`root` 会先折算成缩点后的代表 id。
204    pub fn ego_ids(&self, root: i64, depth: usize, limit: usize) -> Vec<i64> {
205        let root = self.canonical.get(&root).copied().unwrap_or(root);
206        let start = match self.index.get(&root) {
207            Some(&i) => i,
208            None => return Vec::new(),
209        };
210        let mut seen: HashSet<NodeIndex> = HashSet::new();
211        seen.insert(start);
212        let mut frontier = vec![start];
213        let mut out: Vec<i64> = Vec::new();
214        for _ in 0..depth {
215            let mut next = Vec::new();
216            for n in &frontier {
217                for nb in self.graph.neighbors(*n) {
218                    if seen.insert(nb) {
219                        out.push(self.graph[nb]);
220                        next.push(nb);
221                        if out.len() >= limit {
222                            return out;
223                        }
224                    }
225                }
226            }
227            if next.is_empty() {
228                break;
229            }
230            frontier = next;
231        }
232        out
233    }
234
235    /// `from` 到 `to` 的最短路径（节点 id 列表，含两端）。用 petgraph 的 `astar`。
236    /// 两端会先折算成缩点后的代表 id。
237    pub fn path_ids(&self, from: i64, to: i64) -> Option<Vec<i64>> {
238        let from = self.canonical.get(&from).copied().unwrap_or(from);
239        let to = self.canonical.get(&to).copied().unwrap_or(to);
240        let start = self.index.get(&from)?;
241        let goal = self.index.get(&to)?;
242        let (_cost, path) = astar(&self.graph, *start, |n| n == *goal, |_e| 1i32, |_n| 0i32)?;
243        Some(path.into_iter().map(|n| self.graph[n]).collect())
244    }
245
246    /// 连通分量数量（无向）。用 petgraph 的 `connected_components`。
247    pub fn component_count(&self) -> usize {
248        connected_components(&self.graph)
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use crate::graph::{EntityInput, GraphBatch, RelationInput};
255    use crate::types::{ReadFilter, RecordInput};
256    use crate::KnowledgeBase;
257    use std::collections::BTreeMap;
258
259    fn ent(name: &str) -> EntityInput {
260        EntityInput {
261            record: RecordInput::default(),
262            name: name.into(),
263            entity_type: "person".into(),
264            aliases: Vec::new(),
265            attributes: BTreeMap::new(),
266            summary: String::new(),
267        }
268    }
269
270    fn rel(s: i64, p: &str, o: i64) -> RelationInput {
271        RelationInput {
272            record: RecordInput::default(),
273            subject_id: s,
274            predicate: p.into(),
275            object_id: o,
276            confidence: 1.0,
277            reason: String::new(),
278        }
279    }
280
281    #[test]
282    fn ego_path_components() {
283        let dir = tempfile::tempdir().unwrap();
284        let kb = KnowledgeBase::open(dir.path()).unwrap();
285        let ents = kb
286            .graph()
287            .apply_batch(&GraphBatch { entities: vec![ent("A"), ent("B"), ent("C")], ..Default::default() })
288            .unwrap()
289            .value
290            .entities;
291        let (a, b, c) = (ents[0].header.id, ents[1].header.id, ents[2].header.id);
292        kb.graph()
293            .apply_batch(&GraphBatch { relations: vec![rel(a, "knows", b), rel(b, "knows", c)], ..Default::default() })
294            .unwrap();
295
296        let view = kb.graph().build_graph(&ReadFilter::default()).unwrap();
297        assert_eq!(view.node_count(), 3);
298        assert_eq!(view.edge_count(), 2);
299        assert_eq!(view.component_count(), 1);
300
301        // 邻域：A 的一跳只有 B，两跳兜住 B、C
302        assert_eq!(view.ego_ids(a, 1, 100), vec![b]);
303        let mut ego = view.ego_ids(a, 2, 100);
304        ego.sort();
305        let mut expect = vec![b, c];
306        expect.sort();
307        assert_eq!(ego, expect);
308
309        // 桥接：A→C 的最短路径
310        assert_eq!(view.path_ids(a, c), Some(vec![a, b, c]));
311
312        // Entity 级查询：返回的实体对象自带 name 字段
313        let names: Vec<String> = kb
314            .graph()
315            .ego(a, 2, &ReadFilter::default(), 100)
316            .unwrap()
317            .into_iter()
318            .map(|e| e.name)
319            .collect();
320        assert!(names.contains(&"B".to_string()) && names.contains(&"C".to_string()));
321        let hop: Vec<String> = kb
322            .graph()
323            .path(a, c, &ReadFilter::default())
324            .unwrap()
325            .unwrap()
326            .into_iter()
327            .map(|e| e.name)
328            .collect();
329        assert_eq!(hop, vec!["A", "B", "C"]);
330
331        // 强连通：补 B→A 形成环，SCC 找到互相可达的 {A,B}
332        kb.graph()
333            .apply_batch(&GraphBatch { relations: vec![rel(b, "knows", a)], ..Default::default() })
334            .unwrap();
335        let scc = kb.graph().strongly_connected(&ReadFilter::default()).unwrap();
336        assert_eq!(scc.len(), 1);
337        let mut ring = scc[0].clone();
338        ring.sort();
339        let mut exp = vec![a, b];
340        exp.sort();
341        assert_eq!(ring, exp);
342
343        // 孤立实体独立成块
344        kb.graph()
345            .apply_batch(&GraphBatch { entities: vec![ent("D")], ..Default::default() })
346            .unwrap();
347        let view = kb.graph().build_graph(&ReadFilter::default()).unwrap();
348        assert_eq!(view.node_count(), 4);
349        assert_eq!(view.component_count(), 2);
350    }
351
352    #[test]
353    fn filter_isolates_namespaces() {
354        let dir = tempfile::tempdir().unwrap();
355        let kb = KnowledgeBase::open(dir.path()).unwrap();
356
357        let mut ea = ent("A");
358        ea.record.namespace = "a".into();
359        let mut eb = ent("B");
360        eb.record.namespace = "a".into();
361        let mut ec = ent("C");
362        ec.record.namespace = "b".into();
363        let got = kb
364            .graph()
365            .apply_batch(&GraphBatch { entities: vec![ea, eb, ec], ..Default::default() })
366            .unwrap()
367            .value
368            .entities;
369        let (a, b) = (got[0].header.id, got[1].header.id);
370        let mut r = rel(a, "knows", b);
371        r.record.namespace = "a".into();
372        kb.graph()
373            .apply_batch(&GraphBatch { relations: vec![r], ..Default::default() })
374            .unwrap();
375
376        let fa = ReadFilter { namespace: "a".into(), scopes: vec!["public".into()], tags: vec![], note_ids: vec![] };
377        let fb = ReadFilter { namespace: "b".into(), scopes: vec!["public".into()], tags: vec![], note_ids: vec![] };
378        let va = kb.graph().build_graph(&fa).unwrap();
379        assert_eq!((va.node_count(), va.edge_count()), (2, 1));
380        let vb = kb.graph().build_graph(&fb).unwrap();
381        assert_eq!((vb.node_count(), vb.edge_count()), (1, 0));
382    }
383
384    #[test]
385    fn virtual_inverse_and_symmetric_edges_enable_reverse_queries() {
386        let dir = tempfile::tempdir().unwrap();
387        let kb = KnowledgeBase::open(dir.path()).unwrap();
388        let filter = ReadFilter::default();
389        let ents = kb
390            .graph()
391            .apply_batch(&GraphBatch { entities: vec![ent("张伟"), ent("张父"), ent("甲"), ent("乙")], ..Default::default() })
392            .unwrap()
393            .value
394            .entities;
395        let (son, father, jia, yi) = (ents[0].header.id, ents[1].header.id, ents[2].header.id, ents[3].header.id);
396        // 物理只写单向边：张伟 --父亲--> 张父；甲 --同事--> 乙
397        kb.graph()
398            .apply_batch(&GraphBatch { relations: vec![rel(son, "父亲", father), rel(jia, "同事", yi)], ..Default::default() })
399            .unwrap();
400
401        // 未登记规则：反向查询天然不通（有向图）
402        let view = kb.graph().build_graph(&filter).unwrap();
403        assert!(view.ego_ids(father, 1, 100).is_empty());
404        assert_eq!(view.path_ids(father, son), None);
405
406        // 登记逆谓词：反向补出「子女」边，反向 ego/path 打通
407        kb.graph().set_predicate_rule("父亲", Some("子女"), false).unwrap();
408        let view = kb.graph().build_graph(&filter).unwrap();
409        assert_eq!(view.ego_ids(father, 1, 100), vec![son]);
410        assert_eq!(view.path_ids(father, son), Some(vec![father, son]));
411
412        // 登记对称谓词：反向即自身
413        kb.graph().set_predicate_rule("同事", None, true).unwrap();
414        let view = kb.graph().build_graph(&filter).unwrap();
415        assert_eq!(view.ego_ids(yi, 1, 100), vec![jia]);
416
417        // 对称谓词声明逆谓词属于非法组合
418        assert!(kb.graph().set_predicate_rule("同事", Some("同事"), true).is_err());
419    }
420
421    #[test]
422    fn same_as_contracts_alias_nodes() {
423        let dir = tempfile::tempdir().unwrap();
424        let kb = KnowledgeBase::open(dir.path()).unwrap();
425        let filter = ReadFilter::default();
426        let ents = kb
427            .graph()
428            .apply_batch(&GraphBatch { entities: vec![ent("林昭"), ent("林先生"), ent("导师")], ..Default::default() })
429            .unwrap()
430            .value
431            .entities;
432        let (lin, alias, mentor) = (ents[0].header.id, ents[1].header.id, ents[2].header.id);
433        // 别名等价用 sys:same_as 显式关系表达，绝不物理合并实体
434        kb.graph()
435            .apply_batch(&GraphBatch { relations: vec![rel(lin, "sys:same_as", alias), rel(alias, "导师", mentor)], ..Default::default() })
436            .unwrap();
437
438        let view = kb.graph().build_graph(&filter).unwrap();
439        // 三个实体折叠成两个节点：别名与主实体是同一超节点
440        assert_eq!(view.node_count(), 2);
441        assert_eq!(view.edge_count(), 1);
442        // 别名不再打断邻域：林昭与林先生的一跳都直接到导师
443        assert_eq!(view.ego_ids(lin, 1, 100), vec![mentor]);
444        assert_eq!(view.ego_ids(alias, 1, 100), vec![mentor]);
445        // 路径里不再夹着别名中间节点：只剩两个节点
446        let path = view.path_ids(lin, mentor).unwrap();
447        assert_eq!(path.len(), 2);
448        assert!(path.contains(&mentor));
449    }
450}
p_memory/graph_search.rs

p_memory/
graph_search.rs
