From f59af0be9339c2384e37065a6da6c17c977ecad0 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 4 Oct 2026 23:44:40 +0800 Subject: [PATCH 01/10] perf(tdigest): route unmergeable collected digests past the merge heap When the combined weight of collected inputs is at most k / 2, no two centroids fit under the K_2 limit, so the k-way heap merge only reproduces the stable-sorted concatenation one centroid at a time. Take the same sort path as mixed compressed/buffered inputs, which already skips the merge scan at these counts. --- datasketches/src/tdigest/sketch.rs | 6 ++- .../tests/tdigest_test/sketch.rs | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index ef39ffd..eafdcd6 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -1160,7 +1160,11 @@ impl FromIterator for TDigestMut { }); } - if !all_compressed { + // Below k / 2 combined weight no two centroids fit under the K_2 limit (the bound in + // compress_sorted_centroids), so the heap would only reproduce the stable-sorted + // concatenation one centroid at a time. That case takes the same sort path as mixed + // inputs, which already skips the merge scan. + if !all_compressed || total_weight <= u64::from(merged.k) / 2 { // Raw tails precede compressed prefixes to preserve stable equal-mean ordering. let mut centroids = Vec::with_capacity(num_centroids); for input in &heap { diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 428d96c..5716c6a 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -401,6 +401,44 @@ fn test_from_iter_matches_single_compression_for_interleaved_runs() { } } +#[test] +fn test_from_iter_concatenates_small_compressed_inputs() { + // A combined weight of at most k / 2 admits no merges, so collection concatenates the + // inputs in stable mean order through the compression fast path instead of the heap scan. + let runs: [&[(f64, u64)]; 3] = [ + &[(2.0, 1), (5.0, 2)], + &[(1.0, 1), (5.0, 1)], + &[(5.0, 1), (9.0, 1)], + ]; + let concatenated = [ + (1.0, 1), + (2.0, 1), + (5.0, 2), // run order breaks mean ties + (5.0, 1), + (5.0, 1), + (9.0, 1), + ]; + for reverse in [false, true] { + let make_digest = |centroids: &[(f64, u64)]| { + let mut digest = deserialize_with_centroids(100, 1.0, 9.0, centroids); + if reverse { + let mut bytes = digest.serialize(); + bytes[5] |= 1 << 2; // reverse-merge flag + digest = TDigestMut::deserialize(&bytes).unwrap(); + } + digest + }; + let mut merged = runs.into_iter().map(make_digest).collect::(); + + let mut expected_bytes = + deserialize_with_centroids(100, 1.0, 9.0, &concatenated).serialize(); + // Collection flips the reverse-merge flag once, like a single compression pass. + expected_bytes[5] |= u8::from(!reverse) << 2; + assert_eq!(merged.serialize(), expected_bytes); + assert_eq!(merged.total_weight(), 7); + } +} + #[test] fn test_from_iter_handles_mixed_compressed_and_buffered_inputs() { let mut expected = TDigestMut::new(100).unwrap(); From ead68b214bdc58ff2e846a81e88213fe29708b04 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 4 Oct 2026 23:44:45 +0800 Subject: [PATCH 02/10] perf(tdigest): share one prefix pass across cdf split points cdf and pmf answered each split point with an independent rank query, rescanning the centroid weights twice per point. Precompute each centroid's starting weight once; integer prefix differences are exact, so results are bit-identical while each point costs two binary searches. --- datasketches/src/tdigest/sketch.rs | 61 +++++++++++++++++++----------- 1 file changed, 38 insertions(+), 23 deletions(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index eafdcd6..0c64c10 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -1633,12 +1633,18 @@ impl TDigestView<'_> { return None; } + // One prefix pass keeps each split point at two binary searches instead of two + // linear weight scans. + let mut prefix_weights = Vec::with_capacity(self.centroids.len()); + let mut weight = 0; + for centroid in self.centroids { + prefix_weights.push(weight); + weight += centroid.weight.get(); + } + let mut ranks = Vec::with_capacity(split_points.len() + 1); for &p in split_points { - match self.rank(p) { - Some(rank) => ranks.push(rank), - None => unreachable!("checked non-empty above"), - } + ranks.push(self.rank_with_weights(p, Some(&prefix_weights))); } ranks.push(1.0); Some(ranks) @@ -1650,14 +1656,20 @@ impl TDigestView<'_> { if self.centroids.is_empty() { return None; } + Some(self.rank_with_weights(value, None)) + } + + /// Rank within a non-empty centroid sequence. `prefix_weights` holds each centroid's + /// starting weight when several queries share one prefix pass. + fn rank_with_weights(&self, value: f64, prefix_weights: Option<&[u64]>) -> f64 { if value < self.min { - return Some(0.0); + return 0.0; } if value > self.max { - return Some(1.0); + return 1.0; } if self.min == self.max { - return Some(0.5); + return 0.5; } let centroids_weight = self.centroids_weight as f64; @@ -1669,7 +1681,7 @@ impl TDigestView<'_> { let first = &self.centroids[0]; if value < first.mean { let center = centroid_center(0, first.weight.get()); - return Some(if value == self.min { + return if value == self.min { 0.5 / centroids_weight } else { interpolate( @@ -1677,14 +1689,14 @@ impl TDigestView<'_> { center, interpolation_fraction(value, self.min, first.mean), ) / centroids_weight - }); + }; } let last = &self.centroids[num_centroids - 1]; if value > last.mean { let center = centroid_center(self.centroids_weight - last.weight.get(), last.weight.get()); - return Some(if value == self.max { + return if value == self.max { // Round the maximum's mass in the same coordinates as the centroid centers. centroid_center(self.centroids_weight - 1, 1) / centroids_weight } else { @@ -1693,7 +1705,7 @@ impl TDigestView<'_> { center.max((self.centroids_weight - 1) as f64), interpolation_fraction(value, last.mean, self.max), ) / centroids_weight - }); + }; } let mut lower = self @@ -1713,11 +1725,16 @@ impl TDigestView<'_> { upper -= 1; } - let weight_below = self.centroids[..lower].iter().map(|c| c.weight.get()).sum(); - let weight_between: u64 = self.centroids[lower..upper] - .iter() - .map(|c| c.weight.get()) - .sum(); + let (weight_below, weight_between): (u64, u64) = match prefix_weights { + Some(prefix) => (prefix[lower], prefix[upper] - prefix[lower]), + None => ( + self.centroids[..lower].iter().map(|c| c.weight.get()).sum(), + self.centroids[lower..upper] + .iter() + .map(|c| c.weight.get()) + .sum(), + ), + }; let left = &self.centroids[lower]; let right = &self.centroids[upper]; let fraction = if left.mean < right.mean { @@ -1725,13 +1742,11 @@ impl TDigestView<'_> { } else { 0.5 }; - Some( - interpolate( - centroid_center(weight_below, left.weight.get()), - centroid_center(weight_below + weight_between, right.weight.get()), - fraction, - ) / centroids_weight, - ) + interpolate( + centroid_center(weight_below, left.weight.get()), + centroid_center(weight_below + weight_between, right.weight.get()), + fraction, + ) / centroids_weight } fn quantile(&self, rank: f64) -> Option { From 2b6925bd2e8ade247c0a283491cadc8d25ebfdff Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 00:02:49 +0800 Subject: [PATCH 03/10] refactor(tdigest): drop redundant empty checks in query wrappers The view layer already returns None for empty centroid sequences, and compressing an empty buffer is a no-op, so the early returns in TDigestMut::cdf, pmf, quantile, and quantiles only duplicated the view contract. TDigestMut::rank keeps its check: the empty digest's sentinel extrema would otherwise answer out-of-range queries instead of None. --- datasketches/src/tdigest/sketch.rs | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 0c64c10..14c80fc 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -424,10 +424,6 @@ impl TDigestMut { pub fn cdf(&mut self, split_points: &[f64]) -> Option> { check_split_points(split_points); - if self.is_empty() { - return None; - } - self.view().cdf(split_points) } @@ -453,10 +449,6 @@ impl TDigestMut { pub fn pmf(&mut self, split_points: &[f64]) -> Option> { check_split_points(split_points); - if self.is_empty() { - return None; - } - self.view().pmf(split_points) } @@ -521,10 +513,6 @@ impl TDigestMut { "rank must be in [0.0, 1.0]; got {rank}" ); - if self.is_empty() { - return None; - } - self.view().quantile(rank) } @@ -536,10 +524,6 @@ impl TDigestMut { pub fn quantiles(&mut self, ranks: &[f64]) -> Option> { check_ranks(ranks); - if self.is_empty() { - return None; - } - self.view().quantiles(ranks) } From cbe5e87b2228d6e7d7b1b620419b93bb72044ad4 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 00:03:01 +0800 Subject: [PATCH 04/10] refactor(tdigest): share the compat deserialization pipeline The COMPAT_DOUBLE and COMPAT_FLOAT arms ran the same validation pipeline with only scalar widths, the capacities field, and message labels differing. Parameterize them behind CompatLayout, keeping every error message byte-identical. --- datasketches/src/tdigest/sketch.rs | 237 ++++++++++++++++------------- 1 file changed, 129 insertions(+), 108 deletions(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 14c80fc..3dd4658 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -788,115 +788,64 @@ impl TDigestMut { let mut cursor = SketchSlice::new(bytes); let ty = cursor.read_u32_be().map_err(make_error("type"))?; - match ty { - COMPAT_DOUBLE => { - fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { - move |_| Error::insufficient_data_of("compat double format", tag) - } - // compatibility with asBytes() - let min = cursor.read_f64_be().map_err(make_error("min"))?; - let max = cursor.read_f64_be().map_err(make_error("max"))?; - check_extrema(min, max, "compat double TDigest")?; - let k = cursor.read_f64_be().map_err(make_error("k"))? as u16; - if k < 10 { - return Err(Error::deserial(format!( - "k must be at least 10 in compat double format, got {k}" - ))); - } - let num_centroids = - cursor.read_u32_be().map_err(make_error("num_centroids"))? as usize; - if num_centroids == 0 { - return Err(Error::deserial( - "malformed data: compat double TDigest must contain a centroid", - )); - } - let mut total_weight = 0u64; - let mut centroids = Vec::with_capacity(num_centroids); - let mut previous_mean = min; - let mut centroid_means_valid = true; - for _ in 0..num_centroids { - let weight = cursor.read_f64_be().map_err(make_error("weight"))?; - let mean = cursor.read_f64_be().map_err(make_error("mean"))?; - let weight = - check_compat_weight(weight, "centroid weight in compat double format")?; - centroid_means_valid &= - mean.is_finite() & (mean >= previous_mean) & (mean <= max); - previous_mean = mean; - total_weight = checked_weight_sum(total_weight, weight.get())?; - centroids.push(Centroid { mean, weight }); - } - if !centroid_means_valid { - return Err(Error::deserial( - "malformed data: centroid means in compat double format must be finite, within extrema, and nondecreasing", - )); - } - check_single_sample_extrema(min, max, total_weight)?; - Ok(TDigestMut::make( - k, - false, - min, - max, - TDigestBuffer::new(centroids, 0), - total_weight, - )) - } - COMPAT_FLOAT => { - fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { - move |_| Error::insufficient_data_of("compat float format", tag) - } - // COMPAT_FLOAT: compatibility with asSmallBytes() - // reference implementation uses doubles for min and max - let min = cursor.read_f64_be().map_err(make_error("min"))?; - let max = cursor.read_f64_be().map_err(make_error("max"))?; - check_extrema(min, max, "compat float TDigest")?; - let k = cursor.read_f32_be().map_err(make_error("k"))? as u16; - if k < 10 { - return Err(Error::deserial(format!( - "k must be at least 10 in compat float format, got {k}" - ))); - } - // reference implementation stores capacities of the array of centroids and the - // buffer as shorts they can be derived from k in the constructor - cursor.read_u32_be().map_err(make_error(""))?; - let num_centroids = - cursor.read_u16_be().map_err(make_error("num_centroids"))? as usize; - if num_centroids == 0 { - return Err(Error::deserial( - "malformed data: compat float TDigest must contain a centroid", - )); - } - let mut total_weight = 0u64; - let mut centroids = Vec::with_capacity(num_centroids); - let mut previous_mean = min; - let mut centroid_means_valid = true; - for _ in 0..num_centroids { - let weight = cursor.read_f32_be().map_err(make_error("weight"))? as f64; - let mean = cursor.read_f32_be().map_err(make_error("mean"))? as f64; - let weight = - check_compat_weight(weight, "centroid weight in compat float format")?; - centroid_means_valid &= - mean.is_finite() & (mean >= previous_mean) & (mean <= max); - previous_mean = mean; - total_weight = checked_weight_sum(total_weight, weight.get())?; - centroids.push(Centroid { mean, weight }); - } - if !centroid_means_valid { - return Err(Error::deserial( - "malformed data: centroid means in compat float format must be finite, within extrema, and nondecreasing", - )); - } - check_single_sample_extrema(min, max, total_weight)?; - Ok(TDigestMut::make( - k, - false, - min, - max, - TDigestBuffer::new(centroids, 0), - total_weight, - )) - } - ty => Err(Error::deserial(format!("unknown TDigest compat type {ty}"))), + let layout = match ty { + // compatibility with asBytes() + COMPAT_DOUBLE => CompatLayout::Double, + // compatibility with asSmallBytes() + COMPAT_FLOAT => CompatLayout::Float, + ty => return Err(Error::deserial(format!("unknown TDigest compat type {ty}"))), + }; + + // The reference implementation stores min and max as doubles in both layouts. + let min = cursor + .read_f64_be() + .map_err(|_| Error::insufficient_data_of(layout.format(), "min"))?; + let max = cursor + .read_f64_be() + .map_err(|_| Error::insufficient_data_of(layout.format(), "max"))?; + check_extrema(min, max, layout.subject())?; + let k = layout.read_k(&mut cursor)?; + if k < 10 { + return Err(Error::deserial(format!( + "k must be at least 10 in {}, got {k}", + layout.format() + ))); + } + let num_centroids = layout.read_num_centroids(&mut cursor)?; + if num_centroids == 0 { + return Err(Error::deserial(format!( + "malformed data: {} must contain a centroid", + layout.subject() + ))); + } + let mut total_weight = 0u64; + let mut centroids = Vec::with_capacity(num_centroids); + let mut previous_mean = min; + let mut centroid_means_valid = true; + for _ in 0..num_centroids { + let weight = layout.read_scalar(&mut cursor, "weight")?; + let mean = layout.read_scalar(&mut cursor, "mean")?; + let weight = check_compat_weight(weight, layout.weight_tag())?; + centroid_means_valid &= mean.is_finite() & (mean >= previous_mean) & (mean <= max); + previous_mean = mean; + total_weight = checked_weight_sum(total_weight, weight.get())?; + centroids.push(Centroid { mean, weight }); + } + if !centroid_means_valid { + return Err(Error::deserial(format!( + "malformed data: centroid means in {} must be finite, within extrema, and nondecreasing", + layout.format() + ))); } + check_single_sample_extrema(min, max, total_weight)?; + Ok(TDigestMut::make( + k, + false, + min, + max, + TDigestBuffer::new(centroids, 0), + total_weight, + )) } /// Processes unmerged values and merges centroids if needed. @@ -2069,6 +2018,78 @@ fn check_finite(value: f64, tag: &'static str) -> Result<(), Error> { Ok(()) } +/// The reference implementation's two compat layouts, written by `asBytes()` and +/// `asSmallBytes()`. Scalar widths differ; the validation pipeline and its messages do not. +#[derive(Clone, Copy)] +enum CompatLayout { + Double, + Float, +} + +impl CompatLayout { + /// Malformed-data subject, e.g. "compat double TDigest". + fn subject(self) -> &'static str { + match self { + CompatLayout::Double => "compat double TDigest", + CompatLayout::Float => "compat float TDigest", + } + } + + /// Truncation context, e.g. "compat double format". + fn format(self) -> &'static str { + match self { + CompatLayout::Double => "compat double format", + CompatLayout::Float => "compat float format", + } + } + + /// Centroid weight tag, e.g. "centroid weight in compat double format". + fn weight_tag(self) -> &'static str { + match self { + CompatLayout::Double => "centroid weight in compat double format", + CompatLayout::Float => "centroid weight in compat float format", + } + } + + fn read_k(self, cursor: &mut SketchSlice) -> Result { + let result = match self { + CompatLayout::Double => cursor.read_f64_be(), + CompatLayout::Float => cursor.read_f32_be().map(f64::from), + }; + result + .map(|k| k as u16) + .map_err(|_| Error::insufficient_data_of(self.format(), "k")) + } + + fn read_num_centroids(self, cursor: &mut SketchSlice) -> Result { + match self { + CompatLayout::Double => cursor + .read_u32_be() + .map(|count| count as usize) + .map_err(|_| Error::insufficient_data_of(self.format(), "num_centroids")), + CompatLayout::Float => { + // The reference implementation stores the centroid and buffer capacities as + // shorts; they are derivable from k in the constructor. + cursor + .read_u32_be() + .map_err(|_| Error::insufficient_data_of(self.format(), ""))?; + cursor + .read_u16_be() + .map(|count| count as usize) + .map_err(|_| Error::insufficient_data_of(self.format(), "num_centroids")) + } + } + } + + fn read_scalar(self, cursor: &mut SketchSlice, tag: &'static str) -> Result { + let result = match self { + CompatLayout::Double => cursor.read_f64_be(), + CompatLayout::Float => cursor.read_f32_be().map(f64::from), + }; + result.map_err(|_| Error::insufficient_data_of(self.format(), tag)) + } +} + #[inline] fn check_extrema(min: f64, max: f64, format: &'static str) -> Result<(), Error> { if !min.is_finite() || !max.is_finite() { From d77c8c38c31da3f2ed78d4b620b7094bb620acaf Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 00:06:05 +0800 Subject: [PATCH 05/10] test(tdigest): cover small combined weights in aggregation benchmarks Shapes with a combined weight of at most k / 2 exercise the no-merge fast paths in both borrowed and collected merging. Also document the faster multi-point cdf and pmf in the changelog. --- CHANGELOG.md | 1 + benchmarks/tdigest/aggregation.rs | 13 +++++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 995aa64..bed7a3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +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, which now share one pass over the centroid weights. * 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/aggregation.rs b/benchmarks/tdigest/aggregation.rs index da17069..465842d 100644 --- a/benchmarks/tdigest/aggregation.rs +++ b/benchmarks/tdigest/aggregation.rs @@ -24,8 +24,17 @@ use super::support::build_mut_digest; use super::support::values; // (number of partial states, values per state). Each partial samples the same distribution. -const SERIALIZED_INPUTS: &[(usize, usize)] = - &[(1, 8), (64, 8), (512, 8), (64, 64), (512, 64), (64, 4_096)]; +// Small shapes keep the combined weight at most k / 2, where no two centroids can merge. +const SERIALIZED_INPUTS: &[(usize, usize)] = &[ + (1, 8), + (16, 2), + (64, 1), + (64, 8), + (512, 8), + (64, 64), + (512, 64), + (64, 4_096), +]; #[divan::bench(args = SERIALIZED_INPUTS, sample_size = 1)] fn serialized_borrowed(bencher: Bencher, shape: (usize, usize)) { From 80701c525a75bcf1f4209b13143bb93034af7f59 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 07:31:36 +0800 Subject: [PATCH 06/10] perf(tdigest): accumulate CDF weights as split points advance Reuse the weight below the previous anchor instead of allocating a full prefix array for every CDF call. Sorted split points make the lower anchors monotonic, so each query only sums the newly crossed weights. Keep the existing rank interpolation shared with scalar queries and scope the changelog claim to queries with many split points. --- CHANGELOG.md | 2 +- datasketches/src/tdigest/sketch.rs | 45 +++++++++++++++--------------- 2 files changed, 23 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bed7a3f..8d4b1de 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, which now share one pass over the centroid weights. +* `TDigest` and `TDigestMut` `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/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 3dd4658..24886f3 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -1566,18 +1566,21 @@ impl TDigestView<'_> { return None; } - // One prefix pass keeps each split point at two binary searches instead of two - // linear weight scans. - let mut prefix_weights = Vec::with_capacity(self.centroids.len()); - let mut weight = 0; - for centroid in self.centroids { - prefix_weights.push(weight); - weight += centroid.weight.get(); - } + // Split points increase, so their lower centroid indexes never move backwards. + // Accumulate integer weights only as far as each query needs. + let mut previous_lower = 0; + let mut weight_below = 0; let mut ranks = Vec::with_capacity(split_points.len() + 1); for &p in split_points { - ranks.push(self.rank_with_weights(p, Some(&prefix_weights))); + ranks.push(self.rank_with_weights(p, |lower| { + weight_below += self.centroids[previous_lower..lower] + .iter() + .map(|c| c.weight.get()) + .sum::(); + previous_lower = lower; + weight_below + })); } ranks.push(1.0); Some(ranks) @@ -1589,12 +1592,13 @@ impl TDigestView<'_> { if self.centroids.is_empty() { return None; } - Some(self.rank_with_weights(value, None)) + Some(self.rank_with_weights(value, |lower| { + self.centroids[..lower].iter().map(|c| c.weight.get()).sum() + })) } - /// Rank within a non-empty centroid sequence. `prefix_weights` holds each centroid's - /// starting weight when several queries share one prefix pass. - fn rank_with_weights(&self, value: f64, prefix_weights: Option<&[u64]>) -> f64 { + /// Rank within a non-empty centroid sequence, using the weight before its lower anchor. + fn rank_with_weights(&self, value: f64, weight_before: impl FnOnce(usize) -> u64) -> f64 { if value < self.min { return 0.0; } @@ -1658,16 +1662,11 @@ impl TDigestView<'_> { upper -= 1; } - let (weight_below, weight_between): (u64, u64) = match prefix_weights { - Some(prefix) => (prefix[lower], prefix[upper] - prefix[lower]), - None => ( - self.centroids[..lower].iter().map(|c| c.weight.get()).sum(), - self.centroids[lower..upper] - .iter() - .map(|c| c.weight.get()) - .sum(), - ), - }; + let weight_below = weight_before(lower); + let weight_between: u64 = self.centroids[lower..upper] + .iter() + .map(|c| c.weight.get()) + .sum(); let left = &self.centroids[lower]; let right = &self.centroids[upper]; let fraction = if left.mean < right.mean { From a41f34e9e9a39e241bfc8577f2ae7f6d20f8a531 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 07:31:36 +0800 Subject: [PATCH 07/10] refactor(tdigest): restore specialized compatibility decoders Restore the separate double and float decoding paths. The shared CompatLayout reader adds calls to each centroid iteration and regresses throughput; keeping the existing loops also removes the layout helpers. Validation rules, conversion order, and error messages stay unchanged. --- datasketches/src/tdigest/sketch.rs | 237 +++++++++++++---------------- 1 file changed, 108 insertions(+), 129 deletions(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 24886f3..581e6a3 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -788,64 +788,115 @@ impl TDigestMut { let mut cursor = SketchSlice::new(bytes); let ty = cursor.read_u32_be().map_err(make_error("type"))?; - let layout = match ty { - // compatibility with asBytes() - COMPAT_DOUBLE => CompatLayout::Double, - // compatibility with asSmallBytes() - COMPAT_FLOAT => CompatLayout::Float, - ty => return Err(Error::deserial(format!("unknown TDigest compat type {ty}"))), - }; - - // The reference implementation stores min and max as doubles in both layouts. - let min = cursor - .read_f64_be() - .map_err(|_| Error::insufficient_data_of(layout.format(), "min"))?; - let max = cursor - .read_f64_be() - .map_err(|_| Error::insufficient_data_of(layout.format(), "max"))?; - check_extrema(min, max, layout.subject())?; - let k = layout.read_k(&mut cursor)?; - if k < 10 { - return Err(Error::deserial(format!( - "k must be at least 10 in {}, got {k}", - layout.format() - ))); - } - let num_centroids = layout.read_num_centroids(&mut cursor)?; - if num_centroids == 0 { - return Err(Error::deserial(format!( - "malformed data: {} must contain a centroid", - layout.subject() - ))); - } - let mut total_weight = 0u64; - let mut centroids = Vec::with_capacity(num_centroids); - let mut previous_mean = min; - let mut centroid_means_valid = true; - for _ in 0..num_centroids { - let weight = layout.read_scalar(&mut cursor, "weight")?; - let mean = layout.read_scalar(&mut cursor, "mean")?; - let weight = check_compat_weight(weight, layout.weight_tag())?; - centroid_means_valid &= mean.is_finite() & (mean >= previous_mean) & (mean <= max); - previous_mean = mean; - total_weight = checked_weight_sum(total_weight, weight.get())?; - centroids.push(Centroid { mean, weight }); - } - if !centroid_means_valid { - return Err(Error::deserial(format!( - "malformed data: centroid means in {} must be finite, within extrema, and nondecreasing", - layout.format() - ))); + match ty { + COMPAT_DOUBLE => { + fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { + move |_| Error::insufficient_data_of("compat double format", tag) + } + // compatibility with asBytes() + let min = cursor.read_f64_be().map_err(make_error("min"))?; + let max = cursor.read_f64_be().map_err(make_error("max"))?; + check_extrema(min, max, "compat double TDigest")?; + let k = cursor.read_f64_be().map_err(make_error("k"))? as u16; + if k < 10 { + return Err(Error::deserial(format!( + "k must be at least 10 in compat double format, got {k}" + ))); + } + let num_centroids = + cursor.read_u32_be().map_err(make_error("num_centroids"))? as usize; + if num_centroids == 0 { + return Err(Error::deserial( + "malformed data: compat double TDigest must contain a centroid", + )); + } + let mut total_weight = 0u64; + let mut centroids = Vec::with_capacity(num_centroids); + let mut previous_mean = min; + let mut centroid_means_valid = true; + for _ in 0..num_centroids { + let weight = cursor.read_f64_be().map_err(make_error("weight"))?; + let mean = cursor.read_f64_be().map_err(make_error("mean"))?; + let weight = + check_compat_weight(weight, "centroid weight in compat double format")?; + centroid_means_valid &= + mean.is_finite() & (mean >= previous_mean) & (mean <= max); + previous_mean = mean; + total_weight = checked_weight_sum(total_weight, weight.get())?; + centroids.push(Centroid { mean, weight }); + } + if !centroid_means_valid { + return Err(Error::deserial( + "malformed data: centroid means in compat double format must be finite, within extrema, and nondecreasing", + )); + } + check_single_sample_extrema(min, max, total_weight)?; + Ok(TDigestMut::make( + k, + false, + min, + max, + TDigestBuffer::new(centroids, 0), + total_weight, + )) + } + COMPAT_FLOAT => { + fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { + move |_| Error::insufficient_data_of("compat float format", tag) + } + // COMPAT_FLOAT: compatibility with asSmallBytes() + // reference implementation uses doubles for min and max + let min = cursor.read_f64_be().map_err(make_error("min"))?; + let max = cursor.read_f64_be().map_err(make_error("max"))?; + check_extrema(min, max, "compat float TDigest")?; + let k = cursor.read_f32_be().map_err(make_error("k"))? as u16; + if k < 10 { + return Err(Error::deserial(format!( + "k must be at least 10 in compat float format, got {k}" + ))); + } + // reference implementation stores capacities of the array of centroids and the + // buffer as shorts they can be derived from k in the constructor + cursor.read_u32_be().map_err(make_error(""))?; + let num_centroids = + cursor.read_u16_be().map_err(make_error("num_centroids"))? as usize; + if num_centroids == 0 { + return Err(Error::deserial( + "malformed data: compat float TDigest must contain a centroid", + )); + } + let mut total_weight = 0u64; + let mut centroids = Vec::with_capacity(num_centroids); + let mut previous_mean = min; + let mut centroid_means_valid = true; + for _ in 0..num_centroids { + let weight = cursor.read_f32_be().map_err(make_error("weight"))? as f64; + let mean = cursor.read_f32_be().map_err(make_error("mean"))? as f64; + let weight = + check_compat_weight(weight, "centroid weight in compat float format")?; + centroid_means_valid &= + mean.is_finite() & (mean >= previous_mean) & (mean <= max); + previous_mean = mean; + total_weight = checked_weight_sum(total_weight, weight.get())?; + centroids.push(Centroid { mean, weight }); + } + if !centroid_means_valid { + return Err(Error::deserial( + "malformed data: centroid means in compat float format must be finite, within extrema, and nondecreasing", + )); + } + check_single_sample_extrema(min, max, total_weight)?; + Ok(TDigestMut::make( + k, + false, + min, + max, + TDigestBuffer::new(centroids, 0), + total_weight, + )) + } + ty => Err(Error::deserial(format!("unknown TDigest compat type {ty}"))), } - check_single_sample_extrema(min, max, total_weight)?; - Ok(TDigestMut::make( - k, - false, - min, - max, - TDigestBuffer::new(centroids, 0), - total_weight, - )) } /// Processes unmerged values and merges centroids if needed. @@ -2017,78 +2068,6 @@ fn check_finite(value: f64, tag: &'static str) -> Result<(), Error> { Ok(()) } -/// The reference implementation's two compat layouts, written by `asBytes()` and -/// `asSmallBytes()`. Scalar widths differ; the validation pipeline and its messages do not. -#[derive(Clone, Copy)] -enum CompatLayout { - Double, - Float, -} - -impl CompatLayout { - /// Malformed-data subject, e.g. "compat double TDigest". - fn subject(self) -> &'static str { - match self { - CompatLayout::Double => "compat double TDigest", - CompatLayout::Float => "compat float TDigest", - } - } - - /// Truncation context, e.g. "compat double format". - fn format(self) -> &'static str { - match self { - CompatLayout::Double => "compat double format", - CompatLayout::Float => "compat float format", - } - } - - /// Centroid weight tag, e.g. "centroid weight in compat double format". - fn weight_tag(self) -> &'static str { - match self { - CompatLayout::Double => "centroid weight in compat double format", - CompatLayout::Float => "centroid weight in compat float format", - } - } - - fn read_k(self, cursor: &mut SketchSlice) -> Result { - let result = match self { - CompatLayout::Double => cursor.read_f64_be(), - CompatLayout::Float => cursor.read_f32_be().map(f64::from), - }; - result - .map(|k| k as u16) - .map_err(|_| Error::insufficient_data_of(self.format(), "k")) - } - - fn read_num_centroids(self, cursor: &mut SketchSlice) -> Result { - match self { - CompatLayout::Double => cursor - .read_u32_be() - .map(|count| count as usize) - .map_err(|_| Error::insufficient_data_of(self.format(), "num_centroids")), - CompatLayout::Float => { - // The reference implementation stores the centroid and buffer capacities as - // shorts; they are derivable from k in the constructor. - cursor - .read_u32_be() - .map_err(|_| Error::insufficient_data_of(self.format(), ""))?; - cursor - .read_u16_be() - .map(|count| count as usize) - .map_err(|_| Error::insufficient_data_of(self.format(), "num_centroids")) - } - } - } - - fn read_scalar(self, cursor: &mut SketchSlice, tag: &'static str) -> Result { - let result = match self { - CompatLayout::Double => cursor.read_f64_be(), - CompatLayout::Float => cursor.read_f32_be().map(f64::from), - }; - result.map_err(|_| Error::insufficient_data_of(self.format(), tag)) - } -} - #[inline] fn check_extrema(min: f64, max: f64, format: &'static str) -> Result<(), Error> { if !min.is_finite() || !max.is_finite() { From b68c957b2f4a5c84580dffa87af9c723ba48b85b Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 07:31:36 +0800 Subject: [PATCH 08/10] test(tdigest): reuse interleaved merge coverage at the no-merge boundary Add k=440 to the existing 220-weight interleaved-input fixture and remove the duplicate small-input test. This covers the exact k/2 boundary while retaining weighted ties, both compression directions, and later updates. --- .../tests/tdigest_test/sketch.rs | 41 +------------------ 1 file changed, 2 insertions(+), 39 deletions(-) diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 5716c6a..22bf978 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -363,7 +363,8 @@ fn test_from_iter_matches_single_compression_for_interleaved_runs() { &[(1.0, 2), (3.0, 6), (4.0, 1), (7.0, 8), (9.0, 20)], &[(0.0, 2), (3.0, 9), (6.0, 7), (9.0, 30)], ]; - for k in [10, 100] { + // The inputs weigh 220 in total; k = 440 exercises the no-merge boundary. + for k in [10, 100, 440] { for reverse in [false, true] { let make_digest = |centroids: &[(f64, u64)]| { let mut digest = deserialize_with_centroids(k, -10.0, 20.0, centroids); @@ -401,44 +402,6 @@ fn test_from_iter_matches_single_compression_for_interleaved_runs() { } } -#[test] -fn test_from_iter_concatenates_small_compressed_inputs() { - // A combined weight of at most k / 2 admits no merges, so collection concatenates the - // inputs in stable mean order through the compression fast path instead of the heap scan. - let runs: [&[(f64, u64)]; 3] = [ - &[(2.0, 1), (5.0, 2)], - &[(1.0, 1), (5.0, 1)], - &[(5.0, 1), (9.0, 1)], - ]; - let concatenated = [ - (1.0, 1), - (2.0, 1), - (5.0, 2), // run order breaks mean ties - (5.0, 1), - (5.0, 1), - (9.0, 1), - ]; - for reverse in [false, true] { - let make_digest = |centroids: &[(f64, u64)]| { - let mut digest = deserialize_with_centroids(100, 1.0, 9.0, centroids); - if reverse { - let mut bytes = digest.serialize(); - bytes[5] |= 1 << 2; // reverse-merge flag - digest = TDigestMut::deserialize(&bytes).unwrap(); - } - digest - }; - let mut merged = runs.into_iter().map(make_digest).collect::(); - - let mut expected_bytes = - deserialize_with_centroids(100, 1.0, 9.0, &concatenated).serialize(); - // Collection flips the reverse-merge flag once, like a single compression pass. - expected_bytes[5] |= u8::from(!reverse) << 2; - assert_eq!(merged.serialize(), expected_bytes); - assert_eq!(merged.total_weight(), 7); - } -} - #[test] fn test_from_iter_handles_mixed_compressed_and_buffered_inputs() { let mut expected = TDigestMut::new(100).unwrap(); From 8b29af0b44ec2158c4d97bfadd8707d21effcf8c Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 08:15:37 +0800 Subject: [PATCH 09/10] perf(tdigest): retain heap merging for compressed collections A combined weight below k / 2 proves that no centroids can merge, but sorting the concatenated inputs is not consistently cheaper than merging their existing sorted runs. Restore the original routing after the sort path regressed two compressed 50-row inputs by about 11%. Remove the benchmark shapes and test parameter added for that routing. The existing interleaved, mixed-input, and weighted collection tests remain unchanged from main. Validation: cargo test --package tests-integration --test tdigest_test (52 tests passed). --- benchmarks/tdigest/aggregation.rs | 13 ++----------- datasketches/src/tdigest/sketch.rs | 6 +----- tests-integration/tests/tdigest_test/sketch.rs | 3 +-- 3 files changed, 4 insertions(+), 18 deletions(-) diff --git a/benchmarks/tdigest/aggregation.rs b/benchmarks/tdigest/aggregation.rs index 465842d..da17069 100644 --- a/benchmarks/tdigest/aggregation.rs +++ b/benchmarks/tdigest/aggregation.rs @@ -24,17 +24,8 @@ use super::support::build_mut_digest; use super::support::values; // (number of partial states, values per state). Each partial samples the same distribution. -// Small shapes keep the combined weight at most k / 2, where no two centroids can merge. -const SERIALIZED_INPUTS: &[(usize, usize)] = &[ - (1, 8), - (16, 2), - (64, 1), - (64, 8), - (512, 8), - (64, 64), - (512, 64), - (64, 4_096), -]; +const SERIALIZED_INPUTS: &[(usize, usize)] = + &[(1, 8), (64, 8), (512, 8), (64, 64), (512, 64), (64, 4_096)]; #[divan::bench(args = SERIALIZED_INPUTS, sample_size = 1)] fn serialized_borrowed(bencher: Bencher, shape: (usize, usize)) { diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 581e6a3..377a81d 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -1144,11 +1144,7 @@ impl FromIterator for TDigestMut { }); } - // Below k / 2 combined weight no two centroids fit under the K_2 limit (the bound in - // compress_sorted_centroids), so the heap would only reproduce the stable-sorted - // concatenation one centroid at a time. That case takes the same sort path as mixed - // inputs, which already skips the merge scan. - if !all_compressed || total_weight <= u64::from(merged.k) / 2 { + if !all_compressed { // Raw tails precede compressed prefixes to preserve stable equal-mean ordering. let mut centroids = Vec::with_capacity(num_centroids); for input in &heap { diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 22bf978..428d96c 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -363,8 +363,7 @@ fn test_from_iter_matches_single_compression_for_interleaved_runs() { &[(1.0, 2), (3.0, 6), (4.0, 1), (7.0, 8), (9.0, 20)], &[(0.0, 2), (3.0, 9), (6.0, 7), (9.0, 30)], ]; - // The inputs weigh 220 in total; k = 440 exercises the no-merge boundary. - for k in [10, 100, 440] { + for k in [10, 100] { for reverse in [false, true] { let make_digest = |centroids: &[(f64, u64)]| { let mut digest = deserialize_with_centroids(k, -10.0, 20.0, centroids); From 9b1bb86aced35eaa3081ec60b9e117bc84c5ad17 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 5 Oct 2026 08:17:34 +0800 Subject: [PATCH 10/10] perf(tdigest): restore fast returns for empty mutable queries Return None from cdf, pmf, quantile, and quantiles before entering the compression and view path for an empty digest. The view-level checks remain necessary for frozen queries but do not make these early returns redundant for performance. Keep argument validation before the early returns, preserving the existing panic contracts even when the digest is empty. Validation: cargo x check, cargo x test (745 tests), and cargo x lint. --- datasketches/src/tdigest/sketch.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 377a81d..4b6a27d 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -424,6 +424,10 @@ impl TDigestMut { pub fn cdf(&mut self, split_points: &[f64]) -> Option> { check_split_points(split_points); + if self.is_empty() { + return None; + } + self.view().cdf(split_points) } @@ -449,6 +453,10 @@ impl TDigestMut { pub fn pmf(&mut self, split_points: &[f64]) -> Option> { check_split_points(split_points); + if self.is_empty() { + return None; + } + self.view().pmf(split_points) } @@ -513,6 +521,10 @@ impl TDigestMut { "rank must be in [0.0, 1.0]; got {rank}" ); + if self.is_empty() { + return None; + } + self.view().quantile(rank) } @@ -524,6 +536,10 @@ impl TDigestMut { pub fn quantiles(&mut self, ranks: &[f64]) -> Option> { check_ranks(ranks); + if self.is_empty() { + return None; + } + self.view().quantiles(ranks) }