## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/functions-aggregate-common/src/tdigest.rs

    Line
    Count
    Source
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    //! An implementation of the [TDigest sketch algorithm] providing approximate
    19
    //! quantile calculations.
    20
    //!
    21
    //! The TDigest code in this module is modified from
    22
    //! <https://github.com/MnO2/t-digest>, itself a rust reimplementation of
    23
    //! [Facebook's Folly TDigest] implementation.
    24
    //!
    25
    //! Alterations include reduction of runtime heap allocations, broader type
    26
    //! support, (de-)serialization support, reduced type conversions and null value
    27
    //! tolerance.
    28
    //!
    29
    //! [TDigest sketch algorithm]: https://arxiv.org/abs/1902.04023
    30
    //! [Facebook's Folly TDigest]: https://github.com/facebook/folly/blob/main/folly/stats/TDigest.h
    31
    32
    use arrow::datatypes::DataType;
    33
    use arrow::datatypes::Float64Type;
    34
    use datafusion_common::cast::as_primitive_array;
    35
    use datafusion_common::Result;
    36
    use datafusion_common::ScalarValue;
    37
    use std::cmp::Ordering;
    38
    use std::mem::{size_of, size_of_val};
    39
    40
    pub const DEFAULT_MAX_SIZE: usize = 100;
    41
    42
    // Cast a non-null [`ScalarValue::Float64`] to an [`f64`], or
    43
    // panic.
    44
    macro_rules! cast_scalar_f64 {
    45
        ($value:expr ) => {
    46
            match &$value {
    47
                ScalarValue::Float64(Some(v)) => *v,
    48
                v => panic!("invalid type {}", v),
    49
            }
    50
        };
    51
    }
    52
    53
    // Cast a non-null [`ScalarValue::UInt64`] to an [`u64`], or
    54
    // panic.
    55
    macro_rules! cast_scalar_u64 {
    56
        ($value:expr ) => {
    57
            match &$value {
    58
                ScalarValue::UInt64(Some(v)) => *v,
    59
                v => panic!("invalid type {}", v),
    60
            }
    61
        };
    62
    }
    63
    64
    /// This trait is implemented for each type a [`TDigest`] can operate on,
    65
    /// allowing it to support both numerical rust types (obtained from
    66
    /// `PrimitiveArray` instances), and [`ScalarValue`] instances.
    67
    pub trait TryIntoF64 {
    68
        /// A fallible conversion of a possibly null `self` into a [`f64`].
    69
        ///
    70
        /// If `self` is null, this method must return `Ok(None)`.
    71
        ///
    72
        /// If `self` cannot be coerced to the desired type, this method must return
    73
        /// an `Err` variant.
    74
        fn try_as_f64(&self) -> Result<Option<f64>>;
    75
    }
    76
    77
    /// Generate an infallible conversion from `type` to an [`f64`].
    78
    macro_rules! impl_try_ordered_f64 {
    79
        ($type:ty) => {
    80
            impl TryIntoF64 for $type {
    81
    0
                fn try_as_f64(&self) -> Result<Option<f64>> {
    82
    0
                    Ok(Some(*self as f64))
    83
    0
                }
    84
            }
    85
        };
    86
    }
    87
    88
    impl_try_ordered_f64!(f64);
    89
    impl_try_ordered_f64!(f32);
    90
    impl_try_ordered_f64!(i64);
    91
    impl_try_ordered_f64!(i32);
    92
    impl_try_ordered_f64!(i16);
    93
    impl_try_ordered_f64!(i8);
    94
    impl_try_ordered_f64!(u64);
    95
    impl_try_ordered_f64!(u32);
    96
    impl_try_ordered_f64!(u16);
    97
    impl_try_ordered_f64!(u8);
    98
    99
    /// Centroid implementation to the cluster mentioned in the paper.
    100
    #[derive(Debug, PartialEq, Clone)]
    101
    pub struct Centroid {
    102
        mean: f64,
    103
        weight: f64,
    104
    }
    105
    106
    impl Centroid {
    107
    0
        pub fn new(mean: f64, weight: f64) -> Self {
    108
    0
            Centroid { mean, weight }
    109
    0
        }
    110
    111
        #[inline]
    112
    0
        pub fn mean(&self) -> f64 {
    113
    0
            self.mean
    114
    0
        }
    115
    116
        #[inline]
    117
    0
        pub fn weight(&self) -> f64 {
    118
    0
            self.weight
    119
    0
        }
    120
    121
    0
        pub fn add(&mut self, sum: f64, weight: f64) -> f64 {
    122
    0
            let new_sum = sum + self.weight * self.mean;
    123
    0
            let new_weight = self.weight + weight;
    124
    0
            self.weight = new_weight;
    125
    0
            self.mean = new_sum / new_weight;
    126
    0
            new_sum
    127
    0
        }
    128
    129
    0
        pub fn cmp_mean(&self, other: &Self) -> Ordering {
    130
    0
            self.mean.total_cmp(&other.mean)
    131
    0
        }
    132
    }
    133
    134
    impl Default for Centroid {
    135
    0
        fn default() -> Self {
    136
    0
            Centroid {
    137
    0
                mean: 0_f64,
    138
    0
                weight: 1_f64,
    139
    0
            }
    140
    0
        }
    141
    }
    142
    143
    /// T-Digest to be operated on.
    144
    #[derive(Debug, PartialEq, Clone)]
    145
    pub struct TDigest {
    146
        centroids: Vec<Centroid>,
    147
        max_size: usize,
    148
        sum: f64,
    149
        count: u64,
    150
        max: f64,
    151
        min: f64,
    152
    }
    153
    154
    impl TDigest {
    155
    0
        pub fn new(max_size: usize) -> Self {
    156
    0
            TDigest {
    157
    0
                centroids: Vec::new(),
    158
    0
                max_size,
    159
    0
                sum: 0_f64,
    160
    0
                count: 0,
    161
    0
                max: f64::NAN,
    162
    0
                min: f64::NAN,
    163
    0
            }
    164
    0
        }
    165
    166
        #[expect(clippy::needless_pass_by_value)]
    167
    0
        pub fn new_with_centroid(max_size: usize, centroid: Centroid) -> Self {
    168
    0
            TDigest {
    169
    0
                centroids: vec![centroid.clone()],
    170
    0
                max_size,
    171
    0
                sum: centroid.mean * centroid.weight,
    172
    0
                count: 1,
    173
    0
                max: centroid.mean,
    174
    0
                min: centroid.mean,
    175
    0
            }
    176
    0
        }
    177
    178
        #[inline]
    179
    0
        pub fn count(&self) -> u64 {
    180
    0
            self.count
    181
    0
        }
    182
    183
        #[inline]
    184
    0
        pub fn max(&self) -> f64 {
    185
    0
            self.max
    186
    0
        }
    187
    188
        #[inline]
    189
    0
        pub fn min(&self) -> f64 {
    190
    0
            self.min
    191
    0
        }
    192
    193
        #[inline]
    194
    0
        pub fn max_size(&self) -> usize {
    195
    0
            self.max_size
    196
    0
        }
    197
    198
        /// Size in bytes including `Self`.
    199
    0
        pub fn size(&self) -> usize {
    200
    0
            size_of_val(self) + (size_of::<Centroid>() * self.centroids.capacity())
    201
    0
        }
    202
    }
    203
    204
    impl Default for TDigest {
    205
    0
        fn default() -> Self {
    206
    0
            TDigest {
    207
    0
                centroids: Vec::new(),
    208
    0
                max_size: 100,
    209
    0
                sum: 0_f64,
    210
    0
                count: 0,
    211
    0
                max: f64::NAN,
    212
    0
                min: f64::NAN,
    213
    0
            }
    214
    0
        }
    215
    }
    216
    217
    impl TDigest {
    218
    0
        fn k_to_q(k: u64, d: usize) -> f64 {
    219
    0
            let k_div_d = k as f64 / d as f64;
    220
    0
            if k_div_d >= 0.5 {
    221
    0
                let base = 1.0 - k_div_d;
    222
    0
                1.0 - 2.0 * base * base
    223
            } else {
    224
    0
                2.0 * k_div_d * k_div_d
    225
            }
    226
    0
        }
    227
    228
    0
        fn clamp(v: f64, lo: f64, hi: f64) -> f64 {
    229
    0
            if lo.is_nan() || hi.is_nan() {
    230
    0
                return v;
    231
    0
            }
    232
    0
            v.clamp(lo, hi)
    233
    0
        }
    234
    235
        // public for testing in other modules
    236
    0
        pub fn merge_unsorted_f64(&self, unsorted_values: Vec<f64>) -> TDigest {
    237
    0
            let mut values = unsorted_values;
    238
    0
            values.sort_by(|a, b| a.total_cmp(b));
    239
    0
            self.merge_sorted_f64(&values)
    240
    0
        }
    241
    242
    0
        pub fn merge_sorted_f64(&self, sorted_values: &[f64]) -> TDigest {
    243
            #[cfg(debug_assertions)]
    244
    0
            debug_assert!(is_sorted(sorted_values), "unsorted input to TDigest");
    245
    246
    0
            if sorted_values.is_empty() {
    247
    0
                return self.clone();
    248
    0
            }
    249
    250
    0
            let mut result = TDigest::new(self.max_size());
    251
    0
            result.count = self.count() + sorted_values.len() as u64;
    252
    253
    0
            let maybe_min = *sorted_values.first().unwrap();
    254
    0
            let maybe_max = *sorted_values.last().unwrap();
    255
    256
    0
            if self.count() > 0 {
    257
    0
                result.min = self.min.min(maybe_min);
    258
    0
                result.max = self.max.max(maybe_max);
    259
    0
            } else {
    260
    0
                result.min = maybe_min;
    261
    0
                result.max = maybe_max;
    262
    0
            }
    263
    264
    0
            let mut compressed: Vec<Centroid> = Vec::with_capacity(self.max_size);
    265
    266
    0
            let mut k_limit: u64 = 1;
    267
    0
            let mut q_limit_times_count =
    268
    0
                Self::k_to_q(k_limit, self.max_size) * result.count() as f64;
    269
    0
            k_limit += 1;
    270
    271
    0
            let mut iter_centroids = self.centroids.iter().peekable();
    272
    0
            let mut iter_sorted_values = sorted_values.iter().peekable();
    273
    274
    0
            let mut curr: Centroid = if let Some(c) = iter_centroids.peek() {
    275
    0
                let curr = **iter_sorted_values.peek().unwrap();
    276
    0
                if c.mean() < curr {
    277
    0
                    iter_centroids.next().unwrap().clone()
    278
                } else {
    279
    0
                    Centroid::new(*iter_sorted_values.next().unwrap(), 1.0)
    280
                }
    281
            } else {
    282
    0
                Centroid::new(*iter_sorted_values.next().unwrap(), 1.0)
    283
            };
    284
    285
    0
            let mut weight_so_far = curr.weight();
    286
    287
    0
            let mut sums_to_merge = 0_f64;
    288
    0
            let mut weights_to_merge = 0_f64;
    289
    290
    0
            while iter_centroids.peek().is_some() || iter_sorted_values.peek().is_some() {
    291
    0
                let next: Centroid = if let Some(c) = iter_centroids.peek() {
    292
    0
                    if iter_sorted_values.peek().is_none()
    293
    0
                        || c.mean() < **iter_sorted_values.peek().unwrap()
    294
                    {
    295
    0
                        iter_centroids.next().unwrap().clone()
    296
                    } else {
    297
    0
                        Centroid::new(*iter_sorted_values.next().unwrap(), 1.0)
    298
                    }
    299
                } else {
    300
    0
                    Centroid::new(*iter_sorted_values.next().unwrap(), 1.0)
    301
                };
    302
    303
    0
                let next_sum = next.mean() * next.weight();
    304
    0
                weight_so_far += next.weight();
    305
    306
    0
                if weight_so_far <= q_limit_times_count {
    307
    0
                    sums_to_merge += next_sum;
    308
    0
                    weights_to_merge += next.weight();
    309
    0
                } else {
    310
    0
                    result.sum += curr.add(sums_to_merge, weights_to_merge);
    311
    0
                    sums_to_merge = 0_f64;
    312
    0
                    weights_to_merge = 0_f64;
    313
    0
    314
    0
                    compressed.push(curr.clone());
    315
    0
                    q_limit_times_count =
    316
    0
                        Self::k_to_q(k_limit, self.max_size) * result.count() as f64;
    317
    0
                    k_limit += 1;
    318
    0
                    curr = next;
    319
    0
                }
    320
            }
    321
    322
    0
            result.sum += curr.add(sums_to_merge, weights_to_merge);
    323
    0
            compressed.push(curr);
    324
    0
            compressed.shrink_to_fit();
    325
    0
            compressed.sort_by(|a, b| a.cmp_mean(b));
    326
    327
    0
            result.centroids = compressed;
    328
    0
            result
    329
    0
        }
    330
    331
    0
        fn external_merge(
    332
    0
            centroids: &mut [Centroid],
    333
    0
            first: usize,
    334
    0
            middle: usize,
    335
    0
            last: usize,
    336
    0
        ) {
    337
    0
            let mut result: Vec<Centroid> = Vec::with_capacity(centroids.len());
    338
    339
    0
            let mut i = first;
    340
    0
            let mut j = middle;
    341
    342
    0
            while i < middle && j < last {
    343
    0
                match centroids[i].cmp_mean(&centroids[j]) {
    344
    0
                    Ordering::Less => {
    345
    0
                        result.push(centroids[i].clone());
    346
    0
                        i += 1;
    347
    0
                    }
    348
    0
                    Ordering::Greater => {
    349
    0
                        result.push(centroids[j].clone());
    350
    0
                        j += 1;
    351
    0
                    }
    352
    0
                    Ordering::Equal => {
    353
    0
                        result.push(centroids[i].clone());
    354
    0
                        i += 1;
    355
    0
                    }
    356
                }
    357
            }
    358
    359
    0
            while i < middle {
    360
    0
                result.push(centroids[i].clone());
    361
    0
                i += 1;
    362
    0
            }
    363
    364
    0
            while j < last {
    365
    0
                result.push(centroids[j].clone());
    366
    0
                j += 1;
    367
    0
            }
    368
    369
    0
            i = first;
    370
    0
            for centroid in result.into_iter() {
    371
    0
                centroids[i] = centroid;
    372
    0
                i += 1;
    373
    0
            }
    374
    0
        }
    375
    376
        // Merge multiple T-Digests
    377
    0
        pub fn merge_digests<'a>(digests: impl IntoIterator<Item = &'a TDigest>) -> TDigest {
    378
    0
            let digests = digests.into_iter().collect::<Vec<_>>();
    379
    0
            let n_centroids: usize = digests.iter().map(|d| d.centroids.len()).sum();
    380
    0
            if n_centroids == 0 {
    381
    0
                return TDigest::default();
    382
    0
            }
    383
    384
    0
            let max_size = digests.first().unwrap().max_size;
    385
    0
            let mut centroids: Vec<Centroid> = Vec::with_capacity(n_centroids);
    386
    0
            let mut starts: Vec<usize> = Vec::with_capacity(digests.len());
    387
    388
    0
            let mut count = 0;
    389
    0
            let mut min = f64::INFINITY;
    390
    0
            let mut max = f64::NEG_INFINITY;
    391
    392
    0
            let mut start: usize = 0;
    393
    0
            for digest in digests.iter() {
    394
    0
                starts.push(start);
    395
    396
    0
                let curr_count = digest.count();
    397
    0
                if curr_count > 0 {
    398
    0
                    min = min.min(digest.min);
    399
    0
                    max = max.max(digest.max);
    400
    0
                    count += curr_count;
    401
    0
                    for centroid in &digest.centroids {
    402
    0
                        centroids.push(centroid.clone());
    403
    0
                        start += 1;
    404
    0
                    }
    405
    0
                }
    406
            }
    407
    408
    0
            let mut digests_per_block: usize = 1;
    409
    0
            while digests_per_block < starts.len() {
    410
    0
                for i in (0..starts.len()).step_by(digests_per_block * 2) {
    411
    0
                    if i + digests_per_block < starts.len() {
    412
    0
                        let first = starts[i];
    413
    0
                        let middle = starts[i + digests_per_block];
    414
    0
                        let last = if i + 2 * digests_per_block < starts.len() {
    415
    0
                            starts[i + 2 * digests_per_block]
    416
                        } else {
    417
    0
                            centroids.len()
    418
                        };
    419
    420
    0
                        debug_assert!(first <= middle && middle <= last);
    421
    0
                        Self::external_merge(&mut centroids, first, middle, last);
    422
    0
                    }
    423
                }
    424
    425
    0
                digests_per_block *= 2;
    426
            }
    427
    428
    0
            let mut result = TDigest::new(max_size);
    429
    0
            let mut compressed: Vec<Centroid> = Vec::with_capacity(max_size);
    430
    431
    0
            let mut k_limit = 1;
    432
    0
            let mut q_limit_times_count = Self::k_to_q(k_limit, max_size) * count as f64;
    433
    434
    0
            let mut iter_centroids = centroids.iter_mut();
    435
    0
            let mut curr = iter_centroids.next().unwrap();
    436
    0
            let mut weight_so_far = curr.weight();
    437
    0
            let mut sums_to_merge = 0_f64;
    438
    0
            let mut weights_to_merge = 0_f64;
    439
    440
    0
            for centroid in iter_centroids {
    441
    0
                weight_so_far += centroid.weight();
    442
    443
    0
                if weight_so_far <= q_limit_times_count {
    444
    0
                    sums_to_merge += centroid.mean() * centroid.weight();
    445
    0
                    weights_to_merge += centroid.weight();
    446
    0
                } else {
    447
    0
                    result.sum += curr.add(sums_to_merge, weights_to_merge);
    448
    0
                    sums_to_merge = 0_f64;
    449
    0
                    weights_to_merge = 0_f64;
    450
    0
                    compressed.push(curr.clone());
    451
    0
                    q_limit_times_count = Self::k_to_q(k_limit, max_size) * count as f64;
    452
    0
                    k_limit += 1;
    453
    0
                    curr = centroid;
    454
    0
                }
    455
            }
    456
    457
    0
            result.sum += curr.add(sums_to_merge, weights_to_merge);
    458
    0
            compressed.push(curr.clone());
    459
    0
            compressed.shrink_to_fit();
    460
    0
            compressed.sort_by(|a, b| a.cmp_mean(b));
    461
    462
    0
            result.count = count;
    463
    0
            result.min = min;
    464
    0
            result.max = max;
    465
    0
            result.centroids = compressed;
    466
    0
            result
    467
    0
        }
    468
    469
        /// To estimate the value located at `q` quantile
    470
    0
        pub fn estimate_quantile(&self, q: f64) -> f64 {
    471
    0
            if self.centroids.is_empty() {
    472
    0
                return 0.0;
    473
    0
            }
    474
    475
    0
            let rank = q * self.count as f64;
    476
    477
            let mut pos: usize;
    478
            let mut t;
    479
    0
            if q > 0.5 {
    480
    0
                if q >= 1.0 {
    481
    0
                    return self.max();
    482
    0
                }
    483
    484
    0
                pos = 0;
    485
    0
                t = self.count as f64;
    486
    487
    0
                for (k, centroid) in self.centroids.iter().enumerate().rev() {
    488
    0
                    t -= centroid.weight();
    489
    490
    0
                    if rank >= t {
    491
    0
                        pos = k;
    492
    0
                        break;
    493
    0
                    }
    494
                }
    495
            } else {
    496
    0
                if q <= 0.0 {
    497
    0
                    return self.min();
    498
    0
                }
    499
    500
    0
                pos = self.centroids.len() - 1;
    501
    0
                t = 0_f64;
    502
    503
    0
                for (k, centroid) in self.centroids.iter().enumerate() {
    504
    0
                    if rank < t + centroid.weight() {
    505
    0
                        pos = k;
    506
    0
                        break;
    507
    0
                    }
    508
    509
    0
                    t += centroid.weight();
    510
                }
    511
            }
    512
    513
    0
            let mut delta = 0_f64;
    514
    0
            let mut min = self.min;
    515
    0
            let mut max = self.max;
    516
    517
    0
            if self.centroids.len() > 1 {
    518
    0
                if pos == 0 {
    519
    0
                    delta = self.centroids[pos + 1].mean() - self.centroids[pos].mean();
    520
    0
                    max = self.centroids[pos + 1].mean();
    521
    0
                } else if pos == (self.centroids.len() - 1) {
    522
    0
                    delta = self.centroids[pos].mean() - self.centroids[pos - 1].mean();
    523
    0
                    min = self.centroids[pos - 1].mean();
    524
    0
                } else {
    525
    0
                    delta = (self.centroids[pos + 1].mean() - self.centroids[pos - 1].mean())
    526
    0
                        / 2.0;
    527
    0
                    min = self.centroids[pos - 1].mean();
    528
    0
                    max = self.centroids[pos + 1].mean();
    529
    0
                }
    530
    0
            }
    531
    532
    0
            let value = self.centroids[pos].mean()
    533
    0
                + ((rank - t) / self.centroids[pos].weight() - 0.5) * delta;
    534
    535
            // In `merge_digests()`: `min` is initialized to Inf, `max` is initialized to -Inf
    536
            // and gets updated according to different `TDigest`s
    537
            // However, `min`/`max` won't get updated if there is only one `NaN` within `TDigest`
    538
            // The following two checks is for such edge case
    539
    0
            if !min.is_finite() && min.is_sign_positive() {
    540
    0
                min = f64::NEG_INFINITY;
    541
    0
            }
    542
    543
    0
            if !max.is_finite() && max.is_sign_negative() {
    544
    0
                max = f64::INFINITY;
    545
    0
            }
    546
    547
    0
            Self::clamp(value, min, max)
    548
    0
        }
    549
    550
        /// This method decomposes the [`TDigest`] and its [`Centroid`] instances
    551
        /// into a series of primitive scalar values.
    552
        ///
    553
        /// First the values of the TDigest are packed, followed by the variable
    554
        /// number of centroids packed into a [`ScalarValue::List`] of
    555
        /// [`ScalarValue::Float64`]:
    556
        ///
    557
        /// ```text
    558
        ///
    559
        ///    ┌────────┬────────┬────────┬───────┬────────┬────────┐
    560
        ///    │max_size│  sum   │ count  │  max  │  min   │centroid│
    561
        ///    └────────┴────────┴────────┴───────┴────────┴────────┘
    562
        ///                                                     │
    563
        ///                               ┌─────────────────────┘
    564
        ///                               ▼
    565
        ///                          ┌ List ───┐
    566
        ///                          │┌ ─ ─ ─ ┐│
    567
        ///                          │  mean   │
    568
        ///                          │├ ─ ─ ─ ┼│─ ─ Centroid 1
    569
        ///                          │ weight  │
    570
        ///                          │└ ─ ─ ─ ┘│
    571
        ///                          │         │
    572
        ///                          │┌ ─ ─ ─ ┐│
    573
        ///                          │  mean   │
    574
        ///                          │├ ─ ─ ─ ┼│─ ─ Centroid 2
    575
        ///                          │ weight  │
    576
        ///                          │└ ─ ─ ─ ┘│
    577
        ///                          │         │
    578
        ///                              ...
    579
        /// ```
    580
        ///
    581
        /// The [`TDigest::from_scalar_state()`] method reverses this processes,
    582
        /// consuming the output of this method and returning an unpacked
    583
        /// [`TDigest`].
    584
    0
        pub fn to_scalar_state(&self) -> Vec<ScalarValue> {
    585
            // Gather up all the centroids
    586
    0
            let centroids: Vec<ScalarValue> = self
    587
    0
                .centroids
    588
    0
                .iter()
    589
    0
                .flat_map(|c| [c.mean(), c.weight()])
    590
    0
                .map(|v| ScalarValue::Float64(Some(v)))
    591
    0
                .collect();
    592
    593
    0
            let arr = ScalarValue::new_list_nullable(&centroids, &DataType::Float64);
    594
    595
    0
            vec![
    596
    0
                ScalarValue::UInt64(Some(self.max_size as u64)),
    597
    0
                ScalarValue::Float64(Some(self.sum)),
    598
    0
                ScalarValue::UInt64(Some(self.count)),
    599
    0
                ScalarValue::Float64(Some(self.max)),
    600
    0
                ScalarValue::Float64(Some(self.min)),
    601
    0
                ScalarValue::List(arr),
    602
            ]
    603
    0
        }
    604
    605
        /// Unpack the serialized state of a [`TDigest`] produced by
    606
        /// [`Self::to_scalar_state()`].
    607
        ///
    608
        /// # Correctness
    609
        ///
    610
        /// Providing input to this method that was not obtained from
    611
        /// [`Self::to_scalar_state()`] results in undefined behaviour and may
    612
        /// panic.
    613
    0
        pub fn from_scalar_state(state: &[ScalarValue]) -> Self {
    614
    0
            assert_eq!(state.len(), 6, "invalid TDigest state");
    615
    616
    0
            let max_size = match &state[0] {
    617
    0
                ScalarValue::UInt64(Some(v)) => *v as usize,
    618
    0
                v => panic!("invalid max_size type {v:?}"),
    619
            };
    620
    621
    0
            let centroids: Vec<_> = match &state[5] {
    622
    0
                ScalarValue::List(arr) => {
    623
    0
                    let array = arr.values();
    624
    625
    0
                    let f64arr =
    626
    0
                        as_primitive_array::<Float64Type>(array).expect("expected f64 array");
    627
    0
                    f64arr
    628
    0
                        .values()
    629
    0
                        .chunks(2)
    630
    0
                        .map(|v| Centroid::new(v[0], v[1]))
    631
    0
                        .collect()
    632
                }
    633
    0
                v => panic!("invalid centroids type {v:?}"),
    634
            };
    635
    636
    0
            let max = cast_scalar_f64!(&state[3]);
    637
    0
            let min = cast_scalar_f64!(&state[4]);
    638
    639
    0
            if min.is_finite() && max.is_finite() {
    640
    0
                assert!(max.total_cmp(&min).is_ge());
    641
    0
            }
    642
    643
            Self {
    644
    0
                max_size,
    645
    0
                sum: cast_scalar_f64!(state[1]),
    646
    0
                count: cast_scalar_u64!(&state[2]),
    647
    0
                max,
    648
    0
                min,
    649
    0
                centroids,
    650
            }
    651
    0
        }
    652
    }
    653
    654
    #[cfg(debug_assertions)]
    655
    0
    fn is_sorted(values: &[f64]) -> bool {
    656
    0
        values.windows(2).all(|w| w[0].total_cmp(&w[1]).is_le())
    657
    0
    }
    658
    659
    #[cfg(test)]
    660
    mod tests {
    661
        use super::*;
    662
    663
        // A macro to assert the specified `quantile` estimated by `t` is within the
    664
        // allowable relative error bound.
    665
        macro_rules! assert_error_bounds {
    666
            ($t:ident, quantile = $quantile:literal, want = $want:literal) => {
    667
                assert_error_bounds!(
    668
                    $t,
    669
                    quantile = $quantile,
    670
                    want = $want,
    671
                    allowable_error = 0.01
    672
                )
    673
            };
    674
            ($t:ident, quantile = $quantile:literal, want = $want:literal, allowable_error = $re:literal) => {
    675
                let ans = $t.estimate_quantile($quantile);
    676
                let expected: f64 = $want;
    677
                let percentage: f64 = (expected - ans).abs() / expected;
    678
                assert!(
    679
                    percentage < $re,
    680
                    "relative error {} is more than {}% (got quantile {}, want {})",
    681
                    percentage,
    682
                    $re,
    683
                    ans,
    684
                    expected
    685
                );
    686
            };
    687
        }
    688
    689
        macro_rules! assert_state_roundtrip {
    690
            ($t:ident) => {
    691
                let state = $t.to_scalar_state();
    692
                let other = TDigest::from_scalar_state(&state);
    693
                assert_eq!($t, other);
    694
            };
    695
        }
    696
    697
        #[test]
    698
        fn test_int64_uniform() {
    699
            let values = (1i64..=1000).map(|v| v as f64).collect();
    700
    701
            let t = TDigest::new(100);
    702
            let t = t.merge_unsorted_f64(values);
    703
    704
            assert_error_bounds!(t, quantile = 0.1, want = 100.0);
    705
            assert_error_bounds!(t, quantile = 0.5, want = 500.0);
    706
            assert_error_bounds!(t, quantile = 0.9, want = 900.0);
    707
            assert_state_roundtrip!(t);
    708
        }
    709
    710
        #[test]
    711
        fn test_centroid_addition_regression() {
    712
            // https://github.com/MnO2/t-digest/pull/1
    713
    714
            let vals = vec![1.0, 1.0, 1.0, 2.0, 1.0, 1.0];
    715
            let mut t = TDigest::new(10);
    716
    717
            for v in vals {
    718
                t = t.merge_unsorted_f64(vec![v]);
    719
            }
    720
    721
            assert_error_bounds!(t, quantile = 0.5, want = 1.0);
    722
            assert_error_bounds!(t, quantile = 0.95, want = 2.0);
    723
            assert_state_roundtrip!(t);
    724
        }
    725
    726
        #[test]
    727
        fn test_merge_unsorted_against_uniform_distro() {
    728
            let t = TDigest::new(100);
    729
            let values: Vec<_> = (1..=1_000_000).map(f64::from).collect();
    730
    731
            let t = t.merge_unsorted_f64(values);
    732
    733
            assert_error_bounds!(t, quantile = 1.0, want = 1_000_000.0);
    734
            assert_error_bounds!(t, quantile = 0.99, want = 990_000.0);
    735
            assert_error_bounds!(t, quantile = 0.01, want = 10_000.0);
    736
            assert_error_bounds!(t, quantile = 0.0, want = 1.0);
    737
            assert_error_bounds!(t, quantile = 0.5, want = 500_000.0);
    738
            assert_state_roundtrip!(t);
    739
        }
    740
    741
        #[test]
    742
        fn test_merge_unsorted_against_skewed_distro() {
    743
            let t = TDigest::new(100);
    744
            let mut values: Vec<_> = (1..=600_000).map(f64::from).collect();
    745
            values.resize(1_000_000, 1_000_000_f64);
    746
    747
            let t = t.merge_unsorted_f64(values);
    748
    749
            assert_error_bounds!(t, quantile = 0.99, want = 1_000_000.0);
    750
            assert_error_bounds!(t, quantile = 0.01, want = 10_000.0);
    751
            assert_error_bounds!(t, quantile = 0.5, want = 500_000.0);
    752
            assert_state_roundtrip!(t);
    753
        }
    754
    755
        #[test]
    756
        fn test_merge_digests() {
    757
            let mut digests: Vec<TDigest> = Vec::new();
    758
    759
            for _ in 1..=100 {
    760
                let t = TDigest::new(100);
    761
                let values: Vec<_> = (1..=1_000).map(f64::from).collect();
    762
                let t = t.merge_unsorted_f64(values);
    763
                digests.push(t)
    764
            }
    765
    766
            let t = TDigest::merge_digests(&digests);
    767
    768
            assert_error_bounds!(t, quantile = 1.0, want = 1000.0);
    769
            assert_error_bounds!(t, quantile = 0.99, want = 990.0);
    770
            assert_error_bounds!(t, quantile = 0.01, want = 10.0, allowable_error = 0.2);
    771
            assert_error_bounds!(t, quantile = 0.0, want = 1.0);
    772
            assert_error_bounds!(t, quantile = 0.5, want = 500.0);
    773
            assert_state_roundtrip!(t);
    774
        }
    775
    776
        #[test]
    777
        fn test_size() {
    778
            let t = TDigest::new(10);
    779
            let t = t.merge_unsorted_f64(vec![0.0, 1.0]);
    780
    781
            assert_eq!(t.size(), 96);
    782
        }
    783
    }