zhuqi-lucas opened a new pull request, #23806:
URL: https://github.com/apache/datafusion/pull/23806
## Which issue does this PR close?
Rebased from #22518. Phase 1 prototype for #22405. Supersedes #22518, which
was 483 commits behind main after the #22710 aggregation-stream migration
(`row_hash.rs` → `grouped_hash_stream.rs` / `hash_stream.rs` /
`skip_partial.rs`).
## Rationale for this change
The current `skip_partial_aggregation_probe_ratio_threshold` (default `0.8`)
is a single fixed knob: when measured `num_groups / input_rows` ≥ 0.8, partial
aggregation skips. This catches the no-reduction case but **misses the
medium-ratio band** (~0.5–0.7) where partial aggregation is still net-negative
because per-row cost is high — heavy variable-length keys, complex aggregates,
etc.
ClickBench Q18 is the motivating example. Measured ratio is 0.565, well
below 0.8, so partial aggregation keeps running and burns significant compute
across 12 partitions for only ~40% reduction. Lowering the global threshold to
0.6 fixes Q18 but likely regresses lower-cost queries that benefit from partial
agg at that ratio.
This PR replaces the static-threshold guess with **measured A/B sampling**:
probe the per-row cost of both partial-agg and passthrough, then pick the
cheaper path via a closed-form cost comparison. No magic constants.
## How it works
After the existing partial probe window (`probe_rows_threshold`, default
100k rows) closes:
1. **Partial probe** — measure `partial_ns_per_row` and `ratio = num_groups
/ input_rows`.
2. **Rule 1 short-circuit** — `ratio ≥ probe_ratio_threshold (0.8)` skips
immediately (existing behaviour, preserved).
3. **A/B sampling** — otherwise route the next `ab_sampling_rows` (default
10k) through the passthrough (`transform_to_states`) path. Hash table is
preserved; the passthrough output is sent downstream and merges naturally in
Final agg.
4. **Cost decision** — measure `passthrough_ns_per_row`, then apply `skip ⇔
ratio > passthrough_ns_per_row / partial_ns_per_row`.
5. Skip → emit partial hash table, continue via `SkippingAggregation`. Keep
→ return to `ReadingInput` and continue accumulating.
Derived from `cost_keep_partial = partial × N + final × N × ratio` vs
`cost_skip = passthrough × N + final × N` assuming `final ≈ partial`.
Automatically adapts to hardware and query shape.
## Original benchmark (from #22518, ClickBench partitioned, ARM Neoverse-V2
12 vCPU)
| Metric | HEAD | This PR | Δ |
|---|---|---|---|
| Total | 20027 ms | 19726 ms | **−1.5%** |
| Q18 | 1291 ms | 1156 ms | **+1.12×** |
| Q19 | 39 ms | 27 ms | **+1.43×** |
| Q39 | 153 ms | 118 ms | **+1.30×** |
| Q29 | 51 ms | 41 ms | **+1.23×** |
| Q36 | 73 ms | 63 ms | **+1.16×** |
| Q35 | 323 ms | 297 ms | **+1.09×** |
Fresh CI run on current main is TODO before flipping this out of draft.
## What changes are included in this PR?
Three commits, mapped to the post-#22710 aggregation-stream structure:
1. **`e4d1cba` config knobs** — `skip_partial_aggregation_use_cost_model` +
`skip_partial_aggregation_ab_sampling_rows` + docs.
2. **`c01757d` SkipAggregationProbe A/B state machine** — extends the probe
(in the standalone `skip_partial.rs` since #22729) with `Partial → AbSampling →
Locked { should_skip }`. Lazily registers four diagnostic gauges
(`partial_agg_probe_partial_ns_per_row`,
`partial_agg_probe_passthrough_ns_per_row`,
`partial_agg_probe_ratio_per_mille`, `partial_agg_probe_cost_decision_skip`).
Callers in `grouped_hash_stream.rs` + `hash_stream.rs` updated for the new
constructor.
3. **`f1182e8` state-machine wiring** — new `ExecutionState::AbSampling`
variant + transitions (`ReadingInput → AbSampling → {SkippingAggregation |
ReadingInput}`).
## Config knobs (both `datafusion.execution.*`)
- `skip_partial_aggregation_use_cost_model` (bool, default **true**) — turns
the A/B path on.
- `skip_partial_aggregation_ab_sampling_rows` (usize, default 10000).
## Are these changes tested?
- **123/123 existing `aggregates::*` unit tests pass locally.**
- Legacy
`test_partial_hash_stream_skip_aggregation_probe_not_locked_until_skip` updated
to explicitly disable the cost model (asserts pre-#22518 semantics).
- **TODO before undraft**: rebase the seven A/B-specific unit tests from
#22518 onto the new `skip_partial.rs` module. SLT snapshots for
`push_down_filter_regression.slt` and `information_schema.slt` will need regen
for the new gauges + config rows.
## Are there any user-facing changes?
Two additive `datafusion.execution.*` config options. Default is on based on
the #22518 benchmark; opt-out via `SET
datafusion.execution.skip_partial_aggregation_use_cost_model = false`.
## Followups (from #22518)
- **Segment-level re-probing** — attempted in #22518 as `44f815a`, reverted
in `a258afe`. Blocked on a `GroupValues` reset semantic. Should be tackled as a
follow-up.
- The simplifying assumption `final_ns ≈ partial_ns` in the cost formula is
reasonable but not exact. A more refined model could track Final-agg cost
separately.
## Why draft
- New unit tests for the A/B path still to be migrated from #22518.
- SLT snapshots need regen.
- Local ClickBench sanity run pending on current main.
--
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]
