diff --git a/CHANGELOG.md b/CHANGELOG.md index f65b267..645b6ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,14 +11,17 @@ All significant changes to this project will be documented in this file. ### New features * Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, totally ordered custom item types, a `KllFloat` adapter for non-NaN floating-point values, and serialization. +* Implement `FromIterator` for batch construction, and add `TDigestMut::quantiles` and `TDigest::quantiles` for querying several ranks in one centroid scan. ### Improvements * Improve truncated-input diagnostics across sketch deserializers. +* T-Digest batch construction from owned partial sketches avoids recompressing intermediate results, and batch quantile queries reuse one traversal for ranks supplied in nondecreasing order. ### Bug fixes * T-Digest deserialization now rejects unknown or conflicting flags, reversed extrema, out-of-range values, unsorted centroids, and non-empty images without stored values. +* T-Digest merging now uses the smaller `k` when sketches have different compression parameters, preserving the size bound of the coarser input. ## v0.5.0 diff --git a/benchmarks/tdigest/merge.rs b/benchmarks/tdigest/merge.rs index 13559c6..6c92ed8 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::>(); @@ -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::>(); @@ -102,6 +102,38 @@ 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::>(); + + 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::()); +} + +#[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::>(); + + bencher + .counter(ItemsCount::new(values.len())) + .with_inputs(|| partials.clone()) + .bench_local_values(|partials| black_box(partials).into_iter().collect::()); +} + #[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); diff --git a/benchmarks/tdigest/query.rs b/benchmarks/tdigest/query.rs index 0981657..62422d3 100644 --- a/benchmarks/tdigest/query.rs +++ b/benchmarks/tdigest/query.rs @@ -56,3 +56,24 @@ fn quantiles_2_sequential(bencher: Bencher) { ] }); } + +#[divan::bench] +fn quantiles_6_sequential(bencher: Bencher) { + let digest = prepared_digest(); + let ranks = [0.25, 0.5, 0.75, 0.9, 0.95, 0.99]; + + bencher + .bench_local(|| black_box(ranks).map(|rank| black_box(&digest).quantile(black_box(rank)))); +} + +#[divan::bench(args = [2, 6, 100])] +fn quantiles_batch(bencher: Bencher, num_ranks: usize) { + let digest = prepared_digest(); + let ranks = (1..=num_ranks) + .map(|rank| rank as f64 / (num_ranks + 1) as f64) + .collect::>(); + + bencher + .counter(ItemsCount::new(num_ranks)) + .bench_local(|| black_box(&digest).quantiles(black_box(&ranks))); +} diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index d496a9c..c96649f 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 { 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; } @@ -322,6 +321,13 @@ impl TDigestMut { /// Merges the given t-digest into this one. /// + /// If the sketches have different `k` values, the merged sketch uses the smaller value because + /// merging cannot recover detail already discarded by the lower-`k` sketch. + /// + /// # Panics + /// + /// Panics if the combined total weight exceeds `u64::MAX`. + /// /// # Examples /// /// ``` @@ -340,8 +346,64 @@ impl TDigestMut { } let self_unmerged_weight = self.buffer.unmerged_len() as u64; + let additional_weight = self_unmerged_weight + .checked_add(other.total_weight()) + .expect("combined t-digest weight exceeds u64::MAX"); + self.compressed_weight + .checked_add(additional_weight) + .expect("combined t-digest weight exceeds u64::MAX"); let centroids = std::mem::take(&mut self.buffer).into_merged_centroids(&other.buffer); - self.compress_sorted_centroids(centroids, self_unmerged_weight + other.total_weight()) + self.k = self.k.min(other.k); + self.min = self.min.min(other.min); + self.max = self.max.max(other.max); + self.compress_sorted_centroids(centroids, additional_weight); + } + + fn from_owned_digests(mut digests: Vec) -> Self { + debug_assert!(digests.len() >= 2); + debug_assert!(digests.iter().all(|digest| !digest.is_empty())); + + let mut total_weight = 0u64; + let mut num_centroids = 0usize; + let mut k = digests[0].k; + let reverse_merge = digests[0].reverse_merge; + let mut min = digests[0].min; + let mut max = digests[0].max; + for digest in &digests { + 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"); + k = k.min(digest.k); + min = min.min(digest.min); + max = max.max(digest.max); + } + + let mut merged = TDigestMut::make(k, reverse_merge, min, max, TDigestBuffer::default(), 0); + let all_compressed = digests + .iter() + .all(|digest| digest.buffer.unmerged_len() == 0); + if all_compressed { + // Compressed buffers are sorted runs, so feed their k-way merge directly into + // compression instead of materializing every input centroid in another vector. + let centroids = KWayMerge::new(&digests, reverse_merge); + merged.compress_merged_centroids(centroids, num_centroids, total_weight); + } else { + // Raw tails are unsorted. Put them before the compressed prefixes so the stable sort + // preserves the same equal-mean tie order as regular updates followed by summaries. + let mut centroids = Vec::with_capacity(num_centroids); + for digest in &mut digests { + let tail_start = digest.buffer.compressed_prefix_len(); + centroids.extend(digest.buffer.centroids.drain(tail_start..)); + } + for digest in digests { + centroids.extend(digest.buffer.centroids); + } + merged.compress_centroids(centroids, total_weight); + } + merged } /// Converts this mutable t-digest into an immutable one. @@ -506,6 +568,21 @@ impl TDigestMut { self.view().quantile(rank) } + /// Returns the quantiles described by [`TDigest::quantiles`]. + /// + /// # Panics + /// + /// Panics if any rank is outside `[0.0, 1.0]`. + pub fn quantiles(&mut self, ranks: &[f64]) -> Option> { + check_ranks(ranks); + + if self.is_empty() { + return None; + } + + self.view().quantiles(ranks) + } + /// Serializes this mutable t-digest to bytes. /// /// # Examples @@ -919,15 +996,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))); - } + let add_this = should_merge_centroid( + current, + len, + weight_so_far, + proposed_weight, + compressed_weight, + normalizer, + ); if add_this { // merge into existing centroid centroids[num_centroids - 1].add(c); @@ -951,6 +1027,57 @@ impl TDigestMut { self.buffer = TDigestBuffer::new(centroids, 0); } + fn compress_merged_centroids( + &mut self, + centroids: impl Iterator, + num_input_centroids: usize, + total_weight: u64, + ) { + debug_assert_ne!(num_input_centroids, 0); + let compressed_weight = total_weight as f64; + let normalizer = scale_function::normalizer(2.0 * f64::from(self.k), compressed_weight); + let mut centroids = centroids.enumerate(); + let (_, first) = centroids.next().expect("non-empty centroid stream"); + let capacity = self.target_retained_capacity().min(num_input_centroids); + let mut retained = Vec::with_capacity(capacity); + retained.push(first); + let mut weight_so_far = 0.; + + for (current, centroid) in centroids { + let proposed_weight = retained.last().unwrap().weight() + centroid.weight(); + let add_this = should_merge_centroid( + current, + num_input_centroids, + weight_so_far, + proposed_weight, + compressed_weight, + normalizer, + ); + if add_this { + retained.last_mut().unwrap().add(centroid); + } else { + weight_so_far += retained.last().unwrap().weight(); + retained.push(centroid); + } + } + + debug_assert!(retained.len() <= self.target_centroids()); + if self.reverse_merge { + retained.reverse(); + } + debug_assert_eq!( + retained + .iter() + .map(|centroid| centroid.weight.get()) + .sum::(), + total_weight, + "compressed centroids must preserve total weight" + ); + self.compressed_weight = total_weight; + self.reverse_merge = !self.reverse_merge; + self.buffer = TDigestBuffer::new(retained, 0); + } + fn reduce_retained_capacity(&self, centroids: &mut Vec) { let target_capacity = self.target_retained_capacity().max(centroids.len()); if centroids.capacity() <= target_capacity { @@ -968,6 +1095,47 @@ impl TDigestMut { } } +/// Collects owned t-digests into one result with a single compression pass. +/// +/// Empty inputs are ignored. The result uses the smallest `k` among the non-empty inputs. A single +/// non-empty input is returned unchanged. Collecting consumes each digest, so callers do not need +/// to clone inputs. Unlike repeated [`TDigestMut::merge`] calls, it temporarily retains all input +/// centroids so it can avoid recompressing intermediate results. Use repeated `merge` calls when +/// inputs must be processed with bounded additional memory. +/// +/// # Panics +/// +/// Panics if the combined total weight exceeds `u64::MAX` or the combined centroid count exceeds +/// `usize::MAX`. +/// +/// # Examples +/// +/// ``` +/// use datasketches::tdigest::TDigestMut; +/// +/// let partials = [1.0, 2.0].map(|value| { +/// let mut digest = TDigestMut::new(100).unwrap(); +/// digest.update(value); +/// digest +/// }); +/// let merged = partials.into_iter().collect::(); +/// assert_eq!(merged.total_weight(), 2); +/// ``` +impl FromIterator for TDigestMut { + fn from_iter>(iter: T) -> Self { + let mut digests = iter.into_iter().filter(|digest| !digest.is_empty()); + let Some(first) = digests.next() else { + return TDigestMut::default(); + }; + let mut owned = digests.collect::>(); + if owned.is_empty() { + return first; + } + owned.insert(0, first); + TDigestMut::from_owned_digests(owned) + } +} + fn serialize_compressed( k: u16, reverse_merge: bool, @@ -1268,6 +1436,35 @@ impl TDigest { self.view().quantile(rank) } + /// Computes approximate quantiles for the given normalized ranks. + /// + /// Ranks in nondecreasing order are answered with one centroid scan. Ranks in any other order + /// are accepted and results are returned in the same order as the input. + /// + /// Returns `None` if this t-digest is empty. + /// + /// # Panics + /// + /// Panics if any rank is outside `[0.0, 1.0]`. + /// + /// # Examples + /// + /// ``` + /// use datasketches::tdigest::TDigestMut; + /// + /// let mut sketch = TDigestMut::new(100).unwrap(); + /// for value in [1.0, 2.0, 3.0] { + /// sketch.update(value); + /// } + /// let digest = sketch.freeze(); + /// let quantiles = digest.quantiles(&[0.25, 0.5, 0.75]).unwrap(); + /// assert_eq!(quantiles.len(), 3); + /// ``` + pub fn quantiles(&self, ranks: &[f64]) -> Option> { + check_ranks(ranks); + self.view().quantiles(ranks) + } + /// Converts this immutable t-digest into a mutable one. /// /// # Examples @@ -1435,82 +1632,135 @@ impl TDigestView<'_> { return None; } + Some(QuantileCursor::new(self).quantile(rank)) + } + + fn quantiles(&self, ranks: &[f64]) -> Option> { + debug_assert!( + ranks.iter().all(|rank| (0.0..=1.0).contains(rank)), + "ranks must be in [0.0, 1.0]" + ); + + if self.centroids.is_empty() { + return None; + } + + let mut quantiles = vec![0.; ranks.len()]; + let mut cursor = QuantileCursor::new(self); + if ranks.windows(2).all(|pair| pair[0] <= pair[1]) { + for (index, &rank) in ranks.iter().enumerate() { + quantiles[index] = cursor.quantile(rank); + } + return Some(quantiles); + } + + // The cursor only moves forward. Sort indices so queries become monotonic without changing + // the caller's output order. + let mut rank_order = (0..ranks.len()).collect::>(); + rank_order.sort_by(|&left, &right| ranks[left].total_cmp(&ranks[right])); + for index in rank_order { + quantiles[index] = cursor.quantile(ranks[index]); + } + Some(quantiles) + } +} + +/// Incrementally answers quantile queries supplied in nondecreasing rank order. +struct QuantileCursor<'a> { + min: f64, + max: f64, + centroids: &'a [Centroid], + centroids_weight: f64, + centroid_index: usize, + weight_so_far: f64, +} + +impl<'a> QuantileCursor<'a> { + fn new(view: &TDigestView<'a>) -> Self { + debug_assert!(!view.centroids.is_empty()); + + QuantileCursor { + min: view.min, + max: view.max, + centroids: view.centroids, + centroids_weight: view.centroids_weight as f64, + centroid_index: 0, + weight_so_far: view.centroids[0].weight() / 2., + } + } + + fn quantile(&mut self, rank: f64) -> f64 { + debug_assert!(!self.centroids.is_empty()); + if self.centroids.len() == 1 { - return Some(self.centroids[0].mean); + return self.centroids[0].mean; } // at least 2 centroids - let centroids_weight = self.centroids_weight as f64; + let centroids_weight = self.centroids_weight; let num_centroids = self.centroids.len(); let weight = rank * centroids_weight; if weight < 1. { - return Some(self.min); + return self.min; } if weight > centroids_weight - 1. { - return Some(self.max); + return self.max; } let first_weight = self.centroids[0].weight(); if first_weight > 1. && weight < first_weight / 2. { - return Some( - self.min - + (((weight - 1.) / ((first_weight / 2.) - 1.)) - * (self.centroids[0].mean - self.min)), - ); + return self.min + + (((weight - 1.) / ((first_weight / 2.) - 1.)) + * (self.centroids[0].mean - self.min)); } let last_weight = self.centroids[num_centroids - 1].weight(); if last_weight > 1. && (centroids_weight - weight <= last_weight / 2.) { if last_weight == 2. { - return Some(self.max); + return self.max; } - return Some( - self.max - - (((centroids_weight - weight - 1.) / ((last_weight / 2.) - 1.)) - * (self.max - self.centroids[num_centroids - 1].mean)), - ); + return self.max + - (((centroids_weight - weight - 1.) / ((last_weight / 2.) - 1.)) + * (self.max - self.centroids[num_centroids - 1].mean)); } // interpolate between extremes - let mut weight_so_far = first_weight / 2.; - for i in 0..(num_centroids - 1) { - let dw = (self.centroids[i].weight() + self.centroids[i + 1].weight()) / 2.; - if weight_so_far + dw > weight { + while self.centroid_index < num_centroids - 1 { + let dw = (self.centroids[self.centroid_index].weight() + + self.centroids[self.centroid_index + 1].weight()) + / 2.; + if self.weight_so_far + dw > weight { // the target weight is between centroids i and i+1 let mut left_weight = 0.; - if self.centroids[i].weight.get() == 1 { - if weight - weight_so_far < 0.5 { - return Some(self.centroids[i].mean); + if self.centroids[self.centroid_index].weight.get() == 1 { + if weight - self.weight_so_far < 0.5 { + return self.centroids[self.centroid_index].mean; } left_weight = 0.5; } let mut right_weight = 0.; - if self.centroids[i + 1].weight.get() == 1 { - if weight_so_far + dw - weight <= 0.5 { - return Some(self.centroids[i + 1].mean); + if self.centroids[self.centroid_index + 1].weight.get() == 1 { + if self.weight_so_far + dw - weight <= 0.5 { + return self.centroids[self.centroid_index + 1].mean; } right_weight = 0.5; } // Each centroid is weighted by the distance from the target to the *other* // centroid, so the estimate approaches the nearer one. - let distance_from_left = weight - weight_so_far - left_weight; - let distance_to_right = weight_so_far + dw - weight - right_weight; - return Some(weighted_average( - self.centroids[i].mean, + let distance_from_left = weight - self.weight_so_far - left_weight; + let distance_to_right = self.weight_so_far + dw - weight - right_weight; + return weighted_average( + self.centroids[self.centroid_index].mean, distance_to_right, - self.centroids[i + 1].mean, + self.centroids[self.centroid_index + 1].mean, distance_from_left, - )); + ); } - weight_so_far += dw; + self.weight_so_far += dw; + self.centroid_index += 1; } let w1 = weight - (centroids_weight) - ((self.centroids[num_centroids - 1].weight()) / 2.); let w2 = (self.centroids[num_centroids - 1].weight() / 2.) - w1; - Some(weighted_average( - self.centroids[num_centroids - 1].mean, - w1, - self.max, - w2, - )) + weighted_average(self.centroids[num_centroids - 1].mean, w1, self.max, w2) } } @@ -1526,6 +1776,14 @@ fn check_split_points(split_points: &[f64]) { } } +#[track_caller] +fn check_ranks(ranks: &[f64]) { + assert!( + ranks.iter().all(|rank| (0.0..=1.0).contains(rank)), + "ranks must be in [0.0, 1.0]" + ); +} + fn centroid_cmp(a: &Centroid, b: &Centroid) -> Ordering { match a.mean.partial_cmp(&b.mean) { Some(order) => order, @@ -1539,6 +1797,137 @@ fn centroids_are_sorted(centroids: &[Centroid]) -> bool { .all(|pair| centroid_cmp(&pair[0], &pair[1]) != Ordering::Greater) } +fn should_merge_centroid( + current: usize, + len: usize, + weight_so_far: f64, + proposed_weight: f64, + compressed_weight: f64, + normalizer: f64, +) -> bool { + if current == 1 || current == len - 1 { + return false; + } + let q0 = weight_so_far / compressed_weight; + let q2 = (weight_so_far + proposed_weight) / compressed_weight; + proposed_weight + <= compressed_weight + * scale_function::max(q0, normalizer).min(scale_function::max(q2, normalizer)) +} + +/// A non-empty sorted run participating in a k-way merge. +struct CentroidRun<'a> { + head: Centroid, + tail: &'a [Centroid], + // Breaks equal-mean ties as if the input runs had been concatenated and stably sorted. + order: usize, +} + +impl<'a> CentroidRun<'a> { + fn new(centroids: &'a [Centroid], order: usize, reverse: bool) -> Option { + let (head, tail) = if reverse { + let (head, tail) = centroids.split_last()?; + (*head, tail) + } else { + let (head, tail) = centroids.split_first()?; + (*head, tail) + }; + Some(CentroidRun { head, tail, order }) + } + + fn head(&self) -> Centroid { + self.head + } + + fn advance(&mut self, reverse: bool) -> bool { + if self.tail.is_empty() { + return false; + } + if reverse { + let (head, tail) = self.tail.split_last().unwrap(); + self.head = *head; + self.tail = tail; + } else { + let (head, tail) = self.tail.split_first().unwrap(); + self.head = *head; + self.tail = tail; + } + true + } + + fn precedes(&self, other: &Self, reverse: bool) -> bool { + match centroid_cmp(&self.head(), &other.head()) { + Ordering::Less => !reverse, + Ordering::Greater => reverse, + Ordering::Equal if reverse => self.order > other.order, + Ordering::Equal => self.order < other.order, + } + } +} + +/// Lazily merges the sorted centroid buffers of fully compressed digests. +struct KWayMerge<'a> { + heap: Vec>, + reverse: bool, +} + +impl<'a> KWayMerge<'a> { + fn new(digests: &'a [TDigestMut], reverse: bool) -> Self { + debug_assert!( + digests + .iter() + .all(|digest| digest.buffer.unmerged_len() == 0) + ); + let mut runs = Vec::with_capacity(digests.len()); + runs.extend(digests.iter().enumerate().filter_map(|(order, digest)| { + CentroidRun::new(&digest.buffer.centroids, order, reverse) + })); + for index in (0..runs.len() / 2).rev() { + sift_down_centroid_runs(&mut runs, index, reverse); + } + KWayMerge { + heap: runs, + reverse, + } + } +} + +impl Iterator for KWayMerge<'_> { + type Item = Centroid; + + fn next(&mut self) -> Option { + if self.heap.is_empty() { + return None; + } + let centroid = self.heap[0].head(); + if !self.heap[0].advance(self.reverse) { + self.heap.swap_remove(0); + } + sift_down_centroid_runs(&mut self.heap, 0, self.reverse); + Some(centroid) + } +} + +fn sift_down_centroid_runs(heap: &mut [CentroidRun<'_>], mut position: usize, reverse: bool) { + loop { + let left = (position * 2) + 1; + if left >= heap.len() { + return; + } + let right = left + 1; + let next = if right < heap.len() && heap[right].precedes(&heap[left], reverse) { + right + } else { + left + }; + if !heap[next].precedes(&heap[position], reverse) { + return; + } + heap.swap(position, next); + position = next; + } +} + fn merge_sorted_centroids(left: &mut Vec, 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 302cc4c..997a2fb 100644 --- a/tests-integration/tests/tdigest_test/sketch.rs +++ b/tests-integration/tests/tdigest_test/sketch.rs @@ -275,6 +275,85 @@ fn test_merge_large() { assert_that!(td1.rank(n as f64).unwrap(), eq(1.0)); } +#[test] +fn test_mixed_k_merge_uses_smaller_k() { + let mut left = TDigestMut::new(200).unwrap(); + let mut right = TDigestMut::new(50).unwrap(); + for value in 0..1_000 { + left.update(value as f64); + right.update((value + 1_000) as f64); + } + + left.merge(&right); + + assert_eq!(left.k(), 50); + assert_eq!(left.total_weight(), 2_000); + assert_eq!(left.min_value(), Some(0.0)); + assert_eq!(left.max_value(), Some(1_999.0)); +} + +#[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::(); + + 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)); + assert_quantiles_are_nondecreasing(&mut merged); +} + +#[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::>(); + let mut merged = partials.into_iter().collect::(); + + 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_handles_empty_and_single_input_without_recompression() { + let empty = std::iter::empty::().collect::(); + 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::(); + + assert_eq!(collected.k(), 50); + assert_eq!(collected.serialize(), serialized); +} + #[test] fn test_invalid_inputs() { let n = 100; @@ -324,6 +403,33 @@ fn test_extreme_values_produce_finite_quantiles() { } } +#[test] +fn test_batch_quantiles_match_scalar_queries_in_input_order() { + let mut tdigest = TDigestMut::new(100).unwrap(); + for value in 0..10_000 { + tdigest.update(((value * 37) % 1_003) as f64); + } + let tdigest = tdigest.freeze(); + + for ranks in [ + vec![0.0, 0.001, 0.25, 0.5, 0.5, 0.99, 1.0], + vec![0.99, 0.0, 0.5, 1.0, 0.001, 0.5, 0.25], + vec![], + ] { + let expected = ranks + .iter() + .map(|&rank| tdigest.quantile(rank).unwrap()) + .collect::>(); + assert_eq!(tdigest.quantiles(&ranks), Some(expected)); + } +} + +fn assert_quantiles_are_nondecreasing(tdigest: &mut TDigestMut) { + let ranks = (0..=100).map(|rank| rank as f64 / 100.).collect::>(); + let quantiles = tdigest.quantiles(&ranks).unwrap(); + assert!(quantiles.windows(2).all(|pair| pair[0] <= pair[1])); +} + #[test] fn test_estimate_repeat_values() { let mut tdigest = TDigestMut::default(); @@ -393,6 +499,15 @@ fn test_quantile_handles_two_sample_last_centroid() { assert_eq!(tdigest.quantile(0.75), Some(100.0)); } +#[test] +#[should_panic(expected = "combined t-digest weight exceeds u64::MAX")] +fn test_merge_panics_on_total_weight_overflow() { + let mut left = deserialize_with_centroids(100, 0.0, 0.0, &[(0.0, u64::MAX)]); + let right = deserialize_with_centroids(50, 1.0, 1.0, &[(1.0, 1)]); + + left.merge(&right); +} + #[test] fn test_rank_left_tail_is_a_fraction_of_the_total_weight() { let mut tdigest =