From 7346058c81716a491256b391e9596ab65d5eb9fe Mon Sep 17 00:00:00 2001 From: jonathan keinan Date: Tue, 13 Oct 2026 19:45:26 +0300 Subject: [PATCH] feat: add absorb for in-place merges of buffered digests Folding many small digests into one means flushing each and calling merge_digests, which sorts, compresses and allocates per digest even when it holds only a few dozen values. TDigest::absorb merges another digest in place. If the other digest has never flushed, its buffered values are pushed into the receiver, exactly as if they had been pushed there directly. Otherwise both are flushed and combined with merge_digests. The result keeps the larger max_size. Adds unit and property tests for both paths and an absorb_small_digests bench comparing absorb with flush plus merge_digests. Co-Authored-By: Claude Opus 5.5 --- README.md | 4 ++ benches/tdigest.rs | 45 +++++++++++++++++ docs/architecture.md | 7 +++ src/lib.rs | 112 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 168 insertions(+) diff --git a/README.md b/README.md index 19ce0c8..b86107d 100644 --- a/README.md +++ b/README.md @@ -112,6 +112,10 @@ before `estimate_quantile`, `quantiles`, `estimate_rank`, `trimmed_mean`, `centroids`, immutable merges, or serialization. Automatic compression happens when the buffer fills, so a partially filled buffer still needs an explicit flush. +`absorb` merges another digest in place. A digest that has only buffered values +is pushed into the receiver, so folding many small digests into one costs the +same as pushing their values directly. + Collecting an iterator creates a digest with the default compression factor and flushes it before returning: diff --git a/benches/tdigest.rs b/benches/tdigest.rs index 5043f1e..e7268b6 100644 --- a/benches/tdigest.rs +++ b/benches/tdigest.rs @@ -119,6 +119,50 @@ fn merge_digests(c: &mut Criterion) { group.finish(); } +fn absorb_small_digests(c: &mut Criterion) { + let input = values(10_000, 70_000); + let buffered: Vec = input + .chunks(64) + .map(|chunk| { + let mut digest = TDigest::new_with_size(MAX_SIZE); + digest.extend_values(chunk.iter().copied()); + digest + }) + .collect(); + let mut group = c.benchmark_group("absorb_small_digests"); + group.sample_size(10).measurement_time(Duration::from_secs(2)); + group.throughput(Throughput::Elements(input.len() as u64)); + group.bench_function("absorb", |b| { + b.iter_batched( + || buffered.clone(), + |parts| { + let mut total = TDigest::new_with_size(MAX_SIZE); + for part in parts { + total.absorb(part); + } + total.flush(); + total + }, + BatchSize::SmallInput, + ) + }); + group.bench_function("flush_and_merge_digests", |b| { + b.iter_batched( + || buffered.clone(), + |parts| { + let mut total = TDigest::new_with_size(MAX_SIZE); + for mut part in parts { + part.flush(); + total = TDigest::merge_digests(vec![total, part]); + } + total + }, + BatchSize::SmallInput, + ) + }); + group.finish(); +} + fn estimate_quantile(c: &mut Criterion) { let digest = TDigest::new_with_size(MAX_SIZE).merge_unsorted(values(100_000, 30_000)); let mut group = c.benchmark_group("estimate_quantile"); @@ -162,6 +206,7 @@ criterion_group!( small_batch_ingest, buffered_flush, merge_digests, + absorb_small_digests, estimate_quantile, max_size_sensitivity ); diff --git a/docs/architecture.md b/docs/architecture.md index 5f5d428..b8bf54d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -212,6 +212,12 @@ assertions to catch an unflushed buffer. Serialization does not check it, and omits pending values while retaining their summary statistics. Violating the flush contract can therefore produce incomplete or inconsistent data. +`absorb` is the mutable counterpart of `merge_digests` for two digests. If the +absorbed digest has never flushed, its buffered values are pushed into the +receiver, which is exactly equivalent to pushing them there directly. Otherwise +both digests are flushed and merged by `merge_digests`. The result keeps the +larger `max_size`. + `FromIterator` creates a default-size digest, ingests the iterator through the buffer, and flushes before returning. `flush` is a no-op when no values are pending. Query methods take `&self` and do not flush automatically; the type @@ -230,6 +236,7 @@ requested quantiles. The final centroid sort is included in these bounds. | `push` / `extend_values` | Amortized `O(log m)` per value | `O(m)` buffer and compression result | | `flush` | `O(m log m)` when the buffer is full | `O(m)` compression result; reuses the buffer | | `merge_digests` | `O(C log C)` | `O(C + m)` collected centroids and result | +| `absorb` | As `push` per value if the other digest is unflushed; otherwise as `merge_digests` | As `push`, or as `merge_digests` | | `estimate_quantile` / `estimate_rank` / `trimmed_mean` | `O(c)` | `O(1)` | | `quantiles` | `O(c + k log c)` | `O(c + k)` cumulative weights and output | | `mean` | `O(1)` normally; `O(c + m)` if the sum is non-finite | `O(1)` | diff --git a/src/lib.rs b/src/lib.rs index d987c8f..a9456eb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -555,6 +555,47 @@ impl TDigest { *self = result; } + /// Merge `other` into this digest in place, using the mutable ingestion API. + /// + /// If `other` holds only buffered values (it has never flushed), they are + /// pushed into `self`. That is exactly equivalent to pushing them here + /// directly, and avoids compressing `other` and merging centroids, which + /// dominates when folding many small digests into one. Otherwise both + /// digests are flushed and combined as by [`TDigest::merge_digests`]. + /// + /// The result uses the larger `max_size` of the two digests. As with + /// [`TDigest::push`], call [`TDigest::flush`] before queries. + /// + /// ``` + /// use tdigest::TDigest; + /// + /// let mut total = TDigest::new_with_size(100); + /// for chunk in [[1.0, 2.0], [3.0, 4.0]] { + /// let mut part = TDigest::new_with_size(100); + /// part.extend_values(chunk); + /// total.absorb(part); + /// } + /// total.flush(); + /// assert_eq!(total.count(), 4.0); + /// assert_eq!(total.estimate_quantile(0.5), Some(2.5)); + /// ``` + pub fn absorb(&mut self, mut other: TDigest) { + if other.count == 0.0 { + return; + } + self.max_size = self.max_size.max(other.max_size); + if other.centroids.is_empty() { + for value in mem::take(&mut other.buffer) { + self.push(value); + } + return; + } + self.flush(); + other.flush(); + let this = mem::replace(self, TDigest::new_with_size(self.max_size)); + *self = TDigest::merge_digests(vec![this, other]); + } + /// Merge unsorted values into a new digest, sorting them internally. /// /// ``` @@ -1403,6 +1444,60 @@ mod tests { assert_eq!(digest.max(), Some((capacity + 1) as f64)); } + fn buffered(values: impl IntoIterator, max_size: usize) -> TDigest { + let mut digest = TDigest::new_with_size(max_size); + digest.extend_values(values); + digest + } + + #[test] + fn test_absorb_buffered_matches_direct_push() { + let values: Vec = (0..10_000).map(|i| ((i * 7919) % 10_007) as f64 - 5_000.0).collect(); + let mut direct = TDigest::new_with_size(50); + let mut absorbed = TDigest::new_with_size(50); + for chunk in values.chunks(63) { + direct.extend_values(chunk.iter().copied()); + absorbed.absorb(buffered(chunk.iter().copied(), 50)); + } + assert_eq!(absorbed, direct); + } + + #[test] + fn test_absorb_flushed_matches_merge_digests() { + let mut left = buffered((0..1_000).map(f64::from), 100); + let mut right = buffered((500..2_500).map(f64::from), 100); + right.flush(); + let mut absorbed = left.clone(); + absorbed.absorb(right.clone()); + left.flush(); + assert_eq!(absorbed, TDigest::merge_digests(vec![left, right])); + } + + #[test] + fn test_absorb_empty_and_into_empty() { + let mut digest = buffered([1.0, 2.0, 3.0], 100); + let expected = digest.clone(); + digest.absorb(TDigest::new_with_size(100)); + assert_eq!(digest, expected); + + let mut empty = TDigest::new_with_size(100); + empty.absorb(expected.clone()); + assert_eq!(empty, expected); + } + + #[test] + fn test_absorb_keeps_larger_max_size() { + let mut small = buffered([1.0, 2.0], 10); + small.absorb(buffered([3.0], 200)); + assert_eq!(small.max_size(), 200); + + let mut large = buffered([1.0, 2.0], 200); + let mut flushed = buffered([3.0], 10); + flushed.flush(); + large.absorb(flushed); + assert_eq!(large.max_size(), 200); + } + #[test] fn test_mutable_collection_traits() { let mut extended = TDigest::default(); @@ -1581,6 +1676,23 @@ mod tests { prop_assert!((estimate - expected).abs() <= scale * 0.1); } + #[test] + fn absorbing_buffered_chunks_matches_direct_push( + values in finite_values(), + // Smaller than the flush threshold, so each part stays buffered. + chunk_size in 1usize..BUFFER_FACTOR * 20, + ) { + let mut direct = TDigest::new_with_size(20); + let mut absorbed = TDigest::new_with_size(20); + for chunk in values.chunks(chunk_size) { + direct.extend_values(chunk.iter().copied()); + let mut part = TDigest::new_with_size(20); + part.extend_values(chunk.iter().copied()); + absorbed.absorb(part); + } + prop_assert_eq!(absorbed, direct); + } + #[test] fn rank_quantile_round_trip( values in finite_values(),