perf: tdigest grows its buffers as they fill

tdigest/include/tdigest.hpp

@@ -299,6 +299,8 @@ class tdigest {
   bool is_single_value() const;
   uint8_t get_preamble_longs() const;
   void merge(vector_centroid& buffer, W weight);
+  // grows centroids_ to hold num centroids, up to centroids_capacity_, before a merge needs them
+  void reserve_centroids(size_t num);
 
   // for deserialize
   tdigest(bool reverse_merge, uint16_t k, T min, T max, vector_centroid&& centroids, uint64_t total_weight_, vector_t&& buffer);

tdigest/include/tdigest_impl.hpp

@@ -62,6 +62,9 @@ void tdigest<T, A>::update(T value) {
   if (std::isnan(value)) return;
   if (std::isinf(value)) return;
   if (buffer_.size() == centroids_capacity_ * BUFFER_MULTIPLIER) compress();
+  if (buffer_.size() == buffer_.capacity()) {
+    buffer_.reserve(std::min(std::max<size_t>(2 * buffer_.size(), 8), centroids_capacity_ * BUFFER_MULTIPLIER));
+  }
   buffer_.push_back(value);
   min_ = std::min(min_, value);
   max_ = std::max(max_, value);
@@ -73,6 +76,7 @@ void tdigest<T, A>::merge(const tdigest& other) {
   // Extreme centroid means are not the true min/max when those centroids have weight > 1.
   min_ = std::min(min_, other.min_);
   max_ = std::max(max_, other.max_);
+  reserve_centroids(buffer_.size() + centroids_.size() + other.buffer_.size() + other.centroids_.size());
   vector_centroid tmp(buffer_.get_allocator());
   tmp.reserve(buffer_.size() + centroids_.size() + other.buffer_.size() + other.centroids_.size());
   for (const T value: buffer_) tmp.push_back(centroid(value, 1));
@@ -84,6 +88,7 @@ void tdigest<T, A>::merge(const tdigest& other) {
 template<typename T, typename A>
 void tdigest<T, A>::compress() {
   if (buffer_.size() == 0) return;
+  reserve_centroids(buffer_.size() + centroids_.size());
   vector_centroid tmp(buffer_.get_allocator());
   tmp.reserve(buffer_.size() + centroids_.size());
   for (const T value: buffer_) tmp.push_back(centroid(value, 1));
@@ -711,8 +716,16 @@ buffer_(std::move(buffer))
   if (k < 10) throw std::invalid_argument("k must be at least 10");
   const size_t fudge = k < 30 ? 30 : 10;
   centroids_capacity_ = 2 * k_ + fudge;
-  centroids_.reserve(centroids_capacity_);
-  buffer_.reserve(centroids_capacity_ * BUFFER_MULTIPLIER);
+  // centroids_ and buffer_ grow as they fill (see update() and reserve_centroids()) instead of being reserved
+  // here: many t-digests (e.g. per-group aggregation states) hold few values
+}
+
+template<typename T, typename A>
+void tdigest<T, A>::reserve_centroids(size_t num) {
+  num = std::min(num, centroids_capacity_);
+  if (centroids_.capacity() < num) {
+    centroids_.reserve(std::max(num, std::min(2 * centroids_.capacity(), centroids_capacity_)));
+  }
 }
 
 template<typename T, typename A>

tdigest/test/tdigest_custom_allocator_test.cpp

@@ -40,4 +40,33 @@ TEST_CASE("tdigest custom allocator", "[tdigest]") {
   REQUIRE(test_allocator_net_allocations == 0);
 }
 
+TEST_CASE("tdigest custom allocator: buffers grow as they fill", "[tdigest]") {
+  test_allocator_total_bytes = 0;
+  test_allocator_net_allocations = 0;
+  {
+    tdigest_d empty(100, alloc_d(0));
+    REQUIRE(test_allocator_total_bytes == 0); // nothing reserved up front
+
+    tdigest_d small(100, alloc_d(0));
+    for (int i = 0; i < 3; ++i) small.update(static_cast<double>(i));
+    // a few values take a small buffer, not room for 4 * (2k + fudge) values
+    REQUIRE(test_allocator_total_bytes <= static_cast<long long>(8 * sizeof(double)));
+    REQUIRE(small.get_total_weight() == 3);
+    REQUIRE(small.get_quantile(0.5) == 1.0);
+
+    // a digest built through the growing buffers gives the same result as before
+    tdigest_d grown(100, alloc_d(0));
+    tdigest<double> reference(100);
+    for (int i = 0; i < 100000; ++i) {
+      grown.update(static_cast<double>(i % 7919));
+      reference.update(static_cast<double>(i % 7919));
+    }
+    for (double rank: {0.01, 0.25, 0.5, 0.75, 0.99}) {
+      REQUIRE(grown.get_quantile(rank) == reference.get_quantile(rank));
+    }
+  }
+  REQUIRE(test_allocator_total_bytes == 0);
+  REQUIRE(test_allocator_net_allocations == 0);
+}
+
 } /* namespace datasketches */