Add update_theta_sketch::get_result() to trim to k in one pass

Returns a compact_theta_sketch bounded to the nominal size k (2^lg_k),
matching the at-most-k guarantee of theta_union::get_result(). It does
the nth_element/erase cutback directly on the output vector compact()
already allocates, avoiding the throwaway 2k rehash of trim()+compact().

Co-authored-by: Isaac

Author: stojkomilos

theta/include/theta_sketch.hpp

@@ -336,6 +336,13 @@ class update_theta_sketch_alloc: public theta_sketch_alloc<Allocator> {
    */
   compact_theta_sketch_alloc<Allocator> compact(bool ordered = true) const;
 
+  /**
+   * Produces a compact sketch trimmed to the nominal size k in a single pass.
+   * Like trim() followed by compact(), but without rebuilding the hash table. Result is unordered.
+   * @return compact sketch with at most k retained entries
+   */
+  compact_theta_sketch_alloc<Allocator> get_result() const;
+
   virtual iterator begin();
   virtual iterator end();
   virtual const_iterator begin() const;
@@ -518,6 +525,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"
@@ -243,6 +244,25 @@ compact_theta_sketch_alloc<A> update_theta_sketch_alloc<A>::compact(bool ordered
   return compact_theta_sketch_alloc<A>(*this, ordered);
 }
 
+template<typename A>
+compact_theta_sketch_alloc<A> update_theta_sketch_alloc<A>::get_result() const {
+  std::vector<uint64_t, A> entries(table_.allocator_);
+  if (is_empty()) {
+    return compact_theta_sketch_alloc<A>(true, true, get_seed_hash(), get_theta64(), std::move(entries));
+  }
+  entries.reserve(get_num_retained());
+  std::copy(begin(), end(), std::back_inserter(entries));
+  uint64_t theta = table_.theta_;
+  const uint32_t nominal_num = 1 << table_.lg_nom_size_;
+  if (entries.size() > nominal_num) {
+    std::nth_element(entries.begin(), entries.begin() + nominal_num, entries.end());
+    theta = entries[nominal_num];
+    entries.erase(entries.begin() + nominal_num, entries.end());
+  }
+  const bool ordered = entries.size() <= 1;
+  return compact_theta_sketch_alloc<A>(false, ordered, get_seed_hash(), theta, std::move(entries));
+}
+
 template<typename A>
 void update_theta_sketch_alloc<A>::print_specifics(std::ostringstream& os) const {
   os << "   lg nominal size      : " << static_cast<int>(table_.lg_nom_size_) << std::endl;

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>
@@ -167,6 +168,49 @@ TEST_CASE("theta sketch: estimation", "[theta_sketch]") {
   REQUIRE(compact_sketch.get_upper_bound(1) > n);
 }
 
+TEST_CASE("theta sketch: get_result trims to k in one pass", "[theta_sketch]") {
+  update_theta_sketch update_sketch = update_theta_sketch::builder().build();
+  const int n = 8000;
+  for (int i = 0; i < n; i++) update_sketch.update(i);
+  const uint32_t k = 1 << theta_constants::DEFAULT_LG_K;
+  REQUIRE(update_sketch.get_num_retained() > k); // over-provisioned before trimming
+
+  compact_theta_sketch result = update_sketch.get_result();
+  REQUIRE_FALSE(result.is_empty());
+  REQUIRE(result.is_estimation_mode());
+  REQUIRE(result.get_num_retained() == k); // trimmed to nominal size
+  REQUIRE_FALSE(result.is_ordered());      // unordered: no sort performed
+
+  // fused get_result() must match trim() + compact()
+  update_theta_sketch trimmed = update_sketch;
+  trimmed.trim();
+  compact_theta_sketch expected = trimmed.compact(false);
+  REQUIRE(result.get_theta64() == expected.get_theta64());
+  REQUIRE(result.get_num_retained() == expected.get_num_retained());
+  REQUIRE(result.get_estimate() == expected.get_estimate());
+
+  // same set of retained hashes (order-independent)
+  std::vector<uint64_t> a(result.begin(), result.end());
+  std::vector<uint64_t> b(expected.begin(), expected.end());
+  std::sort(a.begin(), a.end());
+  std::sort(b.begin(), b.end());
+  REQUIRE(a == b);
+}
+
+TEST_CASE("theta sketch: get_result on empty and below-k sketches", "[theta_sketch]") {
+  compact_theta_sketch empty_result = update_theta_sketch::builder().build().get_result();
+  REQUIRE(empty_result.is_empty());
+  REQUIRE(empty_result.get_num_retained() == 0);
+
+  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.get_result();
+  REQUIRE_FALSE(small_result.is_estimation_mode());
+  REQUIRE(small_result.get_num_retained() == 100); // below k: nothing trimmed
+  REQUIRE(small_result.get_estimate() == Approx(100.0));
+}
+
 TEST_CASE("theta sketch: deserialize compact v1 empty from java", "[theta_sketch]") {
   std::ifstream is;
   is.exceptions(std::ios::failbit | std::ios::badbit);