[comphead]

Which issue does this PR close?

Part of #7341. This is the first item of the solution proposed there. The GroupsAccumulator and the datasketches changes it also lists are not in this PR, so #7341 stays open.

Rationale for this change

HllSketchAccumulator::merge_batch merged each partial sketch through SparkHllSketch::merge_sketch, which built a new SparkHllUnion for the row, merged the accumulator and the partial into it, and turned the union back into a sketch. Once a group's sketch holds the dense register array, datasketches 0.3.0 copies that array into the empty union and rescans all 2^lgConfigK registers to rebuild its estimator. So every merged row costs as much as the whole array, 262,144 registers at lgConfigK 18. The issue has the profile and the numbers from a Spark 4.0 job where the Final aggregation's median task took 40.4 s against Spark's 8.9 s.

What changes are included in this PR?

• HllSketchAccumulator is an enum of a SparkHllSketch and a SparkHllUnion. Input rows are hashed into the sketch as before. The first merge_batch folds the sketch into a union, which merges every partial sketch from then on, as HllUnionAccumulator already does. state() and evaluate() serialize whichever of the two the accumulator holds.
• The accumulator holds one or the other rather than both, so it grows by 8 bytes, from 96 to 104. An empty sketch kept next to the union would roughly double what a small group is charged in a Final aggregation.
• DataFusion gives an accumulator either update_batch or merge_batch calls. If one ever gets rows after a merge, they reach the union through a sketch of their own.
• SparkHllSketch::merge_sketch and the private SparkHllUnion::to_sketch are removed, since nothing else used them. SparkHllUnion::lg_max_k is added for the case above.
• The hll_agg bench gets a Final-mode case, hll_sketch_agg_final_dense: 8,000 partial states at lgConfigK 18, 20 groups of 400 partials with 200 distinct values each, read one partial per group per round as a reducer would.

Estimates can move slightly once a group is dense. Before, every merge into a dense group marked its estimator out of order, which makes the estimate come from the composite estimator instead of the HIP accumulator. Now only merging a dense partial does that, as in datasketches-java, which Spark's HllSketchAgg.merge uses. LIST and SET mode results are unchanged. In the dense probe case below, the sum of the 20 estimates went from 1,600,370 to 1,600,533, against a true 1,600,000.

How are these changes tested?

Two new unit tests in hll_sketch_agg.rs:

• merge_batch_keeps_one_union_across_calls merges 40 partial sketches in batches of 7 and checks after every batch that the accumulator holds the same bytes as one SparkHllUnion fed the same sketches, in SET mode and after the group turns dense.
• rows_before_and_after_a_merge_all_count gives one accumulator rows, then a partial sketch, then more rows, and checks that the estimate counts all three.

expressions/aggregate/hll_sketch_agg.sql compares with Spark 4.0+: every input type with NULLs and bounds, and four files that the Final aggregation merges as one partial per group. Groups whose union stays a coupon set and groups merging partials that already hold the register array match Spark exactly. Where coupon-set partials take the union to the register array, the estimate depends on merge order in Spark too, so that query checks each engine's estimate against the true count, within 1%.

The existing tests cover the update path, the empty group and a Partial/Final plan.

Benchmarks, release build on an Apple M5 Max, cargo bench --bench hll_agg -- --sample-size 10 before and after:

Bench	Before	After
hll_agg/hll_sketch_agg_final_dense (new)	1.682 s	25.8 ms
hll_agg/hll_sketch_agg (Partial)	1.877 ms	1.838 ms
hll_agg/hll_union_agg (Partial)	43.98 ms	43.86 ms

The Final-mode probe from the issue, per partial row:

Partial states, lgConfigK 18	Before	After
1.7M rows, 200K groups, 0 to 8 values each	332 ns	281 ns
200K rows, 2K groups, 50 values each (SET mode)	2.31 µs	1.76 µs
8K rows, 20 groups, 200 values each (dense)	207.7 µs	3.3 µs

[andygrove]

Thanks for this. The idea is sound and the numbers are convincing. Folding the sketch into one SparkHllUnion on the first merge_batch removes the per-row copy and rescan of the register array, and the Partial path is unchanged. I ran the hll tests in datafusion-comet-spark-expr on the PR head and all 38 pass, including the two new ones. I did not run the benches.

I compared against branch-1.2. The only user-visible difference I found is that estimates for dense groups can move slightly. That fits what expressions.md already says about hll_sketch_agg being opt-in via allowIncompatible, so I don't see a support-level change. LIST and SET mode results, fallback behavior, errors and configs are the same.

Two things I'd like to see addressed.

1. The unreachable! in merge_batch. The branch is correct today, but it is the one panic in this accumulator. Could the conversion and the borrow be one step, for example a small fn union_mut(&mut self) -> &mut SparkHllUnion that replaces a Sketch with a Union and returns the union? Then merge_batch and the update_batch fallback could share it, and the compiler would enforce the invariant instead of a comment.

2. The dense-group estimate against real Spark. The description says estimates move slightly once a group is dense, and the 1,600,533 against 1,600,000 probe shows we are inside the error bound. Spark's HllSketchAgg.merge builds a new union per row, so Spark's own estimate also depends on merge order. Do we have a test that compares Comet's Final estimate with Spark 4.0 output for a group that goes dense, for example an hll_sketch_agg SQL file test at a high lgConfigK? The existing Comet-versus-Spark comparisons check exact numbers, so I'd like to know they still pass on that shape.

The comparison with a plain SparkHllUnion in merge_batch_keeps_one_union_across_calls pins the "one union across calls" behavior well. It uses the same code on both sides, so it is the SQL tests above that give the Spark comparison.

CI had only the preflight checks finished when I looked, so the Comet suites haven't reported yet.

[comphead]

Thanks @andygrove, both are addressed in 0d514bb.

1. merge_batch now gets its union from union_mut, which swaps a Sketch for a Union and then calls itself, so it can only return through the Union arm and the unreachable! is gone. The update_batch fallback already holds the union from its own match, so it doesn't need it.

2. No test compared hll_sketch_agg with Spark. The CometExpressionSuite tests check the true count within 5%. The new expressions/aggregate/hll_sketch_agg.sql compares with Spark over four files, each one partial per group that the Final aggregation merges. It expects exact results for every input type with NULLs and bounds, for groups whose union stays a coupon set (lgConfigK 21), and for groups merging partials that already hold the register array (lgConfigK 12, grouped and global). At lgConfigK 18 coupon-set partials take the union to the register array during the third merge. From there the estimate depends on merge order in Spark as well, so that query compares within 1%.

I haven't run the file against Spark locally. Its expectations come from replaying both sides on the same data: datasketches-java 6.1.1 and 6.2.0 (Spark 4.0 and 4.1, which agree everywhere) against HllSketchAgg in a Final AggregateExec. For the fixture's four partials I tried all 24 merge orders.

Shape	This PR vs Spark	Before vs Spark
lgConfigK 21, coupon sets	equal	equal
lgConfigK 12, dense partials	equal	equal
lgConfigK 12, 300 list partials turning dense	equal	6 of 1,500
lgConfigK 18, coupon sets turning dense	up to 54 of 40,000	up to 70
Issue probe, 400 partials of 200 values	up to 5 of 80,000	up to 156

On the lgConfigK 18 shape Spark's own estimate moves by up to 103 between merge orders. What is left there is coupon order. The estimate depends on the order the coupons reach the union, and the two engines read a coupon-set partial back in different orders. With list partials, whose order both keep, the results are equal.

One difference predates this PR. When a group's only partial is already dense, datasketches 0.3.0 marks the new union out of order, so Comet uses the composite estimate where Spark keeps HIP, 684 against 695 for 700 values at lgConfigK 12. Taking a group's first partial as is would match Spark there, since that partial read back gives 695. I'd leave it for a follow-up.

[tisonkun]

I've opened an upstream PR for the DataSketches Rust costs identified here and in #7341: apache/datasketches-rust#317.

It contains three separate commits:

• Copy dense HLL8 inputs, including their estimator state, without rescanning the registers when no downsampling is needed.
• Add HllUnion::serialize() to emit HLL8 sketch bytes directly, avoiding to_sketch(HllType::Hll8).serialize() and its intermediate copy.
• Defer register-cache reconstruction across dense merges until an estimate, bound, or serialized result is needed, while retaining shared-reference queries and thread safety.

In the library benchmarks at lg_k 18, merging 16 dense partials and then querying drops from 4.91 ms to 0.284 ms; recreating a union for each of 400 small partials drops from 79.71 ms to 3.10 ms. These are library measurements, not end-to-end Comet results. Keeping a persistent union as this PR does remains useful.

Reviews and feedback from the Comet workload would be very welcome, and hopefully Comet can pick these up in a subsequent upgrade once they are merged and released. Newer releases already provide estimated_size() and fixes for compact HLL decoding and single-input HIP preservation, which can also remove several 0.3.0 workarounds. Using estimated_size() instead of calling estimate() for memory accounting will help preserve the benefit of deferred reconstruction.

[comphead]

related #7402

[andygrove]

@comphead Should we merge #7402 first?

[comphead]

@comphead Should we merge #7402 first?

Yeah, lets do it, and i resolve conflicts later

[andygrove]

Thanks for this. Folding the sketch into one SparkHllUnion on the first merge_batch removes the per-row copy and rescan of the register array, and union_mut makes the invariant something the compiler enforces. Both asks from my earlier review are addressed.

I merged the PR into current main locally, which includes the datasketches 0.6.0 bump from #7402. The merge is clean. The 40 hll unit tests pass, and hll_sketch_agg.sql passes against Spark 4.0 under -Pspark-4.0. On the hll_agg bench, hll_sketch_agg_final_dense goes from 1.91 s on main to 37 ms. The Partial benches are unchanged.

I compared against branch-1.2. The only user-visible difference is that estimates for dense groups can move slightly, and the PR table shows them closer to Spark than before. LIST and SET results, serialized bytes, fallback, errors, configs and the support level are the same.

The branch was last merged with main before #7402, so CI hasn't seen 0.6.0. Could you merge main into the branch so it does?