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

theta/include/theta_sketch.hpp

@@ -336,6 +336,14 @@ 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.
+   * @param ordered optional flag to specify if an ordered sketch should be produced
+   * @return compact sketch with at most k retained entries
+   */
+  compact_theta_sketch_alloc<Allocator> get_result(bool ordered = true) const;
+
   virtual iterator begin();
   virtual iterator end();
   virtual const_iterator begin() const;
@@ -518,6 +526,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(bool ordered) 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());
+  }
+  if (ordered) std::sort(entries.begin(), entries.end());
+  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,61 @@ 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
+
+  // default is ordered, matching union/intersection get_result
+  compact_theta_sketch ordered_result = update_sketch.get_result();
+  REQUIRE_FALSE(ordered_result.is_empty());
+  REQUIRE(ordered_result.is_estimation_mode());
+  REQUIRE(ordered_result.get_num_retained() == k); // trimmed to nominal size
+  REQUIRE(ordered_result.is_ordered());
+  REQUIRE(std::is_sorted(ordered_result.begin(), ordered_result.end()));
+
+  // fused get_result(true) must match trim() + compact(true)
+  update_theta_sketch trimmed = update_sketch;
+  trimmed.trim();
+  compact_theta_sketch expected = trimmed.compact(true);
+  REQUIRE(ordered_result.get_theta64() == expected.get_theta64());
+  REQUIRE(ordered_result.get_num_retained() == expected.get_num_retained());
+  REQUIRE(ordered_result.get_estimate() == expected.get_estimate());
+  REQUIRE(std::vector<uint64_t>(ordered_result.begin(), ordered_result.end())
+          == std::vector<uint64_t>(expected.begin(), expected.end()));
+
+  // unordered variant: same trimmed set and theta, no sort
+  compact_theta_sketch unordered_result = update_sketch.get_result(false);
+  REQUIRE_FALSE(unordered_result.is_ordered());
+  REQUIRE(unordered_result.get_num_retained() == k);
+  REQUIRE(unordered_result.get_theta64() == expected.get_theta64());
+  std::vector<uint64_t> unordered_hashes(unordered_result.begin(), unordered_result.end());
+  std::sort(unordered_hashes.begin(), unordered_hashes.end());
+  REQUIRE(unordered_hashes == std::vector<uint64_t>(expected.begin(), expected.end()));
+}
+
+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);
+  REQUIRE(empty_result.is_ordered());
+
+  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(); // default ordered
+  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));
+  REQUIRE(small_result.is_ordered());
+  REQUIRE(std::is_sorted(small_result.begin(), small_result.end()));
+
+  REQUIRE_FALSE(small.get_result(false).is_ordered()); // unordered variant
+}
+
 TEST_CASE("theta sketch: deserialize compact v1 empty from java", "[theta_sketch]") {
   std::ifstream is;
   is.exceptions(std::ios::failbit | std::ios::badbit);