diff --git a/theta/include/theta_sketch.hpp b/theta/include/theta_sketch.hpp index 4aab4b92..494df894 100644 --- a/theta/include/theta_sketch.hpp +++ b/theta/include/theta_sketch.hpp @@ -336,6 +336,14 @@ class update_theta_sketch_alloc: public theta_sketch_alloc { */ compact_theta_sketch_alloc 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 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 { 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 304ae64c..5ad971bb 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" @@ -243,6 +244,25 @@ compact_theta_sketch_alloc update_theta_sketch_alloc::compact(bool ordered return compact_theta_sketch_alloc(*this, ordered); } +template +compact_theta_sketch_alloc update_theta_sketch_alloc::get_result(bool ordered) const { + std::vector entries(table_.allocator_); + if (is_empty()) { + return compact_theta_sketch_alloc(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(false, ordered, get_seed_hash(), theta, std::move(entries)); +} + template void update_theta_sketch_alloc::print_specifics(std::ostringstream& os) const { os << " lg nominal size : " << static_cast(table_.lg_nom_size_) << std::endl; diff --git a/theta/test/theta_sketch_test.cpp b/theta/test/theta_sketch_test.cpp index 97c4f14e..3077c7b5 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 @@ -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(ordered_result.begin(), ordered_result.end()) + == std::vector(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 unordered_hashes(unordered_result.begin(), unordered_result.end()); + std::sort(unordered_hashes.begin(), unordered_hashes.end()); + REQUIRE(unordered_hashes == std::vector(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);