From b6da6df41b74e7d8c11f7a0d513e7a4762d89283 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 16 Jul 2026 07:16:06 -0400 Subject: [PATCH 1/3] fix: do not mark GROUP BY derived unique keys as nullable GROUP BY treats NULLs as equal, so every key combination -- NULL included -- occurs in exactly one output row, like a primary key. Marking the dependency nullable (multiple NULL keys may coexist) was overly conservative. Co-Authored-By: Claude Fable 5 --- datafusion/common/src/functional_dependencies.rs | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 8b15c49c565f1..91cfc1bcc48fd 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), ); From 2b62f29e42ced6d727f83ffe3c13e69623ff36d8 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 15 Jul 2026 21:09:24 -0400 Subject: [PATCH 2/3] fix: do not remove DISTINCT based on a nullable UNIQUE constraint `ReplaceDistinctWithAggregate` removes `Distinct::All` when a `Dependency::Single` functional dependence covers all input fields. However a nullable UNIQUE constraint also produces a `Single`-mode dependence, and SQL UNIQUE permits multiple NULL keys. DISTINCT treats NULLs as equal, so it must still collapse those rows: CREATE TABLE t (x INT UNIQUE) AS VALUES (NULL), (NULL), (1); SELECT DISTINCT x FROM t; returned two NULL rows. Only remove the DISTINCT when the dependence is non-nullable or none of its source fields are actually nullable, mirroring the guards already used by `EliminateJoin` and `Filter::is_scalar`. Co-Authored-By: Claude Fable 5 --- .../src/replace_distinct_aggregate.rs | 21 +++++++--- .../sqllogictest/test_files/group_by.slt | 41 +++++++++++++++++++ 2 files changed, 57 insertions(+), 5 deletions(-) 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; From f11e19dc25f174e3ff8977abe81ab29ddd34cb38 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 16 Jul 2026 10:16:41 -0400 Subject: [PATCH 3/3] docs: reformat comment on DISTINCT removal conditions Co-Authored-By: Claude Fable 5 --- .../optimizer/src/replace_distinct_aggregate.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/datafusion/optimizer/src/replace_distinct_aggregate.rs b/datafusion/optimizer/src/replace_distinct_aggregate.rs index 6aa96fa7631d3..fe2332ebdc049 100644 --- a/datafusion/optimizer/src/replace_distinct_aggregate.rs +++ b/datafusion/optimizer/src/replace_distinct_aggregate.rs @@ -104,14 +104,17 @@ impl OptimizerRule for ReplaceDistinctWithAggregate { 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. + // is a no-op and we can simply remove it. // - // 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. + // 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()