Skip to content

perf: hll_sketch_agg merge builds a new union for every merged row #7341

Open

Description

@comphead

What is the problem the feature request solves?

The merge side of hll_sketch_agg gets much slower than Spark's once a group's sketch reaches the dense register array. On a Spark 4.0 job with 17 hll_sketch_agg(CASE WHEN ... THEN id END, 18) columns and 7 grouping keys, the final aggregation merges 454M partial rows into 53M groups over 676 tasks:

Spark Comet, DataFusion 55.2 Comet, DataFusion main
Task time, median 8.9 s 40.4 s 18.1 s
Aggregate time, all tasks 2,229 s 26,608 s 11,732 s

Spark's aggregate is ObjectHashAggregate (time in aggregation build, sort fallback in every task), Comet's is CometHashAggregate.

HllSketchAccumulator::merge_batch calls SparkHllSketch::merge_sketch for every row. merge_sketch creates a new SparkHllUnion, merges the accumulator into it, merges the peer, and converts the union back to a sketch. Once the accumulator is in HLL mode, merging it into the empty union goes through copy_or_downsample in datasketches 0.3.0, which allocates a new 2^lgConfigK register array, takes the register-wise max, and runs Array8::rebuild_cached_values, one f64 division per register. to_sketch then clones the array. So every merged row costs O(2^lgConfigK), 262,144 registers at lgConfigK 18, and 96% of the samples in a profile of the merge are in copy_or_downsample.

Spark's HllSketchAgg.merge also builds a union per row, but datasketches-java copies the array (copyAs), defers the rebuild to getResult behind a flag, and computes each 2^-v from the exponent bits (invPow2) instead of dividing.

Merging 2,000 partial sketches of 200 values each into one accumulator at lgConfigK 18 (release build, Apple M-series laptop):

Merge path Per merged row Estimate (true 400,000)
Comet merge_sketch, today 418 to 450 µs 399,354
Spark HllSketchAgg.merge, datasketches-java 6.1.1 12.7 to 12.9 µs 399,437
One SparkHllUnion kept across merges 2.6 to 2.8 µs 399,441
Snippet
use datafusion_comet_spark_expr::{estimate_from_bytes, SparkHllSketch, SparkHllUnion};
use std::time::Instant;

fn main() {
    let peers: Vec<Vec<u8>> = (0..2000u64)
        .map(|p| {
            let mut s = SparkHllSketch::new(18);
            for v in 0..200 {
                s.update_i64((p * 1_000_000 + v) as i64);
            }
            s.to_sketch_bytes()
        })
        .collect();

    // What HllSketchAccumulator::merge_batch does today.
    let t = Instant::now();
    let mut acc = SparkHllSketch::new(18);
    for b in &peers {
        acc.merge_sketch(&SparkHllSketch::from_bytes(b).unwrap());
    }
    let today = acc.to_sketch_bytes();
    println!("merge_sketch: {:.1} us/row", t.elapsed().as_secs_f64() * 1e6 / 2000.0);

    // One union kept across merges, as HllUnionAccumulator does.
    let t = Instant::now();
    let mut u = SparkHllUnion::new(18);
    for b in &peers {
        u.merge(&SparkHllSketch::from_bytes(b).unwrap());
    }
    let union = u.to_sketch_bytes();
    println!("persistent union: {:.1} us/row", t.elapsed().as_secs_f64() * 1e6 / 2000.0);
    println!("{} {}", estimate_from_bytes(&today).unwrap(), estimate_from_bytes(&union).unwrap());
}

A Final-mode AggregateExec over partial states, interleaved the way a reducer reads map outputs (about one row per group per batch), shows the same split. Small sketches are already on par with Spark, and the cost is in groups that turn dense:

Partial states, lgConfigK 18 Comet today Persistent union GroupsAccumulator, one union per group Spark path
1.7M rows, 200K groups, 0 to 8 values each 291 ns/row 225 ns 137 ns 224 to 326 ns
200K rows, 2K groups, 50 values each (SET mode) 2.2 µs 1.6 µs 1.9 µs 2.2 to 2.5 µs
8K rows, 20 groups, 200 values each (dense) 198 µs 3.1 µs 3.3 µs 13.0 to 13.5 µs

The persistent union gives the same estimates as today in LIST and SET mode, and estimates as close to the true count in HLL mode.

Describe the potential solution

  1. Keep one SparkHllUnion in HllSketchAccumulator for merge_batch, as HllUnionAccumulator already does, and build the sketch only in state() and evaluate(). If one accumulator sees both update_batch and merge_batch, fold the updated sketch into the union there. The intermediate stays Comet-only, because the serdes keep Partial and Final in Comet, and hll_sketch_estimate and hll_union read any HLL_8 form.
  2. Optionally implement GroupsAccumulator for hll_sketch_agg with one union per group. That removes the GroupsAccumulatorAdapter dispatch, which slices the input and calls size() before and after each group in every batch, and each size() computes an estimate. Together with the union, it took small sketches from 291 to 137 ns per row above.
  3. In datasketches, merge_array_same_lgk rebuilds the cached values eagerly on every array merge, so a dense peer still costs O(2^lgConfigK) per merge even with a persistent union, and so do hll_union_agg and hll_union on dense inputs. Deferring the rebuild, computing 2^-v without a division, and copying the cached values in copy_or_downsample when lgConfigK matches would follow datasketches-java. Comet pins datasketches = "=0.3.0".

Additional context

The hll_agg bench runs only AggregateMode::Partial at lgConfigK 12, so a Final-mode case with groups that turn dense at lgConfigK 18 would cover this. The aggregates came in with #4802.

Activity

  1. comphead commented on Nov 14, 2026

    @comphead
    ContributorAuthor

    Data point with #7345 in, on Comet main with DataFusion main and Spark 4.0: the final merge of 17 hll_sketch_agg(CASE WHEN ... THEN id END, 18) columns, 398M partial rows into 257M groups over 216 tasks, 8 tasks per executor with a 3 GiB task share.

    Spark (ObjectHashAggregate, sort fallback) Comet (CometHashAggregate)
    Final aggregate time, all tasks 1,105 s 10,554 s
    Partial aggregate time, all tasks 997 s 4,054 s
    Final stage task time 15 to 21 s 48 to 56 s for the 209 tasks that spilled, 30 to 31 s for the 7 that did not

    A task's groups need about 5.1 GiB of reservation, about 4.6 KB per group for 17 sketches. fair_unified gives the aggregate 1.5 GiB, half of the task's share, because the native shuffle writer above it holds the other half. So almost every task spills, and spilling serializes, sorts and merges every sketch again.

    So with #7345 in, the remaining gap is item 2, the GroupsAccumulatorAdapter dispatch with one boxed accumulator per group per column, and the per-group footprint, which decides whether the merge spills. The split of the task's share between the aggregate and the shuffle writer is in #7628.

  2. self-assigned this
    on Nov 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions