nusy_graph_query/
traversal.rs

1//! Graph traversal — generic BFS/DFS over Arrow edge RecordBatches.
2//!
3//! Parameterized by column indices so the same traversal logic works for:
4//! - nusy-kanban: `depends_on` column in items table
5//! - nusy-codegraph: `source_id`/`target_id` columns in edges table
6//! - Being cognitive graphs: causal chains, learning paths
7
8use arrow::array::{Array, RecordBatch, StringArray};
9use std::collections::{HashMap, HashSet, VecDeque};
10
11/// Configuration for edge column layout in a RecordBatch.
12///
13/// Different consumers have different schemas — this struct lets you
14/// specify which columns contain source IDs, target IDs, and predicates.
15#[derive(Debug, Clone)]
16pub struct EdgeSchema {
17    /// Column index for source node ID.
18    pub source_col: usize,
19    /// Column index for target node ID.
20    pub target_col: usize,
21    /// Column index for edge predicate/type (optional — None means "all edges").
22    pub predicate_col: Option<usize>,
23}
24
25/// Direction of traversal relative to a node.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum Direction {
28    /// Follow edges where the node is the source (find targets).
29    Forward,
30    /// Follow edges where the node is the target (find sources).
31    Reverse,
32}
33
34/// Build an adjacency list from an edge RecordBatch.
35///
36/// If `predicate_filter` is provided, only edges with matching predicate
37/// are included. Direction determines how source/target map to from/to.
38pub fn build_adjacency(
39    edges: &RecordBatch,
40    schema: &EdgeSchema,
41    direction: Direction,
42    predicate_filter: Option<&str>,
43) -> HashMap<String, Vec<String>> {
44    let mut adj: HashMap<String, Vec<String>> = HashMap::new();
45
46    if edges.num_rows() == 0 {
47        return adj;
48    }
49
50    let Some(sources) = edges
51        .column(schema.source_col)
52        .as_any()
53        .downcast_ref::<StringArray>()
54    else {
55        return adj;
56    };
57    let Some(targets) = edges
58        .column(schema.target_col)
59        .as_any()
60        .downcast_ref::<StringArray>()
61    else {
62        return adj;
63    };
64
65    let predicates = schema
66        .predicate_col
67        .and_then(|col| edges.column(col).as_any().downcast_ref::<StringArray>());
68
69    for i in 0..edges.num_rows() {
70        // Filter by predicate if specified
71        if let (Some(filter), Some(pred_col)) = (predicate_filter, predicates)
72            && (pred_col.is_null(i) || pred_col.value(i) != filter)
73        {
74            continue;
75        }
76
77        if sources.is_null(i) || targets.is_null(i) {
78            continue;
79        }
80
81        match direction {
82            Direction::Forward => {
83                adj.entry(sources.value(i).to_string())
84                    .or_default()
85                    .push(targets.value(i).to_string());
86            }
87            Direction::Reverse => {
88                adj.entry(targets.value(i).to_string())
89                    .or_default()
90                    .push(sources.value(i).to_string());
91            }
92        }
93    }
94
95    adj
96}
97
98/// A traversal result — node ID with its depth from the start.
99#[derive(Debug, Clone)]
100pub struct TraversalNode {
101    pub id: String,
102    pub depth: usize,
103}
104
105/// BFS traversal from `start_id` following edges up to `max_depth` hops.
106///
107/// Returns nodes in breadth-first order (excluding the start node).
108/// Works with any edge RecordBatch via `EdgeSchema` configuration.
109pub fn bfs(
110    start_id: &str,
111    edges: &RecordBatch,
112    schema: &EdgeSchema,
113    direction: Direction,
114    predicate_filter: Option<&str>,
115    max_depth: usize,
116) -> Vec<TraversalNode> {
117    let adj = build_adjacency(edges, schema, direction, predicate_filter);
118
119    let mut visited: HashSet<String> = HashSet::new();
120    visited.insert(start_id.to_string());
121    let mut queue: VecDeque<(String, usize)> = VecDeque::new();
122    queue.push_back((start_id.to_string(), 0));
123    let mut result: Vec<TraversalNode> = Vec::new();
124
125    while let Some((current, depth)) = queue.pop_front() {
126        if depth >= max_depth {
127            continue;
128        }
129        if let Some(neighbors) = adj.get(&current) {
130            for neighbor in neighbors {
131                if visited.insert(neighbor.clone()) {
132                    result.push(TraversalNode {
133                        id: neighbor.clone(),
134                        depth: depth + 1,
135                    });
136                    queue.push_back((neighbor.clone(), depth + 1));
137                }
138            }
139        }
140    }
141
142    result
143}
144
145/// Build an adjacency list from a simple string list column.
146///
147/// This variant handles the kanban `depends_on` pattern where dependencies
148/// are stored as a List<Utf8> column on each item, rather than in a
149/// separate edges table.
150pub fn build_adjacency_from_list(
151    batch: &RecordBatch,
152    id_col: usize,
153    list_col: usize,
154    direction: Direction,
155) -> HashMap<String, Vec<String>> {
156    use arrow::array::ListArray;
157
158    let mut adj: HashMap<String, Vec<String>> = HashMap::new();
159
160    if batch.num_rows() == 0 {
161        return adj;
162    }
163
164    let Some(ids) = batch.column(id_col).as_any().downcast_ref::<StringArray>() else {
165        return adj;
166    };
167    let Some(lists) = batch.column(list_col).as_any().downcast_ref::<ListArray>() else {
168        return adj;
169    };
170
171    for i in 0..batch.num_rows() {
172        if lists.is_null(i) {
173            continue;
174        }
175        let values = lists.value(i);
176        let Some(str_arr) = values.as_any().downcast_ref::<StringArray>() else {
177            continue;
178        };
179
180        let id = ids.value(i);
181        for j in 0..str_arr.len() {
182            if str_arr.is_null(j) {
183                continue;
184            }
185            let dep = str_arr.value(j);
186            match direction {
187                Direction::Forward => {
188                    adj.entry(id.to_string()).or_default().push(dep.to_string());
189                }
190                Direction::Reverse => {
191                    adj.entry(dep.to_string()).or_default().push(id.to_string());
192                }
193            }
194        }
195    }
196
197    adj
198}
199
200/// BFS traversal using a pre-built adjacency list.
201///
202/// Useful when the adjacency list comes from `build_adjacency_from_list`
203/// or has been cached.
204pub fn bfs_with_adjacency(
205    start_id: &str,
206    adj: &HashMap<String, Vec<String>>,
207    max_depth: usize,
208) -> Vec<TraversalNode> {
209    let mut visited: HashSet<String> = HashSet::new();
210    visited.insert(start_id.to_string());
211    let mut queue: VecDeque<(String, usize)> = VecDeque::new();
212    queue.push_back((start_id.to_string(), 0));
213    let mut result: Vec<TraversalNode> = Vec::new();
214
215    while let Some((current, depth)) = queue.pop_front() {
216        if depth >= max_depth {
217            continue;
218        }
219        if let Some(neighbors) = adj.get(&current) {
220            for neighbor in neighbors {
221                if visited.insert(neighbor.clone()) {
222                    result.push(TraversalNode {
223                        id: neighbor.clone(),
224                        depth: depth + 1,
225                    });
226                    queue.push_back((neighbor.clone(), depth + 1));
227                }
228            }
229        }
230    }
231
232    result
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use arrow::array::{ListBuilder, StringBuilder};
239    use arrow::datatypes::{DataType, Field, Schema};
240    use std::sync::Arc;
241
242    /// Build a simple edges RecordBatch: source, target, predicate.
243    fn make_edges(triples: &[(&str, &str, &str)]) -> RecordBatch {
244        let schema = Arc::new(Schema::new(vec![
245            Field::new("source", DataType::Utf8, false),
246            Field::new("target", DataType::Utf8, false),
247            Field::new("predicate", DataType::Utf8, false),
248        ]));
249        let sources: Vec<&str> = triples.iter().map(|(s, _, _)| *s).collect();
250        let targets: Vec<&str> = triples.iter().map(|(_, t, _)| *t).collect();
251        let preds: Vec<&str> = triples.iter().map(|(_, _, p)| *p).collect();
252
253        RecordBatch::try_new(
254            schema,
255            vec![
256                Arc::new(StringArray::from(sources)),
257                Arc::new(StringArray::from(targets)),
258                Arc::new(StringArray::from(preds)),
259            ],
260        )
261        .expect("build edges batch")
262    }
263
264    fn edge_schema() -> EdgeSchema {
265        EdgeSchema {
266            source_col: 0,
267            target_col: 1,
268            predicate_col: Some(2),
269        }
270    }
271
272    #[test]
273    fn test_build_adjacency_forward() {
274        let edges = make_edges(&[
275            ("A", "B", "calls"),
276            ("A", "C", "calls"),
277            ("B", "C", "calls"),
278        ]);
279        let adj = build_adjacency(&edges, &edge_schema(), Direction::Forward, Some("calls"));
280
281        assert_eq!(adj.get("A").unwrap().len(), 2);
282        assert_eq!(adj.get("B").unwrap().len(), 1);
283        assert!(adj.get("C").is_none());
284    }
285
286    #[test]
287    fn test_build_adjacency_reverse() {
288        let edges = make_edges(&[
289            ("A", "B", "calls"),
290            ("A", "C", "calls"),
291            ("B", "C", "calls"),
292        ]);
293        let adj = build_adjacency(&edges, &edge_schema(), Direction::Reverse, Some("calls"));
294
295        assert!(adj.get("A").is_none()); // nothing calls A
296        assert_eq!(adj.get("B").unwrap(), &["A"]);
297        assert_eq!(adj.get("C").unwrap().len(), 2); // A and B call C
298    }
299
300    #[test]
301    fn test_build_adjacency_predicate_filter() {
302        let edges = make_edges(&[
303            ("A", "B", "calls"),
304            ("A", "C", "tests"),
305            ("B", "C", "calls"),
306        ]);
307        let adj = build_adjacency(&edges, &edge_schema(), Direction::Forward, Some("calls"));
308
309        assert_eq!(adj.get("A").unwrap(), &["B"]); // only "calls" edges
310    }
311
312    #[test]
313    fn test_build_adjacency_no_filter() {
314        let edges = make_edges(&[
315            ("A", "B", "calls"),
316            ("A", "C", "tests"),
317            ("B", "C", "calls"),
318        ]);
319        let adj = build_adjacency(&edges, &edge_schema(), Direction::Forward, None);
320
321        assert_eq!(adj.get("A").unwrap().len(), 2); // both edges
322    }
323
324    #[test]
325    fn test_bfs_depth_1() {
326        let edges = make_edges(&[("A", "B", "dep"), ("A", "C", "dep"), ("B", "D", "dep")]);
327        let result = bfs(
328            "A",
329            &edges,
330            &edge_schema(),
331            Direction::Forward,
332            Some("dep"),
333            1,
334        );
335
336        assert_eq!(result.len(), 2); // B and C
337        assert!(result.iter().all(|n| n.depth == 1));
338    }
339
340    #[test]
341    fn test_bfs_depth_2() {
342        let edges = make_edges(&[("A", "B", "dep"), ("B", "C", "dep"), ("C", "D", "dep")]);
343        let result = bfs(
344            "A",
345            &edges,
346            &edge_schema(),
347            Direction::Forward,
348            Some("dep"),
349            2,
350        );
351
352        assert_eq!(result.len(), 2); // B (depth 1) and C (depth 2)
353        assert_eq!(result[0].id, "B");
354        assert_eq!(result[0].depth, 1);
355        assert_eq!(result[1].id, "C");
356        assert_eq!(result[1].depth, 2);
357    }
358
359    #[test]
360    fn test_bfs_depth_0() {
361        let edges = make_edges(&[("A", "B", "dep")]);
362        let result = bfs(
363            "A",
364            &edges,
365            &edge_schema(),
366            Direction::Forward,
367            Some("dep"),
368            0,
369        );
370        assert!(result.is_empty());
371    }
372
373    #[test]
374    fn test_bfs_reverse() {
375        let edges = make_edges(&[("A", "C", "dep"), ("B", "C", "dep")]);
376        let result = bfs(
377            "C",
378            &edges,
379            &edge_schema(),
380            Direction::Reverse,
381            Some("dep"),
382            1,
383        );
384
385        assert_eq!(result.len(), 2); // A and B
386    }
387
388    #[test]
389    fn test_bfs_cycle_safe() {
390        let edges = make_edges(&[("A", "B", "dep"), ("B", "A", "dep")]);
391        let result = bfs(
392            "A",
393            &edges,
394            &edge_schema(),
395            Direction::Forward,
396            Some("dep"),
397            10,
398        );
399
400        assert_eq!(result.len(), 1); // only B (A is already visited)
401    }
402
403    #[test]
404    fn test_bfs_empty_edges() {
405        let edges = make_edges(&[]);
406        let result = bfs("A", &edges, &edge_schema(), Direction::Forward, None, 5);
407        assert!(result.is_empty());
408    }
409
410    #[test]
411    fn test_build_adjacency_from_list() {
412        // Build a batch with id + depends_on (list column)
413        let schema = Arc::new(Schema::new(vec![
414            Field::new("id", DataType::Utf8, false),
415            Field::new(
416                "depends_on",
417                DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
418                false,
419            ),
420        ]));
421
422        let ids = StringArray::from(vec!["A", "B", "C"]);
423        let mut list_builder = ListBuilder::new(StringBuilder::new());
424        // A depends on B and C
425        list_builder.values().append_value("B");
426        list_builder.values().append_value("C");
427        list_builder.append(true);
428        // B depends on C
429        list_builder.values().append_value("C");
430        list_builder.append(true);
431        // C depends on nothing
432        list_builder.append(true);
433
434        let batch =
435            RecordBatch::try_new(schema, vec![Arc::new(ids), Arc::new(list_builder.finish())])
436                .unwrap();
437
438        let adj = build_adjacency_from_list(&batch, 0, 1, Direction::Forward);
439        assert_eq!(adj.get("A").unwrap().len(), 2);
440        assert_eq!(adj.get("B").unwrap().len(), 1);
441        assert!(adj.get("C").is_none()); // C has no deps
442
443        // Reverse: "who depends on X?"
444        let adj_rev = build_adjacency_from_list(&batch, 0, 1, Direction::Reverse);
445        assert!(adj_rev.get("A").is_none()); // nobody depends on A
446        assert_eq!(adj_rev.get("B").unwrap(), &["A"]);
447        assert_eq!(adj_rev.get("C").unwrap().len(), 2); // A and B depend on C
448    }
449
450    #[test]
451    fn test_bfs_with_adjacency() {
452        let mut adj = HashMap::new();
453        adj.insert("A".to_string(), vec!["B".to_string(), "C".to_string()]);
454        adj.insert("B".to_string(), vec!["D".to_string()]);
455
456        let result = bfs_with_adjacency("A", &adj, 2);
457        assert_eq!(result.len(), 3); // B, C, D
458    }
459}
nusy_graph_query/traversal.rs

nusy_graph_query/
traversal.rs
