comphead opened a new issue, #7341:
URL: https://github.com/apache/datafusion-comet/issues/7341
### 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 |
<details><summary>Snippet</summary>
```rust
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());
}
```
</details>
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
