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

diskann_benchmark_runner/utils/
percentiles.rs
