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 618ea2e feat(tdigest): Add owned batch merging (#276)
618ea2e is described below
commit 618ea2e842b1029eb77e741a75f923e59cb3018c
Author: Ryu <[email protected]>
AuthorDate: Sun Oct 4 07:43:18 2026 +0200
feat(tdigest): Add owned batch merging (#276)
Co-authored-by: tison <[email protected]>
---
CHANGELOG.md | 1 +
benchmarks/tdigest/aggregation.rs | 130 +++++++++++
benchmarks/tdigest/merge.rs | 125 +++++++++-
benchmarks/tdigest/mod.rs | 1 +
datasketches/src/tdigest/mod.rs | 2 +-
datasketches/src/tdigest/sketch.rs | 303 +++++++++++++++++++++++--
tests-integration/tests/tdigest_test/sketch.rs | 240 ++++++++++++++++++++
7 files changed, 784 insertions(+), 18 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 524755e..24934e7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ All significant changes to this project will be documented in
this file.
* `KllSketch` is now available behind the `kll` feature, with rank, quantile,
PMF, and CDF queries, merging, serialization, custom ordered item types, and a
`KllFloat` adapter for non-NaN floating-point values.
* Add `TDigestMut::quantiles` and `TDigest::quantiles` for querying several
ranks in one centroid scan while preserving input order.
+* Add `FromIterator<TDigestMut>` for batch construction from owned partial
sketches, avoiding repeated compression of intermediate results.
### Improvements
diff --git a/benchmarks/tdigest/aggregation.rs
b/benchmarks/tdigest/aggregation.rs
new file mode 100644
index 0000000..da17069
--- /dev/null
+++ b/benchmarks/tdigest/aggregation.rs
@@ -0,0 +1,130 @@
+// 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.
+
+use datasketches::tdigest::TDigestMut;
+use divan::Bencher;
+use divan::black_box;
+use divan::counter::ItemsCount;
+
+use super::support::build_mut_digest;
+use super::support::values;
+
+// (number of partial states, values per state). Each partial samples the same
distribution.
+const SERIALIZED_INPUTS: &[(usize, usize)] =
+ &[(1, 8), (64, 8), (512, 8), (64, 64), (512, 64), (64, 4_096)];
+
+#[divan::bench(args = SERIALIZED_INPUTS, sample_size = 1)]
+fn serialized_borrowed(bencher: Bencher, shape: (usize, usize)) {
+ let (num_partials, rows_per_partial) = shape;
+ let partials = serialize_partials(&values(num_partials *
rows_per_partial), rows_per_partial);
+
+ bencher
+ .counter(ItemsCount::new(num_partials * rows_per_partial))
+ .bench_local(|| {
+ let mut merged = TDigestMut::default();
+ for bytes in black_box(&partials) {
+ merged.merge(&TDigestMut::deserialize(bytes).unwrap());
+ }
+ final_quantile(merged)
+ });
+}
+
+#[divan::bench(args = SERIALIZED_INPUTS, sample_size = 1)]
+fn serialized_owned(bencher: Bencher, shape: (usize, usize)) {
+ let (num_partials, rows_per_partial) = shape;
+ let partials = serialize_partials(&values(num_partials *
rows_per_partial), rows_per_partial);
+
+ bencher
+ .counter(ItemsCount::new(num_partials * rows_per_partial))
+ .bench_local(|| {
+ let merged = black_box(&partials)
+ .iter()
+ .map(|bytes| TDigestMut::deserialize(bytes))
+ .collect::<Result<TDigestMut, _>>()
+ .unwrap();
+ final_quantile(merged)
+ });
+}
+
+#[divan::bench(args = SERIALIZED_INPUTS, sample_size = 1)]
+fn serialized_in_batches(bencher: Bencher, shape: (usize, usize)) {
+ let (num_partials, rows_per_partial) = shape;
+ let partials = serialize_partials(&values(num_partials *
rows_per_partial), rows_per_partial);
+
+ bencher
+ .counter(ItemsCount::new(num_partials * rows_per_partial))
+ .bench_local(|| {
+ let mut merged = TDigestMut::default();
+ for batch in black_box(&partials).chunks(16) {
+ let batch = batch
+ .iter()
+ .map(|bytes| TDigestMut::deserialize(bytes))
+ .collect::<Result<TDigestMut, _>>()
+ .unwrap();
+ merged.merge(&batch);
+ }
+ final_quantile(merged)
+ });
+}
+
+#[divan::bench(args = [8, 64], sample_size = 1)]
+fn end_to_end_borrowed(bencher: Bencher, rows_per_partial: usize) {
+ let inputs = values(64 * rows_per_partial);
+
+ bencher
+ .counter(ItemsCount::new(inputs.len()))
+ .bench_local(|| {
+ let partials = serialize_partials(black_box(&inputs),
rows_per_partial);
+ let mut merged = TDigestMut::default();
+ for bytes in &partials {
+ merged.merge(&TDigestMut::deserialize(bytes).unwrap());
+ }
+ final_quantile(merged)
+ });
+}
+
+#[divan::bench(args = [8, 64], sample_size = 1)]
+fn end_to_end_owned(bencher: Bencher, rows_per_partial: usize) {
+ let inputs = values(64 * rows_per_partial);
+
+ bencher
+ .counter(ItemsCount::new(inputs.len()))
+ .bench_local(|| {
+ let partials = serialize_partials(black_box(&inputs),
rows_per_partial);
+ let merged = partials
+ .iter()
+ .map(|bytes| TDigestMut::deserialize(bytes))
+ .collect::<Result<TDigestMut, _>>()
+ .unwrap();
+ final_quantile(merged)
+ });
+}
+
+fn serialize_partials(inputs: &[f64], rows_per_partial: usize) -> Vec<Vec<u8>>
{
+ inputs
+ .chunks(rows_per_partial)
+ .map(|values| build_mut_digest(values).serialize())
+ .collect()
+}
+
+fn final_quantile(mut merged: TDigestMut) -> Option<f64> {
+ // A later aggregation stage decodes the serialized merge result.
+ let bytes = merged.serialize();
+ TDigestMut::deserialize(black_box(&bytes))
+ .unwrap()
+ .quantile(black_box(0.9))
+}
diff --git a/benchmarks/tdigest/merge.rs b/benchmarks/tdigest/merge.rs
index 13559c6..12f6f8b 100644
--- a/benchmarks/tdigest/merge.rs
+++ b/benchmarks/tdigest/merge.rs
@@ -65,7 +65,7 @@ fn small_partials(bencher: Bencher) {
let partials = partial_digests(64, SMALL_ROWS_PER_PARTIAL)
.into_iter()
.map(|mut digest| {
- black_box(digest.rank(0.0));
+ black_box(digest.quantile(0.5));
digest
})
.collect::<Vec<_>>();
@@ -86,7 +86,7 @@ fn partials(bencher: Bencher) {
let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL)
.into_iter()
.map(|mut digest| {
- black_box(digest.rank(0.0));
+ black_box(digest.quantile(0.5));
digest
})
.collect::<Vec<_>>();
@@ -102,6 +102,80 @@ fn partials(bencher: Bencher) {
});
}
+#[divan::bench]
+fn partials_from_iter(bencher: Bencher) {
+ let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL)
+ .into_iter()
+ .map(|mut digest| {
+ black_box(digest.quantile(0.5));
+ digest
+ })
+ .collect::<Vec<_>>();
+
+ bencher
+ .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL))
+ // Construct fresh owned inputs outside the measurement, as a real
caller transfers
+ // ownership rather than cloning solely for `FromIterator`.
+ .with_inputs(|| partials.clone())
+ .bench_local_values(|partials|
black_box(partials).into_iter().collect::<TDigestMut>());
+}
+
+#[divan::bench]
+fn partials_from_iter_cloned(bencher: Bencher) {
+ let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL)
+ .into_iter()
+ .map(|mut digest| {
+ black_box(digest.quantile(0.5));
+ digest
+ })
+ .collect::<Vec<_>>();
+
+ bencher
+ .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL))
+ // Include cloning when the caller needs to retain the original inputs.
+ .bench_local(||
black_box(&partials).iter().cloned().collect::<TDigestMut>());
+}
+
+#[divan::bench(args = [0, 64, 4_096])]
+fn partials_from_iter_finalize(bencher: Bencher, additional_rows: usize) {
+ let partials = partial_digests_with(DEFAULT_DIGEST_K, 64, ROWS_PER_PARTIAL)
+ .into_iter()
+ .map(|mut digest| {
+ black_box(digest.quantile(0.5));
+ digest
+ })
+ .collect::<Vec<_>>();
+ let updates = values(additional_rows)
+ .into_iter()
+ .map(|value| value * (64 * ROWS_PER_PARTIAL) as f64)
+ .collect::<Vec<_>>();
+
+ bencher
+ .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL + additional_rows))
+ .with_inputs(|| partials.clone())
+ .bench_local_values(|partials| {
+ let mut merged =
black_box(partials).into_iter().collect::<TDigestMut>();
+ for &value in black_box(&updates) {
+ merged.update(value);
+ }
+ black_box(merged.freeze())
+ });
+}
+
+#[divan::bench]
+fn uncompressed_partials_from_iter(bencher: Bencher) {
+ let values = values(64 * ROWS_PER_PARTIAL);
+ let partials = values
+ .chunks_exact(ROWS_PER_PARTIAL)
+ .map(build_mut_digest)
+ .collect::<Vec<_>>();
+
+ bencher
+ .counter(ItemsCount::new(values.len()))
+ .with_inputs(|| partials.clone())
+ .bench_local_values(|partials|
black_box(partials).into_iter().collect::<TDigestMut>());
+}
+
#[divan::bench(args = [SMALL_ROWS_PER_PARTIAL, ROWS_PER_PARTIAL])]
fn serialized_partials(bencher: Bencher, rows_per_partial: usize) {
let partials = serialized_partial_digests(64, rows_per_partial);
@@ -137,3 +211,50 @@ fn serialized_overlapping_partials(bencher: Bencher) {
black_box(merged.quantile(0.5))
});
}
+
+#[divan::bench(args = [false, true])]
+fn serialized_partials_from_iter(bencher: Bencher, overlapping: bool) {
+ let partials = serialized_merge_inputs(overlapping);
+
+ bencher
+ .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL))
+ .bench_local(|| {
+ let mut merged = black_box(&partials)
+ .iter()
+ .map(|partial| TDigestMut::deserialize(partial))
+ .collect::<Result<TDigestMut, _>>()
+ .unwrap();
+ black_box(merged.quantile(0.5))
+ });
+}
+
+#[divan::bench(args = [false, true])]
+fn serialized_partials_in_batches(bencher: Bencher, overlapping: bool) {
+ let partials = serialized_merge_inputs(overlapping);
+
+ bencher
+ .counter(ItemsCount::new(64 * ROWS_PER_PARTIAL))
+ .bench_local(|| {
+ let mut merged = TDigestMut::new(DEFAULT_DIGEST_K).unwrap();
+ for batch in black_box(&partials).chunks(16) {
+ let batch = batch
+ .iter()
+ .map(|partial| TDigestMut::deserialize(partial))
+ .collect::<Result<TDigestMut, _>>()
+ .unwrap();
+ merged.merge(&batch);
+ }
+ black_box(merged.quantile(0.5))
+ });
+}
+
+fn serialized_merge_inputs(overlapping: bool) -> Vec<Vec<u8>> {
+ if overlapping {
+ values(64 * ROWS_PER_PARTIAL)
+ .chunks_exact(ROWS_PER_PARTIAL)
+ .map(|values| build_mut_digest(values).serialize())
+ .collect()
+ } else {
+ serialized_partial_digests(64, ROWS_PER_PARTIAL)
+ }
+}
diff --git a/benchmarks/tdigest/mod.rs b/benchmarks/tdigest/mod.rs
index adf1883..ce18622 100644
--- a/benchmarks/tdigest/mod.rs
+++ b/benchmarks/tdigest/mod.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+mod aggregation;
mod compress;
mod merge;
mod query;
diff --git a/datasketches/src/tdigest/mod.rs b/datasketches/src/tdigest/mod.rs
index 7d07de8..0d46618 100644
--- a/datasketches/src/tdigest/mod.rs
+++ b/datasketches/src/tdigest/mod.rs
@@ -23,7 +23,7 @@
//! The implementation in this library has a few differences from the
reference implementation
//! associated with that paper:
//!
-//! * Merging does not modify the input.
+//! * [`TDigestMut::merge`] borrows its input without modifying it.
//! * Deserialization is similar to other sketches in this library, although
reading the reference
//! implementation format is supported.
//!
diff --git a/datasketches/src/tdigest/sketch.rs
b/datasketches/src/tdigest/sketch.rs
index 317e61a..eb31e87 100644
--- a/datasketches/src/tdigest/sketch.rs
+++ b/datasketches/src/tdigest/sketch.rs
@@ -119,11 +119,10 @@ impl TDigestBuffer {
/// Combines this buffer with a non-empty borrowed buffer in stable mean
order.
fn into_merged_centroids(mut self, other: &TDigestBuffer) -> Vec<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
- && centroids_are_sorted(&self.centroids)
- && centroids_are_sorted(&other.centroids)
- {
+ if self.unmerged_tail_len == 0 && other.unmerged_tail_len == 0 {
+ // Compression and deserialization both establish this invariant.
+ debug_assert!(centroids_are_sorted(&self.centroids));
+ debug_assert!(centroids_are_sorted(&other.centroids));
merge_sorted_centroids(&mut self.centroids, &other.centroids);
return self.centroids;
}
@@ -327,6 +326,10 @@ impl TDigestMut {
/// Merges the given t-digest into this one.
///
+ /// Retains this digest's `k`, even if the other digest uses a different
value.
+ /// The borrowed input remains available for reuse. Collecting owned
inputs into a
+ /// [`TDigestMut`] with [`Iterator::collect`] combines a batch with one
compression pass.
+ ///
/// # Panics
///
/// Panics without modifying the digest if the combined total weight would
exceed `u64::MAX`.
@@ -943,16 +946,14 @@ impl TDigestMut {
while current < len {
let c = centroids[current];
let proposed_weight = centroids[num_centroids - 1].weight() +
c.weight();
- let mut add_this = false;
- if (current != 1) && (current != (len - 1)) {
- let q0 = weight_so_far / compressed_weight;
- let q2 = (weight_so_far + proposed_weight) / compressed_weight;
- add_this = proposed_weight
- <= (compressed_weight
- * scale_function::max(q0, normalizer)
- .min(scale_function::max(q2, normalizer)));
- }
- if add_this {
+ if should_merge_centroid(
+ current,
+ len,
+ weight_so_far,
+ proposed_weight,
+ compressed_weight,
+ normalizer,
+ ) {
// merge into existing centroid
centroids[num_centroids - 1].add(c);
} else {
@@ -992,6 +993,242 @@ impl TDigestMut {
}
}
+/// 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.
+///
+/// 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;
+ };
+
+ // 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,
+ }
+
+ 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;
+ }
+ }
+
+ 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,
+ });
+ }
+
+ 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 {
+ 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);
+ }
+ }
+
+ 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,
@@ -1667,6 +1904,42 @@ fn centroids_are_sorted(centroids: &[Centroid]) -> bool {
.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));
diff --git a/tests-integration/tests/tdigest_test/sketch.rs
b/tests-integration/tests/tdigest_test/sketch.rs
index e35d2d6..cae816b 100644
--- a/tests-integration/tests/tdigest_test/sketch.rs
+++ b/tests-integration/tests/tdigest_test/sketch.rs
@@ -287,6 +287,246 @@ fn test_merge_large() {
assert_that!(td1.rank(n as f64).unwrap(), eq(1.0));
}
+#[test]
+fn test_mixed_k_merge_retains_receiver_k() {
+ for (left_k, right_k, right_count) in [(200, 50, 1_000), (50, 200, 1_000),
(200, 10, 1)] {
+ let mut left = TDigestMut::new(left_k).unwrap();
+ let mut right = TDigestMut::new(right_k).unwrap();
+ for value in 0..10_000 {
+ left.update(value as f64);
+ }
+ for value in 0..right_count {
+ right.update((value + 10_000) as f64);
+ }
+
+ left.merge(&right);
+
+ assert_eq!(left.k(), left_k);
+ assert_eq!(left.total_weight(), 10_000 + right_count);
+ assert_eq!(left.quantile(0.0), Some(0.0));
+ assert_eq!(left.quantile(1.0), Some((9_999 + right_count) as f64));
+ }
+}
+
+#[test]
+fn test_from_iter_uses_one_result_with_the_smallest_nonempty_k() {
+ let mut first = TDigestMut::new(100).unwrap();
+ let mut second = TDigestMut::new(50).unwrap();
+ let empty = TDigestMut::new(10).unwrap();
+ for value in 0..1_000 {
+ first.update(value as f64);
+ second.update((value + 1_000) as f64);
+ }
+ let _ = first.quantile(0.5);
+ let _ = second.quantile(0.5);
+
+ let mut merged = [first, empty,
second].into_iter().collect::<TDigestMut>();
+
+ assert_eq!(merged.k(), 50);
+ assert_eq!(merged.total_weight(), 2_000);
+ assert_eq!(merged.min_value(), Some(0.0));
+ assert_eq!(merged.max_value(), Some(1_999.0));
+ let quantiles = (0..=100)
+ .map(|rank| merged.quantile(rank as f64 / 100.).unwrap())
+ .collect::<Vec<_>>();
+ assert!(quantiles.windows(2).all(|pair| pair[0] <= pair[1]));
+}
+
+#[test]
+fn test_from_iter_matches_single_digest_for_uncompressed_inputs() {
+ let values = [3.0, 1.0, 2.0, 2.0, 5.0, 4.0];
+ let partials = values
+ .chunks(3)
+ .map(|values| {
+ let mut digest = TDigestMut::new(100).unwrap();
+ for &value in values {
+ digest.update(value);
+ }
+ digest
+ })
+ .collect::<Vec<_>>();
+ let mut merged = partials.into_iter().collect::<TDigestMut>();
+
+ let mut expected = TDigestMut::new(100).unwrap();
+ for value in values {
+ expected.update(value);
+ }
+
+ assert_eq!(merged.serialize(), expected.serialize());
+}
+
+#[test]
+fn test_from_iter_matches_single_compression_for_interleaved_runs() {
+ let runs: [&[(f64, u64)]; 4] = [
+ &[(0.0, 1), (3.0, 2), (3.0, 5), (6.0, 100), (9.0, 1)],
+ &[(0.0, 7), (2.0, 3), (5.0, 4), (8.0, 2), (9.0, 10)],
+ &[(1.0, 2), (3.0, 6), (4.0, 1), (7.0, 8), (9.0, 20)],
+ &[(0.0, 2), (3.0, 9), (6.0, 7), (9.0, 30)],
+ ];
+ for k in [10, 100] {
+ for reverse in [false, true] {
+ let make_digest = |centroids: &[(f64, u64)]| {
+ let mut digest = deserialize_with_centroids(k, -10.0, 20.0,
centroids);
+ if reverse {
+ let mut bytes = digest.serialize();
+ bytes[5] |= 1 << 2; // reverse-merge flag
+ digest = TDigestMut::deserialize(&bytes).unwrap();
+ }
+ digest
+ };
+ let mut ordered = runs
+ .iter()
+ .flat_map(|run| run.iter().copied())
+ .collect::<Vec<_>>();
+ ordered.sort_by(|left, right| left.0.total_cmp(&right.0));
+
+ // Borrowed merge puts the right input first on ties. Keeping the
last centroid on
+ // the left reproduces the concatenated runs' stable order in one
compression pass.
+ let last = ordered.pop().unwrap();
+ let mut expected = make_digest(&[last]);
+ expected.merge(&make_digest(&ordered));
+ let mut merged =
runs.into_iter().map(make_digest).collect::<TDigestMut>();
+
+ assert_eq!(merged.serialize(), expected.serialize());
+ assert_eq!(merged.quantile(0.0), Some(-10.0));
+ assert_eq!(merged.quantile(1.0), Some(20.0));
+
+ // A collected digest must still support enough updates to grow
and compress again.
+ for value in 0..1_000 {
+ merged.update(value as f64);
+ expected.update(value as f64);
+ }
+ assert_eq!(merged.serialize(), expected.serialize());
+ }
+ }
+}
+
+#[test]
+fn test_from_iter_handles_mixed_compressed_and_buffered_inputs() {
+ let mut expected = TDigestMut::new(100).unwrap();
+ let partials = (0..4)
+ .map(|run| {
+ let mut digest = TDigestMut::new(100).unwrap();
+ for row in 0..60 {
+ let value = ((row * 13 + run * 7) % 37) as f64;
+ digest.update(value);
+ expected.update(value);
+ if run == 2 && row == 29 {
+ let _ = digest.quantile(0.5);
+ }
+ }
+ if run == 1 {
+ let _ = digest.quantile(0.5);
+ }
+ digest
+ })
+ .collect::<Vec<_>>();
+ let mut merged = partials.into_iter().collect::<TDigestMut>();
+
+ assert_eq!(merged.total_weight(), 240);
+ assert_eq!(merged.serialize(), expected.serialize());
+}
+
+#[test]
+fn test_serialized_batch_merge_tree_preserves_weight_and_quantiles() {
+ let mut values = Vec::new();
+ let partials = (0..64)
+ .map(|input| {
+ let mut digest = TDigestMut::default();
+ // Unequal weights, empty states, and repeated values exercise
intermediate merges.
+ let rows = [0, 1, 8, 64, 1_024][input % 5];
+ for row in 0..rows {
+ let value = ((row * 37 + input * 113) % 1_024) as f64 - 512.0;
+ values.push(value);
+ digest.update(value);
+ }
+ digest.serialize()
+ })
+ .collect::<Vec<_>>();
+ values.sort_by(f64::total_cmp);
+
+ for fan_in in [2, 7, 64] {
+ let mut states = partials.clone();
+ while states.len() > 1 {
+ states = states
+ .chunks(fan_in)
+ .map(|batch| {
+ let mut merged = batch
+ .iter()
+ .map(|bytes| TDigestMut::deserialize(bytes))
+ .collect::<Result<TDigestMut, _>>()
+ .unwrap();
+ merged.serialize()
+ })
+ .collect();
+ }
+
+ let mut merged = TDigestMut::deserialize(&states[0]).unwrap();
+ assert_eq!(merged.total_weight(), values.len() as u64);
+ assert_eq!(merged.quantile(0.0), values.first().copied());
+ assert_eq!(merged.quantile(1.0), values.last().copied());
+
+ let mut previous = values[0];
+ for rank in [0.01, 0.5, 0.9, 0.99] {
+ let estimate = merged.quantile(rank).unwrap();
+ assert!((previous..=values[values.len() - 1]).contains(&estimate));
+ // Check this fixture's empirical rank, allowing equal values to
span a rank interval.
+ let lower =
+ values.partition_point(|value| *value < estimate) as f64 /
values.len() as f64;
+ let upper =
+ values.partition_point(|value| *value <= estimate) as f64 /
values.len() as f64;
+ assert!(
+ (lower - 0.01..=upper + 0.01).contains(&rank),
+ "fan_in={fan_in}, rank={rank}, estimate={estimate},
observed=[{lower}, {upper}]"
+ );
+ previous = estimate;
+ }
+ }
+}
+
+#[test]
+fn test_from_iter_checks_total_weight_before_compression() {
+ for compressed in [false, true] {
+ let heavy = deserialize_with_centroids(100, 0.0, 0.0, &[(0.0, u64::MAX
- 2)]);
+ let [one, two, extra] = [1.0, 2.0, 3.0].map(|value| {
+ let mut digest = TDigestMut::new(100).unwrap();
+ digest.update(value);
+ if compressed {
+ let _ = digest.quantile(0.5);
+ }
+ digest
+ });
+ let mut at_limit = [heavy.clone(), one.clone(), two.clone()]
+ .into_iter()
+ .collect::<TDigestMut>();
+ assert_eq!(at_limit.total_weight(), u64::MAX);
+ assert_eq!(at_limit.quantile(0.0), Some(0.0));
+ assert_eq!(at_limit.quantile(1.0), Some(2.0));
+
+ assert!(
+ catch_unwind(|| [heavy, one, two,
extra].into_iter().collect::<TDigestMut>()).is_err()
+ );
+ }
+}
+
+#[test]
+fn test_from_iter_handles_empty_and_single_input_without_recompression() {
+ let empty = std::iter::empty::<TDigestMut>().collect::<TDigestMut>();
+ assert!(empty.is_empty());
+
+ let mut input = TDigestMut::new(50).unwrap();
+ for value in 0..1_000 {
+ input.update(value as f64);
+ }
+ let serialized = input.serialize();
+ let mut collected = [TDigestMut::new(10).unwrap(), input]
+ .into_iter()
+ .collect::<TDigestMut>();
+
+ assert_eq!(collected.k(), 50);
+ assert_eq!(collected.serialize(), serialized);
+}
+
#[test]
fn test_invalid_inputs() {
let n = 100;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
