fix: UNIQUE constraint with NULLs incorrectly collapses GROUP BY groups

Author: xiedeyantu

datafusion/common/src/functional_dependencies.rs

@@ -422,7 +422,6 @@ pub fn aggregate_functional_dependencies(
 ) -> FunctionalDependencies {
     let mut aggregate_func_dependencies = vec![];
     let aggr_input_fields = aggr_input_schema.field_names();
-    let aggr_fields = aggr_schema.fields();
     // Association covers the whole table:
     let target_indices = (0..aggr_schema.fields().len()).collect::<Vec<_>>();
     // Get functional dependencies of the schema:
@@ -484,9 +483,12 @@ pub fn aggregate_functional_dependencies(
     if !group_by_expr_names.is_empty() {
         let count = group_by_expr_names.len();
         let source_indices = (0..count).collect::<Vec<_>>();
-        let nullable = source_indices
-            .iter()
-            .any(|idx| aggr_fields[*idx].is_nullable());
+        // Aggregation with GROUP BY always produces unique output rows for
+        // each distinct combination of GROUP BY keys. The nullable flag is
+        // set to false here so that subsequent expansion (e.g. a second
+        // GROUP BY on the aggregate output) is never blocked by source
+        // field nullability.
+        let nullable = false;
         // If GROUP BY expressions do not already act as a determinant:
         if !aggregate_func_dependencies.iter().any(|item| {
             // If `item.source_indices` is a subset of GROUP BY expressions, we shouldn't add
@@ -565,9 +567,17 @@ pub fn get_required_group_by_exprs_indices(
     for FunctionalDependence {
         source_indices,
         target_indices,
+        nullable,
         ..
     } in &dependencies.deps
     {
+        // Skip nullable dependencies: UNIQUE constraints allow NULL values,
+        // and NULLs are not considered equal in SQL, so two rows with NULL in
+        // the source key are NOT in the same group. We therefore cannot use a
+        // nullable FD to eliminate other GROUP BY columns.
+        if *nullable {
+            continue;
+        }
         if source_indices
             .iter()
             .all(|source_idx| groupby_expr_indices.contains(source_idx))

datafusion/expr/src/logical_plan/builder.rs

@@ -54,10 +54,9 @@ use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::metadata::FieldMetadata;
 use datafusion_common::{
-    Column, Constraints, DFSchema, DFSchemaRef, NullEquality, Result, ScalarValue,
-    TableReference, ToDFSchema, UnnestOptions, exec_err,
-    get_target_functional_dependencies, internal_datafusion_err, plan_datafusion_err,
-    plan_err,
+    Column, Constraints, DFSchema, DFSchemaRef, Dependency, NullEquality, Result, ScalarValue,
+    TableReference, ToDFSchema, UnnestOptions, exec_err, internal_datafusion_err,
+    plan_datafusion_err, plan_err,
 };
 use datafusion_expr_common::type_coercion::binary::type_union_resolution;
 
@@ -1853,15 +1852,40 @@ pub fn add_group_by_exprs_from_dependencies(
         .map(|e| e.schema_name().to_string())
         .collect::<Vec<_>>();
 
-    if let Some(target_indices) =
-        get_target_functional_dependencies(schema, &group_by_field_names)
-    {
-        for idx in target_indices {
-            let expr = Expr::Column(Column::from(schema.qualified_field(idx)));
-            let expr_name = expr.schema_name().to_string();
-            if !group_by_field_names.contains(&expr_name) {
-                group_by_field_names.push(expr_name);
-                group_expr.push(expr);
+    let field_names = schema.field_names();
+    for dependence in schema.functional_dependencies().iter() {
+        // SQL UNIQUE constraints (mode == Single) allow multiple NULLs, so a
+        // nullable UNIQUE key on a nullable source column cannot safely be
+        // used to synthesize additional GROUP BY keys. Doing so could split
+        // NULL groups and change aggregate results.
+        // Join-derived dependencies (mode == Multi) represent downgraded PKs
+        // from outer joins and are safe to expand even when nullable.
+        if dependence.nullable
+            && dependence.mode == Dependency::Single
+            && dependence
+                .source_indices
+                .iter()
+                .any(|&source_idx| schema.field(source_idx).is_nullable())
+        {
+            continue;
+        }
+
+        let source_key_names = dependence
+            .source_indices
+            .iter()
+            .map(|idx| &field_names[*idx])
+            .collect::<Vec<_>>();
+        if source_key_names
+            .iter()
+            .all(|source_key_name| group_by_field_names.contains(source_key_name))
+        {
+            for idx in &dependence.target_indices {
+                let expr = Expr::Column(Column::from(schema.qualified_field(*idx)));
+                let expr_name = expr.schema_name().to_string();
+                if !group_by_field_names.contains(&expr_name) {
+                    group_by_field_names.push(expr_name);
+                    group_expr.push(expr);
+                }
             }
         }
     }
@@ -2887,6 +2911,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn plan_builder_aggregate_does_not_expand_nullable_unique_group_by_exprs()
+    -> Result<()> {
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, true),
+            Field::new("state", DataType::Utf8, false),
+            Field::new("salary", DataType::Int32, false),
+        ]);
+        let constraints = Constraints::new_unverified(vec![Constraint::Unique(vec![0])]);
+        let table_source = table_source_with_constraints(&schema, constraints);
+
+        let options =
+            LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
+        let plan = LogicalPlanBuilder::scan("employee_csv", table_source, None)?
+            .with_options(options)
+            .aggregate(vec![col("id")], vec![sum(col("salary"))])?
+            .build()?;
+
+        assert_snapshot!(plan, @r"
+        Aggregate: groupBy=[[employee_csv.id]], aggr=[[sum(employee_csv.salary)]]
+          TableScan: employee_csv
+        ");
+
+        Ok(())
+    }
+
     #[test]
     fn test_join_metadata() -> Result<()> {
         let left_schema = DFSchema::new_with_metadata(

datafusion/expr/src/logical_plan/plan.rs

@@ -351,7 +351,6 @@ impl LogicalPlan {
             LogicalPlan::Ddl(ddl) => ddl.schema(),
             LogicalPlan::Unnest(Unnest { schema, .. }) => schema,
             LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => {
-                // we take the schema of the static term as the schema of the entire recursive query
                 static_term.schema()
             }
         }
@@ -2402,6 +2401,10 @@ impl SubqueryAlias {
         // no field must share the same column name as this would lead to ambiguity when referencing
         // columns in parent logical nodes.
 
+        // Capture whether the input is a RecursiveQuery before `plan` may be
+        // rebound to a wrapping Projection below.
+        let is_recursive_query = matches!(plan.as_ref(), LogicalPlan::RecursiveQuery(_));
+
         // Compute unique aliases, if any, for each column of the input's schema.
         let aliases = unique_field_aliases(plan.schema().fields());
         let is_projection_needed = aliases.iter().any(Option::is_some);
@@ -2431,7 +2434,14 @@ impl SubqueryAlias {
         // Requalify fields with the new `alias`.
         let fields = plan.schema().fields().clone();
         let meta_data = plan.schema().metadata().clone();
-        let func_dependencies = plan.schema().functional_dependencies().clone();
+        // Recursive queries do not expose the anchor's functional dependencies to
+        // the outer schema — the recursive term can produce rows that violate
+        // those dependencies, so they are intentionally dropped here.
+        let func_dependencies = if is_recursive_query {
+            FunctionalDependencies::empty()
+        } else {
+            plan.schema().functional_dependencies().clone()
+        };
 
         let schema = DFSchema::from_unqualified_fields(fields, meta_data)?;
         let schema = schema.as_arrow();

datafusion/expr/src/logical_plan/tree_node.rs

@@ -329,6 +329,7 @@ impl TreeNode for LogicalPlan {
                 static_term,
                 recursive_term,
                 is_distinct,
+                ..
             }) => (static_term, recursive_term).map_elements(f)?.update_data(
                 |(static_term, recursive_term)| {
                     LogicalPlan::RecursiveQuery(RecursiveQuery {

datafusion/sql/src/select.rs

@@ -29,15 +29,16 @@ use crate::utils::{
 
 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, Dependency, Result, not_impl_err, plan_err};
 use datafusion_common::{RecursionUnnestOption, UnnestOptions};
 use datafusion_expr::expr::{PlannedReplaceSelectItem, WildcardOptions};
 use datafusion_expr::expr_rewriter::{
     normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts,
 };
 use datafusion_expr::select_expr::SelectExpr;
 use datafusion_expr::utils::{
-    expr_as_column_expr, expr_to_columns, find_aggregate_exprs, find_window_exprs,
+    expr_as_column_expr, expr_to_columns, find_aggregate_exprs,
+    find_window_exprs,
 };
 use datafusion_expr::{
     Aggregate, Expr, Filter, GroupingSet, LogicalPlan, LogicalPlanBuilder,
@@ -966,12 +967,15 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         group_by_exprs: &[Expr],
         aggr_exprs: &[Expr],
     ) -> Result<AggregatePlanResult> {
+        let group_by_exprs = self.add_required_group_by_exprs_from_dependencies(
+            input,
+            group_by_exprs,
+        )?;
+
         // create the aggregate plan
-        let options =
-            LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
         let plan = LogicalPlanBuilder::from(input.clone())
-            .with_options(options)
-            .aggregate(group_by_exprs.to_vec(), aggr_exprs.to_vec())?
+            .with_options(LogicalPlanBuilderOptions::new())
+            .aggregate(group_by_exprs.clone(), aggr_exprs.to_vec())?
             .build()?;
         let group_by_exprs = if let LogicalPlan::Aggregate(agg) = &plan {
             &agg.group_expr
@@ -1138,6 +1142,74 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         })
     }
 
+    /// Expand `group_by_exprs` with additional columns that are functionally
+    /// determined by the current GROUP BY keys, using the functional
+    /// dependencies recorded on the input schema.
+    ///
+    /// A dependency is skipped when **all** of the following are true:
+    /// - `nullable = true` — the constraint allows NULLs (i.e. SQL `UNIQUE`,
+    ///   not `PRIMARY KEY`)
+    /// - `mode == Single` — the dependency comes from a single-table scan,
+    ///   not from a join.  Join-derived dependencies (`mode == Multi`) are
+    ///   downgraded PKs and remain safe to expand.
+    /// - At least one source field is nullable in the schema.
+    ///
+    /// Together these conditions identify nullable `UNIQUE` keys on nullable
+    /// columns, where multiple NULLs are permitted and `key → other_cols`
+    /// does not hold for the NULL group.  Expanding such a key would split
+    /// the NULL group and produce incorrect aggregates.
+    fn add_required_group_by_exprs_from_dependencies(
+        &self,
+        input: &LogicalPlan,
+        group_by_exprs: &[Expr],
+    ) -> Result<Vec<Expr>> {
+        let schema = input.schema();
+        let field_names = schema.field_names();
+        let mut result = group_by_exprs.to_vec();
+        let mut group_by_field_names: Vec<String> =
+            result.iter().map(|e| e.schema_name().to_string()).collect();
+
+        for dependence in schema.functional_dependencies().iter() {
+            // SQL UNIQUE constraints (mode == Single) allow multiple NULLs,
+            // so a nullable UNIQUE key on a nullable source column cannot
+            // safely synthesise additional GROUP BY keys — it could split
+            // NULL groups and change aggregate results.
+            // Join-derived dependencies (mode == Multi) represent downgraded
+            // PKs from outer joins and are safe to expand even when nullable.
+            if dependence.nullable
+                && dependence.mode == Dependency::Single
+                && dependence
+                    .source_indices
+                    .iter()
+                    .any(|&source_idx| schema.field(source_idx).is_nullable())
+            {
+                continue;
+            }
+
+            let source_key_names = dependence
+                .source_indices
+                .iter()
+                .map(|idx| &field_names[*idx])
+                .collect::<Vec<_>>();
+            if source_key_names
+                .iter()
+                .all(|n| group_by_field_names.contains(n))
+            {
+                for idx in &dependence.target_indices {
+                    let expr =
+                        Expr::Column(Column::from(schema.qualified_field(*idx)));
+                    let expr_name = expr.schema_name().to_string();
+                    if !group_by_field_names.contains(&expr_name) {
+                        group_by_field_names.push(expr_name);
+                        result.push(expr);
+                    }
+                }
+            }
+        }
+
+        Ok(result)
+    }
+
     // If the projection is done over a named window, that window
     // name must be defined. Otherwise, it gives an error.
     fn match_window_definitions(

datafusion/sqllogictest/test_files/group_by.slt

@@ -3565,33 +3565,27 @@ SELECT r.sn, r.amount, SUM(r.amount)
   GROUP BY r.sn
 ORDER BY r.sn
 
-# left semi join should propagate constraint of left side as is.
-query IRR
+# left semi join with a nullable UNIQUE key cannot safely propagate the
+# constraint for expansion, because UNIQUE allows multiple NULLs.
+statement error DataFusion error: Error during planning: Column in SELECT must be in GROUP BY or an aggregate function: While expanding wildcard, column "l\.amount" must appear in the GROUP BY clause or must be part of an aggregate function, currently only "l\.sn, sum\(l\.amount\)" appears in the SELECT clause satisfies this requirement
 SELECT l.sn, l.amount, SUM(l.amount)
   FROM (SELECT *
     FROM sales_global_with_unique as l
     LEFT SEMI JOIN sales_global_with_unique as r
     ON l.amount >= r.amount + 10)
   GROUP BY l.sn
 ORDER BY l.sn
-----
-1 50 50
-2 75 75
-3 200 200
-4 100 100
-NULL 100 100
 
-# Similarly, left anti join should propagate constraint of left side as is.
-query IRR
+# Similarly, left anti join with a nullable UNIQUE key cannot safely propagate
+# the constraint for expansion.
+statement error DataFusion error: Error during planning: Column in SELECT must be in GROUP BY or an aggregate function: While expanding wildcard, column "l\.amount" must appear in the GROUP BY clause or must be part of an aggregate function, currently only "l\.sn, sum\(l\.amount\)" appears in the SELECT clause satisfies this requirement
 SELECT l.sn, l.amount, SUM(l.amount)
   FROM (SELECT *
     FROM sales_global_with_unique as l
     LEFT ANTI JOIN sales_global_with_unique as r
     ON l.amount >= r.amount + 10)
   GROUP BY l.sn
 ORDER BY l.sn
-----
-0 30 30
 
 # Should support grouping by list column
 query ?I
@@ -5641,3 +5635,32 @@ set datafusion.execution.target_partitions = 4;
 
 statement count 0
 drop table t;
+
+# Test that GROUP BY with a UNIQUE constraint does not incorrectly collapse
+# NULL rows. UNIQUE allows multiple NULLs (NULLs are not equal in SQL), so
+# a UNIQUE column cannot be used to eliminate other GROUP BY columns.
+# Regression test for https://github.com/apache/datafusion/issues/21507
+
+statement ok
+CREATE TABLE t_unique_null(a INT, b INT, c INT, UNIQUE(a));
+
+statement ok
+INSERT INTO t_unique_null VALUES (1, 10, 100), (NULL, 20, 200), (NULL, 30, 300);
+
+# The two NULL rows must stay in separate groups (grouped by b as well).
+query II rowsort
+SELECT a, SUM(c) AS total FROM t_unique_null GROUP BY a, b;
+----
+1 100
+NULL 200
+NULL 300
+
+# GROUP BY on the UNIQUE column alone must still merge the NULL rows into one group.
+query II rowsort
+SELECT a, SUM(c) AS total FROM t_unique_null GROUP BY a;
+----
+1 100
+NULL 500
+
+statement ok
+DROP TABLE t_unique_null;