diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 1fbf8c6fb78cd..27115a2dc3a54 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -19,8 +19,8 @@ use std::sync::Arc; use datafusion_common::Result; use datafusion_physical_expr::{ - Distribution, EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, - PhysicalExpr, physical_exprs_equal, + AcrossPartitions, Distribution, EquivalenceProperties, LexOrdering, LexRequirement, + Partitioning, PhysicalExpr, physical_exprs_equal, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; @@ -32,6 +32,25 @@ use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +/// Returns `true` when `expr` is a *globally* constant column in `eq_properties` +/// and can therefore be safely dropped from a required ordering. +/// +/// Only [`AcrossPartitions::Uniform`] qualifies. An +/// [`AcrossPartitions::Heterogeneous`] value is constant *within* each partition +/// but may differ *across* partitions, so it must be kept in the sort key: the +/// `SortExec` these helpers build preserves partitioning, and its output is +/// merged across partitions downstream, where the column is still an ordering +/// discriminator. Dropping it there silently loses the global ordering. +fn is_uniform_constant( + eq_properties: &EquivalenceProperties, + expr: &Arc, +) -> bool { + matches!( + eq_properties.is_expr_constant(expr), + Some(AcrossPartitions::Uniform(_)) + ) +} + /// This utility function adds a `SortExec` above an operator according to the /// given ordering requirements while preserving the original partitioning. /// @@ -45,10 +64,7 @@ pub fn add_sort_above( ) -> PlanContext { let mut sort_reqs: Vec<_> = sort_requirements.into(); sort_reqs.retain(|sort_expr| { - node.plan - .equivalence_properties() - .is_expr_constant(&sort_expr.expr) - .is_none() + !is_uniform_constant(node.plan.equivalence_properties(), &sort_expr.expr) }); let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); let Some(ordering) = LexOrdering::new(sort_exprs) else { @@ -73,10 +89,7 @@ pub fn add_sort_above_with_distribution( ) -> PlanContext { let mut sort_reqs: Vec<_> = sort_requirements.into(); sort_reqs.retain(|sort_expr| { - node.plan - .equivalence_properties() - .is_expr_constant(&sort_expr.expr) - .is_none() + !is_uniform_constant(node.plan.equivalence_properties(), &sort_expr.expr) }); let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); let Some(ordering) = LexOrdering::new(sort_exprs) else { @@ -220,3 +233,77 @@ pub(crate) fn range_partitioning_satisfies_key_partitioning( pub fn is_limit(plan: &Arc) -> bool { plan.is::() || plan.is::() } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::PhysicalSortRequirement; + use datafusion_physical_expr::expressions::{col, lit}; + use datafusion_physical_expr::projection::ProjectionExpr; + use datafusion_physical_plan::displayable; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_physical_plan::projection::ProjectionExec; + + /// A `UnionExec` whose `sample` column is 1 on the left branch and 2 on the + /// right makes `sample` a `Heterogeneous` constant: constant within each + /// partition, but different across them. A required ordering of + /// `[sample, other]` must retain `sample`; the `SortExec` built here + /// preserves partitioning and its output is merged across partitions + /// downstream, where `sample` is still the leading ordering discriminator. + /// Regression test: `add_sort_above` used to drop every constant, uniform + /// or not, and silently discarded the global ordering. + #[test] + fn add_sort_above_keeps_heterogeneous_constant_key() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)])); + + let branch = |value: i32| -> Result> { + let source = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: lit(value), + alias: "sample".to_string(), + }, + ProjectionExpr { + expr: col("c", &schema)?, + alias: "other".to_string(), + }, + ], + source, + )?; + Ok(Arc::new(projection)) + }; + + let union: Arc = + UnionExec::try_new(vec![branch(1)?, branch(2)?])?; + let union_schema = union.schema(); + + // Precondition: the union reports `sample` as a heterogeneous constant. + let sample = col("sample", &union_schema)?; + assert!(matches!( + union.equivalence_properties().is_expr_constant(&sample), + Some(AcrossPartitions::Heterogeneous) + )); + + let asc = SortOptions::default(); + let requirement: LexRequirement = [ + PhysicalSortRequirement::new(sample, Some(asc)), + PhysicalSortRequirement::new(col("other", &union_schema)?, Some(asc)), + ] + .into(); + + let node = PlanContext::<()>::new_default(union); + let result = add_sort_above(node, requirement, None); + + let plan = displayable(result.plan.as_ref()).indent(true).to_string(); + assert!( + plan.contains("sample@0 ASC"), + "add_sort_above dropped the heterogeneous-constant sort key:\n{plan}" + ); + + Ok(()) + } +}