jayzhan211 commented on code in PR #26000:
URL: https://github.com/apache/datafusion/pull/26000#discussion_r4550162203
##########
datafusion/common/src/config.rs:
##########
@@ -1941,15 +1941,22 @@ config_namespace! {
pub enable_unions_to_filter: bool, default = false
/// Prefer native grouped DISTINCT aggregation when target_partitions
is 1
- /// and every aggregate explicitly supports a specialized
GroupsAccumulator
- /// for its DISTINCT argument types. Defaults to false, preserving the
- /// two-phase GROUP BY rewrite.
+ /// and every DISTINCT argument is a fixed-width, non-float type (such
as
+ /// integers, decimals, dates), instead of rewriting into a two-phase
`GROUP BY`.
+ /// Defaults to false, preserving the two-phase GROUP BY rewrite.
///
- /// The native plan can avoid materializing distinct (group, value)
pairs,
- /// but can be slower for many groups with duplicate values. It must
fit
- /// each group's distinct values in memory, so queries that succeed by
- /// spilling with the rewrite may fail with this option enabled. One
target
- /// partition can still have multiple input partitions and merge
native states.
+ /// The native plan avoids materializing intermediate `(group,
distinct value)`
+ /// pairs, lowering peak memory when no memory limit is set. However,
+ /// native DISTINCT must fit each group's distinct values in memory:
under
+ /// a memory limit, the rewrite can spill those pairs, but the native
+ /// accumulator cannot, so queries that succeed by spilling with the
rewrite
+ /// may fail with this option enabled.
Review Comment:
With many small groups, the native plan needs more memory, not less. At one
target partition on 4M rows with 1M groups, `sum(DISTINCT x)` needs a 384M
limit (no spilling) with this setting on vs 96M without it, and runs 3.8x
slower. `sum` goes through `GroupsAccumulatorAdapter`, which keeps one
accumulator per group. On the #25932 table (2,000 groups) it does halve memory
(192M → 96M), so the doc should say when it helps:
```suggestion
/// The native plan avoids materializing intermediate `(group,
distinct value)`
/// pairs, which lowers peak memory when there are few groups with
many
/// distinct values each. With many small groups it can use more
memory and
/// run slower, because functions without a specialized DISTINCT
/// `GroupsAccumulator` (such as `sum`) keep one accumulator per
group.
/// Native DISTINCT must also fit each group's distinct values in
memory:
/// under a memory limit, the rewrite can spill those pairs, but the
native
/// accumulator cannot, so queries that succeed by spilling with the
rewrite
/// may fail with this option enabled.
```
`configs.md` and `information_schema.slt` need regenerating after this.
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -110,7 +111,56 @@ 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 **and** where bit-identity equals
+/// value-equality, so that skipping the GROUP BY rewrite cannot silently
change
+/// query results.
+///
+/// Float types are deliberately excluded: IEEE 754 defines `-0.0 == 0.0`, so a
+/// `GROUP BY` on floats merges those two bit-patterns into one group. The
+/// native distinct accumulators hash raw bits and therefore count them
+/// separately. Keeping floats in this list would let
`single_distinct_prefer_native`
+/// skip the rewrite for a float column, producing a different answer than the
+/// non-rewritten path — a silent correctness regression.
+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::Date32
+ | DataType::Date64
+ | DataType::Time32(_)
+ | DataType::Time64(_)
+ | DataType::Timestamp(_, _)
+ | DataType::Duration(_)
Review Comment:
With the setting on, `avg(DISTINCT d)` over a `Duration` column now fails
with `AVG(DISTINCT) for (Duration(s) --> Duration(s)) not supported`, while the
rewrite returns rows. `avg` has no DISTINCT accumulator for durations. I ran
every built-in aggregate against every type in this list with the setting on
and off, and this is the only mismatch. Could you drop `Duration` here and add
the test below? It fails without the change.
```suggestion
```
<details>
<summary>Test for single_distinct_to_groupby.slt, after <code>DROP TABLE
t_prefer_mem_float;</code></summary>
```sql
# `avg` has no DISTINCT accumulator for durations, so durations keep the
rewrite.
query I?
SELECT g, avg(DISTINCT d) FROM (SELECT g, arrow_cast(x, 'Duration(Second)')
AS d FROM t_prefer_mem) GROUP BY g ORDER BY g;
----
1 0 days 0 hours 0 mins 15 secs
2 0 days 0 hours 0 mins 10 secs
```
</details>
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -1366,4 +1480,220 @@ mod tests {
assert_eq!(plan.schema().as_ref(), optimized_plan.schema().as_ref());
Ok(())
}
+
+ #[derive(Debug, Clone, PartialEq, Eq, Hash)]
+ struct UnsupportedUdaf {
+ signature: Signature,
+ }
+
+ impl UnsupportedUdaf {
+ fn new() -> Self {
+ Self {
+ signature: Signature::exact(
+ vec![DataType::UInt32],
+ Volatility::Immutable,
+ ),
+ }
+ }
+ }
+
+ impl AggregateUDFImpl for UnsupportedUdaf {
+ fn name(&self) -> &str {
+ "unsupported"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt32)
+ }
+
+ fn accumulator(
+ &self,
+ _acc_args: AccumulatorArgs,
+ ) -> Result<Box<dyn Accumulator>> {
+ unimplemented!("not needed for optimizer tests")
+ }
+
+ fn distinct_handling(&self) -> DistinctHandling {
+ DistinctHandling::Unsupported
+ }
+ }
+
+ fn unsupported_distinct(expr: Expr) -> Expr {
+ Expr::AggregateFunction(AggregateFunction::new_udf(
+ Arc::new(AggregateUDF::from(UnsupportedUdaf::new())),
+ vec![expr],
+ true,
+ None,
+ vec![],
+ None,
+ ))
+ }
Review Comment:
Optional cleanup: `crate::test::udfs::DistinctHandlingTestUDAF` already
provides an aggregate with a chosen `DistinctHandling`, so this struct can go.
The snapshot stays the same:
```suggestion
fn unsupported_distinct(expr: Expr) -> Expr {
Expr::AggregateFunction(AggregateFunction::new_udf(
Arc::new(AggregateUDF::from(udfs::DistinctHandlingTestUDAF::new(
"unsupported",
DistinctHandling::Unsupported,
))),
vec![expr],
true,
None,
vec![],
None,
))
}
```
--
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]
