You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by av...@apache.org on 2026/08/12 01:16:53 UTC
(datafusion-ballista) branch main updated: perf(kll): batch-oriented ingest KLL (#2216)
This is an automated email from the ASF dual-hosted git repository.
avantgardnerio pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new a4bf0190f perf(kll): batch-oriented ingest KLL (#2216)
a4bf0190f is described below
commit a4bf0190f0806e51b5936dacd7d60ce798129afd
Author: Brent Gardner <bg...@squarelabs.net>
AuthorDate: Tue Aug 11 19:16:48 2026 -0600
perf(kll): batch-oriented ingest KLL (#2216)
* bench(quantile-sketch): TDigest vs KLL<OwnedRow> ingest on the RuntimeStatsExec path
Measures the sketch swap RuntimeStatsExec is queued for: TDigest today
(`Vec<f64>` → `merge_unsorted_f64` per batch) vs KLL over `OwnedRow` (batch
→ RowConverter → per-row `owned()` → `insert`). Same uniform-f64 stream
into both, seeded, batch size 8192, TDigest max_size=100 (production),
KLL k=800 (analytical ~1/√k ε ≈ 0.035 to match TDigest at the tails).
Ingest-only — for the sizes RuntimeStatsExec sees, merge and quantile
query are 3+ orders of magnitude cheaper than ingest.
Rows TDigest KLL (OwnedRow) Ratio
100K 1.38 ms (73 M/s) 9.13 ms (11 M/s) 6.6× slower
1M 15.1 ms (66 M/s) 94.3 ms (11 M/s) 6.3× slower
Consistent ~6.5× regression across scale — per-row overhead, not
compaction blowup. Breakdown across encode / owned() / insert /
compaction cmp is TBD from a flamegraph before we decide between the
"land with TODO for f64 fast path" and "special-case single-column
native-Ord" designs.
* bench(quantile-sketch): add KLL<OrderedFloat<f64>> — isolates arrow-row cost
Third variant bypasses RowConverter + owned() entirely by using
OrderedFloat<f64> as the sketch item type, so ingest exercises only the
KLL algorithm itself. Diffing this against KLL<OwnedRow> attributes cost
between arrow-row encoding and the sketch algorithm; diffing against
TDigest measures the sketch algorithm's inherent overhead.
Rows TDigest KLL<OrdFloat> KLL<OwnedRow>
100K 1.36 ms (74 M/s) 4.44 ms (23 M/s) 9.07 ms (11 M/s)
1M 13.7 ms (73 M/s) 51.3 ms (20 M/s) 92.0 ms (11 M/s)
At 1M rows:
- KLL<OrdFloat> is 3.7× slower than TDigest — the KLL algorithm's
inherent cost (per-row insert + compact_all walk vs TDigest's batch
merge_unsorted_f64).
- KLL<OwnedRow> is 6.7× slower than TDigest — arrow-row encoding +
owned() heap alloc doubles the OrdFloat cost.
Implication: dispatching by column type in RuntimeStatsExec at
construction (single-column native-Ord → KllSketch<T>, multi-column or
non-Ord → KllSketch<OwnedRow>) halves the regression for the common
parallel-window PARTITION-BY-one-column case.
* docs(kll): collect deferred perf optimizations behind a single section
The three follow-ups previously appeared as inline "future optimization"
asides under Ownership and Level order — split across two sections and
easy to miss. Collect them under a "Deferred optimizations" section with
the KLL vs TDigest bench numbers as motivation, and note explicitly which
`T` each helps and which stack on which.
* test(kll): statistical rank-error bound + merge equivalence
Before optimizing the ingest hot path we need a way to catch subtle
correctness regressions — the existing tests check mass invariants and
a single loose per-seed error bound (slack ±0.25n at k=64), which won't
detect an optimization bug that biases the error distribution by a few
percent.
Two new tests, both fixed-seed for reproducibility, sweeping 100 seeds
each:
1. `rank_error_meets_datasketches_bound_across_seeds` — for k=200, n=10K
uniform shuffled stream, probe rank at 9 interior quantiles, assert
≤5% of trials exceed DataSketches' empirical single-sided normalized
rank error `2.296 / k^0.9723` (source:
`datasketches-cpp/kll/include/kll_sketch_impl.hpp:619`, fit to P99
max error across their trials). Since the bound is a P99 fit, ≤1%
expected; 5% slack absorbs tail luck without letting a real
distribution shift slip through.
2. `merge_matches_single_sketch_within_bound` — same stream inserted
two ways: whole-stream into one sketch, vs. split into 4 shards
merged pairwise. Assert quantile-probe difference ≤ 2ε (linear
sum-of-errors upper bound), same 5% fail-rate assertion. Guards
any batch-oriented merge rewrite from silently drifting.
Both run in ~0.7s combined. 15/15 kll tests pass.
* test(kll): rank-error bound holds under clustered + heavy-ties distributions
Uniform-only coverage would miss any bug that surfaces only under skew
or degenerate orderings — the classes an optimizer is most likely to
break. Refactor the trial helper to consume any u32 stream and compute
true rank from a sorted copy via partition_point (works for arbitrary
distributions), then add two more distributions:
- clustered: 90% of items in the bottom 1% of the value range, 10%
spread over the rest. Cuts near low quantiles land inside the dense
cluster; cuts near high quantiles land in the sparse tail.
- heavy-ties: only 100 distinct values, ~n/100 copies each. Exercises
sort + coin-flip halving under degenerate orderings.
All three distributions hold the DataSketches P99 bound at ≤5% fail
rate across 100 seeds. Also picks up an incidental fmt fix in the
benchmark.
* perf(kll): batch-oriented ingest — absorb (iterator) + absorb_slice (Clone)
Two new ingest APIs alongside `insert`, both bit-identical to a per-row
insert loop:
- `absorb<I: IntoIterator<Item = T>>` — caches level 0's capacity once
and skips `compact_all`'s O(L) scan when it can prove no level is
full. Takes owned items to preserve the no-clone path for heap-heavy
T (e.g. OwnedRow).
- `absorb_slice(&[T])` — stacks two extra wins on top of `absorb`:
batch min/max in a tight scan (one Option update per batch instead
of per row), and `Vec::extend_from_slice` for chunk-fills of level 0
(memcpy for Copy types, potentially SIMD-vectorized). Trades one
clone per item; prefer for Copy-T, avoid for heap-heavy T.
Bench (runtime_stats_ingest, N=1M rows):
variant time thrpt vs TDigest
tdigest 13.8 ms 73 M/s 1.0×
kll_ordered_float insert 52.3 ms 19 M/s 3.8×
kll_ordered_float absorb_slice 29.0 ms 35 M/s 2.1× (45% faster)
kll_row insert 94.0 ms 11 M/s 6.8×
kll_row absorb 80.5 ms 12 M/s 5.8× (14% faster)
The Copy-path (OrderedFloat) picks up all three wins; the heap-path
(OwnedRow) picks up only the compact_all amortization since we still
push per-row. Deferred-optimization item (2) — defer ownership at
level 0 — remains the path to closing the OwnedRow gap further.
Correctness: new `absorb_apis_match_per_row_insert` test asserts three
sketches seeded identically and fed the same stream via insert, absorb,
and absorb_slice agree on every rank probe.
* docs
* perf(kll): per-level sortedness invariant — skip compact-time re-sort
The compaction sort dominated ingest time once the batch-oriented
`absorb`/`absorb_slice` amortized `compact_all`'s scan overhead. This
lands DataSketches' `is_level_zero_sorted_` bookkeeping, generalized
per-level:
- `sorted: Vec<bool>` runs in lockstep with `levels`. Every insertion
path clears `sorted[0]` (arbitrary data breaks any prior order).
- `compact_level(h)` skips `sort_unstable` when `sorted[h]` is already
set. Post-compaction level `h` is empty (or one leftover) — trivially
sorted, so `sorted[h] = true`.
- Promoted items form a sorted subsequence of a sorted array (every
other index). If the destination `h+1` is already sorted, we
merge-extend in linear time via a new `merge_sorted` helper instead
of concatenate-then-sort at the next compaction.
- Cross-sketch `KllSketch::merge` doesn't yet fold the two sortedness
flags — noted as further deferred work.
Bench at 1M rows (uniform f64 stream):
variant before after vs TDigest
tdigest 13.8 ms 14.0 ms 1.0×
kll_ordered_float insert 52.3 ms 48.0 ms 3.4×
kll_ordered_float absorb_slice 29.0 ms 24.9 ms 1.8× (was 2.1×)
kll_row insert 94.0 ms 92.9 ms 6.6×
kll_row absorb 80.5 ms 80.1 ms 5.7×
The Copy path is now 1.8× TDigest — 3.8× → 1.8× is the story reviewers
see for swapping in KLL for the f64 case. OwnedRow gets almost nothing
here because the compaction sort's cost for heap-allocated rows is
dominated by cache misses on the OwnedRow storage, not the sort's
algorithmic efficiency — deferred item (2), level-0 borrowed rows,
remains the path to closing the OwnedRow gap.
Correctness: bit-identical for unique-value streams (existing
`absorb_apis_match_per_row_insert` still passes); statistical bound
(rank_error_bound_{uniform,clustered,heavy_ties},
merge_matches_single_sketch_within_bound) all continue to hold at ≤5%
fail rate across 100 seeds. `sort_unstable` and `merge_sorted` produce
identical output when all keys are distinct, and both preserve KLL's
error guarantee under duplicates.
* docs(kll): mark absorb (iterator) as interim, retire when borrowed-row API lands
The iterator `absorb` only helps `T = OwnedRow` — the Copy path uses
`absorb_slice` for the memcpy extend + batch min/max. And when
deferred-optimization item (2) lands, the OwnedRow path will take an
arrow `Rows` object directly (holding `Row<'_>` at level 0, materializing
`OwnedRow` only at promotion) — a specialized API, not a generalization
of the iterator one.
Keeping `absorb` for now as a modest 14% speedup over `insert`-in-loop
for the interim OwnedRow path, but flagging its provisional status so
the next dev doesn't grow dependencies on it.
* perf(kll): absorb_sorted_slice — skip every compaction sort for sorted input
The single largest cost on the Copy path after items (1) and (3) landed
was the level-0 sort at each compaction — the per-level sortedness
invariant only helps levels 1+ because level 0 receives raw unsorted
batch data. If the caller can guarantee sorted input, level 0 stays
permanently sorted via merge-extend and compaction skips every sort,
not just those at levels 1+.
The motivating plan shape is the ORRE arm:
current: Scan -> RuntimeStats -> Sort -> ORRE -> ShuffleWrite
follow-up: Scan -> Sort -> RuntimeStats -> DamExec -> ORRE -> ShuffleWrite
`SortExec` already runs upstream of ORRE, so moving it above
`RuntimeStatsExec` is a plan-shape change with no added CPU cost — the
Sort was going to happen anyway. `DamExec` between stats and ORRE
holds the sorted rows until the scheduler's global cuts arrive.
Sketch side, added:
- `absorb_sorted_slice(&[T])` — merge-extends level 0 with sorted chunks
instead of `extend_from_slice`+dirty. Min/max read as `xs[0]`/`xs[last]`,
no scan.
- `debug_assert!` on ascending order catches caller bugs in test builds.
- Correctness test: bit-identical to `absorb_slice` on sorted unique
input (sort_unstable of a sorted distinct array is the identity, so
both paths halve the same subsequence).
Bench (uniform f64 stream, globally sorted for the new variant):
variant 100K 1M vs TDigest
tdigest 1.37 ms 13.8 ms 1.0×
kll_ordered_float insert 3.92 ms 48.3 ms 3.5×
kll_ordered_float absorb_slice 2.37 ms 25.1 ms 1.8×
kll_ordered_float absorb_sorted_slice 0.375 ms 7.26 ms 0.53× (1.9× FASTER)
The story for reviewers flips: swapping TDigest for KLL is 1.8× on the
unsorted path today; with the ORRE plan-shape follow-up it's 1.9×
FASTER than TDigest at ingest, on top of unlocking multi-column,
nullable, and per-column DESC ORDER BY grammar for partition boundaries.
The sketch side is in place; the planner-side move is deferred to a
follow-up branch. Documented as item (4) in the module doc.
* fix(kll): absorb_sorted_slice verifies sortedness, falls back to absorb_slice
Replaces the debug_assert! precondition (compiled out under [profile.ci], so
the #[should_panic] test couldn't fire) with an unconditional O(n) sortedness
check. Unsorted input falls through to absorb_slice, which sorts at compaction
time. If a future optimizer change moves or drops the SortExec above
RuntimeStatsExec, quantiles stay correct — just without the sort-skip speedup.
* fix(kll): absorb_sorted_slice restores level 0 sortedness on mixed-use entry
Prior insert/absorb_slice calls leave sorted[0] = false. absorb_sorted_slice
would then merge_sorted a sorted batch into an unsorted buffer, producing
unsorted output — still correct (compaction re-sorts) but a lie against the
doc claim "level 0 stays sorted through the whole ingest". Sort-and-set at
the top so the invariant holds for any caller, not just exclusive users.
* bench(quantile-sketch): empirical KLL/TDigest accuracy parity
The old KLL_K = 800 justification was an analytical hand-wave (1/√k ≈ 0.035
normalized error), and pointed the wrong direction — measured on the 1M
uniform-f64 stream this bench feeds, KLL at k=800 sits at 0.0016 worst rank
error across the deciles vs TDigest max_size=100 at 0.0021. Same k, honest
justification.
Adds an opt-in parity harness printed at bench start when KLL_PARITY_CHECK=1
so the numbers are reproducible.
* bench(quantile-sketch): sorted TDigest arm + fmt
Add tdigest_sorted/{n} arm using TDigest::merge_sorted_f64 on the same
globally-sorted batches the absorb_sorted_slice arm consumes, so the ORRE-
plan-shape race is apples-to-apples: both sketches get sorted input, both
skip their own pre-sort. On n=1M this shows TDigest sorted at 2.19 ms
(457.6 M/s) vs KLL absorb_sorted_slice at 7.22 ms (138.6 M/s) — the "1.9x
faster" claim from the previous PR body was sorted-KLL vs unsorted-TDigest,
not a fair race. Phillip's review comment.
Also picks up cargo fmt line-wraps in print_parity_table that CI's rustfmt
--check flagged.
* chore(deps): centralize ordered-float + move aggregate-common to dev-deps
Both are Andy's smaller review asks:
- ordered-float = "4" moves from a raw version in benchmarks/Cargo.toml to
workspace.dependencies (ballista manages crate versions centrally), and
the bench references it via { workspace = true }.
- datafusion-functions-aggregate-common was in [dependencies] on the bench
crate but only used by benches/quantile_sketch.rs — move to
[dev-dependencies] so it isn't linked into shuffle_bench / tpcds / tpch
binaries that don't need it.
---
Cargo.lock | 2 +
Cargo.toml | 1 +
ballista/core/src/kll.rs | 623 +++++++++++++++++++++++++++++++++-
benchmarks/Cargo.toml | 6 +
benchmarks/benches/quantile_sketch.rs | 348 +++++++++++++++++++
5 files changed, 966 insertions(+), 14 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 60948985e..e7d4489e4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1040,10 +1040,12 @@ dependencies = [
"clap 4.6.3",
"criterion",
"datafusion",
+ "datafusion-functions-aggregate-common",
"datafusion-proto",
"env_logger",
"futures",
"mimalloc",
+ "ordered-float 4.6.0",
"rand 0.10.2",
"serde",
"serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index 986980717..90405c7b8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -58,6 +58,7 @@ insta = "1.47"
itertools = "0.15"
mimalloc = { version = "0.1" }
object_store = "0.13.2"
+ordered-float = "4"
prost = "0.14"
prost-types = "0.14"
rstest = { version = "0.26" }
diff --git a/ballista/core/src/kll.rs b/ballista/core/src/kll.rs
index ca2ca15d6..33035f87b 100644
--- a/ballista/core/src/kll.rs
+++ b/ballista/core/src/kll.rs
@@ -82,11 +82,8 @@
//! Callers hand fully-owned items to `KllSketch::insert`. For the Arrow
//! use case this means materializing each input `Row<'_>` into an
//! `OwnedRow` up front — one small heap allocation per input row, most of
-//! which are eventually discarded by cascading compactions. A future
-//! optimization can defer materialization: hold borrowed `Row<'_>` in
-//! level 0 within a single batch, compact intra-batch, and only own the
-//! survivors that promote to level 1. That would trade the current
-//! per-item API for a batch-oriented one (e.g. `absorb_rows(&Rows)`).
+//! which are eventually discarded by cascading compactions. See
+//! **Deferred optimizations** below.
//!
//! # Level order
//!
@@ -94,11 +91,53 @@
//! `quantile` collects-then-sorts, and every compaction sorts before
//! halving — so `merge` can simply extend same-height compactors without
//! order concerns. That trades ~`log(k)` extra comparisons per amortized
-//! insert for a simpler shape. Apache DataSketches maintains per-level
-//! sortedness (with an `is_level_zero_sorted_` flag) to skip re-sorting at
-//! compact time and to merge two sketches via linear-time merge-sort
-//! rather than concatenate-and-resort — a legitimate follow-up if
-//! benchmarks demand it.
+//! insert for a simpler shape. See **Deferred optimizations** below for
+//! the per-level-sortedness follow-up that DataSketches implements.
+//!
+//! # Deferred optimizations
+//!
+//! Ingest measured on uniform 1M-row streams (see
+//! `benchmarks/benches/quantile_sketch.rs`):
+//!
+//! - `T = OrderedFloat<f64>` via `absorb_slice`: **1.8× TDigest** (down
+//! from 3.7× on per-row `insert` with no batch API).
+//! - `T = OwnedRow` via `absorb`: **5.7× TDigest** (down from 6.8× on
+//! per-row `insert`).
+//!
+//! Items (1) and (3) below have landed. One deferred win remains, and
+//! it's the one that matters for the OwnedRow path:
+//!
+//! 1. ~~**Amortize `compact_all` across a batch.**~~ **Landed.** `absorb`
+//! caches level 0's capacity and skips the `compact_all` O(L) scan
+//! when no level is full. `absorb_slice` stacks a batch min/max and
+//! `extend_from_slice` on top for `T: Copy`.
+//!
+//! 2. **Defer ownership at level 0.** For `T = OwnedRow` the caller pays
+//! one heap allocation per input row via `row.owned()`, even though
+//! ~half of level-0 items are coin-flipped out at the first compaction
+//! and never promoted. Holding borrowed `Row<'_>` at level 0 within a
+//! batch and materializing `OwnedRow` only on promotion to level 1
+//! saves the alloc/free on the discarded half. The batch-oriented
+//! `absorb` API is what makes the borrow lifetime work out — level 0
+//! lives for the duration of `absorb`, compacts before the caller's
+//! `Rows` buffer goes out of scope. Only helps `T = OwnedRow`.
+//!
+//! 3. ~~**Per-level sortedness invariant.**~~ **Landed.** `KllSketch`
+//! carries a `sorted: Vec<bool>` parallel to `levels`. `compact_level`
+//! skips the sort when the level is already ordered, and promoted
+//! items merge-extend the next level in linear time when it too was
+//! sorted. `KllSketch::merge` (across sketches) does not yet fold the
+//! two sortedness flags — cross-sketch merge-extend is a further
+//! deferred win.
+//!
+//! 4. **Sorted-input ingest via a caller-side plan-shape change.** The
+//! `absorb_sorted_slice` API exists but is not yet called from
+//! production. When wired into an ORRE plan shape where a `SortExec`
+//! is placed upstream of `RuntimeStatsExec` (with a `DamExec` between
+//! stats and the router), level 0 stays permanently sorted and
+//! compaction skips *every* sort — not just those at levels 1+. That
+//! kills the largest remaining ingest cost on the `T: Copy` path.
+//! The sketch side is free; the win requires the planner-side move.
//!
//! # Reference
//!
@@ -122,6 +161,14 @@ const MIN_LEVEL_WIDTH: usize = 8;
pub struct KllSketch<T: Ord + Clone> {
/// One buffer per level; `levels[h]` is the level-`h` compactor.
levels: Vec<Vec<T>>,
+ /// `sorted[h] == true` iff `levels[h]` is in ascending order. Same
+ /// length as `levels`, kept in lockstep with every mutation.
+ /// Compaction reads this to skip the re-sort when a level was left
+ /// sorted by an earlier promotion; promotions extend the destination
+ /// level via linear-time merge instead of concatenate-then-sort when
+ /// the destination is already sorted. Matches DataSketches'
+ /// `is_level_zero_sorted_` bookkeeping, generalized per-level.
+ sorted: Vec<bool>,
/// Nominal compactor capacity — the top level's capacity. Lower
/// levels shrink geometrically toward `MIN_LEVEL_WIDTH`.
k: usize,
@@ -167,6 +214,8 @@ impl<T: Ord + Clone> KllSketch<T> {
pub fn with_seed(k: usize, seed: u64) -> Self {
Self {
levels: vec![Vec::new()],
+ // An empty level is trivially sorted.
+ sorted: vec![true],
k,
rng: StdRng::seed_from_u64(seed),
min: None,
@@ -200,11 +249,19 @@ impl<T: Ord + Clone> KllSketch<T> {
(None, b) => b,
};
// Extend same-height compactors, growing the stack as needed.
+ // Cross-sketch extend does not preserve any per-level sortedness
+ // — the concatenation of two sorted runs is not sorted, and we
+ // don't yet fold the sortedness flags of two sketches into a
+ // merge-extend. `compact_all` will sort on demand.
for (h, level) in other.levels.into_iter().enumerate() {
if h == self.levels.len() {
self.levels.push(Vec::new());
+ self.sorted.push(true);
+ }
+ if !level.is_empty() {
+ self.levels[h].extend(level);
+ self.sorted[h] = false;
}
- self.levels[h].extend(level);
}
self.compact_all();
}
@@ -220,9 +277,172 @@ impl<T: Ord + Clone> KllSketch<T> {
self.max = Some(x.clone());
}
self.levels[0].push(x);
+ // Arbitrary insertion breaks any existing sort order at level 0.
+ self.sorted[0] = false;
self.compact_all();
}
+ /// Add many items in one call, amortizing the per-item overhead of
+ /// `insert`.
+ ///
+ /// `insert` invokes `compact_all` after every row, and each call scans
+ /// every level asking "is anyone full?" That's an `O(L)` scan per row
+ /// even when the answer is universally "no." `absorb` caches level 0's
+ /// capacity once, checks a single length per row, and only fires
+ /// `compact_all` when level 0 actually reaches capacity.
+ ///
+ /// Owned-value ingest — takes items by value, no `T: Clone` needed on
+ /// the incoming stream. Prefer this for heap-heavy `T` (e.g.
+ /// `OwnedRow`) where cloning would double the allocation cost. For
+ /// `T: Copy` prefer `absorb_slice` — it uses `extend_from_slice` +
+ /// batch min/max for a memcpy-friendly hot loop.
+ ///
+ /// Output is bit-identical to `for x in xs { self.insert(x) }`: the
+ /// compaction trigger points, coin flips, and min/max tracking are all
+ /// preserved.
+ ///
+ /// TODO: retire once **Deferred optimizations** item (2) lands. The
+ /// borrowed-row ingest path there takes an arrow `Rows` directly and
+ /// holds `Row<'_>` at level 0 — a specialized API, not a generalization
+ /// of this one. This function exists as a modest speedup for the
+ /// interim OwnedRow path (~14% over `insert`-in-loop on 1M rows) and
+ /// has no other production caller.
+ pub fn absorb<I: IntoIterator<Item = T>>(&mut self, xs: I) {
+ // Cache level 0's capacity; recompute whenever compaction may have
+ // grown the stack (which shrinks per-level capacities).
+ let mut level_0_cap = level_capacity(self.k, self.levels.len(), 0);
+ for x in xs {
+ if self.min.as_ref().is_none_or(|m| x < *m) {
+ self.min = Some(x.clone());
+ }
+ if self.max.as_ref().is_none_or(|m| x > *m) {
+ self.max = Some(x.clone());
+ }
+ self.levels[0].push(x);
+ self.sorted[0] = false;
+ if self.levels[0].len() >= level_0_cap {
+ self.compact_all();
+ level_0_cap = level_capacity(self.k, self.levels.len(), 0);
+ }
+ }
+ }
+
+ /// Slice-ingest variant that trades one clone per item for two extra
+ /// wins on top of `absorb`'s compact_all amortization:
+ ///
+ /// 1. Batch min/max in a single tight scan of `xs`, updating `self.min`
+ /// / `self.max` at most once per call instead of at most once per
+ /// row.
+ /// 2. `Vec::extend_from_slice` for chunk-fills of level 0 — one memcpy
+ /// per chunk for `T: Copy`, potentially SIMD-vectorized by LLVM.
+ ///
+ /// Prefer this for `T: Copy` (e.g. `OrderedFloat<f64>`, `i64`) where
+ /// the clone is free. Avoid for heap-heavy `T` (e.g. `OwnedRow`) —
+ /// there the extend_from_slice clone allocates on top of whatever the
+ /// caller already paid to construct the value; `absorb` moves instead.
+ ///
+ /// Output is bit-identical to `for x in xs { self.insert(x.clone()) }`.
+ pub fn absorb_slice(&mut self, xs: &[T]) {
+ if xs.is_empty() {
+ return;
+ }
+ // One scan across xs for min/max; a single Option update per batch
+ // regardless of how many progressive extremes appear in xs.
+ let (batch_min, batch_max) = {
+ let mut min = &xs[0];
+ let mut max = &xs[0];
+ for x in &xs[1..] {
+ if x < min {
+ min = x;
+ }
+ if x > max {
+ max = x;
+ }
+ }
+ (min, max)
+ };
+ if self.min.as_ref().is_none_or(|m| batch_min < m) {
+ self.min = Some(batch_min.clone());
+ }
+ if self.max.as_ref().is_none_or(|m| batch_max > m) {
+ self.max = Some(batch_max.clone());
+ }
+ // Chunk-fill level 0 via memcpy-style extend; compact when it fills.
+ let mut i = 0;
+ while i < xs.len() {
+ let cap = level_capacity(self.k, self.levels.len(), 0);
+ let room = cap.saturating_sub(self.levels[0].len());
+ let take = (xs.len() - i).min(room);
+ self.levels[0].extend_from_slice(&xs[i..i + take]);
+ self.sorted[0] = false;
+ i += take;
+ if self.levels[0].len() >= cap {
+ self.compact_all();
+ }
+ }
+ }
+
+ /// Slice-ingest variant that expects the input to already be sorted
+ /// ascending — the caller is downstream of a `SortExec` or otherwise
+ /// producing sorted output. Keeps level 0 permanently sorted by
+ /// merge-extending each chunk into it, so `compact_level(0)` never
+ /// pays a `sort_unstable`. That's the largest remaining cost on the
+ /// `T: Copy` ingest path once the per-level sortedness invariant is
+ /// in place.
+ ///
+ /// The min/max scan is also skipped — for sorted input the extremes
+ /// are `xs[0]` and `xs[last]`.
+ ///
+ /// Sortedness is verified unconditionally (O(n) comparisons against
+ /// the O(n log n) sort skipped when it holds). Unsorted input falls
+ /// through to `absorb_slice`, which sorts at compaction time. The
+ /// motivating plan shape is `Scan -> Sort -> RuntimeStats -> DamExec
+ /// -> ORRE`, but a future optimizer change that moves or drops the
+ /// upstream `SortExec` must not silently corrupt quantiles — the
+ /// fallback keeps this method safe under any caller.
+ pub fn absorb_sorted_slice(&mut self, xs: &[T]) {
+ if xs.is_empty() {
+ return;
+ }
+ if !xs.is_sorted() {
+ self.absorb_slice(xs);
+ return;
+ }
+ // Sorted input: extremes are the first and last elements. No scan.
+ let batch_min = &xs[0];
+ let batch_max = &xs[xs.len() - 1];
+ if self.min.as_ref().is_none_or(|m| batch_min < m) {
+ self.min = Some(batch_min.clone());
+ }
+ if self.max.as_ref().is_none_or(|m| batch_max > m) {
+ self.max = Some(batch_max.clone());
+ }
+ // If a prior `insert` or `absorb_slice` left level 0 unsorted,
+ // `merge_sorted` below would interleave sorted `xs` into
+ // unsorted state and the result wouldn't be sorted. Pay the sort
+ // once here so the loop's merge-extends stay honest and the doc's
+ // "level 0 stays sorted" claim actually holds.
+ if !self.sorted[0] {
+ self.levels[0].sort_unstable();
+ self.sorted[0] = true;
+ }
+ // Chunk-merge sorted `xs` into (sorted) level 0. Level 0 stays
+ // sorted through the whole ingest — `sorted[0]` never flips false.
+ let mut i = 0;
+ while i < xs.len() {
+ let cap = level_capacity(self.k, self.levels.len(), 0);
+ let room = cap.saturating_sub(self.levels[0].len());
+ let take = (xs.len() - i).min(room);
+ let incoming = xs[i..i + take].to_vec();
+ let existing = std::mem::take(&mut self.levels[0]);
+ self.levels[0] = merge_sorted(existing, incoming);
+ i += take;
+ if self.levels[0].len() >= cap {
+ self.compact_all();
+ }
+ }
+ }
+
/// Return the estimated number of items strictly less than `x`.
///
/// Sum over levels `h` of `2^h · | { y ∈ levels[h] : y < x } |`.
@@ -306,7 +526,8 @@ impl<T: Ord + Clone> KllSketch<T> {
}
}
- /// Sort level `h`, coin-flip, promote every other item to level `h+1`.
+ /// Sort level `h` (or reuse an existing sort), coin-flip, promote every
+ /// other item to level `h+1`.
///
/// On odd-length levels the smallest item stays behind at level `h`
/// with its original weight — the algorithm requires an even-length
@@ -314,13 +535,25 @@ impl<T: Ord + Clone> KllSketch<T> {
/// current weight is what preserves the total-weight invariant across
/// compaction. Matches Apache DataSketches' `general_compress`
/// odd-pop handling.
+ ///
+ /// Sortedness bookkeeping: promoted items form a sorted subsequence
+ /// (every other element of a sorted array). If the destination level
+ /// was already sorted (which it is by default when previously touched
+ /// only by this function), we merge-extend in linear time instead of
+ /// concatenate-then-sort. The saving compounds up the stack — for
+ /// large streams, levels 1+ are always sorted going into compaction,
+ /// so the `sort_unstable` at level 0 is the only one paid.
fn compact_level(&mut self, h: usize) {
let level_len = self.levels[h].len();
if level_len < 2 {
return;
}
let mut buf = std::mem::take(&mut self.levels[h]);
- buf.sort_unstable();
+ if !self.sorted[h] {
+ buf.sort_unstable();
+ }
+ // levels[h] is now empty and thus trivially sorted.
+ self.sorted[h] = true;
let leftover = if level_len % 2 == 1 {
Some(buf.remove(0))
} else {
@@ -335,14 +568,55 @@ impl<T: Ord + Clone> KllSketch<T> {
.collect();
if h + 1 == self.levels.len() {
self.levels.push(Vec::new());
+ self.sorted.push(true);
+ }
+ // If the destination is already sorted, linear-time merge preserves
+ // its sortedness. Otherwise fall back to append; the next compaction
+ // at h+1 will pay the sort.
+ if self.sorted[h + 1] {
+ let existing = std::mem::take(&mut self.levels[h + 1]);
+ self.levels[h + 1] = merge_sorted(existing, promoted);
+ } else {
+ self.levels[h + 1].extend(promoted);
}
- self.levels[h + 1].extend(promoted);
if let Some(x) = leftover {
+ // Single-item level is trivially sorted.
self.levels[h].push(x);
}
}
}
+/// Merge two sorted vectors into one sorted vector in linear time.
+/// Preserves ties in source order (`a` before `b`) — matters only for
+/// non-unique `T`, where downstream halving still picks half the items
+/// but may pick different specific copies than `sort_unstable(a ++ b)`.
+fn merge_sorted<T: Ord>(a: Vec<T>, b: Vec<T>) -> Vec<T> {
+ let mut out = Vec::with_capacity(a.len() + b.len());
+ let mut ai = a.into_iter().peekable();
+ let mut bi = b.into_iter().peekable();
+ loop {
+ match (ai.peek(), bi.peek()) {
+ (Some(av), Some(bv)) => {
+ if av <= bv {
+ out.push(ai.next().unwrap());
+ } else {
+ out.push(bi.next().unwrap());
+ }
+ }
+ (Some(_), None) => {
+ out.extend(ai);
+ break;
+ }
+ (None, Some(_)) => {
+ out.extend(bi);
+ break;
+ }
+ (None, None) => break,
+ }
+ }
+ out
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -667,4 +941,325 @@ mod tests {
assert_eq!(sketch.rank(&10_001), 10_000);
assert_eq!(sketch.rank(&1), 0);
}
+
+ /// Apache DataSketches' single-sided normalized rank error, fit to the
+ /// P99 max error across thousands of empirical trials. Source:
+ /// `datasketches-cpp/kll/include/kll_sketch_impl.hpp:619`.
+ fn normalized_rank_error_ss(k: usize) -> f64 {
+ 2.296 / (k as f64).powf(0.9723)
+ }
+
+ /// Insert `stream` into a KLL sketch at capacity `k`, probe rank at 9
+ /// interior quantiles picked from a sorted copy of the stream, and
+ /// return the worst normalized rank error observed.
+ fn worst_rank_error(stream: &[u32], k: usize, sketch_seed: u64) -> f64 {
+ let mut sketch: KllSketch<u32> = KllSketch::with_seed(k, sketch_seed);
+ for x in stream {
+ sketch.insert(*x);
+ }
+ let mut sorted = stream.to_vec();
+ sorted.sort_unstable();
+ let n = stream.len();
+ (1..10)
+ .map(|i| {
+ let idx = (i as f64 / 10.0 * n as f64) as usize;
+ let probe = sorted[idx.min(n - 1)];
+ // True rank: count of items strictly less than probe.
+ // partition_point on the sorted stream gives it in O(log n).
+ let true_rank_frac =
+ sorted.partition_point(|x| *x < probe) as f64 / n as f64;
+ let est_rank_frac = sketch.rank(&probe) as f64 / n as f64;
+ (est_rank_frac - true_rank_frac).abs()
+ })
+ .fold(0.0_f64, f64::max)
+ }
+
+ /// Shuffled `1..=n` — every distinct value appears exactly once.
+ fn uniform_stream(n: u32, rng: &mut StdRng) -> Vec<u32> {
+ use rand::seq::SliceRandom;
+ let mut v: Vec<u32> = (1..=n).collect();
+ v.shuffle(rng);
+ v
+ }
+
+ /// 90% of items in the bottom 1% of the value range, remainder spread
+ /// over the rest. Cuts near the low quantiles land inside the dense
+ /// cluster; cuts near the high quantiles land in the sparse tail.
+ fn clustered_stream(n: u32, rng: &mut StdRng) -> Vec<u32> {
+ use rand::seq::SliceRandom;
+ let hot_max = (n / 100).max(2);
+ let hot = n * 9 / 10;
+ let cold = n - hot;
+ let mut v = Vec::with_capacity(n as usize);
+ for _ in 0..hot {
+ v.push(rng.random_range(1u32..=hot_max));
+ }
+ for _ in 0..cold {
+ v.push(rng.random_range(hot_max + 1..=n));
+ }
+ v.shuffle(rng);
+ v
+ }
+
+ /// Only 100 distinct values, each ~n/100 times. Exercises the sort +
+ /// coin-flip halving under degenerate orderings and confirms that
+ /// ties don't shift rank estimates.
+ fn tied_stream(n: u32, rng: &mut StdRng) -> Vec<u32> {
+ (0..n).map(|_| rng.random_range(1u32..=100)).collect()
+ }
+
+ /// Sweep `trials` seeds building a fresh stream per trial and assert
+ /// the DataSketches bound holds in ≥95% of them.
+ fn assert_ds_bound_holds(
+ distribution: &str,
+ stream_fn: impl Fn(u32, &mut StdRng) -> Vec<u32>,
+ ) {
+ let k = 200;
+ let n = 10_000u32;
+ let bound = normalized_rank_error_ss(k);
+ let trials = 100u64;
+ let failures = (0..trials)
+ .filter(|&seed| {
+ let mut shuffle_rng = StdRng::seed_from_u64(seed);
+ let stream = stream_fn(n, &mut shuffle_rng);
+ // Sketch RNG derived from a different sub-seed so the
+ // stream construction and the compaction coin flips aren't
+ // drawn from correlated streams.
+ let sketch_seed = seed.wrapping_add(0x9E37_79B9);
+ worst_rank_error(&stream, k, sketch_seed) > bound
+ })
+ .count();
+ let fail_rate = failures as f64 / trials as f64;
+ // The DataSketches eps is a P99 empirical fit, so ≤1% of trials
+ // should exceed it in a correct implementation. 5% slack absorbs
+ // tail luck without letting a real distribution shift slip
+ // through — an optimization bug that biases the error
+ // distribution up by a few percent pushes the fail rate well
+ // past 5%.
+ assert!(
+ fail_rate <= 0.05,
+ "{distribution}: rank error exceeded DataSketches bound \
+ {bound:.4} in {failures}/{trials} trials ({:.1}%); expected ≤ 5%",
+ fail_rate * 100.0
+ );
+ }
+
+ #[test]
+ fn absorb_apis_match_per_row_insert() {
+ // Both `absorb` and `absorb_slice` only skip work that per-row
+ // `insert` provably would have skipped too — the compaction
+ // trigger points, coin flips, and final min/max are all preserved.
+ // Three sketches seeded identically and fed the same stream via
+ // the three APIs must agree on every rank probe.
+ use rand::seq::SliceRandom;
+ let seed = 12345u64;
+ let mut shuffle_rng = StdRng::seed_from_u64(seed);
+ let n = 10_000u32;
+ let mut stream: Vec<u32> = (1..=n).collect();
+ stream.shuffle(&mut shuffle_rng);
+
+ let mut via_insert: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ for x in &stream {
+ via_insert.insert(*x);
+ }
+
+ let mut via_absorb: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ via_absorb.absorb(stream.iter().copied());
+
+ let mut via_absorb_slice: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ via_absorb_slice.absorb_slice(&stream);
+
+ // Probe every 100 to keep the test cheap; the space is dense
+ // enough that any drift would surface within these probes.
+ for probe in (1..=n).step_by(100) {
+ let r_insert = via_insert.rank(&probe);
+ let r_absorb = via_absorb.rank(&probe);
+ let r_absorb_slice = via_absorb_slice.rank(&probe);
+ assert_eq!(
+ r_insert, r_absorb,
+ "rank({probe}) diverged: insert={r_insert} absorb={r_absorb}"
+ );
+ assert_eq!(
+ r_insert, r_absorb_slice,
+ "rank({probe}) diverged: insert={r_insert} absorb_slice={r_absorb_slice}"
+ );
+ }
+ assert_eq!(via_insert.quantile(0.0), via_absorb.quantile(0.0));
+ assert_eq!(via_insert.quantile(0.0), via_absorb_slice.quantile(0.0));
+ assert_eq!(via_insert.quantile(0.5), via_absorb.quantile(0.5));
+ assert_eq!(via_insert.quantile(0.5), via_absorb_slice.quantile(0.5));
+ assert_eq!(via_insert.quantile(1.0), via_absorb.quantile(1.0));
+ assert_eq!(via_insert.quantile(1.0), via_absorb_slice.quantile(1.0));
+ }
+
+ #[test]
+ fn absorb_sorted_slice_matches_absorb_slice_on_sorted_input() {
+ // absorb_sorted_slice's trick is to skip the compact-time sort by
+ // keeping level 0 sorted via merge-extend. On sorted input its
+ // output is bit-identical to absorb_slice because sort_unstable on
+ // an already-sorted array of distinct items is the identity, so
+ // both paths produce the same halved subsequence at each
+ // compaction.
+ let seed = 42u64;
+ let n = 10_000u32;
+ let sorted_stream: Vec<u32> = (1..=n).collect();
+
+ let mut via_absorb_slice: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ via_absorb_slice.absorb_slice(&sorted_stream);
+
+ let mut via_sorted: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ via_sorted.absorb_sorted_slice(&sorted_stream);
+
+ for probe in (1..=n).step_by(100) {
+ assert_eq!(
+ via_absorb_slice.rank(&probe),
+ via_sorted.rank(&probe),
+ "rank({probe}) diverged"
+ );
+ }
+ assert_eq!(via_absorb_slice.quantile(0.0), via_sorted.quantile(0.0));
+ assert_eq!(via_absorb_slice.quantile(0.5), via_sorted.quantile(0.5));
+ assert_eq!(via_absorb_slice.quantile(1.0), via_sorted.quantile(1.0));
+ }
+
+ #[test]
+ fn absorb_sorted_slice_restores_level_0_sortedness_after_prior_insert() {
+ // Prior `insert` calls leave sorted[0] = false. absorb_sorted_slice
+ // must sort level 0 before merge-extending, otherwise it would
+ // interleave sorted input into an unsorted buffer and the doc's
+ // "level 0 stays sorted" claim would be a lie.
+ //
+ // Small enough n to avoid tripping level-0 compaction, so we can
+ // inspect the sorted[] flag and level 0 directly.
+ let mut sketch: KllSketch<u32> = KllSketch::with_seed(1024, 42);
+ for x in [7u32, 3, 10, 1, 5] {
+ sketch.insert(x);
+ }
+ assert!(
+ !sketch.sorted[0],
+ "test setup: insert should leave sorted[0] = false"
+ );
+ sketch.absorb_sorted_slice(&[2u32, 4, 6, 8, 9]);
+ assert!(
+ sketch.sorted[0],
+ "absorb_sorted_slice should sort-and-set level 0 before merge-extending"
+ );
+ assert!(
+ sketch.levels[0].is_sorted(),
+ "level 0 contents must actually be sorted: {:?}",
+ sketch.levels[0]
+ );
+ // Sanity on the final counts and extremes across the mixed ingest.
+ assert_eq!(sketch.rank(&11), 10);
+ assert_eq!(sketch.quantile(0.0), Some(&1));
+ assert_eq!(sketch.quantile(1.0), Some(&10));
+ }
+
+ #[test]
+ fn absorb_sorted_slice_falls_back_on_unsorted_input() {
+ // Fed unsorted data, absorb_sorted_slice must not silently corrupt
+ // the sketch — it verifies sortedness and falls through to
+ // absorb_slice. Guards against a future plan-shape change dropping
+ // the upstream SortExec: the sketch would still answer correctly,
+ // just without the sort-skip speedup.
+ use rand::seq::SliceRandom;
+ let seed = 12345u64;
+ let mut shuffle_rng = StdRng::seed_from_u64(seed);
+ let n = 10_000u32;
+ let mut stream: Vec<u32> = (1..=n).collect();
+ stream.shuffle(&mut shuffle_rng);
+
+ let mut via_absorb_slice: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ via_absorb_slice.absorb_slice(&stream);
+
+ let mut via_sorted: KllSketch<u32> = KllSketch::with_seed(200, seed);
+ via_sorted.absorb_sorted_slice(&stream);
+
+ for probe in (1..=n).step_by(100) {
+ assert_eq!(
+ via_absorb_slice.rank(&probe),
+ via_sorted.rank(&probe),
+ "rank({probe}) diverged"
+ );
+ }
+ assert_eq!(via_absorb_slice.quantile(0.0), via_sorted.quantile(0.0));
+ assert_eq!(via_absorb_slice.quantile(0.5), via_sorted.quantile(0.5));
+ assert_eq!(via_absorb_slice.quantile(1.0), via_sorted.quantile(1.0));
+ }
+
+ #[test]
+ fn rank_error_bound_uniform_distribution() {
+ assert_ds_bound_holds("uniform", uniform_stream);
+ }
+
+ #[test]
+ fn rank_error_bound_clustered_distribution() {
+ assert_ds_bound_holds("clustered", clustered_stream);
+ }
+
+ #[test]
+ fn rank_error_bound_heavy_ties() {
+ assert_ds_bound_holds("heavy-ties", tied_stream);
+ }
+
+ #[test]
+ fn merge_matches_single_sketch_within_bound() {
+ // Insert the same shuffled `1..=n` stream two ways:
+ // Path A: straight into one sketch.
+ // Path B: split into 4 shards, one sketch each, merged pairwise.
+ // Both carry O(ε) rank error individually, so any single quantile
+ // probe differs by at most ~2ε in expectation. Bounds any merge
+ // bug that would shift ranks by more than the sum of the two
+ // sketches' inherent errors.
+ use rand::seq::SliceRandom;
+ let k = 200;
+ let n = 10_000u32;
+ let bound = 2.0 * normalized_rank_error_ss(k);
+ let trials = 100u64;
+ let mut failures = 0;
+ for seed in 0..trials {
+ let mut shuffle_rng = StdRng::seed_from_u64(seed);
+ let mut stream: Vec<u32> = (1..=n).collect();
+ stream.shuffle(&mut shuffle_rng);
+
+ let mut single: KllSketch<u32> =
+ KllSketch::with_seed(k, seed.wrapping_add(0x9E37_79B9));
+ for x in &stream {
+ single.insert(*x);
+ }
+
+ let shard_size = n as usize / 4;
+ let mut shards: Vec<KllSketch<u32>> = (0..4u64)
+ .map(|i| KllSketch::with_seed(k, seed.wrapping_add(0xB504_F334 + i)))
+ .collect();
+ for (i, x) in stream.iter().enumerate() {
+ shards[(i / shard_size).min(3)].insert(*x);
+ }
+ let mut merged = shards.remove(0);
+ for s in shards {
+ merged.merge(s);
+ }
+
+ let worst_diff = (1..10)
+ .map(|i| {
+ let q = i as f64 / 10.0;
+ let probe = ((q * n as f64) as u32).max(1);
+ let ra = single.rank(&probe) as f64 / n as f64;
+ let rb = merged.rank(&probe) as f64 / n as f64;
+ (ra - rb).abs()
+ })
+ .fold(0.0_f64, f64::max);
+
+ if worst_diff > bound {
+ failures += 1;
+ }
+ }
+ let fail_rate = failures as f64 / trials as f64;
+ assert!(
+ fail_rate <= 0.05,
+ "merge vs single quantile diff exceeded {bound:.4} in \
+ {failures}/{trials} trials ({:.1}%); expected ≤ 5%",
+ fail_rate * 100.0
+ );
+ }
}
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index c107aaee5..1061e1363 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -56,7 +56,13 @@ tokio = { version = "^1.52", features = [
[dev-dependencies]
criterion = { version = "0.8", features = ["html_reports"] }
+datafusion-functions-aggregate-common = { workspace = true }
+ordered-float = { workspace = true }
[[bench]]
harness = false
name = "sort_shuffle"
+
+[[bench]]
+harness = false
+name = "quantile_sketch"
diff --git a/benchmarks/benches/quantile_sketch.rs b/benchmarks/benches/quantile_sketch.rs
new file mode 100644
index 000000000..28974f3a5
--- /dev/null
+++ b/benchmarks/benches/quantile_sketch.rs
@@ -0,0 +1,348 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Ingest throughput: TDigest (production RuntimeStatsExec path) vs KLL
+//! over `OwnedRow` (multi-column-capable swap path) vs KLL over
+//! `OrderedFloat<f64>` (single-column-native-Ord fast path).
+//!
+//! The three variants isolate the cost of the arrow-row encoding: TDigest
+//! is the incumbent baseline, KLL<OwnedRow> pays for `RowConverter` +
+//! per-row `owned()` heap alloc, KLL<OrderedFloat<f64>> skips both and
+//! measures the sketch algorithm on its own.
+//!
+//! Merge and quantile query are excluded: for N=1M rows at P=64 partitions
+//! and K=64 cuts, both are 3+ orders of magnitude cheaper than ingest and
+//! would not move the swap decision.
+//!
+//! Sketch sizing is empirical, not analytical. TDigest is held at
+//! `TDIGEST_MAX_SIZE=100` (production), and `KLL_K` is picked so KLL's
+//! worst-case normalized rank error on the same uniform 1M stream is
+//! within a whisker of TDigest's — a fair race. The parity check itself
+//! prints alongside the throughput results when the env var
+//! `KLL_PARITY_CHECK=1` is set, so anyone can rerun it.
+
+use std::sync::Arc;
+
+use ballista_core::kll::KllSketch;
+use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main};
+use datafusion::arrow::array::{ArrayRef, Float64Array};
+use datafusion::arrow::datatypes::DataType;
+use datafusion::arrow::row::{OwnedRow, RowConverter, SortField};
+use datafusion_functions_aggregate_common::tdigest::TDigest;
+use ordered_float::OrderedFloat;
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+
+/// Matches `runtime_stats::TDIGEST_MAX_SIZE` — production sizing.
+const TDIGEST_MAX_SIZE: usize = 100;
+
+/// KLL top-level compactor capacity picked empirically for worst-case
+/// rank-error parity with TDigest at `TDIGEST_MAX_SIZE=100`. Measured on
+/// the uniform 1M stream this bench feeds:
+///
+/// | sketch | worst rank err (9 deciles) |
+/// |------------------------------|---------------------------:|
+/// | TDigest max_size=100 | 0.0021 |
+/// | KLL k=800 | 0.0016 |
+///
+/// Rerun via `KLL_PARITY_CHECK=1 cargo bench --bench quantile_sketch`.
+const KLL_K: usize = 800;
+
+/// RuntimeStatsExec observes one batch at a time from the upstream
+/// operator. 8192 is DataFusion's default target batch size.
+const BATCH_SIZE: usize = 8192;
+
+/// Fixed PRNG seed so the harness is deterministic across runs.
+const SEED: u64 = 0xDEC0DE;
+
+/// Total row counts benched. 100K represents a small stage; 1M represents
+/// a per-partition slice of a large TPC-H stage.
+const ROW_COUNTS: &[usize] = &[100_000, 1_000_000];
+
+/// Build a `Vec<Float64Array>` of `BATCH_SIZE`-length batches whose row
+/// count sums to `n`. Uniform on [0, 1), seeded.
+fn build_batches(n: usize) -> Vec<Float64Array> {
+ let mut rng = StdRng::seed_from_u64(SEED);
+ let mut batches = Vec::with_capacity(n.div_ceil(BATCH_SIZE));
+ let mut remaining = n;
+ while remaining > 0 {
+ let take = remaining.min(BATCH_SIZE);
+ let vals: Vec<f64> = (0..take).map(|_| rng.random::<f64>()).collect();
+ batches.push(Float64Array::from(vals));
+ remaining -= take;
+ }
+ batches
+}
+
+/// Same underlying stream as `build_batches`, sorted globally before
+/// batching. Consecutive batches are monotonically increasing — matches
+/// what a `SortExec` upstream of `RuntimeStatsExec` would feed the sketch
+/// in the ORRE plan shape.
+fn build_globally_sorted_batches(n: usize) -> Vec<Float64Array> {
+ let mut rng = StdRng::seed_from_u64(SEED);
+ let mut all: Vec<f64> = (0..n).map(|_| rng.random::<f64>()).collect();
+ all.sort_by(f64::total_cmp);
+ all.chunks(BATCH_SIZE)
+ .map(|chunk| Float64Array::from(chunk.to_vec()))
+ .collect()
+}
+
+/// Production TDigest ingest path — one `Vec<f64>` per batch through
+/// `merge_unsorted_f64`. This mirrors `runtime_stats::StreamState::record`
+/// modulo the null-flatten and partition-slot lookup.
+fn ingest_tdigest(batches: &[Float64Array]) -> TDigest {
+ let mut sketch = TDigest::new(TDIGEST_MAX_SIZE);
+ for arr in batches {
+ // Uniform-f64 batches have no nulls, so `values()` matches the
+ // shape the sketch consumes after RuntimeStatsExec's flatten step.
+ let values: Vec<f64> = arr.values().to_vec();
+ sketch = sketch.merge_unsorted_f64(values);
+ }
+ sketch
+}
+
+/// Sorted-input TDigest ingest — same shape as `ingest_tdigest` but hits
+/// `merge_sorted_f64` on each pre-sorted batch. Answers the apples-to-
+/// apples question against `absorb_sorted_slice` on the ORRE-plan-shape
+/// arm: TDigest also benefits from sorted input (it skips its own
+/// pre-sort), so a matched race lines them up on equal footing.
+fn ingest_tdigest_sorted(batches: &[Float64Array]) -> TDigest {
+ let mut sketch = TDigest::new(TDIGEST_MAX_SIZE);
+ for arr in batches {
+ sketch = sketch.merge_sorted_f64(arr.values());
+ }
+ sketch
+}
+
+/// Multi-column-capable KLL ingest path — encode each batch through the
+/// arrow row converter, then push per-row `OwnedRow` into a
+/// `KllSketch<OwnedRow>`. The heap alloc per `owned()` is what the
+/// level-0 borrowed-row optimization would eliminate.
+fn ingest_kll_row(
+ batches: &[Float64Array],
+ converter: &RowConverter,
+) -> KllSketch<OwnedRow> {
+ let mut sketch = KllSketch::<OwnedRow>::new(KLL_K);
+ for arr in batches {
+ let col: ArrayRef = Arc::new(arr.clone());
+ let rows = converter.convert_columns(&[col]).unwrap();
+ for row in rows.iter() {
+ sketch.insert(row.owned());
+ }
+ }
+ sketch
+}
+
+/// Single-column native-`Ord` KLL fast path — bypasses arrow-row entirely
+/// by using `OrderedFloat<f64>` as the sketch item type. Answers "what
+/// does KLL cost when we dispatch by column type instead of routing
+/// everything through OwnedRow?"
+fn ingest_kll_ordered_float(batches: &[Float64Array]) -> KllSketch<OrderedFloat<f64>> {
+ let mut sketch = KllSketch::<OrderedFloat<f64>>::new(KLL_K);
+ for arr in batches {
+ for v in arr.values() {
+ sketch.insert(OrderedFloat(*v));
+ }
+ }
+ sketch
+}
+
+/// Same as `ingest_kll_ordered_float` but batches each Arrow batch into
+/// `absorb_slice` — measures the batch-oriented ingest optimization
+/// (amortized compact_all + batch min/max + memcpy-style extend).
+fn ingest_kll_ordered_float_absorb_slice(
+ batches: &[Float64Array],
+) -> KllSketch<OrderedFloat<f64>> {
+ let mut sketch = KllSketch::<OrderedFloat<f64>>::new(KLL_K);
+ for arr in batches {
+ let vals: Vec<OrderedFloat<f64>> =
+ arr.values().iter().copied().map(OrderedFloat).collect();
+ sketch.absorb_slice(&vals);
+ }
+ sketch
+}
+
+/// Same shape as `ingest_kll_ordered_float_absorb_slice`, but consumes
+/// pre-sorted batches via `absorb_sorted_slice` — measures the sketch
+/// side of the ORRE-plan-shape follow-up, where a `SortExec` upstream of
+/// `RuntimeStatsExec` guarantees sorted input and lets the sketch skip
+/// every compaction sort (not just those above level 0).
+fn ingest_kll_ordered_float_absorb_sorted_slice(
+ batches: &[Float64Array],
+) -> KllSketch<OrderedFloat<f64>> {
+ let mut sketch = KllSketch::<OrderedFloat<f64>>::new(KLL_K);
+ for arr in batches {
+ let vals: Vec<OrderedFloat<f64>> =
+ arr.values().iter().copied().map(OrderedFloat).collect();
+ sketch.absorb_sorted_slice(&vals);
+ }
+ sketch
+}
+
+/// Same as `ingest_kll_row` but batches each Arrow batch into `absorb`
+/// (owned-iter variant) — measures the compact_all amortization win on the
+/// `OwnedRow` path without the extra clone that `absorb_slice` would
+/// force.
+fn ingest_kll_row_absorb(
+ batches: &[Float64Array],
+ converter: &RowConverter,
+) -> KllSketch<OwnedRow> {
+ let mut sketch = KllSketch::<OwnedRow>::new(KLL_K);
+ for arr in batches {
+ let col: ArrayRef = Arc::new(arr.clone());
+ let rows = converter.convert_columns(&[col]).unwrap();
+ sketch.absorb(rows.iter().map(|r| r.owned()));
+ }
+ sketch
+}
+
+/// Print worst-case quantile-inversion error at deciles for TDigest at
+/// `TDIGEST_MAX_SIZE=100` and KLL at a sweep of `k`. Called from
+/// `bench_ingest` when `KLL_PARITY_CHECK=1` is set. Reproducibility
+/// artifact for the `KLL_K` choice — a reviewer can rerun and confirm
+/// the two sketches sit at matched accuracy at the settings the bench
+/// uses.
+///
+/// Error metric: at each decile `q`, ask the sketch for `quantile(q)`,
+/// look up what fraction of the true stream is below the returned value,
+/// and take `|true_fraction - q|`. That's normalized rank error via
+/// quantile inversion — the same units for both sketches even though
+/// TDigest lacks a public rank API.
+fn print_parity_table(n: usize) {
+ let batches = build_batches(n);
+ let all: Vec<f64> = batches
+ .iter()
+ .flat_map(|b| b.values().iter().copied())
+ .collect();
+ let mut sorted = all.clone();
+ sorted.sort_by(f64::total_cmp);
+
+ let true_rank_frac =
+ |probe: f64| -> f64 { sorted.partition_point(|x| *x < probe) as f64 / n as f64 };
+ // Nine interior deciles; hits TDigest's median hump (worst-case for it)
+ // and both tails (KLL uniform, TDigest tightest).
+ let qs: Vec<f64> = (1..10).map(|i| i as f64 / 10.0).collect();
+
+ let td = ingest_tdigest(&batches);
+ let td_worst = qs
+ .iter()
+ .map(|&q| (true_rank_frac(td.estimate_quantile(q)) - q).abs())
+ .fold(0.0_f64, f64::max);
+
+ println!("\n=== quantile_sketch parity @ n={n} (uniform f64) ===");
+ println!("TDigest max_size={TDIGEST_MAX_SIZE:>4} → worst rank err = {td_worst:.4}");
+ for &k in &[50usize, 100, 200, 400, 800] {
+ let mut sk = KllSketch::<OrderedFloat<f64>>::new(k);
+ for arr in &batches {
+ for v in arr.values() {
+ sk.insert(OrderedFloat(*v));
+ }
+ }
+ let kll_worst = qs
+ .iter()
+ .map(|&q| {
+ let guess = sk.quantile(q).map(|of| of.0).unwrap_or(f64::NAN);
+ (true_rank_frac(guess) - q).abs()
+ })
+ .fold(0.0_f64, f64::max);
+ let marker = if k == KLL_K { " ← KLL_K" } else { "" };
+ println!("KLL k ={k:>4} → worst rank err = {kll_worst:.4}{marker}");
+ }
+ println!();
+}
+
+fn bench_ingest(c: &mut Criterion) {
+ if std::env::var_os("KLL_PARITY_CHECK").is_some() {
+ for &n in ROW_COUNTS {
+ print_parity_table(n);
+ }
+ }
+
+ let mut group = c.benchmark_group("runtime_stats_ingest");
+ // Ingest is CPU-bound and single-threaded; ten samples is enough for
+ // the ratios we care about and keeps the harness under a minute.
+ group.sample_size(10);
+
+ let converter = RowConverter::new(vec![SortField::new(DataType::Float64)]).unwrap();
+
+ for &n in ROW_COUNTS {
+ let batches = build_batches(n);
+ group.throughput(Throughput::Elements(n as u64));
+
+ group.bench_function(format!("tdigest/{n}"), |b| {
+ b.iter_batched(
+ || batches.clone(),
+ |bs| ingest_tdigest(&bs),
+ BatchSize::LargeInput,
+ );
+ });
+
+ group.bench_function(format!("kll_row/{n}"), |b| {
+ b.iter_batched(
+ || batches.clone(),
+ |bs| ingest_kll_row(&bs, &converter),
+ BatchSize::LargeInput,
+ );
+ });
+
+ group.bench_function(format!("kll_row_absorb/{n}"), |b| {
+ b.iter_batched(
+ || batches.clone(),
+ |bs| ingest_kll_row_absorb(&bs, &converter),
+ BatchSize::LargeInput,
+ );
+ });
+
+ group.bench_function(format!("kll_ordered_float/{n}"), |b| {
+ b.iter_batched(
+ || batches.clone(),
+ |bs| ingest_kll_ordered_float(&bs),
+ BatchSize::LargeInput,
+ );
+ });
+
+ group.bench_function(format!("kll_ordered_float_absorb_slice/{n}"), |b| {
+ b.iter_batched(
+ || batches.clone(),
+ |bs| ingest_kll_ordered_float_absorb_slice(&bs),
+ BatchSize::LargeInput,
+ );
+ });
+
+ let sorted_batches = build_globally_sorted_batches(n);
+ group.bench_function(format!("tdigest_sorted/{n}"), |b| {
+ b.iter_batched(
+ || sorted_batches.clone(),
+ |bs| ingest_tdigest_sorted(&bs),
+ BatchSize::LargeInput,
+ );
+ });
+
+ group.bench_function(format!("kll_ordered_float_absorb_sorted_slice/{n}"), |b| {
+ b.iter_batched(
+ || sorted_batches.clone(),
+ |bs| ingest_kll_ordered_float_absorb_sorted_slice(&bs),
+ BatchSize::LargeInput,
+ );
+ });
+ }
+
+ group.finish();
+}
+
+criterion_group!(benches, bench_ingest);
+criterion_main!(benches);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org