diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d4b1de..4f366e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,7 +19,7 @@ All significant changes to this project will be documented in this file. ### Improvements * Serializing buffered `TDigestMut` inputs is faster when their total weight is at most `k / 2`. -* `TDigest` and `TDigestMut` `cdf` and `pmf` queries are faster when answering many split points. +* `TDigest` and `TDigestMut` `rank` queries are faster for values between centroid means; `cdf` and `pmf` queries are faster when answering many split points. * Improve the readability of `Debug` output for HLL and CPC sketches and unions. * `KllSketch::quantile` and `ReqSketch::quantile` are faster at ranks `0.0` and `1.0`, with query work independent of the number of retained samples. * `BloomFilter::insert` is faster for integer and raw-byte inputs. `BloomFilter::contains_and_insert` is also faster when checking already-present integer values. diff --git a/benchmarks/tdigest/query.rs b/benchmarks/tdigest/query.rs index ae36977..338179e 100644 --- a/benchmarks/tdigest/query.rs +++ b/benchmarks/tdigest/query.rs @@ -22,6 +22,7 @@ use rand::SeedableRng; use rand::rngs::StdRng; use rand::seq::SliceRandom; +use super::support::build_digest; use super::support::prepared_digest; #[divan::bench] @@ -31,6 +32,14 @@ fn rank(bencher: Bencher) { bencher.bench_local(|| black_box(&digest).rank(black_box(0.531_25))); } +#[divan::bench(args = [0.0, 3.0, 3.5, 7.0])] +fn rank_repeated_means(bencher: Bencher, value: f64) { + let values = (0..100_000).map(|i| (i % 8) as f64).collect::>(); + let digest = build_digest(&values); + + bencher.bench_local(|| black_box(&digest).rank(black_box(value))); +} + #[divan::bench] fn quantile(bencher: Bencher) { let digest = prepared_digest(); diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 4b6a27d..810c86b 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -16,7 +16,6 @@ // under the License. use std::cmp::Ordering; -use std::convert::identity; use std::num::NonZeroU64; use crate::codec::SketchBytes; @@ -1708,22 +1707,15 @@ impl TDigestView<'_> { }; } - let mut lower = self - .centroids - .binary_search_by(|c| centroid_lower_bound(c, value)) - .unwrap_or_else(identity); - assert_ne!(lower, num_centroids, "get_rank: lower == end"); - let mut upper = self - .centroids - .binary_search_by(|c| centroid_upper_bound(c, value)) - .unwrap_or_else(identity); - assert_ne!(upper, 0, "get_rank: upper == begin"); - if value < self.centroids[lower].mean { - lower -= 1; - } - if (upper == num_centroids) || (self.centroids[upper - 1].mean >= value) { - upper -= 1; - } + // The tail checks place value within the centroid means. Between means, the first + // mean >= value and its predecessor bracket it; at a mean, use the entire equal run. + let index = self.centroids.partition_point(|c| c.mean < value); + let (lower, upper) = if self.centroids[index].mean == value { + let end = self.centroids.partition_point(|c| c.mean <= value); + (index, end - 1) + } else { + (index - 1, index) + }; let weight_below = weight_before(lower); let weight_between: u64 = self.centroids[lower..upper] @@ -2001,22 +1993,6 @@ fn merge_sorted_centroids(left: &mut Vec, right: &[Centroid]) { } } -fn centroid_lower_bound(c: &Centroid, value: f64) -> Ordering { - if c.mean < value { - Ordering::Less - } else { - Ordering::Greater - } -} - -fn centroid_upper_bound(c: &Centroid, value: f64) -> Ordering { - if c.mean > value { - Ordering::Greater - } else { - Ordering::Less - } -} - #[derive(Debug, Clone, Copy, PartialEq)] struct Centroid { mean: f64, diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 428d96c..b5bb67e 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -244,6 +244,49 @@ fn test_repeated_blocks() { assert_eq!(tdigest.rank(3.01), Some(1.0)); } +#[test] +fn test_rank_brackets_weighted_equal_mean_runs() { + let mut digest = deserialize_with_centroids( + 100, + 0., + 20., + &[ + (0., 2), + (0., 6), + (10., 4), + (10., 8), + (10., 12), + (20., 14), + (20., 18), + ], + ); + // The centroid centers are 1, 5, 10, 16, 26, 39, and 55. At an equal-mean run, + // rank averages the first and last centers; between runs it uses the adjacent centers. + let points = [-1., 0., 5., 10., 15., 20., 21.]; + let expected = [0., 3., 7.5, 18., 32.5, 47., 64.].map(|mass| mass / 64.); + let frozen = digest.clone().freeze(); + for (&point, &rank) in points.iter().zip(&expected) { + assert_eq!(digest.rank(point), Some(rank)); + assert_eq!(frozen.rank(point), Some(rank)); + } + let mut cdf = expected.to_vec(); + cdf.push(1.); + assert_eq!(digest.cdf(&points), Some(cdf.clone())); + assert_eq!(frozen.cdf(&points), Some(cdf)); + let pmf = vec![ + 0., + 3. / 64., + 4.5 / 64., + 10.5 / 64., + 14.5 / 64., + 14.5 / 64., + 17. / 64., + 0., + ]; + assert_eq!(digest.pmf(&points), Some(pmf.clone())); + assert_eq!(frozen.pmf(&points), Some(pmf)); +} + #[test] fn test_merge_small() { let mut td1 = TDigestMut::new(10).unwrap();