SingleDistinctToGroupBy fires where it costs far more memory: adding count(*) makes a query ~1000x cheaper
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
Adding `count(*)` to a grouped query that computes `min(DISTINCT x)` makes the query need 1000x less memory. The extra aggregate makes the query compute strictly more work. It also takes the plan out of `SingleDistinctToGroupBy`, and that rewrite is what costs the memory.
`SingleDistinctToGroupBy` turns `AGG(DISTINCT x)` into an inner `GROUP BY (group_keys, x)` plus an outer aggregate. The inner aggregate holds one row per distinct `(group_keys, x)` pair instead of one row per group. Every other aggregate in the query moves down to that finer grain.
### To Reproduce
Everything below runs in a stock `datafusion-cli`. There is no patch, no feature flag and no custom build.
Build the table:
```sql
CREATE TABLE t AS
SELECT v % 2000 AS g, (v * 48271) % 999983 AS x, v % 1000 AS y
FROM (SELECT unnest(generate_series(0, 3999999)) AS v);
```
That is 4,000,000 rows, 2,000 groups, 999,983 distinct values of `x`, and 4,000,000 distinct `(g, x)` pairs.
Run each query in its own session, under a bounded memory pool:
```
datafusion-cli -m -d 0 --mem-pool-type greedy -f
```
Start the file with `SET datafusion.execution.target_partitions = 1;`, then the `CREATE TABLE` above, then one query. One partition makes the result the same on any machine. `-d 0` turns spilling off, so `` measures the memory the query needs, not how far it can spill.
| query | rule fires | smallest limit that completes |
| --- | --- | --- |
| `SELECT g, min(DISTINCT x) FROM t GROUP BY g` | yes | 192M, fails at 160M |
| `SELECT g, count(*), min(DISTINCT x) FROM t GROUP BY g` | no | 192K, fails at 160K |
At `-m 32M` the first query fails. The second one, which computes strictly more, returns its 2,000 rows in 10 ms.
The two plans differ only by the `count(*)`:
```
> EXPLAIN FORMAT indent SELECT g, min(DISTINCT x) FROM t GROUP BY g;
Projection: t.g, min(alias1) AS min(DISTINCT t.x)
Aggregate: groupBy=[[t.g]], aggr=[[min(alias1)]]
Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[]]
TableScan: t projection=[g, x]
> EXPLAIN FORMAT indent SELECT g, count(*), min(DISTINCT x) FROM t GROUP BY g;
Projection: t.g, count(Int64(1)) AS count(*), min(DISTINCT t.x)
Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), min(DISTINCT t.x)]]
TableScan: t projection=[g, x]
```
The rewritten plan builds a hash table of 4,000,000 rows. The plan that also computes `count(*)` keeps 2,000 groups.
A second pair of queries behaves the same way:
| query | rule fires | smallest limit that completes |
| --- | --- | --- |
| `SELECT g, sum(y), sum(DISTINCT x) FROM t GROUP BY g` | yes | 224M, fails at 192M |
| `SELECT g, count(*), sum(y), sum(DISTINCT x) FROM t GROUP BY g` | no | 96M, fails at 64M |
If you leave spilling on and drop `-d 0`, the first query still completes at a small limit. It spills to disk to get there. The second query never spills.
#### Peak process memory
The memory limit is one instrument. Process memory is a second one, and it needs no limit and no build. Run the same two queries with no `-m` at all, under `/usr/bin/time -l` on macOS or `/usr/bin/time -v` on Linux.
| run | peak RSS |
| --- | --- |
| build the table, then `SELECT 1` | 289 MiB |
| build the table, then `SELECT g, min(DISTINCT x) FROM t GROUP BY g` | 525 MiB |
| build the table, then `SELECT g, count(*), min(DISTINCT x) FROM t GROUP BY g` | 291 MiB |
The table itself accounts for the 289 MiB baseline. On top of that baseline the first query costs about 236 MiB and the second costs about 1 MiB. Three runs of each gave 525, 556 and 525 MiB for the first query, and 291 MiB every time for the second.
`datafusion-cli` sets mimalloc as its global allocator, and mimalloc reports the same thing from inside the process. `MIMALLOC_SHOW_STATS=1` prints `peak rss: 524.8 MiB, peak commit: 538.0 MiB` for the first query and `peak rss: 290.6 MiB, peak commit: 325.8 MiB` for the second. This is committed memory, not an accounting artifact.
### Expected behavior
Adding an aggregate to a query should not cut its memory requirement by three orders of magnitude. The rewrite should not fire on plans where it costs more than the plan it replaces.
What the rule fires on today is which aggregate functions appear in the query. It accepts the plan when every non-distinct aggregate is `sum`, `min` or `max`, and rejects it otherwise. That condition says nothing about whether the rewrite is cheaper, so the rule fires where it is a pure loss. `count` is outside the accepted set, which is why a `count(*)` makes the query cheap.
### Additional context
#### The rewrite is a large win on other plans
The same measurement across five distinct aggregates, two group counts and two column types. The table, the method and the companion aggregate are the same in every row. The `rule on` column is `SELECT g, sum(y), AGG(DISTINCT x) FROM t GROUP BY g`. The `rule off` column is the same query with `count(*)` added, which is the only difference. `EXPLAIN` confirms the rewrite in every `rule on` row and its absence in every `rule off` row.
| distinct aggregate | groups | `x` type | rule on | rule off | |
| --- | --- | --- | --- | --- | --- |
| `min(DISTINCT x)` | 2,000 | BIGINT | 224M | 192K | 1195x worse |
| `sum(DISTINCT x)` | 2,000 | BIGINT | 224M | 96M | 2.3x worse |
| `avg(DISTINCT x)` | 2,000 | BIGINT | 224M | 96M | 2.3x worse |
| `count(DISTINCT x)` | 2,000 | BIGINT | 224M | 192M | 1.2x worse |
| `array_agg(DISTINCT x)` | 2,000 | BIGINT | 224M | 384M | 1.7x better |
| `count(DISTINCT x)` | 500,000 | BIGINT | 224M | 224M | no change |
| `count(DISTINCT x)` | 500,000 | VARCHAR | 256M | 14G | 56x better |
So the rule is worth keeping. It is 56x better for `count(DISTINCT )` at high group cardinality, and 1.7x better for `array_agg(DISTINCT)`. Deleting it outright would trade one regression for another. A gate that reflects when the rewrite helps looks more promising than either the current function-name condition or a blanket removal.
Using `sum`, `min` or `max` as the non-distinct companion gives the same numbers. All three put the `min(DISTINCT x)` pair at 224M with the rule on and 192K with it off. The companion aggregate does not affect the mechanism. It only decides whether the rule fires.
#### Why the outcome varies
What decides the result is the storage cost per distinct value on each side.
The rewrite materializes one hash table row per distinct `(group_keys, x)` pair, plus an accumulator slot per other aggregate at that grain. It wins when the accumulator it replaces costs more than that per value, such as a per-group hash set of strings or an `array_agg`. It loses when the accumulator it replaces costs less.
`min` and `max` are the extreme case, because `min(DISTINCT x)` equals `min(x)`. DataFusion already knows this: `min` and `max` ignore the `DISTINCT` flag, and on the table above the two expressions agree on all 2,000 groups. So the plan without the rewrite holds one scalar per group, and the rewrite builds a hash table of 4,000,000 rows to reach the same answer.
#### Related work
- #11360 asks whether the rule is still needed now that distinct accumulators exist. Its ClickBench comparison found no clear advantage from deleting it. That thread measures planning and runtime, and asks whether the rule is *unnecessary*. This report is about the rule being *actively harmful* on a class of plans, on a trigger condition unrelated to its benefit. It also shows the rule is genuinely valuable on other plans.
- #15099 proposed removing the rule outright. The table above is evidence against that: two of its seven rows get worse without the rule.
- #8266, closing #8123, added the tolerance for non-distinct `sum`, `min` and `max`. That tolerance is what lets the rule fire on the `min(DISTINCT)` and `sum(DISTINCT)` cases above.
- #20782 proposed skipping the rewrite for `count(DISTINCT)` with no `GROUP BY`, on the same grounds that the direct distinct accumulator is cheaper. It was closed without merging.
- #11686 proposes eliminating `DISTINCT` on `min` and `max` early, which would remove the worst case here at its source.
- #20942 and #21087 cover the multiple-distinct case, which this rule does not handle at all.
- #24704 tracks blocked and chunked memory management in hash aggregation.
#### Method
Each memory figure is the smallest value at which the query completes. I found it by bisecting a fixed ladder that runs from 64K to 12G in steps of 14% to 50%. The next value down the ladder fails in every row of every table above, and repeated runs at both values always give the same result. The one figure off that ladder is the 14G rule-off arm of the VARCHAR row, which needs more than the ladder covers. I bracketed it by hand: 14G completes and 13G fails.
Every run uses `-d 0`, `--mem-pool-type greedy` and `SET datafusion.execution.target_partitions = 1`. The `CREATE TABLE` takes no memory from the pool: it succeeds at `-m 64K`.
The three datasets differ only in the `g` and `x` expressions. `v % 2000` and `v % 500000` give the two group counts. `CAST((v * 48271) % 999983 AS VARCHAR)` gives the string column. In all three, `x` is drawn from 999,983 values and every one of the 4,000,000 rows is a distinct `(g, x)` pair.
#### Note on heap profiling
`datafusion-cli` cannot produce a jemalloc heap profile as shipped. `datafusion-cli/src/main.rs` installs mimalloc as the global allocator with no feature flag behind it, and no crate in the workspace depends on jemalloc, so `MALLOC_CONF` has no effect. Getting a jemalloc profile needs a source change and a rebuild. `MIMALLOC_SHOW_STATS=1` and `/usr/bin/time` both work on the shipped binary, which is why the numbers above use those.
#### Code
`is_single_distinct_agg` in `datafusion/optimizer/src/single_distinct_to_groupby.rs` holds the trigger condition. `datafusion/functions-aggregate/src/min_max.rs` never reads `is_distinct`, which is correct and is what makes the `min` case the extreme one.
#### Version
Reproduced with `datafusion-cli 55.0.0`, built from `main` at `b9a0053677`, on macOS arm64. `single_distinct_to_groupby.rs` and `min_max.rs` are unchanged at `9fc7a4d5d`, the current `main` head at the time of writing.
Contributor guide
Assessment
This issue has not been assessed yet.