Add weighted update to kll_sketch

update(item, weight) is equivalent to calling update(item) weight times.
A weight that fits into the free space of level zero is applied as plain
updates; a larger one builds an exact sketch holding one copy of the item
at each level whose bit is set in the weight and merges it, so the cost is
logarithmic in the weight. This mirrors the weighted update in
datasketches-java.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: shanielh

kll/include/kll_sketch.hpp

@@ -237,6 +237,18 @@ class kll_sketch {
     template<typename FwdT>
     void update(FwdT&& item);
 
+    /**
+     * Updates this sketch with the given data item repeated the given number of times.
+     * The result is equivalent to calling update(item) weight times, at a cost that grows
+     * with the logarithm of the weight rather than with the weight itself.
+     * If cross-language portability is required, callers should ensure that
+     * the input string uses a compatible encoding (valid UTF-8).
+     * @param item from a stream of items
+     * @param weight number of times the item is repeated, must be positive
+     */
+    template<typename FwdT>
+    void update(FwdT&& item, uint64_t weight);
+
     /**
      * Merges another sketch into this one.
      * If sketches contain strings, callers are responsible for ensuring that
@@ -570,6 +582,9 @@ class kll_sketch {
         std::unique_ptr<T, items_deleter> items, uint32_t items_size, optional<T>&& min_item,
         optional<T>&& max_item, bool is_level_zero_sorted, const C& comparator);
 
+    // for weighted update
+    kll_sketch(uint16_t k, const T& item, uint64_t weight, const C& comparator, const A& allocator);
+
     // common update code
     inline void update_min_max(const T& item);
     inline uint32_t internal_update();

kll/include/kll_sketch_impl.hpp

@@ -186,6 +186,51 @@ void kll_sketch<T, C, A>::update(FwdT&& item) {
   reset_sorted_view();
 }
 
+template<typename T, typename C, typename A>
+template<typename FwdT>
+void kll_sketch<T, C, A>::update(FwdT&& item, uint64_t weight) {
+  if (weight == 0) throw std::invalid_argument("weight must be positive");
+  if (!check_update_item(item)) { return; }
+  if (weight < levels_[0]) { // fits into level zero without compaction
+    update_min_max(static_cast<const T&>(item));
+    for (uint64_t i = 1; i < weight; ++i) new (&items_[internal_update()]) T(static_cast<const T&>(item));
+    new (&items_[internal_update()]) T(std::forward<FwdT>(item));
+    reset_sorted_view();
+  } else {
+    merge(kll_sketch(k_, static_cast<const T&>(item), weight, comparator_, allocator_));
+  }
+}
+
+// An item at level h carries weight 2^h, so one copy of the item at each level
+// whose bit is set in the weight is an exact sketch of that many copies.
+// Level capacities are defined up to 61 levels, so bits above the top level
+// fold into weight >> 60 copies at level 60.
+template<typename T, typename C, typename A>
+kll_sketch<T, C, A>::kll_sketch(uint16_t k, const T& item, uint64_t weight, const C& comparator, const A& allocator):
+comparator_(comparator),
+allocator_(allocator),
+k_(k),
+m_(kll_constants::DEFAULT_M),
+min_k_(k),
+num_levels_(std::min<uint8_t>(64 - count_leading_zeros_in_u64(weight), 61)),
+is_level_zero_sorted_(true),
+n_(weight),
+levels_(num_levels_ + 1, 0, allocator),
+items_(nullptr),
+items_size_(0),
+min_item_(item),
+max_item_(item),
+sorted_view_(nullptr)
+{
+  for (uint8_t level = 0; level < num_levels_; ++level) {
+    const uint64_t count = level + 1 < num_levels_ ? (weight >> level) & 1 : weight >> level;
+    levels_[level + 1] = levels_[level] + static_cast<uint32_t>(count);
+  }
+  items_size_ = levels_[num_levels_];
+  items_ = allocator_.allocate(items_size_);
+  for (uint32_t i = 0; i < items_size_; ++i) new (&items_[i]) T(item);
+}
+
 template<typename T, typename C, typename A>
 void kll_sketch<T, C, A>::update_min_max(const T& item) {
   if (is_empty()) {

kll/test/kll_sketch_test.cpp

@@ -443,6 +443,9 @@ TEST_CASE("kll sketch", "[kll_sketch]") {
     REQUIRE(kll_helper::floor_of_log2_of_fraction(6, 2) == 1);
     REQUIRE(kll_helper::floor_of_log2_of_fraction(7, 2) == 1);
     REQUIRE(kll_helper::floor_of_log2_of_fraction(8, 2) == 2);
+    REQUIRE(kll_helper::floor_of_log2_of_fraction(1ULL << 63, 1) == 63);
+    REQUIRE(kll_helper::floor_of_log2_of_fraction(std::numeric_limits<uint64_t>::max(), 1) == 63);
+    REQUIRE(kll_helper::floor_of_log2_of_fraction(std::numeric_limits<uint64_t>::max(), 3) == 62);
   }
 
   SECTION("out of order split points, float") {
@@ -831,6 +834,123 @@ TEST_CASE("kll sketch", "[kll_sketch]") {
     REQUIRE(sb.get_n() == 3);
   }
 
+  SECTION("weighted update: zero weight") {
+    kll_float_sketch sketch(200, std::less<float>(), 0);
+    REQUIRE_THROWS_AS(sketch.update(1.0f, 0), std::invalid_argument);
+    REQUIRE(sketch.is_empty());
+  }
+
+  SECTION("weighted update: NaN") {
+    kll_float_sketch sketch(200, std::less<float>(), 0);
+    sketch.update(std::numeric_limits<float>::quiet_NaN(), 1000);
+    REQUIRE(sketch.is_empty());
+  }
+
+  SECTION("weighted update: same as repeated updates in exact mode") {
+    kll_float_sketch weighted(200, std::less<float>(), 0);
+    kll_float_sketch repeated(200, std::less<float>(), 0);
+    for (int i = 1; i <= 15; ++i) {
+      weighted.update(static_cast<float>(i), i);
+      for (int j = 0; j < i; ++j) repeated.update(static_cast<float>(i));
+    }
+    REQUIRE(weighted.get_n() == repeated.get_n());
+    REQUIRE(weighted.get_num_retained() == repeated.get_num_retained());
+    REQUIRE_FALSE(weighted.is_estimation_mode());
+    for (int i = 0; i <= 16; ++i) {
+      REQUIRE(weighted.get_rank(static_cast<float>(i)) == repeated.get_rank(static_cast<float>(i)));
+      REQUIRE(weighted.get_rank(static_cast<float>(i), false) == repeated.get_rank(static_cast<float>(i), false));
+    }
+  }
+
+  SECTION("weighted update: one item with large weight") {
+    kll_float_sketch sketch(200, std::less<float>(), 0);
+    const uint64_t weight = (1ULL << 40) + 12345;
+    sketch.update(7.0f, weight);
+    REQUIRE(sketch.get_n() == weight);
+    REQUIRE(sketch.get_num_retained() == 7); // number of bits set in the weight
+    REQUIRE(sketch.get_min_item() == 7.0f);
+    REQUIRE(sketch.get_max_item() == 7.0f);
+    REQUIRE(sketch.get_rank(7.0f) == 1);
+    REQUIRE(sketch.get_rank(7.0f, false) == 0);
+    REQUIRE(sketch.get_quantile(0.5) == 7.0f);
+
+    uint64_t total_weight = 0;
+    for (auto pair: sketch) total_weight += pair.second;
+    REQUIRE(total_weight == weight);
+  }
+
+  SECTION("weighted update: maximum weight") {
+    kll_float_sketch sketch(200, std::less<float>(), 0);
+    sketch.update(1.0f);
+    sketch.update(2.0f, std::numeric_limits<uint64_t>::max() - 1);
+    REQUIRE(sketch.get_n() == std::numeric_limits<uint64_t>::max());
+    REQUIRE(sketch.get_min_item() == 1.0f);
+    REQUIRE(sketch.get_max_item() == 2.0f);
+    REQUIRE(sketch.get_quantile(0.5) == 2.0f);
+
+    uint64_t total_weight = 0;
+    for (auto pair: sketch) total_weight += pair.second;
+    REQUIRE(total_weight == sketch.get_n());
+
+    auto bytes = sketch.serialize();
+    auto sketch2 = kll_float_sketch::deserialize(bytes.data(), bytes.size(), serde<float>(), std::less<float>(), 0);
+    REQUIRE(sketch2.get_n() == sketch.get_n());
+    REQUIRE(sketch2.get_num_retained() == sketch.get_num_retained());
+  }
+
+  SECTION("weighted update: large weights into a full sketch") {
+    kll_float_sketch sketch(200, std::less<float>(), 0);
+    for (int i = 0; i < 1000; ++i) sketch.update(static_cast<float>(i));
+    sketch.update(-1.0f, 1000);
+    sketch.update(2000.0f, 3000);
+    REQUIRE(sketch.get_n() == 5000);
+    REQUIRE(sketch.get_min_item() == -1.0f);
+    REQUIRE(sketch.get_max_item() == 2000.0f);
+    REQUIRE(sketch.get_rank(-1.0f) == Approx(0.2).margin(RANK_EPS_FOR_K_200));
+    REQUIRE(sketch.get_rank(2000.0f, false) == Approx(0.4).margin(RANK_EPS_FOR_K_200));
+  }
+
+  SECTION("weighted update: estimation mode") {
+    kll_float_sketch sketch(200, std::less<float>(), 0);
+    const int n = 10000;
+    std::vector<uint64_t> weights(n);
+    uint64_t total = 0;
+    for (int i = 0; i < n; ++i) {
+      weights[i] = static_cast<uint64_t>(i % 13) * 1000 + 1;
+      total += weights[i];
+      sketch.update(static_cast<float>(i), weights[i]);
+    }
+    REQUIRE(sketch.get_n() == total);
+    REQUIRE(sketch.is_estimation_mode());
+    REQUIRE(sketch.get_min_item() == 0);
+    REQUIRE(sketch.get_max_item() == n - 1);
+
+    uint64_t weight_below = 0;
+    for (int i = 0; i < n; ++i) {
+      const double true_rank = static_cast<double>(weight_below) / total;
+      REQUIRE(sketch.get_rank(static_cast<float>(i), false) == Approx(true_rank).margin(RANK_EPS_FOR_K_200));
+      weight_below += weights[i];
+    }
+
+    auto bytes = sketch.serialize();
+    auto sketch2 = kll_float_sketch::deserialize(bytes.data(), bytes.size(), serde<float>(), std::less<float>(), 0);
+    REQUIRE(sketch2.get_n() == sketch.get_n());
+    REQUIRE(sketch2.get_num_retained() == sketch.get_num_retained());
+    REQUIRE(sketch2.get_quantile(0.5) == sketch.get_quantile(0.5));
+  }
+
+  SECTION("weighted update: strings") {
+    kll_string_sketch sketch(200, std::less<std::string>(), 0);
+    sketch.update(std::string("b"), 3);
+    const std::string a("a");
+    sketch.update(a, 1000);
+    sketch.update(std::string("c"), 1ULL << 20);
+    REQUIRE(sketch.get_n() == 3 + 1000 + (1ULL << 20));
+    REQUIRE(sketch.get_min_item() == "a");
+    REQUIRE(sketch.get_max_item() == "c");
+    REQUIRE(sketch.get_rank("a") == Approx(1000.0 / sketch.get_n()).margin(RANK_EPS_FOR_K_200));
+  }
+
   // cleanup
   REQUIRE(test_allocator_total_bytes == 0);
 }