Skip to content

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

Open
tisonkun wants to merge 6 commits into
apache:mainfrom
tisonkun:codex/tdigest-batch-merge-validation
Open

feat(tdigest): add owned batch construction and quantile APIs#261
tisonkun wants to merge 6 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 clone-cost probe on the same machine measured:

Workload Borrowed inputs Clone before merge Difference
one merge 1.956 us 2.038 us +4.2% and one 4.288 KB allocation
64-partial batch engine 21.52 us 22.53 us +4.7%, with allocation count increasing from 3 to 68

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

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 separates cardinality and storage cases:

  1. no non-empty inputs returns TDigestMut::default(); one non-empty input is returned unchanged
  2. multiple inputs first compute the resulting weight, extrema, centroid count, and minimum k without mutating the result
  3. fully compressed inputs use a stable heap-based k-way merge of their already-sorted centroid buffers; if any input has an unmerged tail, owned values are moved into stable tie order and sorted once
  4. the combined centroids are compressed once

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

One-pass batch construction must retain the input centroids until the global merge/compression step. It therefore trades higher peak live input memory for less recompression. Callers that need bounded additional memory can continue to deserialize and call borrowed merge one state at a time, or merge bounded chunks.

Performance

Representative local Divan medians after merging the current main:

Workload Existing API New API
merge 64 compressed partials 98.91 us, repeated merge 23.34 us, owned collect
query 6 ranks 1.26 us, repeated quantile 174.8 ns, quantiles

The owned benchmark creates fresh inputs with Divan with_inputs, so cloning test fixtures is outside the measured interval. The measured interval does include releasing the 64 consumed source buffers. The owned path performs 2 allocations (131 KB) versus 53 allocations (234.5 KB) for repeated merge.

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::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
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