fix: Do not mark GROUP BY derived unique keys as nullable

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::<Vec<_>>();
     // 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::<Vec<_>>();
-        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),
             );

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)

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;