fix: have orderings include constants which are heterogenius across p…

…artitions
apache · Feb 26, 2025 · 5a4ef63 · 5a4ef63
1 parent dffeaac
commit 5a4ef63
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs
@@ -211,9 +211,10 @@ pub fn const_exprs_contains(
     const_exprs: &[ConstExpr],
     expr: &Arc<dyn PhysicalExpr>,
 ) -> bool {
-    const_exprs
-        .iter()
-        .any(|const_expr| const_expr.expr.eq(expr))
+    const_exprs.iter().any(|const_expr| {
+        const_expr.expr.eq(expr)
+            && const_expr.across_partitions() != AcrossPartitions::Heterogeneous
+    })
 }
 
 /// An `EquivalenceClass` is a set of [`Arc<dyn PhysicalExpr>`]s that are known

diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs
@@ -438,7 +438,7 @@ impl EquivalenceProperties {
         let filtered_exprs = LexOrdering::new(
             sort_exprs
                 .into_iter()
-                .filter(|expr| !self.is_expr_constant(&expr.expr))
+                .filter(|expr| !self.is_expr_constant_across_partitions(&expr.expr))
                 .collect(),
         );
 

diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs
@@ -45,7 +45,7 @@ pub fn add_sort_above<T: Clone + Default>(
         !node
             .plan
             .equivalence_properties()
-            .is_expr_constant(&sort_expr.expr)
+            .is_expr_constant_across_partitions(&sort_expr.expr)
     });
     let mut new_sort = SortExec::new(sort_expr, Arc::clone(&node.plan)).with_fetch(fetch);
     if node.plan.output_partitioning().partition_count() > 1 {