Skip to main content

ruvector_graph/
index.rs

1//! Index structures for fast node and edge lookups
2//!
3//! Provides label indexes, property indexes, and edge type indexes for efficient querying
4
5use crate::edge::Edge;
6use crate::hyperedge::{Hyperedge, HyperedgeId};
7use crate::node::Node;
8use crate::types::{EdgeId, NodeId, PropertyValue};
9use dashmap::DashMap;
10use std::collections::HashSet;
11use std::sync::Arc;
12
13/// Label index for nodes (maps labels to node IDs)
14#[derive(Debug, Clone)]
15pub struct LabelIndex {
16    /// Label -> Set of node IDs
17    index: Arc<DashMap<String, HashSet<NodeId>>>,
18}
19
20impl LabelIndex {
21    /// Create a new label index
22    pub fn new() -> Self {
23        Self {
24            index: Arc::new(DashMap::new()),
25        }
26    }
27
28    /// Add a node to the index
29    pub fn add_node(&self, node: &Node) {
30        for label in &node.labels {
31            self.index
32                .entry(label.name.clone())
33                .or_insert_with(HashSet::new)
34                .insert(node.id.clone());
35        }
36    }
37
38    /// Remove a node from the index
39    pub fn remove_node(&self, node: &Node) {
40        for label in &node.labels {
41            if let Some(mut set) = self.index.get_mut(&label.name) {
42                set.remove(&node.id);
43            }
44        }
45    }
46
47    /// Get all nodes with a specific label
48    pub fn get_nodes_by_label(&self, label: &str) -> Vec<NodeId> {
49        self.index
50            .get(label)
51            .map(|set| set.iter().cloned().collect())
52            .unwrap_or_default()
53    }
54
55    /// Get all labels in the index
56    pub fn all_labels(&self) -> Vec<String> {
57        self.index.iter().map(|entry| entry.key().clone()).collect()
58    }
59
60    /// Count nodes with a specific label
61    pub fn count_by_label(&self, label: &str) -> usize {
62        self.index.get(label).map(|set| set.len()).unwrap_or(0)
63    }
64
65    /// Clear the index
66    pub fn clear(&self) {
67        self.index.clear();
68    }
69}
70
71impl Default for LabelIndex {
72    fn default() -> Self {
73        Self::new()
74    }
75}
76
77/// Property index for nodes (maps property keys to values to node IDs)
78#[derive(Debug, Clone)]
79pub struct PropertyIndex {
80    /// Property key -> Property value -> Set of node IDs
81    index: Arc<DashMap<String, DashMap<String, HashSet<NodeId>>>>,
82}
83
84impl PropertyIndex {
85    /// Create a new property index
86    pub fn new() -> Self {
87        Self {
88            index: Arc::new(DashMap::new()),
89        }
90    }
91
92    /// Add a node to the index
93    pub fn add_node(&self, node: &Node) {
94        for (key, value) in &node.properties {
95            let value_str = self.property_value_to_string(value);
96            self.index
97                .entry(key.clone())
98                .or_insert_with(DashMap::new)
99                .entry(value_str)
100                .or_insert_with(HashSet::new)
101                .insert(node.id.clone());
102        }
103    }
104
105    /// Remove a node from the index
106    pub fn remove_node(&self, node: &Node) {
107        for (key, value) in &node.properties {
108            let value_str = self.property_value_to_string(value);
109            if let Some(value_map) = self.index.get(key) {
110                if let Some(mut set) = value_map.get_mut(&value_str) {
111                    set.remove(&node.id);
112                }
113            }
114        }
115    }
116
117    /// Get nodes by property key-value pair
118    pub fn get_nodes_by_property(&self, key: &str, value: &PropertyValue) -> Vec<NodeId> {
119        let value_str = self.property_value_to_string(value);
120        self.index
121            .get(key)
122            .and_then(|value_map| {
123                value_map
124                    .get(&value_str)
125                    .map(|set| set.iter().cloned().collect())
126            })
127            .unwrap_or_default()
128    }
129
130    /// Get all nodes that have a specific property key (regardless of value)
131    pub fn get_nodes_with_property(&self, key: &str) -> Vec<NodeId> {
132        self.index
133            .get(key)
134            .map(|value_map| {
135                let mut result = HashSet::new();
136                for entry in value_map.iter() {
137                    for id in entry.value().iter() {
138                        result.insert(id.clone());
139                    }
140                }
141                result.into_iter().collect()
142            })
143            .unwrap_or_default()
144    }
145
146    /// Get all property keys in the index
147    pub fn all_property_keys(&self) -> Vec<String> {
148        self.index.iter().map(|entry| entry.key().clone()).collect()
149    }
150
151    /// Clear the index
152    pub fn clear(&self) {
153        self.index.clear();
154    }
155
156    /// Convert property value to string for indexing
157    fn property_value_to_string(&self, value: &PropertyValue) -> String {
158        match value {
159            PropertyValue::Null => "null".to_string(),
160            PropertyValue::Boolean(b) => b.to_string(),
161            PropertyValue::Integer(i) => i.to_string(),
162            PropertyValue::Float(f) => f.to_string(),
163            PropertyValue::String(s) => s.clone(),
164            PropertyValue::FloatArray(_) | PropertyValue::Array(_) | PropertyValue::List(_) => format!("{:?}", value),
165            PropertyValue::Map(_) => format!("{:?}", value),
166        }
167    }
168}
169
170impl Default for PropertyIndex {
171    fn default() -> Self {
172        Self::new()
173    }
174}
175
176/// Edge type index (maps edge types to edge IDs)
177#[derive(Debug, Clone)]
178pub struct EdgeTypeIndex {
179    /// Edge type -> Set of edge IDs
180    index: Arc<DashMap<String, HashSet<EdgeId>>>,
181}
182
183impl EdgeTypeIndex {
184    /// Create a new edge type index
185    pub fn new() -> Self {
186        Self {
187            index: Arc::new(DashMap::new()),
188        }
189    }
190
191    /// Add an edge to the index
192    pub fn add_edge(&self, edge: &Edge) {
193        self.index
194            .entry(edge.edge_type.clone())
195            .or_insert_with(HashSet::new)
196            .insert(edge.id.clone());
197    }
198
199    /// Remove an edge from the index
200    pub fn remove_edge(&self, edge: &Edge) {
201        if let Some(mut set) = self.index.get_mut(&edge.edge_type) {
202            set.remove(&edge.id);
203        }
204    }
205
206    /// Get all edges of a specific type
207    pub fn get_edges_by_type(&self, edge_type: &str) -> Vec<EdgeId> {
208        self.index
209            .get(edge_type)
210            .map(|set| set.iter().cloned().collect())
211            .unwrap_or_default()
212    }
213
214    /// Get all edge types
215    pub fn all_edge_types(&self) -> Vec<String> {
216        self.index.iter().map(|entry| entry.key().clone()).collect()
217    }
218
219    /// Count edges of a specific type
220    pub fn count_by_type(&self, edge_type: &str) -> usize {
221        self.index.get(edge_type).map(|set| set.len()).unwrap_or(0)
222    }
223
224    /// Clear the index
225    pub fn clear(&self) {
226        self.index.clear();
227    }
228}
229
230impl Default for EdgeTypeIndex {
231    fn default() -> Self {
232        Self::new()
233    }
234}
235
236/// Adjacency index for fast neighbor lookups
237#[derive(Debug, Clone)]
238pub struct AdjacencyIndex {
239    /// Node ID -> Set of outgoing edge IDs
240    outgoing: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
241    /// Node ID -> Set of incoming edge IDs
242    incoming: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
243}
244
245impl AdjacencyIndex {
246    /// Create a new adjacency index
247    pub fn new() -> Self {
248        Self {
249            outgoing: Arc::new(DashMap::new()),
250            incoming: Arc::new(DashMap::new()),
251        }
252    }
253
254    /// Add an edge to the index
255    pub fn add_edge(&self, edge: &Edge) {
256        self.outgoing
257            .entry(edge.from.clone())
258            .or_insert_with(HashSet::new)
259            .insert(edge.id.clone());
260
261        self.incoming
262            .entry(edge.to.clone())
263            .or_insert_with(HashSet::new)
264            .insert(edge.id.clone());
265    }
266
267    /// Remove an edge from the index
268    pub fn remove_edge(&self, edge: &Edge) {
269        if let Some(mut set) = self.outgoing.get_mut(&edge.from) {
270            set.remove(&edge.id);
271        }
272        if let Some(mut set) = self.incoming.get_mut(&edge.to) {
273            set.remove(&edge.id);
274        }
275    }
276
277    /// Get all outgoing edges from a node
278    pub fn get_outgoing_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
279        self.outgoing
280            .get(node_id)
281            .map(|set| set.iter().cloned().collect())
282            .unwrap_or_default()
283    }
284
285    /// Get all incoming edges to a node
286    pub fn get_incoming_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
287        self.incoming
288            .get(node_id)
289            .map(|set| set.iter().cloned().collect())
290            .unwrap_or_default()
291    }
292
293    /// Get all edges connected to a node (both incoming and outgoing)
294    pub fn get_all_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
295        let mut edges = self.get_outgoing_edges(node_id);
296        edges.extend(self.get_incoming_edges(node_id));
297        edges
298    }
299
300    /// Get outgoing edges for multiple nodes in one call (O(k×avg_degree) vs O(E) for full scan).
301    pub fn get_edges_for_nodes(&self, node_ids: &[NodeId]) -> Vec<EdgeId> {
302        let mut result = Vec::with_capacity(node_ids.len() * 4);
303        for node_id in node_ids {
304            if let Some(guard) = self.outgoing.get(node_id) {
305                result.extend(guard.iter().cloned());
306            }
307        }
308
309        result
310    }
311
312    /// Callback-based pass: avoids intermediate Vec<EdgeId> allocation for hot paths.
313    pub fn for_each_outgoing_edge<F: FnMut(&EdgeId)>(&self, node_ids: &[NodeId], mut f: F) {
314        for node_id in node_ids {
315            if let Some(guard) = self.outgoing.get(node_id) {
316                for edge_id in guard.iter() {
317                    f(edge_id);
318                }
319            }
320        }
321    }
322
323    /// Get degree (number of outgoing edges)
324    pub fn out_degree(&self, node_id: &NodeId) -> usize {
325        self.outgoing.get(node_id).map(|set| set.len()).unwrap_or(0)
326    }
327
328    /// Get in-degree (number of incoming edges)
329    pub fn in_degree(&self, node_id: &NodeId) -> usize {
330        self.incoming.get(node_id).map(|set| set.len()).unwrap_or(0)
331    }
332
333    /// Clear the index
334    pub fn clear(&self) {
335        self.outgoing.clear();
336        self.incoming.clear();
337    }
338}
339
340impl Default for AdjacencyIndex {
341    fn default() -> Self {
342        Self::new()
343    }
344}
345
346/// Hyperedge node index (maps nodes to hyperedges they participate in)
347#[derive(Debug, Clone)]
348pub struct HyperedgeNodeIndex {
349    /// Node ID -> Set of hyperedge IDs
350    index: Arc<DashMap<NodeId, HashSet<HyperedgeId>>>,
351}
352
353impl HyperedgeNodeIndex {
354    /// Create a new hyperedge node index
355    pub fn new() -> Self {
356        Self {
357            index: Arc::new(DashMap::new()),
358        }
359    }
360
361    /// Add a hyperedge to the index
362    pub fn add_hyperedge(&self, hyperedge: &Hyperedge) {
363        for node_id in &hyperedge.nodes {
364            self.index
365                .entry(node_id.clone())
366                .or_insert_with(HashSet::new)
367                .insert(hyperedge.id.clone());
368        }
369    }
370
371    /// Remove a hyperedge from the index
372    pub fn remove_hyperedge(&self, hyperedge: &Hyperedge) {
373        for node_id in &hyperedge.nodes {
374            if let Some(mut set) = self.index.get_mut(node_id) {
375                set.remove(&hyperedge.id);
376            }
377        }
378    }
379
380    /// Get all hyperedges containing a node
381    pub fn get_hyperedges_by_node(&self, node_id: &NodeId) -> Vec<HyperedgeId> {
382        self.index
383            .get(node_id)
384            .map(|set| set.iter().cloned().collect())
385            .unwrap_or_default()
386    }
387
388    /// Clear the index
389    pub fn clear(&self) {
390        self.index.clear();
391    }
392}
393
394impl Default for HyperedgeNodeIndex {
395    fn default() -> Self {
396        Self::new()
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403    use crate::node::NodeBuilder;
404
405    #[test]
406    fn test_label_index() {
407        let index = LabelIndex::new();
408
409        let node1 = NodeBuilder::new().label("Person").label("User").build();
410
411        let node2 = NodeBuilder::new().label("Person").build();
412
413        index.add_node(&node1);
414        index.add_node(&node2);
415
416        let people = index.get_nodes_by_label("Person");
417        assert_eq!(people.len(), 2);
418
419        let users = index.get_nodes_by_label("User");
420        assert_eq!(users.len(), 1);
421
422        assert_eq!(index.count_by_label("Person"), 2);
423    }
424
425    #[test]
426    fn test_property_index() {
427        let index = PropertyIndex::new();
428
429        let node1 = NodeBuilder::new()
430            .property("name", "Alice")
431            .property("age", 30i64)
432            .build();
433
434        let node2 = NodeBuilder::new()
435            .property("name", "Bob")
436            .property("age", 30i64)
437            .build();
438
439        index.add_node(&node1);
440        index.add_node(&node2);
441
442        let alice =
443            index.get_nodes_by_property("name", &PropertyValue::String("Alice".to_string()));
444        assert_eq!(alice.len(), 1);
445
446        let age_30 = index.get_nodes_by_property("age", &PropertyValue::Integer(30));
447        assert_eq!(age_30.len(), 2);
448
449        let with_age = index.get_nodes_with_property("age");
450        assert_eq!(with_age.len(), 2);
451    }
452
453    #[test]
454    fn test_edge_type_index() {
455        let index = EdgeTypeIndex::new();
456
457        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
458        let edge2 = Edge::create("n2".to_string(), "n3".to_string(), "KNOWS");
459        let edge3 = Edge::create("n1".to_string(), "n3".to_string(), "WORKS_WITH");
460
461        index.add_edge(&edge1);
462        index.add_edge(&edge2);
463        index.add_edge(&edge3);
464
465        let knows_edges = index.get_edges_by_type("KNOWS");
466        assert_eq!(knows_edges.len(), 2);
467
468        let works_with_edges = index.get_edges_by_type("WORKS_WITH");
469        assert_eq!(works_with_edges.len(), 1);
470
471        assert_eq!(index.all_edge_types().len(), 2);
472    }
473
474    #[test]
475    fn test_adjacency_index() {
476        let index = AdjacencyIndex::new();
477
478        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
479        let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
480        let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
481
482        index.add_edge(&edge1);
483        index.add_edge(&edge2);
484        index.add_edge(&edge3);
485
486        assert_eq!(index.out_degree(&"n1".to_string()), 2);
487        assert_eq!(index.in_degree(&"n1".to_string()), 1);
488
489        let outgoing = index.get_outgoing_edges(&"n1".to_string());
490        assert_eq!(outgoing.len(), 2);
491
492        let incoming = index.get_incoming_edges(&"n1".to_string());
493        assert_eq!(incoming.len(), 1);
494    }
495
496    #[test]
497    fn test_get_edges_for_nodes() {
498        let index = AdjacencyIndex::new();
499
500        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
501        let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
502        let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
503        let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
504
505        index.add_edge(&edge1);
506        index.add_edge(&edge2);
507        index.add_edge(&edge3);
508        index.add_edge(&edge4);
509
510        let result = index.get_edges_for_nodes(&["n1".to_string(), "n2".to_string()]);
511        assert_eq!(result.len(), 3);
512
513        let result_single = index.get_edges_for_nodes(&["n3".to_string()]);
514        assert_eq!(result_single.len(), 1);
515
516        let result_empty = index.get_edges_for_nodes(&["n5".to_string()]);
517        assert_eq!(result_empty.len(), 0);
518
519        let result_multi_empty = index.get_edges_for_nodes(&[]);
520        assert_eq!(result_multi_empty.len(), 0);
521    }
522
523    #[test]
524    fn test_for_each_outgoing_edge() {
525        let index = AdjacencyIndex::new();
526
527        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
528        let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
529        let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
530        let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
531
532        index.add_edge(&edge1);
533        index.add_edge(&edge2);
534        index.add_edge(&edge3);
535        index.add_edge(&edge4);
536
537        let mut collected = Vec::new();
538        index.for_each_outgoing_edge(&["n1".to_string(), "n2".to_string()], |id| {
539            collected.push(id.clone());
540        });
541        assert_eq!(collected.len(), 3);
542
543        let mut single = Vec::new();
544        index.for_each_outgoing_edge(&["n3".to_string()], |id| {
545            single.push(id.clone());
546        });
547        assert_eq!(single.len(), 1);
548
549        let mut empty = Vec::new();
550        index.for_each_outgoing_edge(&["n5".to_string()], |id| {
551            empty.push(id.clone());
552        });
553        assert_eq!(empty.len(), 0);
554
555        let mut none_called = Vec::new();
556        index.for_each_outgoing_edge(&[], |id| {
557            none_called.push(id.clone());
558        });
559        assert_eq!(none_called.len(), 0);
560    }
561}