diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 24ca33c0c2c90..8b15c49c565f1 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/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 { diff --git a/datafusion/optimizer/src/replace_distinct_aggregate.rs b/datafusion/optimizer/src/replace_distinct_aggregate.rs index 06df61e766615..cc2616379057a 100644 --- a/datafusion/optimizer/src/replace_distinct_aggregate.rs +++ b/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() diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 7b0d8a00d55ce..de493d6e4a2b1 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/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;