arrow_graph/streaming/
sampling.rs

1use crate::error::Result;
2use crate::streaming::incremental::IncrementalGraphProcessor;
3use arrow::array::Array;
4use std::collections::{HashMap, HashSet, VecDeque};
5use rand::{Rng, SeedableRng};
6use rand_pcg::Pcg64;
7
8/// Graph sampling strategies for real-time analytics on large graphs
9/// These strategies provide representative subgraphs for efficient processing
10
11/// Random sampling configuration
12#[derive(Debug, Clone)]
13pub struct SamplingConfig {
14    pub seed: Option<u64>,
15    pub sample_rate: f64,        // Fraction of elements to sample (0.0 - 1.0)
16    pub min_sample_size: usize,  // Minimum number of elements to sample
17    pub max_sample_size: usize,  // Maximum number of elements to sample
18}
19
20impl Default for SamplingConfig {
21    fn default() -> Self {
22        Self {
23            seed: None,
24            sample_rate: 0.1, // 10% sample
25            min_sample_size: 100,
26            max_sample_size: 10000,
27        }
28    }
29}
30
31/// Node sampling strategy that selects representative nodes
32#[derive(Debug)]
33pub struct NodeSampler {
34    rng: Pcg64,
35    config: SamplingConfig,
36    sampled_nodes: HashSet<String>,
37    node_scores: HashMap<String, f64>, // Scores for importance-based sampling
38}
39
40impl NodeSampler {
41    pub fn new(config: SamplingConfig) -> Self {
42        let rng = if let Some(seed) = config.seed {
43            Pcg64::seed_from_u64(seed)
44        } else {
45            Pcg64::from_entropy()
46        };
47
48        Self {
49            rng,
50            config,
51            sampled_nodes: HashSet::new(),
52            node_scores: HashMap::new(),
53        }
54    }
55
56    /// Uniform random node sampling
57    pub fn uniform_sample(&mut self, processor: &IncrementalGraphProcessor) -> Result<HashSet<String>> {
58        let graph = processor.graph();
59        let nodes_batch = &graph.nodes;
60        
61        let mut all_nodes = Vec::new();
62        if nodes_batch.num_rows() > 0 {
63            let node_ids = nodes_batch.column(0)
64                .as_any()
65                .downcast_ref::<arrow::array::StringArray>()
66                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for node IDs"))?;
67
68            for i in 0..node_ids.len() {
69                all_nodes.push(node_ids.value(i).to_string());
70            }
71        }
72
73        let sample_size = self.calculate_sample_size(all_nodes.len());
74        let mut sampled = HashSet::new();
75
76        while sampled.len() < sample_size && sampled.len() < all_nodes.len() {
77            let index = self.rng.gen_range(0..all_nodes.len());
78            sampled.insert(all_nodes[index].clone());
79        }
80
81        self.sampled_nodes = sampled.clone();
82        Ok(sampled)
83    }
84
85    /// Degree-based importance sampling (higher degree = higher probability)
86    pub fn degree_based_sample(&mut self, processor: &IncrementalGraphProcessor) -> Result<HashSet<String>> {
87        let graph = processor.graph();
88        let degree_map = self.calculate_node_degrees(processor)?;
89        
90        // Calculate sampling probabilities based on degree
91        let total_degree: u32 = degree_map.values().sum();
92        let mut cumulative_probs = Vec::new();
93        let mut nodes = Vec::new();
94        let mut cumulative = 0.0;
95        
96        for (node, degree) in &degree_map {
97            nodes.push(node.clone());
98            cumulative += (*degree as f64) / (total_degree as f64);
99            cumulative_probs.push(cumulative);
100        }
101
102        let sample_size = self.calculate_sample_size(nodes.len());
103        let mut sampled = HashSet::new();
104
105        for _ in 0..sample_size {
106            let rand_val = self.rng.gen::<f64>();
107            
108            // Find first cumulative probability >= rand_val
109            for (i, &prob) in cumulative_probs.iter().enumerate() {
110                if rand_val <= prob {
111                    sampled.insert(nodes[i].clone());
112                    break;
113                }
114            }
115        }
116
117        self.sampled_nodes = sampled.clone();
118        Ok(sampled)
119    }
120
121    /// PageRank-based importance sampling
122    pub fn pagerank_based_sample(&mut self, processor: &IncrementalGraphProcessor, pagerank_scores: &HashMap<String, f64>) -> Result<HashSet<String>> {
123        // Calculate sampling probabilities based on PageRank scores
124        let total_score: f64 = pagerank_scores.values().sum();
125        let mut cumulative_probs = Vec::new();
126        let mut nodes = Vec::new();
127        let mut cumulative = 0.0;
128        
129        for (node, score) in pagerank_scores {
130            nodes.push(node.clone());
131            cumulative += score / total_score;
132            cumulative_probs.push(cumulative);
133        }
134
135        let sample_size = self.calculate_sample_size(nodes.len());
136        let mut sampled = HashSet::new();
137
138        for _ in 0..sample_size {
139            let rand_val = self.rng.gen::<f64>();
140            
141            for (i, &prob) in cumulative_probs.iter().enumerate() {
142                if rand_val <= prob {
143                    sampled.insert(nodes[i].clone());
144                    break;
145                }
146            }
147        }
148
149        self.sampled_nodes = sampled.clone();
150        Ok(sampled)
151    }
152
153    /// Calculate node degrees
154    fn calculate_node_degrees(&self, processor: &IncrementalGraphProcessor) -> Result<HashMap<String, u32>> {
155        let graph = processor.graph();
156        let edges_batch = &graph.edges;
157        let mut degrees = HashMap::new();
158
159        if edges_batch.num_rows() > 0 {
160            let source_ids = edges_batch.column(0)
161                .as_any()
162                .downcast_ref::<arrow::array::StringArray>()
163                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for source IDs"))?;
164            let target_ids = edges_batch.column(1)
165                .as_any()
166                .downcast_ref::<arrow::array::StringArray>()
167                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for target IDs"))?;
168
169            for i in 0..source_ids.len() {
170                let source = source_ids.value(i).to_string();
171                let target = target_ids.value(i).to_string();
172                
173                *degrees.entry(source).or_insert(0) += 1;
174                *degrees.entry(target).or_insert(0) += 1;
175            }
176        }
177
178        Ok(degrees)
179    }
180
181    fn calculate_sample_size(&self, total_size: usize) -> usize {
182        let target_size = (total_size as f64 * self.config.sample_rate) as usize;
183        target_size.clamp(self.config.min_sample_size, self.config.max_sample_size.min(total_size))
184    }
185
186    pub fn sampled_nodes(&self) -> &HashSet<String> {
187        &self.sampled_nodes
188    }
189}
190
191/// Edge sampling strategy that selects representative edges
192#[derive(Debug)]
193pub struct EdgeSampler {
194    rng: Pcg64,
195    config: SamplingConfig,
196    sampled_edges: HashSet<(String, String)>,
197}
198
199impl EdgeSampler {
200    pub fn new(config: SamplingConfig) -> Self {
201        let rng = if let Some(seed) = config.seed {
202            Pcg64::seed_from_u64(seed)
203        } else {
204            Pcg64::from_entropy()
205        };
206
207        Self {
208            rng,
209            config,
210            sampled_edges: HashSet::new(),
211        }
212    }
213
214    /// Uniform random edge sampling
215    pub fn uniform_sample(&mut self, processor: &IncrementalGraphProcessor) -> Result<HashSet<(String, String)>> {
216        let graph = processor.graph();
217        let edges_batch = &graph.edges;
218        
219        let mut all_edges = Vec::new();
220        if edges_batch.num_rows() > 0 {
221            let source_ids = edges_batch.column(0)
222                .as_any()
223                .downcast_ref::<arrow::array::StringArray>()
224                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for source IDs"))?;
225            let target_ids = edges_batch.column(1)
226                .as_any()
227                .downcast_ref::<arrow::array::StringArray>()
228                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for target IDs"))?;
229
230            for i in 0..source_ids.len() {
231                all_edges.push((source_ids.value(i).to_string(), target_ids.value(i).to_string()));
232            }
233        }
234
235        let sample_size = self.calculate_sample_size(all_edges.len());
236        let mut sampled = HashSet::new();
237
238        while sampled.len() < sample_size && sampled.len() < all_edges.len() {
239            let index = self.rng.gen_range(0..all_edges.len());
240            sampled.insert(all_edges[index].clone());
241        }
242
243        self.sampled_edges = sampled.clone();
244        Ok(sampled)
245    }
246
247    /// Weight-based edge sampling (higher weight = higher probability)
248    pub fn weight_based_sample(&mut self, processor: &IncrementalGraphProcessor) -> Result<HashSet<(String, String)>> {
249        let graph = processor.graph();
250        let edges_batch = &graph.edges;
251        
252        let mut edges_with_weights = Vec::new();
253        let mut total_weight = 0.0;
254
255        if edges_batch.num_rows() > 0 {
256            let source_ids = edges_batch.column(0)
257                .as_any()
258                .downcast_ref::<arrow::array::StringArray>()
259                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for source IDs"))?;
260            let target_ids = edges_batch.column(1)
261                .as_any()
262                .downcast_ref::<arrow::array::StringArray>()
263                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for target IDs"))?;
264            let weights = edges_batch.column(2)
265                .as_any()
266                .downcast_ref::<arrow::array::Float64Array>()
267                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected float64 array for weights"))?;
268
269            for i in 0..source_ids.len() {
270                let weight = weights.value(i).abs(); // Use absolute weight
271                edges_with_weights.push(((source_ids.value(i).to_string(), target_ids.value(i).to_string()), weight));
272                total_weight += weight;
273            }
274        }
275
276        // Calculate cumulative probabilities
277        let mut cumulative_probs = Vec::new();
278        let mut cumulative = 0.0;
279        
280        for (_, weight) in &edges_with_weights {
281            cumulative += weight / total_weight;
282            cumulative_probs.push(cumulative);
283        }
284
285        let sample_size = self.calculate_sample_size(edges_with_weights.len());
286        let mut sampled = HashSet::new();
287
288        for _ in 0..sample_size {
289            let rand_val = self.rng.gen::<f64>();
290            
291            for (i, &prob) in cumulative_probs.iter().enumerate() {
292                if rand_val <= prob {
293                    sampled.insert(edges_with_weights[i].0.clone());
294                    break;
295                }
296            }
297        }
298
299        self.sampled_edges = sampled.clone();
300        Ok(sampled)
301    }
302
303    fn calculate_sample_size(&self, total_size: usize) -> usize {
304        let target_size = (total_size as f64 * self.config.sample_rate) as usize;
305        target_size.clamp(self.config.min_sample_size, self.config.max_sample_size.min(total_size))
306    }
307
308    pub fn sampled_edges(&self) -> &HashSet<(String, String)> {
309        &self.sampled_edges
310    }
311}
312
313/// Subgraph sampling strategy that extracts connected subgraphs
314#[derive(Debug)]
315pub struct SubgraphSampler {
316    rng: Pcg64,
317    config: SamplingConfig,
318    sampled_subgraph: SampledSubgraph,
319}
320
321/// A sampled subgraph containing nodes and edges
322#[derive(Debug, Clone)]
323pub struct SampledSubgraph {
324    pub nodes: HashSet<String>,
325    pub edges: HashSet<(String, String)>,
326    pub sampling_method: String,
327}
328
329impl SubgraphSampler {
330    pub fn new(config: SamplingConfig) -> Self {
331        let rng = if let Some(seed) = config.seed {
332            Pcg64::seed_from_u64(seed)
333        } else {
334            Pcg64::from_entropy()
335        };
336
337        Self {
338            rng,
339            config,
340            sampled_subgraph: SampledSubgraph {
341                nodes: HashSet::new(),
342                edges: HashSet::new(),
343                sampling_method: "none".to_string(),
344            },
345        }
346    }
347
348    /// Random walk sampling starting from a random node
349    pub fn random_walk_sample(&mut self, processor: &IncrementalGraphProcessor, walk_length: usize) -> Result<SampledSubgraph> {
350        let graph = processor.graph();
351        let adjacency = self.build_adjacency_list(processor)?;
352        
353        // Get all nodes
354        let all_nodes: Vec<String> = adjacency.keys().cloned().collect();
355        if all_nodes.is_empty() {
356            return Ok(SampledSubgraph {
357                nodes: HashSet::new(),
358                edges: HashSet::new(),
359                sampling_method: "random_walk".to_string(),
360            });
361        }
362
363        // Start from random node
364        let start_node = &all_nodes[self.rng.gen_range(0..all_nodes.len())];
365        let mut visited_nodes = HashSet::new();
366        let mut visited_edges = HashSet::new();
367        let mut current_node = start_node.clone();
368
369        visited_nodes.insert(current_node.clone());
370
371        for _ in 0..walk_length {
372            if let Some(neighbors) = adjacency.get(&current_node) {
373                if neighbors.is_empty() {
374                    break; // Dead end
375                }
376                
377                let next_node = &neighbors[self.rng.gen_range(0..neighbors.len())];
378                visited_edges.insert((current_node.clone(), next_node.clone()));
379                visited_nodes.insert(next_node.clone());
380                current_node = next_node.clone();
381            } else {
382                break; // No neighbors
383            }
384        }
385
386        let subgraph = SampledSubgraph {
387            nodes: visited_nodes,
388            edges: visited_edges,
389            sampling_method: "random_walk".to_string(),
390        };
391
392        self.sampled_subgraph = subgraph.clone();
393        Ok(subgraph)
394    }
395
396    /// Breadth-First Search (BFS) sampling from multiple seed nodes
397    pub fn bfs_sample(&mut self, processor: &IncrementalGraphProcessor, num_seeds: usize, max_depth: usize) -> Result<SampledSubgraph> {
398        let adjacency = self.build_adjacency_list(processor)?;
399        let all_nodes: Vec<String> = adjacency.keys().cloned().collect();
400        
401        if all_nodes.is_empty() {
402            return Ok(SampledSubgraph {
403                nodes: HashSet::new(),
404                edges: HashSet::new(),
405                sampling_method: "bfs".to_string(),
406            });
407        }
408
409        let mut visited_nodes = HashSet::new();
410        let mut visited_edges = HashSet::new();
411        let mut queue = VecDeque::new();
412
413        // Select random seed nodes
414        let actual_seeds = num_seeds.min(all_nodes.len());
415        let mut selected_seeds = HashSet::new();
416        
417        while selected_seeds.len() < actual_seeds {
418            let seed = &all_nodes[self.rng.gen_range(0..all_nodes.len())];
419            if selected_seeds.insert(seed.clone()) {
420                queue.push_back((seed.clone(), 0)); // (node, depth)
421                visited_nodes.insert(seed.clone());
422            }
423        }
424
425        // BFS exploration
426        while let Some((current_node, depth)) = queue.pop_front() {
427            if depth >= max_depth {
428                continue;
429            }
430
431            if let Some(neighbors) = adjacency.get(&current_node) {
432                for neighbor in neighbors {
433                    visited_edges.insert((current_node.clone(), neighbor.clone()));
434                    
435                    if visited_nodes.insert(neighbor.clone()) {
436                        queue.push_back((neighbor.clone(), depth + 1));
437                    }
438                }
439            }
440        }
441
442        let subgraph = SampledSubgraph {
443            nodes: visited_nodes,
444            edges: visited_edges,
445            sampling_method: "bfs".to_string(),
446        };
447
448        self.sampled_subgraph = subgraph.clone();
449        Ok(subgraph)
450    }
451
452    /// Forest Fire sampling (spreads like fire with burning probability)
453    pub fn forest_fire_sample(&mut self, processor: &IncrementalGraphProcessor, burn_probability: f64) -> Result<SampledSubgraph> {
454        let adjacency = self.build_adjacency_list(processor)?;
455        let all_nodes: Vec<String> = adjacency.keys().cloned().collect();
456        
457        if all_nodes.is_empty() {
458            return Ok(SampledSubgraph {
459                nodes: HashSet::new(),
460                edges: HashSet::new(),
461                sampling_method: "forest_fire".to_string(),
462            });
463        }
464
465        // Start from random node
466        let start_node = &all_nodes[self.rng.gen_range(0..all_nodes.len())];
467        let mut visited_nodes = HashSet::new();
468        let mut visited_edges = HashSet::new();
469        let mut burning_queue = VecDeque::new();
470
471        visited_nodes.insert(start_node.clone());
472        burning_queue.push_back(start_node.clone());
473
474        let target_size = self.calculate_sample_size(all_nodes.len());
475
476        while !burning_queue.is_empty() && visited_nodes.len() < target_size {
477            let current_node = burning_queue.pop_front().unwrap();
478            
479            if let Some(neighbors) = adjacency.get(&current_node) {
480                for neighbor in neighbors {
481                    // Fire spreads with burn_probability
482                    if self.rng.gen::<f64>() < burn_probability {
483                        visited_edges.insert((current_node.clone(), neighbor.clone()));
484                        
485                        if visited_nodes.insert(neighbor.clone()) {
486                            burning_queue.push_back(neighbor.clone());
487                        }
488                    }
489                }
490            }
491        }
492
493        let subgraph = SampledSubgraph {
494            nodes: visited_nodes,
495            edges: visited_edges,
496            sampling_method: "forest_fire".to_string(),
497        };
498
499        self.sampled_subgraph = subgraph.clone();
500        Ok(subgraph)
501    }
502
503    /// Build adjacency list representation
504    fn build_adjacency_list(&self, processor: &IncrementalGraphProcessor) -> Result<HashMap<String, Vec<String>>> {
505        let graph = processor.graph();
506        let edges_batch = &graph.edges;
507        let mut adjacency = HashMap::new();
508
509        if edges_batch.num_rows() > 0 {
510            let source_ids = edges_batch.column(0)
511                .as_any()
512                .downcast_ref::<arrow::array::StringArray>()
513                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for source IDs"))?;
514            let target_ids = edges_batch.column(1)
515                .as_any()
516                .downcast_ref::<arrow::array::StringArray>()
517                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for target IDs"))?;
518
519            for i in 0..source_ids.len() {
520                let source = source_ids.value(i).to_string();
521                let target = target_ids.value(i).to_string();
522                
523                adjacency.entry(source).or_insert_with(Vec::new).push(target);
524            }
525        }
526
527        Ok(adjacency)
528    }
529
530    fn calculate_sample_size(&self, total_size: usize) -> usize {
531        let target_size = (total_size as f64 * self.config.sample_rate) as usize;
532        target_size.clamp(self.config.min_sample_size, self.config.max_sample_size.min(total_size))
533    }
534
535    pub fn sampled_subgraph(&self) -> &SampledSubgraph {
536        &self.sampled_subgraph
537    }
538}
539
540/// Reservoir sampling for streaming scenarios
541#[derive(Debug)]
542pub struct ReservoirSampler<T> {
543    reservoir: Vec<T>,
544    capacity: usize,
545    count: usize,
546    rng: Pcg64,
547}
548
549impl<T> ReservoirSampler<T> {
550    pub fn new(capacity: usize, seed: Option<u64>) -> Self {
551        let rng = if let Some(seed) = seed {
552            Pcg64::seed_from_u64(seed)
553        } else {
554            Pcg64::from_entropy()
555        };
556
557        Self {
558            reservoir: Vec::with_capacity(capacity),
559            capacity,
560            count: 0,
561            rng,
562        }
563    }
564
565    /// Add an item to the reservoir sample
566    pub fn add(&mut self, item: T) {
567        self.count += 1;
568
569        if self.reservoir.len() < self.capacity {
570            // Reservoir not full, just add
571            self.reservoir.push(item);
572        } else {
573            // Reservoir full, decide whether to replace
574            let j = self.rng.gen_range(0..self.count);
575            if j < self.capacity {
576                self.reservoir[j] = item;
577            }
578        }
579    }
580
581    /// Get the current sample
582    pub fn sample(&self) -> &[T] {
583        &self.reservoir
584    }
585
586    /// Get sample size
587    pub fn sample_size(&self) -> usize {
588        self.reservoir.len()
589    }
590
591    /// Get total items seen
592    pub fn total_count(&self) -> usize {
593        self.count
594    }
595
596    /// Reset the sampler
597    pub fn reset(&mut self) {
598        self.reservoir.clear();
599        self.count = 0;
600    }
601}
602
603/// Combined graph sampling processor
604#[derive(Debug)]
605pub struct GraphSamplingProcessor {
606    node_sampler: NodeSampler,
607    edge_sampler: EdgeSampler,
608    subgraph_sampler: SubgraphSampler,
609    reservoir_nodes: ReservoirSampler<String>,
610    reservoir_edges: ReservoirSampler<(String, String)>,
611    config: SamplingConfig,
612}
613
614impl GraphSamplingProcessor {
615    pub fn new(config: SamplingConfig) -> Self {
616        let reservoir_capacity = config.max_sample_size;
617        
618        Self {
619            node_sampler: NodeSampler::new(config.clone()),
620            edge_sampler: EdgeSampler::new(config.clone()),
621            subgraph_sampler: SubgraphSampler::new(config.clone()),
622            reservoir_nodes: ReservoirSampler::new(reservoir_capacity, config.seed),
623            reservoir_edges: ReservoirSampler::new(reservoir_capacity, config.seed),
624            config,
625        }
626    }
627
628    /// Perform comprehensive sampling of the graph
629    pub fn sample_graph(&mut self, processor: &IncrementalGraphProcessor) -> Result<GraphSample> {
630        // Sample nodes using different strategies
631        let uniform_nodes = self.node_sampler.uniform_sample(processor)?;
632        let degree_nodes = self.node_sampler.degree_based_sample(processor)?;
633        
634        // Sample edges
635        let uniform_edges = self.edge_sampler.uniform_sample(processor)?;
636        let weight_edges = self.edge_sampler.weight_based_sample(processor)?;
637        
638        // Sample subgraphs
639        let random_walk_subgraph = self.subgraph_sampler.random_walk_sample(processor, 100)?;
640        let bfs_subgraph = self.subgraph_sampler.bfs_sample(processor, 5, 3)?;
641        let forest_fire_subgraph = self.subgraph_sampler.forest_fire_sample(processor, 0.7)?;
642
643        Ok(GraphSample {
644            uniform_nodes,
645            degree_based_nodes: degree_nodes,
646            uniform_edges,
647            weight_based_edges: weight_edges,
648            random_walk_subgraph,
649            bfs_subgraph,
650            forest_fire_subgraph,
651            reservoir_nodes: self.reservoir_nodes.sample().to_vec(),
652            reservoir_edges: self.reservoir_edges.sample().to_vec(),
653        })
654    }
655
656    /// Update reservoir samplers with new graph elements
657    pub fn update_reservoirs(&mut self, processor: &IncrementalGraphProcessor) -> Result<()> {
658        let graph = processor.graph();
659        
660        // Update node reservoir
661        let nodes_batch = &graph.nodes;
662        if nodes_batch.num_rows() > 0 {
663            let node_ids = nodes_batch.column(0)
664                .as_any()
665                .downcast_ref::<arrow::array::StringArray>()
666                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for node IDs"))?;
667
668            for i in 0..node_ids.len() {
669                self.reservoir_nodes.add(node_ids.value(i).to_string());
670            }
671        }
672
673        // Update edge reservoir
674        let edges_batch = &graph.edges;
675        if edges_batch.num_rows() > 0 {
676            let source_ids = edges_batch.column(0)
677                .as_any()
678                .downcast_ref::<arrow::array::StringArray>()
679                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for source IDs"))?;
680            let target_ids = edges_batch.column(1)
681                .as_any()
682                .downcast_ref::<arrow::array::StringArray>()
683                .ok_or_else(|| crate::error::GraphError::graph_construction("Expected string array for target IDs"))?;
684
685            for i in 0..source_ids.len() {
686                self.reservoir_edges.add((source_ids.value(i).to_string(), target_ids.value(i).to_string()));
687            }
688        }
689
690        Ok(())
691    }
692
693    /// Get sampling statistics
694    pub fn sampling_stats(&self) -> SamplingStats {
695        SamplingStats {
696            reservoir_node_count: self.reservoir_nodes.sample_size(),
697            reservoir_edge_count: self.reservoir_edges.sample_size(),
698            total_nodes_seen: self.reservoir_nodes.total_count(),
699            total_edges_seen: self.reservoir_edges.total_count(),
700            config: self.config.clone(),
701        }
702    }
703}
704
705/// Comprehensive graph sample containing multiple sampling strategies
706#[derive(Debug, Clone)]
707pub struct GraphSample {
708    pub uniform_nodes: HashSet<String>,
709    pub degree_based_nodes: HashSet<String>,
710    pub uniform_edges: HashSet<(String, String)>,
711    pub weight_based_edges: HashSet<(String, String)>,
712    pub random_walk_subgraph: SampledSubgraph,
713    pub bfs_subgraph: SampledSubgraph,
714    pub forest_fire_subgraph: SampledSubgraph,
715    pub reservoir_nodes: Vec<String>,
716    pub reservoir_edges: Vec<(String, String)>,
717}
718
719/// Statistics about sampling performance
720#[derive(Debug, Clone)]
721pub struct SamplingStats {
722    pub reservoir_node_count: usize,
723    pub reservoir_edge_count: usize,
724    pub total_nodes_seen: usize,
725    pub total_edges_seen: usize,
726    pub config: SamplingConfig,
727}
728
729#[cfg(test)]
730mod tests {
731    use super::*;
732    use crate::graph::ArrowGraph;
733    use arrow::array::{StringArray, Float64Array};
734    use arrow::record_batch::RecordBatch;
735    use arrow::datatypes::{Schema, Field, DataType};
736    use std::sync::Arc;
737
738    fn create_test_graph() -> Result<ArrowGraph> {
739        let nodes_schema = Arc::new(Schema::new(vec![
740            Field::new("id", DataType::Utf8, false),
741        ]));
742        let node_ids = StringArray::from(vec!["A", "B", "C", "D", "E"]);
743        let nodes_batch = RecordBatch::try_new(
744            nodes_schema,
745            vec![Arc::new(node_ids)],
746        )?;
747
748        let edges_schema = Arc::new(Schema::new(vec![
749            Field::new("source", DataType::Utf8, false),
750            Field::new("target", DataType::Utf8, false),
751            Field::new("weight", DataType::Float64, false),
752        ]));
753        let sources = StringArray::from(vec!["A", "B", "C", "D", "A", "B"]);
754        let targets = StringArray::from(vec!["B", "C", "D", "E", "C", "D"]);
755        let weights = Float64Array::from(vec![1.0, 2.0, 3.0, 1.0, 4.0, 2.0]);
756        let edges_batch = RecordBatch::try_new(
757            edges_schema,
758            vec![Arc::new(sources), Arc::new(targets), Arc::new(weights)],
759        )?;
760
761        ArrowGraph::new(nodes_batch, edges_batch)
762    }
763
764    #[test]
765    fn test_node_uniform_sampling() {
766        let graph = create_test_graph().unwrap();
767        let processor = IncrementalGraphProcessor::new(graph).unwrap();
768        
769        let config = SamplingConfig {
770            seed: Some(42),
771            sample_rate: 0.6, // 60%
772            min_sample_size: 1,
773            max_sample_size: 3,
774        };
775        
776        let mut sampler = NodeSampler::new(config);
777        let sampled = sampler.uniform_sample(&processor).unwrap();
778        
779        assert!(!sampled.is_empty());
780        assert!(sampled.len() <= 3); // Max sample size
781        
782        // All sampled nodes should be valid
783        let valid_nodes = ["A", "B", "C", "D", "E"];
784        for node in &sampled {
785            assert!(valid_nodes.contains(&node.as_str()));
786        }
787    }
788
789    #[test]
790    fn test_edge_uniform_sampling() {
791        let graph = create_test_graph().unwrap();
792        let processor = IncrementalGraphProcessor::new(graph).unwrap();
793        
794        let config = SamplingConfig {
795            seed: Some(42),
796            sample_rate: 0.5,
797            min_sample_size: 1,
798            max_sample_size: 3,
799        };
800        
801        let mut sampler = EdgeSampler::new(config);
802        let sampled = sampler.uniform_sample(&processor).unwrap();
803        
804        assert!(!sampled.is_empty());
805        assert!(sampled.len() <= 3);
806    }
807
808    #[test]
809    fn test_degree_based_sampling() {
810        let graph = create_test_graph().unwrap();
811        let processor = IncrementalGraphProcessor::new(graph).unwrap();
812        
813        let config = SamplingConfig {
814            seed: Some(42),
815            sample_rate: 0.6,
816            min_sample_size: 1,
817            max_sample_size: 3,
818        };
819        
820        let mut sampler = NodeSampler::new(config);
821        let sampled = sampler.degree_based_sample(&processor).unwrap();
822        
823        assert!(!sampled.is_empty());
824        assert!(sampled.len() <= 3);
825    }
826
827    #[test]
828    fn test_random_walk_sampling() {
829        let graph = create_test_graph().unwrap();
830        let processor = IncrementalGraphProcessor::new(graph).unwrap();
831        
832        let config = SamplingConfig {
833            seed: Some(42),
834            sample_rate: 0.5,
835            min_sample_size: 1,
836            max_sample_size: 10,
837        };
838        
839        let mut sampler = SubgraphSampler::new(config);
840        let subgraph = sampler.random_walk_sample(&processor, 5).unwrap();
841        
842        assert!(!subgraph.nodes.is_empty());
843        assert_eq!(subgraph.sampling_method, "random_walk");
844        
845        // Edges should connect nodes in the subgraph
846        for (source, target) in &subgraph.edges {
847            assert!(subgraph.nodes.contains(source));
848            assert!(subgraph.nodes.contains(target));
849        }
850    }
851
852    #[test]
853    fn test_bfs_sampling() {
854        let graph = create_test_graph().unwrap();
855        let processor = IncrementalGraphProcessor::new(graph).unwrap();
856        
857        let config = SamplingConfig {
858            seed: Some(42),
859            sample_rate: 0.8,
860            min_sample_size: 1,
861            max_sample_size: 10,
862        };
863        
864        let mut sampler = SubgraphSampler::new(config);
865        let subgraph = sampler.bfs_sample(&processor, 2, 2).unwrap();
866        
867        assert!(!subgraph.nodes.is_empty());
868        assert_eq!(subgraph.sampling_method, "bfs");
869    }
870
871    #[test]
872    fn test_forest_fire_sampling() {
873        let graph = create_test_graph().unwrap();
874        let processor = IncrementalGraphProcessor::new(graph).unwrap();
875        
876        let config = SamplingConfig {
877            seed: Some(42),
878            sample_rate: 0.6,
879            min_sample_size: 1,
880            max_sample_size: 10,
881        };
882        
883        let mut sampler = SubgraphSampler::new(config);
884        let subgraph = sampler.forest_fire_sample(&processor, 0.7).unwrap();
885        
886        assert!(!subgraph.nodes.is_empty());
887        assert_eq!(subgraph.sampling_method, "forest_fire");
888    }
889
890    #[test]
891    fn test_reservoir_sampling() {
892        let mut reservoir = ReservoirSampler::new(3, Some(42));
893        
894        // Add more items than capacity
895        for i in 0..10 {
896            reservoir.add(format!("item_{}", i));
897        }
898        
899        assert_eq!(reservoir.sample_size(), 3); // Should maintain capacity
900        assert_eq!(reservoir.total_count(), 10); // Should track all items seen
901        
902        // Sample should contain valid items
903        for item in reservoir.sample() {
904            assert!(item.starts_with("item_"));
905        }
906    }
907
908    #[test]
909    fn test_weight_based_edge_sampling() {
910        let graph = create_test_graph().unwrap();
911        let processor = IncrementalGraphProcessor::new(graph).unwrap();
912        
913        let config = SamplingConfig {
914            seed: Some(42),
915            sample_rate: 0.5,
916            min_sample_size: 1,
917            max_sample_size: 3,
918        };
919        
920        let mut sampler = EdgeSampler::new(config);
921        let sampled = sampler.weight_based_sample(&processor).unwrap();
922        
923        assert!(!sampled.is_empty());
924        assert!(sampled.len() <= 3);
925    }
926
927    #[test]
928    fn test_comprehensive_graph_sampling() {
929        let graph = create_test_graph().unwrap();
930        let processor = IncrementalGraphProcessor::new(graph).unwrap();
931        
932        let config = SamplingConfig {
933            seed: Some(42),
934            sample_rate: 0.6,
935            min_sample_size: 1,
936            max_sample_size: 5,
937        };
938        
939        let mut sampling_processor = GraphSamplingProcessor::new(config);
940        let sample = sampling_processor.sample_graph(&processor).unwrap();
941        
942        // Check that all sampling methods produced results
943        assert!(!sample.uniform_nodes.is_empty());
944        assert!(!sample.degree_based_nodes.is_empty());
945        assert!(!sample.uniform_edges.is_empty());
946        assert!(!sample.weight_based_edges.is_empty());
947        assert!(!sample.random_walk_subgraph.nodes.is_empty());
948        assert!(!sample.bfs_subgraph.nodes.is_empty());
949        assert!(!sample.forest_fire_subgraph.nodes.is_empty());
950    }
951
952    #[test]
953    fn test_reservoir_updates() {
954        let graph = create_test_graph().unwrap();
955        let processor = IncrementalGraphProcessor::new(graph).unwrap();
956        
957        let config = SamplingConfig {
958            seed: Some(42),
959            sample_rate: 0.5,
960            min_sample_size: 1,
961            max_sample_size: 3,
962        };
963        
964        let mut sampling_processor = GraphSamplingProcessor::new(config);
965        sampling_processor.update_reservoirs(&processor).unwrap();
966        
967        let stats = sampling_processor.sampling_stats();
968        assert!(stats.reservoir_node_count > 0);
969        assert!(stats.reservoir_edge_count > 0);
970        assert!(stats.total_nodes_seen > 0);
971        assert!(stats.total_edges_seen > 0);
972    }
973}
arrow_graph/streaming/sampling.rs

arrow_graph/streaming/
sampling.rs
