Repository navigation
Description
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
- Keep one
SparkHllUnioninHllSketchAccumulatorformerge_batch, asHllUnionAccumulatoralready does, and build the sketch only instate()andevaluate(). If one accumulator sees bothupdate_batchandmerge_batch, fold the updated sketch into the union there. The intermediate stays Comet-only, because the serdes keep Partial and Final in Comet, andhll_sketch_estimateandhll_unionread any HLL_8 form. - Optionally implement
GroupsAccumulatorforhll_sketch_aggwith one union per group. That removes theGroupsAccumulatorAdapterdispatch, which slices the input and callssize()before and after each group in every batch, and eachsize()computes an estimate. Together with the union, it took small sketches from 291 to 137 ns per row above. - In
datasketches,merge_array_same_lgkrebuilds 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 dohll_union_aggandhll_unionon dense inputs. Deferring the rebuild, computing 2^-v without a division, and copying the cached values incopy_or_downsamplewhen lgConfigK matches would follow datasketches-java. Comet pinsdatasketches = "=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