[andygrove]

Stacked on #2664. GitHub cannot base a PR on a fork branch, so all of #2664 shows here until it merges. Review only the last commit, feat: collect HyperLogLog distinct counts on shuffle writes for AQE.

Which issue does this PR close?

Closes #2665.

Rationale for this change

A resolved exchange carries no distinct counts, so DataFusion's join cardinality estimate downstream of a shuffle falls back to row counts and estimates an inner join at the smaller input's size. For a key/foreign-key join whose large side is the foreign key, that undercounts by the join's fan-out, which is what led AQE to broadcast a 151M-row join estimated at 400K rows in TPC-H q21 (#2661, #2662). #2664 already carries per-column statistics from completed stages onto the AQE ExchangeExec, and the protocol already reserves TaskColumnStats.sketches for HLL sketches, so this fills that slot.

What changes are included in this PR?

• ballista_core::hll: a precision-12 HyperLogLog (4 KiB of registers, about 1.6% standard error). Values are hashed with DataFusion's HLL_RANDOM_STATE, as approx_distinct does, and counted with the same estimator. Nulls are skipped.
• Shuffle writers: the shared per-stream null counter from feat: surface shuffle column statistics to the AQE planner #2664 becomes ColumnStatsCollector, which also updates one sketch per hash-partitioning key column (plain column references in the hash expressions). Each concurrently drained stream sketches on its own and merges into the task total when dropped, so there is no per-batch locking. The passthrough writer takes the keys from its child's hash partitioning, the sort writer from its own, and the range writer sketches nothing. Each writer gets a column_stats() that the executor reports as TaskColumnStats, with the registers in a ColumnSketch { type: HLL }.
• Scheduler: RunningStage::output_column_stats merges each column's sketches across successful tasks (element-wise max) and sets distinct_count to Inexact. It stays Absent if any successful task didn't sketch the column, mirroring null_count.
• New config ballista.shuffle.collect_distinct_counts (default true) to turn collection off, and a matching --collect-distinct-counts flag on shuffle_bench for measuring the cost.

Sketching only the partition keys keeps the task status payload at 4 KiB per key column per task. While a task writes, each input partition being drained holds 4 KiB per key column outside the memory pool.

On a local shuffle_bench run (sort writer, 60M TPC-H SF100 lineitem rows hashed on l_orderkey, 8 output partitions, 3 iterations each, alternating), the difference between collection on and off was within run-to-run noise (19.4s to 19.6s for both, plus one 22.2s outlier with collection on caused by a slow spill). That run is spill-dominated, so a TPC-H SF1000 cluster run is still worth doing before relying on the default.

Tests cover the sketch itself (accuracy, nulls, merge, serialization round trip and validation), key column selection, the sort and passthrough writers reporting a sketch only on the hash key, the disable flag, and the scheduler merge, including overlapping keys across tasks and a task without a sketch.

Are there any user-facing changes?

A new config key, ballista.shuffle.collect_distinct_counts, documented in the generated config reference. No breaking API changes.

[villebro]

Very valuable feature 👍 Some random thoughts.

[villebro]

I don't disagree with 12 over 14 here, but it does pose the question, why not 10, or some other less memory intensive number?

[villebro]

Optional consideration: expose ballista.shuffle.collect_distinct_counts_precision that defaults to 12 for those that really want to fine tune their deployment?