perf(tdigest): simplify rank anchor lookup

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.

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::<Vec<_>>();
+    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();

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<Centroid>, 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,

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();