jayzhan211 opened a new issue, #24980:
URL: https://github.com/apache/datafusion/issues/24980
### Describe the bug
`LimitedDistinctAggregation` rewrites `SELECT DISTINCT ... LIMIT n` by
pushing `n` into `AggregateExec` as a soft limit (`LimitOptions`, shown as
`lim=[n]` in `EXPLAIN`). The intent is that the aggregate stops consuming input
as soon as `n` distinct groups exist, and the downstream limit operator trims
the exact row count.
The legacy `GroupedHashAggregateStream` honours this in every mode via
`hit_soft_group_limit()`. Among the new dedicated streams (#22710), only
`PartialHashAggregateStream` and `FinalHashAggregateStream` implement it.
`SingleHashAggregateStream` has no `group_values_soft_limit` at all: given a
`Single` / `SinglePartitioned` aggregate with `lim=[n]`, it reads its whole
input and materialises every group before emitting.
**Current `main` (since #24961):** the soft limit is treated as a hint
during stream selection, and the `limit_options.is_none()` conditions were
dropped from the stream predicates (see the doc comment on
`AggregateExec::limit_options`). Single-stage distinct aggregates therefore
route to `SingleHashAggregateStream` and lose the early termination. The plan
text is unchanged, so this is a performance regression with no visible signal.
**Before #24961:** `should_use_single_hash_stream()` required
`limit_options.is_none()`, so any single-stage aggregate carrying a soft limit
fell back to the legacy `GroupedHashAggregateStream` and kept the early
termination. The gap in `SingleHashAggregateStream` was masked by that routing,
not fixed, and the legacy stream is now deprecated and slated for removal
(#22710), so restoring the fallback is not a long-term option.
Nothing in the test suite catches this, because the plan text is identical
before and after #24961: the node prints `lim=[n]` whether or not the selected
stream acts on it. For example, `aggregate.slt` has
```
AggregateExec: mode=SinglePartitioned, gby=[c3@0 as c3, min(...)@1 as ...],
aggr=[], lim=[5]
```
which passes unchanged while scanning all of its input under the limit. The
same shape appears for `SELECT DISTINCT ... LIMIT n` whenever the aggregate is
planned as a single stage: `target_partitions = 1`, or hash-partitioned input
on the group keys (`SinglePartitioned`).
Streams affected:
| Stream | Soft limit | Notes |
|---|---|---|
| `GroupedHashAggregateStream` (legacy) | yes | all modes |
| `PartialHashAggregateStream` | yes | |
| `FinalHashAggregateStream` | yes | skipped once the stream has spilled |
| `SingleHashAggregateStream` | **no** | selected for single-stage
distinct-with-limit since #24961; previously masked by the legacy fallback |
| `PartialReduceHashAggregateStream` | no | only matters if a limit is ever
pushed into a `PartialReduce` node |
| `OrderedSingleAggregateStream` / `OrderedFinalAggregateStream` | no | low
impact: they emit groups incrementally, so the downstream limit stops polling
after ~one extra batch |
### To Reproduce
```sql
set datafusion.execution.target_partitions = 1;
create table t as select v % 1000 as k from generate_series(1, 10000000) as
s(v);
explain analyze select distinct k from t limit 5;
```
The plan shows `AggregateExec: mode=Single, ..., aggr=[], lim=[5]`. On
current `main` the scan's `output_rows` reports all 10M rows. Before #24961
(legacy stream selected), and after the fix proposed below, it stays at a
single batch.
### Expected behavior
`SingleHashAggregateStream` (and ideally the other new streams that can
accumulate a distinct group count) should stop reading input once `limit <=
building_group_count()`, matching the legacy stream and
`PartialHashAggregateStream`.
## Proposed solution
1. **Implement the soft limit in `SingleHashAggregateStream`.** Mirror
`PartialHashAggregateStream`:
- Add `group_values_soft_limit: Option<usize>`, initialised from
`agg.limit_options().map(|c| c.limit())`.
- In the `ReadingInput` branch, after
`hash_table.aggregate_batch(&batch)`, check `limit <=
hash_table.building_group_count()`. When hit, `close_input()` and take the same
transition as the `Poll::Ready(None)` (input exhausted) arm.
- Reusing the input-exhausted transition keeps spilling correct for free:
if the stream has already spilled, it goes through `PreparingMergeInput` /
`MergingSpills` as usual. Stopping early is still valid there, because the
in-memory table alone already holds at least `limit` distinct groups, and the
soft limit only applies to `aggr=[]` (see
`is_unordered_unfiltered_group_by_distinct`), so there is no accumulator state
to lose.
- `OrderedSingleAggregateStream` does not need it for the reason in the
table above.
- `PartialReduceHashAggregateStream` can get the same treatment if we
ever push limits into `PartialReduce`; otherwise document that it is
unsupported.
2. **Add a test that fails without the fix.** The plan text cannot catch
this, so the test has to observe consumption. A unit test next to the existing
partial/final soft-limit test in `aggregates/mod.rs` that runs
`AggregateMode::Single` with `LimitOptions::new(2)` over several small input
batches, asserts `StreamType::SingleHash`, and checks either (a) the output row
count equals the limit, or (b) via a counting input stream, that not every
input batch was polled.
3. **Keep the `AggregateExec::limit_options` doc comment in sync** so the
list of streams that honour the hint stays accurate.
Fallback if (1) is not wanted right away: restore `limit_options.is_none()`
in `should_use_single_hash_stream()` so the single-stage distinct-with-limit
case routes back to the legacy stream until it is removed. That reintroduces
the legacy dependency #24961 dropped, so I'd prefer implementing the limit
directly.
### Additional context
- Part of #22710 (split aggregation into dedicated streams). Routing change
that exposed this: #24961.
- The soft limit is set by `LimitedDistinctAggregation`
(`datafusion/physical-optimizer/src/limited_distinct_aggregation.rs`), gated by
`datafusion.optimizer.enable_distinct_aggregation_soft_limit`.
- Legacy check: `GroupedHashAggregateStream::hit_soft_group_limit`.
New-stream check to mirror: `PartialHashAggregateStream::hit_soft_group_limit`
("Step 2: Soft limit optimization" in `hash_stream.rs`).
--
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]
