Skip to main content

uqa_graph/operators/
traverse.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Bounded breadth-first graph traversal.
8
9use super::{
10    BTreeMap, BTreeSet, DocId, EdgeId, GraphPayload, GraphPostingList, GraphStore, GraphStoreError,
11    GraphStoreResult, Payload, PostingEntry, PostingList, VertexId, VertexPredicate,
12    DEFAULT_GRAPH_SCORE,
13};
14
15/// `Traverse_{v,l,k}` (Definition 2.2.1): BFS from `start_vertex` along
16/// edges with `label` (any label when `None`) up to `max_hops` hops.
17/// Each visited vertex becomes its own entry in the result.
18pub 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                        // Already explored or about to be — but still record the edge.
76                        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// -------------------------------------------------------------------------
127// VertexMatch
128// -------------------------------------------------------------------------