From 7c1a10cca01046a88bd8b83137eb9b10a006ae36 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 07:11:00 +0800 Subject: [PATCH 1/6] feat(tdigest): add checked batch aggregation APIs --- CHANGELOG.md | 13 + benchmarks/tdigest/merge.rs | 19 + benchmarks/tdigest/query.rs | 21 + datasketches/src/tdigest/sketch.rs | 370 +++++++++++++++--- .../tests/serde_tests/tdigest.rs | 54 +++ .../tests/tdigest_test/sketch.rs | 81 ++++ 6 files changed, 499 insertions(+), 59 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 969a56ae..2a6daa5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,19 @@ All significant changes to this project will be documented in this file. ## Unreleased +### New features + +* Add `TDigestMut::try_merge` and `TDigestMut::merge_many` for checked and batch merging, and add `TDigestMut::quantiles` and `TDigest::quantiles` for querying several ranks in one centroid scan. + +### Performance improvements + +* T-Digest batch merge avoids recompressing each partial sketch, and batch quantile queries reuse one traversal for ranks supplied in nondecreasing order. + +### Bug fixes + +* T-Digest merging now uses the smaller `k` when sketches have different compression parameters, preserving the size bound of the coarser input. +* T-Digest deserialization now rejects unknown or conflicting flags, reversed extrema, out-of-range values, unsorted centroids, and non-empty images without stored values. + ## v0.5.0 ### Breaking changes diff --git a/benchmarks/tdigest/merge.rs b/benchmarks/tdigest/merge.rs index 13559c66..7f67b2c2 100644 --- a/benchmarks/tdigest/merge.rs +++ b/benchmarks/tdigest/merge.rs @@ -102,6 +102,25 @@ fn partials(bencher: Bencher) { }); } +#[divan::bench] +fn partials_batch(bencher: Bencher) { + let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL) + .into_iter() + .map(|mut digest| { + black_box(digest.rank(0.0)); + digest + }) + .collect::>(); + + bencher + .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL)) + .bench_local(|| { + let mut merged = TDigestMut::new(DEFAULT_DIGEST_K).unwrap(); + merged.merge_many(black_box(&partials)).unwrap(); + black_box(merged) + }); +} + #[divan::bench(args = [SMALL_ROWS_PER_PARTIAL, ROWS_PER_PARTIAL])] fn serialized_partials(bencher: Bencher, rows_per_partial: usize) { let partials = serialized_partial_digests(64, rows_per_partial); diff --git a/benchmarks/tdigest/query.rs b/benchmarks/tdigest/query.rs index 0981657d..62422d3f 100644 --- a/benchmarks/tdigest/query.rs +++ b/benchmarks/tdigest/query.rs @@ -56,3 +56,24 @@ fn quantiles_2_sequential(bencher: Bencher) { ] }); } + +#[divan::bench] +fn quantiles_6_sequential(bencher: Bencher) { + let digest = prepared_digest(); + let ranks = [0.25, 0.5, 0.75, 0.9, 0.95, 0.99]; + + bencher + .bench_local(|| black_box(ranks).map(|rank| black_box(&digest).quantile(black_box(rank)))); +} + +#[divan::bench(args = [2, 6, 100])] +fn quantiles_batch(bencher: Bencher, num_ranks: usize) { + let digest = prepared_digest(); + let ranks = (1..=num_ranks) + .map(|rank| rank as f64 / (num_ranks + 1) as f64) + .collect::>(); + + bencher + .counter(ItemsCount::new(num_ranks)) + .bench_local(|| black_box(&digest).quantiles(black_box(&ranks))); +} diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 942c7358..597dc1e3 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -142,6 +142,16 @@ impl TDigestBuffer { self.centroids } + fn append_unmerged_to(&self, output: &mut Vec) { + let compressed_prefix_len = self.compressed_prefix_len(); + output.extend_from_slice(&self.centroids[compressed_prefix_len..]); + } + + fn append_compressed_to(&self, output: &mut Vec) { + let compressed_prefix_len = self.compressed_prefix_len(); + output.extend_from_slice(&self.centroids[..compressed_prefix_len]); + } + fn compressed_centroids(&self) -> &[Centroid] { assert_eq!( self.unmerged_tail_len, 0, @@ -322,6 +332,14 @@ impl TDigestMut { /// Merges the given t-digest into this one. /// + /// If the sketches have different `k` values, the merged sketch uses the smaller value because + /// centroids produced with a smaller `k` cannot be split to satisfy the tighter size bound. + /// + /// # Panics + /// + /// Panics if the combined total weight exceeds `u64::MAX`. Use [`try_merge`](Self::try_merge) + /// to handle that condition. + /// /// # Examples /// /// ``` @@ -335,13 +353,92 @@ impl TDigestMut { /// assert_eq!(left.total_weight(), 2); /// ``` pub fn merge(&mut self, other: &TDigestMut) { + self.try_merge(other) + .expect("combined t-digest weight exceeds u64::MAX"); + } + + /// Merges the given t-digest into this one. + /// + /// If the sketches have different `k` values, the merged sketch uses the smaller value because + /// centroids produced with a smaller `k` cannot be split to satisfy the tighter size bound. + /// + /// # Errors + /// + /// Returns an error if the combined total weight exceeds `u64::MAX`. + pub fn try_merge(&mut self, other: &TDigestMut) -> Result<(), Error> { if other.is_empty() { - return; + return Ok(()); } let self_unmerged_weight = self.buffer.unmerged_len() as u64; + let additional_weight = + checked_merged_weight_sum(self_unmerged_weight, other.total_weight())?; + checked_merged_weight_sum(self.compressed_weight, additional_weight)?; let centroids = std::mem::take(&mut self.buffer).into_merged_centroids(&other.buffer); - self.compress_sorted_centroids(centroids, self_unmerged_weight + other.total_weight()) + self.k = self.k.min(other.k); + self.min = self.min.min(other.min); + self.max = self.max.max(other.max); + self.compress_sorted_centroids(centroids, additional_weight); + Ok(()) + } + + /// Merges several t-digests into this one with one compression pass. + /// + /// Empty inputs are ignored. The merged sketch uses the smallest `k` among the non-empty + /// inputs and this sketch. This method temporarily retains every input centroid to avoid + /// recompressing intermediate results. + /// + /// # Errors + /// + /// Returns an error if the combined total weight exceeds `u64::MAX`. + pub fn merge_many<'a>( + &mut self, + others: impl IntoIterator, + ) -> Result<(), Error> { + let others = others + .into_iter() + .filter(|other| !other.is_empty()) + .collect::>(); + if others.is_empty() { + return Ok(()); + } + + let mut additional_weight = self.buffer.unmerged_len() as u64; + let mut num_centroids = self.buffer.len(); + let mut k = self.k; + let mut min = self.min; + let mut max = self.max; + for other in &others { + additional_weight = checked_merged_weight_sum(additional_weight, other.total_weight())?; + num_centroids = num_centroids + .checked_add(other.buffer.len()) + .ok_or_else(|| { + Error::invalid_argument("combined t-digest centroid count exceeds usize::MAX") + })?; + k = k.min(other.k); + min = min.min(other.min); + max = max.max(other.max); + } + checked_merged_weight_sum(self.compressed_weight, additional_weight)?; + + let own_buffer = std::mem::take(&mut self.buffer); + let mut centroids = Vec::with_capacity(num_centroids); + // Stable sorting keeps raw values before existing summaries when their means are equal. + own_buffer.append_unmerged_to(&mut centroids); + for other in &others { + other.buffer.append_unmerged_to(&mut centroids); + } + own_buffer.append_compressed_to(&mut centroids); + for other in &others { + other.buffer.append_compressed_to(&mut centroids); + } + centroids.sort_by(centroid_cmp); + + self.k = k; + self.min = min; + self.max = max; + self.compress_sorted_centroids(centroids, additional_weight); + Ok(()) } /// Converts this mutable t-digest into an immutable one. @@ -506,6 +603,21 @@ impl TDigestMut { self.view().quantile(rank) } + /// Returns the quantiles described by [`TDigest::quantiles`]. + /// + /// # Panics + /// + /// Panics if any rank is outside `[0.0, 1.0]`. + pub fn quantiles(&mut self, ranks: &[f64]) -> Option> { + check_ranks(ranks); + + if self.is_empty() { + return None; + } + + self.view().quantiles(ranks) + } + /// Serializes this mutable t-digest to bytes. /// /// # Examples @@ -595,8 +707,20 @@ impl TDigestMut { return Err(Error::deserial(format!("k must be at least 10, got {k}"))); } let flags = cursor.read_u8().map_err(insufficient_data("flags"))?; + let known_flags = FLAGS_IS_EMPTY | FLAGS_IS_SINGLE_VALUE | FLAGS_REVERSE_MERGE; + if flags & !known_flags != 0 { + return Err(Error::deserial(format!( + "malformed data: unknown TDigest flags 0x{:02x}", + flags & !known_flags + ))); + } let is_empty = (flags & FLAGS_IS_EMPTY) != 0; let is_single_value = (flags & FLAGS_IS_SINGLE_VALUE) != 0; + if is_empty && is_single_value { + return Err(Error::deserial( + "malformed data: empty and single-value flags are mutually exclusive", + )); + } let expected_preamble_longs = if is_empty || is_single_value { PREAMBLE_LONGS_EMPTY_OR_SINGLE } else { @@ -655,10 +779,7 @@ impl TDigestMut { cursor.read_f64_le().map_err(insufficient_data("max"))?, ) }; - check_non_nan(min, "min")?; - check_non_nan(max, "max")?; - check_finite(min, "min")?; - check_finite(max, "max")?; + check_extrema(min, max, "TDigest")?; let (centroid_bytes, buffered_value_bytes) = if is_f32 { (size_of::() + size_of::(), size_of::()) } else { @@ -686,8 +807,15 @@ impl TDigestMut { let stored_centroids = num_centroids.checked_add(num_buffered).ok_or_else(|| { Error::deserial("num_centroids and num_buffered exceed the supported size") })?; + if stored_centroids == 0 { + return Err(Error::deserial( + "malformed data: non-empty TDigest must contain a centroid or buffered value", + )); + } let mut centroids = Vec::with_capacity(stored_centroids); let mut compressed_weight = 0u64; + let mut previous_mean = min; + let mut centroid_means_valid = true; for bytes in centroid_payload.chunks_exact(centroid_bytes) { let (mean, weight) = if is_f32 { ( @@ -700,26 +828,36 @@ impl TDigestMut { u64::from_le_bytes(bytes[8..].try_into().unwrap()), ) }; - check_non_nan(mean, "centroid mean")?; - check_finite(mean, "centroid")?; + centroid_means_valid &= mean.is_finite() & (mean >= previous_mean) & (mean <= max); + previous_mean = mean; let weight = check_nonzero(weight, "centroid weight")?; compressed_weight = checked_weight_sum(compressed_weight, weight.get())?; centroids.push(Centroid { mean, weight }); } + if !centroid_means_valid { + return Err(Error::deserial( + "malformed data: centroid means must be finite, within extrema, and nondecreasing", + )); + } checked_weight_sum(compressed_weight, num_buffered as u64)?; + let mut buffered_values_valid = true; for bytes in buffered_payload.chunks_exact(buffered_value_bytes) { let value = if is_f32 { f32::from_le_bytes(bytes.try_into().unwrap()) as f64 } else { f64::from_le_bytes(bytes.try_into().unwrap()) }; - check_non_nan(value, "buffered_value mean")?; - check_finite(value, "buffered_value mean")?; + buffered_values_valid &= value.is_finite() & (value >= min) & (value <= max); centroids.push(Centroid { mean: value, weight: DEFAULT_WEIGHT, }); } + if !buffered_values_valid { + return Err(Error::deserial( + "malformed data: buffered values must be finite and within extrema", + )); + } Ok(TDigestMut::make( k, reverse_merge, @@ -748,10 +886,7 @@ impl TDigestMut { // 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_non_nan(min, "min in compat double format")?; - check_non_nan(max, "max in compat double format")?; - check_finite(min, "min in compat double format")?; - check_finite(max, "max in compat double format")?; + 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!( @@ -760,18 +895,31 @@ impl TDigestMut { } 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")?; - check_non_nan(mean, "centroid mean in compat double format")?; - check_finite(mean, "centroid mean 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", + )); + } Ok(TDigestMut::make( k, false, @@ -789,10 +937,7 @@ impl TDigestMut { // 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_non_nan(min, "min in compat float format")?; - check_non_nan(max, "max in compat float format")?; - check_finite(min, "min in compat float format")?; - check_finite(max, "max in compat float format")?; + 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!( @@ -804,18 +949,31 @@ impl TDigestMut { 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")?; - check_non_nan(mean, "centroid mean in compat float format")?; - check_finite(mean, "centroid mean 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", + )); + } Ok(TDigestMut::make( k, false, @@ -1222,6 +1380,35 @@ impl TDigest { self.view().quantile(rank) } + /// Computes approximate quantiles for the given normalized ranks. + /// + /// Ranks in nondecreasing order are answered with one centroid scan. Ranks in any other order + /// are accepted and results are returned in the same order as the input. + /// + /// Returns `None` if this t-digest is empty. + /// + /// # Panics + /// + /// Panics if any rank is outside `[0.0, 1.0]`. + /// + /// # Examples + /// + /// ``` + /// use datasketches::tdigest::TDigestMut; + /// + /// let mut sketch = TDigestMut::new(100).unwrap(); + /// for value in [1.0, 2.0, 3.0] { + /// sketch.update(value); + /// } + /// let digest = sketch.freeze(); + /// let quantiles = digest.quantiles(&[0.25, 0.5, 0.75]).unwrap(); + /// assert_eq!(quantiles.len(), 3); + /// ``` + pub fn quantiles(&self, ranks: &[f64]) -> Option> { + check_ranks(ranks); + self.view().quantiles(ranks) + } + /// Converts this immutable t-digest into a mutable one. /// /// # Examples @@ -1389,8 +1576,51 @@ impl TDigestView<'_> { return None; } + let mut centroid_index = 0; + let mut weight_so_far = self.centroids[0].weight() / 2.; + Some(self.quantile_with_cursor(rank, &mut centroid_index, &mut weight_so_far)) + } + + fn quantiles(&self, ranks: &[f64]) -> Option> { + debug_assert!( + ranks.iter().all(|rank| (0.0..=1.0).contains(rank)), + "ranks must be in [0.0, 1.0]" + ); + + if self.centroids.is_empty() { + return None; + } + + let mut quantiles = vec![0.; ranks.len()]; + let mut centroid_index = 0; + let mut weight_so_far = self.centroids[0].weight() / 2.; + if ranks.windows(2).all(|pair| pair[0] <= pair[1]) { + for (index, &rank) in ranks.iter().enumerate() { + quantiles[index] = + self.quantile_with_cursor(rank, &mut centroid_index, &mut weight_so_far); + } + return Some(quantiles); + } + + let mut rank_order = (0..ranks.len()).collect::>(); + rank_order.sort_by(|&left, &right| ranks[left].total_cmp(&ranks[right])); + for index in rank_order { + quantiles[index] = + self.quantile_with_cursor(ranks[index], &mut centroid_index, &mut weight_so_far); + } + Some(quantiles) + } + + fn quantile_with_cursor( + &self, + rank: f64, + centroid_index: &mut usize, + weight_so_far: &mut f64, + ) -> f64 { + debug_assert!(!self.centroids.is_empty()); + if self.centroids.len() == 1 { - return Some(self.centroids[0].mean); + return self.centroids[0].mean; } // at least 2 centroids @@ -1398,73 +1628,66 @@ impl TDigestView<'_> { let num_centroids = self.centroids.len(); let weight = rank * centroids_weight; if weight < 1. { - return Some(self.min); + return self.min; } if weight > centroids_weight - 1. { - return Some(self.max); + return self.max; } let first_weight = self.centroids[0].weight(); if first_weight > 1. && weight < first_weight / 2. { - return Some( - self.min - + (((weight - 1.) / ((first_weight / 2.) - 1.)) - * (self.centroids[0].mean - self.min)), - ); + return self.min + + (((weight - 1.) / ((first_weight / 2.) - 1.)) + * (self.centroids[0].mean - self.min)); } let last_weight = self.centroids[num_centroids - 1].weight(); if last_weight > 1. && (centroids_weight - weight <= last_weight / 2.) { if last_weight == 2. { - return Some(self.max); + return self.max; } - return Some( - self.max - - (((centroids_weight - weight - 1.) / ((last_weight / 2.) - 1.)) - * (self.max - self.centroids[num_centroids - 1].mean)), - ); + return self.max + - (((centroids_weight - weight - 1.) / ((last_weight / 2.) - 1.)) + * (self.max - self.centroids[num_centroids - 1].mean)); } // interpolate between extremes - let mut weight_so_far = first_weight / 2.; - for i in 0..(num_centroids - 1) { - let dw = (self.centroids[i].weight() + self.centroids[i + 1].weight()) / 2.; - if weight_so_far + dw > weight { + while *centroid_index < num_centroids - 1 { + let dw = (self.centroids[*centroid_index].weight() + + self.centroids[*centroid_index + 1].weight()) + / 2.; + if *weight_so_far + dw > weight { // the target weight is between centroids i and i+1 let mut left_weight = 0.; - if self.centroids[i].weight.get() == 1 { - if weight - weight_so_far < 0.5 { - return Some(self.centroids[i].mean); + if self.centroids[*centroid_index].weight.get() == 1 { + if weight - *weight_so_far < 0.5 { + return self.centroids[*centroid_index].mean; } left_weight = 0.5; } let mut right_weight = 0.; - if self.centroids[i + 1].weight.get() == 1 { - if weight_so_far + dw - weight <= 0.5 { - return Some(self.centroids[i + 1].mean); + if self.centroids[*centroid_index + 1].weight.get() == 1 { + if *weight_so_far + dw - weight <= 0.5 { + return self.centroids[*centroid_index + 1].mean; } right_weight = 0.5; } // Each centroid is weighted by the distance from the target to the *other* // centroid, so the estimate approaches the nearer one. - let distance_from_left = weight - weight_so_far - left_weight; - let distance_to_right = weight_so_far + dw - weight - right_weight; - return Some(weighted_average( - self.centroids[i].mean, + let distance_from_left = weight - *weight_so_far - left_weight; + let distance_to_right = *weight_so_far + dw - weight - right_weight; + return weighted_average( + self.centroids[*centroid_index].mean, distance_to_right, - self.centroids[i + 1].mean, + self.centroids[*centroid_index + 1].mean, distance_from_left, - )); + ); } - weight_so_far += dw; + *weight_so_far += dw; + *centroid_index += 1; } let w1 = weight - (centroids_weight) - ((self.centroids[num_centroids - 1].weight()) / 2.); let w2 = (self.centroids[num_centroids - 1].weight() / 2.) - w1; - Some(weighted_average( - self.centroids[num_centroids - 1].mean, - w1, - self.max, - w2, - )) + weighted_average(self.centroids[num_centroids - 1].mean, w1, self.max, w2) } } @@ -1480,6 +1703,14 @@ fn check_split_points(split_points: &[f64]) { } } +#[track_caller] +fn check_ranks(ranks: &[f64]) { + assert!( + ranks.iter().all(|rank| (0.0..=1.0).contains(rank)), + "ranks must be in [0.0, 1.0]" + ); +} + fn centroid_cmp(a: &Centroid, b: &Centroid) -> Ordering { match a.mean.partial_cmp(&b.mean) { Some(order) => order, @@ -1597,6 +1828,21 @@ fn check_finite(value: f64, tag: &'static str) -> Result<(), Error> { Ok(()) } +#[inline] +fn check_extrema(min: f64, max: f64, format: &'static str) -> Result<(), Error> { + if !min.is_finite() || !max.is_finite() { + return Err(Error::deserial(format!( + "malformed data: {format} extrema must be finite" + ))); + } + if min > max { + return Err(Error::deserial(format!( + "malformed data: {format} min {min} exceeds max {max}" + ))); + } + Ok(()) +} + fn check_nonzero(value: u64, tag: &'static str) -> Result { NonZeroU64::new(value) .ok_or_else(|| Error::deserial(format!("malformed data: {tag} cannot be zero"))) @@ -1624,6 +1870,12 @@ fn checked_weight_sum(total_weight: u64, weight: u64) -> Result { .ok_or_else(|| Error::deserial("malformed data: total weight overflow")) } +fn checked_merged_weight_sum(total_weight: u64, weight: u64) -> Result { + total_weight + .checked_add(weight) + .ok_or_else(|| Error::invalid_argument("combined t-digest weight exceeds u64::MAX")) +} + /// Generates cluster sizes proportional to `q*(1-q)`. /// /// The use of a normalizing function results in a strictly bounded number of clusters no matter diff --git a/tests-integration/tests/serde_tests/tdigest.rs b/tests-integration/tests/serde_tests/tdigest.rs index c5bd0fe1..1b871ef8 100644 --- a/tests-integration/tests/serde_tests/tdigest.rs +++ b/tests-integration/tests/serde_tests/tdigest.rs @@ -350,6 +350,60 @@ fn test_updates_normalize_overfull_deserialized_mixed_buffer() { assert_eq!(roundtrip.max_value(), Some(1_000.0)); } +fn serialized_two_value_digest() -> Vec { + let mut tdigest = TDigestMut::new(100).unwrap(); + tdigest.update(0.0); + tdigest.update(1.0); + tdigest.serialize() +} + +fn assert_invalid_tdigest(bytes: &[u8]) { + let error = TDigestMut::deserialize(bytes).unwrap_err(); + assert_eq!(error.kind(), datasketches::error::ErrorKind::InvalidData); +} + +#[test] +fn test_deserialize_rejects_unknown_or_conflicting_flags() { + let mut unknown = serialized_two_value_digest(); + unknown[5] |= 0x80; + assert_invalid_tdigest(&unknown); + + let mut empty = TDigestMut::new(100).unwrap(); + let mut conflicting = empty.serialize(); + conflicting[5] |= 1 << 1; + assert_invalid_tdigest(&conflicting); +} + +#[test] +fn test_deserialize_rejects_invalid_extrema_and_centroid_ranges() { + let mut reversed_extrema = serialized_two_value_digest(); + reversed_extrema[16..24].copy_from_slice(&2_f64.to_le_bytes()); + assert_invalid_tdigest(&reversed_extrema); + + let mut centroid_outside_extrema = serialized_two_value_digest(); + centroid_outside_extrema[32..40].copy_from_slice(&(-1_f64).to_le_bytes()); + assert_invalid_tdigest(¢roid_outside_extrema); + + let mut buffered_outside_extrema = serialized_two_value_digest(); + buffered_outside_extrema[8..12].copy_from_slice(&1_u32.to_le_bytes()); + buffered_outside_extrema[12..16].copy_from_slice(&1_u32.to_le_bytes()); + buffered_outside_extrema[48..56].copy_from_slice(&2_f64.to_le_bytes()); + assert_invalid_tdigest(&buffered_outside_extrema); +} + +#[test] +fn test_deserialize_rejects_unsorted_or_missing_centroids() { + let mut unsorted = serialized_two_value_digest(); + unsorted[32..40].copy_from_slice(&1_f64.to_le_bytes()); + unsorted[48..56].copy_from_slice(&0_f64.to_le_bytes()); + assert_invalid_tdigest(&unsorted); + + let mut missing = serialized_two_value_digest(); + missing[8..12].copy_from_slice(&0_u32.to_le_bytes()); + missing[12..16].copy_from_slice(&0_u32.to_le_bytes()); + assert_invalid_tdigest(&missing); +} + #[test] fn test_deserialize_rejects_truncated_large_payload_before_allocation() { let mut tdigest = TDigestMut::new(10).unwrap(); diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 302cc4c7..8491ef94 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -275,6 +275,43 @@ fn test_merge_large() { assert_that!(td1.rank(n as f64).unwrap(), eq(1.0)); } +#[test] +fn test_mixed_k_merge_uses_smaller_k() { + let mut left = TDigestMut::new(200).unwrap(); + let mut right = TDigestMut::new(50).unwrap(); + for value in 0..1_000 { + left.update(value as f64); + right.update((value + 1_000) as f64); + } + + left.try_merge(&right).unwrap(); + + assert_eq!(left.k(), 50); + assert_eq!(left.total_weight(), 2_000); + assert_eq!(left.min_value(), Some(0.0)); + assert_eq!(left.max_value(), Some(1_999.0)); +} + +#[test] +fn test_merge_many_uses_one_result_with_the_smallest_nonempty_k() { + let mut merged = TDigestMut::new(200).unwrap(); + let mut first = TDigestMut::new(100).unwrap(); + let mut second = TDigestMut::new(50).unwrap(); + let empty = TDigestMut::new(10).unwrap(); + for value in 0..1_000 { + first.update(value as f64); + second.update((value + 1_000) as f64); + } + + merged.merge_many([&first, &empty, &second]).unwrap(); + + assert_eq!(merged.k(), 50); + assert_eq!(merged.total_weight(), 2_000); + assert_eq!(merged.min_value(), Some(0.0)); + assert_eq!(merged.max_value(), Some(1_999.0)); + assert_quantiles_are_nondecreasing(&mut merged); +} + #[test] fn test_invalid_inputs() { let n = 100; @@ -324,6 +361,33 @@ fn test_extreme_values_produce_finite_quantiles() { } } +#[test] +fn test_batch_quantiles_match_scalar_queries_in_input_order() { + let mut tdigest = TDigestMut::new(100).unwrap(); + for value in 0..10_000 { + tdigest.update(((value * 37) % 1_003) as f64); + } + let tdigest = tdigest.freeze(); + + for ranks in [ + vec![0.0, 0.001, 0.25, 0.5, 0.5, 0.99, 1.0], + vec![0.99, 0.0, 0.5, 1.0, 0.001, 0.5, 0.25], + vec![], + ] { + let expected = ranks + .iter() + .map(|&rank| tdigest.quantile(rank).unwrap()) + .collect::>(); + assert_eq!(tdigest.quantiles(&ranks), Some(expected)); + } +} + +fn assert_quantiles_are_nondecreasing(tdigest: &mut TDigestMut) { + let ranks = (0..=100).map(|rank| rank as f64 / 100.).collect::>(); + let quantiles = tdigest.quantiles(&ranks).unwrap(); + assert!(quantiles.windows(2).all(|pair| pair[0] <= pair[1])); +} + #[test] fn test_estimate_repeat_values() { let mut tdigest = TDigestMut::default(); @@ -393,6 +457,23 @@ fn test_quantile_handles_two_sample_last_centroid() { assert_eq!(tdigest.quantile(0.75), Some(100.0)); } +#[test] +fn test_try_merge_reports_total_weight_overflow_without_mutating_receiver() { + let mut left = deserialize_with_centroids(100, 0.0, 0.0, &[(0.0, u64::MAX)]); + let right = deserialize_with_centroids(50, 1.0, 1.0, &[(1.0, 1)]); + + let error = left.try_merge(&right).unwrap_err(); + + assert_eq!( + error.kind(), + datasketches::error::ErrorKind::InvalidArgument + ); + assert_eq!(left.k(), 100); + assert_eq!(left.total_weight(), u64::MAX); + assert_eq!(left.min_value(), Some(0.0)); + assert_eq!(left.max_value(), Some(0.0)); +} + #[test] fn test_rank_left_tail_is_a_fraction_of_the_total_weight() { let mut tdigest = From 751f7cb27e7986cfd7006292a7ab4445e47d619a Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 07:51:02 +0800 Subject: [PATCH 2/6] perf(tdigest): merge sorted partials without resorting --- datasketches/src/tdigest/sketch.rs | 97 ++++++++++++++++--- .../tests/tdigest_test/sketch.rs | 2 + 2 files changed, 88 insertions(+), 11 deletions(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 597dc1e3..2fe01c2f 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -16,6 +16,7 @@ // under the License. use std::cmp::Ordering; +use std::collections::BinaryHeap; use std::convert::identity; use std::num::NonZeroU64; @@ -422,17 +423,27 @@ impl TDigestMut { checked_merged_weight_sum(self.compressed_weight, additional_weight)?; let own_buffer = std::mem::take(&mut self.buffer); - let mut centroids = Vec::with_capacity(num_centroids); - // Stable sorting keeps raw values before existing summaries when their means are equal. - own_buffer.append_unmerged_to(&mut centroids); - for other in &others { - other.buffer.append_unmerged_to(&mut centroids); - } - own_buffer.append_compressed_to(&mut centroids); - for other in &others { - other.buffer.append_compressed_to(&mut centroids); - } - centroids.sort_by(centroid_cmp); + let centroids = if own_buffer.unmerged_len() == 0 + && others.iter().all(|other| other.buffer.unmerged_len() == 0) + { + let sources = std::iter::once(own_buffer.centroids.as_slice()) + .chain(others.iter().map(|other| other.buffer.centroids.as_slice())) + .collect::>(); + merge_sorted_centroid_slices(&sources, num_centroids) + } else { + let mut centroids = Vec::with_capacity(num_centroids); + // Stable sorting keeps raw values before existing summaries when their means are equal. + own_buffer.append_unmerged_to(&mut centroids); + for other in &others { + other.buffer.append_unmerged_to(&mut centroids); + } + own_buffer.append_compressed_to(&mut centroids); + for other in &others { + other.buffer.append_compressed_to(&mut centroids); + } + centroids.sort_by(centroid_cmp); + centroids + }; self.k = k; self.min = min; @@ -1724,6 +1735,70 @@ fn centroids_are_sorted(centroids: &[Centroid]) -> bool { .all(|pair| centroid_cmp(&pair[0], &pair[1]) != Ordering::Greater) } +#[derive(Clone, Copy)] +struct CentroidCursor { + centroid: Centroid, + source_index: usize, + centroid_index: usize, +} + +impl PartialEq for CentroidCursor { + fn eq(&self, other: &Self) -> bool { + self.source_index == other.source_index && self.centroid_index == other.centroid_index + } +} + +impl Eq for CentroidCursor {} + +impl PartialOrd for CentroidCursor { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for CentroidCursor { + fn cmp(&self, other: &Self) -> Ordering { + // Reverse every key because BinaryHeap is a max-heap. Source order breaks equal-mean ties + // in the same way as concatenating the inputs and applying a stable sort. + centroid_cmp(&other.centroid, &self.centroid) + .then_with(|| other.source_index.cmp(&self.source_index)) + .then_with(|| other.centroid_index.cmp(&self.centroid_index)) + } +} + +fn merge_sorted_centroid_slices(sources: &[&[Centroid]], num_centroids: usize) -> Vec { + debug_assert_eq!( + sources.iter().map(|source| source.len()).sum::(), + num_centroids + ); + debug_assert!(sources.iter().all(|source| centroids_are_sorted(source))); + + let mut heap = BinaryHeap::with_capacity(sources.len()); + for (source_index, source) in sources.iter().enumerate() { + if let Some(¢roid) = source.first() { + heap.push(CentroidCursor { + centroid, + source_index, + centroid_index: 0, + }); + } + } + + let mut centroids = Vec::with_capacity(num_centroids); + while let Some(cursor) = heap.pop() { + centroids.push(cursor.centroid); + let centroid_index = cursor.centroid_index + 1; + if let Some(¢roid) = sources[cursor.source_index].get(centroid_index) { + heap.push(CentroidCursor { + centroid, + source_index: cursor.source_index, + centroid_index, + }); + } + } + centroids +} + fn merge_sorted_centroids(left: &mut Vec, right: &[Centroid]) { debug_assert!(!right.is_empty()); debug_assert!(centroids_are_sorted(left)); diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 8491ef94..707022ab 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -302,6 +302,8 @@ fn test_merge_many_uses_one_result_with_the_smallest_nonempty_k() { first.update(value as f64); second.update((value + 1_000) as f64); } + let _ = first.rank(0.0); + let _ = second.rank(0.0); merged.merge_many([&first, &empty, &second]).unwrap(); From 0270ef6b0d545631b297917f983708815dae5867 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 10:40:14 +0800 Subject: [PATCH 3/6] refactor(tdigest): collect owned partials for batch merge --- CHANGELOG.md | 5 +- benchmarks/tdigest/merge.rs | 11 +- datasketches/src/tdigest/sketch.rs | 225 ++++++------------ .../tests/serde_tests/tdigest.rs | 54 ----- .../tests/tdigest_test/sketch.rs | 39 +-- 5 files changed, 105 insertions(+), 229 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a6daa5f..a2656462 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,16 +6,15 @@ All significant changes to this project will be documented in this file. ### New features -* Add `TDigestMut::try_merge` and `TDigestMut::merge_many` for checked and batch merging, and add `TDigestMut::quantiles` and `TDigest::quantiles` for querying several ranks in one centroid scan. +* Implement `FromIterator` for batch construction, and add `TDigestMut::quantiles` and `TDigest::quantiles` for querying several ranks in one centroid scan. ### Performance improvements -* T-Digest batch merge avoids recompressing each partial sketch, and batch quantile queries reuse one traversal for ranks supplied in nondecreasing order. +* T-Digest batch construction from owned partial sketches avoids recompressing intermediate results, and batch quantile queries reuse one traversal for ranks supplied in nondecreasing order. ### Bug fixes * T-Digest merging now uses the smaller `k` when sketches have different compression parameters, preserving the size bound of the coarser input. -* T-Digest deserialization now rejects unknown or conflicting flags, reversed extrema, out-of-range values, unsorted centroids, and non-empty images without stored values. ## v0.5.0 diff --git a/benchmarks/tdigest/merge.rs b/benchmarks/tdigest/merge.rs index 7f67b2c2..a0a508d1 100644 --- a/benchmarks/tdigest/merge.rs +++ b/benchmarks/tdigest/merge.rs @@ -103,7 +103,7 @@ fn partials(bencher: Bencher) { } #[divan::bench] -fn partials_batch(bencher: Bencher) { +fn partials_from_iter(bencher: Bencher) { let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL) .into_iter() .map(|mut digest| { @@ -114,11 +114,10 @@ fn partials_batch(bencher: Bencher) { bencher .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL)) - .bench_local(|| { - let mut merged = TDigestMut::new(DEFAULT_DIGEST_K).unwrap(); - merged.merge_many(black_box(&partials)).unwrap(); - black_box(merged) - }); + // Construct fresh owned inputs outside the measurement, as a real caller transfers + // ownership rather than cloning solely for `FromIterator`. + .with_inputs(|| partials.clone()) + .bench_local_values(|partials| black_box(partials).into_iter().collect::()); } #[divan::bench(args = [SMALL_ROWS_PER_PARTIAL, ROWS_PER_PARTIAL])] diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 2fe01c2f..85326009 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -143,16 +143,6 @@ impl TDigestBuffer { self.centroids } - fn append_unmerged_to(&self, output: &mut Vec) { - let compressed_prefix_len = self.compressed_prefix_len(); - output.extend_from_slice(&self.centroids[compressed_prefix_len..]); - } - - fn append_compressed_to(&self, output: &mut Vec) { - let compressed_prefix_len = self.compressed_prefix_len(); - output.extend_from_slice(&self.centroids[..compressed_prefix_len]); - } - fn compressed_centroids(&self) -> &[Centroid] { assert_eq!( self.unmerged_tail_len, 0, @@ -338,8 +328,7 @@ impl TDigestMut { /// /// # Panics /// - /// Panics if the combined total weight exceeds `u64::MAX`. Use [`try_merge`](Self::try_merge) - /// to handle that condition. + /// Panics if the combined total weight exceeds `u64::MAX`. /// /// # Examples /// @@ -354,75 +343,49 @@ impl TDigestMut { /// assert_eq!(left.total_weight(), 2); /// ``` pub fn merge(&mut self, other: &TDigestMut) { - self.try_merge(other) - .expect("combined t-digest weight exceeds u64::MAX"); - } - - /// Merges the given t-digest into this one. - /// - /// If the sketches have different `k` values, the merged sketch uses the smaller value because - /// centroids produced with a smaller `k` cannot be split to satisfy the tighter size bound. - /// - /// # Errors - /// - /// Returns an error if the combined total weight exceeds `u64::MAX`. - pub fn try_merge(&mut self, other: &TDigestMut) -> Result<(), Error> { if other.is_empty() { - return Ok(()); + return; } let self_unmerged_weight = self.buffer.unmerged_len() as u64; - let additional_weight = - checked_merged_weight_sum(self_unmerged_weight, other.total_weight())?; - checked_merged_weight_sum(self.compressed_weight, additional_weight)?; + let additional_weight = self_unmerged_weight + .checked_add(other.total_weight()) + .expect("combined t-digest weight exceeds u64::MAX"); + self.compressed_weight + .checked_add(additional_weight) + .expect("combined t-digest weight exceeds u64::MAX"); let centroids = std::mem::take(&mut self.buffer).into_merged_centroids(&other.buffer); self.k = self.k.min(other.k); self.min = self.min.min(other.min); self.max = self.max.max(other.max); self.compress_sorted_centroids(centroids, additional_weight); - Ok(()) } - /// Merges several t-digests into this one with one compression pass. - /// - /// Empty inputs are ignored. The merged sketch uses the smallest `k` among the non-empty - /// inputs and this sketch. This method temporarily retains every input centroid to avoid - /// recompressing intermediate results. - /// - /// # Errors - /// - /// Returns an error if the combined total weight exceeds `u64::MAX`. - pub fn merge_many<'a>( - &mut self, - others: impl IntoIterator, - ) -> Result<(), Error> { - let others = others - .into_iter() - .filter(|other| !other.is_empty()) - .collect::>(); - if others.is_empty() { - return Ok(()); - } - + fn merge_owned(&mut self, mut others: Vec) { + debug_assert!(!self.is_empty()); + debug_assert!(!others.is_empty()); + debug_assert!(others.iter().all(|other| !other.is_empty())); let mut additional_weight = self.buffer.unmerged_len() as u64; let mut num_centroids = self.buffer.len(); let mut k = self.k; let mut min = self.min; let mut max = self.max; for other in &others { - additional_weight = checked_merged_weight_sum(additional_weight, other.total_weight())?; + additional_weight = additional_weight + .checked_add(other.total_weight()) + .expect("combined t-digest weight exceeds u64::MAX"); num_centroids = num_centroids .checked_add(other.buffer.len()) - .ok_or_else(|| { - Error::invalid_argument("combined t-digest centroid count exceeds usize::MAX") - })?; + .expect("combined t-digest centroid count exceeds usize::MAX"); k = k.min(other.k); min = min.min(other.min); max = max.max(other.max); } - checked_merged_weight_sum(self.compressed_weight, additional_weight)?; + self.compressed_weight + .checked_add(additional_weight) + .expect("combined t-digest weight exceeds u64::MAX"); - let own_buffer = std::mem::take(&mut self.buffer); + let mut own_buffer = std::mem::take(&mut self.buffer); let centroids = if own_buffer.unmerged_len() == 0 && others.iter().all(|other| other.buffer.unmerged_len() == 0) { @@ -433,13 +396,15 @@ impl TDigestMut { } else { let mut centroids = Vec::with_capacity(num_centroids); // Stable sorting keeps raw values before existing summaries when their means are equal. - own_buffer.append_unmerged_to(&mut centroids); - for other in &others { - other.buffer.append_unmerged_to(&mut centroids); + let compressed_prefix_len = own_buffer.compressed_prefix_len(); + centroids.extend(own_buffer.centroids.drain(compressed_prefix_len..)); + for other in &mut others { + let compressed_prefix_len = other.buffer.compressed_prefix_len(); + centroids.extend(other.buffer.centroids.drain(compressed_prefix_len..)); } - own_buffer.append_compressed_to(&mut centroids); - for other in &others { - other.buffer.append_compressed_to(&mut centroids); + centroids.extend(own_buffer.centroids); + for other in others { + centroids.extend(other.buffer.centroids); } centroids.sort_by(centroid_cmp); centroids @@ -449,7 +414,6 @@ impl TDigestMut { self.min = min; self.max = max; self.compress_sorted_centroids(centroids, additional_weight); - Ok(()) } /// Converts this mutable t-digest into an immutable one. @@ -718,20 +682,8 @@ impl TDigestMut { return Err(Error::deserial(format!("k must be at least 10, got {k}"))); } let flags = cursor.read_u8().map_err(insufficient_data("flags"))?; - let known_flags = FLAGS_IS_EMPTY | FLAGS_IS_SINGLE_VALUE | FLAGS_REVERSE_MERGE; - if flags & !known_flags != 0 { - return Err(Error::deserial(format!( - "malformed data: unknown TDigest flags 0x{:02x}", - flags & !known_flags - ))); - } let is_empty = (flags & FLAGS_IS_EMPTY) != 0; let is_single_value = (flags & FLAGS_IS_SINGLE_VALUE) != 0; - if is_empty && is_single_value { - return Err(Error::deserial( - "malformed data: empty and single-value flags are mutually exclusive", - )); - } let expected_preamble_longs = if is_empty || is_single_value { PREAMBLE_LONGS_EMPTY_OR_SINGLE } else { @@ -790,7 +742,10 @@ impl TDigestMut { cursor.read_f64_le().map_err(insufficient_data("max"))?, ) }; - check_extrema(min, max, "TDigest")?; + check_non_nan(min, "min")?; + check_non_nan(max, "max")?; + check_finite(min, "min")?; + check_finite(max, "max")?; let (centroid_bytes, buffered_value_bytes) = if is_f32 { (size_of::() + size_of::(), size_of::()) } else { @@ -818,15 +773,8 @@ impl TDigestMut { let stored_centroids = num_centroids.checked_add(num_buffered).ok_or_else(|| { Error::deserial("num_centroids and num_buffered exceed the supported size") })?; - if stored_centroids == 0 { - return Err(Error::deserial( - "malformed data: non-empty TDigest must contain a centroid or buffered value", - )); - } let mut centroids = Vec::with_capacity(stored_centroids); let mut compressed_weight = 0u64; - let mut previous_mean = min; - let mut centroid_means_valid = true; for bytes in centroid_payload.chunks_exact(centroid_bytes) { let (mean, weight) = if is_f32 { ( @@ -839,36 +787,26 @@ impl TDigestMut { u64::from_le_bytes(bytes[8..].try_into().unwrap()), ) }; - centroid_means_valid &= mean.is_finite() & (mean >= previous_mean) & (mean <= max); - previous_mean = mean; + check_non_nan(mean, "centroid mean")?; + check_finite(mean, "centroid")?; let weight = check_nonzero(weight, "centroid weight")?; compressed_weight = checked_weight_sum(compressed_weight, weight.get())?; centroids.push(Centroid { mean, weight }); } - if !centroid_means_valid { - return Err(Error::deserial( - "malformed data: centroid means must be finite, within extrema, and nondecreasing", - )); - } checked_weight_sum(compressed_weight, num_buffered as u64)?; - let mut buffered_values_valid = true; for bytes in buffered_payload.chunks_exact(buffered_value_bytes) { let value = if is_f32 { f32::from_le_bytes(bytes.try_into().unwrap()) as f64 } else { f64::from_le_bytes(bytes.try_into().unwrap()) }; - buffered_values_valid &= value.is_finite() & (value >= min) & (value <= max); + check_non_nan(value, "buffered_value mean")?; + check_finite(value, "buffered_value mean")?; centroids.push(Centroid { mean: value, weight: DEFAULT_WEIGHT, }); } - if !buffered_values_valid { - return Err(Error::deserial( - "malformed data: buffered values must be finite and within extrema", - )); - } Ok(TDigestMut::make( k, reverse_merge, @@ -897,7 +835,10 @@ impl TDigestMut { // 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")?; + check_non_nan(min, "min in compat double format")?; + check_non_nan(max, "max in compat double format")?; + check_finite(min, "min in compat double format")?; + check_finite(max, "max in compat double format")?; let k = cursor.read_f64_be().map_err(make_error("k"))? as u16; if k < 10 { return Err(Error::deserial(format!( @@ -906,31 +847,18 @@ impl TDigestMut { } 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; + check_non_nan(mean, "centroid mean in compat double format")?; + check_finite(mean, "centroid mean in compat double format")?; 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", - )); - } Ok(TDigestMut::make( k, false, @@ -948,7 +876,10 @@ impl TDigestMut { // 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")?; + check_non_nan(min, "min in compat float format")?; + check_non_nan(max, "max in compat float format")?; + check_finite(min, "min in compat float format")?; + check_finite(max, "max in compat float format")?; let k = cursor.read_f32_be().map_err(make_error("k"))? as u16; if k < 10 { return Err(Error::deserial(format!( @@ -960,31 +891,18 @@ impl TDigestMut { 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; + check_non_nan(mean, "centroid mean in compat float format")?; + check_finite(mean, "centroid mean in compat float format")?; 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", - )); - } Ok(TDigestMut::make( k, false, @@ -1091,6 +1009,32 @@ impl TDigestMut { } } +/// Collects owned t-digests into one result with a single compression pass. +/// +/// Empty inputs are ignored. The result uses the smallest `k` among the non-empty inputs. A single +/// non-empty input is returned unchanged. Collecting consumes each digest without cloning its +/// buffer. Unlike repeated [`TDigestMut::merge`] calls, it temporarily retains all input centroids +/// so it can avoid recompressing intermediate results. Use repeated `merge` calls when inputs must +/// be processed with bounded additional memory. +/// +/// # Panics +/// +/// Panics if the combined total weight exceeds `u64::MAX` or the combined centroid count exceeds +/// `usize::MAX`. +impl FromIterator for TDigestMut { + fn from_iter>(iter: T) -> Self { + let mut digests = iter.into_iter().filter(|digest| !digest.is_empty()); + let Some(mut merged) = digests.next() else { + return TDigestMut::default(); + }; + let others = digests.collect::>(); + if !others.is_empty() { + merged.merge_owned(others); + } + merged + } +} + fn serialize_compressed( k: u16, reverse_merge: bool, @@ -1903,21 +1847,6 @@ fn check_finite(value: f64, tag: &'static str) -> Result<(), Error> { Ok(()) } -#[inline] -fn check_extrema(min: f64, max: f64, format: &'static str) -> Result<(), Error> { - if !min.is_finite() || !max.is_finite() { - return Err(Error::deserial(format!( - "malformed data: {format} extrema must be finite" - ))); - } - if min > max { - return Err(Error::deserial(format!( - "malformed data: {format} min {min} exceeds max {max}" - ))); - } - Ok(()) -} - fn check_nonzero(value: u64, tag: &'static str) -> Result { NonZeroU64::new(value) .ok_or_else(|| Error::deserial(format!("malformed data: {tag} cannot be zero"))) @@ -1945,12 +1874,6 @@ fn checked_weight_sum(total_weight: u64, weight: u64) -> Result { .ok_or_else(|| Error::deserial("malformed data: total weight overflow")) } -fn checked_merged_weight_sum(total_weight: u64, weight: u64) -> Result { - total_weight - .checked_add(weight) - .ok_or_else(|| Error::invalid_argument("combined t-digest weight exceeds u64::MAX")) -} - /// Generates cluster sizes proportional to `q*(1-q)`. /// /// The use of a normalizing function results in a strictly bounded number of clusters no matter diff --git a/tests-integration/tests/serde_tests/tdigest.rs b/tests-integration/tests/serde_tests/tdigest.rs index 1b871ef8..c5bd0fe1 100644 --- a/tests-integration/tests/serde_tests/tdigest.rs +++ b/tests-integration/tests/serde_tests/tdigest.rs @@ -350,60 +350,6 @@ fn test_updates_normalize_overfull_deserialized_mixed_buffer() { assert_eq!(roundtrip.max_value(), Some(1_000.0)); } -fn serialized_two_value_digest() -> Vec { - let mut tdigest = TDigestMut::new(100).unwrap(); - tdigest.update(0.0); - tdigest.update(1.0); - tdigest.serialize() -} - -fn assert_invalid_tdigest(bytes: &[u8]) { - let error = TDigestMut::deserialize(bytes).unwrap_err(); - assert_eq!(error.kind(), datasketches::error::ErrorKind::InvalidData); -} - -#[test] -fn test_deserialize_rejects_unknown_or_conflicting_flags() { - let mut unknown = serialized_two_value_digest(); - unknown[5] |= 0x80; - assert_invalid_tdigest(&unknown); - - let mut empty = TDigestMut::new(100).unwrap(); - let mut conflicting = empty.serialize(); - conflicting[5] |= 1 << 1; - assert_invalid_tdigest(&conflicting); -} - -#[test] -fn test_deserialize_rejects_invalid_extrema_and_centroid_ranges() { - let mut reversed_extrema = serialized_two_value_digest(); - reversed_extrema[16..24].copy_from_slice(&2_f64.to_le_bytes()); - assert_invalid_tdigest(&reversed_extrema); - - let mut centroid_outside_extrema = serialized_two_value_digest(); - centroid_outside_extrema[32..40].copy_from_slice(&(-1_f64).to_le_bytes()); - assert_invalid_tdigest(¢roid_outside_extrema); - - let mut buffered_outside_extrema = serialized_two_value_digest(); - buffered_outside_extrema[8..12].copy_from_slice(&1_u32.to_le_bytes()); - buffered_outside_extrema[12..16].copy_from_slice(&1_u32.to_le_bytes()); - buffered_outside_extrema[48..56].copy_from_slice(&2_f64.to_le_bytes()); - assert_invalid_tdigest(&buffered_outside_extrema); -} - -#[test] -fn test_deserialize_rejects_unsorted_or_missing_centroids() { - let mut unsorted = serialized_two_value_digest(); - unsorted[32..40].copy_from_slice(&1_f64.to_le_bytes()); - unsorted[48..56].copy_from_slice(&0_f64.to_le_bytes()); - assert_invalid_tdigest(&unsorted); - - let mut missing = serialized_two_value_digest(); - missing[8..12].copy_from_slice(&0_u32.to_le_bytes()); - missing[12..16].copy_from_slice(&0_u32.to_le_bytes()); - assert_invalid_tdigest(&missing); -} - #[test] fn test_deserialize_rejects_truncated_large_payload_before_allocation() { let mut tdigest = TDigestMut::new(10).unwrap(); diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 707022ab..84233b1f 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -284,7 +284,7 @@ fn test_mixed_k_merge_uses_smaller_k() { right.update((value + 1_000) as f64); } - left.try_merge(&right).unwrap(); + left.merge(&right); assert_eq!(left.k(), 50); assert_eq!(left.total_weight(), 2_000); @@ -293,8 +293,7 @@ fn test_mixed_k_merge_uses_smaller_k() { } #[test] -fn test_merge_many_uses_one_result_with_the_smallest_nonempty_k() { - let mut merged = TDigestMut::new(200).unwrap(); +fn test_from_iter_uses_one_result_with_the_smallest_nonempty_k() { let mut first = TDigestMut::new(100).unwrap(); let mut second = TDigestMut::new(50).unwrap(); let empty = TDigestMut::new(10).unwrap(); @@ -305,7 +304,7 @@ fn test_merge_many_uses_one_result_with_the_smallest_nonempty_k() { let _ = first.rank(0.0); let _ = second.rank(0.0); - merged.merge_many([&first, &empty, &second]).unwrap(); + let mut merged = [first, empty, second].into_iter().collect::(); assert_eq!(merged.k(), 50); assert_eq!(merged.total_weight(), 2_000); @@ -314,6 +313,24 @@ fn test_merge_many_uses_one_result_with_the_smallest_nonempty_k() { assert_quantiles_are_nondecreasing(&mut merged); } +#[test] +fn test_from_iter_handles_empty_and_single_input_without_recompression() { + let empty = std::iter::empty::().collect::(); + assert!(empty.is_empty()); + + let mut input = TDigestMut::new(50).unwrap(); + for value in 0..1_000 { + input.update(value as f64); + } + let serialized = input.serialize(); + let mut collected = [TDigestMut::new(10).unwrap(), input] + .into_iter() + .collect::(); + + assert_eq!(collected.k(), 50); + assert_eq!(collected.serialize(), serialized); +} + #[test] fn test_invalid_inputs() { let n = 100; @@ -460,20 +477,12 @@ fn test_quantile_handles_two_sample_last_centroid() { } #[test] -fn test_try_merge_reports_total_weight_overflow_without_mutating_receiver() { +#[should_panic(expected = "combined t-digest weight exceeds u64::MAX")] +fn test_merge_panics_on_total_weight_overflow() { let mut left = deserialize_with_centroids(100, 0.0, 0.0, &[(0.0, u64::MAX)]); let right = deserialize_with_centroids(50, 1.0, 1.0, &[(1.0, 1)]); - let error = left.try_merge(&right).unwrap_err(); - - assert_eq!( - error.kind(), - datasketches::error::ErrorKind::InvalidArgument - ); - assert_eq!(left.k(), 100); - assert_eq!(left.total_weight(), u64::MAX); - assert_eq!(left.min_value(), Some(0.0)); - assert_eq!(left.max_value(), Some(0.0)); + left.merge(&right); } #[test] From ebcd8486d437fd9783b19e7050dcfab3acd4485b Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 15:50:16 +0800 Subject: [PATCH 4/6] refactor(tdigest): clarify batch aggregation internals --- datasketches/src/tdigest/sketch.rs | 136 +++++++++++------- .../tests/tdigest_test/sketch.rs | 23 +++ 2 files changed, 108 insertions(+), 51 deletions(-) diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 5d814f7c..55c6aa79 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -120,11 +120,10 @@ impl TDigestBuffer { /// Combines this buffer with a non-empty borrowed buffer in stable mean order. fn into_merged_centroids(mut self, other: &TDigestBuffer) -> Vec { debug_assert!(!other.is_empty(), "an empty right-hand buffer is a no-op"); - if self.unmerged_tail_len == 0 - && other.unmerged_tail_len == 0 - && centroids_are_sorted(&self.centroids) - && centroids_are_sorted(&other.centroids) - { + if self.unmerged_tail_len == 0 && other.unmerged_tail_len == 0 { + // Compression and deserialization both establish this invariant. + debug_assert!(centroids_are_sorted(&self.centroids)); + debug_assert!(centroids_are_sorted(&other.centroids)); merge_sorted_centroids(&mut self.centroids, &other.centroids); return self.centroids; } @@ -324,7 +323,7 @@ impl TDigestMut { /// Merges the given t-digest into this one. /// /// If the sketches have different `k` values, the merged sketch uses the smaller value because - /// centroids produced with a smaller `k` cannot be split to satisfy the tighter size bound. + /// merging cannot recover detail already discarded by the lower-`k` sketch. /// /// # Panics /// @@ -361,10 +360,13 @@ impl TDigestMut { self.compress_sorted_centroids(centroids, additional_weight); } - fn merge_owned(&mut self, mut others: Vec) { + fn merge_owned_batch(&mut self, mut others: Vec) { debug_assert!(!self.is_empty()); debug_assert!(!others.is_empty()); debug_assert!(others.iter().all(|other| !other.is_empty())); + + // The receiver's compressed prefix is already included in `compressed_weight`; its + // unmerged tail and every centroid in the other digests are new weight for compression. let mut additional_weight = self.buffer.unmerged_len() as u64; let mut num_centroids = self.buffer.len(); let mut k = self.k; @@ -1058,15 +1060,29 @@ impl TDigestMut { /// Collects owned t-digests into one result with a single compression pass. /// /// Empty inputs are ignored. The result uses the smallest `k` among the non-empty inputs. A single -/// non-empty input is returned unchanged. Collecting consumes each digest without cloning its -/// buffer. Unlike repeated [`TDigestMut::merge`] calls, it temporarily retains all input centroids -/// so it can avoid recompressing intermediate results. Use repeated `merge` calls when inputs must -/// be processed with bounded additional memory. +/// non-empty input is returned unchanged. Collecting consumes each digest, so callers do not need +/// to clone inputs. Unlike repeated [`TDigestMut::merge`] calls, it temporarily retains all input +/// centroids so it can avoid recompressing intermediate results. Use repeated `merge` calls when +/// inputs must be processed with bounded additional memory. /// /// # Panics /// /// Panics if the combined total weight exceeds `u64::MAX` or the combined centroid count exceeds /// `usize::MAX`. +/// +/// # Examples +/// +/// ``` +/// use datasketches::tdigest::TDigestMut; +/// +/// let partials = [1.0, 2.0].map(|value| { +/// let mut digest = TDigestMut::new(100).unwrap(); +/// digest.update(value); +/// digest +/// }); +/// let merged = partials.into_iter().collect::(); +/// assert_eq!(merged.total_weight(), 2); +/// ``` impl FromIterator for TDigestMut { fn from_iter>(iter: T) -> Self { let mut digests = iter.into_iter().filter(|digest| !digest.is_empty()); @@ -1075,7 +1091,7 @@ impl FromIterator for TDigestMut { }; let others = digests.collect::>(); if !others.is_empty() { - merged.merge_owned(others); + merged.merge_owned_batch(others); } merged } @@ -1577,9 +1593,7 @@ impl TDigestView<'_> { return None; } - let mut centroid_index = 0; - let mut weight_so_far = self.centroids[0].weight() / 2.; - Some(self.quantile_with_cursor(rank, &mut centroid_index, &mut weight_so_far)) + Some(QuantileCursor::new(self).quantile(rank)) } fn quantiles(&self, ranks: &[f64]) -> Option> { @@ -1593,31 +1607,50 @@ impl TDigestView<'_> { } let mut quantiles = vec![0.; ranks.len()]; - let mut centroid_index = 0; - let mut weight_so_far = self.centroids[0].weight() / 2.; + let mut cursor = QuantileCursor::new(self); if ranks.windows(2).all(|pair| pair[0] <= pair[1]) { for (index, &rank) in ranks.iter().enumerate() { - quantiles[index] = - self.quantile_with_cursor(rank, &mut centroid_index, &mut weight_so_far); + quantiles[index] = cursor.quantile(rank); } return Some(quantiles); } + // The cursor only moves forward. Sort indices so queries become monotonic without changing + // the caller's output order. let mut rank_order = (0..ranks.len()).collect::>(); rank_order.sort_by(|&left, &right| ranks[left].total_cmp(&ranks[right])); for index in rank_order { - quantiles[index] = - self.quantile_with_cursor(ranks[index], &mut centroid_index, &mut weight_so_far); + quantiles[index] = cursor.quantile(ranks[index]); } Some(quantiles) } +} - fn quantile_with_cursor( - &self, - rank: f64, - centroid_index: &mut usize, - weight_so_far: &mut f64, - ) -> f64 { +/// Incrementally answers quantile queries supplied in nondecreasing rank order. +struct QuantileCursor<'a> { + min: f64, + max: f64, + centroids: &'a [Centroid], + centroids_weight: f64, + centroid_index: usize, + weight_so_far: f64, +} + +impl<'a> QuantileCursor<'a> { + fn new(view: &TDigestView<'a>) -> Self { + debug_assert!(!view.centroids.is_empty()); + + QuantileCursor { + min: view.min, + max: view.max, + centroids: view.centroids, + centroids_weight: view.centroids_weight as f64, + centroid_index: 0, + weight_so_far: view.centroids[0].weight() / 2., + } + } + + fn quantile(&mut self, rank: f64) -> f64 { debug_assert!(!self.centroids.is_empty()); if self.centroids.len() == 1 { @@ -1625,7 +1658,7 @@ impl TDigestView<'_> { } // at least 2 centroids - let centroids_weight = self.centroids_weight as f64; + let centroids_weight = self.centroids_weight; let num_centroids = self.centroids.len(); let weight = rank * centroids_weight; if weight < 1. { @@ -1651,39 +1684,39 @@ impl TDigestView<'_> { } // interpolate between extremes - while *centroid_index < num_centroids - 1 { - let dw = (self.centroids[*centroid_index].weight() - + self.centroids[*centroid_index + 1].weight()) + while self.centroid_index < num_centroids - 1 { + let dw = (self.centroids[self.centroid_index].weight() + + self.centroids[self.centroid_index + 1].weight()) / 2.; - if *weight_so_far + dw > weight { + if self.weight_so_far + dw > weight { // the target weight is between centroids i and i+1 let mut left_weight = 0.; - if self.centroids[*centroid_index].weight.get() == 1 { - if weight - *weight_so_far < 0.5 { - return self.centroids[*centroid_index].mean; + if self.centroids[self.centroid_index].weight.get() == 1 { + if weight - self.weight_so_far < 0.5 { + return self.centroids[self.centroid_index].mean; } left_weight = 0.5; } let mut right_weight = 0.; - if self.centroids[*centroid_index + 1].weight.get() == 1 { - if *weight_so_far + dw - weight <= 0.5 { - return self.centroids[*centroid_index + 1].mean; + if self.centroids[self.centroid_index + 1].weight.get() == 1 { + if self.weight_so_far + dw - weight <= 0.5 { + return self.centroids[self.centroid_index + 1].mean; } right_weight = 0.5; } // Each centroid is weighted by the distance from the target to the *other* // centroid, so the estimate approaches the nearer one. - let distance_from_left = weight - *weight_so_far - left_weight; - let distance_to_right = *weight_so_far + dw - weight - right_weight; + let distance_from_left = weight - self.weight_so_far - left_weight; + let distance_to_right = self.weight_so_far + dw - weight - right_weight; return weighted_average( - self.centroids[*centroid_index].mean, + self.centroids[self.centroid_index].mean, distance_to_right, - self.centroids[*centroid_index + 1].mean, + self.centroids[self.centroid_index + 1].mean, distance_from_left, ); } - *weight_so_far += dw; - *centroid_index += 1; + self.weight_so_far += dw; + self.centroid_index += 1; } let w1 = weight - (centroids_weight) - ((self.centroids[num_centroids - 1].weight()) / 2.); @@ -1726,27 +1759,27 @@ fn centroids_are_sorted(centroids: &[Centroid]) -> bool { } #[derive(Clone, Copy)] -struct CentroidCursor { +struct CentroidMergeCursor { centroid: Centroid, source_index: usize, centroid_index: usize, } -impl PartialEq for CentroidCursor { +impl PartialEq for CentroidMergeCursor { fn eq(&self, other: &Self) -> bool { self.source_index == other.source_index && self.centroid_index == other.centroid_index } } -impl Eq for CentroidCursor {} +impl Eq for CentroidMergeCursor {} -impl PartialOrd for CentroidCursor { +impl PartialOrd for CentroidMergeCursor { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) } } -impl Ord for CentroidCursor { +impl Ord for CentroidMergeCursor { fn cmp(&self, other: &Self) -> Ordering { // Reverse every key because BinaryHeap is a max-heap. Source order breaks equal-mean ties // in the same way as concatenating the inputs and applying a stable sort. @@ -1756,6 +1789,7 @@ impl Ord for CentroidCursor { } } +/// Stably merges sorted centroid slices without sorting their combined contents again. fn merge_sorted_centroid_slices(sources: &[&[Centroid]], num_centroids: usize) -> Vec { debug_assert_eq!( sources.iter().map(|source| source.len()).sum::(), @@ -1766,7 +1800,7 @@ fn merge_sorted_centroid_slices(sources: &[&[Centroid]], num_centroids: usize) - let mut heap = BinaryHeap::with_capacity(sources.len()); for (source_index, source) in sources.iter().enumerate() { if let Some(¢roid) = source.first() { - heap.push(CentroidCursor { + heap.push(CentroidMergeCursor { centroid, source_index, centroid_index: 0, @@ -1779,7 +1813,7 @@ fn merge_sorted_centroid_slices(sources: &[&[Centroid]], num_centroids: usize) - centroids.push(cursor.centroid); let centroid_index = cursor.centroid_index + 1; if let Some(¢roid) = sources[cursor.source_index].get(centroid_index) { - heap.push(CentroidCursor { + heap.push(CentroidMergeCursor { centroid, source_index: cursor.source_index, centroid_index, diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 84233b1f..49e00068 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -313,6 +313,29 @@ fn test_from_iter_uses_one_result_with_the_smallest_nonempty_k() { assert_quantiles_are_nondecreasing(&mut merged); } +#[test] +fn test_from_iter_matches_single_digest_for_uncompressed_inputs() { + let values = [3.0, 1.0, 2.0, 2.0, 5.0, 4.0]; + let partials = values + .chunks(3) + .map(|values| { + let mut digest = TDigestMut::new(100).unwrap(); + for &value in values { + digest.update(value); + } + digest + }) + .collect::>(); + let mut merged = partials.into_iter().collect::(); + + let mut expected = TDigestMut::new(100).unwrap(); + for value in values { + expected.update(value); + } + + assert_eq!(merged.serialize(), expected.serialize()); +} + #[test] fn test_from_iter_handles_empty_and_single_input_without_recompression() { let empty = std::iter::empty::().collect::(); From 8a05f204eba7d4874571dee9097fd979180859be Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 15:56:01 +0800 Subject: [PATCH 5/6] docs: group performance changes under improvements --- CHANGELOG.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ce01cd30..645b6ad9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,9 +16,6 @@ All significant changes to this project will be documented in this file. ### Improvements * Improve truncated-input diagnostics across sketch deserializers. - -### Performance improvements - * T-Digest batch construction from owned partial sketches avoids recompressing intermediate results, and batch quantile queries reuse one traversal for ranks supplied in nondecreasing order. ### Bug fixes From b363b7d70159367273cebe5a9229a956dafb682f Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 19:54:20 +0800 Subject: [PATCH 6/6] refactor(tdigest): stream owned batch merging --- benchmarks/tdigest/merge.rs | 20 +- datasketches/src/tdigest/sketch.rs | 327 ++++++++++++------ .../tests/tdigest_test/sketch.rs | 4 +- 3 files changed, 235 insertions(+), 116 deletions(-) diff --git a/benchmarks/tdigest/merge.rs b/benchmarks/tdigest/merge.rs index a0a508d1..6c92ed8c 100644 --- a/benchmarks/tdigest/merge.rs +++ b/benchmarks/tdigest/merge.rs @@ -65,7 +65,7 @@ fn small_partials(bencher: Bencher) { let partials = partial_digests(64, SMALL_ROWS_PER_PARTIAL) .into_iter() .map(|mut digest| { - black_box(digest.rank(0.0)); + black_box(digest.quantile(0.5)); digest }) .collect::>(); @@ -86,7 +86,7 @@ fn partials(bencher: Bencher) { let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL) .into_iter() .map(|mut digest| { - black_box(digest.rank(0.0)); + black_box(digest.quantile(0.5)); digest }) .collect::>(); @@ -107,7 +107,7 @@ fn partials_from_iter(bencher: Bencher) { let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL) .into_iter() .map(|mut digest| { - black_box(digest.rank(0.0)); + black_box(digest.quantile(0.5)); digest }) .collect::>(); @@ -120,6 +120,20 @@ fn partials_from_iter(bencher: Bencher) { .bench_local_values(|partials| black_box(partials).into_iter().collect::()); } +#[divan::bench] +fn uncompressed_partials_from_iter(bencher: Bencher) { + let values = values(64 * ROWS_PER_PARTIAL); + let partials = values + .chunks_exact(ROWS_PER_PARTIAL) + .map(build_mut_digest) + .collect::>(); + + bencher + .counter(ItemsCount::new(values.len())) + .with_inputs(|| partials.clone()) + .bench_local_values(|partials| black_box(partials).into_iter().collect::()); +} + #[divan::bench(args = [SMALL_ROWS_PER_PARTIAL, ROWS_PER_PARTIAL])] fn serialized_partials(bencher: Bencher, rows_per_partial: usize) { let partials = serialized_partial_digests(64, rows_per_partial); diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 55c6aa79..c96649f9 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::collections::BinaryHeap; use std::convert::identity; use std::num::NonZeroU64; @@ -360,62 +359,51 @@ impl TDigestMut { self.compress_sorted_centroids(centroids, additional_weight); } - fn merge_owned_batch(&mut self, mut others: Vec) { - debug_assert!(!self.is_empty()); - debug_assert!(!others.is_empty()); - debug_assert!(others.iter().all(|other| !other.is_empty())); - - // The receiver's compressed prefix is already included in `compressed_weight`; its - // unmerged tail and every centroid in the other digests are new weight for compression. - let mut additional_weight = self.buffer.unmerged_len() as u64; - let mut num_centroids = self.buffer.len(); - let mut k = self.k; - let mut min = self.min; - let mut max = self.max; - for other in &others { - additional_weight = additional_weight - .checked_add(other.total_weight()) + fn from_owned_digests(mut digests: Vec) -> Self { + debug_assert!(digests.len() >= 2); + debug_assert!(digests.iter().all(|digest| !digest.is_empty())); + + let mut total_weight = 0u64; + let mut num_centroids = 0usize; + let mut k = digests[0].k; + let reverse_merge = digests[0].reverse_merge; + let mut min = digests[0].min; + let mut max = digests[0].max; + for digest in &digests { + total_weight = total_weight + .checked_add(digest.total_weight()) .expect("combined t-digest weight exceeds u64::MAX"); num_centroids = num_centroids - .checked_add(other.buffer.len()) + .checked_add(digest.buffer.len()) .expect("combined t-digest centroid count exceeds usize::MAX"); - k = k.min(other.k); - min = min.min(other.min); - max = max.max(other.max); - } - self.compressed_weight - .checked_add(additional_weight) - .expect("combined t-digest weight exceeds u64::MAX"); - - let mut own_buffer = std::mem::take(&mut self.buffer); - let centroids = if own_buffer.unmerged_len() == 0 - && others.iter().all(|other| other.buffer.unmerged_len() == 0) - { - let sources = std::iter::once(own_buffer.centroids.as_slice()) - .chain(others.iter().map(|other| other.buffer.centroids.as_slice())) - .collect::>(); - merge_sorted_centroid_slices(&sources, num_centroids) + k = k.min(digest.k); + min = min.min(digest.min); + max = max.max(digest.max); + } + + let mut merged = TDigestMut::make(k, reverse_merge, min, max, TDigestBuffer::default(), 0); + let all_compressed = digests + .iter() + .all(|digest| digest.buffer.unmerged_len() == 0); + if all_compressed { + // Compressed buffers are sorted runs, so feed their k-way merge directly into + // compression instead of materializing every input centroid in another vector. + let centroids = KWayMerge::new(&digests, reverse_merge); + merged.compress_merged_centroids(centroids, num_centroids, total_weight); } else { + // Raw tails are unsorted. Put them before the compressed prefixes so the stable sort + // preserves the same equal-mean tie order as regular updates followed by summaries. let mut centroids = Vec::with_capacity(num_centroids); - // Stable sorting keeps raw values before existing summaries when their means are equal. - let compressed_prefix_len = own_buffer.compressed_prefix_len(); - centroids.extend(own_buffer.centroids.drain(compressed_prefix_len..)); - for other in &mut others { - let compressed_prefix_len = other.buffer.compressed_prefix_len(); - centroids.extend(other.buffer.centroids.drain(compressed_prefix_len..)); + for digest in &mut digests { + let tail_start = digest.buffer.compressed_prefix_len(); + centroids.extend(digest.buffer.centroids.drain(tail_start..)); } - centroids.extend(own_buffer.centroids); - for other in others { - centroids.extend(other.buffer.centroids); + for digest in digests { + centroids.extend(digest.buffer.centroids); } - centroids.sort_by(centroid_cmp); - centroids - }; - - self.k = k; - self.min = min; - self.max = max; - self.compress_sorted_centroids(centroids, additional_weight); + merged.compress_centroids(centroids, total_weight); + } + merged } /// Converts this mutable t-digest into an immutable one. @@ -1008,15 +996,14 @@ impl TDigestMut { while current < len { let c = centroids[current]; let proposed_weight = centroids[num_centroids - 1].weight() + c.weight(); - let mut add_this = false; - if (current != 1) && (current != (len - 1)) { - let q0 = weight_so_far / compressed_weight; - let q2 = (weight_so_far + proposed_weight) / compressed_weight; - add_this = proposed_weight - <= (compressed_weight - * scale_function::max(q0, normalizer) - .min(scale_function::max(q2, normalizer))); - } + let add_this = should_merge_centroid( + current, + len, + weight_so_far, + proposed_weight, + compressed_weight, + normalizer, + ); if add_this { // merge into existing centroid centroids[num_centroids - 1].add(c); @@ -1040,6 +1027,57 @@ impl TDigestMut { self.buffer = TDigestBuffer::new(centroids, 0); } + fn compress_merged_centroids( + &mut self, + centroids: impl Iterator, + num_input_centroids: usize, + total_weight: u64, + ) { + debug_assert_ne!(num_input_centroids, 0); + let compressed_weight = total_weight as f64; + let normalizer = scale_function::normalizer(2.0 * f64::from(self.k), compressed_weight); + let mut centroids = centroids.enumerate(); + let (_, first) = centroids.next().expect("non-empty centroid stream"); + let capacity = self.target_retained_capacity().min(num_input_centroids); + let mut retained = Vec::with_capacity(capacity); + retained.push(first); + let mut weight_so_far = 0.; + + for (current, centroid) in centroids { + let proposed_weight = retained.last().unwrap().weight() + centroid.weight(); + let add_this = should_merge_centroid( + current, + num_input_centroids, + weight_so_far, + proposed_weight, + compressed_weight, + normalizer, + ); + if add_this { + retained.last_mut().unwrap().add(centroid); + } else { + weight_so_far += retained.last().unwrap().weight(); + retained.push(centroid); + } + } + + debug_assert!(retained.len() <= self.target_centroids()); + if self.reverse_merge { + retained.reverse(); + } + debug_assert_eq!( + retained + .iter() + .map(|centroid| centroid.weight.get()) + .sum::(), + total_weight, + "compressed centroids must preserve total weight" + ); + self.compressed_weight = total_weight; + self.reverse_merge = !self.reverse_merge; + self.buffer = TDigestBuffer::new(retained, 0); + } + fn reduce_retained_capacity(&self, centroids: &mut Vec) { let target_capacity = self.target_retained_capacity().max(centroids.len()); if centroids.capacity() <= target_capacity { @@ -1086,14 +1124,15 @@ impl TDigestMut { impl FromIterator for TDigestMut { fn from_iter>(iter: T) -> Self { let mut digests = iter.into_iter().filter(|digest| !digest.is_empty()); - let Some(mut merged) = digests.next() else { + let Some(first) = digests.next() else { return TDigestMut::default(); }; - let others = digests.collect::>(); - if !others.is_empty() { - merged.merge_owned_batch(others); + let mut owned = digests.collect::>(); + if owned.is_empty() { + return first; } - merged + owned.insert(0, first); + TDigestMut::from_owned_digests(owned) } } @@ -1758,69 +1797,135 @@ fn centroids_are_sorted(centroids: &[Centroid]) -> bool { .all(|pair| centroid_cmp(&pair[0], &pair[1]) != Ordering::Greater) } -#[derive(Clone, Copy)] -struct CentroidMergeCursor { - centroid: Centroid, - source_index: usize, - centroid_index: usize, +fn should_merge_centroid( + current: usize, + len: usize, + weight_so_far: f64, + proposed_weight: f64, + compressed_weight: f64, + normalizer: f64, +) -> bool { + if current == 1 || current == len - 1 { + return false; + } + let q0 = weight_so_far / compressed_weight; + let q2 = (weight_so_far + proposed_weight) / compressed_weight; + proposed_weight + <= compressed_weight + * scale_function::max(q0, normalizer).min(scale_function::max(q2, normalizer)) } -impl PartialEq for CentroidMergeCursor { - fn eq(&self, other: &Self) -> bool { - self.source_index == other.source_index && self.centroid_index == other.centroid_index - } +/// A non-empty sorted run participating in a k-way merge. +struct CentroidRun<'a> { + head: Centroid, + tail: &'a [Centroid], + // Breaks equal-mean ties as if the input runs had been concatenated and stably sorted. + order: usize, } -impl Eq for CentroidMergeCursor {} +impl<'a> CentroidRun<'a> { + fn new(centroids: &'a [Centroid], order: usize, reverse: bool) -> Option { + let (head, tail) = if reverse { + let (head, tail) = centroids.split_last()?; + (*head, tail) + } else { + let (head, tail) = centroids.split_first()?; + (*head, tail) + }; + Some(CentroidRun { head, tail, order }) + } + + fn head(&self) -> Centroid { + self.head + } -impl PartialOrd for CentroidMergeCursor { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + fn advance(&mut self, reverse: bool) -> bool { + if self.tail.is_empty() { + return false; + } + if reverse { + let (head, tail) = self.tail.split_last().unwrap(); + self.head = *head; + self.tail = tail; + } else { + let (head, tail) = self.tail.split_first().unwrap(); + self.head = *head; + self.tail = tail; + } + true } + + fn precedes(&self, other: &Self, reverse: bool) -> bool { + match centroid_cmp(&self.head(), &other.head()) { + Ordering::Less => !reverse, + Ordering::Greater => reverse, + Ordering::Equal if reverse => self.order > other.order, + Ordering::Equal => self.order < other.order, + } + } +} + +/// Lazily merges the sorted centroid buffers of fully compressed digests. +struct KWayMerge<'a> { + heap: Vec>, + reverse: bool, } -impl Ord for CentroidMergeCursor { - fn cmp(&self, other: &Self) -> Ordering { - // Reverse every key because BinaryHeap is a max-heap. Source order breaks equal-mean ties - // in the same way as concatenating the inputs and applying a stable sort. - centroid_cmp(&other.centroid, &self.centroid) - .then_with(|| other.source_index.cmp(&self.source_index)) - .then_with(|| other.centroid_index.cmp(&self.centroid_index)) +impl<'a> KWayMerge<'a> { + fn new(digests: &'a [TDigestMut], reverse: bool) -> Self { + debug_assert!( + digests + .iter() + .all(|digest| digest.buffer.unmerged_len() == 0) + ); + let mut runs = Vec::with_capacity(digests.len()); + runs.extend(digests.iter().enumerate().filter_map(|(order, digest)| { + CentroidRun::new(&digest.buffer.centroids, order, reverse) + })); + for index in (0..runs.len() / 2).rev() { + sift_down_centroid_runs(&mut runs, index, reverse); + } + KWayMerge { + heap: runs, + reverse, + } } } -/// Stably merges sorted centroid slices without sorting their combined contents again. -fn merge_sorted_centroid_slices(sources: &[&[Centroid]], num_centroids: usize) -> Vec { - debug_assert_eq!( - sources.iter().map(|source| source.len()).sum::(), - num_centroids - ); - debug_assert!(sources.iter().all(|source| centroids_are_sorted(source))); - - let mut heap = BinaryHeap::with_capacity(sources.len()); - for (source_index, source) in sources.iter().enumerate() { - if let Some(¢roid) = source.first() { - heap.push(CentroidMergeCursor { - centroid, - source_index, - centroid_index: 0, - }); +impl Iterator for KWayMerge<'_> { + type Item = Centroid; + + fn next(&mut self) -> Option { + if self.heap.is_empty() { + return None; + } + let centroid = self.heap[0].head(); + if !self.heap[0].advance(self.reverse) { + self.heap.swap_remove(0); } + sift_down_centroid_runs(&mut self.heap, 0, self.reverse); + Some(centroid) } +} - let mut centroids = Vec::with_capacity(num_centroids); - while let Some(cursor) = heap.pop() { - centroids.push(cursor.centroid); - let centroid_index = cursor.centroid_index + 1; - if let Some(¢roid) = sources[cursor.source_index].get(centroid_index) { - heap.push(CentroidMergeCursor { - centroid, - source_index: cursor.source_index, - centroid_index, - }); +fn sift_down_centroid_runs(heap: &mut [CentroidRun<'_>], mut position: usize, reverse: bool) { + loop { + let left = (position * 2) + 1; + if left >= heap.len() { + return; } + let right = left + 1; + let next = if right < heap.len() && heap[right].precedes(&heap[left], reverse) { + right + } else { + left + }; + if !heap[next].precedes(&heap[position], reverse) { + return; + } + heap.swap(position, next); + position = next; } - centroids } fn merge_sorted_centroids(left: &mut Vec, right: &[Centroid]) { diff --git a/tests-integration/tests/tdigest_test/sketch.rs b/tests-integration/tests/tdigest_test/sketch.rs index 49e00068..997a2fb8 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -301,8 +301,8 @@ fn test_from_iter_uses_one_result_with_the_smallest_nonempty_k() { first.update(value as f64); second.update((value + 1_000) as f64); } - let _ = first.rank(0.0); - let _ = second.rank(0.0); + let _ = first.quantile(0.5); + let _ = second.quantile(0.5); let mut merged = [first, empty, second].into_iter().collect::();