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

arrow_graph/streaming/
sampling.rs
