arrow_graph/algorithms/
components.rs

1use arrow::record_batch::RecordBatch;
2use arrow::array::{StringArray, UInt32Array};
3use arrow::datatypes::{DataType, Field, Schema};
4use std::sync::Arc;
5use std::collections::HashMap;
6use crate::algorithms::{GraphAlgorithm, AlgorithmParams};
7use crate::graph::ArrowGraph;
8use crate::error::{GraphError, Result};
9
10/// Union-Find (Disjoint Set) data structure for efficient connected components
11#[derive(Debug)]
12struct UnionFind {
13    parent: HashMap<String, String>,
14    rank: HashMap<String, usize>,
15    component_sizes: HashMap<String, usize>,
16}
17
18impl UnionFind {
19    fn new() -> Self {
20        UnionFind {
21            parent: HashMap::new(),
22            rank: HashMap::new(),
23            component_sizes: HashMap::new(),
24        }
25    }
26    
27    fn make_set(&mut self, node: String) {
28        if !self.parent.contains_key(&node) {
29            self.parent.insert(node.clone(), node.clone());
30            self.rank.insert(node.clone(), 0);
31            self.component_sizes.insert(node.clone(), 1);
32        }
33    }
34    
35    fn find(&mut self, node: &str) -> Option<String> {
36        if !self.parent.contains_key(node) {
37            return None;
38        }
39        
40        // Path compression
41        let parent = self.parent.get(node).unwrap().clone();
42        if parent != node {
43            let root = self.find(&parent)?;
44            self.parent.insert(node.to_string(), root.clone());
45            Some(root)
46        } else {
47            Some(parent)
48        }
49    }
50    
51    fn union(&mut self, node1: &str, node2: &str) -> bool {
52        let root1 = match self.find(node1) {
53            Some(r) => r,
54            None => return false,
55        };
56        
57        let root2 = match self.find(node2) {
58            Some(r) => r,
59            None => return false,
60        };
61        
62        if root1 == root2 {
63            return false; // Already in same component
64        }
65        
66        // Union by rank
67        let rank1 = *self.rank.get(&root1).unwrap_or(&0);
68        let rank2 = *self.rank.get(&root2).unwrap_or(&0);
69        
70        let (new_root, old_root) = if rank1 > rank2 {
71            (root1, root2)
72        } else if rank1 < rank2 {
73            (root2, root1)
74        } else {
75            // Equal ranks, choose root1 and increment its rank
76            self.rank.insert(root1.clone(), rank1 + 1);
77            (root1, root2)
78        };
79        
80        // Update parent
81        self.parent.insert(old_root.clone(), new_root.clone());
82        
83        // Update component sizes
84        let size1 = *self.component_sizes.get(&new_root).unwrap_or(&0);
85        let size2 = *self.component_sizes.get(&old_root).unwrap_or(&0);
86        self.component_sizes.insert(new_root, size1 + size2);
87        self.component_sizes.remove(&old_root);
88        
89        true
90    }
91    
92    fn get_components(&mut self) -> HashMap<String, Vec<String>> {
93        let mut components: HashMap<String, Vec<String>> = HashMap::new();
94        
95        // Get all nodes and their root components
96        let nodes: Vec<String> = self.parent.keys().cloned().collect();
97        for node in nodes {
98            if let Some(root) = self.find(&node) {
99                components.entry(root).or_default().push(node);
100            }
101        }
102        
103        components
104    }
105    
106    fn component_count(&mut self) -> usize {
107        self.get_components().len()
108    }
109}
110
111pub struct WeaklyConnectedComponents;
112
113impl WeaklyConnectedComponents {
114    /// Find weakly connected components using Union-Find
115    fn compute_components(&self, graph: &ArrowGraph) -> Result<HashMap<String, u32>> {
116        let mut uf = UnionFind::new();
117        
118        // Initialize all nodes
119        for node_id in graph.node_ids() {
120            uf.make_set(node_id.clone());
121        }
122        
123        // Union nodes connected by edges (treat as undirected)
124        for node_id in graph.node_ids() {
125            if let Some(neighbors) = graph.neighbors(node_id) {
126                for neighbor in neighbors {
127                    uf.union(node_id, neighbor);
128                }
129            }
130        }
131        
132        // Assign component IDs
133        let components = uf.get_components();
134        let mut node_to_component: HashMap<String, u32> = HashMap::new();
135        
136        for (component_id, (_root, nodes)) in components.into_iter().enumerate() {
137            for node in nodes {
138                node_to_component.insert(node, component_id as u32);
139            }
140        }
141        
142        Ok(node_to_component)
143    }
144}
145
146impl GraphAlgorithm for WeaklyConnectedComponents {
147    fn execute(&self, graph: &ArrowGraph, _params: &AlgorithmParams) -> Result<RecordBatch> {
148        let component_map = self.compute_components(graph)?;
149        
150        if component_map.is_empty() {
151            // Return empty result with proper schema
152            let schema = Arc::new(Schema::new(vec![
153                Field::new("node_id", DataType::Utf8, false),
154                Field::new("component_id", DataType::UInt32, false),
155            ]));
156            
157            return RecordBatch::try_new(
158                schema,
159                vec![
160                    Arc::new(StringArray::from(Vec::<String>::new())),
161                    Arc::new(UInt32Array::from(Vec::<u32>::new())),
162                ],
163            ).map_err(GraphError::from);
164        }
165        
166        // Sort by component ID for consistent output
167        let mut sorted_nodes: Vec<(&String, &u32)> = component_map.iter().collect();
168        sorted_nodes.sort_by_key(|(_, &component_id)| component_id);
169        
170        let node_ids: Vec<String> = sorted_nodes.iter().map(|(node, _)| (*node).clone()).collect();
171        let component_ids: Vec<u32> = sorted_nodes.iter().map(|(_, &comp)| comp).collect();
172        
173        let schema = Arc::new(Schema::new(vec![
174            Field::new("node_id", DataType::Utf8, false),
175            Field::new("component_id", DataType::UInt32, false),
176        ]));
177        
178        RecordBatch::try_new(
179            schema,
180            vec![
181                Arc::new(StringArray::from(node_ids)),
182                Arc::new(UInt32Array::from(component_ids)),
183            ],
184        ).map_err(GraphError::from)
185    }
186    
187    fn name(&self) -> &'static str {
188        "weakly_connected_components"
189    }
190    
191    fn description(&self) -> &'static str {
192        "Find weakly connected components using Union-Find with path compression"
193    }
194}
195
196pub struct StronglyConnectedComponents;
197
198impl StronglyConnectedComponents {
199    /// Tarjan's algorithm for strongly connected components
200    fn tarjan_scc(&self, graph: &ArrowGraph) -> Result<HashMap<String, u32>> {
201        let mut index_counter = 0;
202        let mut stack = Vec::new();
203        let mut indices: HashMap<String, usize> = HashMap::new();
204        let mut lowlinks: HashMap<String, usize> = HashMap::new();
205        let mut on_stack: HashMap<String, bool> = HashMap::new();
206        let mut components: Vec<Vec<String>> = Vec::new();
207        
208        // Initialize
209        for node_id in graph.node_ids() {
210            on_stack.insert(node_id.clone(), false);
211        }
212        
213        // Run DFS from each unvisited node
214        for node_id in graph.node_ids() {
215            if !indices.contains_key(node_id) {
216                self.tarjan_strongconnect(
217                    node_id,
218                    graph,
219                    &mut index_counter,
220                    &mut stack,
221                    &mut indices,
222                    &mut lowlinks,
223                    &mut on_stack,
224                    &mut components,
225                )?;
226            }
227        }
228        
229        // Create component mapping
230        let mut node_to_component: HashMap<String, u32> = HashMap::new();
231        for (comp_id, component) in components.into_iter().enumerate() {
232            for node in component {
233                node_to_component.insert(node, comp_id as u32);
234            }
235        }
236        
237        Ok(node_to_component)
238    }
239    
240    fn tarjan_strongconnect(
241        &self,
242        node: &str,
243        graph: &ArrowGraph,
244        index_counter: &mut usize,
245        stack: &mut Vec<String>,
246        indices: &mut HashMap<String, usize>,
247        lowlinks: &mut HashMap<String, usize>,
248        on_stack: &mut HashMap<String, bool>,
249        components: &mut Vec<Vec<String>>,
250    ) -> Result<()> {
251        // Set the depth index for this node
252        indices.insert(node.to_string(), *index_counter);
253        lowlinks.insert(node.to_string(), *index_counter);
254        *index_counter += 1;
255        
256        // Push node onto stack
257        stack.push(node.to_string());
258        on_stack.insert(node.to_string(), true);
259        
260        // Consider successors
261        if let Some(neighbors) = graph.neighbors(node) {
262            for neighbor in neighbors {
263                if !indices.contains_key(neighbor) {
264                    // Successor has not yet been visited; recurse on it
265                    self.tarjan_strongconnect(
266                        neighbor,
267                        graph,
268                        index_counter,
269                        stack,
270                        indices,
271                        lowlinks,
272                        on_stack,
273                        components,
274                    )?;
275                    
276                    let neighbor_lowlink = *lowlinks.get(neighbor).unwrap_or(&0);
277                    let current_lowlink = *lowlinks.get(node).unwrap_or(&0);
278                    lowlinks.insert(node.to_string(), current_lowlink.min(neighbor_lowlink));
279                } else if *on_stack.get(neighbor).unwrap_or(&false) {
280                    // Successor is in stack and hence in the current SCC
281                    let neighbor_index = *indices.get(neighbor).unwrap_or(&0);
282                    let current_lowlink = *lowlinks.get(node).unwrap_or(&0);
283                    lowlinks.insert(node.to_string(), current_lowlink.min(neighbor_index));
284                }
285            }
286        }
287        
288        // If node is a root node, pop the stack and create an SCC
289        let node_index = *indices.get(node).unwrap_or(&0);
290        let node_lowlink = *lowlinks.get(node).unwrap_or(&0);
291        
292        if node_lowlink == node_index {
293            let mut component = Vec::new();
294            loop {
295                if let Some(w) = stack.pop() {
296                    on_stack.insert(w.clone(), false);
297                    component.push(w.clone());
298                    if w == node {
299                        break;
300                    }
301                } else {
302                    break;
303                }
304            }
305            components.push(component);
306        }
307        
308        Ok(())
309    }
310}
311
312impl GraphAlgorithm for StronglyConnectedComponents {
313    fn execute(&self, graph: &ArrowGraph, _params: &AlgorithmParams) -> Result<RecordBatch> {
314        let component_map = self.tarjan_scc(graph)?;
315        
316        if component_map.is_empty() {
317            // Return empty result with proper schema
318            let schema = Arc::new(Schema::new(vec![
319                Field::new("node_id", DataType::Utf8, false),
320                Field::new("component_id", DataType::UInt32, false),
321            ]));
322            
323            return RecordBatch::try_new(
324                schema,
325                vec![
326                    Arc::new(StringArray::from(Vec::<String>::new())),
327                    Arc::new(UInt32Array::from(Vec::<u32>::new())),
328                ],
329            ).map_err(GraphError::from);
330        }
331        
332        // Sort by component ID for consistent output
333        let mut sorted_nodes: Vec<(&String, &u32)> = component_map.iter().collect();
334        sorted_nodes.sort_by_key(|(_, &component_id)| component_id);
335        
336        let node_ids: Vec<String> = sorted_nodes.iter().map(|(node, _)| (*node).clone()).collect();
337        let component_ids: Vec<u32> = sorted_nodes.iter().map(|(_, &comp)| comp).collect();
338        
339        let schema = Arc::new(Schema::new(vec![
340            Field::new("node_id", DataType::Utf8, false),
341            Field::new("component_id", DataType::UInt32, false),
342        ]));
343        
344        RecordBatch::try_new(
345            schema,
346            vec![
347                Arc::new(StringArray::from(node_ids)),
348                Arc::new(UInt32Array::from(component_ids)),
349            ],
350        ).map_err(GraphError::from)
351    }
352    
353    fn name(&self) -> &'static str {
354        "strongly_connected_components"
355    }
356    
357    fn description(&self) -> &'static str {
358        "Find strongly connected components using Tarjan's algorithm"
359    }
360}
arrow_graph/algorithms/components.rs

arrow_graph/algorithms/
components.rs
