This is an automated email from the ASF dual-hosted git repository.
tisonkun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datasketches-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 0cf3cf6 refactor(tdigest): organize public APIs and internal helpers
(#313)
0cf3cf6 is described below
commit 0cf3cf6a8ca0d55c9a2d25197fcadb56976223a2
Author: tison <[email protected]>
AuthorDate: Wed Oct 28 00:28:06 2026 +0800
refactor(tdigest): organize public APIs and internal helpers (#313)
---
datasketches/src/tdigest/mod.rs | 2 -
datasketches/src/tdigest/serialization.rs | 27 -
datasketches/src/tdigest/sketch.rs | 2703 +++++++++++++++--------------
3 files changed, 1355 insertions(+), 1377 deletions(-)
diff --git a/datasketches/src/tdigest/mod.rs b/datasketches/src/tdigest/mod.rs
index 8c5818b..652ba81 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 debe678..0000000
--- 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 83d8d2d..e665a10 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<Centroid>,
- unmerged_tail_len: usize,
-}
-
-impl TDigestBuffer {
- fn new(centroids: Vec<Centroid>, 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<Centroid> {
- 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<Centroid>
{
- 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<Centroid> {
- 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::<Centroid>()
- }
-}
-
/// 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<Self, Error> {
- 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>"))?; // 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::<f32>() + size_of::<u32>(), size_of::<f32>())
- } else {
- (size_of::<f64>() + size_of::<u64>(), size_of::<f64>())
- };
- 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<Self, Error> {
- 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("<unused>"))?;
- 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::<TDigestMut>();
-/// 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::<Result<TDigestMut, _>>()
-/// .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<TDigestMut> for TDigestMut {
- fn from_iter<T: IntoIterator<Item = TDigestMut>>(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<Centroid>,
- 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<Centroid>,
+ 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<f64> {
+ 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<Centroid> =
- 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<f64> {
+ 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::<u64>(),
- 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<u8> {
- 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::<u64>()
- } else {
- // The short header plus centroid and buffered-value counts.
- size_of::<u64>() * 2
- };
- if is_single_value {
- total_size += size_of::<f64>();
- } else if !is_empty {
- total_size += size_of::<f64>() * 2;
- total_size += centroids.len() * (size_of::<f64>() + size_of::<u64>());
- }
-
- 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<Centroid>,
- 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<f64> {
- 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<f64> {
- 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::<Centroid>()
}
-}
-struct TDigestView<'a> {
- min: f64,
- max: f64,
- centroids: &'a [Centroid],
- centroids_weight: u64,
-}
-
-impl TDigestView<'_> {
- fn pmf(&self, split_points: &[f64]) -> Option<Vec<f64>> {
- 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<Vec<f64>> {
- 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::<TDigestMut>();
+/// 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::<Result<TDigestMut, _>>()
+/// .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<TDigestMut> for TDigestMut {
+ fn from_iter<T: IntoIterator<Item = TDigestMut>>(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<Centroid>,
+ unmerged_len: usize,
+ input_index: usize,
}
- ranks.push(1.0);
- Some(ranks)
- }
-
- #[inline(always)]
- fn rank(&self, value: f64) -> Option<f64> {
- 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<f64>) {
- 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::<u64>();
- 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<f64>) {
- 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<Centroid> =
+ 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::<u64>(),
+ 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<Centroid>,
+ unmerged_tail_len: usize,
+}
+
+impl TDigestBuffer {
+ fn new(centroids: Vec<Centroid>, 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::<u64>();
- 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<f64> {
- 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<Vec<f64>> {
- 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<Centroid> {
+ 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<Centroid>
{
+ 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::<Vec<_>>();
- 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<Item = (usize, f64)>,
- quantiles: &mut [f64],
- ) {
- if self.min == self.max {
- quantiles.fill(self.min);
- return;
- }
+ fn into_compressed_centroids(self) -> Vec<Centroid> {
+ 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::<Centroid>()
+ }
+}
- // 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<Vec<f64>> {
+ 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<Vec<f64>> {
+ 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<f64> {
+ 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<f64>) {
+ 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::<u64>();
+ 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<f64>) {
+ 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<f64> {
+ 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<Centroid>, 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<Vec<f64>> {
+ 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::<Vec<_>>();
+ 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<Item = (usize, f64)>,
+ 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<Centroid>, 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::<u64>();
+ 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:
+/// <https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p0811r3.html>.
+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<u8> {
+ 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::<u64>()
+ } else {
+ // The short header plus centroid and buffered-value counts.
+ size_of::<u64>() * 2
+ };
+ if is_single_value {
+ total_size += size_of::<f64>();
+ } else if !is_empty {
+ total_size += size_of::<f64>() * 2;
+ total_size += centroids.len() * (size_of::<f64>() + size_of::<u64>());
+ }
+
+ 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<Self, Error> {
+ 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>"))?; // 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::<f32>() + size_of::<u32>(), size_of::<f32>())
+ } else {
+ (size_of::<f64>() + size_of::<u64>(), size_of::<f64>())
+ };
+ 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<Self, Error> {
+ 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("<unused>"))?;
+ 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<u64, Error> {
.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:
-/// <https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p0811r3.html>.
-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;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
