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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions