Skip to main content

nodedb_graph/traversal/
subgraph.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Subgraph materialization over CSR adjacency.
4
5use std::collections::HashSet;
6
7#[cfg(test)]
8use super::DEFAULT_MAX_VISITED;
9use crate::csr::{CsrIndex, Direction};
10use crate::overlay_delta::GraphOverlayDelta;
11use crate::traversal_overlay::{OverlaySubgraphParams, subgraph_overlay};
12
13impl CsrIndex {
14    /// Materialize a subgraph as `(src, label, dst)` edge tuples within
15    /// max_depth, expanding in `direction`.
16    ///
17    /// Every expanded node records each of its edges in `direction`. The last
18    /// admitted level is not expanded: it records only its edges in
19    /// `direction` whose other endpoint is an admitted node, and admits no
20    /// node.
21    ///
22    /// An empty `label_filter` keeps every edge. Otherwise an edge whose label
23    /// is any listed label passes.
24    ///
25    /// `max_visited` caps the number of nodes visited to prevent supernode fan-out
26    /// explosion. Pass [`crate::traversal::DEFAULT_MAX_VISITED`] for the standard limit.
27    ///
28    /// `overlay`: when `Some` and non-empty, staged edges are included and
29    /// staged tombstones subtract durable edges (read-your-own-writes),
30    /// including through staged-only intermediate nodes. When `None` or
31    /// empty, the durable-only dense path runs unchanged.
32    pub fn subgraph(
33        &self,
34        start_nodes: &[&str],
35        label_filter: &[&str],
36        direction: Direction,
37        max_depth: usize,
38        max_visited: usize,
39        overlay: Option<&GraphOverlayDelta>,
40    ) -> Vec<(String, String, String)> {
41        Self::subgraph_on(
42            Some(self),
43            start_nodes,
44            label_filter,
45            direction,
46            max_depth,
47            max_visited,
48            overlay,
49        )
50    }
51
52    /// [`Self::subgraph`] over an optional partition. `None` is a tenant
53    /// with no CSR partition: the walk then follows only the staged edges of
54    /// `overlay`, and returns no edge without them.
55    pub fn subgraph_on(
56        partition: Option<&Self>,
57        start_nodes: &[&str],
58        label_filter: &[&str],
59        direction: Direction,
60        max_depth: usize,
61        max_visited: usize,
62        overlay: Option<&GraphOverlayDelta>,
63    ) -> Vec<(String, String, String)> {
64        match (partition, overlay) {
65            (_, Some(ov)) if !ov.is_empty() => subgraph_overlay(
66                partition,
67                OverlaySubgraphParams {
68                    start_nodes,
69                    label_filter,
70                    direction,
71                    max_depth,
72                    max_visited,
73                },
74                ov,
75            ),
76            (Some(csr), _) => {
77                csr.subgraph_dense(start_nodes, label_filter, direction, max_depth, max_visited)
78            }
79            (None, _) => Vec::new(),
80        }
81    }
82
83    /// Durable-only subgraph materialization over the dense u32 CSR ids.
84    fn subgraph_dense(
85        &self,
86        start_nodes: &[&str],
87        label_filter: &[&str],
88        direction: Direction,
89        max_depth: usize,
90        max_visited: usize,
91    ) -> Vec<(String, String, String)> {
92        let labels = self.label_filter(label_filter);
93        let mut visited: HashSet<u32> = HashSet::new();
94        let mut frontier: Vec<u32> = Vec::new();
95        let mut edges = Vec::new();
96        // Each physical edge once: `Both` reaches an edge from both ends, and
97        // one triple can be stored under several collections.
98        let mut seen: HashSet<(u32, u32, u32)> = HashSet::new();
99
100        for &node in start_nodes {
101            if let Some(&id) = self.node_to_id.get(node)
102                && visited.insert(id)
103            {
104                frontier.push(id);
105            }
106        }
107
108        // Level by level, as `traverse_bfs_dense`: every frontier node's edges
109        // are recorded, then the level's new nodes are admitted in name order.
110        for _depth in 0..max_depth {
111            if frontier.is_empty() || visited.len() >= max_visited {
112                break;
113            }
114            let mut candidates: Vec<u32> = Vec::new();
115            for &node_id in &frontier {
116                self.record_access(node_id);
117                if matches!(direction, Direction::Out | Direction::Both) {
118                    for (lid, dst) in self.dense_iter_out(node_id) {
119                        if labels.keeps(lid) {
120                            if seen.insert((node_id, lid, dst)) {
121                                edges.push((
122                                    self.id_to_node[node_id as usize].clone(),
123                                    self.label_name(lid).to_string(),
124                                    self.id_to_node[dst as usize].clone(),
125                                ));
126                            }
127                            if !visited.contains(&dst) {
128                                candidates.push(dst);
129                            }
130                        }
131                    }
132                }
133                if matches!(direction, Direction::In | Direction::Both) {
134                    for (lid, src) in self.dense_iter_in(node_id) {
135                        if labels.keeps(lid) {
136                            if seen.insert((src, lid, node_id)) {
137                                edges.push((
138                                    self.id_to_node[src as usize].clone(),
139                                    self.label_name(lid).to_string(),
140                                    self.id_to_node[node_id as usize].clone(),
141                                ));
142                            }
143                            if !visited.contains(&src) {
144                                candidates.push(src);
145                            }
146                        }
147                    }
148                }
149            }
150            frontier = self.admit_by_name(candidates, &mut visited, max_visited);
151        }
152
153        // The last admitted level is never expanded. Its edges to admitted
154        // nodes, itself included, are still part of the subgraph.
155        for &node_id in &frontier {
156            if matches!(direction, Direction::Out | Direction::Both) {
157                for (lid, dst) in self.dense_iter_out(node_id) {
158                    if labels.keeps(lid)
159                        && visited.contains(&dst)
160                        && seen.insert((node_id, lid, dst))
161                    {
162                        edges.push((
163                            self.id_to_node[node_id as usize].clone(),
164                            self.label_name(lid).to_string(),
165                            self.id_to_node[dst as usize].clone(),
166                        ));
167                    }
168                }
169            }
170            if matches!(direction, Direction::In | Direction::Both) {
171                for (lid, src) in self.dense_iter_in(node_id) {
172                    if labels.keeps(lid)
173                        && visited.contains(&src)
174                        && seen.insert((src, lid, node_id))
175                    {
176                        edges.push((
177                            self.id_to_node[src as usize].clone(),
178                            self.label_name(lid).to_string(),
179                            self.id_to_node[node_id as usize].clone(),
180                        ));
181                    }
182                }
183            }
184        }
185
186        edges
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use crate::test_support::{chain_csr, test_memory};
194
195    #[test]
196    fn subgraph_materialization() {
197        let csr = chain_csr();
198        let edges = csr.subgraph(&["a"], &[], Direction::Out, 2, DEFAULT_MAX_VISITED, None);
199        assert_eq!(edges.len(), 3);
200        assert!(edges.contains(&("a".into(), "KNOWS".into(), "b".into())));
201        assert!(edges.contains(&("a".into(), "WORKS".into(), "e".into())));
202        assert!(edges.contains(&("b".into(), "KNOWS".into(), "c".into())));
203    }
204    #[test]
205    fn a_capped_subgraph_expands_only_admitted_levels() {
206        // `a` points at `z`, `m` and `b`, stored in that order.
207        let mut csr = CsrIndex::new(test_memory());
208        csr.add_edge("a", "L", "z").unwrap();
209        csr.add_edge("a", "L", "m").unwrap();
210        csr.add_edge("a", "L", "b").unwrap();
211        csr.add_edge("z", "L", "y").unwrap();
212        csr.add_edge("b", "L", "c").unwrap();
213        let mut edges = csr.subgraph(&["a"], &[], Direction::Out, 3, 3, None);
214        edges.sort();
215        // Level 1 fills the cap, so no level-1 node expands.
216        let expected: Vec<(String, String, String)> = ["b", "m", "z"]
217            .iter()
218            .map(|dst| ("a".to_string(), "L".to_string(), dst.to_string()))
219            .collect();
220        assert_eq!(edges, expected);
221    }
222
223    #[test]
224    fn both_directions_return_each_physical_edge_once() {
225        let mut csr = CsrIndex::new(test_memory());
226        csr.add_edge("a", "L", "b").unwrap();
227        csr.add_edge("b", "L", "a").unwrap();
228        csr.add_edge("a", "SELF", "a").unwrap();
229        csr.add_edge_in_collection("a", "L", "c", "first").unwrap();
230        csr.add_edge_in_collection("a", "L", "c", "second").unwrap();
231        let mut edges = csr.subgraph(&["a"], &[], Direction::Both, 3, DEFAULT_MAX_VISITED, None);
232        edges.sort();
233        let expected: Vec<(String, String, String)> = [
234            ("a", "L", "b"),
235            ("a", "L", "c"),
236            ("a", "SELF", "a"),
237            ("b", "L", "a"),
238        ]
239        .iter()
240        .map(|(s, l, d)| (s.to_string(), l.to_string(), d.to_string()))
241        .collect();
242        assert_eq!(edges, expected);
243    }
244
245    fn triples(edges: &[(&str, &str, &str)]) -> Vec<(String, String, String)> {
246        let mut out: Vec<(String, String, String)> = edges
247            .iter()
248            .map(|(s, l, d)| (s.to_string(), l.to_string(), d.to_string()))
249            .collect();
250        out.sort();
251        out
252    }
253
254    fn boundary_csr() -> CsrIndex {
255        let mut csr = CsrIndex::new(test_memory());
256        for (src, dst) in [
257            ("a", "b"),
258            ("a", "c"),
259            ("b", "c"),
260            ("c", "b"),
261            ("c", "a"),
262            ("c", "x"),
263        ] {
264            csr.add_edge(src, "L", dst).unwrap();
265        }
266        csr
267    }
268
269    #[test]
270    fn boundary_nodes_record_edges_among_admitted_nodes() {
271        let csr = boundary_csr();
272        // A non-empty overlay runs the string-keyed path.
273        let mut staged = GraphOverlayDelta::new();
274        staged.stage_tombstone("q", "L", "r");
275        for overlay in [None, Some(&staged)] {
276            let mut edges =
277                csr.subgraph(&["a"], &[], Direction::Out, 1, DEFAULT_MAX_VISITED, overlay);
278            edges.sort();
279            assert_eq!(
280                edges,
281                triples(&[
282                    ("a", "L", "b"),
283                    ("a", "L", "c"),
284                    ("b", "L", "c"),
285                    ("c", "L", "a"),
286                    ("c", "L", "b"),
287                ]),
288                "c -> x leaves the admitted set"
289            );
290        }
291    }
292
293    #[test]
294    fn boundary_edges_follow_the_walk_direction() {
295        let mut csr = CsrIndex::new(test_memory());
296        for (src, dst) in [("b", "a"), ("c", "a"), ("b", "c"), ("a", "z")] {
297            csr.add_edge(src, "L", dst).unwrap();
298        }
299        let mut edges = csr.subgraph(&["a"], &[], Direction::In, 1, DEFAULT_MAX_VISITED, None);
300        edges.sort();
301        assert_eq!(
302            edges,
303            triples(&[("b", "L", "a"), ("b", "L", "c"), ("c", "L", "a")])
304        );
305    }
306
307    #[test]
308    fn depth_zero_records_only_the_start_nodes_edges_among_themselves() {
309        let mut csr = CsrIndex::new(test_memory());
310        csr.add_edge("a", "L", "a").unwrap();
311        csr.add_edge("a", "L", "b").unwrap();
312        let edges = csr.subgraph(&["a"], &[], Direction::Out, 0, DEFAULT_MAX_VISITED, None);
313        assert_eq!(edges, triples(&[("a", "L", "a")]));
314    }
315
316    #[test]
317    fn a_label_set_returns_parallel_edges_of_both_labels() {
318        let mut csr = CsrIndex::new(test_memory());
319        csr.add_edge("a", "K", "b").unwrap();
320        csr.add_edge("a", "W", "b").unwrap();
321        csr.add_edge("a", "X", "c").unwrap();
322        let mut edges = csr.subgraph(
323            &["a"],
324            &["K", "W"],
325            Direction::Out,
326            2,
327            DEFAULT_MAX_VISITED,
328            None,
329        );
330        edges.sort();
331        let expected: Vec<(String, String, String)> = [("a", "K", "b"), ("a", "W", "b")]
332            .iter()
333            .map(|(s, l, d)| (s.to_string(), l.to_string(), d.to_string()))
334            .collect();
335        assert_eq!(edges, expected);
336    }
337}