diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 8b15c49c565f1..e0281acd7898f 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -144,6 +144,16 @@ pub struct FunctionalDependence { /// such as after LEFT JOIN or RIGHT JOIN operations, this property may /// change. pub nullable: bool, + /// Whether the source key permits multiple NULL rows with inconsistent + /// dependent values. This is `true` only for dependencies derived from + /// `UNIQUE` constraints (which allow duplicate NULLs). It is `false` for + /// PRIMARY KEY constraints, downgraded PKs (where join padding produces + /// consistent NULL dependents), and GROUP BY derived keys (where at most + /// one NULL group exists). + /// + /// When `true`, the dependency must NOT be used for GROUP BY expansion, + /// GROUP BY/ORDER BY reduction, or DISTINCT removal. + pub duplicate_nulls: bool, // The functional dependency mode: pub mode: Dependency, } @@ -168,6 +178,7 @@ impl FunctionalDependence { source_indices, target_indices, nullable, + duplicate_nulls: false, // Start with the least restrictive mode by default: mode: Dependency::Multi, } @@ -177,6 +188,11 @@ impl FunctionalDependence { self.mode = mode; self } + + pub fn with_duplicate_nulls(mut self, duplicate_nulls: bool) -> Self { + self.duplicate_nulls = duplicate_nulls; + self + } } /// This object encapsulates all functional dependencies in a given relation. @@ -219,7 +235,8 @@ impl FunctionalDependencies { indices.to_vec(), (0..n_field).collect::>(), true, - ), + ) + .with_duplicate_nulls(true), }; // As primary keys are guaranteed to be unique, set the // functional dependency mode to `Dependency::Single`: @@ -301,6 +318,7 @@ impl FunctionalDependencies { source_indices, target_indices, nullable, + duplicate_nulls, mode, } in &self.deps { @@ -321,7 +339,8 @@ impl FunctionalDependencies { new_target_indices, *nullable, ) - .with_mode(*mode); + .with_mode(*mode) + .with_duplicate_nulls(*duplicate_nulls); projected_func_dependencies.push(new_func_dependence); } } @@ -522,9 +541,15 @@ pub fn get_target_functional_dependencies( for FunctionalDependence { source_indices, target_indices, + duplicate_nulls, .. } in &dependencies.deps { + // A dependency that allows duplicate NULLs (from a UNIQUE constraint) + // does not guarantee determination across NULL keys, so skip it. + if *duplicate_nulls { + continue; + } let source_key_names = source_indices .iter() .map(|id_key_idx| &field_names[*id_key_idx]) @@ -546,6 +571,48 @@ pub fn get_target_functional_dependencies( }) } +/// Returns target indices for determinant keys that allow duplicate NULLs +/// (from UNIQUE constraints) and whose source columns are all inside the +/// given group by expressions. These columns are functionally determined by +/// the GROUP BY key when the key is non-NULL, but cannot safely be added to +/// GROUP BY because NULL keys may have multiple distinct target values. +/// They should instead be wrapped in an aggregate like ANY_VALUE. +pub fn get_nullable_unique_target_functional_dependencies( + schema: &DFSchema, + group_by_expr_names: &[String], +) -> Option> { + let mut combined_target_indices = HashSet::new(); + let dependencies = schema.functional_dependencies(); + let field_names = schema.field_names(); + for FunctionalDependence { + source_indices, + target_indices, + duplicate_nulls, + .. + } in &dependencies.deps + { + // Only interested in dependencies that allow duplicate NULLs (UNIQUE). + if !*duplicate_nulls { + continue; + } + let source_key_names = source_indices + .iter() + .map(|id_key_idx| &field_names[*id_key_idx]) + .collect::>(); + if source_key_names + .iter() + .all(|source_key_name| group_by_expr_names.contains(source_key_name)) + { + combined_target_indices.extend(target_indices.iter()); + } + } + (!combined_target_indices.is_empty()).then_some({ + let mut result = combined_target_indices.into_iter().collect::>(); + result.sort(); + result + }) +} + /// Returns indices for the minimal subset of GROUP BY expressions that are /// functionally equivalent to the original set of GROUP BY expressions. pub fn get_required_group_by_exprs_indices( diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index 2eebfe4963057..9e71940a7dea4 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -85,8 +85,10 @@ pub use file_options::file_type::{ }; pub use functional_dependencies::{ Constraint, Constraints, Dependency, FunctionalDependence, FunctionalDependencies, - aggregate_functional_dependencies, get_required_group_by_exprs_indices, - get_required_sort_exprs_indices, get_target_functional_dependencies, + aggregate_functional_dependencies, + get_nullable_unique_target_functional_dependencies, + get_required_group_by_exprs_indices, get_required_sort_exprs_indices, + get_target_functional_dependencies, }; use hashbrown::DefaultHashBuilder; pub use join_type::{JoinConstraint, JoinSide, JoinType}; diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index bdab013144462..5a9eacf0c34ad 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -32,7 +32,10 @@ use crate::utils::{ use arrow::datatypes::DataType; use datafusion_common::error::DataFusionErrorBuilder; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, not_impl_err, plan_err}; +use datafusion_common::{ + Column, DFSchema, DFSchemaRef, Result, + get_nullable_unique_target_functional_dependencies, not_impl_err, plan_err, +}; use datafusion_common::{NullHandling, RecursionUnnestOption, UnnestOptions}; use datafusion_expr::ExprSchemable; use datafusion_expr::builder::get_struct_unnested_columns; @@ -314,6 +317,22 @@ impl SqlToRel<'_, S> { } } + // For UNIQUE constraints with nullable keys, columns that are + // functionally determined cannot be safely added to GROUP BY (because + // NULL keys may map to multiple target values). Instead, wrap them in + // ANY_VALUE so they become valid aggregate expressions and the query + // returns the correct number of groups. + let select_exprs = if !group_by_exprs.is_empty() { + self.wrap_nullable_unique_targets_in_any_value( + &base_plan, + &select_exprs, + &group_by_exprs, + &mut aggr_exprs, + )? + } else { + select_exprs + }; + // Process group by, aggregation or having let AggregatePlanResult { plan, @@ -1146,6 +1165,86 @@ impl SqlToRel<'_, S> { } } + /// For columns that are targets of UNIQUE (nullable) functional + /// dependencies and not already in the GROUP BY or aggregate expressions, + /// wrap them in `ANY_VALUE(col)` so the query returns the correct number + /// of groups instead of incorrectly splitting NULL groups. + /// + /// Returns a (possibly modified) copy of `select_exprs`. Any new aggregate + /// expressions are appended to `aggr_exprs`. + fn wrap_nullable_unique_targets_in_any_value( + &self, + input: &LogicalPlan, + select_exprs: &[Expr], + group_by_exprs: &[Expr], + aggr_exprs: &mut Vec, + ) -> Result> { + let schema = input.schema(); + + // Compute GROUP BY expression names the same way as + // add_group_by_exprs_from_dependencies does. + let group_by_expr_names: Vec = group_by_exprs + .iter() + .map(|e| e.schema_name().to_string()) + .collect(); + + // Find target indices of UNIQUE deps whose sources are in GROUP BY. + let Some(target_indices) = get_nullable_unique_target_functional_dependencies( + schema, + &group_by_expr_names, + ) else { + return Ok(select_exprs.to_vec()); + }; + + // Look up ANY_VALUE aggregate. If not available, fall back to the + // normal validation path (which will produce an error). + let Some(any_value_udf) = self.context_provider.get_aggregate_meta("any_value") + else { + return Ok(select_exprs.to_vec()); + }; + + // Build a set of target field names (qualified) to check against. + let target_field_names: HashSet = target_indices + .iter() + .map(|&idx| { + let (qualifier, field) = schema.qualified_field(idx); + Column::new(qualifier.cloned(), field.name().clone()).flat_name() + }) + .collect(); + + // Rewrite select_exprs: wrap bare column references that are UNIQUE + // dep targets (and not already in GROUP BY) in ANY_VALUE. + let mut new_select_exprs = Vec::with_capacity(select_exprs.len()); + for expr in select_exprs { + if let Expr::Column(col) = expr { + let col_name = col.flat_name(); + if target_field_names.contains(&col_name) + && !group_by_expr_names.contains(&col_name) + { + let any_value_expr = Expr::AggregateFunction( + datafusion_expr::expr::AggregateFunction::new_udf( + Arc::clone(&any_value_udf), + vec![expr.clone()], + false, + None, + vec![], + None, + ), + ); + // Add to aggregate expressions if not already present. + if !aggr_exprs.contains(&any_value_expr) { + aggr_exprs.push(any_value_expr.clone()); + } + // Alias to preserve the original column name in output. + new_select_exprs.push(any_value_expr.alias(col.name())); + continue; + } + } + new_select_exprs.push(expr.clone()); + } + Ok(new_select_exprs) + } + /// Create an aggregate plan. /// /// An aggregate plan consists of grouping expressions, aggregate expressions, an diff --git a/datafusion/sqllogictest/test_files/functional_dependencies.slt b/datafusion/sqllogictest/test_files/functional_dependencies.slt index c49004190dc60..96c2f3d5344cb 100644 --- a/datafusion/sqllogictest/test_files/functional_dependencies.slt +++ b/datafusion/sqllogictest/test_files/functional_dependencies.slt @@ -229,15 +229,15 @@ query II rowsort SELECT x, y FROM t_uniq GROUP BY x; ---- 1 3 -NULL 1 NULL 2 query TT EXPLAIN SELECT x, y FROM t_uniq GROUP BY x; ---- logical_plan -01)Aggregate: groupBy=[[t_uniq.x, t_uniq.y]], aggr=[[]] -02)--TableScan: t_uniq projection=[x, y] +01)Projection: t_uniq.x, any_value(t_uniq.y) AS y +02)--Aggregate: groupBy=[[t_uniq.x]], aggr=[[any_value(t_uniq.y)]] +03)----TableScan: t_uniq projection=[x, y] statement ok