1use arrow::record_batch::RecordBatch;
2use arrow::array::{StringArray, Float64Array, UInt64Array};
3use arrow::datatypes::{DataType, Field, Schema};
4use std::sync::Arc;
5use std::collections::HashMap;
6use crate::graph::ArrowGraph;
7use crate::error::{GraphError, Result};
8
9#[derive(Debug, Clone)]
11pub enum StreamUpdate {
12 AddNode { node_id: String },
14 RemoveNode { node_id: String },
16 AddEdge { source: String, target: String, weight: Option<f64> },
18 RemoveEdge { source: String, target: String },
20 UpdateEdgeWeight { source: String, target: String, weight: f64 },
22 Batch { operations: Vec<StreamUpdate> },
24}
25
26pub struct StreamingGraphProcessor {
28 graph: ArrowGraph,
29 update_count: u64,
30 change_log: Vec<(u64, StreamUpdate)>,
31 enable_change_log: bool,
32}
33
34impl StreamingGraphProcessor {
35 pub fn new(initial_graph: ArrowGraph) -> Self {
37 Self {
38 graph: initial_graph,
39 update_count: 0,
40 change_log: Vec::new(),
41 enable_change_log: false,
42 }
43 }
44
45 pub fn empty() -> Result<Self> {
47 let empty_graph = ArrowGraph::empty()?;
48 Ok(Self::new(empty_graph))
49 }
50
51 pub fn set_change_log_enabled(&mut self, enabled: bool) {
53 self.enable_change_log = enabled;
54 if !enabled {
55 self.change_log.clear();
56 }
57 }
58
59 pub fn graph(&self) -> &ArrowGraph {
61 &self.graph
62 }
63
64 pub fn update_count(&self) -> u64 {
66 self.update_count
67 }
68
69 pub fn apply_update(&mut self, update: StreamUpdate) -> Result<UpdateResult> {
71 let start_time = std::time::Instant::now();
72 let _initial_node_count = self.graph.node_count();
73 let _initial_edge_count = self.graph.edge_count();
74
75 if self.enable_change_log {
77 self.change_log.push((self.update_count, update.clone()));
78 }
79
80 let result = match update {
81 StreamUpdate::AddNode { node_id } => {
82 self.graph.add_node(node_id.clone())?;
83 UpdateResult {
84 operation: "add_node".to_string(),
85 affected_nodes: vec![node_id],
86 affected_edges: Vec::new(),
87 nodes_added: 1,
88 nodes_removed: 0,
89 edges_added: 0,
90 edges_removed: 0,
91 processing_time_ms: start_time.elapsed().as_millis() as f64,
92 }
93 }
94 StreamUpdate::RemoveNode { node_id } => {
95 let edges_to_remove = self.count_node_edges(&node_id);
97 self.graph.remove_node(&node_id)?;
98 UpdateResult {
99 operation: "remove_node".to_string(),
100 affected_nodes: vec![node_id],
101 affected_edges: Vec::new(),
102 nodes_added: 0,
103 nodes_removed: 1,
104 edges_added: 0,
105 edges_removed: edges_to_remove,
106 processing_time_ms: start_time.elapsed().as_millis() as f64,
107 }
108 }
109 StreamUpdate::AddEdge { source, target, weight } => {
110 self.graph.add_edge(source.clone(), target.clone(), weight)?;
111 UpdateResult {
112 operation: "add_edge".to_string(),
113 affected_nodes: vec![source.clone(), target.clone()],
114 affected_edges: vec![format!("{}→{}", source, target)],
115 nodes_added: 0,
116 nodes_removed: 0,
117 edges_added: 1,
118 edges_removed: 0,
119 processing_time_ms: start_time.elapsed().as_millis() as f64,
120 }
121 }
122 StreamUpdate::RemoveEdge { source, target } => {
123 self.graph.remove_edge(&source, &target)?;
124 UpdateResult {
125 operation: "remove_edge".to_string(),
126 affected_nodes: vec![source.clone(), target.clone()],
127 affected_edges: vec![format!("{}→{}", source, target)],
128 nodes_added: 0,
129 nodes_removed: 0,
130 edges_added: 0,
131 edges_removed: 1,
132 processing_time_ms: start_time.elapsed().as_millis() as f64,
133 }
134 }
135 StreamUpdate::UpdateEdgeWeight { source, target, weight } => {
136 self.graph.remove_edge(&source, &target)?;
138 self.graph.add_edge(source.clone(), target.clone(), Some(weight))?;
139 UpdateResult {
140 operation: "update_edge_weight".to_string(),
141 affected_nodes: vec![source.clone(), target.clone()],
142 affected_edges: vec![format!("{}→{}", source, target)],
143 nodes_added: 0,
144 nodes_removed: 0,
145 edges_added: 0,
146 edges_removed: 0,
147 processing_time_ms: start_time.elapsed().as_millis() as f64,
148 }
149 }
150 StreamUpdate::Batch { operations } => {
151 let mut batch_result = UpdateResult {
152 operation: "batch".to_string(),
153 affected_nodes: Vec::new(),
154 affected_edges: Vec::new(),
155 nodes_added: 0,
156 nodes_removed: 0,
157 edges_added: 0,
158 edges_removed: 0,
159 processing_time_ms: 0.0,
160 };
161
162 for op in operations {
163 let op_result = self.apply_update(op)?;
164 batch_result.merge(op_result);
165 }
166
167 batch_result.processing_time_ms = start_time.elapsed().as_millis() as f64;
168 batch_result
169 }
170 };
171
172 self.update_count += 1;
173 Ok(result)
174 }
175
176 pub fn apply_updates(&mut self, updates: Vec<StreamUpdate>) -> Result<Vec<UpdateResult>> {
178 let mut results = Vec::new();
179 for update in updates {
180 results.push(self.apply_update(update)?);
181 }
182 Ok(results)
183 }
184
185 pub fn get_change_log_since(&self, since_update: u64) -> Vec<(u64, StreamUpdate)> {
187 self.change_log
188 .iter()
189 .filter(|(update_count, _)| *update_count >= since_update)
190 .cloned()
191 .collect()
192 }
193
194 pub fn get_statistics(&self) -> StreamingStatistics {
196 StreamingStatistics {
197 total_updates: self.update_count,
198 current_node_count: self.graph.node_count(),
199 current_edge_count: self.graph.edge_count(),
200 change_log_size: self.change_log.len(),
201 change_log_enabled: self.enable_change_log,
202 }
203 }
204
205 pub fn create_snapshot(&self) -> Result<GraphSnapshot> {
207 Ok(GraphSnapshot {
208 update_count: self.update_count,
209 graph: self.graph.clone(),
210 timestamp: std::time::SystemTime::now(),
211 })
212 }
213
214 pub fn restore_from_snapshot(&mut self, snapshot: GraphSnapshot) {
216 self.graph = snapshot.graph;
217 self.update_count = snapshot.update_count;
218 self.change_log.clear();
220 }
221
222 pub fn compact_change_log(&mut self, keep_since: u64) {
224 if self.enable_change_log {
225 self.change_log.retain(|(update_count, _)| *update_count >= keep_since);
226 }
227 }
228
229 fn count_node_edges(&self, node_id: &str) -> usize {
231 let mut count = 0;
232
233 if let Some(neighbors) = self.graph.neighbors(node_id) {
235 count += neighbors.len();
236 }
237
238 for other_node in self.graph.node_ids() {
240 if other_node != node_id {
241 if let Some(neighbors) = self.graph.neighbors(other_node) {
242 count += neighbors.iter().filter(|&n| n == node_id).count();
243 }
244 }
245 }
246
247 count
248 }
249}
250
251#[derive(Debug, Clone)]
253pub struct UpdateResult {
254 pub operation: String,
255 pub affected_nodes: Vec<String>,
256 pub affected_edges: Vec<String>,
257 pub nodes_added: usize,
258 pub nodes_removed: usize,
259 pub edges_added: usize,
260 pub edges_removed: usize,
261 pub processing_time_ms: f64,
262}
263
264impl UpdateResult {
265 pub fn merge(&mut self, other: UpdateResult) {
267 self.affected_nodes.extend(other.affected_nodes);
268 self.affected_edges.extend(other.affected_edges);
269 self.nodes_added += other.nodes_added;
270 self.nodes_removed += other.nodes_removed;
271 self.edges_added += other.edges_added;
272 self.edges_removed += other.edges_removed;
273 }
275
276 pub fn to_record_batch(&self) -> Result<RecordBatch> {
278 let schema = Arc::new(Schema::new(vec![
279 Field::new("operation", DataType::Utf8, false),
280 Field::new("nodes_added", DataType::UInt64, false),
281 Field::new("nodes_removed", DataType::UInt64, false),
282 Field::new("edges_added", DataType::UInt64, false),
283 Field::new("edges_removed", DataType::UInt64, false),
284 Field::new("processing_time_ms", DataType::Float64, false),
285 ]));
286
287 RecordBatch::try_new(
288 schema,
289 vec![
290 Arc::new(StringArray::from(vec![self.operation.clone()])),
291 Arc::new(UInt64Array::from(vec![self.nodes_added as u64])),
292 Arc::new(UInt64Array::from(vec![self.nodes_removed as u64])),
293 Arc::new(UInt64Array::from(vec![self.edges_added as u64])),
294 Arc::new(UInt64Array::from(vec![self.edges_removed as u64])),
295 Arc::new(Float64Array::from(vec![self.processing_time_ms])),
296 ],
297 ).map_err(GraphError::from)
298 }
299}
300
301#[derive(Debug, Clone)]
303pub struct StreamingStatistics {
304 pub total_updates: u64,
305 pub current_node_count: usize,
306 pub current_edge_count: usize,
307 pub change_log_size: usize,
308 pub change_log_enabled: bool,
309}
310
311#[derive(Debug, Clone)]
313pub struct GraphSnapshot {
314 pub update_count: u64,
315 pub graph: ArrowGraph,
316 pub timestamp: std::time::SystemTime,
317}
318
319pub struct IncrementalAlgorithmProcessor {
321 cached_results: HashMap<String, (u64, RecordBatch)>,
322 invalidation_threshold: u64,
323}
324
325impl IncrementalAlgorithmProcessor {
326 pub fn new() -> Self {
328 Self {
329 cached_results: HashMap::new(),
330 invalidation_threshold: 10, }
332 }
333
334 pub fn set_invalidation_threshold(&mut self, threshold: u64) {
336 self.invalidation_threshold = threshold;
337 }
338
339 pub fn is_cache_valid(&self, algorithm_name: &str, current_update_count: u64) -> bool {
341 if let Some((cached_update_count, _)) = self.cached_results.get(algorithm_name) {
342 current_update_count - cached_update_count <= self.invalidation_threshold
343 } else {
344 false
345 }
346 }
347
348 pub fn cache_result(&mut self, algorithm_name: String, update_count: u64, result: RecordBatch) {
350 self.cached_results.insert(algorithm_name, (update_count, result));
351 }
352
353 pub fn get_cached_result(&self, algorithm_name: &str, current_update_count: u64) -> Option<RecordBatch> {
355 if self.is_cache_valid(algorithm_name, current_update_count) {
356 self.cached_results.get(algorithm_name).map(|(_, result)| result.clone())
357 } else {
358 None
359 }
360 }
361
362 pub fn invalidate_cache(&mut self, algorithm_name: &str) {
364 self.cached_results.remove(algorithm_name);
365 }
366
367 pub fn clear_cache(&mut self) {
369 self.cached_results.clear();
370 }
371
372 pub fn get_cache_statistics(&self) -> CacheStatistics {
374 CacheStatistics {
375 cached_algorithms: self.cached_results.len(),
376 invalidation_threshold: self.invalidation_threshold,
377 total_cache_size: self.cached_results.values()
378 .map(|(_, batch)| batch.get_array_memory_size())
379 .sum(),
380 }
381 }
382}
383
384impl Default for IncrementalAlgorithmProcessor {
385 fn default() -> Self {
386 Self::new()
387 }
388}
389
390#[derive(Debug, Clone)]
392pub struct CacheStatistics {
393 pub cached_algorithms: usize,
394 pub invalidation_threshold: u64,
395 pub total_cache_size: usize,
396}
397
398pub struct StreamingGraphSystem {
400 graph_processor: StreamingGraphProcessor,
401 algorithm_processor: IncrementalAlgorithmProcessor,
402}
403
404impl StreamingGraphSystem {
405 pub fn new(initial_graph: ArrowGraph) -> Self {
407 Self {
408 graph_processor: StreamingGraphProcessor::new(initial_graph),
409 algorithm_processor: IncrementalAlgorithmProcessor::new(),
410 }
411 }
412
413 pub fn empty() -> Result<Self> {
415 Ok(Self {
416 graph_processor: StreamingGraphProcessor::empty()?,
417 algorithm_processor: IncrementalAlgorithmProcessor::new(),
418 })
419 }
420
421 pub fn graph_processor(&self) -> &StreamingGraphProcessor {
423 &self.graph_processor
424 }
425
426 pub fn graph_processor_mut(&mut self) -> &mut StreamingGraphProcessor {
428 &mut self.graph_processor
429 }
430
431 pub fn algorithm_processor(&self) -> &IncrementalAlgorithmProcessor {
433 &self.algorithm_processor
434 }
435
436 pub fn algorithm_processor_mut(&mut self) -> &mut IncrementalAlgorithmProcessor {
438 &mut self.algorithm_processor
439 }
440
441 pub fn apply_update_with_cache_invalidation(&mut self, update: StreamUpdate) -> Result<UpdateResult> {
443 let result = self.graph_processor.apply_update(update)?;
444
445 match result.operation.as_str() {
447 "add_node" | "remove_node" => {
448 self.algorithm_processor.clear_cache();
450 }
451 "add_edge" | "remove_edge" | "update_edge_weight" => {
452 self.algorithm_processor.invalidate_cache("pagerank");
454 self.algorithm_processor.invalidate_cache("betweenness_centrality");
455 self.algorithm_processor.invalidate_cache("closeness_centrality");
456 self.algorithm_processor.invalidate_cache("eigenvector_centrality");
457 self.algorithm_processor.invalidate_cache("shortest_path");
458 }
459 _ => {}
460 }
461
462 Ok(result)
463 }
464}