Add optional trim flag to update_theta_sketch::compact()

theta/include/theta_sketch.hpp

@@ -330,11 +330,25 @@ class update_theta_sketch_alloc: public theta_sketch_alloc<Allocator> {
   void reset();
 
   /**
-   * Converts this sketch to a compact sketch (ordered or unordered).
+   * Converts this sketch to a compact sketch (ordered or unordered, trimmed or not trimmed).
+   * This does not modify the source update sketch.
    * @param ordered optional flag to specify if an ordered sketch should be produced
+   * @param trim optional flag to reduce the size of the returned sketch to at most
+   * the nominal size k, if required. An update sketch retains more than k entries
+   * between rebuilds; those extra entries below theta improve the estimate, so the
+   * default is to keep them.
+   * Trimming is lossy and is never required for correctness. It discards retained
+   * entries, and since the relative error scales with 1 / sqrt(retained), it always
+   * degrades accuracy and widens the confidence bounds, whatever mode the source is
+   * in. Worst case, a sketch grown to just under the rebuild threshold of 15/16 * 2k
+   * loses nearly half its entries, widening the bounds by about sqrt(15/8), or
+   * roughly 37%. A sketch in exact mode that retains more than k entries loses
+   * exactness as well: it is returned in estimation mode, so get_estimate() carries
+   * error where it would otherwise have returned an exact count.
+   * Only pass true if a bounded result size matters more than that accuracy.
    * @return compact sketch
    */
-  compact_theta_sketch_alloc<Allocator> compact(bool ordered = true) const;
+  compact_theta_sketch_alloc<Allocator> compact(bool ordered = true, bool trim = false) const;
 
   virtual iterator begin();
   virtual iterator end();
@@ -518,6 +532,7 @@ class compact_theta_sketch_alloc: public theta_sketch_alloc<Allocator> {
   template<typename E, typename EK, typename P, typename S, typename CS, typename A> friend class theta_union_base;
   template<typename E, typename EK, typename P, typename S, typename CS, typename A> friend class theta_intersection_base;
   template<typename E, typename EK, typename CS, typename A> friend class theta_set_difference_base;
+  template<typename A> friend class update_theta_sketch_alloc;
   compact_theta_sketch_alloc(bool is_empty, bool is_ordered, uint16_t seed_hash, uint64_t theta, std::vector<uint64_t, Allocator>&& entries);
 };
 

theta/include/theta_sketch_impl.hpp

@@ -23,6 +23,7 @@
 #include <sstream>
 #include <vector>
 #include <stdexcept>
+#include <algorithm>
 
 #include "binomial_bounds.hpp"
 #include "theta_helpers.hpp"
@@ -239,8 +240,27 @@ auto update_theta_sketch_alloc<A>::end() const -> const_iterator {
 }
 
 template<typename A>
-compact_theta_sketch_alloc<A> update_theta_sketch_alloc<A>::compact(bool ordered) const {
-  return compact_theta_sketch_alloc<A>(*this, ordered);
+compact_theta_sketch_alloc<A> update_theta_sketch_alloc<A>::compact(bool ordered, bool trim) const {
+  if (!trim) return compact_theta_sketch_alloc<A>(*this, ordered);
+
+  std::vector<uint64_t, A> entries(table_.allocator_);
+  if (this->is_empty()) {
+    return compact_theta_sketch_alloc<A>(true, true, this->get_seed_hash(), this->get_theta64(), std::move(entries));
+  }
+  // copy first: this method is const, and quick select would permute the source table
+  entries.reserve(this->get_num_retained());
+  std::copy(this->begin(), this->end(), std::back_inserter(entries));
+  uint64_t theta = this->get_theta64();
+  const uint32_t nominal_size = 1 << table_.lg_nom_size_;
+  if (entries.size() > nominal_size) {
+    // partial sort so that entries[nominal_size] is the (nominal_size + 1)-th smallest hash;
+    // it becomes the new theta, and the nominal_size entries below it are all we keep
+    std::nth_element(entries.begin(), entries.begin() + nominal_size, entries.end());
+    theta = entries[nominal_size];
+    entries.erase(entries.begin() + nominal_size, entries.end());
+  }
+  if (ordered) std::sort(entries.begin(), entries.end());
+  return compact_theta_sketch_alloc<A>(false, ordered, this->get_seed_hash(), theta, std::move(entries));
 }
 
 template<typename A>

theta/test/theta_sketch_test.cpp

@@ -22,6 +22,7 @@
 #include <sstream>
 #include <vector>
 #include <stdexcept>
+#include <algorithm>
 
 #include <catch2/catch.hpp>
 #include <theta_sketch.hpp>
@@ -151,6 +152,130 @@ TEST_CASE("theta sketch: single item", "[theta_sketch]") {
   REQUIRE(update_sketch.compact(false).is_ordered());
 }
 
+TEST_CASE("theta sketch: compact with trim, all four cases", "[theta_sketch]") {
+  update_theta_sketch update_sketch = update_theta_sketch::builder().build();
+  for (int i = 0; i < 8000; i++) update_sketch.update(i);
+  const uint32_t k = 1 << theta_constants::DEFAULT_LG_K;
+  // an update sketch retains more than k between rebuilds, so trimming has work to do
+  REQUIRE(update_sketch.get_num_retained() > k);
+  const uint32_t retained_before = update_sketch.get_num_retained();
+  const uint64_t theta_before = update_sketch.get_theta64();
+
+  // the trimmed result is what trim() + compact() would produce
+  update_theta_sketch trimmed = update_sketch;
+  trimmed.trim();
+  compact_theta_sketch expected = trimmed.compact(true);
+  const std::vector<uint64_t> expected_entries(expected.begin(), expected.end());
+
+  // case 1: ordered, not trimmed (the default) keeps every retained entry
+  compact_theta_sketch c1 = update_sketch.compact(true, false);
+  REQUIRE(c1.is_ordered());
+  REQUIRE(c1.get_num_retained() == retained_before);
+  REQUIRE(c1.get_theta64() == theta_before);
+  REQUIRE(std::is_sorted(c1.begin(), c1.end()));
+
+  // case 2: unordered, not trimmed
+  compact_theta_sketch c2 = update_sketch.compact(false, false);
+  REQUIRE_FALSE(c2.is_ordered());
+  REQUIRE(c2.get_num_retained() == retained_before);
+  REQUIRE(c2.get_theta64() == theta_before);
+
+  // case 3: ordered and trimmed
+  compact_theta_sketch c3 = update_sketch.compact(true, true);
+  REQUIRE(c3.is_ordered());
+  REQUIRE(c3.get_num_retained() == k);
+  REQUIRE(c3.get_theta64() < theta_before); // new theta is the k-th smallest hash
+  REQUIRE(c3.get_theta64() == expected.get_theta64());
+  REQUIRE(std::is_sorted(c3.begin(), c3.end()));
+  REQUIRE(std::vector<uint64_t>(c3.begin(), c3.end()) == expected_entries);
+  // every retained hash is strictly below the new theta
+  for (auto h: c3) REQUIRE(h < c3.get_theta64());
+
+  // case 4: unordered and trimmed: same set and theta, no sort
+  compact_theta_sketch c4 = update_sketch.compact(false, true);
+  REQUIRE_FALSE(c4.is_ordered());
+  REQUIRE(c4.get_num_retained() == k);
+  REQUIRE(c4.get_theta64() == expected.get_theta64());
+  std::vector<uint64_t> c4_entries(c4.begin(), c4.end());
+  std::sort(c4_entries.begin(), c4_entries.end());
+  REQUIRE(c4_entries == expected_entries);
+
+  // the source sketch must be untouched by any of the four
+  REQUIRE(update_sketch.get_num_retained() == retained_before);
+  REQUIRE(update_sketch.get_theta64() == theta_before);
+}
+
+TEST_CASE("theta sketch: compact with trim widens the bounds in estimation mode", "[theta_sketch]") {
+  // trimming is lossy even when the source is already estimating: it discards
+  // retained entries, and the relative error scales with 1 / sqrt(retained)
+  const uint32_t k = 1 << theta_constants::DEFAULT_LG_K;
+  update_theta_sketch sketch = update_theta_sketch::builder().build();
+  for (int i = 0; i < 40000; i++) sketch.update(i);
+  REQUIRE(sketch.is_estimation_mode());
+  REQUIRE(sketch.get_num_retained() > k);
+
+  compact_theta_sketch plain = sketch.compact(true, false);
+  compact_theta_sketch trimmed = sketch.compact(true, true);
+  REQUIRE(trimmed.get_num_retained() == k);
+  REQUIRE(plain.get_num_retained() > trimmed.get_num_retained());
+
+  // fewer retained entries -> strictly wider confidence interval
+  const double plain_width = plain.get_upper_bound(2) - plain.get_lower_bound(2);
+  const double trimmed_width = trimmed.get_upper_bound(2) - trimmed.get_lower_bound(2);
+  REQUIRE(trimmed_width > plain_width);
+}
+
+TEST_CASE("theta sketch: compact with trim converts exact mode to estimation", "[theta_sketch]") {
+  // an update sketch can retain far more than k entries while still in exact mode:
+  // nothing has been evicted yet, so theta is still 1.0 and the count is exact
+  const uint32_t k = 1 << theta_constants::DEFAULT_LG_K;
+  update_theta_sketch sketch = update_theta_sketch::builder().build();
+  const int n = 5000;
+  for (int i = 0; i < n; i++) sketch.update(i);
+  REQUIRE(sketch.get_num_retained() > k);
+  REQUIRE_FALSE(sketch.is_estimation_mode());
+  REQUIRE(sketch.get_theta() == 1.0);
+
+  // not trimming keeps every entry and the exact count
+  compact_theta_sketch exact = sketch.compact(true, false);
+  REQUIRE_FALSE(exact.is_estimation_mode());
+  REQUIRE(exact.get_num_retained() == static_cast<uint32_t>(n));
+  REQUIRE(exact.get_estimate() == Approx(n));
+
+  // trimming is lossy: it discards real data, lowers theta below 1.0 and the
+  // result is an estimate carrying error where the source held an exact count
+  compact_theta_sketch trimmed = sketch.compact(true, true);
+  REQUIRE(trimmed.is_estimation_mode());
+  REQUIRE(trimmed.get_num_retained() == k);
+  REQUIRE(trimmed.get_theta() < 1.0);
+  REQUIRE(trimmed.get_estimate() != Approx(n));
+  // the estimate is still sound: n must lie inside the 3-sigma bounds
+  REQUIRE(trimmed.get_lower_bound(3) <= n);
+  REQUIRE(trimmed.get_upper_bound(3) >= n);
+}
+
+TEST_CASE("theta sketch: compact with trim, empty and below k", "[theta_sketch]") {
+  // empty: trimming changes nothing
+  update_theta_sketch empty_sketch = update_theta_sketch::builder().build();
+  compact_theta_sketch empty_result = empty_sketch.compact(true, true);
+  REQUIRE(empty_result.is_empty());
+  REQUIRE(empty_result.get_num_retained() == 0);
+  REQUIRE(empty_result.is_ordered());
+
+  // below k: nothing to trim, theta and entries are preserved
+  update_theta_sketch small = update_theta_sketch::builder().build();
+  for (int i = 0; i < 100; i++) small.update(i);
+  REQUIRE_FALSE(small.is_estimation_mode());
+
+  compact_theta_sketch small_result = small.compact(true, true);
+  REQUIRE_FALSE(small_result.is_estimation_mode());
+  REQUIRE(small_result.get_num_retained() == 100);
+  REQUIRE(small_result.get_theta64() == small.get_theta64());
+  REQUIRE(small_result.get_estimate() == Approx(100.0));
+  REQUIRE(std::is_sorted(small_result.begin(), small_result.end()));
+  REQUIRE_FALSE(small.compact(false, true).is_ordered());
+}
+
 TEST_CASE("theta sketch: resize exact", "[theta_sketch]") {
   update_theta_sketch update_sketch = update_theta_sketch::builder().build();
   for (int i = 0; i < 2000; i++) update_sketch.update(i);