andygrove opened a new pull request, #2671:
URL: https://github.com/apache/datafusion-ballista/pull/2671
> **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 #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.
--
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]
