Skip to content

feat(tdigest): add owned batch construction and quantile APIs - #261

Closed
tisonkun wants to merge 7 commits into
apache:mainfrom
tisonkun:codex/tdigest-batch-merge-validation
Closed

feat(tdigest): add owned batch construction and quantile APIs#261
tisonkun wants to merge 7 commits into
apache:mainfrom
tisonkun:codex/tdigest-batch-merge-validation

Conversation

@tisonkun

@tisonkun tisonkun commented Sep 1, 2026

Copy link
Copy Markdown
Member

Summary

  • implement FromIterator<TDigestMut> so owned partial digests can be combined with one compression pass
  • keep TDigestMut::merge(&TDigestMut) as the single borrowed merge API, while using the smaller k when inputs differ
  • add TDigestMut::quantiles and TDigest::quantiles; nondecreasing ranks share one centroid scan, while arbitrary input order is preserved

Deserialization 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 merge does 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:

  • use merge(&other) when the source must remain available
  • use partials.into_iter().collect::<TDigestMut>() when the partials can be consumed

The 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 k among non-empty inputs. FromIterator is infallible, matching the existing update and merge APIs; 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:

  1. no non-empty inputs returns TDigestMut::default(); one non-empty input is returned unchanged
  2. multiple owned inputs are retained together while the result weight, extrema, centroid count, and minimum k are computed
  3. fully compressed buffers are already sorted runs; a vendored lazy k-way merge heapifies their heads once, then advances the root and performs one sift-down per centroid
  4. that iterator feeds compression directly, so the implementation materializes only the run heap and retained output rather than another vector containing every merged centroid
  5. if any input has an unsorted update tail, its centroids are moved into one vector, stably sorted once, and passed to the existing in-place compressor; putting raw tails before summaries preserves regular update-path tie semantics

Both 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, QuantileCursor owns 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_assert checks 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 borrowed merge one 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:

Workload Existing API New API
merge 64 compressed partials 44.90 us, repeated merge 17.49 us, owned collect
query 6 ranks 1.27 us, repeated quantile 206.2 ns, quantiles

With 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-testdata
  • cargo x check
  • cargo x lint
  • cargo x test
  • cargo bench --package benchmarks --bench benchmarks -- tdigest::merge::partials --sample-count 300
  • cargo bench --package benchmarks --bench benchmarks -- tdigest::merge::uncompressed_partials_from_iter --sample-count 300
  • cargo bench --package benchmarks --bench benchmarks -- tdigest::query::quantiles --sample-count 500

@tisonkun
tisonkun marked this pull request as draft September 1, 2026 23:31
@tisonkun tisonkun changed the title feat(tdigest): add checked batch merge and quantile APIs feat(tdigest): add owned batch construction and quantile APIs Sep 2, 2026
@tisonkun
tisonkun marked this pull request as ready for review September 2, 2026 03:00
@tisonkun
tisonkun marked this pull request as draft September 2, 2026 11:38
@tisonkun

tisonkun commented Sep 2, 2026

Copy link
Copy Markdown
Member Author

... should have a vendored k-way merge and review allocations more carefully.

@tisonkun

tisonkun commented Sep 2, 2026

Copy link
Copy Markdown
Member Author

.. should split into two PRs.

@tisonkun tisonkun closed this Sep 2, 2026
@tisonkun
tisonkun deleted the codex/tdigest-batch-merge-validation branch September 2, 2026 12:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant