fix: do not remove DISTINCT based on a nullable UNIQUE constraint

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)

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;