diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 8b15c49c565f..91cfc1bcc48f 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -424,7 +424,6 @@ pub fn aggregate_functional_dependencies( ) -> FunctionalDependencies { let mut aggregate_func_dependencies = vec![]; let aggr_input_fields = aggr_input_schema.field_names(); - let aggr_fields = aggr_schema.fields(); // Association covers the whole table: let target_indices = (0..aggr_schema.fields().len()).collect::>(); // Get functional dependencies of the schema: @@ -486,9 +485,6 @@ pub fn aggregate_functional_dependencies( if !group_by_expr_names.is_empty() { let count = group_by_expr_names.len(); let source_indices = (0..count).collect::>(); - let nullable = source_indices - .iter() - .any(|idx| aggr_fields[*idx].is_nullable()); // If GROUP BY expressions do not already act as a determinant: if !aggregate_func_dependencies.iter().any(|item| { // If `item.source_indices` is a subset of GROUP BY expressions, we shouldn't add @@ -498,10 +494,15 @@ pub fn aggregate_functional_dependencies( // GROUP BY expressions come here as a prefix. item.source_indices.iter().all(|idx| idx < &count) }) { - // Add a new functional dependency associated with the whole table: - // Use nullable property of the GROUP BY expression: + // Add a new functional dependency associated with the whole table. + // + // `nullable` is `false`: a nullable dependence means multiple NULL + // keys may coexist (as under a SQL UNIQUE constraint, where NULLs + // compare distinct). That cannot happen after grouping: GROUP BY + // treats NULLs as equal, so every key combination -- NULL included + // -- occurs in exactly one output row, like a primary key. + let nullable = false; aggregate_func_dependencies.push( - // Use nullable property of the GROUP BY expression: FunctionalDependence::new(source_indices, target_indices, nullable) .with_mode(Dependency::Single), ); diff --git a/datafusion/optimizer/src/replace_distinct_aggregate.rs b/datafusion/optimizer/src/replace_distinct_aggregate.rs index cc2616379057..fe2332ebdc04 100644 --- a/datafusion/optimizer/src/replace_distinct_aggregate.rs +++ b/datafusion/optimizer/src/replace_distinct_aggregate.rs @@ -99,17 +99,31 @@ 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] + // 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. + // + // The grouping columns must also not contain NULLs because + // 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 de493d6e4a2b..f7906c096efc 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;