Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy

[adriangb]

Which issue does this PR close?

No existing issue. We found this during an investigation of a production out of memory. We can file one if you want a changelog entry.

Rationale for this change

Reproduce it

Run this in datafusion-cli. It writes 4,000,000 rows in 500,000 groups, with 2,000,000 distinct (g, x) pairs.

COPY (
  SELECT
    value % 500000 AS g,
    'id-' || CAST(value % 2000000 AS VARCHAR) AS x
  FROM generate_series(1, 4000000)
) TO 'repro.parquet' STORED AS PARQUET;

CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'repro.parquet';

EXPLAIN FORMAT INDENT SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g;

On main the rule leaves the aggregate alone:

Projection: t.g, count(Int64(1)) AS count(*), count(DISTINCT t.x)
  Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), count(DISTINCT t.x)]]
    TableScan: t projection=[g, x]

Now put the CREATE EXTERNAL TABLE statement and the query into run.sql, and watch the process:

/usr/bin/time -l datafusion-cli -f run.sql     # macOS; use time -v on Linux

On main this machine reports 16.4 GiB peak RSS and 4.8 s for 500,000 output rows. The count(*) is the only reason the rule stops. Delete it and main rewrites the query.

What this PR changes for that query

The rule now accepts the count(*) and rewrites the aggregate:

Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(*), count(alias1) AS count(DISTINCT t.x)
  Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]]
    Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2]]
      TableScan: t projection=[g, x]

The same run reports 335 MiB peak RSS and 0.07 s. The results are identical.

Read that 50x with care. It is the number on main today, and main does not carry apache#24857. apache#24857 removes a per-group pre-allocation of about 33.8 KiB from the path this rewrite exists to avoid. Stacked on apache#24857 the same comparison gives 1.82x, not 50x. The table below gives both columns.

datafusion.execution.parquet.schema_force_view_types defaults to true, so x above arrives as Utf8View.

Why the rewrite helps here

SingleDistinctToGroupBy rewrites AGG(DISTINCT x) into a two phase group by. The rule accepts a non-distinct sum, min or max next to the distinct aggregate. It rejects a non-distinct count. One count(*) is therefore enough to keep the unrewritten plan.

We hit this in production. A query of this shape drove a process to 10.98 GB and to death.

The rewrite is not free. Every other aggregate moves down into the inner group by. That group by holds one row for each (group, distinct value) pair, not one row for each group. Each aggregate then keeps its state at that finer grain. The rewrite pays for this cost when it takes the distinct aggregate off GroupsAccumulatorAdapter. The adapter keeps one boxed Accumulator for each group, which is the expensive shape.

count(DISTINCT x) has a specialized GroupsAccumulator for each integer type, and for no other type. An integer distinct count never reaches the adapter, so the rewrite buys nothing for it.

The measurements

The harness reports the peak MemoryPool reservation for SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g over 4,000,000 rows. The first column is today's main. The second column is the same measurement on a tree that also carries apache#24857.

groups	distinct argument	on main	with apache#24857
500,000	BIGINT	1.14x worse	1.14x worse
500,000	Utf8	15.1x better	1.67x better
500,000	Utf8View	44.1x better	1.82x better
2,000	BIGINT	1.33x worse	1.33x worse
2,000	Utf8	1.04x worse	1.11x better
2,000	Utf8View	1.18x better	1.30x better

The second column is the honest one, and it is much smaller than the first. The absolute saving in the motivating cell falls from 13.4 GiB to 181 MiB, which is about 75 times less. An earlier version of this description claimed 55x. That number is an artifact of a pre-allocation that goes away when apache#24857 lands.

Two notes on the table.

The type labels in the earlier description were wrong. A string column that comes from a Parquet file is Utf8View, not Utf8, so the rows labelled VARCHAR were Utf8View rows. The two types are far apart on main. At 500,000 groups the unrewritten arm costs 4.00 GiB for Utf8 and 13.67 GiB for Utf8View.

apache#24857 also closes the hole at 2,000 groups, and for an unexpected reason. The unrewritten arm gets larger there, from 261 MiB to 303 MiB. apache#24857 removes the pre-allocation, and it also makes size() report the real hashbrown allocation. At 2,000 distinct values per group the honest accounting is larger than the pre-allocation it removes.

Where the rewrite stops paying

There is no crossover in group count for a string argument. The peak of the rewritten arm follows the number of distinct pairs, which does not change with the group count. The peak of the unrewritten arm grows with the group count.

The crossover is in the density of distinct values. It arrives only when every row holds a distinct value. There the residual Utf8 loss is 1.06x to 1.08x, and Utf8View does not cross at all.

The gate is a proxy

The gate asks which accumulator the distinct aggregate gets. That is not the true discriminator.

What decides the outcome is the cost per distinct value on each side. The rewrite materializes one hash table row for each distinct (group keys, x) pair, which is about 48 bytes, plus one accumulator slot for each companion aggregate at that grain. The rewrite wins when the unrewritten accumulator costs more than that for each value. It loses when the accumulator costs less.

The proxy and the true discriminator agree for count(DISTINCT x), which is the only case this PR opens. They disagree elsewhere, and the pre-existing arms of this rule carry that disagreement:

• sum(DISTINCT int_col) and avg(DISTINCT int_col) have no distinct groups accumulator. They still regress 3.15x, at 71.0 MiB unrewritten against 223.4 MiB rewritten, over 4,000,000 rows in 2,000 groups.
• min(DISTINCT x) is worse. min(DISTINCT x) is the same value as min(x), and min_max correctly ignores is_distinct. The unrewritten plan holds one scalar for each group, and the rewrite builds a 4,000,000 row hash table. We measured up to 1113x more peak memory.

That regression predates this PR, and this PR does not extend it. The gate keeps every one of those functions out of the path this PR adds. We report the broader regression upstream separately. A cost model is out of scope here.

What changes are included in this PR?

This PR allows a non-distinct count next to the distinct aggregate, when the distinct aggregate reports that it has no specialized GroupsAccumulator for its argument types.

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. The outer phase adds those partial counts with sum, because a count over a group is the sum of the counts of any partition of that group.

Two details follow from that substitution.

count and sum come from the session function registry, as replace_distinct_aggregate already does for first_value. The rewrite fires only for that exact count, compared by identity and not by name. A session without a registry, or with its own count, keeps the previous behaviour.

count returns a non-null 0 over an empty input, and sum of no rows returns NULL. An aggregate without a GROUP BY reaches that case. The inner aggregate emits no rows and the outer aggregate still emits one row, 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. That restores the 0, and it keeps the type and the nullability that count had.

FILTER and ORDER BY still block the rewrite.

The gate

An optimizer rule has no AccumulatorArgs to ask AggregateUDFImpl::groups_accumulator_supported with. datafusion-optimizer also cannot depend on datafusion-functions-aggregate to read count's list of types. datafusion/optimizer/Cargo.toml names the intended way out:

If you want to add special handling for a specific function, use the methods on the ScalarUDFImpl or AggregateUDFImpl traits (or add a new method to those traits).

So this PR adds AggregateUDFImpl::groups_accumulator_supported_for_types(&[DataType], is_distinct) -> Option<bool>. Count is the only implementor. Count::groups_accumulator_supported now delegates to it, so there is one list of supported types and not two that can drift.

The default is None, which means the implementation does not answer the question. None is not a third answer, and a caller must not read it as either Some(true) or Some(false). This rule rewrites only on Some(false), which is the one answer that reports a call on the adapter. An aggregate that answers None, which today is every aggregate except count, keeps the previous behaviour.

The gate covers only the count this PR adds. A plan that already qualifies through a non-distinct sum, min or max is rewritten as before, over any distinct argument type. A narrower rule there would change plans that this repository has always rewritten, and nothing measured here calls for that change.

Files touched beyond the rule itself:

• datafusion/expr/src/udaf.rs and datafusion/functions-aggregate/src/count.rs: the new trait method and its one implementation.
• datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt: the new coverage described below.
• datafusion/substrait/tests/cases/roundtrip_logical_plan.rs: aggregate_distinct_with_having now builds its session without this rule, so it keeps round tripping the plan shape that the test was written for. Its distinct argument is a DECIMAL, so the gate does not exclude it.

No existing snapshot in the repository changes.

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. A null-handling error or a type error therefore appears as a result mismatch, and not only as a plan difference.

The table carries the same values in a VARCHAR column and in an INT column, which are the two sides of the gate. The file asserts both sides. The VARCHAR distinct rewrites with a count next to it. The INT distinct does not. The INT distinct still rewrites when it qualifies through sum. The results agree with the unoptimized plan in every case.

The file also asserts that an aggregate which answers None stays unrewritten. sum(DISTINCT v) and min(DISTINCT v) next to a count(*) both keep the plan they have on main.

The remaining coverage is count(*) against count(1) against count(col), grouped and ungrouped, a group whose distinct column is entirely NULL, a group with NULLs in both the distinct column and the summed column, empty input in three shapes, HAVING with ORDER BY on the rewritten count, and the production join shape.

The unit tests of the rule cover both sides of the gate directly. They also cover the None answer twice: once with sum(DISTINCT b), and once with a test aggregate that leaves the new method at its default.

Benchmarks

Q22 is the only ClickBench query whose plan this PR can change, and under the gate it no longer does. Its count(DISTINCT "UserID") is over an Int64. Q9 (RegionID, SUM, COUNT(*), AVG, COUNT(DISTINCT UserID)) still bails out, because AVG disqualifies it. Q8, Q10, Q11 and Q13 have a lone distinct aggregate and main already rewrites them. The gate does not touch them.

Latency. Measured on clickbench_partitioned (100 files, about 100M rows) before the gate, when Q22 did change. A run-level A/B cannot resolve a change this small. A comparison of the base binary against itself with compare.py reported 9 queries faster, 28 slower and 6 unchanged, with swings up to 1.58x. A paired comparison over 40 repetitions, with the 42 unchanged queries as an in-experiment control, put Q22 at -2.03%, 95% CI [-5.48%, +1.39%]. There is no measurable latency difference, and this setup resolves effects of about 3% and no better. That agrees with apache#11360, which found the removal of the whole rule to be a wash.

Memory. A memory-limited run (DATAFUSION_RUNTIME_MEMORY_LIMIT: 4G) of the ungated version produced the gate. The peak MemoryPool reservation of Q22 went from 3.1 MiB to 7.3 MiB, which is +132.2%, and Q22 was the only plan that changed. That is the integer case above. The base plan already had PrimitiveDistinctCountGroupsAccumulator, and the rewrite replaced it with an inner group by on (SearchPhrase, UserID) that also carries min("URL") and min("Title") at that grain.

Re-benchmarked with the gate, same benchmark and same 4G limit, three independent runs from one trigger against the same merge base (1, 2, 3):

Run	Q22 base	Q22 changed	Change
1	4.2 MiB	3.1 MiB	-24.4%
2	3.0 MiB	2.6 MiB	-13.8%
3	4.8 MiB	3.0 MiB	-37.3%

The +132.2% is gone. The plan of Q22 is now the base plan byte for byte, so the two sides differ only by measurement noise. Do not read any of it as an improvement. The noise on a 3 MiB high-water mark reported to 0.1 MiB is large, and the base side alone spans 3.0 MiB to 4.8 MiB across three runs of identical code. For scale, Q17 (GROUP BY "UserID", "SearchPhrase", no DISTINCT, so this PR cannot change its plan) moved +6.3%, +26.7% and +6.9% on a 2 GiB reservation in those same three runs. No query moved consistently in the regressing direction.

Those runs are on d510dd4. The commits after it change a doc comment, the treatment of an unanswered gate question, and tests. No ClickBench query holds a distinct sum, min, max or avg, so none of those commits can change a ClickBench plan.

Are there any user-facing changes?

There is no public API change and no change to query results. AggregateUDFImpl gains one method with a default, which is not a breaking change for implementors.

Plans for SELECT ..., count(...), count(DISTINCT x) ... GROUP BY ... change shape when x has no specialized GroupsAccumulator. EXPLAIN output for that shape therefore differs, and such queries should use less memory. Read the size of that win from the second column of the table above, and not from the first.

The size of the win also depends on the group count, because the cost of the adapter is per group. The rule has always had that property. This PR does not change it, and the optimizer has no group cardinality estimate to gate on.

[macroscopeapp]

🟠 High src/single_distinct_to_groupby.rs:288

The rewrite now fails for queries that group by a column named alias2, such as SELECT alias2, count(*), count(DISTINCT b) ... GROUP BY alias2, instead of producing an executable plan. The non-distinct aggregate is always assigned alias2, so the inner aggregate contains both the grouped alias2 field and count(...) AS alias2; DFSchema::check_names rejects these duplicate or ambiguous names. Generate an internal alias that is collision-free with the grouping fields and other generated aliases.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @datafusion/optimizer/src/single_distinct_to_groupby.rs around line 288:

The rewrite now fails for queries that group by a column named `alias2`, such as `SELECT alias2, count(*), count(DISTINCT b) ... GROUP BY alias2`, instead of producing an executable plan. The non-distinct aggregate is always assigned `alias2`, so the inner aggregate contains both the grouped `alias2` field and `count(...) AS alias2`; `DFSchema::check_names` rejects these duplicate or ambiguous names. Generate an internal alias that is collision-free with the grouping fields and other generated aliases.

[macroscopeapp]

🟠 High src/single_distinct_to_groupby.rs:109

A custom UDAF registered under sum causes this rewrite to replace a normal count with that UDAF over partial counts, producing non-count results. try_new verifies only the registered count, so it must also ensure the selected rollup has the built-in additive sum semantics before enabling the rewrite.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @datafusion/optimizer/src/single_distinct_to_groupby.rs around line 109:

A custom UDAF registered under `sum` causes this rewrite to replace a normal `count` with that UDAF over partial counts, producing non-`count` results. `try_new` verifies only the registered `count`, so it must also ensure the selected rollup has the built-in additive `sum` semantics before enabling the rewrite.

[adriangb]

Rebased from the DataFusion 55 fork branch (friendlymatthew/pydantic-main-df55) onto main at da89c7c8, which is byte-identical to current apache/datafusion main, and retargeted this PR to main. That is why the history changed. The four fork-only commits that sat between main and the old base are dropped; the diff is now only this change.

Revalidated on the new base:

• Upstream has not touched single_distinct_to_groupby.rs or the ClickBench Q22 snapshot since the old base, so the rule is unchanged under us.
• All 24 single_distinct optimizer unit tests pass with no snapshot updates, and the full datafusion-optimizer package is green.
• Re-running the sqllogictest runner with --complete over single_distinct_to_groupby.slt and clickbench.slt produced no diff, so the committed Q22 plan is exactly what the new base emits.
• Swept every .slt file that mentions count(distinct) (35 files) plus the full core_integration suite (1103 tests): green.

One new commit was needed. The substrait roundtrip test aggregate_distinct_with_having passed only because the non-distinct count used to make the rule bail out, leaving the plan alias-free. Substrait carries no names for an aggregate's grouping and measure expressions, so the alias1 and alias2 names the rule introduces are lost on the way back and the plan no longer displays identically. This is pre-existing and not specific to this change: a plain SELECT a, count(distinct b) FROM data GROUP BY a, c, which this PR does not affect, fails the same assertion on unmodified main. The test now runs against a session with the rule removed, so it keeps covering the un-rewritten aggregate it was written for.

[adriangb]

Superseded by the upstream pull request: apache#24859

This copy existed only to run the change through review here before sending it upstream. That is done, so closing this one. The branch is unchanged and still backs the upstream pull request.