kll: add native Rust KLL quantile sketch#2180
Merged
Merged
Conversation
New module ballista_core::kll. Trivial Vec-backed impl (no compaction yet) that satisfies the base case: for n <= capacity, rank(x) matches an exact linear scan. Test is red-green with u32; the same shape against OwnedRow is next. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Same 10 items, same probe, same expected rank — but the items live as OwnedRow bytes produced by RowConverter, exercising the production item type. No sketch-side changes: KllSketch<T: Ord> composes with OwnedRow via its Ord impl (lex byte compare over the encoded row). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds a third test that stitches NULLs into the input array and probes with a non-null value. Under SortField's default (ASC NULLS FIRST), encoded NULLs sort strictly below any encoded Some(_), so rank(Some(5)) picks up the null count on top of the four values below 5. Sketch unchanged: OwnedRow's Ord impl carries the policy. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Fourth honesty test: two-column key (a UInt32, b UInt32), lex ASC/ASC. Same 6 stream rows produce OwnedRow items; the probe (2, 10) picks up (1,10), (1,20), and (2,5). Excluded: (2,15) same a with larger b, and the (3,*) pair with larger a. Cements that the Vec<Option<dyn Ord>> logical model is fully realized as `T = OwnedRow` — arbitrary column count and per-column SortOptions flow through the encoding, sketch stays generic. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds module-level rustdoc: why this sketch exists here (the scheduler needs multi-column, nullable, mixed-direction sort keys for global partition points, which composes as T = OwnedRow), the compactor-stack layout with per-level weights 2^h, the compaction procedure (sort, fair coin, keep half, promote), the weighted rank formula, and the role of k. Includes a link to the Karnin-Lang-Liberty FOCS 2016 paper. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replaces the single Vec<T> with a Vec<Vec<T>> compactor stack indexed by level. When level h fills to k, its contents are sorted, halved by a fair coin flip (odd or even indices), and the surviving k/2 items are promoted to level h+1 with double weight; the cascade continues upward while any level remains full. Rank sums 2^h · (count-below-x) over levels. Constructor becomes new(k) plus with_seed(k, seed) for deterministic tests, backed by an owned StdRng — the coin flip cannot be non-reproducible. The zero-arg new() and Default are gone; k has no natural default. The four earlier tests move to with_seed(1024, 42), below capacity and thus unchanged in behavior. New test locks two exact invariants that hold regardless of coin flips: for a stream of 1..=32 with k=4, rank(33) == 32 (total-weight conservation across cascades) and rank(1) == 0 (nothing below the min). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The current per-item insert API materializes each Arrow Row<'_> into an OwnedRow on the way in — most of those allocations are later discarded by cascading compactions and are effectively wasted. Documents the tradeoff and sketches the shape of the future optimization (a batch-oriented ingest API that holds borrowed rows in level 0 for the duration of a batch), so the choice is visible before real callers accumulate against the current shape. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Ports the level-capacity formula from Apache DataSketches (kll_helper): each level's capacity is ceil(k · (2/3)^depth) floored at m = 8, where depth counts down from the top. The top level always has capacity k, and every level below it is 2/3 the size of the one above — so the geometric sum converges to ~3k and total retained items is bounded constant in n rather than growing O(k · log(n/k)). Compaction is now a fixed-point loop: find the lowest level whose count exceeds its (dynamically shrinking) capacity, compact it, and loop until none is over. This handles the case where adding a new top level shrinks the effective capacities of every existing level, and the same driver will serve merge (multiple levels may be over capacity at once, unlike after an insert). Odd-length compaction is not skipped: the smallest item is retained at the current level with unchanged weight, and the even remainder is sorted, coin-flipped, and halved. Necessary because the (2/3) formula produces odd capacities, and any odd count on any level would otherwise drift the total-weight invariant by one item's worth per compaction. New test inserts 10K items with k=64 and verifies total retained items stays under the shrinking capacity sum (much tighter than the naive k · num_levels bound), while the mass invariants still hold. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Introduces KllSketch::quantile(q) — returns the smallest retained item
whose cumulative weight (across all levels, weighted by 2^h) reaches
q · total_weight. Non-interpolated: the returned reference is always
to an item actually present in the sketch, matching the shape the
scheduler needs for partition-boundary picking. Empty sketch returns
None; q is clamped to [0, 1].
Three tests, layered:
- quantile_on_empty_returns_none: the None branch.
- quantile_is_exact_below_capacity: n < k, single-level sketch,
exact 0.0/0.5/1.0 answers.
- quantile_range_and_rank_roundtrip_under_compactions: n = 10K
with k = 64 to force multi-level state; asserts every returned
quantile lies in the input range and that rank(quantile(q))
tracks q · n within KLL's normalized error bound (~0.23 at
k = 64). This third test exercises the multi-weight cumulative
walk and the sort-across-levels step that the exact test does
not touch.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Compactions can coin-flip the true stream extremes out of the compactor levels, so quantile(0.0) and quantile(1.0) drifted from the input's actual min and max — an unacceptable shape for partition-boundary picks. Matches Apache DataSketches' separately maintained min_item_/max_item_ pattern. Adds a Clone bound to T (satisfied by u32 and OwnedRow, the two item types the sketch is used with today) and two Option<T> fields updated on insert; the clone fires only when x is a new extreme, which is rare after warmup. quantile(0.0)/quantile(1.0) now short-circuit to those fields before touching the compactors. The new test inserts 10K items with k=64 (deep cascades) and asserts quantile(0.0) == Some(&1) and quantile(1.0) == Some(&10_000). Without the fix it returned Some(&3) at the low end — direct evidence the true minimum had been coin-flipped away. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
New KllSketch::merge(&mut self, other: Self): folds one sketch into another and preserves the standard invariants. Same-height compactors concatenate (items at level h from either sketch carry weight 2^h), extremes cross-fold via explicit min/max on T (dodging the trap that Rust's Option::cmp puts None below any Some), and compact_all walks the full stack afterwards because a merge can leave several levels over their shrinking capacities simultaneously — which is exactly why compact_all is a fixed-point loop instead of a cascade from level 0. Mismatched k is a caller error and caught with debug_assert; tracking a min-of-observed-k like DataSketches is future work. Also documents the "levels are not sorted between operations" choice as a legitimate perf follow-up: DataSketches carries per-level sortedness plus an is_level_zero_sorted_ flag to skip re-sorts and to merge via linear-time merge-sort; we re-sort at every compaction, trading ~log(k) comparisons per amortized insert for shape simplicity. Test: merge two disjoint u32 streams (1..=32 and 33..=64) into one sketch, then verify combined mass (rank(65) == 64), below-min (rank(1) == 0), and that quantile(0.0)/quantile(1.0) return the folded extremes. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@phillipleblanc I was going to go e2e with TDigest, but since I had a bit of time I figured I would get started on KLL. |
RAT (Release Audit Tool) requires every source file to carry the standard Apache 2.0 header; PR CI flagged its absence. Copied the verbatim boilerplate used by every other .rs file in the crate. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The existing merge test used two same-sized streams, so both sketches
end up with roughly the same number of levels — never exercising the
height-asymmetry cases that matter in production (executors observe
wildly different item counts before their sketches are combined).
Two new tests, each in one orientation:
merge_grows_stack_when_other_is_taller — small sketch absorbs a
large one. The loop hits `push(Vec::new())` on every iteration
past small's original top; asserts the stack grew to at least
large's height so no promoted item is silently lost.
merge_preserves_upper_levels_when_other_is_shorter — large sketch
absorbs a small one. The loop terminates before large's upper
levels; asserts they weren't dropped and that mass/extremes at
the high end are intact.
Both also assert the standard invariants (mass, below-min, exact
min/max), so a bug that only shows up when heights differ would
surface here rather than in the symmetric test.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
andygrove
approved these changes
Jul 25, 2026
andygrove
left a comment
andygrove
left a comment
Member
There was a problem hiding this comment.
Purely additive. LGTM. Thanks @avantgardnerio
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
2 participants
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
ballista_core::kll::KllSketch<T: Ord + Clone>— native Rust KLL quantile sketch, generic over ordered items, composes witharrow::row::OwnedRowfor the full ORDER BY grammar (multi-column, nullable, per-column ASC/DESC).kll_helper.hpp): shrinking-k per level so total retained items is bounded ~3kregardless of stream size, keep-smallest odd-length compaction preserves the total-weight invariant, min/max side-tracked for exactquantile(0.0)/quantile(1.0), walk-all-levels compaction driver somergefolds cleanly when several levels land over their (shrinking) capacities at once.RuntimeStatsExecswap (per its existing TODO onruntime_stats.rs:39) is a follow-up PR.Test plan
cargo test --lib -p ballista-core kll::— 11 tests passing locallycargo clippy --all-targetsandcargo fmt --checkcleanRuntimeStatsExec::quantile_sketch()andmerged_quantile_sketch(), replacing the TDigest🤖 Generated with Claude Code