arrow_graph/graph/
indexes.rs

1use arrow::record_batch::RecordBatch;
2use arrow::array::{Array, StringArray, Float64Array};
3use hashbrown::HashMap;
4use crate::error::{GraphError, Result};
5
6#[derive(Debug, Clone)]
7pub struct GraphIndexes {
8    pub adjacency_list: HashMap<String, Vec<String>>,
9    pub reverse_adjacency_list: HashMap<String, Vec<String>>,
10    pub node_index: HashMap<String, usize>,
11    pub edge_weights: HashMap<(String, String), f64>,
12    pub node_count: usize,
13    pub edge_count: usize,
14}
15
16impl GraphIndexes {
17    pub fn build(nodes: &RecordBatch, edges: &RecordBatch) -> Result<Self> {
18        let mut adjacency_list: HashMap<String, Vec<String>> = HashMap::new();
19        let mut reverse_adjacency_list: HashMap<String, Vec<String>> = HashMap::new();
20        let mut node_index: HashMap<String, usize> = HashMap::new();
21        let mut edge_weights: HashMap<(String, String), f64> = HashMap::new();
22        
23        // Build node index from nodes RecordBatch
24        // Expected schema: id (String), [label (String)], [properties (JSON)]
25        if nodes.num_columns() > 0 {
26            let id_column = nodes.column(0);
27            if let Some(string_array) = id_column.as_any().downcast_ref::<StringArray>() {
28                for (idx, node_id_opt) in string_array.iter().enumerate() {
29                    if let Some(node_id) = node_id_opt {
30                        node_index.insert(node_id.to_string(), idx);
31                        adjacency_list.insert(node_id.to_string(), Vec::new());
32                        reverse_adjacency_list.insert(node_id.to_string(), Vec::new());
33                    }
34                }
35            } else {
36                return Err(GraphError::graph_construction(
37                    "First column of nodes table must be String (node ID)"
38                ));
39            }
40        }
41        
42        // Build adjacency lists from edges RecordBatch
43        // Expected schema: source (String), target (String), [weight (Float64)], [label (String)]
44        if edges.num_columns() >= 2 {
45            let source_column = edges.column(0);
46            let target_column = edges.column(1);
47            
48            let source_array = source_column.as_any().downcast_ref::<StringArray>()
49                .ok_or_else(|| GraphError::graph_construction(
50                    "First column of edges table must be String (source)"
51                ))?;
52            
53            let target_array = target_column.as_any().downcast_ref::<StringArray>()
54                .ok_or_else(|| GraphError::graph_construction(
55                    "Second column of edges table must be String (target)"
56                ))?;
57            
58            // Handle optional weight column
59            let weight_array = if edges.num_columns() >= 3 {
60                edges.column(2).as_any().downcast_ref::<Float64Array>()
61            } else {
62                None
63            };
64            
65            for i in 0..edges.num_rows() {
66                let source_opt = source_array.value(i);
67                let target_opt = target_array.value(i);
68                let (source, target) = (source_opt, target_opt);
69                let source_str = source.to_string();
70                let target_str = target.to_string();
71                
72                // Add to adjacency lists (create nodes if they don't exist)
73                adjacency_list.entry(source_str.clone())
74                    .or_default()
75                    .push(target_str.clone());
76                
77                reverse_adjacency_list.entry(target_str.clone())
78                    .or_default()
79                    .push(source_str.clone());
80                
81                // Add to node index if not exists
82                if !node_index.contains_key(&source_str) {
83                    let idx = node_index.len();
84                    node_index.insert(source_str.clone(), idx);
85                }
86                if !node_index.contains_key(&target_str) {
87                    let idx = node_index.len();
88                    node_index.insert(target_str.clone(), idx);
89                }
90                
91                // Handle edge weights
92                let weight = if let Some(weights) = weight_array {
93                    weights.value(i)
94                } else {
95                    1.0 // Default weight
96                };
97                
98                edge_weights.insert((source_str, target_str), weight);
99            }
100        }
101        
102        let node_count = node_index.len();
103        let edge_count = edges.num_rows();
104        
105        Ok(GraphIndexes {
106            adjacency_list,
107            reverse_adjacency_list,
108            node_index,
109            edge_weights,
110            node_count,
111            edge_count,
112        })
113    }
114    
115    pub fn neighbors(&self, node_id: &str) -> Option<&Vec<String>> {
116        self.adjacency_list.get(node_id)
117    }
118    
119    pub fn predecessors(&self, node_id: &str) -> Option<&Vec<String>> {
120        self.reverse_adjacency_list.get(node_id)
121    }
122    
123    pub fn has_node(&self, node_id: &str) -> bool {
124        self.node_index.contains_key(node_id)
125    }
126    
127    pub fn edge_weight(&self, source: &str, target: &str) -> Option<f64> {
128        self.edge_weights.get(&(source.to_string(), target.to_string())).copied()
129    }
130    
131    pub fn all_nodes(&self) -> impl Iterator<Item = &String> {
132        self.node_index.keys()
133    }
134}
arrow_graph/graph/indexes.rs

arrow_graph/graph/
indexes.rs
