fix: UNIQUE constraint with NULLs incorrectly collapses GROUP BY groups

Author: xiedeyantu

datafusion/common/src/functional_dependencies.rs

@@ -196,15 +196,25 @@ impl FunctionalDependencies {
     }
 
     /// Creates a new `FunctionalDependencies` object from the given constraints.
+    ///
+    /// `nullable_flags` must contain one entry per field in the relation,
+    /// indicating whether that field is nullable.  A `UNIQUE` constraint whose
+    /// source columns include any nullable field is **not** a functional
+    /// dependency — because SQL treats `NULL` values as distinct, multiple rows
+    /// may carry `NULL` in a unique-key column without violating the constraint.
+    /// Such constraints are therefore omitted entirely.  When all source columns
+    /// are non-nullable a `UNIQUE` constraint is equivalent to a primary key and
+    /// is recorded with `nullable = false`.
     pub fn new_from_constraints(
         constraints: Option<&Constraints>,
-        n_field: usize,
+        nullable_flags: &[bool],
     ) -> Self {
+        let n_field = nullable_flags.len();
         if let Some(Constraints { inner: constraints }) = constraints {
             // Construct dependency objects based on each individual constraint:
             let dependencies = constraints
                 .iter()
-                .map(|constraint| {
+                .filter_map(|constraint| {
                     // All the field indices are associated with the whole table
                     // since we are dealing with table level constraints:
                     let dependency = match constraint {
@@ -213,15 +223,27 @@ impl FunctionalDependencies {
                             (0..n_field).collect::<Vec<_>>(),
                             false,
                         ),
-                        Constraint::Unique(indices) => FunctionalDependence::new(
-                            indices.to_vec(),
-                            (0..n_field).collect::<Vec<_>>(),
-                            true,
-                        ),
+                        Constraint::Unique(indices) => {
+                            // A UNIQUE constraint where any source column is
+                            // nullable is not a functional dependency: SQL does
+                            // not consider NULLs equal, so two rows may both
+                            // have NULL in the key and still satisfy the
+                            // constraint.  Only emit an FD when all source
+                            // columns are non-nullable, in which case it is
+                            // equivalent to a primary key.
+                            if indices.iter().any(|&i| nullable_flags[i]) {
+                                return None;
+                            }
+                            FunctionalDependence::new(
+                                indices.to_vec(),
+                                (0..n_field).collect::<Vec<_>>(),
+                                false,
+                            )
+                        }
                     };
                     // As primary keys are guaranteed to be unique, set the
                     // functional dependency mode to `Dependency::Single`:
-                    dependency.with_mode(Dependency::Single)
+                    Some(dependency.with_mode(Dependency::Single))
                 })
                 .collect::<Vec<_>>();
             Self::new(dependencies)
@@ -422,7 +444,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 +505,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

datafusion/expr/src/logical_plan/builder.rs

@@ -2887,6 +2887,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

@@ -354,7 +354,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()
             }
         }
@@ -2405,6 +2404,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);
@@ -2434,7 +2437,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();
@@ -2856,9 +2866,11 @@ impl TableScan {
             return plan_err!("table_name cannot be empty");
         }
         let schema = table_source.schema();
+        let nullable_flags: Vec<bool> =
+            schema.fields().iter().map(|f| f.is_nullable()).collect();
         let func_dependencies = FunctionalDependencies::new_from_constraints(
             table_source.constraints(),
-            schema.fields.len(),
+            &nullable_flags,
         );
         let projected_schema = projection
             .as_ref()
@@ -5146,7 +5158,7 @@ mod tests {
                         Some(&Constraints::new_unverified(vec![Constraint::Unique(
                             vec![0],
                         )])),
-                        1,
+                        &[false],
                     ),
                 )
                 .unwrap(),

datafusion/optimizer/src/eliminate_duplicated_expr.rs

@@ -145,7 +145,12 @@ mod tests {
     use crate::OptimizerContext;
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::test::*;
-    use datafusion_expr::{col, logical_plan::builder::LogicalPlanBuilder};
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion_common::{Constraint, Constraints};
+    use datafusion_expr::{
+        col, logical_plan::builder::LogicalPlanBuilder,
+        logical_plan::builder::table_source_with_constraints,
+    };
     use std::sync::Arc;
 
     macro_rules! assert_optimized_plan_equal {
@@ -200,4 +205,47 @@ mod tests {
             TableScan: test
         ")
     }
+
+    #[test]
+    fn eliminate_sort_exprs_pk_removes_dependent_key() -> Result<()> {
+        // When `id` is a PRIMARY KEY (non-nullable), it uniquely determines
+        // `val`, so `ORDER BY id, val` can safely be reduced to `ORDER BY id`.
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("val", DataType::Int32, false),
+        ]);
+        let constraints = Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
+        let source = table_source_with_constraints(&schema, constraints);
+        let plan = LogicalPlanBuilder::scan("t", source, None)?
+            .sort_by(vec![col("t.id"), col("t.val")])?
+            .build()?;
+
+        assert_optimized_plan_equal!(plan, @r"
+        Sort: t.id ASC NULLS LAST
+          TableScan: t
+        ")
+    }
+
+    #[test]
+    fn eliminate_sort_exprs_nullable_unique_keeps_dependent_key() -> Result<()> {
+        // When `id` is a nullable UNIQUE column, SQL allows multiple NULL
+        // values in `id`.  Because NULLs are not considered equal, multiple
+        // rows may share `id = NULL` with different `val` values, so `id`
+        // does NOT functionally determine `val`.  `ORDER BY id, val` must
+        // therefore keep both keys.
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, true), // nullable
+            Field::new("val", DataType::Int32, false),
+        ]);
+        let constraints = Constraints::new_unverified(vec![Constraint::Unique(vec![0])]);
+        let source = table_source_with_constraints(&schema, constraints);
+        let plan = LogicalPlanBuilder::scan("t", source, None)?
+            .sort_by(vec![col("t.id"), col("t.val")])?
+            .build()?;
+
+        assert_optimized_plan_equal!(plan, @r"
+        Sort: t.id ASC NULLS LAST, t.val ASC NULLS LAST
+          TableScan: t
+        ")
+    }
 }

datafusion/sqllogictest/test_files/cte.slt

@@ -1319,3 +1319,34 @@ RESET datafusion.execution.enable_recursive_ctes;
 
 statement ok
 RESET datafusion.sql_parser.enable_ident_normalization;
+
+# Regression test: functional dependencies from the static (anchor) term of a
+# recursive CTE must NOT be propagated to the outer SubqueryAlias.  The
+# recursive term can produce rows that violate any uniqueness constraint that
+# holds for the anchor alone.  Without this guard, Filter(pk = const) on the
+# CTE result would be mis-identified as scalar (at most 1 row) and return only
+# one row instead of all matching rows.
+statement ok
+CREATE TABLE pk_table(id INT NOT NULL, val INT NOT NULL, PRIMARY KEY(id));
+
+statement ok
+INSERT INTO pk_table VALUES (1, 100), (2, 200);
+
+# The recursive term produces a second row with id=1 (val=300).  Without the
+# FD fix, Filter(nodes.id = 1) would be deemed scalar and return only the
+# first matching row.
+query II rowsort
+WITH RECURSIVE nodes AS (
+    SELECT id, val FROM pk_table
+    UNION ALL
+    SELECT 1 AS id, 300 AS val
+    FROM nodes
+    WHERE nodes.id = 2
+)
+SELECT id, val FROM nodes WHERE id = 1
+----
+1 100
+1 300
+
+statement ok
+DROP TABLE pk_table;

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;