diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f366e11..2a2e68a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,7 @@ All significant changes to this project will be documented in this file. ### New features -* `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. +* `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. `KllSketch::update_with_weight` adds an item repeated a given number of times, at a cost logarithmic in the weight. * Add `TDigestMut::quantiles` and `TDigest::quantiles` for querying several ranks in one centroid scan while preserving input order. * Add `FromIterator` as an alternative to incremental `TDigestMut::merge` calls for combining owned partial sketches in a single compression pass. diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 9fb3a62a..a2820b0f 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -143,6 +143,72 @@ impl KllSketch { self.internal_update(item); } + /// Updates the sketch with an item repeated `weight` times. + /// + /// The result is equivalent to calling [`update`](Self::update) `weight` times, at a cost + /// that grows with the logarithm of `weight` rather than with `weight` itself. A weight of + /// zero is a no-op. + /// + /// # Panics + /// + /// Panics without modifying the sketch if the stream weight would exceed [`u64::MAX`]. + /// + /// # Examples + /// + /// ``` + /// # use datasketches::common::SearchCriteria; + /// # use datasketches::kll::KllSketch; + /// let mut sketch = KllSketch::::new(200).unwrap(); + /// sketch.update_with_weight(1, 3); + /// sketch.update_with_weight(2, 1_000_000); + /// assert_eq!(sketch.n(), 1_000_003); + /// assert_eq!( + /// sketch.rank(&1, SearchCriteria::Inclusive).unwrap(), + /// 3.0 / 1_000_003.0 + /// ); + /// ``` + pub fn update_with_weight(&mut self, item: T, weight: u64) { + if weight == 0 { + return; + } + let final_n = self + .n + .checked_add(weight) + .expect("total stream weight overflow"); + if weight < (self.capacity - self.num_retained) as u64 { + self.update_min_max(&item); + for _ in 1..weight { + self.internal_update(item.clone()); + } + self.internal_update(item); + } else { + // An item at level `h` carries weight `2^h`, so one copy at each level whose bit is + // set in `weight` is an exact sketch of `weight` copies. Bits above the top level + // fold into `weight >> top` copies at the top level. + let num_levels = ((u64::BITS - weight.leading_zeros()) as usize).min(MAX_NUM_LEVELS); + let levels = (0..num_levels) + .map(|level| { + let count = if level + 1 < num_levels { + (weight >> level) & 1 + } else { + weight >> level + }; + vec![item.clone(); count as usize] + }) + .collect(); + let weighted = Self::make( + self.k, + self.k, + weight, + levels, + Some(item.clone()), + Some(item), + true, + ); + self.merge_unchecked(&weighted, final_n); + } + } + /// Resets this sketch to its empty state while retaining its configuration. pub fn reset(&mut self) { self.min_k = self.k; @@ -180,7 +246,11 @@ impl KllSketch { other.n )) })?; + self.merge_unchecked(other, final_n); + Ok(()) + } + fn merge_unchecked(&mut self, other: &KllSketch, final_n: u64) { self.update_min_max_from_other(other); for item in &other.levels[0] { @@ -197,7 +267,6 @@ impl KllSketch { } debug_assert_eq!(self.total_weight(), self.n, "total weight does not match n"); - Ok(()) } /// Returns the normalized rank of the given item. diff --git a/tests-integration/tests/kll_test/main.rs b/tests-integration/tests/kll_test/main.rs index 26441bc2..7754cc0c 100644 --- a/tests-integration/tests/kll_test/main.rs +++ b/tests-integration/tests/kll_test/main.rs @@ -19,3 +19,4 @@ mod core; mod generic; mod merge; mod query; +mod update; diff --git a/tests-integration/tests/kll_test/update.rs b/tests-integration/tests/kll_test/update.rs new file mode 100644 index 00000000..789deb4d --- /dev/null +++ b/tests-integration/tests/kll_test/update.rs @@ -0,0 +1,148 @@ +// 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 std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + +use datasketches::common::SearchCriteria; +use datasketches::kll::KllSketch; + +const DEFAULT_K: u16 = 200; +const RANK_EPS_FOR_K_200: f64 = 0.0133; + +#[test] +fn zero_weight_is_a_no_op() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + sketch.update_with_weight(1, 0); + assert!(sketch.is_empty()); + assert_eq!(sketch.min_item(), None); +} + +#[test] +fn weighted_update_matches_repeated_updates_in_exact_mode() { + let mut weighted = KllSketch::::new(DEFAULT_K).unwrap(); + let mut repeated = KllSketch::::new(DEFAULT_K).unwrap(); + for item in 1..=15 { + weighted.update_with_weight(item, item as u64); + for _ in 0..item { + repeated.update(item); + } + } + assert!(!weighted.is_estimation_mode()); + assert_eq!(weighted, repeated); +} + +#[test] +fn single_item_with_large_weight() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + let weight = (1u64 << 40) + 12_345; + sketch.update_with_weight(7, weight); + assert_eq!(sketch.n(), weight); + assert_eq!(sketch.num_retained(), weight.count_ones() as usize); + assert_eq!(sketch.min_item(), Some(&7)); + assert_eq!(sketch.max_item(), Some(&7)); + assert_eq!(sketch.rank(&7, SearchCriteria::Inclusive).unwrap(), 1.0); + assert_eq!(sketch.rank(&7, SearchCriteria::Exclusive).unwrap(), 0.0); + assert_eq!(sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), 7); +} + +#[test] +fn maximum_weight() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + sketch.update(1); + sketch.update_with_weight(2, u64::MAX - 1); + assert_eq!(sketch.n(), u64::MAX); + assert_eq!(sketch.min_item(), Some(&1)); + assert_eq!(sketch.max_item(), Some(&2)); + assert_eq!(sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), 2); + + let decoded = KllSketch::::deserialize(&sketch.serialize()).unwrap(); + assert_eq!(decoded, sketch); +} + +#[test] +fn weighted_updates_into_a_full_sketch() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + for item in 0..1_000 { + sketch.update(item); + } + sketch.update_with_weight(-1, 1_000); + sketch.update_with_weight(2_000, 3_000); + assert_eq!(sketch.n(), 5_000); + assert_eq!(sketch.min_item(), Some(&-1)); + assert_eq!(sketch.max_item(), Some(&2_000)); + let rank = sketch.rank(&-1, SearchCriteria::Inclusive).unwrap(); + assert!((rank - 0.2).abs() <= RANK_EPS_FOR_K_200, "rank {rank}"); + let rank = sketch.rank(&2_000, SearchCriteria::Exclusive).unwrap(); + assert!((rank - 0.4).abs() <= RANK_EPS_FOR_K_200, "rank {rank}"); +} + +#[test] +fn weighted_updates_in_estimation_mode() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + let weights: Vec = (0..10_000).map(|item| (item % 13) * 1_000 + 1).collect(); + for (item, &weight) in weights.iter().enumerate() { + sketch.update_with_weight(item as i64, weight); + } + let total: u64 = weights.iter().sum(); + assert_eq!(sketch.n(), total); + assert!(sketch.is_estimation_mode()); + assert_eq!(sketch.min_item(), Some(&0)); + assert_eq!(sketch.max_item(), Some(&9_999)); + + let mut weight_below = 0u64; + for (item, &weight) in weights.iter().enumerate() { + let true_rank = weight_below as f64 / total as f64; + let rank = sketch + .rank(&(item as i64), SearchCriteria::Exclusive) + .unwrap(); + assert!( + (rank - true_rank).abs() <= RANK_EPS_FOR_K_200, + "item {item}: rank {rank}, true rank {true_rank}" + ); + weight_below += weight; + } + + let decoded = KllSketch::::deserialize(&sketch.serialize()).unwrap(); + assert_eq!(decoded.n(), sketch.n()); + assert_eq!(decoded.num_retained(), sketch.num_retained()); +} + +#[test] +fn weighted_update_of_strings() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + sketch.update_with_weight("b".to_string(), 3); + sketch.update_with_weight("a".to_string(), 1_000); + sketch.update_with_weight("c".to_string(), 1 << 20); + assert_eq!(sketch.n(), 3 + 1_000 + (1 << 20)); + assert_eq!(sketch.min_item().map(String::as_str), Some("a")); + assert_eq!(sketch.max_item().map(String::as_str), Some("c")); +} + +#[test] +fn weight_overflow_preserves_state() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + sketch.update_with_weight(0, u64::MAX - 1); + let before = sketch.serialize(); + + assert!(catch_unwind(AssertUnwindSafe(|| sketch.update_with_weight(1, 2))).is_err()); + assert!(sketch.serialize() == before, "overflow changed the sketch"); + + sketch.update_with_weight(1, 1); + assert_eq!(sketch.n(), u64::MAX); + assert_eq!(sketch.max_item(), Some(&1)); +}