feat: support approx_percentile / percentile_approx aggregate by andygrove · Pull Request #4801 · apache/datafusion-comet

andygrove · 2026-07-02T21:24:48Z

Which issue does this PR close?

There is no dedicated tracking issue for the feature itself. Prior to this PR the compatibility guide explicitly listed approx_percentile / percentile_approx as not planned; this PR implements them and removes that note.

This PR also fixes #4813, a pre-existing mixed-execution bug that approx_percentile surfaced (see below).

Rationale for this change

approx_percentile / percentile_approx is a common Spark aggregate that had no Comet support. It was previously considered out of scope on the assumption that its approximate results could not be made bit-identical to Spark, because Spark uses the Greenwald-Khanna (GK) QuantileSummaries algorithm while DataFusion's built-in approx_percentile_cont uses t-digest, which produces different values.

This PR takes the faithful-port approach instead: it reimplements Spark's GK QuantileSummaries in Rust. Because Comet feeds values to the accumulator in the same intra-partition order as Spark and the GK arithmetic/traversal is reproduced exactly, results are bit-identical to Spark for deterministic plans, and otherwise stay within Spark's own accuracy guarantee (Spark's own result already varies across partitionings for this aggregate). The expression is therefore marked Compatible, with no need for spark.comet.expr.allowIncompatible.

What changes are included in this PR?

Native, all-native (partial and final) implementation wired through the standard aggregate layers:

Native GK core (native/spark-expr/src/agg_funcs/quantile_summaries.rs): a bit-for-bit port of Spark's QuantileSummaries (insert with the 50000-value head-buffer flush, backward compress traversal that preserves the minimum, the non-commutative merge with asymmetric delta adjustment and tie-break, and the ceil-rank query with the ascending multi-percentile sweep). Serialization uses a Comet-internal little-endian layout.
DataFusion UDAF + accumulator (native/spark-expr/src/agg_funcs/approx_percentile.rs): wraps one summary per group, carries the digest across partial/final as a single Binary state value, and casts each result back to the input type (exact, since GK returns an actual inserted value).
Protobuf + planner (native/proto/src/proto/expr.proto, native/core/src/execution/planner.rs): a new ApproxPercentile message and planner arm.
Scala serde (spark/src/main/scala/org/apache/comet/serde/aggregates.scala, QueryPlanSerde.scala): CometApproxPercentile, marked Compatible for byte/short/int/long/float/double, supporting the scalar, array-of-percentiles, and explicit-accuracy signatures. A single percentage maps to a scalar result; an array maps to an array result.
Docs: removes the "not planned" note and adds approx_percentile to the aggregate function tables.

Scope: numeric input types byte/short/int/long/float/double. Decimal, date, timestamp, TimestampNTZ, and interval inputs are left for a follow-on PR.

Mixed-execution fix (#4813). approx_percentile surfaced a pre-existing bug: in a distinct-aggregate rewrite the non-distinct aggregate's Partial is separated from its Final by intermediate PartialMerge stages, and the existing #1389 guards only covered a direct Partial → Final pair. So an aggregate with an incompatible intermediate buffer could run part of the chain in Comet and part in Spark, handing a Comet-encoded buffer to a Spark aggregate (or vice versa) — a hard crash for percentile_approx / collect_set, and a silent-wrong-result risk for others. Fixed with two dual guards:

CometExecRule: findPartialAggInPlan now walks through the intermediate PartialMerge stages to tag the bottom Partial, so the whole chain falls back to Spark.
CometBaseAggregate.doConvert: the sparkFinalMode guard is generalized so a Comet PartialMerge that merges buffers (not just a Final) requires a Comet producer below it, covering the Spark-Partial → Comet-Merge direction.

How are these changes tested?

Rust unit tests on QuantileSummaries (query returns actual inserted values, results within the relative-error bound, multi-percentile matches single, merge within bound, serialization round-trip) and on the accumulator (scalar/array output, empty input yields null, split-then-merge matches single-shot).
End-to-end SQL file tests (spark/src/test/resources/sql-tests/expressions/aggregate/approx_percentile.sql) covering scalar, array of percentiles, non-default accuracy, group-by, and byte/short/int/long/float/double inputs, plus nulls-ignored and empty-input-yields-null. These run each query in Spark with and without Comet and diff the results, and assert native execution (a fallback fails the test), so they are the authoritative Spark-compatibility check. Passing on the Spark 3.5 and 4.1 profiles.
Mixed-execution regression tests for Distinct-aggregate rewrite can split an incompatible-buffer aggregate across Comet and Spark, causing crash or wrong results #4813: CometExecRuleSuite asserts the whole distinct-rewrite chain falls back to Spark when either the partial or the final is forced non-native (via the comet.exec.*HashAggregate.enabled debug configs), and CometAggregateSuite runs percentile_approx + count(DISTINCT ...) across the partial/final split matrix and checks results match Spark.

Benchmark results

Added approx_percentile cases to CometAggregateExpressionBenchmark. Because
approx_percentile is a TypedImperativeAggregate (planned as
ObjectHashAggregate), which only runs natively when Comet shuffle is enabled,
the shared CometBenchmarkBase now sets the Comet shuffle manager at startup so
the benchmark measures native execution.

Local run, single node (local[1]), 1,048,576 rows, GROUP BY ~1000 groups
unless noted. Best time, lower is better.

Case	Spark (ms)	Comet (ms)	Speedup
`approx_percentile(int, 0.5)`	347	80	4.4x
`approx_percentile(long, 0.5)`	400	96	4.2x
`approx_percentile(double, 0.5)`	332	73	4.6x
`approx_percentile(double, 0.9)`	384	76	5.1x
`approx_percentile(double, array(0.25,0.5,0.75))`	354	74	4.8x
`approx_percentile(double, 0.5, 100)`	316	60	5.2x
`approx_percentile(double, 0.5)` (no group by)	239	31	7.7x
`approx_percentile(double, 0.5)` (high cardinality)	470	305	1.5x

Wrap the QuantileSummaries (Greenwald-Khanna) sketch in a DataFusion AggregateUDFImpl and per-row Accumulator for Spark's approx_percentile / percentile_approx. State is carried as a single Binary ScalarValue using the Comet-internal little-endian digest format. Supports byte, short, int, long, float, and double inputs, casting the query result back to the original input type. groups_accumulator_supported is false; this runs as a row accumulator per group.

Add a SQL file test that runs approx_percentile scalar, array, group-by, accuracy, doubles, floats, null, and empty-input forms through both Spark and Comet natively, then update the compatibility docs to reflect that approx_percentile now runs natively for the six supported numeric types.

…ounting Add heap_size() to QuantileSummaries so the accumulator's size() reflects actual sampled/head buffer memory instead of only the struct footprint. Document the intentional i64-vs-toInt delta deviation from Spark. Extend the SQL file test with byte and short input coverage, a non-default accuracy case, and fix a comment describing range().id as int.

- Use SingleRowListArrayBuilder for the array output branch. - Inline the trivial output_element_type wrapper. - Fast path in update_batch when the batch has no nulls. - Honor the compressed flag in QuantileSummaries::compress to avoid redundant recompression, matching Spark's compress-once semantics. - Move the peer digest into an empty accumulator instead of cloning. - Pre-size the merged sampled vector. - Keep the QuantileSummaries GK helper private (mirror welford).

ObjectHashAggregate-based aggregates (e.g. approx_percentile) only run natively when Comet shuffle is enabled, which requires the Comet shuffle manager. Set it on the SparkConf at startup so all benchmarks using this base measure native execution instead of falling back to Spark.

…t and Spark in distinct rewrite A distinct-aggregate rewrite separates a non-distinct aggregate's Partial from its Final by intermediate PartialMerge stages. The existing mixed- execution guards (apache#1389) only covered a direct Partial -> Final pair, so an aggregate with an incompatible intermediate buffer (e.g. percentile_approx) could run part of the chain in Comet and part in Spark, handing a Comet- encoded buffer to a Spark aggregate (or vice versa) and crashing. - CometExecRule: walk findPartialAggInPlan through intermediate PartialMerge stages to tag the bottom Partial so the whole chain falls back to Spark. - CometBaseAggregate.doConvert: generalize the sparkFinalMode guard so a Comet PartialMerge that merges buffers (not just a Final) requires a Comet producer below it, covering the Spark-Partial -> Comet-Merge direction. Fixes the ObjectHashAggregateSuite "[typed, with distinct]" crash. See apache#4813.

…eason

mbutrovich

Thanks for the PR, @andygrove! First round of feedback, mostly idiomatic Rust, performance, test stuff, so all minor. Correctness looks good.

mbutrovich · 2026-07-09T14:45:02Z

+            {
+                0
+            } else {
+                // Spark uses `.toInt` here (unlike `.toLong` in merge/compress); we use i64


The comment says Spark uses .toInt here and diverges from Comet past Int.MAX rows. In the Spark source I checked (QuantileSummaries.scala:112, 4.0.0) this line is math.floor(2 * relativeError * currentCount).toLong, so Spark uses .toLong too and the divergence does not exist. The i64 choice is correct, but the comment reads as if the two engines differ when they match. Worth correcting or dropping. Could you confirm 3.4 / 3.5 / 4.1 also use .toLong (this file rarely changes)?

Confirmed: 3.4.3 / 3.5.8 / 4.0.1 / 4.1.0-preview2 all use .toLong on this line, matching our i64, so there is no divergence. Replaced the comment with that fact in cc1ab34.

mbutrovich · 2026-07-09T14:45:18Z

+// specific language governing permissions and limitations
+// under the License.
+
+//! A bit-for-bit port of Spark's `QuantileSummaries` (Greenwald-Khanna), the


The module header calls this a bit-for-bit port, which discourages any change. It might help to state instead what is actually load-bearing for bit-identity: the 50000-value flush boundary, the integer arithmetic (floor(2 * relative_error * count) as i64, the /2 for target_error, ceil(percentile * count)), the backward compress traversal with the < merge_threshold test, the merge interleave and delta adjustment, and the query ceil-rank sweep. The data-structure choices (deque vs vec, clone vs copy, immutable rebuild vs in-place) do not affect results and are free to be optimized. That contract lets the code be idiomatic without risking correctness.

Rewrote the module header to spell out what is load-bearing for bit-identity (the 50000-value flush boundary, the integer arithmetic, the backward compress traversal, the merge interleave and delta adjustment, and the query ceil-rank sweep) and to state that the data-structure choices are free to be optimized.

mbutrovich · 2026-07-09T14:45:26Z

+
+/// A single sampled statistic: the value, its minimum rank jump `g`, and the
+/// maximum span of the rank `delta`.
+#[derive(Debug, Clone, PartialEq)]


Stats is f64 + i64 + i64, all Copy, so deriving Copy would let you drop the roughly eight .clone() calls in with_head_buffer_inserted, compress_immut, and merge. No result change, more idiomatic.

Derived Copy on Stats and dropped the clones in with_head_buffer_inserted, compress_immut, and merge.

mbutrovich · 2026-07-09T14:45:40Z

+        // Spark relies on `Array[Double].sorted`; use total ordering so the port
+        // is deterministic even in the presence of NaN (Spark's typical inputs
+        // are NaN-free).
+        sorted.sort_by(|a, b| a.total_cmp(b));


sort_unstable_by(|a, b| a.total_cmp(b)) would be faster with identical output here. Stability is irrelevant because the sort key is the value itself, so equal keys are bit-equal elements, and total_cmp orders -0.0 / 0.0 / NaN deterministically regardless of stability.

Separately: since float and double are marked Compatible with no caveat, is the NaN and -0.0 ordering worth a test? total_cmp and Scala's Array[Double].sorted can differ across Scala versions, and Comet still cross-builds 2.12.

Switched to sort_unstable_by. On the NaN / -0.0 question: I added a signed_zero_ordering_is_deterministic unit test as a regression anchor. NaN and -0.0 in percentile inputs are pathological, and I did not want to overclaim cross-Scala-version bit-identity for NaN ordering, so I left the Compatible marking as-is. Happy to take a stronger stance if you think it warrants one.

mbutrovich · 2026-07-09T14:45:55Z

+        // are NaN-free).
+        sorted.sort_by(|a, b| a.total_cmp(b));
+
+        let mut new_samples: Vec<Stats> = Vec::new();


new_samples grows to sampled.len() + sorted.len(). Vec::with_capacity(...) here removes the reallocation on every flush. Same idea for reserving head_sampled in update_batch from the input array length.

Preallocated new_samples with Vec::with_capacity(sampled.len() + sorted.len()), and added a reserve method on the summary that update_batch calls with the batch's non-null count.

mbutrovich · 2026-07-09T14:46:18Z

+        res.into()
+    }
+
+    pub fn merge(&self, other: &QuantileSummaries) -> QuantileSummaries {


Not for this PR. merge takes &self / &other and allocates a fresh summary, so merge_batch reallocates on every incoming digest, and with_head_buffer_inserted rebuilds the whole sampled vector on every flush. A double-buffer for the flush and a consuming merge(self, other: &Self) -> Self (or merge_into(&mut self, ...)) would cut allocations if a future profile shows this path is hot. Noting it here as a possible follow-up, not a change to make now.

Agreed this is out of scope for this PR. Noted as a follow-up; I will track the consuming-merge and double-buffer optimization separately rather than expand this PR.

Filed as #4874 to track the consuming-merge and flush double-buffer optimization.

mbutrovich · 2026-07-09T14:46:31Z

+            DataType::Int32 => ScalarValue::Int32(Some(d as i32)),
+            DataType::Int64 => ScalarValue::Int64(Some(d as i64)),
+            DataType::Float32 => ScalarValue::Float32(Some(d as f32)),
+            _ => ScalarValue::Float64(Some(d)),


The _ => Float64 arm silently absorbs any type outside the five explicit ones. Since the serde restricts inputs to six types, making Float64 explicit and using debug_assert! / unreachable! for the rest would enforce the serde contract instead of masking a future mismatch.

Made the Float64 arm explicit and changed the fallthrough to unreachable! with the input type, so a future serde mismatch fails loudly instead of being silently coerced to double.

mbutrovich · 2026-07-09T14:46:41Z

+                QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD,
+                digests.value(i),
+            );
+            if self.summary.count() == 0 {


The count() == 0 branch duplicates logic already inside merge (which handles count == 0). The only gain is avoiding one clone of the already-owned peer. If that is intentional for large digests a one-line note would help, otherwise self.summary = self.summary.merge(&peer) is simpler.

Kept the branch: peer is already owned (built by from_bytes), so this moves it in and avoids merge cloning a potentially large digest. Clarified the comment to say exactly that.

mbutrovich · 2026-07-09T14:46:49Z

+
+    fn evaluate(&mut self) -> Result<ScalarValue> {
+        self.summary.compress();
+        match self.summary.query(&self.percentiles) {


Does an empty percentage array reach here, or does Spark's analyzer reject it first? Spark returns null for empty percentages, but query(&[]) returns Some(vec![]) and would try to build an empty list. If it is unreachable a short comment would help, otherwise a percentiles.is_empty() short-circuit to null matches Spark.

Good catch, this was a real bug. Spark's ApproximatePercentile.eval returns null whenever result.length == 0, so an empty percentage array yields null even for the array signature. evaluate now routes both the no-rows case and the empty-percentiles case through a shared null_result. Added empty_percentiles_is_null and array_output_empty_input_is_null Rust tests plus a SQL case.

mbutrovich · 2026-07-09T14:47:08Z

+
+-- scalar percentile over a bigint (long) column
+query
+SELECT approx_percentile(id, 0.5) FROM range(1000)


Coverage is solid. A few paths look untested and would be quick to add:

int input (cast(id AS int)), the most common case (long, double, float, byte, short are covered).

Extreme percentiles 0.0 and 1.0, which hit the percentile <= relative_error and percentile >= 1 - relative_error short-circuits in query that nothing else exercises.

Negative values (every case uses range(), which is non-negative).

A low-cardinality / duplicate-heavy column (id % 5), which exercises the g accumulation in compress_immut.

Array output with empty input (the return_array empty path in evaluate is untested at both SQL and Rust level).

Added all five: int input, extreme 0.0 / 1.0 percentiles, negative values (cast(id AS int) - 500), a duplicate-heavy id % 5 column, and array output over empty input. The Rust suite also gained extreme_percentiles_hit_short_circuits, negative_values_are_ordered, duplicate_heavy_column_accumulates_g, and array_output_empty_input_is_null.

Return null for an empty percentage array to match Spark, which returns null whenever the percentile result is empty. Previously the array path built an empty list instead. Also apply reviewer suggestions: - Correct the delta comment: Spark uses `.toLong` (verified in 3.4/3.5/ 4.0/4.1), matching our i64, so there is no divergence. - Document the invariants that are load-bearing for bit-identity and note the data-structure choices are free to be optimized. - Derive `Copy` on `Stats` and drop the redundant clones. - Use `sort_unstable_by`, preallocate the flush buffer, and reserve the head buffer from the input batch length. - Build the compressed samples forward into a `Vec` and reverse once instead of prepending into a `VecDeque`. - Make the `Float64` cast arm explicit and `unreachable!` on any other input type to enforce the serde contract. - Add coverage for int input, extreme percentiles, negative values, duplicate-heavy columns, signed-zero ordering, and empty array output.

mbutrovich

LGTM, thanks for the quick revision @andygrove!

…clippy

The same clippy `useless_borrows_in_formatting` lint that fs-hdfs tripped also fires on the `eval_mode` argument in `Cast`'s `Display` impl. Clippy reports it only after the fs-hdfs fix lets it proceed to `spark-expr`, so both must be fixed together to get `rust-test` green.

andygrove added 8 commits July 2, 2026 14:36

feat: add QuantileSummaries GK port for approx_percentile

c9960aa

feat: wire ApproxPercentile through proto and native planner

de80d43

feat: add CometApproxPercentile serde and register it

35b5f28

test: add approx_percentile microbenchmarks

b47fbab

andygrove marked this pull request as ready for review July 2, 2026 21:40

andygrove added 2 commits July 2, 2026 16:24

style: apply cargo fmt to approx_percentile

d1ed88f

andygrove mentioned this pull request Jul 3, 2026

Distinct-aggregate rewrite can split an incompatible-buffer aggregate across Comet and Spark, causing crash or wrong results #4813

Closed

andygrove added this to the 1.0.0 milestone Jul 3, 2026

test: regenerate TPC-DS golden files for updated aggregate fallback r…

edb7679

…eason

andygrove mentioned this pull request Jul 6, 2026

Comet 1.0.0 Release #4082

Open

27 tasks

andygrove requested review from comphead, mbutrovich and parthchandra July 9, 2026 14:15

Merge remote-tracking branch 'apache/main' into feat/approx-percentile

10f40be

mbutrovich requested changes Jul 9, 2026

View reviewed changes

mbutrovich approved these changes Jul 9, 2026

View reviewed changes

andygrove mentioned this pull request Jul 9, 2026

Reduce allocations in approx_percentile QuantileSummaries merge and flush paths #4874

Open

andygrove added 2 commits July 9, 2026 10:18

fix: remove redundant reference in fs-hdfs info! argument to satisfy …

5023d00

…clippy

andygrove merged commit eb5b761 into apache:main Jul 9, 2026
135 of 136 checks passed

andygrove deleted the feat/approx-percentile branch July 9, 2026 18:47

Uh oh!

Conversation

andygrove commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Benchmark results

Uh oh!

mbutrovich left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mbutrovich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

andygrove commented Jul 2, 2026 •

edited

Loading

mbutrovich left a comment •

edited

Loading