diff --git a/datafusion/optimizer/src/replace_distinct_aggregate.rs b/datafusion/optimizer/src/replace_distinct_aggregate.rs index cc2616379057a..6aa96fa7631d3 100644 --- a/datafusion/optimizer/src/replace_distinct_aggregate.rs +++ b/datafusion/optimizer/src/replace_distinct_aggregate.rs @@ -99,17 +99,28 @@ impl OptimizerRule for ReplaceDistinctWithAggregate { }))); } - let field_count = input.schema().fields().len(); - for dep in input.schema().functional_dependencies().iter() { + let schema = input.schema(); + let field_count = schema.fields().len(); + for dep in schema.functional_dependencies().iter() { // If the input is already unique on all of its columns (e.g. // it is a GROUP BY over exactly these columns), the DISTINCT // is a no-op and we can simply remove it. The dependency mode // must be `Single`: a `Multi` dependence (e.g. a former key // downgraded by a join) means equal rows may occur multiple // times, so the DISTINCT still has work to do. - if dep.mode == Dependency::Single - && dep.source_indices.len() >= field_count - && dep.source_indices[..field_count] + // + // The determinant key must also not contain NULLs: a nullable + // UNIQUE constraint permits multiple NULL keys, but DISTINCT + // treats NULLs as equal and must still collapse them. + let source_indices = &dep.source_indices; + let any_source_field_nullable = source_indices + .iter() + .any(|&idx| schema.field(idx).is_nullable()); + let nullable = dep.nullable && any_source_field_nullable; + if !nullable + && dep.mode == Dependency::Single + && source_indices.len() >= field_count + && source_indices[..field_count] .iter() .enumerate() .all(|(idx, f_idx)| idx == *f_idx) diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index de493d6e4a2b1..f7906c096efc0 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -5694,3 +5694,44 @@ drop table users_with_pk; statement ok drop table user_orders; + +# DISTINCT must not be removed based on a nullable UNIQUE constraint: SQL +# UNIQUE permits multiple NULL keys, but DISTINCT treats NULLs as equal and +# must still collapse them. +statement ok +CREATE TABLE t_nullable_unique (x INT UNIQUE) AS VALUES (NULL), (NULL), (1); + +query I +SELECT DISTINCT x FROM t_nullable_unique ORDER BY x NULLS LAST; +---- +1 +NULL + +query TT +EXPLAIN SELECT DISTINCT x FROM t_nullable_unique; +---- +logical_plan +01)Aggregate: groupBy=[[t_nullable_unique.x]], aggr=[[]] +02)--TableScan: t_nullable_unique projection=[x] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[x@0 as x], aggr=[] +02)--RepartitionExec: partitioning=Hash([x@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[x@0 as x], aggr=[] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +drop table t_nullable_unique; + +# A non-nullable key (PRIMARY KEY) guarantees unique rows, so the DISTINCT +# can still be removed: no Aggregate appears in the plan. +statement ok +CREATE TABLE t_pk (x INT, PRIMARY KEY (x)) AS VALUES (1), (2); + +query TT +EXPLAIN SELECT DISTINCT x FROM t_pk; +---- +logical_plan TableScan: t_pk projection=[x] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +drop table t_pk;