jayzhan211 commented on code in PR #25988:
URL: https://github.com/apache/datafusion/pull/25988#discussion_r4504451656
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -154,12 +158,47 @@ fn is_single_distinct_agg(
if aggregate_count != aggr_expr.len() || fields_set.len() != 1 {
return Ok(false);
}
+ // With one target partition, a native GroupsAccumulator avoids building an
+ // extra aggregate with a row per (group, distinct value) pair. With
multiple
+ // partitions, the rewrite also distributes deduplication by both keys;
+ // keeping DISTINCT instead repartitions lists of values by group and
merges
+ // them in the final accumulator. That can concentrate work on a few
groups.
+ // Even one target partition can have multiple input partitions (e.g. UNION
+ // ALL or an existing partitioned table), requiring Partial/Final state
merging.
+ // This is opt-in: the rule cannot see memory limits or cost the input, and
+ // native DISTINCT must fit each group's distinct values in memory.
+ if use_native_grouped_distinct
+ && distinct_aggs.len() == aggregate_count
+ // An unknown answer must preserve the existing rewrite.
+ && distinct_groups_support(&distinct_aggs, input_schema)?
+ .iter()
+ .all(|support| *support == Some(true))
Review Comment:
Optional: two of these conditions aren't covered by any test. Dropping
`distinct_aggs.len() == aggregate_count`, or changing `== Some(true)` to `!=
Some(false)`, still passes `grouped_count_distinct.slt`. These two EXPLAINs pin
both behaviors from the description (mixed aggregates and unknown support keep
the rewrite). I checked that each fails under its mutation and passes as
written. They can go before `# The global aggregate does not use a
GroupsAccumulator.`:
<details>
<summary>Two EXPLAIN tests (flag on, one target partition)</summary>
```sql
# A non-DISTINCT aggregate alongside keeps the rewrite.
query TT
EXPLAIN SELECT g, count(DISTINCT x), sum(x) FROM numeric_counts GROUP BY g;
----
logical_plan
01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT
numeric_counts.x), sum(alias2) AS sum(numeric_counts.x)
02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1),
sum(alias2)]]
03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.x AS alias1]],
aggr=[[sum(CAST(numeric_counts.x AS Int64)) AS alias2]]
04)------TableScan: numeric_counts projection=[g, x]
# `sum` does not report specialized DISTINCT support, so it keeps the
rewrite.
query TT
EXPLAIN SELECT g, sum(DISTINCT x) FROM numeric_counts GROUP BY g;
----
logical_plan
01)Projection: numeric_counts.g, sum(alias1) AS sum(DISTINCT
numeric_counts.x)
02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[sum(alias1)]]
03)----Aggregate: groupBy=[[numeric_counts.g, CAST(numeric_counts.x AS
Int64) AS alias1]], aggr=[[]]
04)------TableScan: numeric_counts projection=[g, x]
```
</details>
--
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]
