adriangb opened a new pull request, #24859: URL: https://github.com/apache/datafusion/pull/24859
## Which issue does this PR close? No existing issue. This was found while investigating a production out of memory. Happy to file one if you would like it tracked for the changelog. ## Rationale for this change `SingleDistinctToGroupBy` rewrites `AGG(DISTINCT x)` into a two phase group by, which keeps a high cardinality distinct off the one-boxed-accumulator-per-group path in `GroupsAccumulatorAdapter`. The rule already tolerates a non-distinct `sum`, `min` or `max` next to the distinct aggregate, but bails out on `count`, so the very common ```sql SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g ``` shape keeps the unrewritten plan and its memory profile. We hit this in production: a query of exactly that shape drove a process running DataFusion to 10.98 GB and death, and the single `count(*)` was the only reason the rewrite did not apply. ## What changes are included in this PR? This allows a non-distinct `count` as well. `count` is the one supported function whose outer phase is a *different* function: the inner group by counts the rows of each `(group, distinct value)` partition, and the outer phase adds those partial counts up with `sum`, since `count` over a group is the sum of the counts of any partition of that group. Two details follow from the substitution: - `count` and `sum` come from the session function registry (as `replace_distinct_aggregate` already does for `first_value`), and the rewrite only fires when the aggregate is that exact `count`, compared by identity rather than by name. A session without a registry, or with its own `count`, is left alone. - `count` returns a non-null `0` over an empty input while `sum` of no rows is NULL, which is reachable for an aggregate with no `GROUP BY`: the inner aggregate emits no rows and the outer still emits one, so `SELECT count(*), count(DISTINCT x) FROM empty` would return `NULL, 0` instead of `0, 0`. The projection selects `CASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END`, restoring the `0` and keeping the column's type and nullability as `count` had them. `FILTER` and `ORDER BY` still block the rewrite. Files touched beyond the rule itself: - `datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt`: the new coverage described below. - `datafusion/sqllogictest/test_files/clickbench.slt`: the one existing snapshot in the repository that changes, discussed below. - `datafusion/substrait/tests/cases/roundtrip_logical_plan.rs`: `aggregate_distinct_with_having` now builds its session with this rule removed, so it keeps round tripping the plan shape the test was written for. ## What is the testing strategy for this PR? `single_distinct_to_groupby.slt` asserts every result twice, once under `datafusion.optimizer.max_passes = 0` and once under the default, with identical expected blocks, so a null-handling or type error surfaces as a result mismatch rather than only a plan diff. It covers `count(*)` vs `count(1)` vs `count(col)` grouped and ungrouped, a group whose distinct column is entirely NULL, a group with NULLs in both the distinct and summed columns, empty input three ways, `HAVING` plus `ORDER BY` on the rewritten count, and the production join shape. Exactly one existing snapshot in the repository changes: the ClickBench Q22 `EXPLAIN`, which is this shape verbatim. Its result block directly beneath, running on real ClickBench parquet, is unchanged. The physical `SortExec: TopK(fetch=10)` moves from below the projection to above it, because the sort key is now a `CASE` output rather than a raw aggregate column. That is order-equivalent, since the `CASE` is the identity on every non-NULL input and the `sum` is never NULL in a grouped aggregate. Run locally: the full sqllogictest suite, plus `datafusion --test core_integration` (1079), `--test tpcds_planning` (198) and `-p datafusion --lib` (444). `cargo clippy -p datafusion-optimizer --all-targets` is clean. ### Benchmarks Q22 is the only ClickBench query whose plan changes. Q9 (`RegionID, SUM, COUNT(*), AVG, COUNT(DISTINCT UserID)`) still bails out, because `AVG` disqualifies it. Measured on `clickbench_partitioned` (100 files, ~100M rows), release builds of this branch and of the base commit it sat on at the time of the run, on a 12-core machine. A run-level A/B could not resolve a change this small here. Comparing the base binary against *itself* with `compare.py` reported 9 queries faster, 28 slower and 6 unchanged, with swings up to 1.58x, and two runs of the same base-vs-branch comparison gave opposite verdicts (11 faster / 21 slower, then 30 faster / 6 slower, with a 1.97x swing). Those tables measure background load, not the patch, because one arm is a full 43-query pass of about four minutes and load drifts between the arms. Instead the arms were paired per query, running base and branch back to back and alternating which goes first, over 40 repetitions. The 42 queries whose plans are unchanged then serve as an in-experiment control for residual bias. Q22, net of the control bias (difference in differences, bootstrap CI): ``` -2.03% 95% CI [-5.48%, +1.39%] ``` The interval includes zero, so there is no measurable latency difference, and the 95% upper bound excludes a Q22 regression larger than about 1.5%. Pooled controls moved +0.36% [-0.64%, +1.05%], confirming the setup resolves effects of roughly 3% and no better. This is latency-neutral on ClickBench, consistent with #11360, which found removing the rule entirely to be a wash. The case for the change rests on the memory behaviour of the rewritten plan, not on latency. **Not covered: memory.** These runs used no `--memory-limit`, so they do not exercise the spilling behaviour that motivates the rewrite. That is a separate experiment. ## Are there any user-facing changes? No public API change and no change to query results. Plans for `SELECT ..., count(...), count(DISTINCT x) ... GROUP BY ...` change shape, so `EXPLAIN` output for that shape differs, and such queries should use substantially less memory. The ClickBench Q22 plan change above is the visible example. -- 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]
