migrate tests for counter & tdigest

Author: xx01cyx

optd-cost-model/Cargo.lock

@@ -16,3 +16,6 @@ chrono = "0.4"
 itertools = "0.13"
 lazy_static = "1.5"
 
+[dev-dependencies]
+crossbeam = "0.8"
+rand = "0.8"

optd-cost-model/Cargo.toml

@@ -69,3 +69,128 @@ where
         self.counts.contains_key(key)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+    use std::sync::{Arc, Mutex};
+
+    use crossbeam::thread;
+    use rand::rngs::StdRng;
+    use rand::seq::SliceRandom;
+    use rand::SeedableRng;
+
+    use super::Counter;
+
+    // Generates hardcoded frequencies and returns them,
+    // along with a flattened randomized array containing those frequencies.
+    fn generate_frequencies() -> (HashMap<i32, i32>, Vec<i32>) {
+        let mut frequencies = HashMap::new();
+
+        frequencies.insert(0, 2);
+        frequencies.insert(1, 4);
+        frequencies.insert(2, 9);
+        frequencies.insert(3, 8);
+        frequencies.insert(4, 50);
+        frequencies.insert(5, 6);
+
+        let mut flattened = Vec::new();
+        for (key, &value) in &frequencies {
+            for _ in 0..value {
+                flattened.push(*key);
+            }
+        }
+
+        let mut rng = StdRng::seed_from_u64(0);
+        flattened.shuffle(&mut rng);
+
+        (frequencies, flattened)
+    }
+
+    #[test]
+    fn aggregate() {
+        let to_track = vec![0, 1, 2, 3];
+        let mut mcv = Counter::<i32>::new(&to_track);
+
+        let (frequencies, flattened) = generate_frequencies();
+
+        mcv.aggregate(&flattened);
+
+        let mcv_freq = mcv.frequencies();
+        assert_eq!(mcv_freq.len(), to_track.len());
+
+        to_track.iter().for_each(|item| {
+            assert!(mcv_freq.contains_key(item));
+            assert_eq!(
+                mcv_freq.get(item),
+                frequencies
+                    .get(item)
+                    .map(|e| (*e as f64 / flattened.len() as f64))
+                    .as_ref()
+            );
+        });
+    }
+
+    #[test]
+    fn merge() {
+        let to_track = vec![0, 1, 2, 3];
+        let n_jobs = 16;
+
+        let total_frequencies = Arc::new(Mutex::new(HashMap::<i32, i32>::new()));
+        let total_count = Arc::new(Mutex::new(0));
+        let result_mcv = Arc::new(Mutex::new(Counter::<i32>::new(&to_track)));
+        thread::scope(|s| {
+            for _ in 0..n_jobs {
+                s.spawn(|_| {
+                    let mut local_mcv = Counter::<i32>::new(&to_track);
+
+                    let (local_frequencies, flattened) = generate_frequencies();
+                    let mut total_frequencies = total_frequencies.lock().unwrap();
+                    let mut total_count = total_count.lock().unwrap();
+                    for (&key, &value) in &local_frequencies {
+                        *total_frequencies.entry(key).or_insert(0) += value;
+                        *total_count += value;
+                    }
+
+                    local_mcv.aggregate(&flattened);
+
+                    let mcv_local_freq = local_mcv.frequencies();
+                    assert_eq!(mcv_local_freq.len(), to_track.len());
+
+                    to_track.iter().for_each(|item| {
+                        assert!(mcv_local_freq.contains_key(item));
+                        assert_eq!(
+                            mcv_local_freq.get(item),
+                            local_frequencies
+                                .get(item)
+                                .map(|e| (*e as f64 / flattened.len() as f64))
+                                .as_ref()
+                        );
+                    });
+
+                    let mut result = result_mcv.lock().unwrap();
+                    result.merge(&local_mcv);
+                });
+            }
+        })
+        .unwrap();
+
+        let mcv = result_mcv.lock().unwrap();
+        let total_count = total_count.lock().unwrap();
+        let mcv_freq = mcv.frequencies();
+
+        assert_eq!(*total_count, mcv.total_count);
+        to_track.iter().for_each(|item| {
+            assert!(mcv_freq.contains_key(item));
+            assert_eq!(
+                mcv_freq.get(item),
+                total_frequencies
+                    .lock()
+                    .unwrap()
+                    .get(item)
+                    .map(|e| (*e as f64 / *total_count as f64))
+                    .as_ref()
+            );
+        });
+    }
+}

optd-cost-model/src/stats/counter.rs

@@ -248,3 +248,148 @@ where
 fn lerp(a: f64, b: f64, f: f64) -> f64 {
     (a * (1.0 - f)) + (b * f)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::{Arc, Mutex};
+
+    use crossbeam::thread;
+    use ordered_float::OrderedFloat;
+    use rand::distributions::{Distribution, Uniform, WeightedIndex};
+    use rand::rngs::StdRng;
+    use rand::SeedableRng;
+
+    use super::{IntoFloat, TDigest};
+
+    impl IntoFloat for OrderedFloat<f64> {
+        fn to_float(&self) -> f64 {
+            self.0
+        }
+    }
+
+    // Whether obtained = expected +/- error
+    fn is_close(obtained: f64, expected: f64, error: f64) -> bool {
+        ((expected - error) < obtained) && (obtained < (expected + error))
+    }
+
+    // Checks whether the tdigest follows a uniform distribution.
+    fn check_tdigest_uniform(
+        tdigest: &TDigest<OrderedFloat<f64>>,
+        buckets: i32,
+        max: f64,
+        min: f64,
+        error: f64,
+    ) {
+        for k in 0..buckets {
+            let expected_cdf = (k as f64) / (buckets as f64);
+            let expected_quantile = (max - min) * expected_cdf + min;
+
+            let obtained_cdf = tdigest.cdf(&OrderedFloat(expected_quantile));
+            let obtained_quantile = tdigest.quantile(expected_cdf);
+
+            assert!(is_close(obtained_cdf, expected_cdf, error));
+            assert!(is_close(
+                obtained_quantile,
+                expected_quantile,
+                (max - min) * error,
+            ));
+        }
+    }
+
+    #[test]
+    fn uniform_merge_sequential() {
+        let buckets = 200;
+        let error = 0.03; // 3% absolute error on each quantile; error gets worse near the median.
+        let mut tdigest = TDigest::new(buckets as f64);
+
+        let (min, max) = (-1000.0, 1000.0);
+        let uniform_distr = Uniform::new(min, max);
+        let mut rng = StdRng::seed_from_u64(0);
+
+        let batch_size = 1024;
+        let batch_numbers = 64;
+
+        for _ in 0..batch_numbers {
+            let mut random_numbers = Vec::with_capacity(batch_size);
+            for _ in 0..batch_size {
+                let num: f64 = uniform_distr.sample(&mut rng);
+                random_numbers.push(OrderedFloat(num));
+            }
+            tdigest.merge_values(&random_numbers);
+        }
+
+        check_tdigest_uniform(&tdigest, buckets, max, min, error);
+    }
+
+    #[test]
+    fn uniform_merge_parallel() {
+        let buckets = 200;
+        let error = 0.03; // 3% absolute error on each quantile, note error is worse near the median.
+
+        let (min, max) = (-1000.0, 1000.0);
+
+        let batch_size = 65536;
+        let batch_numbers = 64;
+
+        let result_tdigest = Arc::new(Mutex::new(TDigest::new(buckets as f64)));
+        thread::scope(|s| {
+            for _ in 0..batch_numbers {
+                s.spawn(|_| {
+                    let mut local_tdigest = TDigest::new(buckets as f64);
+
+                    let mut random_numbers = Vec::with_capacity(batch_size);
+                    let uniform_distr = Uniform::new(min, max);
+                    let mut rng = StdRng::seed_from_u64(0);
+
+                    for _ in 0..batch_size {
+                        let num: f64 = uniform_distr.sample(&mut rng);
+                        random_numbers.push(OrderedFloat(num));
+                    }
+                    local_tdigest.merge_values(&random_numbers);
+
+                    let mut result = result_tdigest.lock().unwrap();
+                    result.merge(&local_tdigest);
+                });
+            }
+        })
+        .unwrap();
+
+        let tdigest = result_tdigest.lock().unwrap();
+        check_tdigest_uniform(&tdigest, buckets, max, min, error);
+    }
+
+    #[test]
+    fn weighted_merge() {
+        let buckets = 200;
+        let error = 0.05; // 5% absolute error on each quantile, note error is worse near the median.
+
+        let mut tdigest = TDigest::new(buckets as f64);
+
+        let choices = [9.0, 900.0, 990.0, 9990.0, 190000.0, 990000.0];
+        let weights = [1, 2, 1, 3, 4, 5]; // Total of 16.
+        let total_weight: i32 = weights.iter().sum();
+
+        let weighted_distr = WeightedIndex::new(weights).unwrap();
+        let mut rng = StdRng::seed_from_u64(0);
+
+        let batch_size = 128;
+        let batch_numbers = 16;
+
+        for _ in 0..batch_numbers {
+            let mut random_numbers = Vec::with_capacity(batch_size);
+            for _ in 0..batch_size {
+                let num: f64 = choices[weighted_distr.sample(&mut rng)];
+                random_numbers.push(OrderedFloat(num));
+            }
+            tdigest.merge_values(&random_numbers);
+        }
+
+        let mut curr_weight = 0;
+        for (c, w) in choices.iter().zip(weights) {
+            curr_weight += w;
+            let estimate_cdf = tdigest.cdf(&OrderedFloat(*c));
+            let obtained_cdf = (curr_weight as f64) / (total_weight as f64);
+            assert!(is_close(obtained_cdf, estimate_cdf, error));
+        }
+    }
+}