jayzhan211 commented on code in PR #25988:
URL: https://github.com/apache/datafusion/pull/25988#discussion_r4499130540
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -154,12 +155,41 @@ 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.
+ // Native accumulator support alone does not establish that this is
cheaper.
+ if use_native_grouped_distinct
Review Comment:
@QinXi-ai Thanks for working on this. With spilling on, this turns queries
that finish under a memory limit into `ResourcesExhausted`: the rewrite dedupes
through sorted `(g, x)` spill runs, but the native accumulator has to hold
every distinct value of the largest group in memory. ClickBench Q8 at one
partition finishes at `-m 16M` on main and fails at `-m 128M` here (it passes
at 256M). The equal-memory table in #25936 used `-d 0`, where neither plan can
spill.
```sql
SET datafusion.execution.target_partitions = 1;
SET datafusion.execution.batch_size = 128;
SET datafusion.runtime.memory_limit = '1M';
-- main: 2, 100000. This PR: Resources exhausted
SELECT count(*), sum(d) FROM (
SELECT v % 2 AS k, count(DISTINCT v) AS d
FROM generate_series(1, 100000) AS t(v)
GROUP BY v % 2
);
```
The rule can't see the memory pool, so it can't choose per query. Could you
make the native plan opt-in so the rewrite stays the default? #26000 adds a
very similar flag to this rule, so it's worth settling on one.
```rs
!group_expr.is_empty()
&& config.options().execution.target_partitions == 1
&& config.options().optimizer.single_distinct_prefer_native
```
Please add the query above to `aggregate_memory_spill.slt` as the flag-off
test, plus an `EXPLAIN` with the flag on showing the native plan. If a
maintainer would rather keep it on by default, it needs an upgrade guide note,
and the "Query results ... are unchanged" line in the description needs
correcting.
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -154,12 +155,41 @@ 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.
+ // Native accumulator support alone does not establish that this is
cheaper.
+ if use_native_grouped_distinct
+ && distinct_aggs.len() == aggregate_count
+ && all_have_groups_accumulators(&distinct_aggs, input_schema)?
+ {
+ return Ok(false);
+ }
if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs,
input_schema)? {
return Ok(false);
}
Ok(true)
}
+fn all_have_groups_accumulators(
Review Comment:
Optional cleanup: this repeats the loop in `rewrite_pays_for_count`, and the
two only differ in which answer they look for. One helper covers both:
```diff
-fn all_have_groups_accumulators(
+/// Each distinct aggregate's answer to whether a specialized
`GroupsAccumulator`
+/// exists for its argument types.
+fn distinct_groups_support(
distinct_aggs: &[(&Arc<AggregateUDF>, &[Expr])],
input_schema: &DFSchema,
-) -> Result<bool> {
- for (func, args) in distinct_aggs {
- let arg_types = args
- .iter()
- .map(|arg| arg.get_type(input_schema))
- .collect::<Result<Vec<_>>>()?;
- // An unknown answer must preserve the existing rewrite.
- if func.groups_accumulator_supported_for_types(&arg_types, true) !=
Some(true) {
- return Ok(false);
- }
- }
- Ok(true)
+) -> Result<Vec<Option<bool>>> {
+ distinct_aggs
+ .iter()
+ .map(|(func, args)| {
+ let arg_types = args
+ .iter()
+ .map(|arg| arg.get_type(input_schema))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(func.groups_accumulator_supported_for_types(&arg_types,
true))
+ })
+ .collect()
}
```
The call here becomes `distinct_groups_support(&distinct_aggs,
input_schema)?.iter().all(|support| *support == Some(true))`, and
`rewrite_pays_for_count` becomes `Ok(distinct_groups_support(distinct_aggs,
input_schema)?.contains(&Some(false)))`.
##########
benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql:
##########
@@ -0,0 +1,11 @@
+CREATE TABLE grouped_distinct AS
+SELECT value % 2000 AS g,
+ value % 8 AS few_g,
+ (value * 48271) % 999983 AS x,
+ (value / 2000) % 8 AS low_x,
+ arrow_cast((value * 48271) % 999983, 'UInt64') AS unsigned_x,
+ CAST((value * 48271) % 999983 AS VARCHAR) AS string_x,
+ CAST((value * 48271) % 999983 AS DOUBLE) AS float_x,
+ CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS skew_g,
+ CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR)
AS string_g
Review Comment:
The suite doesn't include the shape where the native plan loses: many groups
with duplicate rows. With 1M groups and four rows per group, I measured it
about 1.4x slower than main (about 130 ms vs 93 ms median, on a busy laptop),
because every row probes both the group table and `seen`, while the rewrite's
outer aggregate only sees each pair once. The rule has no statistics to avoid
this, so could you add a `q10` over these columns and state this worst case in
the description?
```suggestion
CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR)
AS string_g,
value % 1000000 AS many_g,
value % 2 AS dup_x
```
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -231,6 +261,8 @@ impl OptimizerRule for SingleDistinctToGroupBy {
&aggr_expr,
input.schema(),
count_rollup.as_ref(),
+ !group_expr.is_empty()
+ && config.options().execution.target_partitions == 1,
Review Comment:
`target_partitions == 1` doesn't guarantee one input partition: a `UNION
ALL` or a table created at higher parallelism still plans Partial, Coalesce,
Final, which is the list-merge path this PR avoids for parallel plans. A
4-partition table with a hot group ran about 1.3x slower than main in my runs,
since Partial dedupes first and Final rebuilds the set. Could you add this test
before the switch to four partitions, so that path is covered, and mention the
case in the comment above `use_native_grouped_distinct`?
```sql
query II rowsort
SELECT g, count(DISTINCT x)
FROM (SELECT g, x FROM numeric_counts UNION ALL SELECT g, x FROM
numeric_counts)
GROUP BY g;
----
1 2
2 0
3 1
```
--
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]
