llm_optimizer_processor/aggregation/
composite.rs

1use anyhow::anyhow;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5/// A composite aggregator that runs multiple aggregations in parallel
6///
7/// This is useful when you want to compute multiple statistics at once
8/// (e.g., count, avg, p95, stddev) without making multiple passes over the data.
9///
10/// # Examples
11///
12/// ```
13/// use processor::aggregation::{CompositeAggregator, AggregationResult};
14///
15/// let mut composite = CompositeAggregator::new();
16/// composite.add_count("request_count");
17/// composite.add_avg("avg_latency");
18/// composite.add_p95("p95_latency");
19/// composite.add_stddev("stddev_latency");
20///
21/// // Feed values
22/// for latency in vec![100.0, 150.0, 200.0, 250.0, 300.0] {
23///     composite.update(latency).unwrap();
24/// }
25///
26/// // Get all results
27/// let results = composite.finalize().unwrap();
28/// println!("Count: {}", results.get_count("request_count").unwrap());
29/// println!("Avg: {}", results.get_value("avg_latency").unwrap());
30/// println!("P95: {}", results.get_value("p95_latency").unwrap());
31/// println!("StdDev: {}", results.get_value("stddev_latency").unwrap());
32/// ```
33#[derive(Debug)]
34pub struct CompositeAggregator {
35    aggregations: Vec<AggregationSpec>,
36}
37
38#[derive(Debug, Clone)]
39enum AggregationSpec {
40    Count {
41        name: String,
42        count: u64,
43    },
44    Sum {
45        name: String,
46        sum: f64,
47        count: u64,
48    },
49    Avg {
50        name: String,
51        sum: f64,
52        count: u64,
53    },
54    Min {
55        name: String,
56        min: Option<f64>,
57        count: u64,
58    },
59    Max {
60        name: String,
61        max: Option<f64>,
62        count: u64,
63    },
64    Percentile {
65        name: String,
66        values: Vec<f64>,
67        percentile: f64,
68    },
69    StdDev {
70        name: String,
71        count: u64,
72        mean: f64,
73        m2: f64,
74    },
75}
76
77/// Results from a composite aggregation
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct AggregationResult {
80    counts: HashMap<String, u64>,
81    values: HashMap<String, f64>,
82}
83
84impl AggregationResult {
85    pub fn new() -> Self {
86        Self {
87            counts: HashMap::new(),
88            values: HashMap::new(),
89        }
90    }
91
92    /// Set a count result
93    pub fn set_count(&mut self, name: impl Into<String>, count: u64) {
94        self.counts.insert(name.into(), count);
95    }
96
97    /// Set a numeric result
98    pub fn set_value(&mut self, name: impl Into<String>, value: f64) {
99        self.values.insert(name.into(), value);
100    }
101
102    /// Get a count result by name
103    pub fn get_count(&self, name: &str) -> Option<u64> {
104        self.counts.get(name).copied()
105    }
106
107    /// Get a numeric result by name
108    pub fn get_value(&self, name: &str) -> Option<f64> {
109        self.values.get(name).copied()
110    }
111
112    /// Get all count results
113    pub fn counts(&self) -> &HashMap<String, u64> {
114        &self.counts
115    }
116
117    /// Get all numeric results
118    pub fn values(&self) -> &HashMap<String, f64> {
119        &self.values
120    }
121}
122
123impl CompositeAggregator {
124    /// Create a new composite aggregator
125    pub fn new() -> Self {
126        Self {
127            aggregations: Vec::new(),
128        }
129    }
130
131    /// Add a count aggregation
132    pub fn add_count(&mut self, name: impl Into<String>) -> &mut Self {
133        self.aggregations.push(AggregationSpec::Count {
134            name: name.into(),
135            count: 0,
136        });
137        self
138    }
139
140    /// Add a sum aggregation
141    pub fn add_sum(&mut self, name: impl Into<String>) -> &mut Self {
142        self.aggregations.push(AggregationSpec::Sum {
143            name: name.into(),
144            sum: 0.0,
145            count: 0,
146        });
147        self
148    }
149
150    /// Add an average aggregation
151    pub fn add_avg(&mut self, name: impl Into<String>) -> &mut Self {
152        self.aggregations.push(AggregationSpec::Avg {
153            name: name.into(),
154            sum: 0.0,
155            count: 0,
156        });
157        self
158    }
159
160    /// Add a min aggregation
161    pub fn add_min(&mut self, name: impl Into<String>) -> &mut Self {
162        self.aggregations.push(AggregationSpec::Min {
163            name: name.into(),
164            min: None,
165            count: 0,
166        });
167        self
168    }
169
170    /// Add a max aggregation
171    pub fn add_max(&mut self, name: impl Into<String>) -> &mut Self {
172        self.aggregations.push(AggregationSpec::Max {
173            name: name.into(),
174            max: None,
175            count: 0,
176        });
177        self
178    }
179
180    /// Add a percentile aggregation
181    pub fn add_percentile(&mut self, name: impl Into<String>, percentile: f64) -> anyhow::Result<&mut Self> {
182        if !(0.0..=100.0).contains(&percentile) {
183            return Err(anyhow!("Percentile must be between 0 and 100"));
184        }
185        self.aggregations.push(AggregationSpec::Percentile {
186            name: name.into(),
187            values: Vec::new(),
188            percentile,
189        });
190        Ok(self)
191    }
192
193    /// Add a p50 (median) aggregation
194    pub fn add_p50(&mut self, name: impl Into<String>) -> &mut Self {
195        self.aggregations.push(AggregationSpec::Percentile {
196            name: name.into(),
197            values: Vec::new(),
198            percentile: 50.0,
199        });
200        self
201    }
202
203    /// Add a p95 aggregation
204    pub fn add_p95(&mut self, name: impl Into<String>) -> &mut Self {
205        self.aggregations.push(AggregationSpec::Percentile {
206            name: name.into(),
207            values: Vec::new(),
208            percentile: 95.0,
209        });
210        self
211    }
212
213    /// Add a p99 aggregation
214    pub fn add_p99(&mut self, name: impl Into<String>) -> &mut Self {
215        self.aggregations.push(AggregationSpec::Percentile {
216            name: name.into(),
217            values: Vec::new(),
218            percentile: 99.0,
219        });
220        self
221    }
222
223    /// Add a standard deviation aggregation
224    pub fn add_stddev(&mut self, name: impl Into<String>) -> &mut Self {
225        self.aggregations.push(AggregationSpec::StdDev {
226            name: name.into(),
227            count: 0,
228            mean: 0.0,
229            m2: 0.0,
230        });
231        self
232    }
233
234    /// Update all aggregations with a new value
235    pub fn update(&mut self, value: f64) -> anyhow::Result<()> {
236        for agg in &mut self.aggregations {
237            match agg {
238                AggregationSpec::Count { count, .. } => {
239                    *count += 1;
240                }
241                AggregationSpec::Sum { sum, count, .. } => {
242                    *sum += value;
243                    *count += 1;
244                }
245                AggregationSpec::Avg { sum, count, .. } => {
246                    *sum += value;
247                    *count += 1;
248                }
249                AggregationSpec::Min { min, count, .. } => {
250                    *min = Some(match *min {
251                        Some(current) => current.min(value),
252                        None => value,
253                    });
254                    *count += 1;
255                }
256                AggregationSpec::Max { max, count, .. } => {
257                    *max = Some(match *max {
258                        Some(current) => current.max(value),
259                        None => value,
260                    });
261                    *count += 1;
262                }
263                AggregationSpec::Percentile { values, .. } => {
264                    values.push(value);
265                }
266                AggregationSpec::StdDev { count, mean, m2, .. } => {
267                    *count += 1;
268                    let delta = value - *mean;
269                    *mean += delta / *count as f64;
270                    let delta2 = value - *mean;
271                    *m2 += delta * delta2;
272                }
273            }
274        }
275        Ok(())
276    }
277
278    /// Update all aggregations with multiple values
279    pub fn update_batch(&mut self, values: &[f64]) -> anyhow::Result<()> {
280        for value in values {
281            self.update(*value)?;
282        }
283        Ok(())
284    }
285
286    /// Compute final results for all aggregations
287    pub fn finalize(&self) -> anyhow::Result<AggregationResult> {
288        let mut result = AggregationResult::new();
289
290        for agg in &self.aggregations {
291            match agg {
292                AggregationSpec::Count { name, count } => {
293                    result.counts.insert(name.clone(), *count);
294                }
295                AggregationSpec::Sum { name, sum, .. } => {
296                    result.values.insert(name.clone(), *sum);
297                }
298                AggregationSpec::Avg { name, sum, count } => {
299                    if *count == 0 {
300                        return Err(anyhow!("Cannot compute average '{}' of zero values", name));
301                    }
302                    result.values.insert(name.clone(), *sum / *count as f64);
303                }
304                AggregationSpec::Min { name, min, .. } => {
305                    let value = min.ok_or_else(|| anyhow!("Cannot compute min '{}' of zero values", name))?;
306                    result.values.insert(name.clone(), value);
307                }
308                AggregationSpec::Max { name, max, .. } => {
309                    let value = max.ok_or_else(|| anyhow!("Cannot compute max '{}' of zero values", name))?;
310                    result.values.insert(name.clone(), value);
311                }
312                AggregationSpec::Percentile { name, values, percentile } => {
313                    if values.is_empty() {
314                        return Err(anyhow!("Cannot compute percentile '{}' of zero values", name));
315                    }
316                    let mut sorted = values.clone();
317                    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
318                    let value = calculate_percentile(&sorted, *percentile)
319                        .ok_or_else(|| anyhow!("Failed to calculate percentile '{}'", name))?;
320                    result.values.insert(name.clone(), value);
321                }
322                AggregationSpec::StdDev { name, count, m2, .. } => {
323                    if *count < 2 {
324                        return Err(anyhow!("Cannot compute stddev '{}' with less than 2 values", name));
325                    }
326                    let variance = *m2 / (*count - 1) as f64;
327                    result.values.insert(name.clone(), variance.sqrt());
328                }
329            }
330        }
331
332        Ok(result)
333    }
334
335    /// Reset all aggregations to initial state
336    pub fn reset(&mut self) {
337        for agg in &mut self.aggregations {
338            match agg {
339                AggregationSpec::Count { count, .. } => {
340                    *count = 0;
341                }
342                AggregationSpec::Sum { sum, count, .. } => {
343                    *sum = 0.0;
344                    *count = 0;
345                }
346                AggregationSpec::Avg { sum, count, .. } => {
347                    *sum = 0.0;
348                    *count = 0;
349                }
350                AggregationSpec::Min { min, count, .. } => {
351                    *min = None;
352                    *count = 0;
353                }
354                AggregationSpec::Max { max, count, .. } => {
355                    *max = None;
356                    *count = 0;
357                }
358                AggregationSpec::Percentile { values, .. } => {
359                    values.clear();
360                }
361                AggregationSpec::StdDev { count, mean, m2, .. } => {
362                    *count = 0;
363                    *mean = 0.0;
364                    *m2 = 0.0;
365                }
366            }
367        }
368    }
369
370    /// Get the number of aggregations
371    pub fn len(&self) -> usize {
372        self.aggregations.len()
373    }
374
375    /// Check if there are no aggregations
376    pub fn is_empty(&self) -> bool {
377        self.aggregations.is_empty()
378    }
379}
380
381impl Default for CompositeAggregator {
382    fn default() -> Self {
383        Self::new()
384    }
385}
386
387fn calculate_percentile(sorted_values: &[f64], percentile: f64) -> Option<f64> {
388    if sorted_values.is_empty() {
389        return None;
390    }
391
392    if sorted_values.len() == 1 {
393        return Some(sorted_values[0]);
394    }
395
396    let rank = (percentile / 100.0) * (sorted_values.len() - 1) as f64;
397    let lower_idx = rank.floor() as usize;
398    let upper_idx = rank.ceil() as usize;
399
400    if lower_idx == upper_idx {
401        Some(sorted_values[lower_idx])
402    } else {
403        let fraction = rank - lower_idx as f64;
404        Some(sorted_values[lower_idx] * (1.0 - fraction)
405            + sorted_values[upper_idx] * fraction)
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412
413    #[test]
414    fn test_composite_basic() {
415        let mut composite = CompositeAggregator::new();
416        composite
417            .add_count("count")
418            .add_sum("sum")
419            .add_avg("avg")
420            .add_min("min")
421            .add_max("max");
422
423        composite.update(10.0).unwrap();
424        composite.update(20.0).unwrap();
425        composite.update(30.0).unwrap();
426
427        let results = composite.finalize().unwrap();
428
429        assert_eq!(results.get_count("count").unwrap(), 3);
430        assert_eq!(results.get_value("sum").unwrap(), 60.0);
431        assert_eq!(results.get_value("avg").unwrap(), 20.0);
432        assert_eq!(results.get_value("min").unwrap(), 10.0);
433        assert_eq!(results.get_value("max").unwrap(), 30.0);
434    }
435
436    #[test]
437    fn test_composite_percentiles() {
438        let mut composite = CompositeAggregator::new();
439        composite
440            .add_p50("p50")
441            .add_p95("p95")
442            .add_p99("p99");
443
444        for i in 1..=100 {
445            composite.update(i as f64).unwrap();
446        }
447
448        let results = composite.finalize().unwrap();
449
450        assert!((results.get_value("p50").unwrap() - 50.5).abs() < 0.1);
451        assert!((results.get_value("p95").unwrap() - 95.05).abs() < 0.1);
452        assert!((results.get_value("p99").unwrap() - 99.01).abs() < 0.1);
453    }
454
455    #[test]
456    fn test_composite_stddev() {
457        let mut composite = CompositeAggregator::new();
458        composite
459            .add_avg("avg")
460            .add_stddev("stddev");
461
462        composite.update(10.0).unwrap();
463        composite.update(12.0).unwrap();
464        composite.update(14.0).unwrap();
465        composite.update(16.0).unwrap();
466        composite.update(18.0).unwrap();
467
468        let results = composite.finalize().unwrap();
469
470        assert!((results.get_value("avg").unwrap() - 14.0).abs() < 0.01);
471        assert!((results.get_value("stddev").unwrap() - 3.162).abs() < 0.01);
472    }
473
474    #[test]
475    fn test_composite_batch() {
476        let mut composite = CompositeAggregator::new();
477        composite
478            .add_count("count")
479            .add_avg("avg");
480
481        composite.update_batch(&[10.0, 20.0, 30.0, 40.0, 50.0]).unwrap();
482
483        let results = composite.finalize().unwrap();
484
485        assert_eq!(results.get_count("count").unwrap(), 5);
486        assert_eq!(results.get_value("avg").unwrap(), 30.0);
487    }
488
489    #[test]
490    fn test_composite_reset() {
491        let mut composite = CompositeAggregator::new();
492        composite.add_count("count").add_sum("sum");
493
494        composite.update(10.0).unwrap();
495        composite.update(20.0).unwrap();
496
497        composite.reset();
498
499        composite.update(5.0).unwrap();
500
501        let results = composite.finalize().unwrap();
502
503        assert_eq!(results.get_count("count").unwrap(), 1);
504        assert_eq!(results.get_value("sum").unwrap(), 5.0);
505    }
506
507    #[test]
508    fn test_composite_empty() {
509        let composite = CompositeAggregator::new();
510        assert!(composite.is_empty());
511        assert_eq!(composite.len(), 0);
512    }
513
514    #[test]
515    fn test_composite_all_aggregations() {
516        let mut composite = CompositeAggregator::new();
517        composite
518            .add_count("count")
519            .add_sum("sum")
520            .add_avg("avg")
521            .add_min("min")
522            .add_max("max")
523            .add_p50("p50")
524            .add_p95("p95")
525            .add_stddev("stddev");
526
527        for i in 1..=100 {
528            composite.update(i as f64).unwrap();
529        }
530
531        let results = composite.finalize().unwrap();
532
533        assert_eq!(results.get_count("count").unwrap(), 100);
534        assert_eq!(results.get_value("sum").unwrap(), 5050.0);
535        assert_eq!(results.get_value("avg").unwrap(), 50.5);
536        assert_eq!(results.get_value("min").unwrap(), 1.0);
537        assert_eq!(results.get_value("max").unwrap(), 100.0);
538        assert!((results.get_value("p50").unwrap() - 50.5).abs() < 0.1);
539        assert!((results.get_value("p95").unwrap() - 95.05).abs() < 0.1);
540        assert!((results.get_value("stddev").unwrap() - 29.01).abs() < 0.5);
541    }
542
543    #[test]
544    fn test_composite_serialization() {
545        let mut composite = CompositeAggregator::new();
546        composite.add_count("count").add_avg("avg");
547
548        composite.update(10.0).unwrap();
549        composite.update(20.0).unwrap();
550
551        let results = composite.finalize().unwrap();
552        let serialized = serde_json::to_string(&results).unwrap();
553        let deserialized: AggregationResult = serde_json::from_str(&serialized).unwrap();
554
555        assert_eq!(deserialized.get_count("count").unwrap(), 2);
556        assert_eq!(deserialized.get_value("avg").unwrap(), 15.0);
557    }
558
559    #[test]
560    fn test_composite_custom_percentile() {
561        let mut composite = CompositeAggregator::new();
562        composite.add_percentile("p75", 75.0).unwrap();
563
564        for i in 1..=100 {
565            composite.update(i as f64).unwrap();
566        }
567
568        let results = composite.finalize().unwrap();
569        assert!((results.get_value("p75").unwrap() - 75.25).abs() < 0.5);
570    }
571
572    #[test]
573    fn test_composite_invalid_percentile() {
574        let mut composite = CompositeAggregator::new();
575        assert!(composite.add_percentile("invalid", 150.0).is_err());
576    }
577
578    #[test]
579    fn test_real_world_latency_tracking() {
580        // Simulate tracking request latencies
581        let mut composite = CompositeAggregator::new();
582        composite
583            .add_count("request_count")
584            .add_avg("avg_latency_ms")
585            .add_min("min_latency_ms")
586            .add_max("max_latency_ms")
587            .add_p50("p50_latency_ms")
588            .add_p95("p95_latency_ms")
589            .add_p99("p99_latency_ms")
590            .add_stddev("stddev_latency_ms");
591
592        // Simulate some request latencies
593        let latencies = vec![
594            45.0, 52.0, 48.0, 51.0, 49.0, // Normal requests
595            150.0, 200.0, // Slower requests
596            46.0, 50.0, 47.0, // More normal
597            500.0, // One outlier
598        ];
599
600        for latency in latencies {
601            composite.update(latency).unwrap();
602        }
603
604        let results = composite.finalize().unwrap();
605
606        println!("Latency Statistics:");
607        println!("  Total Requests: {}", results.get_count("request_count").unwrap());
608        println!("  Avg: {:.2}ms", results.get_value("avg_latency_ms").unwrap());
609        println!("  Min: {:.2}ms", results.get_value("min_latency_ms").unwrap());
610        println!("  Max: {:.2}ms", results.get_value("max_latency_ms").unwrap());
611        println!("  P50: {:.2}ms", results.get_value("p50_latency_ms").unwrap());
612        println!("  P95: {:.2}ms", results.get_value("p95_latency_ms").unwrap());
613        println!("  P99: {:.2}ms", results.get_value("p99_latency_ms").unwrap());
614        println!("  StdDev: {:.2}ms", results.get_value("stddev_latency_ms").unwrap());
615
616        assert_eq!(results.get_count("request_count").unwrap(), 11);
617        assert!(results.get_value("avg_latency_ms").unwrap() > 50.0);
618        assert_eq!(results.get_value("min_latency_ms").unwrap(), 45.0);
619        assert_eq!(results.get_value("max_latency_ms").unwrap(), 500.0);
620    }
621}
llm_optimizer_processor/aggregation/composite.rs

llm_optimizer_processor/aggregation/
composite.rs
