feat(tdigest): add owned batch construction and quantile APIs - #261
Closed
tisonkun wants to merge 7 commits into
Closed
feat(tdigest): add owned batch construction and quantile APIs#261tisonkun wants to merge 7 commits into
tisonkun wants to merge 7 commits into
Conversation
tisonkun
marked this pull request as draft
September 1, 2026 23:31
tisonkun
marked this pull request as ready for review
September 2, 2026 03:00
tisonkun
marked this pull request as draft
September 2, 2026 11:38
Member
Author
|
... should have a vendored k-way merge and review allocations more carefully. |
Member
Author
|
.. should split into two PRs. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
FromIterator<TDigestMut>so owned partial digests can be combined with one compression passTDigestMut::merge(&TDigestMut)as the single borrowed merge API, while using the smallerkwhen inputs differTDigestMut::quantilesandTDigest::quantiles; nondecreasing ranks share one centroid scan, while arbitrary input order is preservedDeserialization invariant validation landed independently in #262 and is now part of the base branch rather than this PR's diff.
Ownership and clone cost
The existing borrowed
mergedoes not clone the complete right-hand digest. It moves the receiver's buffer into the result and copies the right-hand centroids into the new contiguous result buffer. Requiring an owned argument when the caller must retain it would add a whole-buffer allocation and copy before that merge work.A temporary single-merge probe on the same machine measured 1.956 us for borrowed input and 2.038 us when cloning first: +4.2% and one 4.288 KB allocation. Cloning a batch would likewise add one allocation and a complete centroid-buffer copy for every non-empty source digest.
The public APIs therefore cover the two ownership cases without forcing a clone:
merge(&other)when the source must remain availablepartials.into_iter().collect::<TDigestMut>()when the partials can be consumedThe owned benchmark creates fresh inputs with Divan
with_inputs, so cloning test fixtures is outside the measured interval; releasing the consumed source buffers remains inside it.Collection ignores empty inputs, returns a single non-empty input unchanged, and uses the smallest
kamong non-empty inputs.FromIteratoris infallible, matching the existingupdateandmergeAPIs; representation overflow is documented as a panic.Implementation shape
Batch construction has one ownership model and two storage algorithms selected by a concrete sortedness invariant:
TDigestMut::default(); one non-empty input is returned unchangedkare computedBoth compression paths use the same centroid-merging predicate. The unsorted fallback is intentional: independently sorting every raw tail and then streaming all runs reduced temporary memory but made the representative uncompressed workload about 69% slower.
For batch queries,
QuantileCursorowns the monotonic scan state. Already-sorted ranks use it directly. Arbitrary ranks sort indices rather than values, drive the same cursor in rank order, and place answers back in the caller's original order.The separate deserialization validation in #262 establishes that fully compressed buffers are sorted. Pairwise merge now treats that as an internal invariant, retaining
debug_assertchecks without rescanning both buffers in release builds.Memory behavior
For fully compressed inputs, owned collection retains the source digests, an
O(number of digests)run heap, and the final retained centroid buffer. It no longer allocates a full combined-centroid vector.Inputs with unsorted update tails require an
O(total centroids)temporary vector for the stable sort. Callers that cannot retain the input states plus this temporary storage can continue to deserialize and call borrowedmergeone state at a time, or merge bounded chunks.On the 64-partial benchmark, repeated borrowed merge used 8.192 KB of allocation/growth during the measured operation. Owned collection used 35.36 KB across two allocations for fully compressed inputs and 131 KB across two allocations for uncompressed inputs. The owned path uses more peak scratch than repeated merge but avoids recompressing intermediate results.
Performance
Representative local Divan medians after merging the current
main:mergecollectquantilequantilesWith the corrected compressed fixture, the previous batch implementation took about 63.0 us and allocated 68.6 KB in three allocations. The lazy k-way implementation takes 17.5 us and allocates 35.4 KB in two allocations. The uncompressed stable-sort fallback remains about 48.1 us, matching the previous implementation rather than paying the cost of sorting every raw tail separately.
Validation
cargo x prepare-testdatacargo x checkcargo x lintcargo x testcargo bench --package benchmarks --bench benchmarks -- tdigest::merge::partials --sample-count 300cargo bench --package benchmarks --bench benchmarks -- tdigest::merge::uncompressed_partials_from_iter --sample-count 300cargo bench --package benchmarks --bench benchmarks -- tdigest::query::quantiles --sample-count 500