uqa_graph/operators/
traverse.rs1use super::{
10 BTreeMap, BTreeSet, DocId, EdgeId, GraphPayload, GraphPostingList, GraphStore, GraphStoreError,
11 GraphStoreResult, Payload, PostingEntry, PostingList, VertexId, VertexPredicate,
12 DEFAULT_GRAPH_SCORE,
13};
14
15pub struct Traverse<'a> {
19 pub start_vertex: VertexId,
20 pub graph: &'a str,
21 pub label: Option<&'a str>,
22 pub max_hops: u32,
23 pub vertex_predicate: Option<VertexPredicate>,
24 pub score: f64,
25}
26
27impl<'a> Traverse<'a> {
28 pub fn new(start: VertexId, graph: &'a str) -> Self {
29 Self {
30 start_vertex: start,
31 graph,
32 label: None,
33 max_hops: 1,
34 vertex_predicate: None,
35 score: DEFAULT_GRAPH_SCORE,
36 }
37 }
38
39 pub fn label(mut self, label: &'a str) -> Self {
40 self.label = Some(label);
41 self
42 }
43
44 pub fn max_hops(mut self, hops: u32) -> Self {
45 self.max_hops = hops;
46 self
47 }
48
49 pub fn predicate(mut self, p: VertexPredicate) -> Self {
50 self.vertex_predicate = Some(p);
51 self
52 }
53
54 pub fn execute<G: GraphStore>(&self, store: &G) -> GraphStoreResult<GraphPostingList> {
55 store.require_vertex_in_graph(self.start_vertex, self.graph)?;
56 let mut visited: BTreeSet<VertexId> = BTreeSet::new();
57 let mut frontier: BTreeSet<VertexId> = BTreeSet::new();
58 frontier.insert(self.start_vertex);
59 let mut all_edges: BTreeSet<EdgeId> = BTreeSet::new();
60
61 for _ in 0..self.max_hops {
62 let mut next_frontier: BTreeSet<VertexId> = BTreeSet::new();
63 for v in &frontier {
64 for eid in store.out_edge_ids(*v, self.graph)? {
65 let edge = store.get_edge(eid).ok_or_else(|| {
66 GraphStoreError::CorruptGraph(format!("missing traversal edge {eid}"))
67 })?;
68 if let Some(want) = self.label {
69 if edge.label != want {
70 continue;
71 }
72 }
73 let neighbor = edge.target_id;
74 if visited.contains(&neighbor) || frontier.contains(&neighbor) {
75 all_edges.insert(eid);
77 continue;
78 }
79 if let Some(pred) = &self.vertex_predicate {
80 let vtx = store.get_vertex(neighbor).ok_or_else(|| {
81 GraphStoreError::CorruptGraph(format!(
82 "traversal edge {eid} references missing vertex {neighbor}"
83 ))
84 })?;
85 if !pred.matches(vtx) {
86 continue;
87 }
88 }
89 next_frontier.insert(neighbor);
90 all_edges.insert(eid);
91 }
92 }
93 visited.append(&mut frontier.clone());
94 frontier = next_frontier;
95 if frontier.is_empty() {
96 break;
97 }
98 }
99 visited.append(&mut frontier);
100
101 let visited_vec: Vec<VertexId> = visited.iter().copied().collect();
102 let edges_vec: Vec<EdgeId> = all_edges.iter().copied().collect();
103
104 let mut entries: Vec<PostingEntry> = Vec::with_capacity(visited_vec.len());
105 let mut graph_payloads: BTreeMap<DocId, GraphPayload> = BTreeMap::new();
106 for vid in &visited_vec {
107 entries.push(PostingEntry::new(*vid, Payload::with_score(self.score)));
108 graph_payloads.insert(
109 *vid,
110 GraphPayload {
111 subgraph_vertices: visited_vec.clone(),
112 subgraph_edges: edges_vec.clone(),
113 graph_name: self.graph.to_string(),
114 score_override: Some(self.score),
115 },
116 );
117 }
118 GraphPostingList::try_from_parts(
119 PostingList::from_sorted_unchecked(entries),
120 graph_payloads,
121 )
122 .map_err(Into::into)
123 }
124}
125
126