perf: kll_sketch allocates its items on first use

kll/include/kll_sketch.hpp

@@ -521,6 +521,7 @@ class kll_sketch {
 
     /**
      * Gets the sorted view of this sketch
+     * An empty sketch returns an empty view without allocating item storage.
      * @return the sorted view of this sketch
      */
     quantiles_sorted_view<T, C, A> get_sorted_view() const;
@@ -558,7 +559,7 @@ class kll_sketch {
     bool is_level_zero_sorted_;
     uint64_t n_;
     vector_u32 levels_;
-    T* items_;
+    T* items_; // nullptr until the sketch first receives items, see allocate_items()
     uint32_t items_size_;
     optional<T> min_item_;
     optional<T> max_item_;
@@ -572,6 +573,13 @@ class kll_sketch {
 
     // common update code
     inline void update_min_max(const T& item);
+    // items_ is allocated lazily, and these are the only places that deal with it being nullptr:
+    // allocates items_ (items_size_ items) unless it is allocated already
+    inline void allocate_items();
+    // copy-constructs the items of other, if other has any, into this newly constructed sketch
+    template<typename O> void copy_items(const O& other);
+    // &items_[index], or nullptr without item storage, which makes every item range empty
+    inline T* item_ptr(uint32_t index) const;
     inline uint32_t internal_update();
 
     // The following code is only valid in the special case of exactly reaching capacity while updating.

kll/include/kll_sketch_impl.hpp

@@ -55,7 +55,8 @@ sorted_view_(nullptr)
         + std::to_string(kll_constants::MAX_K) + ": " + std::to_string(k));
   }
   levels_[0] = levels_[1] = k;
-  items_ = allocator_.allocate(items_size_);
+  // items_ is allocated by the first update or merge: many sketches (e.g. per-group aggregation states)
+  // stay empty or are only ever merged into
 }
 
 template<typename T, typename C, typename A>
@@ -75,8 +76,7 @@ min_item_(other.min_item_),
 max_item_(other.max_item_),
 sorted_view_(nullptr)
 {
-  items_ = allocator_.allocate(items_size_);
-  for (auto i = levels_[0]; i < levels_[num_levels_]; ++i) new (&items_[i]) T(other.items_[i]);
+  copy_items(other);
 }
 
 template<typename T, typename C, typename A>
@@ -171,8 +171,7 @@ sorted_view_(nullptr)
     std::is_constructible<T, TT>::value,
     "Type converting constructor requires new type to be constructible from existing type"
   );
-  items_ = allocator_.allocate(items_size_);
-  for (auto i = levels_[0]; i < levels_[num_levels_]; ++i) new (&items_[i]) T(other.items_[i]);
+  copy_items(other);
   check_sorting();
 }
 
@@ -197,8 +196,27 @@ void kll_sketch<T, C, A>::update_min_max(const T& item) {
   }
 }
 
+template<typename T, typename C, typename A>
+void kll_sketch<T, C, A>::allocate_items() {
+  if (items_ == nullptr) items_ = allocator_.allocate(items_size_);
+}
+
+template<typename T, typename C, typename A>
+template<typename O>
+void kll_sketch<T, C, A>::copy_items(const O& other) {
+  if (other.items_ == nullptr) return;
+  allocate_items();
+  for (auto i = levels_[0]; i < levels_[num_levels_]; ++i) new (&items_[i]) T(other.items_[i]);
+}
+
+template<typename T, typename C, typename A>
+T* kll_sketch<T, C, A>::item_ptr(uint32_t index) const {
+  return items_ == nullptr ? nullptr : items_ + index;
+}
+
 template<typename T, typename C, typename A>
 uint32_t kll_sketch<T, C, A>::internal_update() {
+  allocate_items();
   if (levels_[0] == 0) { compress_while_updating(); }
   n_++;
   is_level_zero_sorted_ = false;
@@ -212,6 +230,7 @@ void kll_sketch<T, C, A>::merge(FwdSk&& other) {
   if (m_ != other.m_) {
     throw std::invalid_argument("incompatible M: " + std::to_string(m_) + " and " + std::to_string(other.m_));
   }
+  allocate_items(); // the higher levels are merged into items_ even if other's level zero is empty
   if (is_empty()) {
     min_item_.emplace(conditional_forward<FwdSk>(*other.min_item_));
     max_item_.emplace(conditional_forward<FwdSk>(*other.max_item_));
@@ -748,7 +767,7 @@ void kll_sketch<T, C, A>::add_empty_top_level_to_completely_full_sketch() {
 template<typename T, typename C, typename A>
 void kll_sketch<T, C, A>::sort_level_zero() {
   if (!is_level_zero_sorted_) {
-    std::sort(items_ + levels_[0], items_ + levels_[1], comparator_);
+    std::sort(item_ptr(levels_[0]), item_ptr(levels_[1]), comparator_);
     is_level_zero_sorted_ = true;
   }
 }
@@ -757,8 +776,8 @@ template<typename T, typename C, typename A>
 void kll_sketch<T, C, A>::check_sorting() const {
   // not checking level 0
   for (uint8_t level = 1; level < num_levels_; ++level) {
-    const auto from = items_ + levels_[level];
-    const auto to = items_ + levels_[level + 1];
+    const auto from = item_ptr(levels_[level]);
+    const auto to = item_ptr(levels_[level + 1]);
     if (!std::is_sorted(from, to, comparator_)) {
       throw std::logic_error("levels must be sorted");
     }
@@ -770,8 +789,8 @@ quantiles_sorted_view<T, C, A> kll_sketch<T, C, A>::get_sorted_view() const {
   const_cast<kll_sketch*>(this)->sort_level_zero(); // allow this side effect
   quantiles_sorted_view<T, C, A> view(get_num_retained(), comparator_, allocator_);
   for (uint8_t level = 0; level < num_levels_; ++level) {
-    const auto from = items_ + levels_[level];
-    const auto to = items_ + levels_[level + 1]; // exclusive
+    const auto from = item_ptr(levels_[level]);
+    const auto to = item_ptr(levels_[level + 1]); // exclusive
     view.add(from, to, 1ULL << level);
   }
   view.convert_to_cummulative();

kll/test/kll_sketch_custom_type_test.cpp

@@ -153,4 +153,109 @@ TEST_CASE("kll sketch custom type", "[kll_sketch]") {
   REQUIRE(test_allocator_net_allocations == 0);
 }
 
+TEST_CASE("kll sketch custom type: items are allocated on first use", "[kll_sketch]") {
+  test_allocator_total_bytes = 0;
+  test_allocator_net_allocations = 0;
+  {
+    kll_test_type_sketch sketch(200, test_type_less(), 0);
+    // only the levels: an empty sketch has no item storage
+    const long long empty_bytes = test_allocator_total_bytes;
+    REQUIRE(empty_bytes < static_cast<long long>(200 * sizeof(test_type)));
+
+    kll_test_type_sketch copy(sketch);
+    kll_test_type_sketch converted(sketch, test_type_less(), alloc(0));
+    REQUIRE(test_allocator_total_bytes == 3 * empty_bytes);
+
+    kll_test_type_sketch target(200, test_type_less(), 0);
+    target.merge(sketch); // merging an empty sketch allocates nothing either
+    REQUIRE(test_allocator_total_bytes == 4 * empty_bytes);
+    REQUIRE(target.is_empty());
+
+    sketch.update(1);
+    REQUIRE(test_allocator_total_bytes == static_cast<long long>(4 * empty_bytes + 200 * sizeof(test_type)));
+    REQUIRE(sketch.get_n() == 1);
+    REQUIRE(sketch.get_min_item().get_value() == 1);
+  }
+  REQUIRE(test_allocator_total_bytes == 0);
+  REQUIRE(test_allocator_net_allocations == 0);
+}
+
+TEST_CASE("kll sketch custom type: merge into a sketch without items", "[kll_sketch]") {
+  test_allocator_total_bytes = 0;
+  test_allocator_net_allocations = 0;
+  {
+    // estimation mode sizes, some of which leave level zero empty after a compaction
+    for (int n = 1; n < 2000; n += 37) {
+      kll_test_type_sketch source(8, test_type_less(), 0);
+      for (int i = 0; i < n; ++i) source.update(i);
+      kll_test_type_sketch target(8, test_type_less(), 0);
+      target.merge(source);
+      REQUIRE(target.get_n() == source.get_n());
+      REQUIRE(target.get_min_item().get_value() == 0);
+      REQUIRE(target.get_max_item().get_value() == n - 1);
+      kll_test_type_sketch moved_target(8, test_type_less(), 0);
+      moved_target.merge(std::move(source));
+      REQUIRE(moved_target.get_n() == target.get_n());
+    }
+  }
+  REQUIRE(test_allocator_total_bytes == 0);
+  REQUIRE(test_allocator_net_allocations == 0);
+}
+
+TEST_CASE("kll sketch custom type: empty sorted view", "[kll_sketch]") {
+  test_allocator_total_bytes = 0;
+  test_allocator_net_allocations = 0;
+  {
+    kll_test_type_sketch empty(200, test_type_less(), 0);
+    kll_test_type_sketch copy(empty);
+    kll_test_type_sketch converted(empty, test_type_less(), alloc(0));
+    kll_test_type_sketch moved(std::move(copy));
+
+    kll_test_type_sketch assigned(200, test_type_less(), 0);
+    assigned.update(99);
+    REQUIRE(assigned.get_quantile(0).get_value() == 99); // populate the cached view
+    assigned = empty;
+
+    kll_test_type_sketch move_assigned(200, test_type_less(), 0);
+    move_assigned.update(99);
+    REQUIRE(move_assigned.get_quantile(0).get_value() == 99);
+    move_assigned = kll_test_type_sketch(empty);
+
+    const auto bytes = empty.serialize(0, test_type_serde());
+    auto from_bytes = kll_test_type_sketch::deserialize(bytes.data(), bytes.size(),
+        test_type_serde(), test_type_less(), alloc(0));
+    std::stringstream stream;
+    empty.serialize(stream, test_type_serde());
+    auto from_stream = kll_test_type_sketch::deserialize(stream,
+        test_type_serde(), test_type_less(), alloc(0));
+
+    for (auto* sketch: {&empty, &converted, &moved, &assigned, &move_assigned, &from_bytes, &from_stream}) {
+      REQUIRE(sketch->is_empty());
+      REQUIRE(sketch->begin() == sketch->end());
+      const auto bytes_before = test_allocator_total_bytes;
+      const auto allocations_before = test_allocator_net_allocations;
+      for (int i = 0; i < 2; ++i) {
+        const auto view = sketch->get_sorted_view();
+        REQUIRE(view.size() == 0);
+        REQUIRE(view.begin() == view.end());
+        REQUIRE_THROWS_AS(view.get_quantile(0.5), std::runtime_error);
+        REQUIRE(test_allocator_total_bytes == bytes_before);
+        REQUIRE(test_allocator_net_allocations == allocations_before);
+      }
+
+      // Reading an empty view must not interfere with allocation or sorting on the first updates.
+      sketch->update(2);
+      sketch->update(1);
+      REQUIRE(sketch->get_n() == 2);
+      const auto view = sketch->get_sorted_view();
+      REQUIRE(view.size() == 2);
+      REQUIRE(view.get_quantile(0).get_value() == 1);
+      REQUIRE(view.get_quantile(1).get_value() == 2);
+      REQUIRE(sketch->get_quantile(0).get_value() == 1);
+    }
+  }
+  REQUIRE(test_allocator_total_bytes == 0);
+  REQUIRE(test_allocator_net_allocations == 0);
+}
+
 } /* namespace datasketches */

kll/test/kll_sketch_test.cpp

@@ -75,6 +75,9 @@ TEST_CASE("kll sketch", "[kll_sketch]") {
       unused(pair); // to suppress "unused" warning
       FAIL("should be no iterations over an empty sketch");
     }
+    const auto empty_bytes = test_allocator_total_bytes;
+    REQUIRE(sketch.get_sorted_view().size() == 0);
+    REQUIRE(test_allocator_total_bytes == empty_bytes);
   }
 
   SECTION("get bad quantile") {
@@ -115,6 +118,7 @@ TEST_CASE("kll sketch", "[kll_sketch]") {
     kll_float_sketch sketch(200, std::less<float>(), 0);
     sketch.update(std::numeric_limits<float>::quiet_NaN());
     REQUIRE(sketch.is_empty());
+    REQUIRE(sketch.get_sorted_view().size() == 0);
 
     sketch.update(0);
     sketch.update(std::numeric_limits<float>::quiet_NaN());