llm_optimizer_processor/aggregation/
count.rs

1use super::trait_::Aggregator;
2
3use serde::{Deserialize, Serialize};
4
5/// Accumulator for count aggregation
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
7pub struct CountAccumulator {
8    count: u64,
9}
10
11/// Count aggregator - counts the number of values
12///
13/// This is the simplest aggregator, it just counts how many values have been processed.
14/// Useful for tracking event counts, request counts, etc.
15///
16/// # Examples
17///
18/// ```
19/// use processor::aggregation::{Aggregator, CountAggregator};
20///
21/// let mut agg = CountAggregator::new();
22/// agg.update(1.0).unwrap();
23/// agg.update(2.0).unwrap();
24/// agg.update(3.0).unwrap();
25///
26/// assert_eq!(agg.finalize().unwrap(), 3);
27/// ```
28#[derive(Debug, Clone)]
29pub struct CountAggregator<T> {
30    count: u64,
31    _phantom: std::marker::PhantomData<T>,
32}
33
34impl<T: Clone> CountAggregator<T> {
35    /// Create a new count aggregator
36    pub fn new() -> Self {
37        Self {
38            count: 0,
39            _phantom: std::marker::PhantomData,
40        }
41    }
42}
43
44impl<T: Clone> Default for CountAggregator<T> {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl<T: Clone + Send + Sync + std::fmt::Debug> Aggregator for CountAggregator<T> {
51    type Input = T;
52    type Output = u64;
53    type Accumulator = CountAccumulator;
54
55    fn new() -> Self {
56        Self::new()
57    }
58
59    fn update(&mut self, _value: <Self as Aggregator>::Input) -> anyhow::Result<()> {
60        self.count += 1;
61        Ok(())
62    }
63
64    fn finalize(&self) -> anyhow::Result<<Self as Aggregator>::Output> {
65        Ok(self.count)
66    }
67
68    fn accumulator(&self) -> <Self as Aggregator>::Accumulator {
69        CountAccumulator {
70            count: self.count,
71        }
72    }
73
74    fn merge(&mut self, other: <Self as Aggregator>::Accumulator) -> anyhow::Result<()> {
75        self.count += other.count;
76        Ok(())
77    }
78
79    fn reset(&mut self) {
80        self.count = 0;
81    }
82
83    fn count(&self) -> u64 {
84        self.count
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_count_basic() {
94        let mut agg = CountAggregator::<f64>::new();
95        assert_eq!(agg.count(), 0);
96        assert!(agg.is_empty());
97
98        agg.update(1.0).unwrap();
99        agg.update(2.0).unwrap();
100        agg.update(3.0).unwrap();
101
102        assert_eq!(agg.count(), 3);
103        assert!(!agg.is_empty());
104        assert_eq!(agg.finalize().unwrap(), 3);
105    }
106
107    #[test]
108    fn test_count_batch() {
109        let mut agg = CountAggregator::<f64>::new();
110        agg.update_batch(&[1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
111
112        assert_eq!(agg.count(), 5);
113        assert_eq!(agg.finalize().unwrap(), 5);
114    }
115
116    #[test]
117    fn test_count_merge() {
118        let mut agg1 = CountAggregator::<i32>::new();
119        agg1.update(1).unwrap();
120        agg1.update(2).unwrap();
121
122        let mut agg2 = CountAggregator::<i32>::new();
123        agg2.update(3).unwrap();
124        agg2.update(4).unwrap();
125        agg2.update(5).unwrap();
126
127        let acc2 = agg2.accumulator();
128        agg1.merge(acc2).unwrap();
129
130        assert_eq!(agg1.finalize().unwrap(), 5);
131    }
132
133    #[test]
134    fn test_count_reset() {
135        let mut agg = CountAggregator::<f64>::new();
136        agg.update(1.0).unwrap();
137        agg.update(2.0).unwrap();
138
139        assert_eq!(agg.count(), 2);
140
141        agg.reset();
142        assert_eq!(agg.count(), 0);
143        assert!(agg.is_empty());
144    }
145
146    #[test]
147    fn test_count_serialization() {
148        let mut agg = CountAggregator::<f64>::new();
149        agg.update(1.0).unwrap();
150        agg.update(2.0).unwrap();
151
152        let acc = agg.accumulator();
153        let serialized = serde_json::to_string(&acc).unwrap();
154        let deserialized: CountAccumulator = serde_json::from_str(&serialized).unwrap();
155
156        assert_eq!(acc, deserialized);
157        assert_eq!(deserialized.count, 2);
158    }
159
160    #[test]
161    fn test_count_with_different_types() {
162        // Test with strings
163        let mut agg = CountAggregator::<String>::new();
164        agg.update("hello".to_string()).unwrap();
165        agg.update("world".to_string()).unwrap();
166        assert_eq!(agg.finalize().unwrap(), 2);
167
168        // Test with tuples
169        let mut agg = CountAggregator::<(i32, String)>::new();
170        agg.update((1, "a".to_string())).unwrap();
171        agg.update((2, "b".to_string())).unwrap();
172        assert_eq!(agg.finalize().unwrap(), 2);
173    }
174}
llm_optimizer_processor/aggregation/count.rs

llm_optimizer_processor/aggregation/
count.rs
