Skip to main content

radiate_core/stats/
metric.rs

1use crate::{ProjectExpr, stats::{MetricView, Tag, TagType, defaults, metric_fields}};
2use radiate_error::{RadiateError, radiate_err};
3use radiate_expr::SelectOp;
4use radiate_utils::{
5    AnyValue, DataType, SmallStr, Statistic
6};
7#[cfg(feature = "serde")]
8use serde::{Deserialize, Serialize};
9use std::{hash::Hash, time::Duration};
10
11const DTYPE_NULL: u8 = 0;
12const DTYPE_FLOAT32: u8 = 1;
13const DTYPE_DURATION: u8 = 2;
14const DTYPE_LIST: u8 = 3;
15
16#[macro_export]
17macro_rules! metric {
18    ($name:expr, $update:expr) => {{
19        let mut metric = $crate::Metric::new($name);
20        metric.apply_update($update);
21        metric
22    }};
23    ($name:expr) => {{ $crate::Metric::new($name).upsert(1) }};
24}
25
26
27#[derive(Clone, PartialEq, Default)]
28#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
29pub(super) struct Meta {
30    pub(super) update_count: usize,
31    pub(super) generation: usize,
32}
33
34#[derive(Clone, PartialEq, Default)]
35#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
36pub struct Metric {
37    name: SmallStr,
38    inner: Statistic,
39    samples: Option<Vec<f32>>,
40    meta: Meta,
41    tags: Tag,
42    dtype: u8,
43}
44
45impl Metric {
46    pub fn new(name: impl Into<SmallStr>) -> Self {
47        let name = name.into();
48        let tags = defaults::default_tags(&name);
49
50        Self {
51            name,
52            inner: Statistic::default(),
53            meta: Meta::default(),
54            samples: None,
55            tags,
56            dtype: DTYPE_NULL,
57        }
58    }
59
60    pub fn is_empty(&self) -> bool {
61        self.meta.update_count == 0 && self.inner.count() == 0
62    }
63
64    #[inline(always)]
65    pub fn update_count(&self) -> usize {
66        self.meta.update_count
67    }
68
69    #[inline(always)]
70    pub fn generation(&self) -> usize {
71        self.meta.generation
72    }
73
74    #[inline(always)]
75    pub fn set_generation(&mut self, generation: usize) {
76        if generation != self.meta.generation {
77            self.meta.update_count = 0;
78        }
79
80        self.meta.generation = generation;
81    }
82
83    pub fn dtype(&self) -> DataType {
84        match self.dtype {
85            DTYPE_NULL => DataType::Null,
86            DTYPE_FLOAT32 => DataType::Float32,
87            DTYPE_DURATION => DataType::Duration,
88            DTYPE_LIST => DataType::List(Box::new(DataType::Float32)),
89            _ => DataType::Null,
90        }
91    }
92
93    #[inline(always)]
94    pub fn tags(&self) -> Tag {
95        self.tags
96    }
97
98    #[inline(always)]
99    pub fn add_tag(&mut self, tag: TagType) {
100        self.tags.insert(tag);
101    }
102
103    pub fn iter_tags(&self) -> impl Iterator<Item = TagType> {
104        self.tags.iter()
105    }
106
107    pub fn clear_values(&mut self) {
108        self.inner = Statistic::default();
109        self.samples = None;
110    }
111
112    pub fn stats<'a>(&'a self) -> Option<MetricView<'a, f32>> {
113        if !self.tags.has(TagType::Statistic) {
114            return None;
115        }
116
117        Some(MetricView {
118            name: &self.name,
119            statistic: &self.inner,
120            samples: self.samples.as_deref(),
121            mapper: |v| v,
122        })
123    }
124
125    pub fn times<'a>(&'a self) -> Option<MetricView<'a, Duration>> {
126        if !self.tags.has(TagType::Time) {
127            return None;
128        }
129
130        Some(MetricView {
131            name: &self.name,
132            statistic: &self.inner,
133            samples: self.samples.as_deref(),
134            mapper: |v| Duration::from_secs_f32(v),
135        })
136    }
137
138    pub fn distributions<'a>(&'a self) -> Option<MetricView<'a, f32>> {
139        if !self.tags.has(TagType::Distribution) {
140            return None;
141        }
142
143        Some(MetricView {
144            name: &self.name,
145            statistic: &self.inner,
146            samples: self.samples.as_deref(),
147            mapper: |v| v,
148         })
149    }
150
151    #[inline(always)]
152    pub fn upsert<'a>(mut self, update: impl Into<MetricUpdate<'a>>) -> Self {
153        self.apply_update(update);
154        self
155    }
156
157    #[inline(always)]
158    pub fn update_from(&mut self, other: Metric) {
159        // Kinda a hack to take advantage of the fact that if count == sum,
160        // we can just apply the sum directly instead of merging statistics - keeps things honest
161        // & avoids merging statistics when we don't have to (even though that's a fast operation).
162        if other.count() as f32 == other.sum() && !other.tags.has(TagType::Distribution) {
163            self.apply_update(other.sum());
164        } else {
165            self.apply_update(other.inner);
166        }
167
168        self.tags = self.tags.union(other.tags);
169    }
170
171    #[inline(always)]
172    pub fn apply_update<'a>(&mut self, update: impl Into<MetricUpdate<'a>>) {
173        let update = update.into();
174        match update {
175            MetricUpdate::Float(value) => {
176                self.update_statistic(value);
177            }
178            MetricUpdate::Usize(value) => {
179                self.update_statistic(value as f32);
180            }
181            MetricUpdate::Duration(value) => {
182                self.update_time_statistic(value);
183            }
184            MetricUpdate::UsizeDistribution(values) => {
185                self.update_statistic_from_iter(values.iter().map(|v| *v as f32));
186            }
187            MetricUpdate::Distribution(values) => {
188                self.update_statistic_from_iter(values.iter().cloned());
189            }
190            MetricUpdate::OwnedDistribution(values) => {
191                self.update_statistic_from_iter(values);
192            }
193            MetricUpdate::Statistic(stat) => {
194                self.inner.merge(&stat);
195                self.dtype = DTYPE_FLOAT32;
196                self.meta.update_count += 1;
197            }
198            MetricUpdate::Bool(value) => {
199                self.update_statistic(if value { 1.0 } else { 0.0 });
200            }
201
202        }
203    }
204
205    fn update_statistic(&mut self, value: f32) {
206        self.inner.add(value);
207        self.add_tag(TagType::Statistic);
208
209        self.meta.update_count += 1;
210
211        if self.dtype == DTYPE_NULL {
212            self.dtype = DTYPE_FLOAT32;
213        }
214    }
215
216    fn update_time_statistic(&mut self, value: Duration) {
217        self.inner.add(value.as_secs_f32());
218        self.add_tag(TagType::Time);
219        self.meta.update_count += 1;
220
221        if self.dtype == DTYPE_NULL {
222            self.dtype = DTYPE_DURATION;
223        }
224    }
225
226    fn update_statistic_from_iter<I>(&mut self, values: I)
227    where
228        I: IntoIterator<Item = f32>,
229    {   
230        let samples = self.samples.get_or_insert_with(Vec::new);
231
232        samples.clear();
233        self.inner.clear();
234
235        for val in values {
236            samples.push(val);
237            self.inner.add(val);
238        }
239        
240        self.meta.update_count += self.inner.count() as usize;
241        
242        self.add_tag(TagType::Distribution);
243
244        if self.dtype == DTYPE_NULL {
245            self.dtype = DTYPE_LIST;
246        }
247    }
248
249    pub fn clear_samples(&mut self) {
250        self.samples = None;
251    }
252        
253    pub fn statistic(&self) -> &Statistic {
254        &self.inner
255    }
256
257    pub fn name(&self) -> &SmallStr {
258        &self.name
259    }
260
261    pub fn last_value(&self) -> f32 {
262        self.inner.last_value()
263    }
264
265    pub fn count(&self) -> u32 {
266        self.inner.count()
267    }
268
269    pub fn mean(&self) -> f32 {
270        self.inner.mean()
271    }
272
273    pub fn var(&self) -> f32 {
274        self.inner.variance().unwrap_or(0.0)
275    }
276
277    pub fn stddev(&self) -> f32 {
278        self.inner.std_dev().unwrap_or(0.0)
279    }
280
281    pub fn skew(&self) -> f32 {
282        self.inner.skewness().unwrap_or(0.0)
283    }
284
285    pub fn kurt(&self) -> f32 {
286        self.inner.kurtosis().unwrap_or(0.0)
287    }
288    
289    pub fn min(&self) -> f32 {
290        self.inner.min()
291    }
292
293    pub fn max(&self) -> f32 {
294        self.inner.max()
295    }
296
297    pub fn sum(&self) -> f32 {
298        self.inner.sum()
299    }
300
301    pub fn quantile(&self, q: f32) -> Option<f32> {
302        self.distributions().and_then(|view| view.quantile(q))
303    }
304}
305
306impl<'a> ProjectExpr<'a> for &Metric {
307    #[inline]
308    fn select(&'a self, sel: &SelectOp) -> Result<AnyValue<'a>, RadiateError> {
309        (*self).select(sel)
310    }
311}
312
313impl<'a> ProjectExpr<'a> for Metric {
314    #[inline]
315    fn select(&self, sel: &SelectOp) -> Result<AnyValue<'a>, RadiateError> {
316        let wrap = |v: f32| match self.dtype {
317            DTYPE_FLOAT32 | DTYPE_LIST => AnyValue::Float32(v),
318            DTYPE_DURATION => AnyValue::Duration(Duration::from_secs_f32(v)),
319            _ => AnyValue::Null,
320        };
321
322        let match_field = |metric: &Metric, field: &SmallStr| {
323            match field.as_str() {
324                f if f == metric_fields::LAST_VALUE => wrap(metric.last_value()),
325                f if f == metric_fields::MEAN => wrap(metric.mean()),
326                f if f == metric_fields::STDDEV => wrap(metric.stddev()),
327                f if f == metric_fields::MIN => wrap(metric.min()),
328                f if f == metric_fields::MAX => wrap(metric.max()),
329                f if f == metric_fields::SUM => wrap(metric.sum()),
330                f if f == metric_fields::VARIANCE => wrap(metric.var()),
331                f if f == metric_fields::SKEWNESS => wrap(metric.skew()),
332                f if f == metric_fields::KURTOSIS => wrap(metric.kurt()),
333                f if f == metric_fields::COUNT => AnyValue::UInt64(metric.count() as u64),
334                f if f == metric_fields::GENERATION => AnyValue::UInt64(metric.generation() as u64),
335                f if f == metric_fields::UPDATE_COUNT => AnyValue::UInt64(metric.update_count() as u64),
336                _ => AnyValue::Null,
337            }
338        };
339
340        match sel {
341            SelectOp::Field(field) => {
342                Ok(match_field(self, field))
343            }
344            SelectOp::Identity => {
345                Ok(AnyValue::from(self))
346            }
347            _ => Ok(AnyValue::Null),
348        }
349    }
350}
351
352impl From<&Metric> for AnyValue<'_> {
353    fn from(metric: &Metric) -> Self {
354        use AnyValue::*;
355
356        AnyValue::Struct(metric.name().clone(), Vec::from([
357            (metric_fields::LAST_VALUE, DataType::Float32, Float32(metric.last_value())),
358            (metric_fields::MEAN, DataType::Float32, Float32(metric.mean())),
359            (metric_fields::STDDEV, DataType::Float32, Float32(metric.stddev())),
360            (metric_fields::MIN, DataType::Float32, Float32(metric.min())),
361            (metric_fields::MAX, DataType::Float32, Float32(metric.max())),
362            (metric_fields::SUM, DataType::Float32, Float32(metric.sum())),
363            (metric_fields::VARIANCE, DataType::Float32, Float32(metric.var())),
364            (metric_fields::SKEWNESS, DataType::Float32, Float32(metric.skew())),
365            (metric_fields::KURTOSIS, DataType::Float32, Float32(metric.kurt())),
366            (metric_fields::COUNT, DataType::UInt64, UInt64(metric.count() as u64)),
367            (metric_fields::GENERATION, DataType::UInt64, UInt64(metric.generation() as u64)),
368            (metric_fields::UPDATE_COUNT, DataType::UInt64, UInt64(metric.update_count() as u64)),
369        ]))
370    }
371}
372
373impl Hash for Metric {
374    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
375        self.name.hash(state);
376        self.inner.hash(state);
377        self.tags.hash(state);
378    }
379}
380
381#[derive( PartialEq, Debug)]
382pub enum MetricUpdate<'a> {
383    Float(f32),
384    Usize(usize),
385    Duration(Duration),
386    Distribution(&'a [f32]),
387    OwnedDistribution(Vec<f32>),
388    UsizeDistribution(&'a [usize]),
389    Statistic(Statistic),
390    Bool(bool),
391}
392
393impl From<f32> for MetricUpdate<'_> {
394    fn from(value: f32) -> Self {
395        MetricUpdate::Float(value)
396    }
397}
398
399impl From<usize> for MetricUpdate<'_> {
400    fn from(value: usize) -> Self {
401        MetricUpdate::Usize(value)
402    }
403}
404
405impl From<Duration> for MetricUpdate<'_> {
406    fn from(value: Duration) -> Self {
407        MetricUpdate::Duration(value)
408    }
409}
410
411impl<'a> From<&'a [f32]> for MetricUpdate<'a> {
412    fn from(value: &'a [f32]) -> Self {
413        MetricUpdate::Distribution(value)
414    }
415}
416
417impl<'a> From<&'a Vec<f32>> for MetricUpdate<'a> {
418    fn from(value: &'a Vec<f32>) -> Self {
419        MetricUpdate::Distribution(value)
420    }
421}
422
423impl<'a> From<&'a Vec<usize>> for MetricUpdate<'a> {
424    fn from(value: &'a Vec<usize>) -> Self {
425        MetricUpdate::UsizeDistribution(value)
426    }
427}
428
429impl From<Statistic> for MetricUpdate<'_> {
430    fn from(value: Statistic) -> Self {
431        MetricUpdate::Statistic(value)
432    }
433}
434
435impl From<bool> for MetricUpdate<'_> {
436    fn from(value: bool) -> Self {
437        MetricUpdate::Bool(value)
438    }
439}
440
441impl<'a> TryFrom<AnyValue<'a>> for MetricUpdate<'a> {
442    type Error = RadiateError;
443
444    fn try_from(value: AnyValue<'a>) -> Result<Self, Self::Error> {
445        match value {
446            AnyValue::UInt8(v) => Ok(MetricUpdate::Float(v as f32)),
447            AnyValue::UInt16(v) => Ok(MetricUpdate::Float(v as f32)),
448            AnyValue::UInt32(v) => Ok(MetricUpdate::Float(v as f32)),
449            AnyValue::UInt64(v) => Ok(MetricUpdate::Float(v as f32)),
450            AnyValue::UInt128(v) => Ok(MetricUpdate::Float(v as f32)),
451
452            AnyValue::Int8(v) => Ok(MetricUpdate::Float(v as f32)),
453            AnyValue::Int16(v) => Ok(MetricUpdate::Float(v as f32)),
454            AnyValue::Int32(v) => Ok(MetricUpdate::Float(v as f32)),
455            AnyValue::Int64(v) => Ok(MetricUpdate::Float(v as f32)),
456            AnyValue::Int128(v) => Ok(MetricUpdate::Float(v as f32)),
457
458            AnyValue::Float32(v) => Ok(MetricUpdate::Float(v)),
459            AnyValue::Float64(v) => Ok(MetricUpdate::Float(v as f32)),
460
461            AnyValue::Duration(v) => Ok(MetricUpdate::Duration(v)),
462
463            AnyValue::Slice(values) => {
464                let out = values
465                    .iter()
466                    .enumerate()
467                    .map(|(index, v)| {
468                        v.clone().extract::<f32>().ok_or(
469                            radiate_err!(
470                                Metric: 
471                                "cannot convert AnyValue sequence into Vec<f32>: element at index {index} has non-numeric type `{}`", v.type_name()))
472                            
473                    })
474                    .collect::<Result<Vec<f32>, _>>()?;
475
476                Ok(MetricUpdate::OwnedDistribution(out))
477            }
478
479            AnyValue::Vector(values) => {
480                let out = values
481                    .into_iter()
482                    .enumerate()
483                    .map(|(index, v)| {
484                        let ty = v.type_name();
485                        v.extract::<f32>()
486                            .ok_or(radiate_err!(
487                                Metric: 
488                                "cannot convert AnyValue sequence into Vec<f32>: element at index {index} has non-numeric type `{ty}`"
489                            ))
490                    })
491                    .collect::<Result<Vec<f32>, _>>()?;
492
493                Ok(MetricUpdate::OwnedDistribution(out))
494            }
495
496            AnyValue::Bool(bool) => Ok(MetricUpdate::Bool(bool)),
497
498            other => Err(radiate_err!(Metric: "cannot convert AnyValue of type `{}` into MetricUpdate", other.type_name())),
499        }
500    }
501}
502
503impl std::fmt::Debug for Metric {
504    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
505        write!(f, "Metric {{ name: {}, }}", self.name)
506    }
507}
508
509
510#[cfg(test)]
511mod tests {
512    use super::*;
513
514    const EPSILON: f32 = 1e-5;
515
516    fn approx_eq(a: f32, b: f32, eps: f32) -> bool {
517        (a - b).abs() <= eps
518    }
519
520    fn assert_stat_eq(m: &Metric, count: u32, mean: f32, var: f32, min: f32, max: f32) {
521        assert_eq!(m.count(), count);
522        assert!(approx_eq(m.mean(), mean, EPSILON), "mean");
523        assert!(approx_eq(m.var(), var, EPSILON), "var");
524        assert!(approx_eq(m.min(), min, EPSILON), "min");
525        assert!(approx_eq(m.max(), max, EPSILON), "max");
526    }
527
528    fn stats_of(values: &[f32]) -> (u32, f32, f32, f32, f32) {
529        // sample variance (n-1), matches your Statistic::variance
530        let n = values.len() as u32;
531        if n == 0 {
532            return (0, 0.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY);
533        }
534        let mean = values.iter().sum::<f32>() / values.len() as f32;
535
536        let mut m2 = 0.0_f32;
537        for &v in values {
538            let d = v - mean;
539            m2 += d * d;
540        }
541
542        let var = if n == 1 { 0.0 } else { m2 / (n as f32 - 1.0) };
543
544        let min = values.iter().cloned().fold(f32::INFINITY, f32::min);
545        let max = values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
546
547        (n, mean, var, min, max)
548    }
549
550    #[test]
551    fn test_metric() {
552        let mut metric = Metric::new("test");
553        metric.apply_update(1.0);
554        metric.apply_update(2.0);
555        metric.apply_update(3.0);
556        metric.apply_update(4.0);
557        metric.apply_update(5.0);
558
559        assert_eq!(metric.count(), 5);
560        assert_eq!(metric.last_value(), 5.0);
561        assert_eq!(metric.mean(), 3.0);
562        assert_eq!(metric.var(), 2.5);
563        assert_eq!(metric.stddev(), 1.5811388);
564        assert_eq!(metric.min(), 1.0);
565        assert_eq!(metric.max(), 5.0);
566        assert_eq!(metric.name(), "test");
567    }
568
569    #[test]
570    fn test_metric_labels() {
571        let mut metric = Metric::new("test");
572
573        metric.apply_update(1.0);
574        metric.apply_update(2.0);
575        metric.apply_update(3.0);
576        metric.apply_update(4.0);
577        metric.apply_update(5.0);
578
579        assert_eq!(metric.count(), 5);
580        assert_eq!(metric.last_value(), 5.0);
581        assert_eq!(metric.mean(), 3.0);
582        assert_eq!(metric.var(), 2.5);
583        assert_eq!(metric.stddev(), 1.5811388);
584        assert_eq!(metric.min(), 1.0);
585        assert_eq!(metric.max(), 5.0);
586    }
587
588    #[test]
589    fn distribution_tag_is_applied_on_any_slice_update() {
590        let mut m = Metric::new("scores");
591
592        // seed with scalar samples first (creates Statistic but not Distribution tag)
593        m.apply_update(1.0);
594        m.apply_update(2.0);
595        assert!(m.tags().has(TagType::Statistic));
596        assert!(!m.tags().has(TagType::Distribution));
597
598        // now apply a slice update - we expect Distribution tag to appear
599        m.apply_update(&[3.0, 4.0][..]);
600
601        assert!(
602            m.tags().has(TagType::Distribution),
603            "expected Distribution tag after slice update"
604        );
605    }
606
607    #[test]
608    fn metric_merge_matches_streaming_samples() {
609        let a = [1.0, 2.0, 3.0, 4.0];
610        let b = [10.0, 20.0, 30.0];
611
612        let mut m1 = Metric::new("x");
613        m1.apply_update(&a[..]);
614
615        let mut m2 = Metric::new("x");
616        m2.apply_update(&b[..]);
617
618        m1.update_from(m2);
619
620        let combined = [1.0, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0];
621        let (n, mean, var, min, max) = stats_of(&combined);
622        assert_stat_eq(&m1, n, mean, var, min, max);
623    }
624}