perf: kll_sketch grows its items buffer by 8x up to k

kll/include/kll_sketch.hpp

@@ -183,7 +183,9 @@ class kll_sketch {
     using quantile_return_type = typename quantiles_sorted_view<T, C, A>::quantile_return_type;
 
     /**
-     * Constructor
+     * Constructor.
+     * The items buffer starts with space for min(k, DEFAULT_K) items (k items if RESIZE_FACTOR is X1),
+     * and grows by RESIZE_FACTOR up to k items before the first compaction.
      * @param k affects the size of the sketch and its estimation error
      * @param comparator strict weak ordering function (see C++ named requirements: Compare)
      * @param allocator used by this sketch to allocate memory
@@ -549,6 +551,10 @@ class kll_sketch {
     static const uint8_t PREAMBLE_INTS_SHORT = 2; // for empty and single item
     static const uint8_t PREAMBLE_INTS_FULL = 5;
 
+    // growth factor of the items buffer until it holds k items
+    static const resize_factor RESIZE_FACTOR = resize_factor::X8;
+    static_assert(RESIZE_FACTOR >= resize_factor::X1, "RESIZE_FACTOR must be at least X1");
+
     C comparator_;
     A allocator_;
     uint16_t k_;
@@ -580,6 +586,7 @@ class kll_sketch {
 
     uint8_t find_level_to_compact() const;
     void add_empty_top_level_to_completely_full_sketch();
+    void grow_items();
     void sort_level_zero();
 
     template<typename O> void merge_higher_levels(O&& other, uint64_t final_n);

kll/include/kll_sketch_impl.hpp

@@ -45,7 +45,7 @@ is_level_zero_sorted_(false),
 n_(0),
 levels_(2, 0, allocator),
 items_(nullptr),
-items_size_(k_),
+items_size_(RESIZE_FACTOR == resize_factor::X1 ? k_ : std::min(k_, kll_constants::DEFAULT_K)),
 min_item_(),
 max_item_(),
 sorted_view_(nullptr)
@@ -54,7 +54,7 @@ sorted_view_(nullptr)
     throw std::invalid_argument("K must be >= " + std::to_string(kll_constants::MIN_K) + " and <= "
         + std::to_string(kll_constants::MAX_K) + ": " + std::to_string(k));
   }
-  levels_[0] = levels_[1] = k;
+  levels_[0] = levels_[1] = items_size_;
   items_ = allocator_.allocate(items_size_);
 }
 
@@ -199,7 +199,13 @@ void kll_sketch<T, C, A>::update_min_max(const T& item) {
 
 template<typename T, typename C, typename A>
 uint32_t kll_sketch<T, C, A>::internal_update() {
-  if (levels_[0] == 0) { compress_while_updating(); }
+  if (levels_[0] == 0) {
+    if (items_size_ < k_) {
+      grow_items();
+    } else {
+      compress_while_updating();
+    }
+  }
   n_++;
   is_level_zero_sorted_ = false;
   return --levels_[0];
@@ -389,7 +395,11 @@ void kll_sketch<T, C, A>::serialize(std::ostream& os, const SerDe& sd) const {
     write(os, min_k_);
     write(os, num_levels_);
     write(os, unused);
-    write(os, levels_.data(), sizeof(levels_[0]) * num_levels_);
+    if (items_size_ < k_) {
+      write(os, levels_[0] + (k_ - items_size_));
+    } else {
+      write(os, levels_.data(), sizeof(levels_[0]) * num_levels_);
+    }
     sd.serialize(os, &*min_item_, 1);
     sd.serialize(os, &*max_item_, 1);
   }
@@ -425,7 +435,11 @@ auto kll_sketch<T, C, A>::serialize(unsigned header_size_bytes, const SerDe& sd)
       ptr += copy_to_mem(min_k_, ptr);
       ptr += copy_to_mem(num_levels_, ptr);
       ptr += sizeof(uint8_t); // unused
-      ptr += copy_to_mem(levels_.data(), ptr, sizeof(levels_[0]) * num_levels_);
+      if (items_size_ < k_) {
+        ptr += copy_to_mem(levels_[0] + (k_ - items_size_), ptr);
+      } else {
+        ptr += copy_to_mem(levels_.data(), ptr, sizeof(levels_[0]) * num_levels_);
+      }
       ptr += sd.serialize(ptr, end_ptr - ptr, &*min_item_, 1);
       ptr += sd.serialize(ptr, end_ptr - ptr, &*max_item_, 1);
     }
@@ -753,6 +767,21 @@ void kll_sketch<T, C, A>::add_empty_top_level_to_completely_full_sketch() {
   levels_[num_levels_] = new_total_cap; // initialize the new "extra" index at the top
 }
 
+// Grows the items buffer of a one-level sketch when it is full and smaller than k.
+// Compaction starts only when the buffer is full at k items, so growing does not change the result.
+template<typename T, typename C, typename A>
+void kll_sketch<T, C, A>::grow_items() {
+  const uint32_t new_size = std::min<uint32_t>(items_size_ << static_cast<uint8_t>(RESIZE_FACTOR), k_);
+  const uint32_t delta = new_size - items_size_;
+  T* new_items = allocator_.allocate(new_size);
+  kll_helper::move_construct<T>(items_, levels_[0], levels_[1], new_items, levels_[0] + delta, true);
+  allocator_.deallocate(items_, items_size_);
+  items_ = new_items;
+  items_size_ = new_size;
+  levels_[0] += delta;
+  levels_[1] += delta;
+}
+
 template<typename T, typename C, typename A>
 void kll_sketch<T, C, A>::sort_level_zero() {
   if (!is_level_zero_sorted_) {
@@ -960,7 +989,8 @@ string<A> kll_sketch<T, C, A>::to_string(bool print_levels, bool print_items) co
   os << "   Estimation mode: " << (is_estimation_mode() ? "true" : "false") << std::endl;
   os << "   Levels         : " << (unsigned int) num_levels_ << std::endl;
   os << "   Sorted         : " << (is_level_zero_sorted_ ? "true" : "false") << std::endl;
-  os << "   Capacity items : " << items_size_ << std::endl;
+  // a buffer still growing to k items shows the capacity of a full one
+  os << "   Capacity items : " << std::max<uint32_t>(items_size_, k_) << std::endl;
   os << "   Retained items : " << get_num_retained() << std::endl;
   if (!is_empty()) {
     os << "   Min item      : " << *min_item_ << std::endl;

kll/test/kll_sketch_test.cpp

@@ -18,11 +18,13 @@
  */
 
 #include <catch2/catch.hpp>
+#include <algorithm>
 #include <cmath>
 #include <cstring>
 #include <sstream>
 #include <fstream>
 #include <stdexcept>
+#include <vector>
 
 #include <kll_sketch.hpp>
 #include <test_allocator.hpp>
@@ -905,4 +907,215 @@ TEST_CASE("kll sketch: deserialize rejects inconsistent levels", "[kll_sketch]")
   }
 }
 
+// A deserialized sketch has a buffer of the full size.
+template<typename Sketch>
+static Sketch with_full_buffer(const Sketch& sketch) {
+  auto bytes = sketch.serialize();
+  return Sketch::deserialize(bytes.data(), bytes.size(), serde<typename Sketch::value_type>(),
+      sketch.get_comparator(), sketch.get_allocator());
+}
+
+// Applies fn to both sketches with the same random bits, so that they compact the same way.
+template<typename Sketch, typename Fn>
+static void apply_to_both(Sketch& growing, Sketch& full, const Fn& fn) {
+  const auto random_bit = random_utils::random_bit;
+  fn(growing);
+  random_utils::random_bit = random_bit;
+  fn(full);
+}
+
+template<typename Sketch>
+static std::vector<std::pair<typename Sketch::value_type, uint64_t>> items_and_weights(const Sketch& sketch) {
+  std::vector<std::pair<typename Sketch::value_type, uint64_t>> result;
+  for (const auto pair: sketch) { result.emplace_back(pair.first, pair.second); }
+  return result;
+}
+
+// Requires the same retained items and weights, the same query results and the same bytes.
+template<typename Sketch>
+static void require_identical(const Sketch& growing, const Sketch& full) {
+  REQUIRE(growing.get_n() == full.get_n());
+  REQUIRE(growing.get_num_retained() == full.get_num_retained());
+  REQUIRE(growing.is_estimation_mode() == full.is_estimation_mode());
+  REQUIRE(growing.get_normalized_rank_error(false) == full.get_normalized_rank_error(false));
+  REQUIRE(growing.get_normalized_rank_error(true) == full.get_normalized_rank_error(true));
+  REQUIRE(items_and_weights(growing) == items_and_weights(full));
+  if (!growing.is_empty()) {
+    REQUIRE(growing.get_min_item() == full.get_min_item());
+    REQUIRE(growing.get_max_item() == full.get_max_item());
+    std::vector<typename Sketch::value_type> split_points;
+    for (int i = 0; i <= 100; ++i) {
+      const double rank = i / 100.0;
+      for (const bool inclusive: {false, true}) {
+        const auto quantile = growing.get_quantile(rank, inclusive);
+        REQUIRE(quantile == full.get_quantile(rank, inclusive));
+        REQUIRE(growing.get_rank(quantile, inclusive) == full.get_rank(quantile, inclusive));
+      }
+      split_points.push_back(growing.get_quantile(rank));
+    }
+    split_points.erase(std::unique(split_points.begin(), split_points.end()), split_points.end());
+    for (const bool inclusive: {false, true}) {
+      REQUIRE(growing.get_PMF(split_points.data(), split_points.size(), inclusive)
+          == full.get_PMF(split_points.data(), split_points.size(), inclusive));
+      REQUIRE(growing.get_CDF(split_points.data(), split_points.size(), inclusive)
+          == full.get_CDF(split_points.data(), split_points.size(), inclusive));
+    }
+  }
+  REQUIRE(growing.serialize() == full.serialize());
+  std::stringstream growing_stream, full_stream;
+  growing.serialize(growing_stream);
+  full.serialize(full_stream);
+  REQUIRE(growing_stream.str() == full_stream.str());
+  REQUIRE(growing.to_string(true, false) == full.to_string(true, false));
+}
+
+// Sizes of the items buffer as it grows by 8x up to k.
+static std::vector<uint32_t> buffer_sizes(uint16_t k) {
+  std::vector<uint32_t> sizes = {std::min(k, kll_constants::DEFAULT_K)};
+  while (sizes.back() < k) {
+    sizes.push_back(std::min<uint32_t>(sizes.back() * 8, k));
+  }
+  return sizes;
+}
+
+// Updates a sketch and a copy of it with a full buffer the same way, and requires the same results.
+template<typename Sketch>
+static void require_same_as_full_buffer(Sketch& growing, uint32_t num_items) {
+  auto full = with_full_buffer(growing);
+  random_utils::random_bit.seed(4);
+  for (uint32_t i = 0; i < num_items; ++i) {
+    const auto item = static_cast<typename Sketch::value_type>(i % 7777);
+    apply_to_both(growing, full, [&item](Sketch& sketch) { sketch.update(item); });
+    if (i == 1000 || i == 20000) { require_identical(growing, full); }
+  }
+  require_identical(growing, full);
+}
+
+TEST_CASE("kll sketch: lazy growth of the items buffer", "[kll_sketch]") {
+  test_allocator_total_bytes = 0;
+  const long long levels_bytes = 2 * sizeof(uint32_t);
+
+  SECTION("k below DEFAULT_K allocates k items") {
+    const uint16_t k = kll_constants::DEFAULT_K / 2;
+    kll_float_sketch sketch(k, std::less<float>(), 0);
+    sketch.update(0);
+    REQUIRE(test_allocator_total_bytes == static_cast<long long>(levels_bytes + k * sizeof(float)));
+    for (int i = 1; i < k; ++i) { sketch.update(static_cast<float>(i)); }
+    REQUIRE(test_allocator_total_bytes == static_cast<long long>(levels_bytes + k * sizeof(float)));
+  }
+
+  SECTION("buffer grows by 8x up to k") {
+    REQUIRE(buffer_sizes(kll_constants::MAX_K) == std::vector<uint32_t>{200, 1600, 12800, 65535});
+    kll_float_sketch sketch(kll_constants::MAX_K, std::less<float>(), 0);
+    REQUIRE(test_allocator_total_bytes
+        == static_cast<long long>(levels_bytes + kll_constants::DEFAULT_K * sizeof(float)));
+    uint32_t n = 0;
+    for (const uint32_t size: buffer_sizes(kll_constants::MAX_K)) {
+      while (n < size) {
+        sketch.update(static_cast<float>(n++));
+        // the buffer grows when an item does not fit
+        REQUIRE(test_allocator_total_bytes == static_cast<long long>(levels_bytes + size * sizeof(float)));
+      }
+    }
+    REQUIRE_FALSE(sketch.is_estimation_mode());
+    sketch.update(0);
+    REQUIRE(sketch.is_estimation_mode());
+  }
+
+  SECTION("merge grows the buffer like updates") {
+    kll_float_sketch other(kll_constants::MAX_K, std::less<float>(), 0);
+    for (int i = 0; i < 1000; ++i) { other.update(static_cast<float>(i)); }
+    const long long other_bytes = test_allocator_total_bytes;
+    kll_float_sketch sketch(kll_constants::MAX_K, std::less<float>(), 0);
+    sketch.merge(other);
+    REQUIRE(test_allocator_total_bytes - other_bytes == static_cast<long long>(levels_bytes + 1600 * sizeof(float)));
+    REQUIRE(sketch.get_n() == 1000);
+  }
+
+  SECTION("same as a full buffer through compactions") {
+    // the buffer grows 1, 2 and 3 times before it reaches these k
+    const uint16_t ks[] = {201, 5000, 20000};
+    for (const uint16_t k: ks) {
+      random_utils::random_bit.seed(k);
+      kll_float_sketch growing(k, std::less<float>(), 0);
+      growing.update(0);
+      const long long growing_bytes = test_allocator_total_bytes;
+      kll_float_sketch full = with_full_buffer(growing);
+      REQUIRE(test_allocator_total_bytes - growing_bytes == static_cast<long long>(levels_bytes + k * sizeof(float)));
+      const auto sizes = buffer_sizes(k);
+      for (uint32_t n = 2; n <= 10u * k; ++n) {
+        const float item = static_cast<float>((n * 7919) % 10007) / 3;
+        apply_to_both(growing, full, [item](kll_float_sketch& sketch) { sketch.update(item); });
+        const bool at_growth = std::find(sizes.begin(), sizes.end(), n) != sizes.end()
+            || std::find(sizes.begin(), sizes.end(), n - 1) != sizes.end();
+        if (at_growth || n % 997 == 0 || n % k <= 1) { require_identical(growing, full); }
+      }
+      REQUIRE(growing.is_estimation_mode());
+      require_identical(growing, full);
+    }
+  }
+
+  SECTION("same as a full buffer when merging") {
+    random_utils::random_bit.seed(3);
+    // exact mode sketches first, so that merging grows the buffer before compactions start
+    const std::vector<std::pair<uint16_t, uint32_t>> k_and_n = {
+        {30000, 1000}, {30000, 3000}, {20000, 15000}, {2000, 10000}, {8, 100}, {20000, 60000}};
+    std::vector<kll_float_sketch> others;
+    for (const auto& shape: k_and_n) {
+      others.emplace_back(shape.first, std::less<float>(), 0);
+      for (uint32_t i = 0; i < shape.second; ++i) { others.back().update(static_cast<float>((i * 31) % 4999)); }
+    }
+    for (const bool reversed: {false, true}) {
+      kll_float_sketch growing(20000, std::less<float>(), 0);
+      growing.update(0);
+      kll_float_sketch full = with_full_buffer(growing);
+      for (size_t i = 0; i < others.size(); ++i) {
+        const auto& other = others[reversed ? others.size() - 1 - i : i];
+        apply_to_both(growing, full, [&other](kll_float_sketch& sketch) { sketch.merge(other); });
+        require_identical(growing, full);
+      }
+      for (int i = 0; i < 50000; ++i) {
+        apply_to_both(growing, full, [i](kll_float_sketch& sketch) { sketch.update(static_cast<float>(i)); });
+      }
+      require_identical(growing, full);
+    }
+  }
+
+  SECTION("copies, moves and converted sketches keep growing the same") {
+    kll_float_sketch sketch(kll_constants::MAX_K, std::less<float>(), 0);
+    for (int i = 0; i < 500; ++i) { sketch.update(static_cast<float>(i)); }
+    kll_float_sketch copied(sketch);
+    kll_float_sketch assigned(8, std::less<float>(), 0);
+    assigned = sketch;
+    kll_float_sketch to_move(sketch);
+    kll_float_sketch moved(std::move(to_move));
+    kll_float_sketch move_assigned(8, std::less<float>(), 0);
+    kll_float_sketch to_move_assign(sketch);
+    move_assigned = std::move(to_move_assign);
+    for (auto* growing: {&sketch, &copied, &assigned, &moved, &move_assigned}) {
+      require_same_as_full_buffer(*growing, 200000);
+    }
+    kll_float_sketch source(kll_constants::MAX_K, std::less<float>(), 0);
+    for (int i = 0; i < 500; ++i) { source.update(static_cast<float>(i)); }
+    kll_sketch<double, std::less<double>, test_allocator<double>> converted(source, std::less<double>(), 0);
+    require_same_as_full_buffer(converted, 200000);
+  }
+
+  SECTION("same as a full buffer with strings") {
+    random_utils::random_bit.seed(5);
+    kll_string_sketch growing(5000, std::less<std::string>(), 0);
+    growing.update("0");
+    kll_string_sketch full = with_full_buffer(growing);
+    for (int i = 1; i < 30000; ++i) {
+      const std::string item = std::to_string((i * 7919) % 10007);
+      apply_to_both(growing, full, [&item](kll_string_sketch& sketch) { sketch.update(item); });
+      if (i == 199 || i == 200 || i == 1599 || i == 1600 || i == 4999 || i == 10000) { require_identical(growing, full); }
+    }
+    require_identical(growing, full);
+  }
+
+  // cleanup
+  REQUIRE(test_allocator_total_bytes == 0);
+}
+
 } /* namespace datasketches */