arrow_graph/streaming/
incremental.rs

1use crate::error::Result;
2use crate::graph::ArrowGraph;
3use arrow::array::{Array, StringArray, Float64Array, RecordBatch};
4use arrow::datatypes::{Schema, Field, DataType};
5use std::collections::{HashMap, HashSet};
6use std::sync::Arc;
7
8/// Incremental graph update operations for streaming graph processing
9/// Supports efficient add/remove of vertices and edges with minimal recomputation
10#[derive(Debug, Clone)]
11pub struct IncrementalGraphProcessor {
12    graph: ArrowGraph,
13    vertex_cache: HashMap<String, usize>, // vertex_id -> internal_index
14    edge_cache: HashMap<(String, String), f64>, // (source, target) -> weight
15    pending_vertex_additions: Vec<String>,
16    pending_vertex_removals: HashSet<String>,
17    pending_edge_additions: Vec<(String, String, f64)>,
18    pending_edge_removals: HashSet<(String, String)>,
19    batch_size: usize, // Number of operations to batch before applying
20}
21
22/// Types of incremental updates
23#[derive(Debug, Clone, PartialEq)]
24pub enum UpdateOperation {
25    AddVertex(String),
26    RemoveVertex(String),
27    AddEdge(String, String, f64), // source, target, weight
28    RemoveEdge(String, String),   // source, target
29}
30
31/// Result of applying incremental updates
32#[derive(Debug, Clone)]
33pub struct UpdateResult {
34    pub vertices_added: usize,
35    pub vertices_removed: usize,
36    pub edges_added: usize,
37    pub edges_removed: usize,
38    pub affected_components: Vec<String>, // Components that changed
39    pub recomputation_needed: bool, // Whether algorithms need full recomputation
40}
41
42impl IncrementalGraphProcessor {
43    /// Create a new incremental processor from an existing graph
44    pub fn new(graph: ArrowGraph) -> Result<Self> {
45        let mut vertex_cache = HashMap::new();
46        let mut edge_cache = HashMap::new();
47
48        // Build initial caches
49        let nodes_batch = &graph.nodes;
50        if nodes_batch.num_rows() > 0 {
51            let node_ids = nodes_batch
52                .column(0)
53                .as_any()
54                .downcast_ref::<StringArray>()
55                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for node IDs"))?;
56
57            for (idx, node_id) in node_ids.iter().enumerate() {
58                if let Some(id) = node_id {
59                    vertex_cache.insert(id.to_string(), idx);
60                }
61            }
62        }
63
64        let edges_batch = &graph.edges;
65        if edges_batch.num_rows() > 0 {
66            let source_ids = edges_batch
67                .column(0)
68                .as_any()
69                .downcast_ref::<StringArray>()
70                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for source IDs"))?;
71            let target_ids = edges_batch
72                .column(1)
73                .as_any()
74                .downcast_ref::<StringArray>()
75                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for target IDs"))?;
76            let weights = edges_batch
77                .column(2)
78                .as_any()
79                .downcast_ref::<Float64Array>()
80                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected float64 array for weights"))?;
81
82            for i in 0..source_ids.len() {
83                let source = source_ids.value(i);
84                let target = target_ids.value(i);
85                let weight = weights.value(i);
86                edge_cache.insert((source.to_string(), target.to_string()), weight);
87            }
88        }
89
90        Ok(Self {
91            graph,
92            vertex_cache,
93            edge_cache,
94            pending_vertex_additions: Vec::new(),
95            pending_vertex_removals: HashSet::new(),
96            pending_edge_additions: Vec::new(),
97            pending_edge_removals: HashSet::new(),
98            batch_size: 1000, // Default batch size
99        })
100    }
101
102    /// Set the batch size for operations
103    pub fn set_batch_size(&mut self, size: usize) {
104        self.batch_size = size;
105    }
106
107    /// Add a vertex to the graph (batched)
108    pub fn add_vertex(&mut self, vertex_id: String) -> Result<()> {
109        // Check if vertex already exists
110        if self.vertex_cache.contains_key(&vertex_id) || 
111           self.pending_vertex_additions.contains(&vertex_id) {
112            return Ok(()); // Already exists, no-op
113        }
114
115        // Remove from pending removals if present
116        self.pending_vertex_removals.remove(&vertex_id);
117        
118        // Add to pending additions
119        self.pending_vertex_additions.push(vertex_id);
120
121        // Auto-flush if batch size reached
122        if self.pending_vertex_additions.len() >= self.batch_size {
123            self.flush_vertex_operations()?;
124        }
125
126        Ok(())
127    }
128
129    /// Remove a vertex from the graph (batched)
130    pub fn remove_vertex(&mut self, vertex_id: String) -> Result<()> {
131        // Check if vertex exists
132        if !self.vertex_cache.contains_key(&vertex_id) && 
133           !self.pending_vertex_additions.contains(&vertex_id) {
134            return Ok(()); // Doesn't exist, no-op
135        }
136
137        // Remove from pending additions if present
138        self.pending_vertex_additions.retain(|v| v != &vertex_id);
139        
140        // Add to pending removals
141        self.pending_vertex_removals.insert(vertex_id.clone());
142
143        // Also remove all edges involving this vertex
144        let edges_to_remove: Vec<(String, String)> = self.edge_cache
145            .keys()
146            .filter(|(source, target)| source == &vertex_id || target == &vertex_id)
147            .cloned()
148            .collect();
149
150        for (source, target) in edges_to_remove {
151            self.pending_edge_removals.insert((source, target));
152        }
153
154        // Auto-flush if batch size reached
155        if self.pending_vertex_removals.len() >= self.batch_size {
156            self.flush_all_operations()?;
157        }
158
159        Ok(())
160    }
161
162    /// Add an edge to the graph (batched)
163    pub fn add_edge(&mut self, source: String, target: String, weight: f64) -> Result<()> {
164        let edge_key = (source.clone(), target.clone());
165
166        // Check if edge already exists
167        if self.edge_cache.contains_key(&edge_key) {
168            // Update weight if different
169            if let Some(existing_weight) = self.edge_cache.get(&edge_key) {
170                if (existing_weight - weight).abs() > f64::EPSILON {
171                    self.edge_cache.insert(edge_key.clone(), weight);
172                }
173            }
174            return Ok(());
175        }
176
177        // Remove from pending removals if present
178        self.pending_edge_removals.remove(&edge_key);
179        
180        // Add vertices if they don't exist
181        self.add_vertex(source.clone())?;
182        self.add_vertex(target.clone())?;
183
184        // Add to pending additions
185        self.pending_edge_additions.push((source, target, weight));
186
187        // Auto-flush if batch size reached
188        if self.pending_edge_additions.len() >= self.batch_size {
189            self.flush_edge_operations()?;
190        }
191
192        Ok(())
193    }
194
195    /// Remove an edge from the graph (batched)
196    pub fn remove_edge(&mut self, source: String, target: String) -> Result<()> {
197        let edge_key = (source.clone(), target.clone());
198
199        // Check if edge exists
200        if !self.edge_cache.contains_key(&edge_key) {
201            return Ok(()); // Doesn't exist, no-op
202        }
203
204        // Remove from pending additions if present
205        self.pending_edge_additions.retain(|(s, t, _)| (s, t) != (&source, &target));
206        
207        // Add to pending removals
208        self.pending_edge_removals.insert(edge_key);
209
210        // Auto-flush if batch size reached
211        if self.pending_edge_removals.len() >= self.batch_size {
212            self.flush_edge_operations()?;
213        }
214
215        Ok(())
216    }
217
218    /// Apply a batch of update operations
219    pub fn apply_updates(&mut self, operations: Vec<UpdateOperation>) -> Result<UpdateResult> {
220        for operation in operations {
221            match operation {
222                UpdateOperation::AddVertex(vertex_id) => self.add_vertex(vertex_id)?,
223                UpdateOperation::RemoveVertex(vertex_id) => self.remove_vertex(vertex_id)?,
224                UpdateOperation::AddEdge(source, target, weight) => self.add_edge(source, target, weight)?,
225                UpdateOperation::RemoveEdge(source, target) => self.remove_edge(source, target)?,
226            }
227        }
228
229        // Flush all pending operations
230        self.flush_all_operations()
231    }
232
233    /// Flush all pending operations and rebuild the graph
234    pub fn flush_all_operations(&mut self) -> Result<UpdateResult> {
235        let mut result = UpdateResult {
236            vertices_added: 0,
237            vertices_removed: 0,
238            edges_added: 0,
239            edges_removed: 0,
240            affected_components: Vec::new(),
241            recomputation_needed: false,
242        };
243
244        // Flush vertices first
245        let vertex_result = self.flush_vertex_operations()?;
246        result.vertices_added += vertex_result.vertices_added;
247        result.vertices_removed += vertex_result.vertices_removed;
248
249        // Then flush edges
250        let edge_result = self.flush_edge_operations()?;
251        result.edges_added += edge_result.edges_added;
252        result.edges_removed += edge_result.edges_removed;
253
254        // Determine if recomputation is needed
255        result.recomputation_needed = result.vertices_added > 0 || 
256                                     result.vertices_removed > 0 || 
257                                     result.edges_added > 0 || 
258                                     result.edges_removed > 0;
259
260        Ok(result)
261    }
262
263    /// Flush pending vertex operations
264    fn flush_vertex_operations(&mut self) -> Result<UpdateResult> {
265        let mut result = UpdateResult {
266            vertices_added: 0,
267            vertices_removed: 0,
268            edges_added: 0,
269            edges_removed: 0,
270            affected_components: Vec::new(),
271            recomputation_needed: false,
272        };
273
274        // Process additions
275        for vertex_id in &self.pending_vertex_additions {
276            if !self.vertex_cache.contains_key(vertex_id) {
277                let new_index = self.vertex_cache.len();
278                self.vertex_cache.insert(vertex_id.clone(), new_index);
279                result.vertices_added += 1;
280            }
281        }
282
283        // Process removals
284        for vertex_id in &self.pending_vertex_removals {
285            if self.vertex_cache.remove(vertex_id).is_some() {
286                result.vertices_removed += 1;
287            }
288        }
289
290        // Rebuild graph if there were changes
291        if result.vertices_added > 0 || result.vertices_removed > 0 {
292            self.rebuild_graph()?;
293        }
294
295        // Clear pending operations
296        self.pending_vertex_additions.clear();
297        self.pending_vertex_removals.clear();
298
299        Ok(result)
300    }
301
302    /// Flush pending edge operations
303    fn flush_edge_operations(&mut self) -> Result<UpdateResult> {
304        let mut result = UpdateResult {
305            vertices_added: 0,
306            vertices_removed: 0,
307            edges_added: 0,
308            edges_removed: 0,
309            affected_components: Vec::new(),
310            recomputation_needed: false,
311        };
312
313        // Process additions
314        for (source, target, weight) in &self.pending_edge_additions {
315            let edge_key = (source.clone(), target.clone());
316            if !self.edge_cache.contains_key(&edge_key) {
317                self.edge_cache.insert(edge_key, *weight);
318                result.edges_added += 1;
319            }
320        }
321
322        // Process removals
323        for edge_key in &self.pending_edge_removals {
324            if self.edge_cache.remove(edge_key).is_some() {
325                result.edges_removed += 1;
326            }
327        }
328
329        // Rebuild graph if there were changes
330        if result.edges_added > 0 || result.edges_removed > 0 {
331            self.rebuild_graph()?;
332        }
333
334        // Clear pending operations
335        self.pending_edge_additions.clear();
336        self.pending_edge_removals.clear();
337
338        Ok(result)
339    }
340
341    /// Rebuild the ArrowGraph from current caches
342    fn rebuild_graph(&mut self) -> Result<()> {
343        // Build nodes RecordBatch
344        let node_ids: Vec<String> = self.vertex_cache.keys().cloned().collect();
345        let nodes_schema = Arc::new(Schema::new(vec![
346            Field::new("id", DataType::Utf8, false),
347        ]));
348        
349        let nodes_batch = if node_ids.is_empty() {
350            None
351        } else {
352            let node_id_array = StringArray::from(node_ids);
353            Some(RecordBatch::try_new(
354                nodes_schema.clone(),
355                vec![Arc::new(node_id_array)],
356            )?)
357        };
358
359        // Build edges RecordBatch
360        let edges_schema = Arc::new(Schema::new(vec![
361            Field::new("source", DataType::Utf8, false),
362            Field::new("target", DataType::Utf8, false),
363            Field::new("weight", DataType::Float64, false),
364        ]));
365
366        let edges_batch = if self.edge_cache.is_empty() {
367            None
368        } else {
369            let mut sources = Vec::new();
370            let mut targets = Vec::new();
371            let mut weights = Vec::new();
372
373            for ((source, target), weight) in &self.edge_cache {
374                sources.push(source.clone());
375                targets.push(target.clone());
376                weights.push(*weight);
377            }
378
379            let source_array = StringArray::from(sources);
380            let target_array = StringArray::from(targets);
381            let weight_array = Float64Array::from(weights);
382
383            Some(RecordBatch::try_new(
384                edges_schema.clone(),
385                vec![Arc::new(source_array), Arc::new(target_array), Arc::new(weight_array)],
386            )?)
387        };
388
389        // Create new graph
390        match (nodes_batch, edges_batch) {
391            (Some(nodes), Some(edges)) => {
392                self.graph = ArrowGraph::new(nodes, edges)?;
393            }
394            (Some(nodes), None) => {
395                // Create empty edges batch
396                let edges_schema = Arc::new(Schema::new(vec![
397                    Field::new("source", DataType::Utf8, false),
398                    Field::new("target", DataType::Utf8, false),
399                    Field::new("weight", DataType::Float64, false),
400                ]));
401                let empty_edges = RecordBatch::new_empty(edges_schema);
402                self.graph = ArrowGraph::new(nodes, empty_edges)?;
403            }
404            (None, Some(edges)) => {
405                self.graph = ArrowGraph::from_edges(edges)?;
406            }
407            (None, None) => {
408                self.graph = ArrowGraph::empty()?;
409            }
410        }
411
412        Ok(())
413    }
414
415    /// Get the current graph
416    pub fn graph(&self) -> &ArrowGraph {
417        &self.graph
418    }
419
420    /// Get the number of pending operations
421    pub fn pending_operations_count(&self) -> usize {
422        self.pending_vertex_additions.len() + 
423        self.pending_vertex_removals.len() + 
424        self.pending_edge_additions.len() + 
425        self.pending_edge_removals.len()
426    }
427
428    /// Check if there are pending operations
429    pub fn has_pending_operations(&self) -> bool {
430        self.pending_operations_count() > 0
431    }
432
433    /// Get statistics about the current state
434    pub fn statistics(&self) -> IncrementalStats {
435        IncrementalStats {
436            vertex_count: self.vertex_cache.len(),
437            edge_count: self.edge_cache.len(),
438            pending_vertex_additions: self.pending_vertex_additions.len(),
439            pending_vertex_removals: self.pending_vertex_removals.len(),
440            pending_edge_additions: self.pending_edge_additions.len(),
441            pending_edge_removals: self.pending_edge_removals.len(),
442            batch_size: self.batch_size,
443        }
444    }
445}
446
447/// Statistics about the incremental processor state
448#[derive(Debug, Clone)]
449pub struct IncrementalStats {
450    pub vertex_count: usize,
451    pub edge_count: usize,
452    pub pending_vertex_additions: usize,
453    pub pending_vertex_removals: usize,
454    pub pending_edge_additions: usize,
455    pub pending_edge_removals: usize,
456    pub batch_size: usize,
457}
458
459#[cfg(test)]
460mod tests {
461    use super::*;
462    use arrow::datatypes::{Schema, Field, DataType};
463    use std::sync::Arc;
464
465    fn create_test_graph() -> Result<ArrowGraph> {
466        // Create nodes
467        let nodes_schema = Arc::new(Schema::new(vec![
468            Field::new("id", DataType::Utf8, false),
469        ]));
470        let node_ids = StringArray::from(vec!["A", "B", "C"]);
471        let nodes_batch = RecordBatch::try_new(
472            nodes_schema,
473            vec![Arc::new(node_ids)],
474        )?;
475
476        // Create edges
477        let edges_schema = Arc::new(Schema::new(vec![
478            Field::new("source", DataType::Utf8, false),
479            Field::new("target", DataType::Utf8, false),
480            Field::new("weight", DataType::Float64, false),
481        ]));
482        let sources = StringArray::from(vec!["A", "B"]);
483        let targets = StringArray::from(vec!["B", "C"]);
484        let weights = Float64Array::from(vec![1.0, 2.0]);
485        let edges_batch = RecordBatch::try_new(
486            edges_schema,
487            vec![Arc::new(sources), Arc::new(targets), Arc::new(weights)],
488        )?;
489
490        ArrowGraph::new(nodes_batch, edges_batch)
491    }
492
493    #[test]
494    fn test_incremental_processor_creation() {
495        let graph = create_test_graph().unwrap();
496        let processor = IncrementalGraphProcessor::new(graph).unwrap();
497        
498        let stats = processor.statistics();
499        assert_eq!(stats.vertex_count, 3);
500        assert_eq!(stats.edge_count, 2);
501        assert_eq!(stats.pending_vertex_additions, 0);
502        assert_eq!(stats.pending_edge_additions, 0);
503    }
504
505    #[test]
506    fn test_add_vertex() {
507        let graph = create_test_graph().unwrap();
508        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
509        processor.set_batch_size(1); // Force immediate flush
510        
511        processor.add_vertex("D".to_string()).unwrap();
512        
513        let stats = processor.statistics();
514        assert_eq!(stats.vertex_count, 4);
515        assert!(processor.vertex_cache.contains_key("D"));
516    }
517
518    #[test]
519    fn test_remove_vertex() {
520        let graph = create_test_graph().unwrap();
521        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
522        processor.set_batch_size(1); // Force immediate flush
523        
524        processor.remove_vertex("C".to_string()).unwrap();
525        
526        let stats = processor.statistics();
527        assert_eq!(stats.vertex_count, 2);
528        assert!(!processor.vertex_cache.contains_key("C"));
529        // Edge B->C should also be removed
530        assert!(!processor.edge_cache.contains_key(&("B".to_string(), "C".to_string())));
531    }
532
533    #[test]
534    fn test_add_edge() {
535        let graph = create_test_graph().unwrap();
536        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
537        processor.set_batch_size(1); // Force immediate flush
538        
539        processor.add_edge("A".to_string(), "C".to_string(), 3.0).unwrap();
540        
541        let stats = processor.statistics();
542        assert_eq!(stats.edge_count, 3);
543        assert_eq!(
544            processor.edge_cache.get(&("A".to_string(), "C".to_string())),
545            Some(&3.0)
546        );
547    }
548
549    #[test]
550    fn test_remove_edge() {
551        let graph = create_test_graph().unwrap();
552        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
553        processor.set_batch_size(1); // Force immediate flush
554        
555        processor.remove_edge("A".to_string(), "B".to_string()).unwrap();
556        
557        let stats = processor.statistics();
558        assert_eq!(stats.edge_count, 1);
559        assert!(!processor.edge_cache.contains_key(&("A".to_string(), "B".to_string())));
560    }
561
562    #[test]
563    fn test_batch_operations() {
564        let graph = create_test_graph().unwrap();
565        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
566        processor.set_batch_size(10); // Large batch size to test batching
567        
568        // Add multiple operations without flushing
569        processor.add_vertex("D".to_string()).unwrap();
570        processor.add_vertex("E".to_string()).unwrap();
571        processor.add_edge("D".to_string(), "E".to_string(), 4.0).unwrap();
572        
573        // Operations should be pending
574        assert!(processor.has_pending_operations());
575        assert_eq!(processor.pending_operations_count(), 3);
576        
577        // Flush all operations
578        let result = processor.flush_all_operations().unwrap();
579        
580        assert_eq!(result.vertices_added, 2);
581        assert_eq!(result.edges_added, 1);
582        assert!(!processor.has_pending_operations());
583        
584        let stats = processor.statistics();
585        assert_eq!(stats.vertex_count, 5);
586        assert_eq!(stats.edge_count, 3);
587    }
588
589    #[test]
590    fn test_apply_updates() {
591        let graph = create_test_graph().unwrap();
592        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
593        
594        let operations = vec![
595            UpdateOperation::AddVertex("D".to_string()),
596            UpdateOperation::AddVertex("E".to_string()),
597            UpdateOperation::AddEdge("D".to_string(), "E".to_string(), 5.0),
598            UpdateOperation::RemoveEdge("A".to_string(), "B".to_string()),
599            UpdateOperation::RemoveVertex("C".to_string()),
600        ];
601        
602        let result = processor.apply_updates(operations).unwrap();
603        
604        assert_eq!(result.vertices_added, 2);
605        assert_eq!(result.vertices_removed, 1);
606        assert_eq!(result.edges_added, 1);
607        assert_eq!(result.edges_removed, 2); // A->B and B->C (removed with vertex C)
608        assert!(result.recomputation_needed);
609        
610        let stats = processor.statistics();
611        assert_eq!(stats.vertex_count, 4); // A, B, D, E (C removed)
612        assert_eq!(stats.edge_count, 1); // Only D->E remains
613    }
614
615    #[test]
616    fn test_duplicate_operations() {
617        let graph = create_test_graph().unwrap();
618        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
619        processor.set_batch_size(1); // Force immediate flush
620        
621        // Add same vertex multiple times
622        processor.add_vertex("D".to_string()).unwrap();
623        processor.add_vertex("D".to_string()).unwrap();
624        
625        let stats = processor.statistics();
626        assert_eq!(stats.vertex_count, 4); // Should only add once
627        
628        // Remove non-existent vertex
629        processor.remove_vertex("Z".to_string()).unwrap();
630        
631        let stats = processor.statistics();
632        assert_eq!(stats.vertex_count, 4); // Should remain unchanged
633    }
634
635    #[test]
636    fn test_edge_with_new_vertices() {
637        let graph = create_test_graph().unwrap();
638        let mut processor = IncrementalGraphProcessor::new(graph).unwrap();
639        processor.set_batch_size(1); // Force immediate flush
640        
641        // Add edge with new vertices
642        processor.add_edge("X".to_string(), "Y".to_string(), 10.0).unwrap();
643        
644        let stats = processor.statistics();
645        assert_eq!(stats.vertex_count, 5); // Original 3 + X + Y
646        assert_eq!(stats.edge_count, 3); // Original 2 + X->Y
647        
648        assert!(processor.vertex_cache.contains_key("X"));
649        assert!(processor.vertex_cache.contains_key("Y"));
650        assert_eq!(
651            processor.edge_cache.get(&("X".to_string(), "Y".to_string())),
652            Some(&10.0)
653        );
654    }
655}
arrow_graph/streaming/incremental.rs

arrow_graph/streaming/
incremental.rs
