feat(kll): add weighted update

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<TDigestMut>` as an alternative to incremental `TDigestMut::merge` calls for combining owned partial sketches in a single compression pass.
 

datasketches/src/kll/sketch.rs

@@ -143,6 +143,72 @@ impl<T: Clone + Ord> KllSketch<T> {
         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::<i64>::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<T: Clone + Ord> KllSketch<T> {
                 other.n
             ))
         })?;
+        self.merge_unchecked(other, final_n);
+        Ok(())
+    }
 
+    fn merge_unchecked(&mut self, other: &KllSketch<T>, final_n: u64) {
         self.update_min_max_from_other(other);
 
         for item in &other.levels[0] {
@@ -197,7 +267,6 @@ impl<T: Clone + Ord> KllSketch<T> {
         }
 
         debug_assert_eq!(self.total_weight(), self.n, "total weight does not match n");
-        Ok(())
     }
 
     /// Returns the normalized rank of the given item.

tests-integration/tests/kll_test/main.rs

@@ -19,3 +19,4 @@ mod core;
 mod generic;
 mod merge;
 mod query;
+mod update;

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::<i64>::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::<i64>::new(DEFAULT_K).unwrap();
+    let mut repeated = KllSketch::<i64>::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::<i64>::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::<i64>::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::<i64>::deserialize(&sketch.serialize()).unwrap();
+    assert_eq!(decoded, sketch);
+}
+
+#[test]
+fn weighted_updates_into_a_full_sketch() {
+    let mut sketch = KllSketch::<i64>::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::<i64>::new(DEFAULT_K).unwrap();
+    let weights: Vec<u64> = (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::<i64>::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::<String>::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::<i64>::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));
+}