fix: do not drop GROUP BY keys using nullable UNIQUE constraints

datafusion/common/src/functional_dependencies.rs

@@ -472,9 +472,14 @@ pub fn aggregate_functional_dependencies(
     // The loop below only re-expresses input dependencies. Skip it when the
     // input has none. The GROUP BY-key dependency below always runs.
     if !func_dependencies.is_empty() {
-        // Compute once: this does not change in the loop.
-        let existing_target_indices =
-            get_target_functional_dependencies(aggr_input_schema, group_by_input_indices);
+        // Compare the columns covered by every dependency, whatever its NULL
+        // semantics: leaving out the ones that do not hold across NULLs would
+        // hide GROUP BY expressions this determinant does not cover and
+        // upgrade its mode to `Single` wrongly.
+        let existing_target_indices = get_target_indices_from_dependencies(
+            func_dependencies.iter(),
+            group_by_input_indices,
+        );
         for FunctionalDependence {
             source_indices,
             nullable,
@@ -500,8 +505,8 @@ pub fn aggregate_functional_dependencies(
                     new_source_input_indices.push(Some(*input_idx));
                 }
             }
-            let new_target_indices = get_target_functional_dependencies(
-                aggr_input_schema,
+            let new_target_indices = get_target_indices_from_dependencies(
+                func_dependencies.iter(),
                 &new_source_input_indices,
             );
             let mode = if existing_target_indices == new_target_indices
@@ -572,6 +577,12 @@ pub fn aggregate_functional_dependencies(
 /// Returns target indices, for the determinant keys that are inside
 /// group by expressions.
 ///
+/// Only dependencies that hold when NULL determinant values are grouped
+/// together are considered (see
+/// [`FunctionalDependence::is_valid_across_nulls`]): a nullable `UNIQUE`
+/// column permits multiple NULL rows that may differ in the other columns,
+/// so it does not determine them across the single NULL group.
+///
 /// `group_by_input_indices` holds, for each GROUP BY expression, the index of
 /// the `schema` field it references, or `None` for a computed expression.
 pub fn get_target_functional_dependencies(
@@ -582,21 +593,27 @@ pub fn get_target_functional_dependencies(
     if dependencies.is_empty() {
         return None;
     }
+    let dependencies = dependencies
+        .iter()
+        .filter(|dependency| dependency.is_valid_across_nulls(schema));
+    get_target_indices_from_dependencies(dependencies, group_by_input_indices)
+}
+
+fn get_target_indices_from_dependencies<'a>(
+    dependencies: impl IntoIterator<Item = &'a FunctionalDependence>,
+    group_by_input_indices: &[Option<usize>],
+) -> Option<Vec<usize>> {
     let mut combined_target_indices = HashSet::new();
-    for FunctionalDependence {
-        source_indices,
-        target_indices,
-        ..
-    } in &dependencies.deps
-    {
+    for dependency in dependencies {
         // If the GROUP BY expression contains a determinant key, we can use
         // the associated fields after aggregation even if they are not part
         // of the GROUP BY expression.
-        if source_indices
+        if dependency
+            .source_indices
             .iter()
             .all(|source_idx| group_by_input_indices.contains(&Some(*source_idx)))
         {
-            combined_target_indices.extend(target_indices.iter());
+            combined_target_indices.extend(dependency.target_indices.iter());
         }
     }
     (!combined_target_indices.is_empty()).then_some({
@@ -623,22 +640,21 @@ pub fn get_required_group_by_exprs_indices(
         .collect::<Option<Vec<_>>>()?;
 
     groupby_expr_indices.sort_unstable();
-    for FunctionalDependence {
-        source_indices,
-        target_indices,
-        ..
-    } in &dependencies.deps
-    {
-        if source_indices
+    for dependency in &dependencies.deps {
+        // If all source indices are among GROUP BY expression indices, we
+        // can remove target indices from GROUP BY expression indices and
+        // use source indices instead. The dependency must also hold across
+        // NULL keys, which GROUP BY collapses into a single group.
+        if dependency
+            .source_indices
             .iter()
             .all(|source_idx| groupby_expr_indices.contains(source_idx))
+            && dependency.is_valid_across_nulls(schema)
         {
-            // If all source indices are among GROUP BY expression indices, we
-            // can remove target indices from GROUP BY expression indices and
-            // use source indices instead.
-            groupby_expr_indices = set_difference(&groupby_expr_indices, target_indices);
             groupby_expr_indices =
-                merge_and_order_indices(groupby_expr_indices, source_indices);
+                set_difference(&groupby_expr_indices, &dependency.target_indices);
+            groupby_expr_indices =
+                merge_and_order_indices(groupby_expr_indices, &dependency.source_indices);
         }
     }
     groupby_expr_indices

datafusion/expr/src/logical_plan/builder.rs

@@ -3138,6 +3138,63 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn plan_builder_aggregate_nullable_unique_adds_no_implicit_group_by_exprs()
+    -> Result<()> {
+        // A nullable UNIQUE column permits multiple NULL rows that may differ
+        // in the other columns, so it does not determine them across the
+        // single NULL group and they must not be added to the GROUP BY.
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, true),
+            Field::new("state", DataType::Utf8, true),
+            Field::new("salary", DataType::Int32, true),
+        ]);
+        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 plan_builder_aggregate_not_null_unique_implicit_group_by_exprs() -> Result<()> {
+        // A UNIQUE column declared NOT NULL has no NULL rows, so it determines
+        // the other columns like a PRIMARY KEY and they are added to the
+        // GROUP BY.
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("state", DataType::Utf8, true),
+            Field::new("salary", DataType::Int32, true),
+        ]);
+        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, employee_csv.state, employee_csv.salary]], aggr=[[sum(employee_csv.salary)]]
+          TableScan: employee_csv
+        ");
+
+        Ok(())
+    }
+
     #[test]
     fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> {
         // https://github.com/apache/datafusion/issues/23812

datafusion/optimizer/src/replace_distinct_aggregate.rs

@@ -102,8 +102,11 @@ impl OptimizerRule for ReplaceDistinctWithAggregate {
                     // 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.
+                    // times, so the DISTINCT still has work to do. Uniqueness
+                    // must also hold across NULLs: a nullable UNIQUE key can
+                    // contain multiple NULL rows that DISTINCT must collapse.
                     if dep.mode == Dependency::Single
+                        && dep.is_valid_across_nulls(input.schema())
                         && dep.source_indices.len() >= field_count
                         && dep.source_indices[..field_count]
                             .iter()