diff --git a/datasketches/src/tdigest/mod.rs b/datasketches/src/tdigest/mod.rs index 8c5818b4..652ba811 100644 --- a/datasketches/src/tdigest/mod.rs +++ b/datasketches/src/tdigest/mod.rs @@ -65,8 +65,6 @@ //! # Ok::<(), datasketches::error::Error>(()) //! ``` -mod serialization; - mod sketch; pub use self::sketch::TDigest; pub use self::sketch::TDigestMut; diff --git a/datasketches/src/tdigest/serialization.rs b/datasketches/src/tdigest/serialization.rs deleted file mode 100644 index debe6786..00000000 --- a/datasketches/src/tdigest/serialization.rs +++ /dev/null @@ -1,27 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -pub const PREAMBLE_LONGS_EMPTY_OR_SINGLE: u8 = 1; -pub const PREAMBLE_LONGS_MULTIPLE: u8 = 2; -pub const SERIAL_VERSION: u8 = 1; -pub const FLAGS_IS_EMPTY: u8 = 1 << 0; -pub const FLAGS_IS_SINGLE_VALUE: u8 = 1 << 1; -pub const FLAGS_REVERSE_MERGE: u8 = 1 << 2; -/// the format of the reference implementation is using double (f64) precision -pub const COMPAT_DOUBLE: u32 = 1; -/// the format of the reference implementation is using float (f32) precision -pub const COMPAT_FLOAT: u32 = 2; diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 83d8d2d2..e665a103 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -25,14 +25,6 @@ use crate::codec::assert::ensure_serial_version_is; use crate::codec::assert::insufficient_data; use crate::codec::family::Family; use crate::error::Error; -use crate::tdigest::serialization::COMPAT_DOUBLE; -use crate::tdigest::serialization::COMPAT_FLOAT; -use crate::tdigest::serialization::FLAGS_IS_EMPTY; -use crate::tdigest::serialization::FLAGS_IS_SINGLE_VALUE; -use crate::tdigest::serialization::FLAGS_REVERSE_MERGE; -use crate::tdigest::serialization::PREAMBLE_LONGS_EMPTY_OR_SINGLE; -use crate::tdigest::serialization::PREAMBLE_LONGS_MULTIPLE; -use crate::tdigest::serialization::SERIAL_VERSION; /// The default value of K if one is not specified. const DEFAULT_K: u16 = 200; @@ -43,126 +35,6 @@ const INITIAL_UNMERGED_CAPACITY: usize = 8; /// Default weight for single values. const DEFAULT_WEIGHT: NonZeroU64 = NonZeroU64::new(1).unwrap(); -// Centroids are stored as `[compressed prefix | unmerged unit-weight tail]`. Carrying a unit weight -// for each unmerged value lets updates, compression, and merge reuse one allocation instead of -// converting raw values into a second vector during compression. -#[derive(Debug, Clone, Default)] -struct TDigestBuffer { - centroids: Vec, - unmerged_tail_len: usize, -} - -impl TDigestBuffer { - fn new(centroids: Vec, unmerged_tail_len: usize) -> Self { - debug_assert!(unmerged_tail_len <= centroids.len()); - TDigestBuffer { - centroids, - unmerged_tail_len, - } - } - - fn len(&self) -> usize { - self.centroids.len() - } - - fn is_empty(&self) -> bool { - self.centroids.is_empty() - } - - fn unmerged_len(&self) -> usize { - self.unmerged_tail_len - } - - fn compressed_prefix_len(&self) -> usize { - self.centroids.len() - self.unmerged_tail_len - } - - fn push_unmerged(&mut self, value: f64, max_unmerged: usize) { - debug_assert!(self.unmerged_tail_len < max_unmerged); - if self.centroids.len() == self.centroids.capacity() { - let target_unmerged = if self.unmerged_tail_len == 0 { - INITIAL_UNMERGED_CAPACITY - } else if self.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_UNMERGED_CAPACITY * UNMERGED_MULTIPLIER * UNMERGED_MULTIPLIER) - .min(max_unmerged) - } else { - self.unmerged_tail_len - .saturating_mul(UNMERGED_MULTIPLIER) - .min(max_unmerged) - }; - let target_capacity = self.compressed_prefix_len().saturating_add(target_unmerged); - self.centroids - .reserve_exact(target_capacity.saturating_sub(self.centroids.len())); - } - - self.centroids.push(Centroid { - mean: value, - weight: DEFAULT_WEIGHT, - }); - self.unmerged_tail_len += 1; - } - - /// Returns all centroids in the tie order expected by stable compression sorting. - /// - /// The buffer is rotated from `[compressed | unmerged]` to `[unmerged | compressed]`, so new - /// values stay before existing centroids when their means are equal. - fn into_centroids_for_compression(mut self) -> Vec { - debug_assert_ne!(self.unmerged_tail_len, 0); - let compressed_prefix_len = self.compressed_prefix_len(); - self.centroids.rotate_left(compressed_prefix_len); - self.centroids - } - - /// 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 { - // 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; - } - - let compressed_prefix_len = self.compressed_prefix_len(); - self.centroids.reserve(other.len()); - let other_prefix_len = other.compressed_prefix_len(); - self.centroids - .extend_from_slice(&other.centroids[other_prefix_len..]); - self.centroids - .extend_from_slice(&other.centroids[..other_prefix_len]); - // Preserve the stable tie order: left unmerged, right unmerged and compressed, then the - // left compressed prefix. - self.centroids.rotate_left(compressed_prefix_len); - self.centroids.sort_by(centroid_cmp); - self.centroids - } - - #[inline] - fn compressed_centroids(&self) -> &[Centroid] { - assert_eq!( - self.unmerged_tail_len, 0, - "t-digest buffer must be compressed before reading centroids" - ); - &self.centroids - } - - fn into_compressed_centroids(self) -> Vec { - assert_eq!( - self.unmerged_tail_len, 0, - "t-digest buffer must be compressed before reading centroids" - ); - self.centroids - } - - /// Returns an estimate of memory usage in bytes. - fn estimated_size(&self) -> usize { - self.centroids.capacity() * size_of::() - } -} - /// T-Digest sketch for estimating quantiles and ranks. /// /// See the [module level documentation](super) for more. @@ -182,7 +54,7 @@ pub struct TDigestMut { impl Default for TDigestMut { fn default() -> Self { - Self::make( + Self::from_parts( DEFAULT_K, false, f64::INFINITY, @@ -215,7 +87,7 @@ impl TDigestMut { ))); } - Ok(Self::make( + Ok(Self::from_parts( k, false, f64::INFINITY, @@ -225,41 +97,6 @@ impl TDigestMut { )) } - fn make( - k: u16, - reverse_merge: bool, - min: f64, - max: f64, - buffer: TDigestBuffer, - compressed_weight: u64, - ) -> Self { - debug_assert!(k >= 10, "k must be at least 10"); - debug_assert!(buffer.unmerged_tail_len <= buffer.centroids.len()); - debug_assert!(buffer.compressed_prefix_len() != 0 || compressed_weight == 0); - - TDigestMut { - k, - reverse_merge, - min, - max, - 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. @@ -393,17 +230,6 @@ impl TDigestMut { } } - #[inline] - fn view(&mut self) -> TDigestView<'_> { - self.compress(); // side effect - TDigestView { - min: self.min, - max: self.max, - centroids: self.buffer.compressed_centroids(), - centroids_weight: self.compressed_weight, - } - } - /// Returns the cumulative distribution approximation described by [`TDigest::cdf`]. /// /// Returns `Ok(None)` if this t-digest is empty. @@ -629,310 +455,55 @@ impl TDigestMut { Self::deserialize_impl(bytes, true) } - fn deserialize_impl(bytes: &[u8], is_f32: bool) -> Result { - let mut cursor = SketchSlice::new(bytes); + /// Returns an estimate of the sketch's memory usage in bytes. + pub fn estimated_size(&self) -> usize { + self.buffer.estimated_size() + } - let preamble_longs = cursor - .read_u8() - .map_err(insufficient_data("preamble_longs"))?; - let serial_version = cursor - .read_u8() - .map_err(insufficient_data("serial_version"))?; - let family_id = cursor.read_u8().map_err(insufficient_data("family_id"))?; - if let Err(err) = Family::TDIGEST.validate_id(family_id) { - return if preamble_longs == 0 && serial_version == 0 && family_id == 0 { - Self::deserialize_compat(bytes) - } else { - Err(err) - }; - } - ensure_serial_version_is(SERIAL_VERSION, serial_version)?; - let k = cursor.read_u16_le().map_err(insufficient_data("k"))?; - if k < 10 { - 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", - )); + fn from_parts( + k: u16, + reverse_merge: bool, + min: f64, + max: f64, + buffer: TDigestBuffer, + compressed_weight: u64, + ) -> Self { + debug_assert!(k >= 10, "k must be at least 10"); + debug_assert!(buffer.unmerged_tail_len <= buffer.centroids.len()); + debug_assert!(buffer.compressed_prefix_len() != 0 || compressed_weight == 0); + + TDigestMut { + k, + reverse_merge, + min, + max, + buffer, + compressed_weight, } - let expected_preamble_longs = if is_empty || is_single_value { - PREAMBLE_LONGS_EMPTY_OR_SINGLE - } else { - PREAMBLE_LONGS_MULTIPLE - }; - ensure_preamble_longs_in(&[expected_preamble_longs], preamble_longs)?; - cursor - .read_u16_le() - .map_err(insufficient_data(""))?; // unused - if is_empty { - return TDigestMut::new(k); + } + + #[inline] + fn view(&mut self) -> TDigestView<'_> { + self.compress(); // side effect + TDigestView { + min: self.min, + max: self.max, + centroids: self.buffer.compressed_centroids(), + centroids_weight: self.compressed_weight, } + } - let reverse_merge = (flags & FLAGS_REVERSE_MERGE) != 0; - if is_single_value { - let value = if is_f32 { - cursor - .read_f32_le() - .map_err(insufficient_data("single_value"))? as f64 - } else { - cursor - .read_f64_le() - .map_err(insufficient_data("single_value"))? - }; - check_non_nan(value, "single_value")?; - check_finite(value, "single_value")?; - return Ok(TDigestMut::make( - k, - reverse_merge, - value, - value, - TDigestBuffer::new( - vec![Centroid { - mean: value, - weight: DEFAULT_WEIGHT, - }], - 0, - ), - 1, - )); - } - let num_centroids = cursor - .read_u32_le() - .map_err(insufficient_data("num_centroids"))? as usize; - let num_buffered = cursor - .read_u32_le() - .map_err(insufficient_data("num_buffered"))? as usize; - let (min, max) = if is_f32 { - ( - cursor.read_f32_le().map_err(insufficient_data("min"))? as f64, - cursor.read_f32_le().map_err(insufficient_data("max"))? as f64, - ) - } else { - ( - cursor.read_f64_le().map_err(insufficient_data("min"))?, - cursor.read_f64_le().map_err(insufficient_data("max"))?, - ) - }; - check_extrema(min, max, "TDigest")?; - let (centroid_bytes, buffered_value_bytes) = if is_f32 { - (size_of::() + size_of::(), size_of::()) - } else { - (size_of::() + size_of::(), size_of::()) - }; - let centroid_payload_bytes = num_centroids - .checked_mul(centroid_bytes) - .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; - let buffered_payload_bytes = num_buffered - .checked_mul(buffered_value_bytes) - .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; - let required_payload_bytes = centroid_payload_bytes - .checked_add(buffered_payload_bytes) - .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; - // Check the whole payload once so fixed-width records can be decoded without per-field I/O. - let available_bytes = cursor.remaining().len(); - if available_bytes < required_payload_bytes { - return Err(Error::insufficient_data_of( - "TDigest payload", - format_args!("expected {required_payload_bytes} bytes, got {available_bytes}"), - )); - } - let payload = &cursor.remaining()[..required_payload_bytes]; - let (centroid_payload, buffered_payload) = payload.split_at(centroid_payload_bytes); - 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 { - ( - f32::from_le_bytes(bytes[..4].try_into().unwrap()) as f64, - u32::from_le_bytes(bytes[4..].try_into().unwrap()) as u64, - ) - } else { - ( - f64::from_le_bytes(bytes[..8].try_into().unwrap()), - u64::from_le_bytes(bytes[8..].try_into().unwrap()), - ) - }; - 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", - )); - } - let total_weight = checked_weight_sum(compressed_weight, num_buffered as u64)?; - check_single_sample_extrema(min, max, total_weight)?; - 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); - 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, - min, - max, - TDigestBuffer::new(centroids, num_buffered), - compressed_weight, - )) + fn target_centroids(&self) -> usize { + let fudge = if self.k < 30 { 30 } else { 10 }; + (usize::from(self.k) * 2) + fudge } - // compatibility with the format of the reference implementation - // default byte order of ByteBuffer is used there, which is big endian - fn deserialize_compat(bytes: &[u8]) -> Result { - fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { - move |_| Error::insufficient_data_of("compat format", tag) - } - - let mut cursor = SketchSlice::new(bytes); + fn max_unmerged(&self) -> usize { + self.target_centroids() * UNMERGED_MULTIPLIER + } - let ty = cursor.read_u32_be().map_err(make_error("type"))?; - match ty { - COMPAT_DOUBLE => { - fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { - move |_| Error::insufficient_data_of("compat double format", tag) - } - // compatibility with asBytes() - let min = cursor.read_f64_be().map_err(make_error("min"))?; - let max = cursor.read_f64_be().map_err(make_error("max"))?; - check_extrema(min, max, "compat double TDigest")?; - let k = cursor.read_f64_be().map_err(make_error("k"))? as u16; - if k < 10 { - return Err(Error::deserial(format!( - "k must be at least 10 in compat double format, got {k}" - ))); - } - let num_centroids = - cursor.read_u32_be().map_err(make_error("num_centroids"))? as usize; - if num_centroids == 0 { - return Err(Error::deserial( - "malformed data: compat double TDigest must contain a centroid", - )); - } - let mut total_weight = 0u64; - let mut centroids = Vec::with_capacity(num_centroids); - let mut previous_mean = min; - let mut centroid_means_valid = true; - for _ in 0..num_centroids { - let weight = cursor.read_f64_be().map_err(make_error("weight"))?; - let mean = cursor.read_f64_be().map_err(make_error("mean"))?; - let weight = - check_compat_weight(weight, "centroid weight in compat double format")?; - centroid_means_valid &= - mean.is_finite() & (mean >= previous_mean) & (mean <= max); - previous_mean = mean; - total_weight = checked_weight_sum(total_weight, weight.get())?; - centroids.push(Centroid { mean, weight }); - } - if !centroid_means_valid { - return Err(Error::deserial( - "malformed data: centroid means in compat double format must be finite, within extrema, and nondecreasing", - )); - } - check_single_sample_extrema(min, max, total_weight)?; - Ok(TDigestMut::make( - k, - false, - min, - max, - TDigestBuffer::new(centroids, 0), - total_weight, - )) - } - COMPAT_FLOAT => { - fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { - move |_| Error::insufficient_data_of("compat float format", tag) - } - // COMPAT_FLOAT: compatibility with asSmallBytes() - // reference implementation uses doubles for min and max - let min = cursor.read_f64_be().map_err(make_error("min"))?; - let max = cursor.read_f64_be().map_err(make_error("max"))?; - check_extrema(min, max, "compat float TDigest")?; - let k = cursor.read_f32_be().map_err(make_error("k"))? as u16; - if k < 10 { - return Err(Error::deserial(format!( - "k must be at least 10 in compat float format, got {k}" - ))); - } - // reference implementation stores capacities of the array of centroids and the - // buffer as shorts they can be derived from k in the constructor - cursor.read_u32_be().map_err(make_error(""))?; - let num_centroids = - cursor.read_u16_be().map_err(make_error("num_centroids"))? as usize; - if num_centroids == 0 { - return Err(Error::deserial( - "malformed data: compat float TDigest must contain a centroid", - )); - } - let mut total_weight = 0u64; - let mut centroids = Vec::with_capacity(num_centroids); - let mut previous_mean = min; - let mut centroid_means_valid = true; - for _ in 0..num_centroids { - let weight = cursor.read_f32_be().map_err(make_error("weight"))? as f64; - let mean = cursor.read_f32_be().map_err(make_error("mean"))? as f64; - let weight = - check_compat_weight(weight, "centroid weight in compat float format")?; - centroid_means_valid &= - mean.is_finite() & (mean >= previous_mean) & (mean <= max); - previous_mean = mean; - total_weight = checked_weight_sum(total_weight, weight.get())?; - centroids.push(Centroid { mean, weight }); - } - if !centroid_means_valid { - return Err(Error::deserial( - "malformed data: centroid means in compat float format must be finite, within extrema, and nondecreasing", - )); - } - check_single_sample_extrema(min, max, total_weight)?; - Ok(TDigestMut::make( - k, - false, - min, - max, - TDigestBuffer::new(centroids, 0), - total_weight, - )) - } - ty => Err(Error::deserial(format!("unknown TDigest compat type {ty}"))), - } + fn target_retained_capacity(&self) -> usize { + self.target_centroids() + self.max_unmerged() } /// Processes unmerged values and merges centroids if needed. @@ -984,7 +555,7 @@ impl TDigestMut { normalizer, ) { // merge into existing centroid - centroids[num_centroids - 1].add(c); + centroids[num_centroids - 1].merge(c); } else { // copy to a new centroid weight_so_far += centroids[num_centroids - 1].weight(); @@ -1016,365 +587,54 @@ impl TDigestMut { // unusually large input does not pin that peak capacity for the rest of the digest's life. centroids.shrink_to(target_capacity); } - - /// Returns an estimate of the sketch's memory usage in bytes. - pub fn estimated_size(&self) -> usize { - self.buffer.estimated_size() - } } -/// Collects owned t-digests into one result with a single compression pass. -/// -/// Empty inputs are ignored; if all inputs are empty, returns [`TDigestMut::default()`]. The result -/// uses the smallest `k` among the non-empty inputs. A single non-empty input is returned -/// unchanged. +/// Immutable (frozen) T-Digest sketch for estimating quantiles and ranks. /// -/// Collecting consumes each digest without cloning it. Callers that need to retain their inputs -/// can use [`TDigestMut::merge`] or explicitly clone them before collection. -/// -/// Batch collection can save repeated compression work when combining many small partials, but -/// is not always faster than repeated [`TDigestMut::merge`] calls. Larger compressed inputs can -/// favor incremental merging. Input count, retained centroid count, and value distribution all -/// affect performance; the number of original samples alone does not determine the cost. -/// -/// Collection retains all non-empty input buffers before compression, even with a lazy iterator. -/// Fully compressed inputs are merged directly into the result; inputs with buffered updates also -/// require a combined centroid buffer and sorting workspace. Repeated `merge` calls or smaller -/// batches limit the number of inputs held at once, at the cost of additional compression passes. -/// Different merge groupings can produce different estimates. -/// -/// # 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); -/// ``` -/// -/// Serialized states can be decoded in batches, retaining only one batch of decoded inputs at a -/// time. A batch size limits the number of digests, not their byte size; a strict memory budget -/// needs to account for varying input sizes and merge workspace. -/// -/// ``` -/// use datasketches::tdigest::TDigestMut; -/// -/// let partial_states = [1.0, 2.0, 3.0, 4.0].map(|value| { -/// let mut digest = TDigestMut::new(100).unwrap(); -/// digest.update(value); -/// digest.serialize() -/// }); -/// -/// let mut merged = TDigestMut::new(100).unwrap(); -/// for batch in partial_states.chunks(2) { -/// let batch = batch -/// .iter() -/// .map(|bytes| TDigestMut::deserialize(bytes)) -/// .collect::>() -/// .unwrap(); -/// merged.merge(&batch); -/// } -/// assert_eq!(merged.total_weight(), 4); -/// assert_eq!(merged.min_value(), Some(1.0)); -/// assert_eq!(merged.max_value(), Some(4.0)); -/// ``` -impl FromIterator for TDigestMut { - fn from_iter>(iter: T) -> Self { - let mut digests = iter.into_iter(); - let Some(first) = digests.find(|digest| !digest.is_empty()) else { - return TDigestMut::default(); - }; - let Some(second) = digests.find(|digest| !digest.is_empty()) else { - return first; - }; +/// See the [module level documentation](super) for more. +pub struct TDigest { + k: u16, - // Each entry owns its input buffer and releases it when that input is exhausted. - struct MergeInput { - // Cache the next centroid so heap comparisons only touch heap storage. - next_centroid: Centroid, - centroids: std::vec::IntoIter, - unmerged_len: usize, - input_index: usize, - } + reverse_merge: bool, + min: f64, + max: f64, - fn sift_down_merge_heap(heap: &mut [MergeInput], mut position: usize, reverse: bool) { - // Equal means follow input order, as in a stable sort of concatenated inputs. - // Reverse compression reverses that entire order, including ties. - let precedes = |left: &MergeInput, right: &MergeInput| match centroid_cmp( - &left.next_centroid, - &right.next_centroid, - ) { - Ordering::Less => !reverse, - Ordering::Greater => reverse, - Ordering::Equal if reverse => left.input_index > right.input_index, - Ordering::Equal => left.input_index < right.input_index, - }; - loop { - let left = (position * 2) + 1; - if left >= heap.len() { - return; - } - let right = left + 1; - let next = if right < heap.len() && precedes(&heap[right], &heap[left]) { - right - } else { - left - }; - if !precedes(&heap[next], &heap[position]) { - return; - } - heap.swap(position, next); - position = next; - } - } + centroids: Vec, + centroids_weight: u64, +} - let reverse = first.reverse_merge; - let mut merged = TDigestMut::make( - first.k, - reverse, - first.min, - first.max, - TDigestBuffer::default(), - 0, - ); - let digests = [first, second].into_iter().chain(digests); - let mut heap = Vec::with_capacity(digests.size_hint().0); - let mut total_weight = 0u64; - let mut num_centroids = 0usize; - let mut all_compressed = true; - for (input_index, digest) in digests.enumerate() { - if digest.is_empty() { - continue; - } - total_weight = total_weight - .checked_add(digest.total_weight()) - .expect("combined t-digest weight exceeds u64::MAX"); - num_centroids = num_centroids - .checked_add(digest.buffer.len()) - .expect("combined t-digest centroid count exceeds usize::MAX"); - merged.k = merged.k.min(digest.k); - merged.min = merged.min.min(digest.min); - merged.max = merged.max.max(digest.max); - let unmerged_len = digest.buffer.unmerged_len(); - all_compressed &= unmerged_len == 0; - let next_centroid = if reverse { - *digest.buffer.centroids.last().unwrap() - } else { - digest.buffer.centroids[0] - }; - heap.push(MergeInput { - next_centroid, - centroids: digest.buffer.centroids.into_iter(), - unmerged_len, - input_index, - }); - } +impl TDigest { + /// Returns the compression parameter `k` used to configure this t-digest. + pub fn k(&self) -> u16 { + self.k + } - if !all_compressed { - // Raw tails precede compressed prefixes to preserve stable equal-mean ordering. - let mut centroids = Vec::with_capacity(num_centroids); - for input in &heap { - let prefix_len = input.centroids.len() - input.unmerged_len; - centroids.extend_from_slice(&input.centroids.as_slice()[prefix_len..]); - } - for input in heap { - let prefix_len = input.centroids.len() - input.unmerged_len; - centroids.extend(input.centroids.take(prefix_len)); - } - centroids.sort_by(centroid_cmp); - merged.compress_sorted_centroids(centroids, total_weight); - return merged; - } + /// Returns `true` if this t-digest has not seen any data. + pub fn is_empty(&self) -> bool { + self.centroids.is_empty() + } - for index in (0..heap.len() / 2).rev() { - sift_down_merge_heap(&mut heap, index, reverse); + /// Returns the minimum value seen by this t-digest, or `None` if it is empty. + pub fn min_value(&self) -> Option { + if self.is_empty() { + None + } else { + Some(self.min) } - let compressed_weight = total_weight as f64; - let normalizer = scale_function::normalizer(2.0 * f64::from(merged.k), compressed_weight); - // Reserve compressed output only; later updates can grow their own workspace. - let mut retained: Vec = - Vec::with_capacity(merged.target_centroids().min(num_centroids)); - let mut weight_so_far = 0.; - // Advancing the selected input can only move it down the heap; the other inputs stay put. - for current in 0..num_centroids { - let input = &mut heap[0]; - let centroid = input.next_centroid; - let _ = if reverse { - input.centroids.next_back() - } else { - input.centroids.next() - }; - if input.centroids.len() == 0 { - heap.swap_remove(0); - } else { - input.next_centroid = if reverse { - *input.centroids.as_slice().last().unwrap() - } else { - input.centroids.as_slice()[0] - }; - } - sift_down_merge_heap(&mut heap, 0, reverse); + } - if let Some(last) = retained.last_mut() { - let proposed_weight = last.weight() + centroid.weight(); - if should_merge_centroid( - current, - num_centroids, - weight_so_far, - proposed_weight, - compressed_weight, - normalizer, - ) { - last.add(centroid); - } else { - weight_so_far += last.weight(); - retained.push(centroid); - } - } else { - retained.push(centroid); - } + /// Returns the maximum value seen by this t-digest, or `None` if it is empty. + pub fn max_value(&self) -> Option { + if self.is_empty() { + None + } else { + Some(self.max) } + } - debug_assert!(retained.len() <= merged.target_centroids()); - if reverse { - retained.reverse(); - } - debug_assert_eq!( - retained - .iter() - .map(|centroid| centroid.weight.get()) - .sum::(), - total_weight, - "compressed centroids must preserve total weight" - ); - merged.compressed_weight = total_weight; - merged.reverse_merge = !reverse; - merged.buffer = TDigestBuffer::new(retained, 0); - merged - } -} - -fn serialize_compressed( - k: u16, - reverse_merge: bool, - min: f64, - max: f64, - centroids: &[Centroid], - total_weight: u64, -) -> Vec { - let is_empty = centroids.is_empty(); - let is_single_value = total_weight == 1; - let mut total_size = if is_empty || is_single_value { - // Preamble, serial version, family, k, flags, and two unused bytes. - size_of::() - } else { - // The short header plus centroid and buffered-value counts. - size_of::() * 2 - }; - if is_single_value { - total_size += size_of::(); - } else if !is_empty { - total_size += size_of::() * 2; - total_size += centroids.len() * (size_of::() + size_of::()); - } - - let mut bytes = SketchBytes::with_capacity(total_size); - bytes.write_u8(if is_empty || is_single_value { - PREAMBLE_LONGS_EMPTY_OR_SINGLE - } else { - PREAMBLE_LONGS_MULTIPLE - }); - bytes.write_u8(SERIAL_VERSION); - bytes.write_u8(Family::TDIGEST.id); - bytes.write_u16_le(k); - bytes.write_u8({ - let mut flags = 0; - if is_empty { - flags |= FLAGS_IS_EMPTY; - } - if is_single_value { - flags |= FLAGS_IS_SINGLE_VALUE; - } - if reverse_merge { - flags |= FLAGS_REVERSE_MERGE; - } - flags - }); - bytes.write_u16_le(0); // unused - if is_empty { - return bytes.into_bytes(); - } - if is_single_value { - bytes.write_f64_le(min); - return bytes.into_bytes(); - } - bytes.write_u32_le(centroids.len() as u32); - bytes.write_u32_le(0); // no buffered values - bytes.write_f64_le(min); - bytes.write_f64_le(max); - for centroid in centroids { - bytes.write_f64_le(centroid.mean); - bytes.write_u64_le(centroid.weight.get()); - } - bytes.into_bytes() -} - -/// Immutable (frozen) T-Digest sketch for estimating quantiles and ranks. -/// -/// See the [module level documentation](super) for more. -pub struct TDigest { - k: u16, - - reverse_merge: bool, - min: f64, - max: f64, - - centroids: Vec, - centroids_weight: u64, -} - -impl TDigest { - /// Returns the compression parameter `k` used to configure this t-digest. - pub fn k(&self) -> u16 { - self.k - } - - /// Returns `true` if this t-digest has not seen any data. - pub fn is_empty(&self) -> bool { - self.centroids.is_empty() - } - - /// Returns the minimum value seen by this t-digest, or `None` if it is empty. - pub fn min_value(&self) -> Option { - if self.is_empty() { - None - } else { - Some(self.min) - } - } - - /// Returns the maximum value seen by this t-digest, or `None` if it is empty. - pub fn max_value(&self) -> Option { - if self.is_empty() { - None - } else { - Some(self.max) - } - } - - /// Returns the total weight. - pub fn total_weight(&self) -> u64 { - self.centroids_weight + /// Returns the total weight. + pub fn total_weight(&self) -> u64 { + self.centroids_weight } /// Serializes this immutable t-digest to bytes in double precision. @@ -1432,15 +692,6 @@ impl TDigest { Ok(TDigestMut::deserialize_f32(bytes)?.freeze()) } - fn view(&self) -> TDigestView<'_> { - TDigestView { - min: self.min, - max: self.max, - centroids: &self.centroids, - centroids_weight: self.centroids_weight, - } - } - /// Returns an approximation to the Cumulative Distribution Function (CDF), which is the /// cumulative analog of the PMF, of the input stream given a set of split points. /// @@ -1627,7 +878,7 @@ impl TDigest { /// assert_eq!(mutable.total_weight(), 2); /// ``` pub fn unfreeze(self) -> TDigestMut { - TDigestMut::make( + TDigestMut::from_parts( self.k, self.reverse_merge, self.min, @@ -1641,532 +892,1353 @@ impl TDigest { pub fn estimated_size(&self) -> usize { self.centroids.capacity() * size_of::() } -} -struct TDigestView<'a> { - min: f64, - max: f64, - centroids: &'a [Centroid], - centroids_weight: u64, -} - -impl TDigestView<'_> { - fn pmf(&self, split_points: &[f64]) -> Option> { - let mut buckets = self.cdf(split_points)?; - let mut previous = buckets[0]; - for bucket in &mut buckets[1..] { - let rank = *bucket; - *bucket -= previous; - previous = rank; + fn view(&self) -> TDigestView<'_> { + TDigestView { + min: self.min, + max: self.max, + centroids: &self.centroids, + centroids_weight: self.centroids_weight, } - Some(buckets) } +} - fn cdf(&self, split_points: &[f64]) -> Option> { - if self.centroids.is_empty() { - return None; - } - if split_points.is_empty() { - return Some(vec![1.0]); - } +/// Collects owned t-digests into one result with a single compression pass. +/// +/// Empty inputs are ignored; if all inputs are empty, returns [`TDigestMut::default()`]. 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 it. Callers that need to retain their inputs +/// can use [`TDigestMut::merge`] or explicitly clone them before collection. +/// +/// Batch collection can save repeated compression work when combining many small partials, but +/// is not always faster than repeated [`TDigestMut::merge`] calls. Larger compressed inputs can +/// favor incremental merging. Input count, retained centroid count, and value distribution all +/// affect performance; the number of original samples alone does not determine the cost. +/// +/// Collection retains all non-empty input buffers before compression, even with a lazy iterator. +/// Fully compressed inputs are merged directly into the result; inputs with buffered updates also +/// require a combined centroid buffer and sorting workspace. Repeated `merge` calls or smaller +/// batches limit the number of inputs held at once, at the cost of additional compression passes. +/// Different merge groupings can produce different estimates. +/// +/// # 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); +/// ``` +/// +/// Serialized states can be decoded in batches, retaining only one batch of decoded inputs at a +/// time. A batch size limits the number of digests, not their byte size; a strict memory budget +/// needs to account for varying input sizes and merge workspace. +/// +/// ``` +/// use datasketches::tdigest::TDigestMut; +/// +/// let partial_states = [1.0, 2.0, 3.0, 4.0].map(|value| { +/// let mut digest = TDigestMut::new(100).unwrap(); +/// digest.update(value); +/// digest.serialize() +/// }); +/// +/// let mut merged = TDigestMut::new(100).unwrap(); +/// for batch in partial_states.chunks(2) { +/// let batch = batch +/// .iter() +/// .map(|bytes| TDigestMut::deserialize(bytes)) +/// .collect::>() +/// .unwrap(); +/// merged.merge(&batch); +/// } +/// assert_eq!(merged.total_weight(), 4); +/// assert_eq!(merged.min_value(), Some(1.0)); +/// assert_eq!(merged.max_value(), Some(4.0)); +/// ``` +impl FromIterator for TDigestMut { + fn from_iter>(iter: T) -> Self { + let mut digests = iter.into_iter(); + let Some(first) = digests.find(|digest| !digest.is_empty()) else { + return TDigestMut::default(); + }; + let Some(second) = digests.find(|digest| !digest.is_empty()) else { + return first; + }; - let binary_search_cost = self.centroids.len().ilog2() as usize; - let use_scan = split_points.len().saturating_mul(binary_search_cost) > self.centroids.len(); - let mut ranks = Vec::with_capacity(split_points.len() + 1); - if use_scan || self.min == self.max { - // Constant sketches and scans can fill the sorted outer tails directly. - let start = split_points.partition_point(|&p| p < self.min); - let end = start + split_points[start..].partition_point(|&p| p <= self.max); - ranks.resize(start, 0.0); - if self.min == self.max { - ranks.resize(end, 0.5); - } else { - self.fill_cdf_scan(&split_points[start..end], &mut ranks); - } - ranks.resize(split_points.len(), 1.0); - } else { - self.fill_cdf_binary_search(split_points, &mut ranks); + // Each entry owns its input buffer and releases it when that input is exhausted. + struct MergeInput { + // Cache the next centroid so heap comparisons only touch heap storage. + next_centroid: Centroid, + centroids: std::vec::IntoIter, + unmerged_len: usize, + input_index: usize, } - ranks.push(1.0); - Some(ranks) - } - - #[inline(always)] - fn rank(&self, value: f64) -> Option { - debug_assert!(!value.is_nan(), "value must not be NaN"); - if self.centroids.is_empty() { - return None; - } - if value < self.min { - return Some(0.0); - } - if value > self.max { - return Some(1.0); - } - if self.min == self.max { - return Some(0.5); - } - let total_weight = self.centroids_weight as f64; - let index = if value <= self.centroids[0].mean { - if value < self.centroids[0].mean { - return Some(self.mass_in_left_tail(value) / total_weight); - } - 0 - } else { - if value > self.centroids.last().unwrap().mean { - return Some(self.mass_in_right_tail(value) / total_weight); + fn sift_down_merge_heap(heap: &mut [MergeInput], mut position: usize, reverse: bool) { + // Equal means follow input order, as in a stable sort of concatenated inputs. + // Reverse compression reverses that entire order, including ties. + let precedes = |left: &MergeInput, right: &MergeInput| match centroid_cmp( + &left.next_centroid, + &right.next_centroid, + ) { + Ordering::Less => !reverse, + Ordering::Greater => reverse, + Ordering::Equal if reverse => left.input_index > right.input_index, + Ordering::Equal => left.input_index < right.input_index, + }; + loop { + let left = (position * 2) + 1; + if left >= heap.len() { + return; + } + let right = left + 1; + let next = if right < heap.len() && precedes(&heap[right], &heap[left]) { + right + } else { + left + }; + if !precedes(&heap[next], &heap[position]) { + return; + } + heap.swap(position, next); + position = next; } - self.centroids.partition_point(|c| c.mean < value) - }; - let (lower, anchors) = self.rank_anchors(value, index); - let weight_below = self.centroids[..lower].iter().map(|c| c.weight.get()).sum(); - Some(Self::interpolated_mass(value, anchors, weight_below) / total_weight) - } + } - #[inline(always)] - fn fill_cdf_binary_search(&self, split_points: &[f64], ranks: &mut Vec) { - let total_weight = self.centroids_weight as f64; - let mut previous_lower = 0; - let mut weight_below = 0; - for &value in split_points { - if value < self.min { - ranks.push(0.0); - continue; - } - if value > self.max { - ranks.push(1.0); + let reverse = first.reverse_merge; + let mut merged = TDigestMut::from_parts( + first.k, + reverse, + first.min, + first.max, + TDigestBuffer::default(), + 0, + ); + let digests = [first, second].into_iter().chain(digests); + let mut heap = Vec::with_capacity(digests.size_hint().0); + let mut total_weight = 0u64; + let mut num_centroids = 0usize; + let mut all_compressed = true; + for (input_index, digest) in digests.enumerate() { + if digest.is_empty() { continue; } - let index = if value <= self.centroids[0].mean { - if value < self.centroids[0].mean { - ranks.push(self.mass_in_left_tail(value) / total_weight); - continue; - } - 0 + total_weight = total_weight + .checked_add(digest.total_weight()) + .expect("combined t-digest weight exceeds u64::MAX"); + num_centroids = num_centroids + .checked_add(digest.buffer.len()) + .expect("combined t-digest centroid count exceeds usize::MAX"); + merged.k = merged.k.min(digest.k); + merged.min = merged.min.min(digest.min); + merged.max = merged.max.max(digest.max); + let unmerged_len = digest.buffer.unmerged_len(); + all_compressed &= unmerged_len == 0; + let next_centroid = if reverse { + *digest.buffer.centroids.last().unwrap() } else { - if value > self.centroids.last().unwrap().mean { - ranks.push(self.mass_in_right_tail(value) / total_weight); - continue; - } - self.centroids.partition_point(|c| c.mean < value) + digest.buffer.centroids[0] }; - let (lower, anchors) = self.rank_anchors(value, index); - weight_below += self.centroids[previous_lower..lower] - .iter() - .map(|c| c.weight.get()) - .sum::(); - previous_lower = lower; - ranks.push(Self::interpolated_mass(value, anchors, weight_below) / total_weight); + heap.push(MergeInput { + next_centroid, + centroids: digest.buffer.centroids.into_iter(), + unmerged_len, + input_index, + }); } - } - #[inline(always)] - fn fill_cdf_scan(&self, split_points: &[f64], ranks: &mut Vec) { - let total_weight = self.centroids_weight as f64; - let mut index = 0; - let mut weight_before_index = 0; - for &value in split_points { - let lower_bound = if value <= self.centroids[0].mean { - if value < self.centroids[0].mean { - ranks.push(self.mass_in_left_tail(value) / total_weight); - continue; - } - 0 + if !all_compressed { + // Raw tails precede compressed prefixes to preserve stable equal-mean ordering. + let mut centroids = Vec::with_capacity(num_centroids); + for input in &heap { + let prefix_len = input.centroids.len() - input.unmerged_len; + centroids.extend_from_slice(&input.centroids.as_slice()[prefix_len..]); + } + for input in heap { + let prefix_len = input.centroids.len() - input.unmerged_len; + centroids.extend(input.centroids.take(prefix_len)); + } + centroids.sort_by(centroid_cmp); + merged.compress_sorted_centroids(centroids, total_weight); + return merged; + } + + for index in (0..heap.len() / 2).rev() { + sift_down_merge_heap(&mut heap, index, reverse); + } + let compressed_weight = total_weight as f64; + let normalizer = scale_function::normalizer(2.0 * f64::from(merged.k), compressed_weight); + // Reserve compressed output only; later updates can grow their own workspace. + let mut retained: Vec = + Vec::with_capacity(merged.target_centroids().min(num_centroids)); + let mut weight_so_far = 0.; + // Advancing the selected input can only move it down the heap; the other inputs stay put. + for current in 0..num_centroids { + let input = &mut heap[0]; + let centroid = input.next_centroid; + let _ = if reverse { + input.centroids.next_back() } else { - if value > self.centroids.last().unwrap().mean { - ranks.push(self.mass_in_right_tail(value) / total_weight); - continue; - } - while index < self.centroids.len() && self.centroids[index].mean < value { - weight_before_index += self.centroids[index].weight.get(); - index += 1; - } - index + input.centroids.next() }; - let (lower, anchors) = self.rank_anchors(value, lower_bound); - // Between means, the lower anchor precedes index. Its weight belongs to the - // interpolation interval rather than the prefix below that interval. - let weight_below = weight_before_index - - if lower < lower_bound { - self.centroids[lower].weight.get() + if input.centroids.len() == 0 { + heap.swap_remove(0); + } else { + input.next_centroid = if reverse { + *input.centroids.as_slice().last().unwrap() } else { - 0 + input.centroids.as_slice()[0] }; - ranks.push(Self::interpolated_mass(value, anchors, weight_below) / total_weight); - } - } + } + sift_down_merge_heap(&mut heap, 0, reverse); - #[inline] - fn mass_in_left_tail(&self, value: f64) -> f64 { - let center = centroid_center(0, self.centroids[0].weight.get()); - if value == self.min { - 0.5 - } else { - interpolate( - center.min(1.), - center, - interpolation_fraction(value, self.min, self.centroids[0].mean), - ) + if let Some(last) = retained.last_mut() { + let proposed_weight = last.weight() + centroid.weight(); + if should_merge_centroid( + current, + num_centroids, + weight_so_far, + proposed_weight, + compressed_weight, + normalizer, + ) { + last.merge(centroid); + } else { + weight_so_far += last.weight(); + retained.push(centroid); + } + } else { + retained.push(centroid); + } } - } - #[inline] - fn mass_in_right_tail(&self, value: f64) -> f64 { - let last = self.centroids.last().unwrap(); - let center = centroid_center(self.centroids_weight - last.weight.get(), last.weight.get()); - if value == self.max { - // Use the same rounded mass as the adjacent centroid center. - centroid_center(self.centroids_weight - 1, 1) - } else { - interpolate( - center, - center.max((self.centroids_weight - 1) as f64), - interpolation_fraction(value, last.mean, self.max), - ) + debug_assert!(retained.len() <= merged.target_centroids()); + if reverse { + retained.reverse(); } + debug_assert_eq!( + retained + .iter() + .map(|centroid| centroid.weight64()) + .sum::(), + total_weight, + "compressed centroids must preserve total weight" + ); + merged.compressed_weight = total_weight; + merged.reverse_merge = !reverse; + merged.buffer = TDigestBuffer::new(retained, 0); + merged } +} - #[inline(always)] - fn rank_anchors(&self, value: f64, index: usize) -> (usize, &[Centroid]) { - // An exact mean spans its entire equal run; otherwise adjacent centroids bracket it. - if self.centroids[index].mean == value { - let end = self.centroids.partition_point(|c| c.mean <= value); - (index, &self.centroids[index..end]) - } else { - (index - 1, &self.centroids[index - 1..=index]) +// Centroids are stored as `[compressed prefix | unmerged unit-weight tail]`. Carrying a unit weight +// for each unmerged value lets updates, compression, and merge reuse one allocation instead of +// converting raw values into a second vector during compression. +#[derive(Debug, Clone, Default)] +struct TDigestBuffer { + centroids: Vec, + unmerged_tail_len: usize, +} + +impl TDigestBuffer { + fn new(centroids: Vec, unmerged_tail_len: usize) -> Self { + debug_assert!(unmerged_tail_len <= centroids.len()); + TDigestBuffer { + centroids, + unmerged_tail_len, } } - #[inline(always)] - fn interpolated_mass(value: f64, anchors: &[Centroid], weight_below: u64) -> f64 { - let left = &anchors[0]; - let Some((right, middle)) = anchors[1..].split_last() else { - return centroid_center(weight_below, left.weight.get()); - }; - let weight_between = left.weight.get() + middle.iter().map(|c| c.weight.get()).sum::(); - let fraction = if left.mean < right.mean { - interpolation_fraction(value, left.mean, right.mean) - } else { - 0.5 - }; - interpolate( - centroid_center(weight_below, left.weight.get()), - centroid_center(weight_below + weight_between, right.weight.get()), - fraction, - ) + fn len(&self) -> usize { + self.centroids.len() } - #[inline] - fn quantile(&self, rank: f64) -> Option { - debug_assert!((0.0..=1.0).contains(&rank), "rank must be in [0.0, 1.0]"); + fn is_empty(&self) -> bool { + self.centroids.is_empty() + } - if self.centroids.is_empty() { - return None; + fn unmerged_len(&self) -> usize { + self.unmerged_tail_len + } + + fn compressed_prefix_len(&self) -> usize { + self.centroids.len() - self.unmerged_tail_len + } + + fn push_unmerged(&mut self, value: f64, max_unmerged: usize) { + debug_assert!(self.unmerged_tail_len < max_unmerged); + if self.centroids.len() == self.centroids.capacity() { + let target_unmerged = if self.unmerged_tail_len == 0 { + INITIAL_UNMERGED_CAPACITY + } else if self.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_UNMERGED_CAPACITY * UNMERGED_MULTIPLIER * UNMERGED_MULTIPLIER) + .min(max_unmerged) + } else { + self.unmerged_tail_len + .saturating_mul(UNMERGED_MULTIPLIER) + .min(max_unmerged) + }; + let target_capacity = self.compressed_prefix_len().saturating_add(target_unmerged); + self.centroids + .reserve_exact(target_capacity.saturating_sub(self.centroids.len())); } - let mut quantile = [0.]; - self.fill_quantiles(std::iter::once((0, rank)), &mut quantile); - Some(quantile[0]) + self.centroids.push(Centroid { + mean: value, + weight: DEFAULT_WEIGHT, + }); + self.unmerged_tail_len += 1; } - 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]" - ); + /// Returns all centroids in the tie order expected by stable compression sorting. + /// + /// The buffer is rotated from `[compressed | unmerged]` to `[unmerged | compressed]`, so new + /// values stay before existing centroids when their means are equal. + fn into_centroids_for_compression(mut self) -> Vec { + debug_assert_ne!(self.unmerged_tail_len, 0); + let compressed_prefix_len = self.compressed_prefix_len(); + self.centroids.rotate_left(compressed_prefix_len); + self.centroids + } - if self.centroids.is_empty() { - return None; + /// 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 { + // 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; } - let mut quantiles = vec![0.; ranks.len()]; - if ranks.is_sorted() { - self.fill_quantiles(ranks.iter().copied().enumerate(), &mut quantiles); - return Some(quantiles); - } + let compressed_prefix_len = self.compressed_prefix_len(); + self.centroids.reserve(other.len()); + let other_prefix_len = other.compressed_prefix_len(); + self.centroids + .extend_from_slice(&other.centroids[other_prefix_len..]); + self.centroids + .extend_from_slice(&other.centroids[..other_prefix_len]); + // Preserve the stable tie order: left unmerged, right unmerged and compressed, then the + // left compressed prefix. + self.centroids.rotate_left(compressed_prefix_len); + self.centroids.sort_by(centroid_cmp); + self.centroids + } - // The scan 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_unstable_by(|&left, &right| { - // Ranks have already been validated, so neither value can be NaN. - ranks[left].partial_cmp(&ranks[right]).unwrap() - }); - self.fill_quantiles( - rank_order.into_iter().map(|index| (index, ranks[index])), - &mut quantiles, + #[inline] + fn compressed_centroids(&self) -> &[Centroid] { + assert_eq!( + self.unmerged_tail_len, 0, + "t-digest buffer must be compressed before reading centroids" ); - Some(quantiles) + &self.centroids } - /// Answers nondecreasing ranks at their original output indices. - fn fill_quantiles( - &self, - ranks: impl DoubleEndedIterator, - quantiles: &mut [f64], - ) { - if self.min == self.max { - quantiles.fill(self.min); - return; - } + fn into_compressed_centroids(self) -> Vec { + assert_eq!( + self.unmerged_tail_len, 0, + "t-digest buffer must be compressed before reading centroids" + ); + self.centroids + } - let centroids_weight = self.centroids_weight as f64; - let mut queries = ranks - .map(|(index, rank)| (index, rank * centroids_weight)) - .rev() - .peekable(); + /// Returns an estimate of memory usage in bytes. + fn estimated_size(&self) -> usize { + self.centroids.capacity() * size_of::() + } +} - // Consume the right tail from the back, so the centroid scan can omit all tail checks. - while let Some((index, _)) = queries.next_if(|&(_, weight)| centroids_weight - weight < 1.) - { - quantiles[index] = self.max; - } - if queries.peek().is_none() { - return; - } +struct TDigestView<'a> { + min: f64, + max: f64, + centroids: &'a [Centroid], + centroids_weight: u64, +} - let last = self.centroids.last().unwrap(); - let last_weight = last.weight.get(); - let last_center = centroid_center(self.centroids_weight - last_weight, last_weight); - if last_weight > 1 { - let tail_end = (self.centroids_weight - 1) as f64; - while let Some((index, weight)) = queries.next_if(|&(_, weight)| weight >= last_center) - { - // Integer counts can round a wider tail to zero width as well. At this - // boundary prefer the maximum, as for an exact two-sample tail. - quantiles[index] = if last_center == tail_end { - self.max - } else { - interpolate( - last.mean, - self.max, - (weight - last_center) / (tail_end - last_center), - ) - }; - } +impl TDigestView<'_> { + fn pmf(&self, split_points: &[f64]) -> Option> { + let mut buckets = self.cdf(split_points)?; + let mut previous = buckets[0]; + for bucket in &mut buckets[1..] { + let rank = *bucket; + *bucket -= previous; + previous = rank; } + Some(buckets) + } - let mut queries = queries.rev().peekable(); - while let Some((index, _)) = queries.next_if(|&(_, weight)| weight < 1.) { - quantiles[index] = self.min; + fn cdf(&self, split_points: &[f64]) -> Option> { + if self.centroids.is_empty() { + return None; } - if queries.peek().is_none() { - return; + if split_points.is_empty() { + return Some(vec![1.0]); } - let first = &self.centroids[0]; - let first_weight = first.weight.get(); - let first_center = centroid_center(0, first_weight); - if first_weight > 1 { - // At first_weight == 2, the weight < 1 pre-pass already consumed this interval. - while let Some((index, weight)) = queries.next_if(|&(_, weight)| weight < first_center) - { - quantiles[index] = - interpolate(self.min, first.mean, (weight - 1.) / (first_center - 1.)); + let binary_search_cost = self.centroids.len().ilog2() as usize; + let use_scan = split_points.len().saturating_mul(binary_search_cost) > self.centroids.len(); + let mut ranks = Vec::with_capacity(split_points.len() + 1); + if use_scan || self.min == self.max { + // Constant sketches and scans can fill the sorted outer tails directly. + let start = split_points.partition_point(|&p| p < self.min); + let end = start + split_points[start..].partition_point(|&p| p <= self.max); + ranks.resize(start, 0.0); + if self.min == self.max { + ranks.resize(end, 0.5); + } else { + self.fill_cdf_scan(&split_points[start..end], &mut ranks); } + ranks.resize(split_points.len(), 1.0); + } else { + self.fill_cdf_binary_search(split_points, &mut ranks); } - if queries.peek().is_none() { - return; - } + ranks.push(1.0); + Some(ranks) + } - // Answer the remaining queries by centroid interval; both streams advance only forward. - let mut weight_before = 0; - let mut left_center = first_center; - for pair in self.centroids.windows(2) { - let left = &pair[0]; - let right = &pair[1]; - let right_before = weight_before + left.weight.get(); - let right_center = centroid_center(right_before, right.weight.get()); - loop { - let Some(&(index, weight)) = queries.peek() else { - return; - }; - if right_center <= weight { - break; - } - queries.next(); - let mut start = left_center; - if left.weight.get() == 1 { - start = right_before as f64; - if weight < start { - quantiles[index] = left.mean; - continue; - } - } - let mut end = right_center; - if right.weight.get() == 1 { - end = right_before as f64; - if weight >= end { - quantiles[index] = right.mean; - continue; - } - } - quantiles[index] = - interpolate(left.mean, right.mean, (weight - start) / (end - start)); - } - weight_before = right_before; - left_center = right_center; - } + #[inline(always)] + fn rank(&self, value: f64) -> Option { + debug_assert!(!value.is_nan(), "value must not be NaN"); - // Rounding at large total weights can exhaust the scan near the maximum. - for (index, _) in queries { - quantiles[index] = self.max; + if self.centroids.is_empty() { + return None; } + if value < self.min { + return Some(0.0); + } + if value > self.max { + return Some(1.0); + } + if self.min == self.max { + return Some(0.5); + } + let total_weight = self.centroids_weight as f64; + let index = if value <= self.centroids[0].mean { + if value < self.centroids[0].mean { + return Some(self.left_tail_weight(value) / total_weight); + } + 0 + } else { + if value > self.centroids.last().unwrap().mean { + return Some(self.right_tail_weight(value) / total_weight); + } + self.centroids.partition_point(|c| c.mean < value) + }; + let (lower, interval) = self.interpolation_interval(value, index); + let weight_below = self.centroids[..lower].iter().map(|c| c.weight64()).sum(); + Some(interpolate_weight(value, interval, weight_below) / total_weight) } -} - -/// Checks the sequential validity of the given array of double values. -/// They must be unique, monotonically increasing and not NaN. -fn check_split_points(split_points: &[f64]) -> Result<(), Error> { - if split_points.iter().any(|split_point| split_point.is_nan()) { - return Err(Error::invalid_argument( - "split_points must not contain NaN values", - )); - } - if !split_points.windows(2).all(|pair| pair[0] < pair[1]) { - return Err(Error::invalid_argument( - "split_points must be unique and monotonically increasing", - )); + + #[inline(always)] + fn fill_cdf_binary_search(&self, split_points: &[f64], ranks: &mut Vec) { + let total_weight = self.centroids_weight as f64; + let mut previous_lower = 0; + let mut weight_below = 0; + for &value in split_points { + if value < self.min { + ranks.push(0.0); + continue; + } + if value > self.max { + ranks.push(1.0); + continue; + } + let index = if value <= self.centroids[0].mean { + if value < self.centroids[0].mean { + ranks.push(self.left_tail_weight(value) / total_weight); + continue; + } + 0 + } else { + if value > self.centroids.last().unwrap().mean { + ranks.push(self.right_tail_weight(value) / total_weight); + continue; + } + self.centroids.partition_point(|c| c.mean < value) + }; + let (lower, interval) = self.interpolation_interval(value, index); + weight_below += self.centroids[previous_lower..lower] + .iter() + .map(|c| c.weight64()) + .sum::(); + previous_lower = lower; + ranks.push(interpolate_weight(value, interval, weight_below) / total_weight); + } } - Ok(()) -} -fn check_rank(rank: f64) -> Result<(), Error> { - if !(0.0..=1.0).contains(&rank) { - return Err(Error::invalid_argument(format!( - "rank must be in [0.0, 1.0], got {rank}" - ))); + #[inline(always)] + fn fill_cdf_scan(&self, split_points: &[f64], ranks: &mut Vec) { + let total_weight = self.centroids_weight as f64; + let mut index = 0; + let mut weight_before_index = 0; + for &value in split_points { + let lower_bound = if value <= self.centroids[0].mean { + if value < self.centroids[0].mean { + ranks.push(self.left_tail_weight(value) / total_weight); + continue; + } + 0 + } else { + if value > self.centroids.last().unwrap().mean { + ranks.push(self.right_tail_weight(value) / total_weight); + continue; + } + while index < self.centroids.len() && self.centroids[index].mean < value { + weight_before_index += self.centroids[index].weight64(); + index += 1; + } + index + }; + let (lower, interval) = self.interpolation_interval(value, lower_bound); + // Between means, the lower anchor precedes index. Its weight belongs to the + // interpolation interval rather than the prefix below that interval. + let weight_below = weight_before_index + - if lower < lower_bound { + self.centroids[lower].weight64() + } else { + 0 + }; + ranks.push(interpolate_weight(value, interval, weight_below) / total_weight); + } } - Ok(()) -} -fn check_ranks(ranks: &[f64]) -> Result<(), Error> { - for &rank in ranks { - check_rank(rank)?; + #[inline] + fn left_tail_weight(&self, value: f64) -> f64 { + let center = centroid_weight_center(0, self.centroids[0].weight64()); + if value == self.min { + 0.5 + } else { + interpolate( + center.min(1.), + center, + interpolation_fraction(value, self.min, self.centroids[0].mean), + ) + } } - Ok(()) -} -fn centroid_cmp(a: &Centroid, b: &Centroid) -> Ordering { - match a.mean.partial_cmp(&b.mean) { - Some(order) => order, - None => unreachable!("NaN values should never be present in centroids"), + #[inline] + fn right_tail_weight(&self, value: f64) -> f64 { + let last = self.centroids.last().unwrap(); + let center = + centroid_weight_center(self.centroids_weight - last.weight64(), last.weight64()); + if value == self.max { + // Use the same rounded weight as the adjacent centroid center. + centroid_weight_center(self.centroids_weight - 1, 1) + } else { + interpolate( + center, + center.max((self.centroids_weight - 1) as f64), + interpolation_fraction(value, last.mean, self.max), + ) + } } -} -fn centroids_are_sorted(centroids: &[Centroid]) -> bool { - centroids - .windows(2) - .all(|pair| centroid_cmp(&pair[0], &pair[1]) != Ordering::Greater) -} + #[inline(always)] + fn interpolation_interval(&self, value: f64, index: usize) -> (usize, &[Centroid]) { + // An exact mean spans its entire equal run; otherwise adjacent centroids bracket it. + if self.centroids[index].mean == value { + let end = self.centroids.partition_point(|c| c.mean <= value); + (index, &self.centroids[index..end]) + } else { + (index - 1, &self.centroids[index - 1..=index]) + } + } -/// Whether the last output centroid can absorb the next input under the K_2 weight limit. -/// -/// `input_index` (zero-based) and `input_len` count input centroids in scan order. -/// `weight_before` sums the finalized output weights, excluding the last output centroid; -/// `proposed_weight` combines that centroid and the next input. `total_weight` sums all -/// input weights, and `normalizer` is the K_2 normalizer for compression `2 * sketch.k`. -/// -/// The [paper] limits a centroid with weight > 1 to a scale span of at most one (Section 2.3, -/// Eq. 4). As in the [reference implementation], we use a conservative weight limit instead of -/// evaluating that span: K_2's reciprocal slope, `q * (1 - q) / normalizer`, bounds the -/// allowed fraction of total weight. Both ends of the proposed rank interval must allow -/// the merge, so we take the smaller limit. This keeps centroids finer near the tails; -/// symmetry under `q -> 1 - q` also permits scanning in reverse. See Section 2.8, Eq. 8 for K_2. -/// -/// [paper]: https://arxiv.org/abs/1902.04023 -/// [reference implementation]: https://github.com/tdunning/t-digest/blob/eca1125a39c13e918d7054105dc89bcda11cfa44/core/src/main/java/com/tdunning/math/stats/MergingDigest.java#L421-L436 -fn should_merge_centroid( - input_index: usize, - input_len: usize, - weight_before: f64, - proposed_weight: f64, - total_weight: f64, - normalizer: f64, -) -> bool { - // Index 0 already seeded the output: refusing index 1 and the last index preserves - // the first and last input centroids without further merging. - if input_index == 1 || input_index == input_len - 1 { - return false; + #[inline] + fn quantile(&self, rank: f64) -> Option { + debug_assert!((0.0..=1.0).contains(&rank), "rank must be in [0.0, 1.0]"); + + if self.centroids.is_empty() { + return None; + } + + let mut quantile = [0.]; + self.fill_quantiles(std::iter::once((0, rank)), &mut quantile); + Some(quantile[0]) } - let q_start = weight_before / total_weight; - let q_end = (weight_before + proposed_weight) / total_weight; - let weight_limit = total_weight - * scale_function::max(q_start, normalizer).min(scale_function::max(q_end, normalizer)); - proposed_weight <= weight_limit -} -fn merge_sorted_centroids(left: &mut Vec, right: &[Centroid]) { - debug_assert!(!right.is_empty()); - debug_assert!(centroids_are_sorted(left)); - debug_assert!(centroids_are_sorted(right)); + 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]" + ); - let mut left_index = left.len(); - let mut right_index = right.len(); - let mut output_index = left_index + right_index; - left.reserve(right.len()); - left.resize(output_index, right[0]); + if self.centroids.is_empty() { + return None; + } - while left_index > 0 && right_index > 0 { - let left_centroid = left[left_index - 1]; - let right_centroid = right[right_index - 1]; - output_index -= 1; - // Taking the left side on ties while filling backward keeps the right side first in the - // final order, matching a stable sort after rotating the compressed left prefix. - if centroid_cmp(&left_centroid, &right_centroid) != Ordering::Less { - left_index -= 1; - left[output_index] = left_centroid; - } else { - right_index -= 1; - left[output_index] = right_centroid; + let mut quantiles = vec![0.; ranks.len()]; + if ranks.is_sorted() { + self.fill_quantiles(ranks.iter().copied().enumerate(), &mut quantiles); + return Some(quantiles); + } + + // The scan 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_unstable_by(|&left, &right| { + // Ranks have already been validated, so neither value can be NaN. + ranks[left].partial_cmp(&ranks[right]).unwrap() + }); + self.fill_quantiles( + rank_order.into_iter().map(|index| (index, ranks[index])), + &mut quantiles, + ); + Some(quantiles) + } + + /// Answers nondecreasing ranks at their original output indices. + fn fill_quantiles( + &self, + ranks: impl DoubleEndedIterator, + quantiles: &mut [f64], + ) { + if self.min == self.max { + quantiles.fill(self.min); + return; + } + + let centroids_weight = self.centroids_weight as f64; + let mut queries = ranks + .map(|(index, rank)| (index, rank * centroids_weight)) + .rev() + .peekable(); + + // Consume the right tail from the back, so the centroid scan can omit all tail checks. + while let Some((index, _)) = queries.next_if(|&(_, weight)| centroids_weight - weight < 1.) + { + quantiles[index] = self.max; + } + if queries.peek().is_none() { + return; + } + + let last = self.centroids.last().unwrap(); + let last_weight = last.weight64(); + let last_center = centroid_weight_center(self.centroids_weight - last_weight, last_weight); + if last_weight > 1 { + let tail_end = (self.centroids_weight - 1) as f64; + while let Some((index, weight)) = queries.next_if(|&(_, weight)| weight >= last_center) + { + // Integer counts can round a wider tail to zero width as well. At this + // boundary prefer the maximum, as for an exact two-sample tail. + quantiles[index] = if last_center == tail_end { + self.max + } else { + interpolate( + last.mean, + self.max, + (weight - last_center) / (tail_end - last_center), + ) + }; + } + } + + let mut queries = queries.rev().peekable(); + while let Some((index, _)) = queries.next_if(|&(_, weight)| weight < 1.) { + quantiles[index] = self.min; + } + if queries.peek().is_none() { + return; + } + + let first = &self.centroids[0]; + let first_weight = first.weight64(); + let first_center = centroid_weight_center(0, first_weight); + if first_weight > 1 { + // At first_weight == 2, the weight < 1 pre-pass already consumed this interval. + while let Some((index, weight)) = queries.next_if(|&(_, weight)| weight < first_center) + { + quantiles[index] = + interpolate(self.min, first.mean, (weight - 1.) / (first_center - 1.)); + } + } + if queries.peek().is_none() { + return; + } + + // Answer the remaining queries by centroid interval; both streams advance only forward. + let mut weight_before = 0; + let mut left_center = first_center; + for pair in self.centroids.windows(2) { + let left = &pair[0]; + let right = &pair[1]; + let right_before = weight_before + left.weight64(); + let right_center = centroid_weight_center(right_before, right.weight64()); + loop { + let Some(&(index, weight)) = queries.peek() else { + return; + }; + if right_center <= weight { + break; + } + queries.next(); + let mut start = left_center; + if left.weight64() == 1 { + start = right_before as f64; + if weight < start { + quantiles[index] = left.mean; + continue; + } + } + let mut end = right_center; + if right.weight64() == 1 { + end = right_before as f64; + if weight >= end { + quantiles[index] = right.mean; + continue; + } + } + quantiles[index] = + interpolate(left.mean, right.mean, (weight - start) / (end - start)); + } + weight_before = right_before; + left_center = right_center; + } + + // Rounding at large total weights can exhaust the scan near the maximum. + for (index, _) in queries { + quantiles[index] = self.max; + } + } +} + +/// Checks the sequential validity of the given array of double values. +/// They must be unique, monotonically increasing and not NaN. +fn check_split_points(split_points: &[f64]) -> Result<(), Error> { + if split_points.iter().any(|split_point| split_point.is_nan()) { + return Err(Error::invalid_argument( + "split_points must not contain NaN values", + )); + } + if !split_points.windows(2).all(|pair| pair[0] < pair[1]) { + return Err(Error::invalid_argument( + "split_points must be unique and monotonically increasing", + )); + } + Ok(()) +} + +fn check_rank(rank: f64) -> Result<(), Error> { + if !(0.0..=1.0).contains(&rank) { + return Err(Error::invalid_argument(format!( + "rank must be in [0.0, 1.0], got {rank}" + ))); + } + Ok(()) +} + +fn check_ranks(ranks: &[f64]) -> Result<(), Error> { + for &rank in ranks { + check_rank(rank)?; + } + Ok(()) +} + +fn centroid_cmp(a: &Centroid, b: &Centroid) -> Ordering { + match a.mean.partial_cmp(&b.mean) { + Some(order) => order, + None => unreachable!("NaN values should never be present in centroids"), + } +} + +fn centroids_are_sorted(centroids: &[Centroid]) -> bool { + centroids + .windows(2) + .all(|pair| centroid_cmp(&pair[0], &pair[1]) != Ordering::Greater) +} + +/// Whether the last output centroid can absorb the next input under the K_2 weight limit. +/// +/// `input_index` (zero-based) and `input_len` count input centroids in scan order. +/// `weight_before` sums the finalized output weights, excluding the last output centroid; +/// `proposed_weight` combines that centroid and the next input. `total_weight` sums all +/// input weights, and `normalizer` is the K_2 normalizer for compression `2 * sketch.k`. +/// +/// The [paper] limits a centroid with weight > 1 to a scale span of at most one (Section 2.3, +/// Eq. 4). As in the [reference implementation], we use a conservative weight limit instead of +/// evaluating that span: K_2's reciprocal slope, `q * (1 - q) / normalizer`, bounds the +/// allowed fraction of total weight. Both ends of the proposed rank interval must allow +/// the merge, so we take the smaller limit. This keeps centroids finer near the tails; +/// symmetry under `q -> 1 - q` also permits scanning in reverse. See Section 2.8, Eq. 8 for K_2. +/// +/// [paper]: https://arxiv.org/abs/1902.04023 +/// [reference implementation]: https://github.com/tdunning/t-digest/blob/eca1125a39c13e918d7054105dc89bcda11cfa44/core/src/main/java/com/tdunning/math/stats/MergingDigest.java#L421-L436 +fn should_merge_centroid( + input_index: usize, + input_len: usize, + weight_before: f64, + proposed_weight: f64, + total_weight: f64, + normalizer: f64, +) -> bool { + // Index 0 already seeded the output: refusing index 1 and the last index preserves + // the first and last input centroids without further merging. + if input_index == 1 || input_index == input_len - 1 { + return false; + } + let q_start = weight_before / total_weight; + let q_end = (weight_before + proposed_weight) / total_weight; + let weight_limit = total_weight + * scale_function::max(q_start, normalizer).min(scale_function::max(q_end, normalizer)); + proposed_weight <= weight_limit +} + +fn merge_sorted_centroids(left: &mut Vec, right: &[Centroid]) { + debug_assert!(!right.is_empty()); + debug_assert!(centroids_are_sorted(left)); + debug_assert!(centroids_are_sorted(right)); + + let mut left_index = left.len(); + let mut right_index = right.len(); + let mut output_index = left_index + right_index; + left.reserve(right.len()); + left.resize(output_index, right[0]); + + while left_index > 0 && right_index > 0 { + let left_centroid = left[left_index - 1]; + let right_centroid = right[right_index - 1]; + output_index -= 1; + // Taking the left side on ties while filling backward keeps the right side first in the + // final order, matching a stable sort after rotating the compressed left prefix. + if centroid_cmp(&left_centroid, &right_centroid) != Ordering::Less { + left_index -= 1; + left[output_index] = left_centroid; + } else { + right_index -= 1; + left[output_index] = right_centroid; + } + } + if right_index > 0 { + left[..right_index].copy_from_slice(&right[..right_index]); + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +struct Centroid { + mean: f64, + weight: NonZeroU64, +} + +impl Centroid { + fn merge(&mut self, other: Centroid) { + let (self_weight, other_weight) = (self.weight(), other.weight()); + let total_weight = self_weight + other_weight; + self.weight = self + .weight + .checked_add(other.weight64()) + .expect("weight overflow"); + + // Start at the heavier centroid so a small contribution is not rounded away in 1 - ratio. + let (start, end, fraction) = if self_weight >= other_weight { + (self.mean, other.mean, other_weight / total_weight) + } else { + (other.mean, self.mean, self_weight / total_weight) + }; + debug_assert!(start.is_finite() && end.is_finite() && fraction > 0. && fraction <= 0.5); + + // Same-sign means within a factor of two in magnitude subtract exactly. For wider + // gaps, a step of at most 1/2 cannot cross either endpoint even if subtraction rounds. + // Queries can interpolate farther and still need the general bounds. + self.mean = if start.is_sign_positive() == end.is_sign_positive() { + (end - start).mul_add(fraction, start) + } else if 1. - fraction == 1. { + // Both corrections can matter even when 1 - fraction rounds to 1. + // Halving before subtraction avoids overflow; the FMA combines the + // correction with start before rounding the result. + (end * 0.5 - start * 0.5).mul_add(2. * fraction, start) + } else { + interpolate(start, end, fraction) + }; + } + + fn weight(&self) -> f64 { + self.weight64() as f64 + } + + fn weight64(&self) -> u64 { + self.weight.get() + } +} + +/// Generates cluster sizes proportional to `q*(1-q)`. +/// +/// The use of a normalizing function results in a strictly bounded number of clusters no matter +/// how many samples. +/// +/// Corresponds to K_2 in the reference implementation +mod scale_function { + pub fn max(q: f64, normalizer: f64) -> f64 { + q * (1. - q) / normalizer + } + + pub fn normalizer(compression: f64, n: f64) -> f64 { + let z = 4. * (n / compression).ln() + 24.; + compression / z + } +} + +#[inline(always)] +fn interpolate_weight(value: f64, interval: &[Centroid], weight_below: u64) -> f64 { + let left = &interval[0]; + let Some((right, middle)) = interval[1..].split_last() else { + return centroid_weight_center(weight_below, left.weight64()); + }; + let weight_between = left.weight64() + middle.iter().map(|c| c.weight64()).sum::(); + let fraction = if left.mean < right.mean { + interpolation_fraction(value, left.mean, right.mean) + } else { + 0.5 + }; + interpolate( + centroid_weight_center(weight_below, left.weight64()), + centroid_weight_center(weight_below + weight_between, right.weight64()), + fraction, + ) +} + +/// Keeps cumulative counts exact until a centroid's midpoint is needed for interpolation. +/// All intervals and tails use this same center to avoid inconsistent boundary rounding. +fn centroid_weight_center(weight_before: u64, weight: u64) -> f64 { + (weight_before + weight / 2) as f64 + (weight % 2) as f64 * 0.5 +} + +/// Interpolates finite endpoints at a fraction in `[0, 1]`, preserving their bounds. +/// +/// The sign split follows WG21 P0811R3: +/// . +fn interpolate(start: f64, end: f64, fraction: f64) -> f64 { + debug_assert!(start.is_finite() && end.is_finite() && (0.0..=1.0).contains(&fraction)); + // Callers guarantee the fraction's range, but division may round it to an endpoint. + // Even an FMA need not reproduce `end` at 1: its subtracted difference already rounded. + if fraction <= 0. { + return start; + } + if fraction >= 1. { + return end; + } + if start.is_sign_positive() != end.is_sign_positive() { + // Avoid the potentially overflowing difference of opposite-signed endpoints. + start * (1. - fraction) + end * fraction + } else { + // Same-sign subtraction is finite; rounding must not overshoot the end. + let value = (end - start).mul_add(fraction, start); + if start < end { + value.min(end) + } else { + value.max(end) + } + } +} + +/// Locates a value between distinct finite endpoints. +fn interpolation_fraction(value: f64, start: f64, end: f64) -> f64 { + let width = end - start; + if width.is_finite() { + (value - start) / width + } else { + // Scaling is needed only when opposite-signed endpoints overflow their difference. + (value * 0.5 - start * 0.5) / (end * 0.5 - start * 0.5) + } +} + +const PREAMBLE_LONGS_EMPTY_OR_SINGLE: u8 = 1; +const PREAMBLE_LONGS_MULTIPLE: u8 = 2; +const SERIAL_VERSION: u8 = 1; +const FLAGS_IS_EMPTY: u8 = 1 << 0; +const FLAGS_IS_SINGLE_VALUE: u8 = 1 << 1; +const FLAGS_REVERSE_MERGE: u8 = 1 << 2; +/// the format of the reference implementation is using double (f64) precision +const COMPAT_DOUBLE: u32 = 1; +/// the format of the reference implementation is using float (f32) precision +const COMPAT_FLOAT: u32 = 2; + +fn serialize_compressed( + k: u16, + reverse_merge: bool, + min: f64, + max: f64, + centroids: &[Centroid], + total_weight: u64, +) -> Vec { + let is_empty = centroids.is_empty(); + let is_single_value = total_weight == 1; + let mut total_size = if is_empty || is_single_value { + // Preamble, serial version, family, k, flags, and two unused bytes. + size_of::() + } else { + // The short header plus centroid and buffered-value counts. + size_of::() * 2 + }; + if is_single_value { + total_size += size_of::(); + } else if !is_empty { + total_size += size_of::() * 2; + total_size += centroids.len() * (size_of::() + size_of::()); + } + + let mut bytes = SketchBytes::with_capacity(total_size); + bytes.write_u8(if is_empty || is_single_value { + PREAMBLE_LONGS_EMPTY_OR_SINGLE + } else { + PREAMBLE_LONGS_MULTIPLE + }); + bytes.write_u8(SERIAL_VERSION); + bytes.write_u8(Family::TDIGEST.id); + bytes.write_u16_le(k); + bytes.write_u8({ + let mut flags = 0; + if is_empty { + flags |= FLAGS_IS_EMPTY; + } + if is_single_value { + flags |= FLAGS_IS_SINGLE_VALUE; + } + if reverse_merge { + flags |= FLAGS_REVERSE_MERGE; + } + flags + }); + bytes.write_u16_le(0); // unused + if is_empty { + return bytes.into_bytes(); + } + if is_single_value { + bytes.write_f64_le(min); + return bytes.into_bytes(); + } + bytes.write_u32_le(centroids.len() as u32); + bytes.write_u32_le(0); // no buffered values + bytes.write_f64_le(min); + bytes.write_f64_le(max); + for centroid in centroids { + bytes.write_f64_le(centroid.mean); + bytes.write_u64_le(centroid.weight64()); + } + bytes.into_bytes() +} + +impl TDigestMut { + fn deserialize_impl(bytes: &[u8], is_f32: bool) -> Result { + let mut cursor = SketchSlice::new(bytes); + + let preamble_longs = cursor + .read_u8() + .map_err(insufficient_data("preamble_longs"))?; + let serial_version = cursor + .read_u8() + .map_err(insufficient_data("serial_version"))?; + let family_id = cursor.read_u8().map_err(insufficient_data("family_id"))?; + if let Err(err) = Family::TDIGEST.validate_id(family_id) { + return if preamble_longs == 0 && serial_version == 0 && family_id == 0 { + Self::deserialize_compat(bytes) + } else { + Err(err) + }; + } + ensure_serial_version_is(SERIAL_VERSION, serial_version)?; + let k = cursor.read_u16_le().map_err(insufficient_data("k"))?; + if k < 10 { + 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 { + PREAMBLE_LONGS_MULTIPLE + }; + ensure_preamble_longs_in(&[expected_preamble_longs], preamble_longs)?; + cursor + .read_u16_le() + .map_err(insufficient_data(""))?; // unused + if is_empty { + return TDigestMut::new(k); + } + + let reverse_merge = (flags & FLAGS_REVERSE_MERGE) != 0; + if is_single_value { + let value = if is_f32 { + cursor + .read_f32_le() + .map_err(insufficient_data("single_value"))? as f64 + } else { + cursor + .read_f64_le() + .map_err(insufficient_data("single_value"))? + }; + check_non_nan(value, "single_value")?; + check_finite(value, "single_value")?; + return Ok(TDigestMut::from_parts( + k, + reverse_merge, + value, + value, + TDigestBuffer::new( + vec![Centroid { + mean: value, + weight: DEFAULT_WEIGHT, + }], + 0, + ), + 1, + )); + } + let num_centroids = cursor + .read_u32_le() + .map_err(insufficient_data("num_centroids"))? as usize; + let num_buffered = cursor + .read_u32_le() + .map_err(insufficient_data("num_buffered"))? as usize; + let (min, max) = if is_f32 { + ( + cursor.read_f32_le().map_err(insufficient_data("min"))? as f64, + cursor.read_f32_le().map_err(insufficient_data("max"))? as f64, + ) + } else { + ( + cursor.read_f64_le().map_err(insufficient_data("min"))?, + cursor.read_f64_le().map_err(insufficient_data("max"))?, + ) + }; + check_extrema(min, max, "TDigest")?; + let (centroid_bytes, buffered_value_bytes) = if is_f32 { + (size_of::() + size_of::(), size_of::()) + } else { + (size_of::() + size_of::(), size_of::()) + }; + let centroid_payload_bytes = num_centroids + .checked_mul(centroid_bytes) + .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; + let buffered_payload_bytes = num_buffered + .checked_mul(buffered_value_bytes) + .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; + let required_payload_bytes = centroid_payload_bytes + .checked_add(buffered_payload_bytes) + .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; + // Check the whole payload once so fixed-width records can be decoded without per-field I/O. + let available_bytes = cursor.remaining().len(); + if available_bytes < required_payload_bytes { + return Err(Error::insufficient_data_of( + "TDigest payload", + format_args!("expected {required_payload_bytes} bytes, got {available_bytes}"), + )); + } + let payload = &cursor.remaining()[..required_payload_bytes]; + let (centroid_payload, buffered_payload) = payload.split_at(centroid_payload_bytes); + 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 { + ( + f32::from_le_bytes(bytes[..4].try_into().unwrap()) as f64, + u32::from_le_bytes(bytes[4..].try_into().unwrap()) as u64, + ) + } else { + ( + f64::from_le_bytes(bytes[..8].try_into().unwrap()), + u64::from_le_bytes(bytes[8..].try_into().unwrap()), + ) + }; + 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", + )); + } + let total_weight = checked_weight_sum(compressed_weight, num_buffered as u64)?; + check_single_sample_extrema(min, max, total_weight)?; + 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); + 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::from_parts( + k, + reverse_merge, + min, + max, + TDigestBuffer::new(centroids, num_buffered), + compressed_weight, + )) } - if right_index > 0 { - left[..right_index].copy_from_slice(&right[..right_index]); - } -} - -#[derive(Debug, Clone, Copy, PartialEq)] -struct Centroid { - mean: f64, - weight: NonZeroU64, -} - -impl Centroid { - fn add(&mut self, other: Centroid) { - let (self_weight, other_weight) = (self.weight(), other.weight()); - let total_weight = self_weight + other_weight; - self.weight = self - .weight - .checked_add(other.weight.get()) - .expect("weight overflow"); - // Start at the heavier centroid so a small contribution is not rounded away in 1 - ratio. - let (start, end, fraction) = if self_weight >= other_weight { - (self.mean, other.mean, other_weight / total_weight) - } else { - (other.mean, self.mean, self_weight / total_weight) - }; - debug_assert!(start.is_finite() && end.is_finite() && fraction > 0. && fraction <= 0.5); + // compatibility with the format of the reference implementation + // default byte order of ByteBuffer is used there, which is big endian + fn deserialize_compat(bytes: &[u8]) -> Result { + fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { + move |_| Error::insufficient_data_of("compat format", tag) + } - // Same-sign means within a factor of two in magnitude subtract exactly. For wider - // gaps, a step of at most 1/2 cannot cross either endpoint even if subtraction rounds. - // Queries can interpolate farther and still need the general bounds. - self.mean = if start.is_sign_positive() == end.is_sign_positive() { - (end - start).mul_add(fraction, start) - } else if 1. - fraction == 1. { - // Both corrections can matter even when 1 - fraction rounds to 1. - // Halving before subtraction avoids overflow; the FMA combines the - // correction with start before rounding the result. - (end * 0.5 - start * 0.5).mul_add(2. * fraction, start) - } else { - interpolate(start, end, fraction) - }; - } + let mut cursor = SketchSlice::new(bytes); - fn weight(&self) -> f64 { - self.weight.get() as f64 + let ty = cursor.read_u32_be().map_err(make_error("type"))?; + match ty { + COMPAT_DOUBLE => { + fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { + move |_| Error::insufficient_data_of("compat double format", tag) + } + // compatibility with asBytes() + let min = cursor.read_f64_be().map_err(make_error("min"))?; + let max = cursor.read_f64_be().map_err(make_error("max"))?; + check_extrema(min, max, "compat double TDigest")?; + let k = cursor.read_f64_be().map_err(make_error("k"))? as u16; + if k < 10 { + return Err(Error::deserial(format!( + "k must be at least 10 in compat double format, got {k}" + ))); + } + let num_centroids = + cursor.read_u32_be().map_err(make_error("num_centroids"))? as usize; + if num_centroids == 0 { + return Err(Error::deserial( + "malformed data: compat double TDigest must contain a centroid", + )); + } + let mut total_weight = 0u64; + let mut centroids = Vec::with_capacity(num_centroids); + let mut previous_mean = min; + let mut centroid_means_valid = true; + for _ in 0..num_centroids { + let weight = cursor.read_f64_be().map_err(make_error("weight"))?; + let mean = cursor.read_f64_be().map_err(make_error("mean"))?; + let weight = + check_compat_weight(weight, "centroid weight in compat double format")?; + centroid_means_valid &= + mean.is_finite() & (mean >= previous_mean) & (mean <= max); + previous_mean = mean; + total_weight = checked_weight_sum(total_weight, weight.get())?; + centroids.push(Centroid { mean, weight }); + } + if !centroid_means_valid { + return Err(Error::deserial( + "malformed data: centroid means in compat double format must be finite, within extrema, and nondecreasing", + )); + } + check_single_sample_extrema(min, max, total_weight)?; + Ok(TDigestMut::from_parts( + k, + false, + min, + max, + TDigestBuffer::new(centroids, 0), + total_weight, + )) + } + COMPAT_FLOAT => { + fn make_error(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { + move |_| Error::insufficient_data_of("compat float format", tag) + } + // COMPAT_FLOAT: compatibility with asSmallBytes() + // reference implementation uses doubles for min and max + let min = cursor.read_f64_be().map_err(make_error("min"))?; + let max = cursor.read_f64_be().map_err(make_error("max"))?; + check_extrema(min, max, "compat float TDigest")?; + let k = cursor.read_f32_be().map_err(make_error("k"))? as u16; + if k < 10 { + return Err(Error::deserial(format!( + "k must be at least 10 in compat float format, got {k}" + ))); + } + // reference implementation stores capacities of the array of centroids and the + // buffer as shorts they can be derived from k in the constructor + cursor.read_u32_be().map_err(make_error(""))?; + let num_centroids = + cursor.read_u16_be().map_err(make_error("num_centroids"))? as usize; + if num_centroids == 0 { + return Err(Error::deserial( + "malformed data: compat float TDigest must contain a centroid", + )); + } + let mut total_weight = 0u64; + let mut centroids = Vec::with_capacity(num_centroids); + let mut previous_mean = min; + let mut centroid_means_valid = true; + for _ in 0..num_centroids { + let weight = cursor.read_f32_be().map_err(make_error("weight"))? as f64; + let mean = cursor.read_f32_be().map_err(make_error("mean"))? as f64; + let weight = + check_compat_weight(weight, "centroid weight in compat float format")?; + centroid_means_valid &= + mean.is_finite() & (mean >= previous_mean) & (mean <= max); + previous_mean = mean; + total_weight = checked_weight_sum(total_weight, weight.get())?; + centroids.push(Centroid { mean, weight }); + } + if !centroid_means_valid { + return Err(Error::deserial( + "malformed data: centroid means in compat float format must be finite, within extrema, and nondecreasing", + )); + } + check_single_sample_extrema(min, max, total_weight)?; + Ok(TDigestMut::from_parts( + k, + false, + min, + max, + TDigestBuffer::new(centroids, 0), + total_weight, + )) + } + ty => Err(Error::deserial(format!("unknown TDigest compat type {ty}"))), + } } } @@ -2242,71 +2314,6 @@ fn checked_weight_sum(total_weight: u64, weight: u64) -> Result { .ok_or_else(|| Error::deserial("malformed data: total weight overflow")) } -/// Generates cluster sizes proportional to `q*(1-q)`. -/// -/// The use of a normalizing function results in a strictly bounded number of clusters no matter -/// how many samples. -/// -/// Corresponds to K_2 in the reference implementation -mod scale_function { - pub fn max(q: f64, normalizer: f64) -> f64 { - q * (1. - q) / normalizer - } - - pub fn normalizer(compression: f64, n: f64) -> f64 { - compression / z(compression, n) - } - - pub fn z(compression: f64, n: f64) -> f64 { - 4. * (n / compression).ln() + 24. - } -} - -/// Keeps cumulative counts exact until a centroid's midpoint is needed for interpolation. -/// All intervals and tails use this same center to avoid inconsistent boundary rounding. -fn centroid_center(weight_before: u64, weight: u64) -> f64 { - (weight_before + weight / 2) as f64 + (weight % 2) as f64 * 0.5 -} - -/// Interpolates finite endpoints at a fraction in `[0, 1]`, preserving their bounds. -/// -/// The sign split follows WG21 P0811R3: -/// . -fn interpolate(start: f64, end: f64, fraction: f64) -> f64 { - debug_assert!(start.is_finite() && end.is_finite() && (0.0..=1.0).contains(&fraction)); - // Callers guarantee the fraction's range, but division may round it to an endpoint. - // Even an FMA need not reproduce `end` at 1: its subtracted difference already rounded. - if fraction <= 0. { - return start; - } - if fraction >= 1. { - return end; - } - if start.is_sign_positive() != end.is_sign_positive() { - // Avoid the potentially overflowing difference of opposite-signed endpoints. - start * (1. - fraction) + end * fraction - } else { - // Same-sign subtraction is finite; rounding must not overshoot the end. - let value = (end - start).mul_add(fraction, start); - if start < end { - value.min(end) - } else { - value.max(end) - } - } -} - -/// Locates a value between distinct finite endpoints. -fn interpolation_fraction(value: f64, start: f64, end: f64) -> f64 { - let width = end - start; - if width.is_finite() { - (value - start) / width - } else { - // Scaling is needed only when opposite-signed endpoints overflow their difference. - (value * 0.5 - start * 0.5) / (end * 0.5 - start * 0.5) - } -} - #[cfg(test)] mod tests { use super::interpolate;