fix: do not prune ORDER BY keys using nullable UNIQUE constraints

datafusion/common/src/functional_dependencies.rs

@@ -23,7 +23,7 @@ use std::ops::Deref;
 use std::vec::IntoIter;
 
 use crate::utils::{merge_and_order_indices, set_difference};
-use crate::{DFSchema, HashSet, JoinType};
+use crate::{DFSchema, HashSet, JoinType, NullEquality};
 
 /// This object defines a constraint on a table.
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
@@ -144,6 +144,13 @@ pub struct FunctionalDependence {
     /// such as after LEFT JOIN or RIGHT JOIN operations, this property may
     /// change.
     pub nullable: bool,
+    /// The NULL-comparison semantics under which this dependency holds. The
+    /// conservative default, [`NullEquality::NullEqualsNothing`], means it
+    /// holds only across rows whose determinant contains no NULLs; e.g. a
+    /// nullable `UNIQUE` constraint permits multiple NULL rows that may differ.
+    /// [`NullEquality::NullEqualsNull`] means it also holds when NULL
+    /// determinant values are treated as equal; e.g. a `GROUP BY` key.
+    pub null_equality: NullEquality,
     // The functional dependency mode:
     pub mode: Dependency,
 }
@@ -168,6 +175,8 @@ impl FunctionalDependence {
             source_indices,
             target_indices,
             nullable,
+            // Assume the dependency does not hold across NULL rows by default:
+            null_equality: NullEquality::NullEqualsNothing,
             // Start with the least restrictive mode by default:
             mode: Dependency::Multi,
         }
@@ -177,6 +186,25 @@ impl FunctionalDependence {
         self.mode = mode;
         self
     }
+
+    pub fn with_null_equality(mut self, null_equality: NullEquality) -> Self {
+        self.null_equality = null_equality;
+        self
+    }
+
+    /// Returns `true` if this dependency remains usable for operations that
+    /// treat NULL determinant values as equal (`GROUP BY`, `DISTINCT` and
+    /// sorting, which place all NULL keys together): it must hold under
+    /// NULLs-are-equal semantics, have a non-nullable determinant (e.g. a
+    /// `PRIMARY KEY`), or have no nullable source field in the given `schema`.
+    pub fn is_valid_across_nulls(&self, schema: &DFSchema) -> bool {
+        self.null_equality == NullEquality::NullEqualsNull
+            || !self.nullable
+            || self
+                .source_indices
+                .iter()
+                .all(|&source_idx| !schema.field(source_idx).is_nullable())
+    }
 }
 
 /// This object encapsulates all functional dependencies in a given relation.
@@ -301,6 +329,7 @@ impl FunctionalDependencies {
             source_indices,
             target_indices,
             nullable,
+            null_equality,
             mode,
         } in &self.deps
         {
@@ -321,7 +350,8 @@ impl FunctionalDependencies {
                     new_target_indices,
                     *nullable,
                 )
-                .with_mode(*mode);
+                .with_mode(*mode)
+                .with_null_equality(*null_equality);
                 projected_func_dependencies.push(new_func_dependence);
             }
         }
@@ -386,7 +416,12 @@ impl FunctionalDependencies {
     fn downgrade_dependencies(&mut self) {
         // Delete nullable dependencies, since they are no longer valid:
         self.deps.retain(|item| !item.nullable);
-        self.deps.iter_mut().for_each(|item| item.nullable = true);
+        // Survivors had a non-null determinant, so every new NULL comes from a
+        // padded row whose targets are all NULL too:
+        self.deps.iter_mut().for_each(|item| {
+            item.nullable = true;
+            item.null_equality = NullEquality::NullEqualsNull;
+        });
     }
 
     /// This function ensures that functional dependencies involving uniquely
@@ -432,6 +467,7 @@ pub fn aggregate_functional_dependencies(
     for FunctionalDependence {
         source_indices,
         nullable,
+        null_equality,
         mode,
         ..
     } in &func_dependencies.deps
@@ -470,13 +506,22 @@ pub fn aggregate_functional_dependencies(
         };
         // All of the composite indices occur in the GROUP BY expression:
         if new_source_indices.len() == source_indices.len() {
+            // GROUP BY treats NULLs as equal: a determinant covering the
+            // complete grouping key gets at most one output row per NULL too.
+            let output_null_equality =
+                if new_source_indices.len() == group_by_expr_names.len() {
+                    NullEquality::NullEqualsNull
+                } else {
+                    *null_equality
+                };
             aggregate_func_dependencies.push(
                 FunctionalDependence::new(
                     new_source_indices,
                     target_indices.clone(),
                     *nullable,
                 )
-                .with_mode(mode),
+                .with_mode(mode)
+                .with_null_equality(output_null_equality),
             );
         }
     }
@@ -497,13 +542,17 @@ pub fn aggregate_functional_dependencies(
             // The following simple comparison is working well because
             // GROUP BY expressions come here as a prefix.
             item.source_indices.iter().all(|idx| idx < &count)
+                // A dependency that is not valid across NULLs cannot replace
+                // the whole GROUP BY key dependency:
+                && item.is_valid_across_nulls(aggr_schema)
         }) {
             // Add a new functional dependency associated with the whole table:
             // Use nullable property of the GROUP BY expression:
             aggregate_func_dependencies.push(
-                // Use nullable property of the GROUP BY expression:
                 FunctionalDependence::new(source_indices, target_indices, nullable)
-                    .with_mode(Dependency::Single),
+                    .with_mode(Dependency::Single)
+                    // Grouping collapses NULL keys into a single group:
+                    .with_null_equality(NullEquality::NullEqualsNull),
             );
         }
     }
@@ -617,15 +666,15 @@ pub fn get_required_sort_exprs_indices(
         };
 
         // A sort expression is removable if its value is functionally determined
-        // by fields that already appear earlier in the sort order: if the earlier
-        // fields are fixed, this one's value is fixed too, so it adds no ordering
-        // information.
+        // by fields that already appear earlier in the sort order (and the
+        // dependency remains valid across NULL rows).
         let removable = dependencies.deps.iter().any(|dependency| {
             dependency.target_indices.contains(&field_idx)
                 && dependency
                     .source_indices
                     .iter()
                     .all(|source_idx| known_field_indices.contains(source_idx))
+                && dependency.is_valid_across_nulls(schema)
         });
 
         if removable {

datafusion/expr/src/logical_plan/builder.rs

@@ -3235,7 +3235,8 @@ mod tests {
                 FunctionalDependence::new(vec![0], vec![0, 1, 2, 3], false)
                     .with_mode(Dependency::Single),
                 FunctionalDependence::new(vec![2], vec![2, 3], true)
-                    .with_mode(Dependency::Multi),
+                    .with_mode(Dependency::Multi)
+                    .with_null_equality(NullEquality::NullEqualsNull),
             ])
         );
         Ok(())

datafusion/sqllogictest/test_files/functional_dependencies.slt

@@ -136,24 +136,35 @@ logical_plan
 # 2.2 Nullable UNIQUE: `x` does NOT determine `y` across the two NULL rows,
 # so the `y` sort key must be kept.
 #
-# BUG:
-# Expected: `1 3`, `NULL 1`, `NULL 2`.
-# Issue: https://github.com/apache/datafusion/issues/23818
 query II
 SELECT x, y FROM t_uniq ORDER BY x NULLS LAST, y;
 ----
 1 3
-NULL 2
 NULL 1
+NULL 2
 
 query TT
 EXPLAIN SELECT x, y FROM t_uniq ORDER BY x NULLS LAST, y;
 ----
 logical_plan
-01)Sort: t_uniq.x ASC NULLS LAST
+01)Sort: t_uniq.x ASC NULLS LAST, t_uniq.y ASC NULLS LAST
 02)--TableScan: t_uniq projection=[x, y]
 
-# 2.3 After `GROUP BY x` the `x` does determine `cnt`, so can drop `cnt` from sort
+# 2.3 A UNIQUE determinant remains usable when its source column is NOT NULL.
+statement ok
+CREATE TABLE t_uniq_not_null (x INT NOT NULL UNIQUE, y INT) AS VALUES (1, 10), (2, 20);
+
+query TT
+EXPLAIN SELECT x, y FROM t_uniq_not_null ORDER BY x, y;
+----
+logical_plan
+01)Sort: t_uniq_not_null.x ASC NULLS LAST
+02)--TableScan: t_uniq_not_null projection=[x, y]
+
+statement ok
+drop table t_uniq_not_null;
+
+# 2.4 After `GROUP BY x` the `x` does determine `cnt`, so can drop `cnt` from sort
 query TT
 EXPLAIN SELECT x, cnt FROM (SELECT x, count(*) AS cnt FROM t_uniq GROUP BY x) ORDER BY x, cnt;
 ----
@@ -163,6 +174,57 @@ logical_plan
 03)----Aggregate: groupBy=[[t_uniq.x]], aggr=[[count(Int64(1))]]
 04)------TableScan: t_uniq projection=[x]
 
+# 2.5 A PRIMARY KEY on the NULL-padded side of a LEFT JOIN: every new NULL in
+# `p.a` comes from a padded row, where `p.b` is NULL too. So `p.a` still
+# determines `p.b` and the `p.b` sort key can be dropped.
+statement ok
+CREATE TABLE t_left (k INT) AS VALUES (1), (3), (4);
+
+statement ok
+CREATE TABLE t_right_pk (a INT, b INT, PRIMARY KEY (a)) AS VALUES (1, 10), (2, 20);
+
+query II
+SELECT p.a, p.b FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a ORDER BY p.a, p.b DESC;
+----
+1 10
+NULL NULL
+NULL NULL
+
+query TT
+EXPLAIN SELECT l.k, p.a, p.b FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a ORDER BY p.a, p.b;
+----
+logical_plan
+01)Sort: p.a ASC NULLS LAST
+02)--Left Join: l.k = p.a
+03)----SubqueryAlias: l
+04)------TableScan: t_left projection=[k]
+05)----SubqueryAlias: p
+06)------TableScan: t_right_pk projection=[a, b]
+
+statement ok
+drop table t_left;
+
+statement ok
+drop table t_right_pk;
+
+# 2.6 After `GROUP BY x, y` the pair `(x, y)` is unique (NULLs included), so
+# it determines `cnt` even though the nullable UNIQUE `x` alone does not.
+query III
+SELECT x, y, cnt FROM (SELECT x, y, count(*) AS cnt FROM t_uniq GROUP BY x, y) ORDER BY x, y DESC, cnt;
+----
+1 3 1
+NULL 2 1
+NULL 1 1
+
+query TT
+EXPLAIN SELECT x, y, cnt FROM (SELECT x, y, count(*) AS cnt FROM t_uniq GROUP BY x, y) ORDER BY x, y, cnt;
+----
+logical_plan
+01)Sort: t_uniq.x ASC NULLS LAST, t_uniq.y ASC NULLS LAST
+02)--Projection: t_uniq.x, t_uniq.y, count(Int64(1)) AS cnt
+03)----Aggregate: groupBy=[[t_uniq.x, t_uniq.y]], aggr=[[count(Int64(1))]]
+04)------TableScan: t_uniq projection=[x, y]
+
 
 # 3.1 PRIMARY KEY: `x` determines `y`, and `y` is not selected, so grouping
 # by `x, y` is the same as grouping by `x`.

docs/source/library-user-guide/upgrading/56.0.0.md

@@ -668,6 +668,29 @@ Wire compatibility is directional:
   ignores the unknown field but cannot preserve the qualifier required by the
   expressions, which can cause resolution failure or incorrect rebinding.
 
+### `FunctionalDependence` has a new `null_equality` field
+
+`FunctionalDependence` now records the NULL semantics its dependency holds under.
+Build it with `FunctionalDependence::new(..)` and, when needed,
+`.with_null_equality(NullEquality::NullEqualsNull)` instead of a struct literal.
+
+**Migration guide:**
+
+```rust,ignore
+// Before
+let dep = FunctionalDependence {
+    source_indices,
+    target_indices,
+    nullable,
+    mode: Dependency::Single,
+};
+
+// After
+let dep = FunctionalDependence::new(source_indices, target_indices, nullable)
+    .with_mode(Dependency::Single)
+    .with_null_equality(NullEquality::NullEqualsNull);
+```
+
 ### `FileSource::exact_filter` controls scan equivalence properties
 
 `FileScanConfig` no longer derives equivalence properties (constant columns