diff --git a/theta/include/theta_union_base_impl.hpp b/theta/include/theta_union_base_impl.hpp index bc8f4902..76f6f035 100644 --- a/theta/include/theta_union_base_impl.hpp +++ b/theta/include/theta_union_base_impl.hpp @@ -41,7 +41,7 @@ void theta_union_base::update(SS&& sketch) { if (sketch.is_empty()) return; if (sketch.get_seed_hash() != compute_seed_hash(table_.seed_)) throw std::invalid_argument("seed hash mismatch"); table_.is_empty_ = false; - if (sketch.get_theta64() < union_theta_) union_theta_ = sketch.get_theta64(); + union_theta_ = std::min(union_theta_, sketch.get_theta64()); for (auto& entry: sketch) { const uint64_t hash = EK()(entry); if (hash < union_theta_ && hash < table_.theta_) { @@ -55,7 +55,7 @@ void theta_union_base::update(SS&& sketch) { if (sketch.is_ordered()) break; // early stop } } - if (table_.theta_ < union_theta_) union_theta_ = table_.theta_; + union_theta_ = std::min(union_theta_, table_.theta_); } template @@ -65,16 +65,16 @@ CS theta_union_base::get_result(bool ordered) const { entries.reserve(table_.num_entries_); uint64_t theta = std::min(union_theta_, table_.theta_); const uint32_t nominal_num = 1 << table_.lg_nom_size_; - if (union_theta_ >= theta && table_.num_entries_ <= nominal_num) { + if (union_theta_ >= table_.theta_) { std::copy_if(table_.begin(), table_.end(), std::back_inserter(entries), key_not_zero()); } 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(); - } + } + 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(); } if (ordered) std::sort(entries.begin(), entries.end(), comparator()); return CS(table_.is_empty_, ordered, compute_seed_hash(table_.seed_), theta, std::move(entries)); diff --git a/theta/include/theta_update_sketch_base_impl.hpp b/theta/include/theta_update_sketch_base_impl.hpp index 0a73c166..0899d061 100644 --- a/theta/include/theta_update_sketch_base_impl.hpp +++ b/theta/include/theta_update_sketch_base_impl.hpp @@ -188,7 +188,7 @@ auto theta_update_sketch_base::begin() const -> iterator { template auto theta_update_sketch_base::end() const -> iterator { - return &entries_[1ULL << lg_cur_size_]; + return entries_ + (1ULL << lg_cur_size_); } template diff --git a/theta/test/theta_union_test.cpp b/theta/test/theta_union_test.cpp index 8135efc1..aef1c312 100644 --- a/theta/test/theta_union_test.cpp +++ b/theta/test/theta_union_test.cpp @@ -128,4 +128,29 @@ TEST_CASE("theta union: seed mismatch", "[theta_union]") { REQUIRE_THROWS_AS(u.update(sketch), std::invalid_argument); } +TEST_CASE("theta union: larger K", "[theta_union]") { + auto update_sketch1 = datasketches::update_theta_sketch::builder().set_lg_k(14).build(); + for(int i = 0; i < 16384; ++i) update_sketch1.update(i); + + auto update_sketch2 = datasketches::update_theta_sketch::builder().set_lg_k(14).build(); + for(int i = 0; i < 26384; ++i) update_sketch2.update(i); + + auto update_sketch3 = datasketches::update_theta_sketch::builder().set_lg_k(14).build(); + for(int i = 0; i < 86384; ++i) update_sketch3.update(i); + + auto union1 = datasketches::theta_union::builder().set_lg_k(16).build(); + union1.update(update_sketch2); + union1.update(update_sketch1); + union1.update(update_sketch3); + auto result1 = union1.get_result(); + REQUIRE(result1.get_estimate() == update_sketch3.get_estimate()); + + auto union2 = datasketches::theta_union::builder().set_lg_k(16).build(); + union2.update(update_sketch1); + union2.update(update_sketch3); + union2.update(update_sketch2); + auto result2 = union2.get_result(); + REQUIRE(result2.get_estimate() == update_sketch3.get_estimate()); +} + } /* namespace datasketches */