From b3f2c9b84e4c754eba2381f87a0bb52bf88fc562 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Sun, 6 Sep 2026 21:56:27 -0700 Subject: [PATCH 1/2] Add optional trim flag to update_theta_sketch::compact() An update sketch retains more than the nominal size k between rebuilds. Callers who need a result bounded by k currently have to copy the sketch, call trim() on the copy, then compact() it, which mutates state and allocates twice. This adds an optional second flag to the existing compact(): compact_theta_sketch_alloc compact(bool ordered = true, bool trim = false) const; The trim path copies the retained entries (the method is const, and quick select would permute the live table), partitions with std::nth_element at 0-based index k so that entries[k] becomes the new theta, keeps the k entries below it, and sorts only when an ordered result is requested. This matches the pivot convention already used by rebuild(), so trimming and rebuilding cannot drift apart. Trimming stays opt-in because it is lossy. Relative error scales with 1 / sqrt(retained), so discarding entries always widens the confidence bounds: measured 1.03x to 1.27x here, and up to about sqrt(15/8) (~37%) worst case for a sketch grown to just under the 15/16 * 2k rebuild threshold. A sketch in exact mode that retains more than k entries also loses exactness and is returned in estimation mode. Both effects are documented on the parameter and pinned by tests. Tests cover the four ordered/trim combinations, that the source sketch is unmodified, that trimmed output matches trim() + compact() exactly, that every retained hash is below the new theta, the exact-to-estimation conversion, the bounds widening, and the empty and below-k cases. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_017EDHa7UhfW4eSj5L82panJ --- theta/include/theta_sketch.hpp | 19 ++++- theta/include/theta_sketch_impl.hpp | 24 +++++- theta/test/theta_sketch_test.cpp | 125 ++++++++++++++++++++++++++++ 3 files changed, 164 insertions(+), 4 deletions(-) diff --git a/theta/include/theta_sketch.hpp b/theta/include/theta_sketch.hpp index 4aab4b92..142588e7 100644 --- a/theta/include/theta_sketch.hpp +++ b/theta/include/theta_sketch.hpp @@ -330,11 +330,25 @@ class update_theta_sketch_alloc: public theta_sketch_alloc { 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 compact(bool ordered = true) const; + compact_theta_sketch_alloc 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 { template friend class theta_union_base; template friend class theta_intersection_base; template friend class theta_set_difference_base; + template 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&& entries); }; diff --git a/theta/include/theta_sketch_impl.hpp b/theta/include/theta_sketch_impl.hpp index b976ade9..b92b6ca1 100644 --- a/theta/include/theta_sketch_impl.hpp +++ b/theta/include/theta_sketch_impl.hpp @@ -23,6 +23,7 @@ #include #include #include +#include #include "binomial_bounds.hpp" #include "theta_helpers.hpp" @@ -239,8 +240,27 @@ auto update_theta_sketch_alloc::end() const -> const_iterator { } template -compact_theta_sketch_alloc update_theta_sketch_alloc::compact(bool ordered) const { - return compact_theta_sketch_alloc(*this, ordered); +compact_theta_sketch_alloc update_theta_sketch_alloc::compact(bool ordered, bool trim) const { + if (!trim) return compact_theta_sketch_alloc(*this, ordered); + + std::vector entries(table_.allocator_); + if (this->is_empty()) { + return compact_theta_sketch_alloc(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(false, ordered, this->get_seed_hash(), theta, std::move(entries)); } template diff --git a/theta/test/theta_sketch_test.cpp b/theta/test/theta_sketch_test.cpp index 95df9a92..d981f077 100644 --- a/theta/test/theta_sketch_test.cpp +++ b/theta/test/theta_sketch_test.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -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 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(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 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(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); From e13672e8a19930a0a7faa5dd1623178f224f5004 Mon Sep 17 00:00:00 2001 From: Lee Rhodes Date: Wed, 9 Sep 2026 16:20:58 -0700 Subject: [PATCH 2/2] Share one trim implementation between union and compact Review feedback: compact(ordered, trim) omitted the shrink_to_fit that theta_union_base::get_result performs after the same erase. The two were copies of one algorithm and the newer copy had already drifted, so share it rather than patch the copy. Adds trim_to_nominal to theta_helpers.hpp, generic over ExtractKey, and calls it from both sites. theta_union_base is instantiated by both theta_union and tuple_union, so all three paths now run identical code; theta's compact is the degenerate case where the entry is the key. The missing shrink matters most where it was missing. An update sketch retains up to 15/16 * 2k entries between rebuilds, and the entries vector is reserved to that before being erased down to k, so a trimmed result kept the untrimmed allocation. Measured at lg_k=12: 7674 retained against 4096 kept, holding 1.87x the memory it reports, 46.6% of the allocation dead. The whole point of trim is a bounded result, so the caller was paying the accuracy cost without getting the size back. The invariant is documented once now: entries[nominal_size] becomes theta and is itself discarded, so everything kept is strictly below it. An off-by-one there does not fail loudly, it biases every estimate. Verified behaviour-preserving: serialized bytes of both the trimmed compact and the union result are identical before and after, across ordered and unordered at 5000, 26989 and 40000 updates. theta 89 cases / 20250056 assertions and tuple 55 cases / 126786 assertions pass. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BKM9nrFVMhH2gmFaag2JuC --- theta/include/theta_helpers.hpp | 32 +++++++++++++++++++++++++ theta/include/theta_sketch_impl.hpp | 8 +------ theta/include/theta_union_base_impl.hpp | 8 ++----- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/theta/include/theta_helpers.hpp b/theta/include/theta_helpers.hpp index 8d83d341..f03e8c99 100644 --- a/theta/include/theta_helpers.hpp +++ b/theta/include/theta_helpers.hpp @@ -20,13 +20,45 @@ #ifndef THETA_HELPERS_HPP_ #define THETA_HELPERS_HPP_ +#include +#include #include #include +#include #include "theta_constants.hpp" +#include "theta_comparators.hpp" namespace datasketches { +/** + * Trims a vector of theta entries down to at most nominal_size and returns the theta the + * result must carry. + * + * The entry at index nominal_size becomes the new theta and is itself discarded, so every + * entry kept is strictly below the returned value. That is what keeps the estimator + * unbiased: theta must be an exclusive upper bound on the retained hashes. Getting this + * off by one does not fail loudly, it quietly biases every estimate the sketch produces. + * + * Capacity is released as well as size. Callers reserve an upper bound before filling, so + * without the shrink a trimmed result keeps the untrimmed allocation, which for an update + * sketch just under the rebuild threshold is nearly twice what it reports. + * + * @param entries entries to trim in place; reordered even when nothing is removed + * @param nominal_size the most entries to keep + * @param theta returned unchanged when there is nothing to trim + * @return the theta of the trimmed result + */ +template +static uint64_t trim_to_nominal(std::vector& entries, uint32_t nominal_size, uint64_t theta) { + if (entries.size() <= nominal_size) return theta; + std::nth_element(entries.begin(), entries.begin() + nominal_size, entries.end(), compare_by_key()); + const uint64_t new_theta = ExtractKey()(entries[nominal_size]); + entries.erase(entries.begin() + nominal_size, entries.end()); + entries.shrink_to_fit(); + return new_theta; +} + template static void check_value(T actual, T expected, const char* description) { if (actual != expected) { diff --git a/theta/include/theta_sketch_impl.hpp b/theta/include/theta_sketch_impl.hpp index b92b6ca1..97fd9938 100644 --- a/theta/include/theta_sketch_impl.hpp +++ b/theta/include/theta_sketch_impl.hpp @@ -252,13 +252,7 @@ compact_theta_sketch_alloc update_theta_sketch_alloc::compact(bool ordered 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()); - } + theta = trim_to_nominal(entries, nominal_size, theta); if (ordered) std::sort(entries.begin(), entries.end()); return compact_theta_sketch_alloc(false, ordered, this->get_seed_hash(), theta, std::move(entries)); } diff --git a/theta/include/theta_union_base_impl.hpp b/theta/include/theta_union_base_impl.hpp index 99a5bbf7..513f1ea3 100644 --- a/theta/include/theta_union_base_impl.hpp +++ b/theta/include/theta_union_base_impl.hpp @@ -24,6 +24,7 @@ #include #include "conditional_forward.hpp" +#include "theta_helpers.hpp" namespace datasketches { @@ -70,12 +71,7 @@ CS theta_union_base::get_result(bool ordered) const { } else { std::copy_if(table_.begin(), table_.end(), std::back_inserter(entries), key_not_zero_less_than(theta)); } - if (entries.size() > nominal_num) { - std::nth_element(entries.begin(), entries.begin() + nominal_num, entries.end(), comparator()); - theta = EK()(entries[nominal_num]); - entries.erase(entries.begin() + nominal_num, entries.end()); - entries.shrink_to_fit(); - } + theta = trim_to_nominal(entries, nominal_num, theta); if (ordered) std::sort(entries.begin(), entries.end(), comparator()); return CS(table_.is_empty_, ordered, compute_seed_hash(table_.seed_), theta, std::move(entries)); }