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                if 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        
103        let node_count = node_index.len();
104        let edge_count = edges.num_rows();
105        
106        Ok(GraphIndexes {
107            adjacency_list,
108            reverse_adjacency_list,
109            node_index,
110            edge_weights,
111            node_count,
112            edge_count,
113        })
114    }
115    
116    pub fn neighbors(&self, node_id: &str) -> Option<&Vec<String>> {
117        self.adjacency_list.get(node_id)
118    }
119    
120    pub fn predecessors(&self, node_id: &str) -> Option<&Vec<String>> {
121        self.reverse_adjacency_list.get(node_id)
122    }
123    
124    pub fn has_node(&self, node_id: &str) -> bool {
125        self.node_index.contains_key(node_id)
126    }
127    
128    pub fn edge_weight(&self, source: &str, target: &str) -> Option<f64> {
129        self.edge_weights.get(&(source.to_string(), target.to_string())).copied()
130    }
131    
132    pub fn all_nodes(&self) -> impl Iterator<Item = &String> {
133        self.node_index.keys()
134    }
135}
arrow_graph/graph/indexes.rs

arrow_graph/graph/
indexes.rs
