diff --git a/hll/include/HllUnion-internal.hpp b/hll/include/HllUnion-internal.hpp
index d83340de..d102b792 100644
--- a/hll/include/HllUnion-internal.hpp
+++ b/hll/include/HllUnion-internal.hpp
@@ -54,7 +54,9 @@ void hll_union_alloc::update(const hll_sketch_alloc& sketch) {
template
void hll_union_alloc::update(hll_sketch_alloc&& sketch) {
if (sketch.is_empty()) { return; }
- if (gadget_.is_empty() && sketch.get_target_type() == HLL_8 && sketch.get_lg_config_k() <= lg_max_k_) {
+ // taking over the sketch is only possible if its memory can be released with the union's allocator
+ if (gadget_.is_empty() && sketch.get_target_type() == HLL_8 && sketch.get_lg_config_k() <= lg_max_k_
+ && sketch.sketch_impl->getAllocator() == gadget_.sketch_impl->getAllocator()) {
if (sketch.get_current_mode() == HLL || sketch.get_lg_config_k() == lg_max_k_) {
gadget_ = std::move(sketch);
}
@@ -208,17 +210,20 @@ double hll_union_alloc::get_rel_err(bool upper_bound, bool unioned,
}
template
-HllSketchImpl* hll_union_alloc::copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k) {
+HllSketchImpl* hll_union_alloc::copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k,
+ const A& allocator) {
if (src_impl->getCurMode() != HLL) {
throw std::logic_error("Attempt to downsample non-HLL sketch");
}
const HllArray* src = static_cast*>(src_impl);
- const int src_lg_k = src->getLgConfigK();
- if (src_lg_k <= tgt_lg_k) {
+ const uint8_t src_lg_k = src->getLgConfigK();
+ // a copy keeps the source's allocator, so it is only made if that is the union's allocator
+ if (src_lg_k <= tgt_lg_k && src->getAllocator() == allocator) {
return src->copyAs(HLL_8);
}
typedef typename std::allocator_traits::template rebind_alloc> hll8Alloc;
- Hll8Array* tgtHllArr = new (hll8Alloc(src->getAllocator()).allocate(1)) Hll8Array(tgt_lg_k, src->getAllocator());
+ const uint8_t lg_k = std::min(src_lg_k, tgt_lg_k);
+ Hll8Array* tgtHllArr = new (hll8Alloc(allocator).allocate(1)) Hll8Array(lg_k, allocator);
tgtHllArr->mergeHll(*src);
//both of these are required for isomorphism
tgtHllArr->putHipAccum(src->getHipAccum());
@@ -239,8 +244,12 @@ template
void hll_union_alloc::union_impl(const hll_sketch_alloc& sketch, uint8_t lg_max_k) {
const HllSketchImpl* src_impl = sketch.sketch_impl; //default
HllSketchImpl* dst_impl = gadget_.sketch_impl; //default
+ // the gadget allocates with the union's allocator: the input sketch may use another allocator
+ // instance (e.g. a stateful one), and copies of it would keep that allocator
+ const A allocator = dst_impl->getAllocator();
if (src_impl->getCurMode() == LIST || src_impl->getCurMode() == SET) {
- if (dst_impl->isEmpty() && src_impl->getLgConfigK() == dst_impl->getLgConfigK()) {
+ if (dst_impl->isEmpty() && src_impl->getLgConfigK() == dst_impl->getLgConfigK()
+ && src_impl->getAllocator() == allocator) {
dst_impl = src_impl->copyAs(HLL_8);
gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
} else {
@@ -254,12 +263,12 @@ void hll_union_alloc::union_impl(const hll_sketch_alloc& sketch, uint8_t l
// swap so that src is LIST or SET, tgt is HLL
// use lg_max_k because LIST has effective K of 2^26
const CouponList* src = static_cast*>(dst_impl);
- dst_impl = copy_or_downsample(src_impl, lg_max_k);
+ dst_impl = copy_or_downsample(src_impl, lg_max_k, allocator);
static_cast*>(dst_impl)->mergeList(*src);
gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
} else { // gadget is HLL
if (src_impl->getLgConfigK() < dst_impl->getLgConfigK()) {
- dst_impl = copy_or_downsample(dst_impl, sketch.get_lg_config_k());
+ dst_impl = copy_or_downsample(dst_impl, sketch.get_lg_config_k(), allocator);
gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
}
const HllArray* src = static_cast*>(src_impl);
@@ -268,7 +277,7 @@ void hll_union_alloc::union_impl(const hll_sketch_alloc& sketch, uint8_t l
static_cast*>(dst_impl)->putHipAccum(0);
}
} else { // src is HLL, gadget is empty
- dst_impl = copy_or_downsample(src_impl, lg_max_k);
+ dst_impl = copy_or_downsample(src_impl, lg_max_k, allocator);
gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
}
gadget_.sketch_impl = dst_impl; // gadget replaced
diff --git a/hll/include/hll.hpp b/hll/include/hll.hpp
index 9950eec0..f160fb14 100644
--- a/hll/include/hll.hpp
+++ b/hll/include/hll.hpp
@@ -640,7 +640,7 @@ class hll_union_alloc {
*/
inline void union_impl(const hll_sketch_alloc& sketch, uint8_t lg_max_k);
- static HllSketchImpl* copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k);
+ static HllSketchImpl* copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k, const A& allocator);
void coupon_update(uint32_t coupon);
diff --git a/hll/test/HllUnionTest.cpp b/hll/test/HllUnionTest.cpp
index 1d065d19..1dba6f62 100644
--- a/hll/test/HllUnionTest.cpp
+++ b/hll/test/HllUnionTest.cpp
@@ -384,4 +384,77 @@ TEST_CASE("hll union: serialization is grouping independent", "[hll_union]") {
REQUIRE(serialize_nested_union(a, b, c) == serialize_flat_union(a, b, c));
}
+// A stateful allocator: instances with different ids are different allocators, and the bytes each one
+// has outstanding are counted.
+static long long pool_bytes[2] = {0, 0};
+
+template
+struct pool_allocator {
+ using value_type = T;
+ int id;
+ explicit pool_allocator(int id): id(id) {}
+ template
+ pool_allocator(const pool_allocator& other): id(other.id) {}
+ T* allocate(std::size_t n) {
+ pool_bytes[id] += n * sizeof(T);
+ return static_cast(::operator new(n * sizeof(T)));
+ }
+ void deallocate(T* p, std::size_t n) {
+ pool_bytes[id] -= n * sizeof(T);
+ ::operator delete(p);
+ }
+};
+template
+bool operator==(const pool_allocator& a, const pool_allocator& b) { return a.id == b.id; }
+template
+bool operator!=(const pool_allocator& a, const pool_allocator& b) { return a.id != b.id; }
+
+using pool_hll_sketch = hll_sketch_alloc>;
+using pool_hll_union = hll_union_alloc>;
+
+static pool_hll_sketch make_pool_sketch(int id, uint8_t lg_k, target_hll_type type, uint64_t n) {
+ pool_hll_sketch sketch(lg_k, type, false, pool_allocator(id));
+ for (uint64_t i = 0; i < n; ++i) sketch.update(i);
+ return sketch;
+}
+
+TEST_CASE("hll union: keeps no memory of an input sketch's allocator", "[hll_union]") {
+ const uint8_t lg_max_k = 10;
+ // LIST, SET and HLL inputs, with a smaller, the same and a larger lg_k than the union
+ const uint64_t sizes[] = {5, 100, 10000};
+ const uint8_t lg_ks[] = {8, 10, 12};
+ for (uint64_t n: sizes) {
+ for (uint8_t lg_k: lg_ks) {
+ for (target_hll_type type: {HLL_4, HLL_6, HLL_8}) {
+ for (bool gadget_empty: {true, false}) {
+ for (bool move: {false, true}) {
+ hll_union u(lg_max_k);
+ pool_hll_union pool_union(lg_max_k, pool_allocator(0));
+ if (!gadget_empty) {
+ for (uint64_t i = 0; i < 3 * n; ++i) {
+ u.update(i + n);
+ pool_union.update(i + n);
+ }
+ }
+ {
+ hll_sketch control(lg_k, type);
+ for (uint64_t i = 0; i < n; ++i) control.update(i);
+ u.update(control);
+ pool_hll_sketch input = make_pool_sketch(1, lg_k, type, n);
+ if (move) {
+ pool_union.update(std::move(input));
+ } else {
+ pool_union.update(input);
+ }
+ }
+ REQUIRE(pool_bytes[1] == 0);
+ REQUIRE(pool_union.get_estimate() == u.get_estimate());
+ }
+ }
+ }
+ }
+ }
+ REQUIRE(pool_bytes[0] == 0);
+}
+
} /* namespace datasketches */