fix: do not remove DISTINCT when a unique key was downgraded by a join

datafusion/common/src/functional_dependencies.rs

@@ -151,8 +151,10 @@ pub struct FunctionalDependence {
 /// Describes functional dependency mode.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Dependency {
-    Single, // A determinant key may occur only once.
-    Multi,  // A determinant key may occur multiple times (in multiple rows).
+    /// A determinant key may occur only once.
+    Single,
+    /// A determinant key may occur multiple times (in multiple rows).
+    Multi,
 }
 
 impl FunctionalDependence {

datafusion/optimizer/src/replace_distinct_aggregate.rs

@@ -22,7 +22,7 @@ use crate::{OptimizerConfig, OptimizerRule};
 use std::sync::Arc;
 
 use datafusion_common::tree_node::Transformed;
-use datafusion_common::{Column, Result};
+use datafusion_common::{Column, Dependency, Result};
 use datafusion_expr::expr_rewriter::normalize_cols;
 use datafusion_expr::utils::expand_wildcard;
 use datafusion_expr::{Aggregate, Distinct, DistinctOn, Expr, LogicalPlan};
@@ -101,9 +101,14 @@ impl OptimizerRule for ReplaceDistinctWithAggregate {
 
                 let field_count = input.schema().fields().len();
                 for dep in input.schema().functional_dependencies().iter() {
-                    // If distinct is exactly the same with a previous GROUP BY, we can
-                    // simply remove it:
-                    if dep.source_indices.len() >= field_count
+                    // 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]
                             .iter()
                             .enumerate()

datafusion/sqllogictest/test_files/group_by.slt

@@ -5641,3 +5641,56 @@ set datafusion.execution.target_partitions = 4;
 
 statement count 0
 drop table t;
+
+# DISTINCT must not be removed when a unique key is downgraded to a
+# non-unique functional dependency by a join: `u.id` is a primary key, but
+# after the LEFT JOIN each `u` row can occur once per matching order.
+statement ok
+CREATE TABLE users_with_pk (id INT, name VARCHAR, primary key(id)) AS VALUES
+  (1, 'alice'),
+  (2, 'bob');
+
+statement ok
+CREATE TABLE user_orders (user_id INT, amount INT) AS VALUES
+  (1, 10),
+  (1, 20),
+  (2, 30);
+
+query I
+SELECT DISTINCT u.id
+  FROM users_with_pk u
+  LEFT JOIN user_orders o ON u.id = o.user_id
+  ORDER BY u.id;
+----
+1
+2
+
+# The DISTINCT must be planned as an Aggregate; it cannot be removed based
+# on the (join-downgraded) primary key of `users_with_pk`.
+query TT
+EXPLAIN SELECT DISTINCT u.id
+  FROM users_with_pk u
+  LEFT JOIN user_orders o ON u.id = o.user_id;
+----
+logical_plan
+01)Aggregate: groupBy=[[u.id]], aggr=[[]]
+02)--Projection: u.id
+03)----Left Join: u.id = o.user_id
+04)------SubqueryAlias: u
+05)--------TableScan: users_with_pk projection=[id]
+06)------SubqueryAlias: o
+07)--------TableScan: user_orders projection=[user_id]
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[]
+02)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=4
+03)----AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[]
+04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+05)--------HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, user_id@0)], projection=[id@0]
+06)----------DataSourceExec: partitions=1, partition_sizes=[1]
+07)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+drop table users_with_pk;
+
+statement ok
+drop table user_orders;