Skip to main content

diskann_benchmark_runner/utils/
percentiles.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use serde::{Deserialize, Serialize};
7use thiserror::Error;
8
9#[derive(Debug, Error)]
10#[error("input slice cannot be empty")]
11pub struct CannotBeEmpty;
12
13#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
14#[non_exhaustive]
15pub struct Percentiles<T> {
16    pub minimum: T,
17    pub mean: f64,
18    pub median: f64,
19    pub p90: T,
20    pub p99: T,
21}
22
23pub trait AsF64Lossy: Copy {
24    fn as_f64_lossy(self) -> f64;
25}
26
27macro_rules! impl_as_f64_lossy {
28    ($T:ty) => {
29        impl AsF64Lossy for $T {
30            fn as_f64_lossy(self) -> f64 {
31                self as f64
32            }
33        }
34    };
35    ($($T:ty),* $(,)?) => {
36        $(impl_as_f64_lossy!($T);)*
37    }
38}
39
40impl_as_f64_lossy!(i8, i16, i32, i64, u8, u16, u32, u64, usize, f32, f64);
41
42pub fn mean<T>(x: &[T]) -> Result<f64, CannotBeEmpty>
43where
44    T: AsF64Lossy + std::iter::Sum,
45{
46    if x.is_empty() {
47        return Err(CannotBeEmpty);
48    }
49
50    let s: T = x.iter().copied().sum();
51    Ok(s.as_f64_lossy() / x.len() as f64)
52}
53
54/// Find the maximum value of the sequence of `f64`.
55pub fn max_f64(x: &[f64]) -> Result<f64, CannotBeEmpty> {
56    x.iter().copied().reduce(f64::max).ok_or(CannotBeEmpty)
57}
58
59/// Return the mean, median, 90th and 99th percentile of the input vector.
60///
61/// NOTE: This is implemented by sorting the input slice.
62pub fn compute_percentiles<T>(x: &mut [T]) -> Result<Percentiles<T>, CannotBeEmpty>
63where
64    T: std::cmp::Ord + std::ops::Add<Output = T> + AsF64Lossy + std::iter::Sum,
65{
66    let mean = mean(x)?;
67
68    x.sort_unstable();
69
70    let len = x.len();
71    let half = len / 2;
72    let median = if len % 2 == 1 {
73        x[half].as_f64_lossy()
74    } else {
75        (x[half - 1].as_f64_lossy() + x[half].as_f64_lossy()) / 2.0
76    };
77
78    let p90 = x[((9 * len) / 10).min(len - 1)];
79    let p99 = x[((99 * len) / 100).min(len - 1)];
80
81    Ok(Percentiles {
82        minimum: x[0],
83        mean,
84        median,
85        p90,
86        p99,
87    })
88}
89
90///////////
91// Tests //
92///////////
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn test_mean() {
100        let empty: &[f32] = &[];
101        assert!(matches!(mean(empty).unwrap_err(), CannotBeEmpty));
102
103        let input = [
104            -2.049918,
105            0.12130953,
106            -0.17400686,
107            0.7511493,
108            0.26361275,
109            1.2661924,
110            1.023522,
111            -2.8727458,
112            -1.0132318,
113            0.531649,
114            -0.8730961,
115            1.0494779,
116            1.8957608,
117            0.45292637,
118            0.1296239,
119            0.06079646,
120            -1.3347862,
121            0.122092366,
122            -0.82615733,
123            1.3791777,
124            1.5189241,
125            -0.8614088,
126            -0.62131107,
127            -2.0626633,
128            -0.49564686,
129        ];
130
131        let r = mean(&input).unwrap();
132        assert!((r - -0.10475030175999998).abs() <= 3.0e-17);
133    }
134
135    #[test]
136    fn test_max() {
137        assert_eq!(max_f64(&[1.0, -1.0, f64::NEG_INFINITY]).unwrap(), 1.0);
138        assert!(matches!(max_f64(&[]).unwrap_err(), CannotBeEmpty));
139        assert_eq!(
140            max_f64(&[1.0, -1.0, f64::NAN, f64::NEG_INFINITY]).unwrap(),
141            1.0
142        );
143    }
144
145    #[test]
146    fn test_compute_percentils() {
147        // Size 0
148        {
149            let empty: &mut [u64] = &mut [];
150            assert!(matches!(
151                compute_percentiles(empty).unwrap_err(),
152                CannotBeEmpty
153            ));
154        }
155
156        // Size 1
157        {
158            let v: &mut [u64] = &mut [10];
159            let p = compute_percentiles(v).unwrap();
160            let e = Percentiles {
161                minimum: 10,
162                mean: 10.0,
163                median: 10.0,
164                p90: 10,
165                p99: 10,
166            };
167            assert_eq!(p, e);
168        }
169
170        // Size 2
171        {
172            let v: &mut [u64] = &mut [2, 1];
173            let p = compute_percentiles(v).unwrap();
174            let e = Percentiles {
175                minimum: 1,
176                mean: 1.5,
177                median: 1.5,
178                p90: 2,
179                p99: 2,
180            };
181            assert_eq!(p, e);
182        }
183
184        // Size 3
185        {
186            let v: &mut [u64] = &mut [2, 1, 3];
187            let p = compute_percentiles(v).unwrap();
188            let e = Percentiles {
189                minimum: 1,
190                mean: 2.0,
191                median: 2.0,
192                p90: 3,
193                p99: 3,
194            };
195            assert_eq!(p, e);
196        }
197
198        // Size 9
199        {
200            let v: &mut [u64] = &mut [2, 1, 3, 4, 9, 6, 7, 5, 8];
201            let p = compute_percentiles(v).unwrap();
202            let e = Percentiles {
203                minimum: 1,
204                mean: 5.0,
205                median: 5.0,
206                p90: 9,
207                p99: 9,
208            };
209            assert_eq!(p, e);
210        }
211
212        // Size 10
213        {
214            let v: &mut [u64] = &mut [2, 10, 1, 3, 4, 9, 6, 7, 5, 8];
215            let p = compute_percentiles(v).unwrap();
216            let e = Percentiles {
217                minimum: 1,
218                mean: 5.5,
219                median: 5.5,
220                p90: 10,
221                p99: 10,
222            };
223            assert_eq!(p, e);
224        }
225
226        // Size 10
227        {
228            let v: &mut [u64] = &mut [2, 10, 1, 3, 4, 9, 6, 11, 7, 5, 8];
229            let p = compute_percentiles(v).unwrap();
230            let e = Percentiles {
231                minimum: 1,
232                mean: 6.0,
233                median: 6.0,
234                p90: 10,
235                p99: 11,
236            };
237            assert_eq!(p, e);
238        }
239    }
240}