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(_) => {
165                format!("{:?}", value)
166            }
167            PropertyValue::Map(_) => format!("{:?}", value),
168        }
169    }
170}
171
172impl Default for PropertyIndex {
173    fn default() -> Self {
174        Self::new()
175    }
176}
177
178/// Edge type index (maps edge types to edge IDs)
179#[derive(Debug, Clone)]
180pub struct EdgeTypeIndex {
181    /// Edge type -> Set of edge IDs
182    index: Arc<DashMap<String, HashSet<EdgeId>>>,
183}
184
185impl EdgeTypeIndex {
186    /// Create a new edge type index
187    pub fn new() -> Self {
188        Self {
189            index: Arc::new(DashMap::new()),
190        }
191    }
192
193    /// Add an edge to the index
194    pub fn add_edge(&self, edge: &Edge) {
195        self.index
196            .entry(edge.edge_type.clone())
197            .or_insert_with(HashSet::new)
198            .insert(edge.id.clone());
199    }
200
201    /// Remove an edge from the index
202    pub fn remove_edge(&self, edge: &Edge) {
203        if let Some(mut set) = self.index.get_mut(&edge.edge_type) {
204            set.remove(&edge.id);
205        }
206    }
207
208    /// Get all edges of a specific type
209    pub fn get_edges_by_type(&self, edge_type: &str) -> Vec<EdgeId> {
210        self.index
211            .get(edge_type)
212            .map(|set| set.iter().cloned().collect())
213            .unwrap_or_default()
214    }
215
216    /// Get all edge types
217    pub fn all_edge_types(&self) -> Vec<String> {
218        self.index.iter().map(|entry| entry.key().clone()).collect()
219    }
220
221    /// Count edges of a specific type
222    pub fn count_by_type(&self, edge_type: &str) -> usize {
223        self.index.get(edge_type).map(|set| set.len()).unwrap_or(0)
224    }
225
226    /// Clear the index
227    pub fn clear(&self) {
228        self.index.clear();
229    }
230}
231
232impl Default for EdgeTypeIndex {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238/// Adjacency index for fast neighbor lookups
239#[derive(Debug, Clone)]
240pub struct AdjacencyIndex {
241    /// Node ID -> Set of outgoing edge IDs
242    outgoing: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
243    /// Node ID -> Set of incoming edge IDs
244    incoming: Arc<DashMap<NodeId, HashSet<EdgeId>>>,
245}
246
247impl AdjacencyIndex {
248    /// Create a new adjacency index
249    pub fn new() -> Self {
250        Self {
251            outgoing: Arc::new(DashMap::new()),
252            incoming: Arc::new(DashMap::new()),
253        }
254    }
255
256    /// Add an edge to the index
257    pub fn add_edge(&self, edge: &Edge) {
258        self.outgoing
259            .entry(edge.from.clone())
260            .or_insert_with(HashSet::new)
261            .insert(edge.id.clone());
262
263        self.incoming
264            .entry(edge.to.clone())
265            .or_insert_with(HashSet::new)
266            .insert(edge.id.clone());
267    }
268
269    /// Remove an edge from the index
270    pub fn remove_edge(&self, edge: &Edge) {
271        if let Some(mut set) = self.outgoing.get_mut(&edge.from) {
272            set.remove(&edge.id);
273        }
274        if let Some(mut set) = self.incoming.get_mut(&edge.to) {
275            set.remove(&edge.id);
276        }
277    }
278
279    /// Get all outgoing edges from a node
280    pub fn get_outgoing_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
281        self.outgoing
282            .get(node_id)
283            .map(|set| set.iter().cloned().collect())
284            .unwrap_or_default()
285    }
286
287    /// Get all incoming edges to a node
288    pub fn get_incoming_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
289        self.incoming
290            .get(node_id)
291            .map(|set| set.iter().cloned().collect())
292            .unwrap_or_default()
293    }
294
295    /// Get all edges connected to a node (both incoming and outgoing)
296    pub fn get_all_edges(&self, node_id: &NodeId) -> Vec<EdgeId> {
297        let mut edges = self.get_outgoing_edges(node_id);
298        edges.extend(self.get_incoming_edges(node_id));
299        edges
300    }
301
302    /// Get outgoing edges for multiple nodes in one call (O(k×avg_degree) vs O(E) for full scan).
303    pub fn get_edges_for_nodes(&self, node_ids: &[NodeId]) -> Vec<EdgeId> {
304        let mut result = Vec::with_capacity(node_ids.len() * 4);
305        for node_id in node_ids {
306            if let Some(guard) = self.outgoing.get(node_id) {
307                result.extend(guard.iter().cloned());
308            }
309        }
310
311        result
312    }
313
314    /// Callback-based pass: avoids intermediate Vec<EdgeId> allocation for hot paths.
315    pub fn for_each_outgoing_edge<F: FnMut(&EdgeId)>(&self, node_ids: &[NodeId], mut f: F) {
316        for node_id in node_ids {
317            if let Some(guard) = self.outgoing.get(node_id) {
318                for edge_id in guard.iter() {
319                    f(edge_id);
320                }
321            }
322        }
323    }
324
325    /// Get degree (number of outgoing edges)
326    pub fn out_degree(&self, node_id: &NodeId) -> usize {
327        self.outgoing.get(node_id).map(|set| set.len()).unwrap_or(0)
328    }
329
330    /// Get in-degree (number of incoming edges)
331    pub fn in_degree(&self, node_id: &NodeId) -> usize {
332        self.incoming.get(node_id).map(|set| set.len()).unwrap_or(0)
333    }
334
335    /// Clear the index
336    pub fn clear(&self) {
337        self.outgoing.clear();
338        self.incoming.clear();
339    }
340}
341
342impl Default for AdjacencyIndex {
343    fn default() -> Self {
344        Self::new()
345    }
346}
347
348/// Hyperedge node index (maps nodes to hyperedges they participate in)
349#[derive(Debug, Clone)]
350pub struct HyperedgeNodeIndex {
351    /// Node ID -> Set of hyperedge IDs
352    index: Arc<DashMap<NodeId, HashSet<HyperedgeId>>>,
353}
354
355impl HyperedgeNodeIndex {
356    /// Create a new hyperedge node index
357    pub fn new() -> Self {
358        Self {
359            index: Arc::new(DashMap::new()),
360        }
361    }
362
363    /// Add a hyperedge to the index
364    pub fn add_hyperedge(&self, hyperedge: &Hyperedge) {
365        for node_id in &hyperedge.nodes {
366            self.index
367                .entry(node_id.clone())
368                .or_insert_with(HashSet::new)
369                .insert(hyperedge.id.clone());
370        }
371    }
372
373    /// Remove a hyperedge from the index
374    pub fn remove_hyperedge(&self, hyperedge: &Hyperedge) {
375        for node_id in &hyperedge.nodes {
376            if let Some(mut set) = self.index.get_mut(node_id) {
377                set.remove(&hyperedge.id);
378            }
379        }
380    }
381
382    /// Get all hyperedges containing a node
383    pub fn get_hyperedges_by_node(&self, node_id: &NodeId) -> Vec<HyperedgeId> {
384        self.index
385            .get(node_id)
386            .map(|set| set.iter().cloned().collect())
387            .unwrap_or_default()
388    }
389
390    /// Clear the index
391    pub fn clear(&self) {
392        self.index.clear();
393    }
394}
395
396impl Default for HyperedgeNodeIndex {
397    fn default() -> Self {
398        Self::new()
399    }
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use crate::node::NodeBuilder;
406
407    #[test]
408    fn test_label_index() {
409        let index = LabelIndex::new();
410
411        let node1 = NodeBuilder::new().label("Person").label("User").build();
412
413        let node2 = NodeBuilder::new().label("Person").build();
414
415        index.add_node(&node1);
416        index.add_node(&node2);
417
418        let people = index.get_nodes_by_label("Person");
419        assert_eq!(people.len(), 2);
420
421        let users = index.get_nodes_by_label("User");
422        assert_eq!(users.len(), 1);
423
424        assert_eq!(index.count_by_label("Person"), 2);
425    }
426
427    #[test]
428    fn test_property_index() {
429        let index = PropertyIndex::new();
430
431        let node1 = NodeBuilder::new()
432            .property("name", "Alice")
433            .property("age", 30i64)
434            .build();
435
436        let node2 = NodeBuilder::new()
437            .property("name", "Bob")
438            .property("age", 30i64)
439            .build();
440
441        index.add_node(&node1);
442        index.add_node(&node2);
443
444        let alice =
445            index.get_nodes_by_property("name", &PropertyValue::String("Alice".to_string()));
446        assert_eq!(alice.len(), 1);
447
448        let age_30 = index.get_nodes_by_property("age", &PropertyValue::Integer(30));
449        assert_eq!(age_30.len(), 2);
450
451        let with_age = index.get_nodes_with_property("age");
452        assert_eq!(with_age.len(), 2);
453    }
454
455    #[test]
456    fn test_edge_type_index() {
457        let index = EdgeTypeIndex::new();
458
459        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
460        let edge2 = Edge::create("n2".to_string(), "n3".to_string(), "KNOWS");
461        let edge3 = Edge::create("n1".to_string(), "n3".to_string(), "WORKS_WITH");
462
463        index.add_edge(&edge1);
464        index.add_edge(&edge2);
465        index.add_edge(&edge3);
466
467        let knows_edges = index.get_edges_by_type("KNOWS");
468        assert_eq!(knows_edges.len(), 2);
469
470        let works_with_edges = index.get_edges_by_type("WORKS_WITH");
471        assert_eq!(works_with_edges.len(), 1);
472
473        assert_eq!(index.all_edge_types().len(), 2);
474    }
475
476    #[test]
477    fn test_adjacency_index() {
478        let index = AdjacencyIndex::new();
479
480        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
481        let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
482        let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
483
484        index.add_edge(&edge1);
485        index.add_edge(&edge2);
486        index.add_edge(&edge3);
487
488        assert_eq!(index.out_degree(&"n1".to_string()), 2);
489        assert_eq!(index.in_degree(&"n1".to_string()), 1);
490
491        let outgoing = index.get_outgoing_edges(&"n1".to_string());
492        assert_eq!(outgoing.len(), 2);
493
494        let incoming = index.get_incoming_edges(&"n1".to_string());
495        assert_eq!(incoming.len(), 1);
496    }
497
498    #[test]
499    fn test_get_edges_for_nodes() {
500        let index = AdjacencyIndex::new();
501
502        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
503        let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
504        let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
505        let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
506
507        index.add_edge(&edge1);
508        index.add_edge(&edge2);
509        index.add_edge(&edge3);
510        index.add_edge(&edge4);
511
512        let result = index.get_edges_for_nodes(&["n1".to_string(), "n2".to_string()]);
513        assert_eq!(result.len(), 3);
514
515        let result_single = index.get_edges_for_nodes(&["n3".to_string()]);
516        assert_eq!(result_single.len(), 1);
517
518        let result_empty = index.get_edges_for_nodes(&["n5".to_string()]);
519        assert_eq!(result_empty.len(), 0);
520
521        let result_multi_empty = index.get_edges_for_nodes(&[]);
522        assert_eq!(result_multi_empty.len(), 0);
523    }
524
525    #[test]
526    fn test_for_each_outgoing_edge() {
527        let index = AdjacencyIndex::new();
528
529        let edge1 = Edge::create("n1".to_string(), "n2".to_string(), "KNOWS");
530        let edge2 = Edge::create("n1".to_string(), "n3".to_string(), "KNOWS");
531        let edge3 = Edge::create("n2".to_string(), "n1".to_string(), "KNOWS");
532        let edge4 = Edge::create("n3".to_string(), "n4".to_string(), "KNOWS");
533
534        index.add_edge(&edge1);
535        index.add_edge(&edge2);
536        index.add_edge(&edge3);
537        index.add_edge(&edge4);
538
539        let mut collected = Vec::new();
540        index.for_each_outgoing_edge(&["n1".to_string(), "n2".to_string()], |id| {
541            collected.push(id.clone());
542        });
543        assert_eq!(collected.len(), 3);
544
545        let mut single = Vec::new();
546        index.for_each_outgoing_edge(&["n3".to_string()], |id| {
547            single.push(id.clone());
548        });
549        assert_eq!(single.len(), 1);
550
551        let mut empty = Vec::new();
552        index.for_each_outgoing_edge(&["n5".to_string()], |id| {
553            empty.push(id.clone());
554        });
555        assert_eq!(empty.len(), 0);
556
557        let mut none_called = Vec::new();
558        index.for_each_outgoing_edge(&[], |id| {
559            none_called.push(id.clone());
560        });
561        assert_eq!(none_called.len(), 0);
562    }
563}