This is an automated email from the ASF dual-hosted git repository.
tisonkun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datasketches-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 974a6e9 perf(tdigest): single centroid buffer (#210)
974a6e9 is described below
commit 974a6e9b03d2f417398f5a135cda86a5bb90fe42
Author: Yijun Zhao <[email protected]>
AuthorDate: Tue Aug 25 09:13:25 2026 +0800
perf(tdigest): single centroid buffer (#210)
Co-authored-by: tison <[email protected]>
---
datasketches/benches/tdigest.rs | 42 ++-
datasketches/src/tdigest/sketch.rs | 526 ++++++++++++++++++++++--------
datasketches/tests/serde_tests/tdigest.rs | 124 +++++++
datasketches/tests/tdigest_test/sketch.rs | 57 ++++
4 files changed, 607 insertions(+), 142 deletions(-)
diff --git a/datasketches/benches/tdigest.rs b/datasketches/benches/tdigest.rs
index b68bd59..b291770 100644
--- a/datasketches/benches/tdigest.rs
+++ b/datasketches/benches/tdigest.rs
@@ -88,11 +88,37 @@ fn partial_digest_lifecycle_by_k(bencher: Bencher, k: u16) {
}
#[divan::bench]
-fn compress(bencher: Bencher) {
+fn compress_initial_buffer(bencher: Bencher) {
// The default k=200 digest buffers 1,640 values before automatic
compression.
let values = values(1_640);
let digest = build_mut_digest(&values);
+ bencher
+ .counter(ItemsCount::new(values.len()))
+ .with_inputs(|| digest.clone())
+ .bench_local_values(|mut digest| black_box(digest.rank(0.5)));
+}
+
+#[divan::bench]
+fn compress_unmerged_tail(bencher: Bencher) {
+ let values = values(3_280);
+ let mut digest = build_mut_digest(&values[..1_640]);
+ black_box(digest.rank(0.5));
+ for &value in &values[1_640..] {
+ digest.update(value);
+ }
+
+ bencher
+ .counter(ItemsCount::new(1_640_usize))
+ .with_inputs(|| digest.clone())
+ .bench_local_values(|mut digest| black_box(digest.rank(0.5)));
+}
+
+#[divan::bench]
+fn freeze_initial_buffer(bencher: Bencher) {
+ let values = values(1_640);
+ let digest = build_mut_digest(&values);
+
bencher
.counter(ItemsCount::new(values.len()))
.with_inputs(|| digest.clone())
@@ -116,6 +142,20 @@ fn merge(bencher: Bencher) {
});
}
+#[divan::bench(args = [8, 1_640])]
+fn unmerged_merge(bencher: Bencher, left_rows: usize) {
+ let left = build_mut_digest(&values(left_rows));
+ let right = build_mut_digest(&values(1_640));
+
+ bencher
+ .counter(ItemsCount::new(left_rows + 1_640))
+ .with_inputs(|| left.clone())
+ .bench_local_values(|mut left| {
+ left.merge(black_box(&right));
+ black_box(left)
+ });
+}
+
#[divan::bench]
fn rank(bencher: Bencher) {
let digest = prepared_digest();
diff --git a/datasketches/src/tdigest/sketch.rs
b/datasketches/src/tdigest/sketch.rs
index 7e216f3..e24a99f 100644
--- a/datasketches/src/tdigest/sketch.rs
+++ b/datasketches/src/tdigest/sketch.rs
@@ -37,13 +37,93 @@ use crate::tdigest::serialization::SERIAL_VERSION;
/// The default value of K if one is not specified.
const DEFAULT_K: u16 = 200;
-/// Multiplier for buffer size relative to centroids capacity.
-const BUFFER_MULTIPLIER: usize = 4;
-/// Buffer capacity allocated by the first update to a digest.
-const INITIAL_BUFFER_CAPACITY: usize = 8;
+/// Multiplier for unmerged values relative to the target number of centroids.
+const UNMERGED_MULTIPLIER: usize = 4;
+/// Unmerged-value capacity allocated by the first update to a digest.
+const INITIAL_UNMERGED_CAPACITY: usize = 8;
/// Default weight for single values.
const DEFAULT_WEIGHT: NonZeroU64 = NonZeroU64::new(1).unwrap();
+// The update buffer has two physical representations:
+//
+// * `Staging` stores raw `f64` values compactly before the first compression.
+// * `Centroids` stores `[compressed prefix | unmerged unit-weight tail]` in
one allocation. The
+// tail length identifies the boundary between the two regions.
+//
+// Compression permanently transitions a non-empty buffer from `Staging` to
`Centroids`.
+#[derive(Debug, Clone)]
+enum TDigestBuffer {
+ Staging(Vec<f64>),
+ Centroids {
+ centroids: Vec<Centroid>,
+ unmerged_tail_len: usize,
+ },
+}
+
+impl Default for TDigestBuffer {
+ fn default() -> Self {
+ TDigestBuffer::Staging(vec![])
+ }
+}
+
+impl TDigestBuffer {
+ fn len(&self) -> usize {
+ match self {
+ TDigestBuffer::Staging(values) => values.len(),
+ TDigestBuffer::Centroids { centroids, .. } => centroids.len(),
+ }
+ }
+
+ fn is_empty(&self) -> bool {
+ self.len() == 0
+ }
+
+ fn unmerged_len(&self) -> usize {
+ match self {
+ TDigestBuffer::Staging(values) => values.len(),
+ TDigestBuffer::Centroids {
+ unmerged_tail_len, ..
+ } => *unmerged_tail_len,
+ }
+ }
+
+ /// Returns compressed centroids after the caller has processed staged
values.
+ fn compressed_centroids(&self) -> &[Centroid] {
+ match self {
+ TDigestBuffer::Staging(values) if values.is_empty() => &[],
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ } => centroids,
+ _ => unreachable!(
+ "t-digest buffer must be compressed before reading centroids:
{self:?}"
+ ),
+ }
+ }
+
+ fn into_compressed_centroids(self) -> Vec<Centroid> {
+ match self {
+ TDigestBuffer::Staging(values) if values.is_empty() => vec![],
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ } => centroids,
+ _ => unreachable!(
+ "t-digest buffer must be compressed before reading centroids:
{self:?}"
+ ),
+ }
+ }
+
+ fn estimated_size(&self) -> usize {
+ match self {
+ TDigestBuffer::Staging(values) => values.capacity() *
size_of::<f64>(),
+ TDigestBuffer::Centroids { centroids, .. } => {
+ centroids.capacity() * size_of::<Centroid>()
+ }
+ }
+ }
+}
+
/// T-Digest sketch for estimating quantiles and ranks.
///
/// See the [module level documentation](super) for more.
@@ -55,10 +135,10 @@ pub struct TDigestMut {
min: f64,
max: f64,
- centroids: Vec<Centroid>,
- centroids_weight: u64,
- centroids_capacity: usize,
- buffer: Vec<f64>,
+ buffer: TDigestBuffer,
+ // Weight represented by the compressed prefix. Staged values or the
unmerged tail contribute
+ // one each and are counted separately by `TDigestBuffer::unmerged_len`.
+ compressed_weight: u64,
}
impl Default for TDigestMut {
@@ -90,9 +170,8 @@ impl TDigestMut {
false,
f64::INFINITY,
f64::NEG_INFINITY,
- vec![],
+ TDigestBuffer::Staging(vec![]),
0,
- vec![],
)
}
@@ -124,9 +203,8 @@ impl TDigestMut {
false,
f64::INFINITY,
f64::NEG_INFINITY,
- vec![],
+ TDigestBuffer::Staging(vec![]),
0,
- vec![],
))
}
@@ -136,27 +214,41 @@ impl TDigestMut {
reverse_merge: bool,
min: f64,
max: f64,
- centroids: Vec<Centroid>,
- centroids_weight: u64,
- buffer: Vec<f64>,
+ buffer: TDigestBuffer,
+ compressed_weight: u64,
) -> Self {
assert!(k >= 10, "k must be at least 10");
-
- let fudge = if k < 30 { 30 } else { 10 };
- let centroids_capacity = (k as usize * 2) + fudge;
+ debug_assert!(match &buffer {
+ TDigestBuffer::Staging(_) => compressed_weight == 0,
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len,
+ } => *unmerged_tail_len <= centroids.len(),
+ });
TDigestMut {
k,
reverse_merge,
min,
max,
- centroids,
- centroids_weight,
- centroids_capacity,
buffer,
+ compressed_weight,
}
}
+ fn target_centroids(&self) -> usize {
+ let fudge = if self.k < 30 { 30 } else { 10 };
+ (usize::from(self.k) * 2) + fudge
+ }
+
+ fn max_unmerged(&self) -> usize {
+ self.target_centroids() * UNMERGED_MULTIPLIER
+ }
+
+ fn target_retained_capacity(&self) -> usize {
+ self.target_centroids() + self.max_unmerged()
+ }
+
/// Updates this t-digest with the given value.
///
/// [f64::NAN], [f64::INFINITY], and [f64::NEG_INFINITY] values are
ignored.
@@ -175,29 +267,71 @@ impl TDigestMut {
return;
}
- let full_buffer_capacity = self.centroids_capacity * BUFFER_MULTIPLIER;
- if self.buffer.len() == full_buffer_capacity {
+ let max_unmerged = self.max_unmerged();
+ if let TDigestBuffer::Staging(values) = &mut self.buffer {
+ if values.len() < max_unmerged {
+ if values.len() == values.capacity() {
+ let target_capacity = if values.capacity() == 0 {
+ INITIAL_UNMERGED_CAPACITY
+ } else if values.capacity() == INITIAL_UNMERGED_CAPACITY {
+ // Once a digest outgrows a tiny group, skip an extra
allocator round trip
+ // while keeping the first allocation small.
+ (INITIAL_UNMERGED_CAPACITY * UNMERGED_MULTIPLIER *
UNMERGED_MULTIPLIER)
+ .min(max_unmerged)
+ } else {
+ values
+ .capacity()
+ .saturating_mul(UNMERGED_MULTIPLIER)
+ .min(max_unmerged)
+ };
+
values.reserve_exact(target_capacity.saturating_sub(values.len()));
+ }
+
+ values.push(value);
+ self.min = self.min.min(value);
+ self.max = self.max.max(value);
+ return;
+ }
self.compress();
}
- if self.buffer.len() == self.buffer.capacity() {
- let target_capacity = if self.buffer.capacity() == 0 {
- INITIAL_BUFFER_CAPACITY
- } else if self.buffer.capacity() == INITIAL_BUFFER_CAPACITY {
+
+ if matches!(
+ &self.buffer,
+ TDigestBuffer::Centroids { unmerged_tail_len, .. } if
*unmerged_tail_len >= max_unmerged
+ ) {
+ self.compress();
+ }
+
+ let TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len,
+ } = &mut self.buffer
+ else {
+ unreachable!("a full staging buffer must become centroid-backed
after compression");
+ };
+ if centroids.len() == centroids.capacity() {
+ let target_unmerged = if *unmerged_tail_len == 0 {
+ INITIAL_UNMERGED_CAPACITY
+ } else if *unmerged_tail_len == INITIAL_UNMERGED_CAPACITY {
// Once a digest outgrows a tiny group, skip an extra
allocator round trip while
// keeping the first allocation small.
- (INITIAL_BUFFER_CAPACITY * BUFFER_MULTIPLIER *
BUFFER_MULTIPLIER)
- .min(full_buffer_capacity)
+ (INITIAL_UNMERGED_CAPACITY * UNMERGED_MULTIPLIER *
UNMERGED_MULTIPLIER)
+ .min(max_unmerged)
} else {
- self.buffer
- .capacity()
- .saturating_mul(BUFFER_MULTIPLIER)
- .min(full_buffer_capacity)
+ unmerged_tail_len
+ .saturating_mul(UNMERGED_MULTIPLIER)
+ .min(max_unmerged)
};
- self.buffer
-
.reserve_exact(target_capacity.saturating_sub(self.buffer.len()));
+ let num_merged = centroids.len() - *unmerged_tail_len;
+ let target_capacity = num_merged.saturating_add(target_unmerged);
+
centroids.reserve_exact(target_capacity.saturating_sub(centroids.len()));
}
- self.buffer.push(value);
+ centroids.push(Centroid {
+ mean: value,
+ weight: DEFAULT_WEIGHT,
+ });
+ *unmerged_tail_len += 1;
self.min = self.min.min(value);
self.max = self.max.max(value);
}
@@ -209,7 +343,7 @@ impl TDigestMut {
/// Returns `true` if this t-digest has not seen any data.
pub fn is_empty(&self) -> bool {
- self.centroids.is_empty() && self.buffer.is_empty()
+ self.buffer.is_empty()
}
/// Returns the minimum value seen by this t-digest, or `None` if it is
empty.
@@ -232,7 +366,7 @@ impl TDigestMut {
/// Returns the total weight.
pub fn total_weight(&self) -> u64 {
- self.centroids_weight + self.buffer.len() as u64
+ self.compressed_weight + self.buffer.unmerged_len() as u64
}
/// Merges the given t-digest into this one.
@@ -254,29 +388,47 @@ impl TDigestMut {
return;
}
- let additional_centroids = self.buffer.len() + other.centroids.len() +
other.buffer.len();
- let mut tmp = std::mem::take(&mut self.centroids);
- let num_existing_centroids = tmp.len();
- tmp.reserve(additional_centroids);
- for &v in &self.buffer {
- tmp.push(Centroid {
- mean: v,
- weight: DEFAULT_WEIGHT,
- });
- }
- for &v in &other.buffer {
- tmp.push(Centroid {
- mean: v,
- weight: DEFAULT_WEIGHT,
- });
- }
- for &c in &other.centroids {
- tmp.push(c);
+ let buffer = std::mem::take(&mut self.buffer);
+ let (mut merge_buffer, existing_prefix_len, self_unmerged_weight) =
match buffer {
+ TDigestBuffer::Staging(values) => {
+ let self_unmerged_weight = values.len() as u64;
+ let mut merge_buffer = Vec::with_capacity(values.len() +
other.buffer.len());
+ merge_buffer.extend(values.into_iter().map(|mean| Centroid {
+ mean,
+ weight: DEFAULT_WEIGHT,
+ }));
+ (merge_buffer, 0, self_unmerged_weight)
+ }
+ TDigestBuffer::Centroids {
+ mut centroids,
+ unmerged_tail_len,
+ } => {
+ let existing_prefix_len = centroids.len() - unmerged_tail_len;
+ centroids.reserve(other.buffer.len());
+ (centroids, existing_prefix_len, unmerged_tail_len as u64)
+ }
+ };
+ match &other.buffer {
+ TDigestBuffer::Staging(values) => {
+ merge_buffer.extend(values.iter().copied().map(|mean| Centroid
{
+ mean,
+ weight: DEFAULT_WEIGHT,
+ }));
+ }
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len,
+ } => {
+ let other_prefix_len = centroids.len() - unmerged_tail_len;
+ merge_buffer.extend_from_slice(¢roids[other_prefix_len..]);
+ merge_buffer.extend_from_slice(¢roids[..other_prefix_len]);
+ }
}
// Preserve the original insertion order for equal means because
t-digest requires a stable
- // sort: buffered values, the other digest's centroids, then this
digest's centroids.
- tmp.rotate_left(num_existing_centroids);
- self.do_merge(tmp, self.buffer.len() as u64 + other.total_weight())
+ // sort: this digest's buffered/unmerged values, the other digest's
buffered/unmerged values
+ // and compressed centroids, then this digest's compressed centroids.
+ merge_buffer.rotate_left(existing_prefix_len);
+ self.compress_centroids(merge_buffer, self_unmerged_weight +
other.total_weight())
}
/// Converts this mutable t-digest into an immutable one.
@@ -293,13 +445,17 @@ impl TDigestMut {
/// ```
pub fn freeze(mut self) -> TDigest {
self.compress();
+ let mut centroids = self.buffer.into_compressed_centroids();
+ // A mutable digest retains update workspace for reuse. The immutable
form cannot use that
+ // spare capacity, so release it at this consuming boundary.
+ centroids.shrink_to_fit();
TDigest {
k: self.k,
reverse_merge: self.reverse_merge,
min: self.min,
max: self.max,
- centroids: self.centroids,
- centroids_weight: self.centroids_weight,
+ centroids,
+ centroids_weight: self.compressed_weight,
}
}
@@ -308,8 +464,8 @@ impl TDigestMut {
TDigestView {
min: self.min,
max: self.max,
- centroids: &self.centroids,
- centroids_weight: self.centroids_weight,
+ centroids: self.buffer.compressed_centroids(),
+ centroids_weight: self.compressed_weight,
}
}
@@ -388,7 +544,7 @@ impl TDigestMut {
return Some(1.0);
}
// one centroid and value == min == max
- if self.centroids.len() + self.buffer.len() == 1 {
+ if self.buffer.len() == 1 {
return Some(0.5);
}
@@ -434,6 +590,7 @@ impl TDigestMut {
/// ```
pub fn serialize(&mut self) -> Vec<u8> {
self.compress();
+ let centroids = self.buffer.compressed_centroids();
let mut total_size = 0;
if self.is_empty() || self.is_single_value() {
@@ -460,7 +617,7 @@ impl TDigestMut {
// + 8 bytes max
total_size += size_of::<f64>() * 2;
// + (8+8) bytes per centroid
- total_size += self.centroids.len() * (size_of::<f64>() +
size_of::<u64>());
+ total_size += centroids.len() * (size_of::<f64>() +
size_of::<u64>());
}
let mut bytes = SketchBytes::with_capacity(total_size);
@@ -493,11 +650,11 @@ impl TDigestMut {
bytes.write_f64_le(self.min);
return bytes.into_bytes();
}
- bytes.write_u32_le(self.centroids.len() as u32);
+ bytes.write_u32_le(centroids.len() as u32);
bytes.write_u32_le(0); // unused
bytes.write_f64_le(self.min);
bytes.write_f64_le(self.max);
- for centroid in &self.centroids {
+ for centroid in centroids {
bytes.write_f64_le(centroid.mean);
bytes.write_u64_le(centroid.weight.get());
}
@@ -582,12 +739,14 @@ impl TDigestMut {
reverse_merge,
value,
value,
- vec![Centroid {
- mean: value,
- weight: DEFAULT_WEIGHT,
- }],
+ TDigestBuffer::Centroids {
+ centroids: vec![Centroid {
+ mean: value,
+ weight: DEFAULT_WEIGHT,
+ }],
+ unmerged_tail_len: 0,
+ },
1,
- vec![],
));
}
let num_centroids = cursor
@@ -611,8 +770,56 @@ impl TDigestMut {
check_non_nan(max, "max")?;
check_finite(min, "min")?;
check_finite(max, "max")?;
- let mut centroids = Vec::with_capacity(num_centroids);
- let mut centroids_weight = 0u64;
+ let (centroid_bytes, buffered_value_bytes) = if is_f32 {
+ (size_of::<f32>() + size_of::<u32>(), size_of::<f32>())
+ } else {
+ (size_of::<f64>() + size_of::<u64>(), size_of::<f64>())
+ };
+ let required_payload_bytes = num_centroids
+ .checked_mul(centroid_bytes)
+ .and_then(|bytes| {
+ num_buffered
+ .checked_mul(buffered_value_bytes)
+ .and_then(|buffered_bytes|
bytes.checked_add(buffered_bytes))
+ })
+ .ok_or_else(|| Error::deserial("TDigest payload size exceeds the
supported size"))?;
+ if cursor.remaining().len() < required_payload_bytes {
+ return Err(Error::insufficient_data(format!(
+ "TDigest payload requires {required_payload_bytes} bytes, got
{}",
+ cursor.remaining().len()
+ )));
+ }
+ if num_centroids == 0 {
+ checked_weight_sum(0, num_buffered as u64)?;
+ let mut initial_buffer = Vec::with_capacity(num_buffered);
+ for _ in 0..num_buffered {
+ let value = if is_f32 {
+ cursor
+ .read_f32_le()
+ .map_err(insufficient_data("buffered_value"))? as f64
+ } else {
+ cursor
+ .read_f64_le()
+ .map_err(insufficient_data("buffered_value"))?
+ };
+ check_non_nan(value, "buffered_value mean")?;
+ check_finite(value, "buffered_value mean")?;
+ initial_buffer.push(value);
+ }
+ return Ok(TDigestMut::make(
+ k,
+ reverse_merge,
+ min,
+ max,
+ TDigestBuffer::Staging(initial_buffer),
+ 0,
+ ));
+ }
+ let stored_centroids =
num_centroids.checked_add(num_buffered).ok_or_else(|| {
+ Error::deserial("num_centroids and num_buffered exceed the
supported size")
+ })?;
+ let mut centroids = Vec::with_capacity(stored_centroids);
+ let mut compressed_weight = 0u64;
for _ in 0..num_centroids {
let (mean, weight) = if is_f32 {
(
@@ -628,11 +835,10 @@ impl TDigestMut {
check_non_nan(mean, "centroid mean")?;
check_finite(mean, "centroid")?;
let weight = check_nonzero(weight, "centroid weight")?;
- centroids_weight = checked_weight_sum(centroids_weight,
weight.get())?;
+ compressed_weight = checked_weight_sum(compressed_weight,
weight.get())?;
centroids.push(Centroid { mean, weight });
}
- checked_weight_sum(centroids_weight, num_buffered as u64)?;
- let mut buffer = Vec::with_capacity(num_buffered);
+ checked_weight_sum(compressed_weight, num_buffered as u64)?;
for _ in 0..num_buffered {
let value = if is_f32 {
cursor
@@ -645,16 +851,21 @@ impl TDigestMut {
};
check_non_nan(value, "buffered_value mean")?;
check_finite(value, "buffered_value mean")?;
- buffer.push(value);
+ centroids.push(Centroid {
+ mean: value,
+ weight: DEFAULT_WEIGHT,
+ });
}
Ok(TDigestMut::make(
k,
reverse_merge,
min,
max,
- centroids,
- centroids_weight,
- buffer,
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: num_buffered,
+ },
+ compressed_weight,
))
}
@@ -705,9 +916,11 @@ impl TDigestMut {
false,
min,
max,
- centroids,
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ },
total_weight,
- vec![],
))
}
COMPAT_FLOAT => {
@@ -750,9 +963,11 @@ impl TDigestMut {
false,
min,
max,
- centroids,
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ },
total_weight,
- vec![],
))
}
ty => Err(Error::deserial(format!("unknown TDigest compat type
{ty}"))),
@@ -763,96 +978,123 @@ impl TDigestMut {
self.total_weight() == 1
}
- /// Process buffered values and merge centroids if needed.
+ /// Processes unmerged values and merges centroids if needed.
fn compress(&mut self) {
- if self.buffer.is_empty() {
- return;
- }
- let mut tmp = std::mem::take(&mut self.centroids);
- let full_buffer_capacity = self.centroids_capacity * BUFFER_MULTIPLIER;
- let required_capacity = tmp.len() + self.buffer.len();
- let target_capacity = if self.buffer.len() == full_buffer_capacity {
- required_capacity.max(full_buffer_capacity +
self.centroids_capacity)
- } else {
- required_capacity
- };
- tmp.reserve_exact(target_capacity.saturating_sub(tmp.len()));
- let num_buffered = self.buffer.len();
- for &v in &self.buffer {
- tmp.push(Centroid {
- mean: v,
- weight: DEFAULT_WEIGHT,
- });
+ let buffer = std::mem::take(&mut self.buffer);
+ match buffer {
+ TDigestBuffer::Staging(values) if values.is_empty() => {
+ self.buffer = TDigestBuffer::Staging(values);
+ }
+ TDigestBuffer::Staging(values) => {
+ debug_assert_eq!(self.compressed_weight, 0);
+ let weight = values.len() as u64;
+ let mut centroids = Vec::with_capacity(values.len());
+ centroids.extend(values.into_iter().map(|mean| Centroid {
+ mean,
+ weight: DEFAULT_WEIGHT,
+ }));
+ self.compress_centroids(centroids, weight);
+ }
+ TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ } => {
+ // Preserve compact deserialized images verbatim, including
images with more
+ // centroids than this implementation would normally produce.
+ self.buffer = TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ };
+ }
+ TDigestBuffer::Centroids {
+ mut centroids,
+ unmerged_tail_len,
+ } => {
+ let compressed_prefix_len = centroids.len() -
unmerged_tail_len;
+ // Preserve the original insertion order for equal means
because t-digest requires
+ // a stable sort: unmerged values before existing centroids.
+ centroids.rotate_left(compressed_prefix_len);
+ self.compress_centroids(centroids, unmerged_tail_len as u64);
+ }
}
- // Preserve the original insertion order for equal means because
t-digest requires a stable
- // sort: buffered values before existing centroids.
- tmp.rotate_right(num_buffered);
- self.do_merge(tmp, self.buffer.len() as u64)
}
- /// Merges the given buffer of centroids into this TDigest.
+ /// Compresses the given centroids into this t-digest.
///
/// # Contract
///
- /// * `buffer` must contain at least one centroid.
- /// * `buffer` contains every centroid to be merged, including all
centroids previously stored
- /// in `self`.
- /// * `weight` is the total weight not yet included in
`self.centroids_weight`.
- /// * Every centroid mean in `buffer` is finite.
- /// * `self.buffer` is cleared before returning.
- fn do_merge(&mut self, mut buffer: Vec<Centroid>, weight: u64) {
- buffer.sort_by(centroid_cmp);
+ /// * `centroids` must contain at least one centroid.
+ /// * `centroids` contains every centroid to be merged, including all
centroids previously
+ /// stored in `self`.
+ /// * `additional_weight` is the total weight not yet included in
`self.compressed_weight`.
+ /// * Every centroid mean in `centroids` is finite.
+ /// * `self.buffer` becomes centroid-backed with no unmerged values before
returning.
+ fn compress_centroids(&mut self, mut centroids: Vec<Centroid>,
additional_weight: u64) {
+ debug_assert!(!centroids.is_empty());
+ centroids.sort_by(centroid_cmp);
if self.reverse_merge {
- buffer.reverse();
+ centroids.reverse();
}
- self.centroids_weight += weight;
+ self.compressed_weight += additional_weight;
let mut num_centroids = 1;
- let len = buffer.len();
- let centroids_weight = self.centroids_weight as f64;
- let normalizer = scale_function::normalizer(2.0 * f64::from(self.k),
centroids_weight);
+ let len = centroids.len();
+ let compressed_weight = self.compressed_weight as f64;
+ let normalizer = scale_function::normalizer(2.0 * f64::from(self.k),
compressed_weight);
let mut current = 1;
let mut weight_so_far = 0.;
while current < len {
- let c = buffer[current];
- let proposed_weight = buffer[num_centroids - 1].weight() +
c.weight();
+ 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 / centroids_weight;
- let q2 = (weight_so_far + proposed_weight) / centroids_weight;
+ let q0 = weight_so_far / compressed_weight;
+ let q2 = (weight_so_far + proposed_weight) / compressed_weight;
add_this = proposed_weight
- <= (centroids_weight
+ <= (compressed_weight
* scale_function::max(q0, normalizer)
.min(scale_function::max(q2, normalizer)));
}
if add_this {
// merge into existing centroid
- buffer[num_centroids - 1].add(c);
+ centroids[num_centroids - 1].add(c);
} else {
// copy to a new centroid
- weight_so_far += buffer[num_centroids - 1].weight();
- buffer[num_centroids] = c;
+ weight_so_far += centroids[num_centroids - 1].weight();
+ centroids[num_centroids] = c;
num_centroids += 1;
}
current += 1;
}
- buffer.truncate(num_centroids);
+ centroids.truncate(num_centroids);
if self.reverse_merge {
- buffer.reverse();
+ centroids.reverse();
}
- self.min = self.min.min(buffer[0].mean);
- self.max = self.max.max(buffer[num_centroids - 1].mean);
- self.centroids = buffer;
+ self.min = self.min.min(centroids[0].mean);
+ self.max = self.max.max(centroids[num_centroids - 1].mean);
self.reverse_merge = !self.reverse_merge;
- self.buffer.clear();
+ self.reduce_retained_capacity(&mut centroids);
+ self.buffer = TDigestBuffer::Centroids {
+ centroids,
+ unmerged_tail_len: 0,
+ };
+ }
+
+ fn reduce_retained_capacity(&self, centroids: &mut Vec<Centroid>) {
+ let target_capacity =
self.target_retained_capacity().max(centroids.len());
+ if centroids.capacity() <= target_capacity {
+ return;
+ }
+
+ // A merge can temporarily exceed the update-path target. Shrink after
compaction so one
+ // unusually large input does not pin that peak capacity for the rest
of the digest's life.
+ centroids.shrink_to(target_capacity);
}
/// Returns the estimated size of the sketch in bytes.
pub fn estimated_size(&self) -> usize {
- size_of::<Self>()
- + self.centroids.capacity() * size_of::<Centroid>()
- + self.buffer.capacity() * size_of::<f64>()
+ size_of::<Self>() + self.buffer.estimated_size()
}
}
@@ -1062,9 +1304,11 @@ impl TDigest {
self.reverse_merge,
self.min,
self.max,
- self.centroids,
+ TDigestBuffer::Centroids {
+ centroids: self.centroids,
+ unmerged_tail_len: 0,
+ },
self.centroids_weight,
- vec![],
)
}
diff --git a/datasketches/tests/serde_tests/tdigest.rs
b/datasketches/tests/serde_tests/tdigest.rs
index 7c860d7..34463fc 100644
--- a/datasketches/tests/serde_tests/tdigest.rs
+++ b/datasketches/tests/serde_tests/tdigest.rs
@@ -25,6 +25,21 @@ use googletest::prelude::near;
use crate::serialization_test_data;
+fn fnv1a(bytes: &[u8]) -> u64 {
+ bytes.iter().fold(0xcbf2_9ce4_8422_2325, |hash, byte| {
+ (hash ^ u64::from(*byte)).wrapping_mul(0x0000_0100_0000_01b3)
+ })
+}
+
+fn patterned_digest(k: u16, len: usize, salt: usize) -> TDigestMut {
+ let mut tdigest = TDigestMut::new(k);
+ for index in 0..len {
+ let value = (((index * 37 + salt * 17) % 101) as f64) - 50.0;
+ tdigest.update(value);
+ }
+ tdigest
+}
+
fn test_sketch_file(path: PathBuf, n: u64, with_buffer: bool, is_f32: bool) {
let bytes = fs::read(&path).unwrap();
let td = TDigestMut::deserialize(&bytes, is_f32).unwrap();
@@ -202,6 +217,115 @@ fn test_many_values() {
assert_eq!(td.quantile(0.5), deserialized_td.quantile(0.5));
}
+#[test]
+fn test_serialized_bytes_stable_for_full_and_merged_digests() {
+ let mut full_buffer = patterned_digest(200, 1_641, 0);
+ let bytes = full_buffer.serialize();
+ assert_eq!(bytes.len(), 2_864);
+ assert_eq!(fnv1a(&bytes), 0x5c01_c50d_d1c8_fdbb);
+
+ let mut left = patterned_digest(10, 201, 2);
+ let mut right = patterned_digest(10, 199, 3);
+ right.rank(0.0);
+ left.merge(&right);
+ let bytes = left.serialize();
+ assert_eq!(bytes.len(), 272);
+ assert_eq!(fnv1a(&bytes), 0x7d2e_a927_9b9e_f559);
+
+ for &(left_len, right_len, expected_len, expected_hash) in &[
+ (8, 201, 272, 0x8522_1f3f_152f_24e5),
+ (201, 8, 256, 0xe60d_1f6f_f4b0_73e0),
+ (201, 401, 288, 0x4cb8_4037_5e68_ca4b),
+ (401, 201, 288, 0x6f6d_e965_77a7_a53f),
+ ] {
+ let mut left = patterned_digest(10, left_len, 2);
+ let right = patterned_digest(10, right_len, 3);
+ left.merge(&right);
+ let bytes = left.serialize();
+ assert_eq!(bytes.len(), expected_len);
+ assert_eq!(fnv1a(&bytes), expected_hash);
+ }
+}
+
+#[test]
+fn test_updates_normalize_overfull_deserialized_staging_buffer() {
+ let path = serialization_test_data("cpp_generated_files",
"tdigest_double_buf_n10_cpp.sk");
+ let mut bytes = fs::read(path).unwrap();
+ assert_eq!(&bytes[8..12], &0_u32.to_le_bytes()); // num centroids
+ assert_eq!(&bytes[12..16], &10_u32.to_le_bytes()); // num buffered
+
+ // k=100 normally compresses at 840 buffered values. Extend a real C++
staging image just past
+ // that producer threshold while keeping every added value within the
recorded min/max range.
+ bytes[12..16].copy_from_slice(&841_u32.to_le_bytes());
+ for _ in 0..831 {
+ bytes.extend_from_slice(&10_f64.to_le_bytes());
+ }
+
+ let mut tdigest = TDigestMut::deserialize(&bytes, false).unwrap();
+ for _ in 0..10_000 {
+ tdigest.update(10.0);
+ }
+
+ assert_eq!(tdigest.total_weight(), 10_841);
+ assert_eq!(tdigest.min_value(), Some(1.0));
+ assert_eq!(tdigest.max_value(), Some(10.0));
+ // The overfull image must not disable future compression and let the
staging buffer grow with
+ // every subsequent value.
+ assert!(tdigest.estimated_size() < 32_768);
+ let serialized = tdigest.serialize();
+ assert_eq!(&serialized[12..16], &0_u32.to_le_bytes());
+
+ let roundtrip = TDigestMut::deserialize(&serialized, false).unwrap();
+ assert_eq!(roundtrip.total_weight(), 10_841);
+ assert_eq!(roundtrip.min_value(), Some(1.0));
+ assert_eq!(roundtrip.max_value(), Some(10.0));
+}
+
+#[test]
+fn test_updates_normalize_overfull_deserialized_centroid_tail() {
+ let path = serialization_test_data("cpp_generated_files",
"tdigest_double_buf_n1000_cpp.sk");
+ let mut bytes = fs::read(path).unwrap();
+ assert_eq!(&bytes[8..12], &89_u32.to_le_bytes()); // num centroids
+ assert_eq!(&bytes[12..16], &160_u32.to_le_bytes()); // num buffered
+
+ // Extend the buffered tail of a real mixed C++ image just past the k=100
producer threshold.
+ bytes[12..16].copy_from_slice(&841_u32.to_le_bytes());
+ for _ in 0..681 {
+ bytes.extend_from_slice(&1_000_f64.to_le_bytes());
+ }
+
+ let mut tdigest = TDigestMut::deserialize(&bytes, false).unwrap();
+ for _ in 0..10_000 {
+ tdigest.update(1_000.0);
+ }
+
+ assert_eq!(tdigest.total_weight(), 11_681);
+ assert_eq!(tdigest.min_value(), Some(1.0));
+ assert_eq!(tdigest.max_value(), Some(1_000.0));
+ // The overfull image must not disable future compression and let the
centroid tail grow with
+ // every subsequent value.
+ assert!(tdigest.estimated_size() < 32_768);
+ let serialized = tdigest.serialize();
+ assert_eq!(&serialized[12..16], &0_u32.to_le_bytes());
+
+ let roundtrip = TDigestMut::deserialize(&serialized, false).unwrap();
+ assert_eq!(roundtrip.total_weight(), 11_681);
+ assert_eq!(roundtrip.min_value(), Some(1.0));
+ assert_eq!(roundtrip.max_value(), Some(1_000.0));
+}
+
+#[test]
+fn test_deserialize_rejects_truncated_large_payload_before_allocation() {
+ let mut tdigest = TDigestMut::new(10);
+ tdigest.update(0.0);
+ tdigest.update(1.0);
+ let mut bytes = tdigest.serialize();
+ bytes[8..12].copy_from_slice(&u32::MAX.to_le_bytes());
+ bytes[12..16].copy_from_slice(&u32::MAX.to_le_bytes());
+
+ assert!(TDigestMut::deserialize(&bytes, false).is_err());
+}
+
#[test]
fn test_large_weights_produce_finite_extreme_quantile() {
let lower = f64::from_bits(f64::MAX.to_bits() - 1);
diff --git a/datasketches/tests/tdigest_test/sketch.rs
b/datasketches/tests/tdigest_test/sketch.rs
index 3094194..56be846 100644
--- a/datasketches/tests/tdigest_test/sketch.rs
+++ b/datasketches/tests/tdigest_test/sketch.rs
@@ -15,6 +15,8 @@
// specific language governing permissions and limitations
// under the License.
+use std::mem::size_of;
+
use datasketches::tdigest::TDigestMut;
use googletest::assert_that;
use googletest::prelude::eq;
@@ -75,6 +77,61 @@ fn test_maximum_k() {
assert_eq!(tdigest.quantile(0.5), Some(1.0));
}
+#[test]
+fn test_estimated_size_releases_initial_staging_after_compression() {
+ const K: u16 = 200;
+ const TARGET_CENTROIDS: usize = 410;
+ const MAX_UNMERGED: usize = TARGET_CENTROIDS * 4;
+
+ let inline_size = size_of::<TDigestMut>();
+ let mut tdigest = TDigestMut::new(K);
+ assert_eq!(tdigest.estimated_size(), inline_size);
+
+ for value in 0..MAX_UNMERGED {
+ tdigest.update(value as f64);
+ }
+ let size_before_compression = tdigest.estimated_size();
+ assert!(size_before_compression > inline_size);
+ tdigest.rank(0.5);
+ // Unit-weight centroids are twice the size of staged f64 values. Allow
that representation
+ // change while guarding against retaining both backing allocations after
compression.
+ assert!(tdigest.estimated_size() <= size_before_compression * 2 +
inline_size);
+
+ for value in MAX_UNMERGED..10_000 {
+ tdigest.update(value as f64);
+ }
+ let size_before_compression = tdigest.estimated_size();
+ tdigest.rank(0.5);
+ assert!(tdigest.estimated_size() <= size_before_compression);
+
+ let mut left = TDigestMut::new(K);
+ for value in 0..8 {
+ left.update(value as f64);
+ }
+ let mut right = TDigestMut::new(K);
+ for value in 0..MAX_UNMERGED {
+ right.update(value as f64);
+ }
+ let right_size = right.estimated_size();
+ left.merge(&right);
+ assert_eq!(left.total_weight(), (MAX_UNMERGED + 8) as u64);
+ assert_eq!(right.total_weight(), MAX_UNMERGED as u64);
+ assert_eq!(right.estimated_size(), right_size);
+
+ let mut full_left = TDigestMut::new(K);
+ for value in 0..MAX_UNMERGED {
+ full_left.update(value as f64);
+ }
+ let combined_size = full_left.estimated_size() + right_size;
+ full_left.merge(&right);
+ assert_eq!(full_left.total_weight(), (MAX_UNMERGED * 2) as u64);
+ assert!(full_left.estimated_size() <= combined_size + combined_size / 2);
+
+ let mutable_size = full_left.estimated_size();
+ let frozen = full_left.freeze();
+ assert!(frozen.estimated_size() <= mutable_size);
+}
+
#[test]
fn test_many_values() {
let n = 10000;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
