From 15024ac6f86ce879cd44ec70b7a60f73fd38b35f Mon Sep 17 00:00:00 2001
From: Rusty Conover <rusty@conover.me>
Date: Wed, 28 Oct 2026 14:06:52 -0400
Subject: [PATCH] perf: tdigest grows its buffers as they fill

The constructor reserved room for 2k + fudge centroids and four times as many
buffered values (about 10 KB for k = 100 with doubles), also when
deserializing. Code that holds many small t-digests paid for all of it.
Grow the buffer geometrically up to its limit on update, and the centroids
before a compress or merge needs them. Results are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tdigest/include/tdigest.hpp                   |  2 ++
 tdigest/include/tdigest_impl.hpp              | 17 +++++++++--
 .../test/tdigest_custom_allocator_test.cpp    | 29 +++++++++++++++++++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/tdigest/include/tdigest.hpp b/tdigest/include/tdigest.hpp
index 095752e9..1603e5db 100644
--- a/tdigest/include/tdigest.hpp
+++ b/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);
diff --git a/tdigest/include/tdigest_impl.hpp b/tdigest/include/tdigest_impl.hpp
index 958c8c40..6addf32c 100644
--- a/tdigest/include/tdigest_impl.hpp
+++ b/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>
diff --git a/tdigest/test/tdigest_custom_allocator_test.cpp b/tdigest/test/tdigest_custom_allocator_test.cpp
index b0687d68..06c5e53b 100644
--- a/tdigest/test/tdigest_custom_allocator_test.cpp
+++ b/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 */