jayzhan211 commented on code in PR #26000:
URL: https://github.com/apache/datafusion/pull/26000#discussion_r4512432631
##########
datafusion/common/src/config.rs:
##########
@@ -1940,6 +1940,23 @@ config_namespace! {
/// projections or aliases.
pub enable_unions_to_filter: bool, default = false
+ /// When `true`, the `SingleDistinctToGroupBy` optimizer rule applies
its cost guard
+ /// universally — including to lone-distinct queries that have no
companion non-distinct
+ /// aggregates. The rewrite is skipped whenever the distinct function
already has a
+ /// native `GroupsAccumulator`, reducing peak memory by avoiding the
larger intermediate
+ /// `(group, distinct value)` result set that the two-phase plan
produces.
+ ///
+ /// The trade-off is reduced parallelism: skipping the rewrite causes
the engine to fall
+ /// back to serial per-group accumulation, which is typically 2–3×
slower for queries
+ /// with high-cardinality distinct values.
+ ///
+ /// When `false` (the default), lone-distinct queries are always
rewritten into the
+ /// two-phase `GROUP BY` so they can fan out across partitions and
achieve maximum
+ /// throughput.
+ ///
+ /// See <https://github.com/apache/datafusion/issues/25932>
Review Comment:
Under a memory limit, this setting needs more memory, not less: at 4 target
partitions, `sum(DISTINCT x)` on the #25932 table needs 16M with the setting on
and finishes within 4M without it (`count`: 48M vs 4M). The rewrite can spill
its `(group, value)` pairs, but the native accumulator keeps each group's
distinct values in memory, so with a few large groups it fails where the
default finishes. Your numbers are right for a pool with no limit, so the doc
should say that:
```suggestion
/// When `true`, the `SingleDistinctToGroupBy` optimizer rule keeps
a grouped
/// DISTINCT aggregate native, instead of rewriting it into a
two-phase `GROUP BY`,
/// when every DISTINCT argument is a fixed-width, non-float type.
This lowers peak
/// memory when no memory limit is set, because the intermediate
/// `(group, distinct value)` result is never built.
///
/// Under a memory limit, the rewrite can spill those pairs, but the
native
/// accumulator keeps each group's distinct values in memory. With
this option
/// enabled, a query can need more memory, or fail where the rewrite
finishes.
///
/// Queries without `GROUP BY`, and other argument types, keep the
rewrite.
///
/// See <https://github.com/apache/datafusion/issues/25932>
```
<details>
<summary>Smallest memory limit that finishes (release datafusion-cli, greedy
pool, spilling on)</summary>
| query on the #25932 table (4M rows, 2,000 groups) | target partitions |
setting off | setting on |
|---|---|---|---|
| `sum(DISTINCT x)` | 1 | ≤4M | ≤4M (also 1M) |
| `sum(DISTINCT x)` | 4 | ≤4M | 16M |
| `count(DISTINCT x)` | 1 | ≤4M | ≤4M |
| `count(DISTINCT x)` | 4 | ≤4M | 48M |
Two large groups at a 1M limit (`SELECT v % 2, count(DISTINCT v) FROM
generate_series(1, 100000) AS t(v) GROUP BY v % 2`, `batch_size = 128`): the
setting fails with `Resources exhausted` at 1 target partition for `count` and
`sum`, and at 4 for `count`. The default finishes in every case. The same query
is in `aggregate_memory_spill.slt`.
</details>
`configs.md` and `information_schema.slt` need regenerating afterwards.
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -110,12 +111,50 @@ impl CountRollup {
}
}
-/// Check whether all aggregate exprs are distinct on a single field.
+/// Returns `true` for fixed-width primitive types where single-phase distinct
+/// accumulation avoids heap allocation churn.
+fn is_fixed_width_type(dt: &DataType) -> bool {
+ matches!(
+ dt,
+ DataType::Boolean
+ | DataType::Int8
+ | DataType::Int16
+ | DataType::Int32
+ | DataType::Int64
+ | DataType::UInt8
+ | DataType::UInt16
+ | DataType::UInt32
+ | DataType::UInt64
+ | DataType::Float16
+ | DataType::Float32
+ | DataType::Float64
Review Comment:
With the setting on, float `DISTINCT` gives different answers: on `(0.0,
-0.0, 1.5)`, `count(DISTINCT f)` goes from 2 to 3 and `avg(DISTINCT f)` from
0.75 to 0.5. The native accumulators compare raw bits (`Hashable` in
`functions-aggregate-common/src/utils.rs`), so `-0.0` and `0.0` count as two
values, while the rewrite's `GROUP BY` treats them as one. Could you keep
floats on the rewrite and add the test below? It fails without this change.
```suggestion
```
<details>
<summary>Regression test for single_distinct_to_groupby.slt (place it while
the setting is on)</summary>
```sql
# Float DISTINCT keeps the rewrite: the native accumulator counts -0.0 and
0.0 separately.
statement ok
CREATE TABLE t_prefer_mem_float AS SELECT * FROM (VALUES (1, CAST('0.0' AS
DOUBLE)), (1, CAST('-0.0' AS DOUBLE)), (1, CAST('1.5' AS DOUBLE))) AS v(g, f);
query IIR
SELECT g, count(DISTINCT f), avg(DISTINCT f) FROM t_prefer_mem_float GROUP
BY g;
----
1 2 0.75
statement ok
DROP TABLE t_prefer_mem_float;
```
</details>
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -158,6 +197,15 @@ fn is_single_distinct_agg(
if aggregate_count != aggr_expr.len() || fields_set.len() != 1 {
return Ok(false);
}
+ // Aggregates with DistinctHandling::Unsupported (e.g. approx_median)
cannot
+ // run natively with DISTINCT; they must always be rewritten.
+ if distinct_aggs
+ .iter()
+ .any(|(func, _)| func.distinct_handling() ==
DistinctHandling::Unsupported)
+ {
+ return Ok(true);
+ }
Review Comment:
With the setting off, this return also changes main's behavior: `SELECT g,
count(*), stddev(DISTINCT x) FROM t GROUP BY g` fails on main with
`STDDEV_POP(DISTINCT) aggregations are not available` and returns the right
rows here. That's a fix, but no test covers it at the default, and the
description says the default is unchanged. Could you add this test before the
`SET … = true` line in `single_distinct_to_groupby.slt` and mention the fix in
the description?
```sql
# stddev cannot run DISTINCT itself, so it is rewritten even next to
count(*).
query IIR rowsort
SELECT g, count(*), stddev(DISTINCT x) FROM t_prefer_mem GROUP BY g;
----
1 2 7.071067811865
2 1 NULL
```
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -176,6 +224,27 @@ fn is_single_distinct_agg(
{
return Ok(false);
}
+
+ // When prefer_memory is set, skip the rewrite only when it actually helps:
+ // - The query has a GROUP BY. Without it, the 2-phase rewrite is the
only
+ // way to parallelize un-grouped distinct work across threads;
disabling it
+ // forces everything onto a single partition (1.6x–3.4x slowdown).
+ // - The distinct argument is a fixed-width primitive type. For
variable-length
+ // types (Utf8, Binary, …) single-phase accumulators suffer heavy
per-value
+ // heap allocation churn that increases both memory and latency.
+ if prefer_memory && has_group_by {
+ // Resolve the data type of each distinct argument.
+ let all_fixed_width = distinct_aggs.iter().all(|(_, args)| {
+ args.iter().all(|arg| {
+ arg.get_type(input_schema)
+ .map(|dt| is_fixed_width_type(&dt))
+ .unwrap_or(false)
+ })
+ });
+ if all_fixed_width {
Review Comment:
Optional cleanup: `.unwrap_or(false)` silently keeps the rewrite when a type
can't be resolved, while `distinct_groups_support` above returns the error with
`?`. This does the same with one pass:
```suggestion
let arg_types = distinct_aggs
.iter()
.flat_map(|(_, args)| args.iter())
.map(|arg| arg.get_type(input_schema))
.collect::<Result<Vec<_>>>()?;
if arg_types.iter().all(is_fixed_width_type) {
```
##########
datafusion/common/src/config.rs:
##########
@@ -1940,6 +1940,23 @@ config_namespace! {
/// projections or aliases.
pub enable_unions_to_filter: bool, default = false
+ /// When `true`, the `SingleDistinctToGroupBy` optimizer rule applies
its cost guard
+ /// universally — including to lone-distinct queries that have no
companion non-distinct
+ /// aggregates. The rewrite is skipped whenever the distinct function
already has a
+ /// native `GroupsAccumulator`, reducing peak memory by avoiding the
larger intermediate
+ /// `(group, distinct value)` result set that the two-phase plan
produces.
+ ///
+ /// The trade-off is reduced parallelism: skipping the rewrite causes
the engine to fall
+ /// back to serial per-group accumulation, which is typically 2–3×
slower for queries
+ /// with high-cardinality distinct values.
+ ///
+ /// When `false` (the default), lone-distinct queries are always
rewritten into the
+ /// two-phase `GROUP BY` so they can fan out across partitions and
achieve maximum
+ /// throughput.
+ ///
+ /// See <https://github.com/apache/datafusion/issues/25932>
+ pub single_distinct_aggregation_prefer_memory: bool, default = false
Review Comment:
#25988 (merged today, not released yet) added
`single_distinct_prefer_native` just below for the same choice, and with both
settings `count(DISTINCT int)` at one target partition skips the rewrite under
either one. Could this PR widen that setting's gate instead of adding a second
setting? Because it hasn't been released, its name and doc can still change,
and users would get one switch for "keep grouped DISTINCT native".
--
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]
