arrow_graph/sql/
recursive_cte.rs

1use datafusion::error::Result as DataFusionResult;
2use std::collections::{HashMap, HashSet, VecDeque};
3
4/// Recursive CTE (Common Table Expression) processor for graph path queries
5/// Supports patterns like: WITH RECURSIVE paths AS (...) SELECT ... FROM paths
6pub struct RecursiveCteProcessor;
7
8/// Represents a path in the graph
9#[derive(Debug, Clone, PartialEq)]
10pub struct GraphPath {
11    pub nodes: Vec<String>,
12    pub total_weight: f64,
13    pub depth: usize,
14}
15
16/// Configuration for recursive path queries
17#[derive(Debug, Clone)]
18pub struct RecursiveQueryConfig {
19    pub max_depth: usize,
20    pub include_cycles: bool,
21    pub weight_threshold: Option<f64>,
22    pub max_paths: Option<usize>,
23}
24
25impl Default for RecursiveQueryConfig {
26    fn default() -> Self {
27        Self {
28            max_depth: 10,
29            include_cycles: false,
30            weight_threshold: None,
31            max_paths: Some(1000),
32        }
33    }
34}
35
36impl RecursiveCteProcessor {
37    pub fn new() -> Self {
38        Self
39    }
40
41    /// Process a recursive path query
42    /// Example: Find all paths from node A to node B with max depth 5
43    pub fn find_recursive_paths(
44        &self,
45        edges: &[(String, String, f64)], // (source, target, weight)
46        start_node: &str,
47        end_node: Option<&str>,
48        config: &RecursiveQueryConfig,
49    ) -> DataFusionResult<Vec<GraphPath>> {
50        let mut paths = Vec::new();
51        let mut queue = VecDeque::new();
52        let visited_in_path = HashSet::new();
53
54        // Build adjacency map for efficient lookup
55        let mut adjacency_map: HashMap<String, Vec<(String, f64)>> = HashMap::new();
56        for (source, target, weight) in edges {
57            adjacency_map
58                .entry(source.clone())
59                .or_insert_with(Vec::new)
60                .push((target.clone(), *weight));
61        }
62
63        // Initialize with starting node
64        let initial_path = GraphPath {
65            nodes: vec![start_node.to_string()],
66            total_weight: 0.0,
67            depth: 0,
68        };
69        queue.push_back((initial_path, visited_in_path.clone()));
70
71        while let Some((current_path, path_visited)) = queue.pop_front() {
72            // Check if we've reached the target (if specified)
73            if let Some(target) = end_node {
74                if current_path.nodes.last() == Some(&target.to_string()) {
75                    paths.push(current_path.clone());
76                    if let Some(max_paths) = config.max_paths {
77                        if paths.len() >= max_paths {
78                            break;
79                        }
80                    }
81                    continue;
82                }
83            } else {
84                // If no specific target, include all paths up to max depth
85                paths.push(current_path.clone());
86                if let Some(max_paths) = config.max_paths {
87                    if paths.len() >= max_paths {
88                        break;
89                    }
90                }
91            }
92
93            // Check depth limit
94            if current_path.depth >= config.max_depth {
95                continue;
96            }
97
98            // Get current node
99            let current_node = current_path.nodes.last().unwrap();
100
101            // Explore neighbors
102            if let Some(neighbors) = adjacency_map.get(current_node) {
103                for (neighbor, weight) in neighbors {
104                    // Check for cycles
105                    if !config.include_cycles && current_path.nodes.contains(neighbor) {
106                        continue;
107                    }
108
109                    // Check weight threshold
110                    if let Some(threshold) = config.weight_threshold {
111                        if current_path.total_weight + weight > threshold {
112                            continue;
113                        }
114                    }
115
116                    // Create new path
117                    let mut new_nodes = current_path.nodes.clone();
118                    new_nodes.push(neighbor.clone());
119
120                    let new_path = GraphPath {
121                        nodes: new_nodes,
122                        total_weight: current_path.total_weight + weight,
123                        depth: current_path.depth + 1,
124                    };
125
126                    // Update visited set for this path
127                    let mut new_visited = path_visited.clone();
128                    new_visited.insert(neighbor.clone());
129
130                    queue.push_back((new_path, new_visited));
131                }
132            }
133        }
134
135        Ok(paths)
136    }
137
138    /// Find shortest paths using recursive approach (for CTE queries)
139    pub fn find_shortest_paths(
140        &self,
141        edges: &[(String, String, f64)],
142        start_node: &str,
143        end_node: &str,
144        max_depth: usize,
145    ) -> DataFusionResult<Vec<GraphPath>> {
146        let config = RecursiveQueryConfig {
147            max_depth,
148            include_cycles: false,
149            weight_threshold: None,
150            max_paths: Some(10), // Limit to top 10 shortest paths
151        };
152
153        let mut all_paths = self.find_recursive_paths(edges, start_node, Some(end_node), &config)?;
154
155        // Sort by total weight to get shortest paths first
156        all_paths.sort_by(|a, b| a.total_weight.partial_cmp(&b.total_weight).unwrap_or(std::cmp::Ordering::Equal));
157
158        Ok(all_paths)
159    }
160
161    /// Find all reachable nodes from a starting node (for CTE queries)
162    pub fn find_reachable_nodes(
163        &self,
164        edges: &[(String, String, f64)],
165        start_node: &str,
166        max_depth: usize,
167    ) -> DataFusionResult<Vec<String>> {
168        let config = RecursiveQueryConfig {
169            max_depth,
170            include_cycles: false,
171            weight_threshold: None,
172            max_paths: None, // No limit on paths for reachability
173        };
174
175        let paths = self.find_recursive_paths(edges, start_node, None, &config)?;
176
177        // Collect all unique nodes reachable from start
178        let mut reachable: HashSet<String> = HashSet::new();
179        for path in paths {
180            for node in path.nodes {
181                reachable.insert(node);
182            }
183        }
184
185        // Remove the starting node itself
186        reachable.remove(start_node);
187
188        Ok(reachable.into_iter().collect())
189    }
190
191    /// Process a WITH RECURSIVE query (simplified parser)
192    /// Example: "WITH RECURSIVE paths(source, target, depth) AS (
193    ///            SELECT source, target, 1 FROM edges WHERE source = 'A'
194    ///            UNION ALL
195    ///            SELECT e.source, e.target, p.depth + 1 
196    ///            FROM edges e JOIN paths p ON e.source = p.target 
197    ///            WHERE p.depth < 5
198    ///          ) SELECT * FROM paths"
199    pub fn parse_and_execute_recursive_cte(
200        &self,
201        _cte_query: &str,
202        edges: &[(String, String, f64)],
203    ) -> DataFusionResult<Vec<GraphPath>> {
204        // This is a simplified implementation
205        // A full implementation would parse the actual SQL CTE syntax
206        
207        // For now, return paths from a default starting node
208        let config = RecursiveQueryConfig::default();
209        
210        // Use first edge's source as starting point if available
211        if let Some((start_node, _, _)) = edges.first() {
212            self.find_recursive_paths(edges, start_node, None, &config)
213        } else {
214            Ok(Vec::new())
215        }
216    }
217}
218
219impl Default for RecursiveCteProcessor {
220    fn default() -> Self {
221        Self::new()
222    }
223}
224
225/// Utility functions for common recursive graph queries
226pub struct RecursiveGraphQueries;
227
228impl RecursiveGraphQueries {
229    /// Standard "find all paths" CTE query
230    pub fn all_paths_query(
231        processor: &RecursiveCteProcessor,
232        edges: &[(String, String, f64)],
233        start: &str,
234        end: &str,
235        max_depth: usize,
236    ) -> DataFusionResult<Vec<GraphPath>> {
237        let config = RecursiveQueryConfig {
238            max_depth,
239            include_cycles: false,
240            weight_threshold: None,
241            max_paths: Some(100),
242        };
243        
244        processor.find_recursive_paths(edges, start, Some(end), &config)
245    }
246
247    /// Standard "find connected component" CTE query
248    pub fn connected_component_query(
249        processor: &RecursiveCteProcessor,
250        edges: &[(String, String, f64)],
251        start: &str,
252    ) -> DataFusionResult<Vec<String>> {
253        processor.find_reachable_nodes(edges, start, 50) // Max depth 50 for components
254    }
255
256    /// Standard "shortest path" CTE query
257    pub fn shortest_path_query(
258        processor: &RecursiveCteProcessor,
259        edges: &[(String, String, f64)],
260        start: &str,
261        end: &str,
262    ) -> DataFusionResult<Option<GraphPath>> {
263        let paths = processor.find_shortest_paths(edges, start, end, 20)?;
264        Ok(paths.into_iter().next())
265    }
266
267    /// Standard "paths within distance" CTE query
268    pub fn paths_within_distance_query(
269        processor: &RecursiveCteProcessor,
270        edges: &[(String, String, f64)],
271        start: &str,
272        max_distance: f64,
273    ) -> DataFusionResult<Vec<GraphPath>> {
274        let config = RecursiveQueryConfig {
275            max_depth: 20,
276            include_cycles: false,
277            weight_threshold: Some(max_distance),
278            max_paths: Some(1000),
279        };
280        
281        processor.find_recursive_paths(edges, start, None, &config)
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    fn create_test_edges() -> Vec<(String, String, f64)> {
290        vec![
291            ("A".to_string(), "B".to_string(), 1.0),
292            ("B".to_string(), "C".to_string(), 2.0),
293            ("C".to_string(), "D".to_string(), 1.0),
294            ("A".to_string(), "C".to_string(), 4.0),
295            ("B".to_string(), "D".to_string(), 3.0),
296        ]
297    }
298
299    #[test]
300    fn test_find_recursive_paths() {
301        let processor = RecursiveCteProcessor::new();
302        let edges = create_test_edges();
303        let config = RecursiveQueryConfig::default();
304
305        let paths = processor.find_recursive_paths(&edges, "A", Some("D"), &config).unwrap();
306        
307        assert!(!paths.is_empty());
308        
309        // Should find multiple paths from A to D
310        let path_endpoints: Vec<_> = paths.iter()
311            .map(|p| (p.nodes.first().unwrap(), p.nodes.last().unwrap()))
312            .collect();
313        
314        for (start, end) in path_endpoints {
315            assert_eq!(start, "A");
316            assert_eq!(end, "D");
317        }
318    }
319
320    #[test]
321    fn test_find_shortest_paths() {
322        let processor = RecursiveCteProcessor::new();
323        let edges = create_test_edges();
324
325        let paths = processor.find_shortest_paths(&edges, "A", "D", 5).unwrap();
326        
327        assert!(!paths.is_empty());
328        
329        // Paths should be sorted by weight
330        for i in 1..paths.len() {
331            assert!(paths[i-1].total_weight <= paths[i].total_weight);
332        }
333    }
334
335    #[test]
336    fn test_find_reachable_nodes() {
337        let processor = RecursiveCteProcessor::new();
338        let edges = create_test_edges();
339
340        let reachable = processor.find_reachable_nodes(&edges, "A", 5).unwrap();
341        
342        assert!(reachable.contains(&"B".to_string()));
343        assert!(reachable.contains(&"C".to_string()));
344        assert!(reachable.contains(&"D".to_string()));
345        assert!(!reachable.contains(&"A".to_string())); // Shouldn't include start node itself
346    }
347
348    #[test]
349    fn test_recursive_query_utilities() {
350        let processor = RecursiveCteProcessor::new();
351        let edges = create_test_edges();
352
353        // Test all paths query
354        let all_paths = RecursiveGraphQueries::all_paths_query(&processor, &edges, "A", "D", 5).unwrap();
355        assert!(!all_paths.is_empty());
356
357        // Test connected component query
358        let component = RecursiveGraphQueries::connected_component_query(&processor, &edges, "A").unwrap();
359        assert!(component.len() >= 3); // Should reach B, C, D
360
361        // Test shortest path query
362        let shortest = RecursiveGraphQueries::shortest_path_query(&processor, &edges, "A", "D").unwrap();
363        assert!(shortest.is_some());
364
365        // Test paths within distance query
366        let within_distance = RecursiveGraphQueries::paths_within_distance_query(&processor, &edges, "A", 5.0).unwrap();
367        assert!(!within_distance.is_empty());
368    }
369
370    #[test]
371    fn test_cycle_detection() {
372        let processor = RecursiveCteProcessor::new();
373        
374        // Create edges with a cycle
375        let edges = vec![
376            ("A".to_string(), "B".to_string(), 1.0),
377            ("B".to_string(), "C".to_string(), 1.0),
378            ("C".to_string(), "A".to_string(), 1.0), // Creates cycle
379        ];
380
381        let config = RecursiveQueryConfig {
382            max_depth: 5,
383            include_cycles: false,
384            weight_threshold: None,
385            max_paths: Some(100),
386        };
387
388        let paths = processor.find_recursive_paths(&edges, "A", None, &config).unwrap();
389        
390        // Without cycles, should not have infinite paths
391        assert!(paths.len() < 100);
392        
393        // All paths should have unique nodes (no cycles)
394        for path in paths {
395            let mut unique_nodes = HashSet::new();
396            for node in &path.nodes {
397                assert!(unique_nodes.insert(node.clone()), "Found cycle in path: {:?}", path.nodes);
398            }
399        }
400    }
401
402    #[test]
403    fn test_max_depth_limit() {
404        let processor = RecursiveCteProcessor::new();
405        let edges = create_test_edges();
406        
407        let config = RecursiveQueryConfig {
408            max_depth: 2,
409            include_cycles: false,
410            weight_threshold: None,
411            max_paths: None,
412        };
413
414        let paths = processor.find_recursive_paths(&edges, "A", None, &config).unwrap();
415        
416        // All paths should respect max depth
417        for path in paths {
418            assert!(path.depth <= 2, "Path exceeded max depth: {:?}", path);
419        }
420    }
421}
arrow_graph/sql/recursive_cte.rs

arrow_graph/sql/
recursive_cte.rs
