Add get_live_bytes() to report the current heap footprint of theta, tuple and kll sketches

kll/include/kll_sketch.hpp

@@ -410,6 +410,18 @@ class kll_sketch {
     template<typename TT = T, typename SerDe = serde<T>, typename std::enable_if<!std::is_arithmetic<TT>::value, int>::type = 0>
     size_t get_serialized_size_bytes(const SerDe& sd = SerDe()) const;
 
+    /**
+     * Returns the number of bytes currently allocated on the heap by this sketch: the items_ buffer
+     * (items_size_ elements) plus the levels_ vector. Unlike a hash-table sketch this is not a table
+     * footprint; it mirrors the sketch's allocate() sites. This is the exact live heap footprint
+     * right now, not an upper bound like get_serialized_size_bytes().
+     * @return the current allocated heap size of the sketch in bytes
+     */
+    size_t get_live_bytes() const {
+      return (items_ == nullptr ? 0 : static_cast<size_t>(items_size_) * sizeof(T))
+          + levels_.capacity() * sizeof(uint32_t);
+    }
+
     /**
      * Returns upper bound on the serialized size of a sketch given a parameter <em>k</em> and stream
      * length. The resulting size is an overestimate to make sure actual sketches don't exceed it.

kll/test/kll_sketch_test.cpp

@@ -835,4 +835,23 @@ TEST_CASE("kll sketch", "[kll_sketch]") {
   REQUIRE(test_allocator_total_bytes == 0);
 }
 
+TEST_CASE("kll sketch: get_live_bytes", "[kll_sketch]") {
+  kll_sketch<float> sketch;
+  // a fresh sketch has already allocated its items buffer and level boundaries
+  const size_t empty_bytes = sketch.get_live_bytes();
+  REQUIRE(empty_bytes > 0);
+
+  // the items buffer and levels only grow as the sketch fills, and grow past the initial size
+  size_t prev_bytes = empty_bytes;
+  bool non_decreasing = true;
+  for (int i = 0; i < 1000000; ++i) {
+    sketch.update(static_cast<float>(i));
+    const size_t bytes = sketch.get_live_bytes();
+    if (bytes < prev_bytes) non_decreasing = false;
+    prev_bytes = bytes;
+  }
+  REQUIRE(non_decreasing);
+  REQUIRE(prev_bytes > empty_bytes);
+}
+
 } /* namespace datasketches */

theta/include/theta_intersection.hpp

@@ -83,6 +83,13 @@ class theta_intersection_alloc {
    */
   bool has_result() const;
 
+  /**
+   * Returns the number of bytes currently allocated for this intersection's internal hash table.
+   * This is the exact live heap footprint right now, not an upper bound.
+   * @return the current allocated size of the internal hash table in bytes
+   */
+  size_t get_live_bytes() const { return state_.get_live_bytes(); }
+
 private:
   State state_;
 };

theta/include/theta_intersection_base.hpp

@@ -46,6 +46,9 @@ class theta_intersection_base {
 
   const Policy& get_policy() const;
 
+  // Live heap footprint of the internal hash table in bytes.
+  size_t get_live_bytes() const { return table_.get_live_bytes(); }
+
 private:
   Policy policy_;
   bool is_valid_;

theta/include/theta_sketch.hpp

@@ -341,6 +341,14 @@ class update_theta_sketch_alloc: public theta_sketch_alloc<Allocator> {
   virtual const_iterator begin() const;
   virtual const_iterator end() const;
 
+  /**
+   * Returns the number of bytes currently allocated for this sketch's internal hash table.
+   * This is the exact live heap footprint of the retained entries right now, not an upper bound
+   * like get_serialized_size_bytes().
+   * @return the current allocated size of the internal hash table in bytes
+   */
+  size_t get_live_bytes() const { return table_.get_live_bytes(); }
+
 private:
   theta_table table_;
 

theta/include/theta_union.hpp

@@ -74,6 +74,13 @@ class theta_union_alloc {
   /// Reset the union to the initial empty state
   void reset();
 
+  /**
+   * Returns the number of bytes currently allocated for this union's internal hash table.
+   * This is the exact live heap footprint right now, not an upper bound.
+   * @return the current allocated size of the internal hash table in bytes
+   */
+  size_t get_live_bytes() const { return state_.get_live_bytes(); }
+
 private:
   State state_;
 

theta/include/theta_union_base.hpp

@@ -49,6 +49,9 @@ class theta_union_base {
 
   void reset();
 
+  // Live heap footprint of the internal hash table in bytes.
+  size_t get_live_bytes() const { return table_.get_live_bytes(); }
+
 private:
   Policy policy_;
   hash_table table_;

theta/include/theta_update_sketch_base.hpp

@@ -62,6 +62,13 @@ struct theta_update_sketch_base {
   iterator begin() const;
   iterator end() const;
 
+  // Bytes currently allocated for the entries_ hash table: (1 << lg_cur_size_) entries, or 0 when
+  // the table is unallocated (lg_cur_size_ == 0 leaves entries_ == nullptr). This is the exact live
+  // heap footprint of the table right now, not an upper bound like the serialized-size estimates.
+  size_t get_live_bytes() const {
+    return entries_ == nullptr ? 0 : (static_cast<size_t>(1) << lg_cur_size_) * sizeof(Entry);
+  }
+
   // resize threshold = 0.5 tuned for speed
   static constexpr double RESIZE_THRESHOLD = 0.5;
   // hash table rebuild threshold = 15/16

theta/test/theta_intersection_test.cpp

@@ -237,4 +237,19 @@ TEST_CASE("theta intersection: seed mismatch", "[theta_intersection]") {
   REQUIRE_THROWS_AS(intersection.update(sketch), std::invalid_argument);
 }
 
+TEST_CASE("theta intersection: get_live_bytes", "[theta_intersection]") {
+  theta_intersection intersection;
+  // before any update the intersection holds no table
+  REQUIRE(intersection.get_live_bytes() == 0);
+
+  update_theta_sketch sketch = update_theta_sketch::builder().build();
+  for (int i = 0; i < 100000; ++i) sketch.update(i);
+  intersection.update(sketch);
+  const size_t bytes = intersection.get_live_bytes();
+  REQUIRE(bytes > 0);
+  REQUIRE(bytes % sizeof(uint64_t) == 0);
+  const size_t entries = bytes / sizeof(uint64_t);
+  REQUIRE((entries & (entries - 1)) == 0);
+}
+
 } /* namespace datasketches */

theta/test/theta_sketch_test.cpp

@@ -625,4 +625,28 @@ TEST_CASE("max serialized size", "[theta_sketch]") {
   REQUIRE(max_size_bytes == compact_theta_sketch::get_max_serialized_size_bytes(lg_k));
 }
 
+TEST_CASE("theta sketch: get_live_bytes", "[theta_sketch]") {
+  update_theta_sketch sketch = update_theta_sketch::builder().build();
+  // a freshly built update sketch has an allocated hash table of 8-byte entries
+  const size_t empty_bytes = sketch.get_live_bytes();
+  REQUIRE(empty_bytes > 0);
+  REQUIRE(empty_bytes % sizeof(uint64_t) == 0);
+  const size_t empty_entries = empty_bytes / sizeof(uint64_t);
+  REQUIRE((empty_entries & (empty_entries - 1)) == 0); // capacity is a power of two
+
+  // the table only grows as distinct keys are inserted, and it grows past the initial size
+  size_t prev_bytes = empty_bytes;
+  bool non_decreasing = true;
+  for (int i = 0; i < 100000; ++i) {
+    sketch.update(i);
+    const size_t bytes = sketch.get_live_bytes();
+    if (bytes < prev_bytes) non_decreasing = false;
+    prev_bytes = bytes;
+  }
+  REQUIRE(non_decreasing);
+  REQUIRE(prev_bytes > empty_bytes);
+  const size_t entries = prev_bytes / sizeof(uint64_t);
+  REQUIRE((entries & (entries - 1)) == 0);
+}
+
 } /* namespace datasketches */

theta/test/theta_union_test.cpp

@@ -153,4 +153,19 @@ TEST_CASE("theta union: larger K", "[theta_union]") {
   REQUIRE(result2.get_estimate() == update_sketch3.get_estimate());
 }
 
+TEST_CASE("theta union: get_live_bytes", "[theta_union]") {
+  update_theta_sketch update_sketch = update_theta_sketch::builder().build();
+  for (int i = 0; i < 100000; ++i) update_sketch.update(i);
+
+  theta_union u = theta_union::builder().build();
+  const size_t empty_bytes = u.get_live_bytes();
+  u.update(update_sketch);
+  const size_t bytes = u.get_live_bytes();
+  // unioning a large sketch grows the internal table past its initial footprint
+  REQUIRE(bytes > empty_bytes);
+  REQUIRE(bytes % sizeof(uint64_t) == 0);
+  const size_t entries = bytes / sizeof(uint64_t);
+  REQUIRE((entries & (entries - 1)) == 0);
+}
+
 } /* namespace datasketches */

tuple/include/tuple_intersection.hpp

@@ -100,6 +100,13 @@ class tuple_intersection {
    */
   bool has_result() const;
 
+  /**
+   * Returns the number of bytes currently allocated for this intersection's internal hash table.
+   * This is the exact live heap footprint right now, not an upper bound.
+   * @return the current allocated size of the internal hash table in bytes
+   */
+  size_t get_live_bytes() const { return state_.get_live_bytes(); }
+
 protected:
   State state_;
 };

tuple/include/tuple_sketch.hpp

@@ -436,6 +436,14 @@ class update_tuple_sketch: public tuple_sketch<Summary, Allocator> {
   virtual const_iterator begin() const;
   virtual const_iterator end() const;
 
+  /**
+   * Returns the number of bytes currently allocated for this sketch's internal hash table.
+   * This is the exact live heap footprint of the retained entries right now, not an upper bound
+   * like get_serialized_size_bytes().
+   * @return the current allocated size of the internal hash table in bytes
+   */
+  size_t get_live_bytes() const { return map_.get_live_bytes(); }
+
 protected:
   Policy policy_;
   tuple_map map_;

tuple/include/tuple_union.hpp

@@ -89,6 +89,13 @@ class tuple_union {
    */
   void reset();
 
+  /**
+   * Returns the number of bytes currently allocated for this union's internal hash table.
+   * This is the exact live heap footprint right now, not an upper bound.
+   * @return the current allocated size of the internal hash table in bytes
+   */
+  size_t get_live_bytes() const { return state_.get_live_bytes(); }
+
 protected:
   State state_;
 

tuple/test/tuple_sketch_test.cpp

@@ -386,4 +386,23 @@ TEST_CASE("tuple sketch: deserialize bounds-checks each entry key", "[tuple_sket
                     std::out_of_range);
 }
 
+TEST_CASE("tuple sketch: get_live_bytes", "[tuple_sketch]") {
+  auto sketch = update_tuple_sketch<float>::builder().build();
+  // a freshly built update sketch has an allocated hash table
+  const size_t empty_bytes = sketch.get_live_bytes();
+  REQUIRE(empty_bytes > 0);
+
+  // the table only grows as distinct keys are inserted, and it grows past the initial size
+  size_t prev_bytes = empty_bytes;
+  bool non_decreasing = true;
+  for (int i = 0; i < 100000; ++i) {
+    sketch.update(i, 1.0f);
+    const size_t bytes = sketch.get_live_bytes();
+    if (bytes < prev_bytes) non_decreasing = false;
+    prev_bytes = bytes;
+  }
+  REQUIRE(non_decreasing);
+  REQUIRE(prev_bytes > empty_bytes);
+}
+
 } /* namespace datasketches */