fix: GROUP BY returns two rows for the NULL group with a nullable UNIQUE constraint

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::<Vec<_>>(),
                             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<Vec<usize>> {
+    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::<Vec<_>>();
+        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::<Vec<_>>();
+        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(

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};

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<S: ContextProvider> 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<S: ContextProvider> 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<Expr>,
+    ) -> Result<Vec<Expr>> {
+        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<String> = 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<String> = 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

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