diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index e9ad978b2e0cb..53aedb6cf1488 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2245,6 +2245,29 @@ async fn test_do_not_pushdown_through_limit() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_limit_preserves_input_order_through_projection() -> Result<()> { + let schema = create_test_schema()?; + let source = memory_exec(&schema); + let input = sort_exec([sort_expr("non_nullable_col", &schema)].into(), source); + let projection = projection_exec( + vec![(col("nullable_col", &schema)?, "nullable_col".to_string())], + input, + )?; + let physical_plan = Arc::new(GlobalLimitExec::new(projection, 1, None)) as _; + + let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); + assert_snapshot!(test.run(), @r" + Input / Optimized Plan: + GlobalLimitExec: skip=1, fetch=None + ProjectionExec: expr=[nullable_col@0 as nullable_col] + SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] + "); + + Ok(()) +} + #[tokio::test] async fn test_remove_unnecessary_spm1() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 2c6c46c82985a..a24743a9f7649 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -299,6 +299,36 @@ fn test_projection_over_multi_partition_sort_limit() { "); } +/// A limit consumes the global ordering established below a projection even +/// when the projection removes the sort key from its output schema. The +/// ordering must survive a second `EnsureRequirements` pass, which sees the +/// `SortPreservingMergeExec` introduced by the first pass. +#[test] +fn test_limit_preserves_hidden_sort_across_distribution_passes() { + let source = Arc::new(MockMultiPartitionExec::new(16)); + let sort_expr = sort_expr_on("a", 0, true, true); + let sort = Arc::new(SortExec::new(sort_expr, source)); + + // Hide the sort key so the projection itself has no output ordering. + let projection = Arc::new( + ProjectionExec::try_new( + vec![(Arc::new(Column::new("b", 1)) as _, "b".to_string())], + sort, + ) + .unwrap(), + ); + let limit: Arc = + Arc::new(GlobalLimitExec::new(projection, 1, None)); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=1, fetch=None + ProjectionExec: expr=[b@1 as b] + SortPreservingMergeExec: [a@0 DESC] + SortExec: expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + // ======================================================================== // Single partition tests (no unnecessary operators) // ======================================================================== diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 087a351b697cc..1e4fd24723fbd 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -50,6 +50,10 @@ pub struct FFI_ExecutionPlan { /// Return a vector of children plans pub children: unsafe extern "C" fn(plan: &Self) -> SVec, + /// Return whether each child input's ordering must be preserved. + pub requires_input_order_preservation: + unsafe extern "C" fn(plan: &Self) -> SVec, + pub with_new_children: unsafe extern "C" fn(plan: &Self, children: SVec) -> FFI_Result, @@ -138,6 +142,15 @@ unsafe extern "C" fn children_fn_wrapper( .collect() } +unsafe extern "C" fn requires_input_order_preservation_fn_wrapper( + plan: &FFI_ExecutionPlan, +) -> SVec { + plan.inner() + .requires_input_order_preservation() + .into_iter() + .collect() +} + unsafe extern "C" fn with_new_children_fn_wrapper( plan: &FFI_ExecutionPlan, children: SVec, @@ -306,6 +319,8 @@ impl FFI_ExecutionPlan { Self { properties: properties_fn_wrapper, children: children_fn_wrapper, + requires_input_order_preservation: + requires_input_order_preservation_fn_wrapper, with_new_children: with_new_children_fn_wrapper, name: name_fn_wrapper, execute: execute_fn_wrapper, @@ -416,6 +431,12 @@ impl ExecutionPlan for ForeignExecutionPlan { self.children.iter().collect() } + fn requires_input_order_preservation(&self) -> Vec { + unsafe { (self.plan.requires_input_order_preservation)(&self.plan) } + .into_iter() + .collect() + } + fn with_new_children( self: Arc, children: Vec>, @@ -485,6 +506,7 @@ pub mod tests { children: Vec>, metrics: Option, statistics: Option, + requires_input_order_preservation: bool, } impl EmptyExec { @@ -499,9 +521,18 @@ pub mod tests { children: Vec::default(), metrics: None, statistics: None, + requires_input_order_preservation: false, } } + pub fn with_input_order_preservation( + mut self, + requires_input_order_preservation: bool, + ) -> Self { + self.requires_input_order_preservation = requires_input_order_preservation; + self + } + pub fn with_metrics(mut self, metrics: MetricsSet) -> Self { self.metrics = Some(metrics); self @@ -545,9 +576,14 @@ pub mod tests { children, metrics: self.metrics.clone(), statistics: self.statistics.clone(), + requires_input_order_preservation: self.requires_input_order_preservation, })) } + fn requires_input_order_preservation(&self) -> Vec { + vec![self.requires_input_order_preservation; self.children.len()] + } + fn execute( &self, _partition: usize, @@ -600,6 +636,25 @@ pub mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_input_order_preservation_round_trip() -> Result<()> { + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Float32, false), + ])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))) as _; + let original_plan = + Arc::new(EmptyExec::new(schema).with_input_order_preservation(true)) + .with_new_children(vec![child])?; + + let mut local_plan = FFI_ExecutionPlan::new(original_plan, None); + local_plan.library_marker_id = crate::mock_foreign_marker_id; + let foreign_plan: Arc = (&local_plan).try_into()?; + + assert_eq!(foreign_plan.requires_input_order_preservation(), vec![true]); + + Ok(()) + } + #[test] fn test_ffi_execution_plan_children() -> Result<()> { let schema = Arc::new(arrow::datatypes::Schema::new(vec![ diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..e5041b1e1b0e1 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -107,6 +107,8 @@ pub struct ForeignLibraryModule { pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan, + pub create_order_sensitive_exec: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan, pub create_table_with_statistics: @@ -161,6 +163,15 @@ pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { FFI_ExecutionPlan::new(plan, None) } +pub(crate) extern "C" fn create_order_sensitive_exec() -> FFI_ExecutionPlan { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))) as _; + let plan = Arc::new(EmptyExec::new(schema).with_input_order_preservation(true)) + .with_new_children(vec![child]) + .expect("order-sensitive test plan should accept one child"); + FFI_ExecutionPlan::new(plan, None) +} + /// Returns canonical statistics used by both the producer and consumer sides of /// the integration tests so round-trips can be asserted without hard-coding /// the values in two places. @@ -259,6 +270,7 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_rank_udwf: create_ffi_rank_func, create_extension_options: config::create_extension_options, create_empty_exec, + create_order_sensitive_exec, create_exec_with_statistics, create_table_with_statistics, create_physical_optimizer_rule: diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 7d04e828bd4a5..8eb6e962b2420 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -28,6 +28,19 @@ mod tests { use datafusion_physical_plan::ExecutionPlan; use std::sync::Arc; + #[test] + fn test_ffi_execution_plan_input_order_preservation_cross_library() + -> Result<(), DataFusionError> { + let module = get_module()?; + let plan = (module.create_order_sensitive_exec)(); + let plan: Arc = (&plan).try_into()?; + + assert!(plan.is::()); + assert_eq!(plan.requires_input_order_preservation(), vec![true]); + + Ok(()) + } + #[test] #[expect(deprecated)] fn test_ffi_execution_plan_partition_statistics_cross_library() diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 952aae9846d0f..99a82a1176ba1 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -839,12 +839,21 @@ struct RemovedDistOps { fn remove_dist_changing_operators( mut distribution_context: DistributionContext, + preserve_output_order: bool, ) -> Result { let mut removed_fetch = None; while is_repartition(&distribution_context.plan) || is_coalesce_partitions(&distribution_context.plan) || is_sort_preserving_merge(&distribution_context.plan) { + // Keep an order-producing distribution operator when an ancestor + // semantically depends on its output sequence. Once such an operator + // is removed, an intervening projection may make its ordering + // impossible to reconstruct. + if preserve_output_order && distribution_context.plan.output_ordering().is_some() + { + break; + } // Preserve fetch from SPM or CoalescePartitions before removing (#14150). if let Some(fetch) = distribution_context.plan.fetch() { removed_fetch = Some( @@ -936,6 +945,7 @@ struct DistributionChildState { context: DistributionContext, required_input_ordering: Option, maintains_input_order: bool, + preserve_output_order: bool, requirement: Distribution, } @@ -1126,13 +1136,46 @@ fn enforce_distribution_relationships( /// This function is intended to be used in a bottom up traversal, as it /// can first repartition (or newly partition) at the datasources -- these /// source partitions may be later repartitioned with additional data exchange operators. +pub fn ensure_distribution( + dist_context: DistributionContext, + config: &ConfigOptions, +) -> Result> { + ensure_distribution_inner(dist_context, config, false) +} + +/// Enforces distribution bottom-up while carrying semantic order-preservation +/// requirements top-down. +pub(super) fn ensure_distribution_with_order_preservation( + mut dist_context: DistributionContext, + config: &ConfigOptions, + preserve_output_order: bool, +) -> Result { + let own_requirements = dist_context.plan.requires_input_order_preservation(); + let maintains_input_order = dist_context.plan.maintains_input_order(); + dist_context.children = izip!( + dist_context.children, + own_requirements, + maintains_input_order + ) + .map(|(child, own_requirement, maintains_input_order)| { + let preserve_child_order = + own_requirement || (preserve_output_order && maintains_input_order); + ensure_distribution_with_order_preservation(child, config, preserve_child_order) + }) + .collect::>>()?; + + ensure_distribution_inner(dist_context, config, preserve_output_order) + .map(|transformed| transformed.data) +} + #[expect( deprecated, reason = "HashPartitioned is accepted during the KeyPartitioned migration" )] -pub fn ensure_distribution( +fn ensure_distribution_inner( dist_context: DistributionContext, config: &ConfigOptions, + preserve_output_order: bool, ) -> Result> { let dist_context = update_children(dist_context)?; @@ -1171,7 +1214,7 @@ pub fn ensure_distribution( children, }, removed_fetch, - } = remove_dist_changing_operators(dist_context)?; + } = remove_dist_changing_operators(dist_context, preserve_output_order)?; if let Some(exec) = plan.downcast_ref::() { if let Some(updated_window) = get_best_fitting_window( @@ -1238,6 +1281,7 @@ pub fn ensure_distribution( children.into_iter(), plan.required_input_ordering(), plan.maintains_input_order(), + plan.requires_input_order_preservation(), repartition_status_flags.into_iter() ) .map( @@ -1246,6 +1290,7 @@ pub fn ensure_distribution( mut child, required_input_ordering, maintains, + own_order_preservation, RepartitionRequirementStatus { requirement, roundrobin_beneficial, @@ -1253,6 +1298,8 @@ pub fn ensure_distribution( hash_necessary, }, )| { + let preserve_child_order = + own_order_preservation || (preserve_output_order && maintains); // Allow subset satisfaction when: // 1. Current partition count >= threshold // 2. Not a partitioned join since must use exact hash matching for joins @@ -1391,6 +1438,7 @@ pub fn ensure_distribution( context: child, required_input_ordering, maintains_input_order: maintains, + preserve_output_order: preserve_child_order, requirement, }) }, @@ -1414,6 +1462,7 @@ pub fn ensure_distribution( mut context, required_input_ordering, maintains_input_order, + preserve_output_order, requirement, }| { let streaming_benefit = if context.data { @@ -1436,6 +1485,7 @@ pub fn ensure_distribution( if (!ordering_satisfied || !order_preserving_variants_desirable) && !streaming_benefit && context.data + && !preserve_output_order { context = replace_order_preserving_variants(context)?; // If ordering requirements were satisfied before repartitioning, @@ -1464,14 +1514,18 @@ pub fn ensure_distribution( // ordering is pointless. However, if it does maintain // input order, we keep order-preserving variants so // ordering can flow through to ancestors that need it. - if !maintains_input_order && !streaming_benefit { + if !maintains_input_order + && !streaming_benefit + && !preserve_output_order + { context = replace_order_preserving_variants(context)?; } } Distribution::UnspecifiedDistribution => { // Since ordering is lost, trying to preserve ordering is pointless - if !maintains_input_order - || plan.is::() + if (!maintains_input_order + || plan.is::()) + && !preserve_output_order { context = replace_order_preserving_variants(context)?; } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 4dce4691f0963..e595ceae02596 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -50,8 +50,8 @@ use std::sync::Arc; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ - add_sort_above, add_sort_above_with_check, is_coalesce_partitions, is_limit, - is_repartition, is_sort, is_sort_preserving_merge, is_window, + add_sort_above, add_sort_above_with_check, is_coalesce_partitions, is_repartition, + is_sort, is_sort_preserving_merge, is_window, }; use datafusion_common::Result; @@ -102,24 +102,27 @@ fn update_sort_ctx_children_data( child_node.data = if is_sort(child_plan) { // child is sort true - } else if is_limit(child_plan) { - // There is no sort linkage for this path, it starts at a limit. - false } else { // If a descendent is a sort, and the child maintains the sort. let is_spm = is_sort_preserving_merge(child_plan); let required_orderings = child_plan.required_input_ordering(); let flags = child_plan.maintains_input_order(); + let preserve_input_order = child_plan.requires_input_order_preservation(); // Add parent node to the tree if there is at least one child with // a sort connection: - izip!(flags, required_orderings).any(|(maintains, required_ordering)| { - let propagates_ordering = - (maintains && required_ordering.is_none()) || is_spm; - // `connected_to_sort` only returns the correct answer with bottom-up traversal - let connected_to_sort = - child_node.children.iter().any(|child| child.data); - propagates_ordering && connected_to_sort - }) + izip!( + flags, + required_orderings, + preserve_input_order, + &child_node.children + ) + .any( + |(maintains, required_ordering, preserve_input_order, child)| { + let propagates_ordering = !preserve_input_order + && ((maintains && required_ordering.is_none()) || is_spm); + propagates_ordering && child.data + }, + ) } } @@ -419,11 +422,12 @@ pub fn ensure_sorting( let plan = &requirements.plan; let mut updated_children = vec![]; - for (idx, (required_ordering, mut child)) in plan - .required_input_ordering() - .into_iter() - .zip(requirements.children) - .enumerate() + for (idx, (required_ordering, preserve_input_order, mut child)) in izip!( + plan.required_input_ordering(), + plan.requires_input_order_preservation(), + requirements.children + ) + .enumerate() { let physical_ordering = child.plan.output_ordering(); @@ -444,7 +448,9 @@ pub fn ensure_sorting( ); child = update_sort_ctx_children_data(child, true)?; } - } else if physical_ordering.is_none() || !plan.maintains_input_order()[idx] { + } else if !preserve_input_order + && (physical_ordering.is_none() || !plan.maintains_input_order()[idx]) + { // We have a `SortExec` whose effect may be neutralized by another // order-imposing operator, remove this sort: child = update_child_to_remove_unnecessary_sort(idx, child, plan)?; diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 5c17ffbd1e7db..5a93352eef0ea 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -54,11 +54,13 @@ use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; /// computational cost is reduced by pushing down `SortExec`s through certain /// executors. The object carries the parent required ordering, the (optional) /// `fetch` value of the parent node, and the parent's distribution requirement -/// (used by the distribution-aware pushdown path) as its data. +/// (used by the distribution-aware pushdown path) as its data. It also tracks +/// whether the input ordering is semantically significant to an ancestor. #[derive(Clone, Debug)] pub struct ParentRequirements { ordering_requirement: Option, fetch: Option, + preserve_input_order: bool, /// The distribution required by the consumer above any SortExec we insert. /// When this is `SinglePartition` and the input has multiple partitions, /// `add_sort_above_with_distribution` wraps the sort in `SortPreservingMergeExec`. @@ -70,6 +72,7 @@ impl Default for ParentRequirements { Self { ordering_requirement: None, fetch: None, + preserve_input_order: false, distribution_requirement: Distribution::UnspecifiedDistribution, } } @@ -80,6 +83,7 @@ pub type SortPushDown = PlanContext; /// Assigns the ordering requirement of the root node to the its children. pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { let reqs = sort_push_down.plan.required_input_ordering(); + let preserve_input_order = sort_push_down.plan.requires_input_order_preservation(); let dists = sort_push_down .plan .input_distribution_requirements() @@ -90,6 +94,7 @@ pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { child.data = ParentRequirements { ordering_requirement: requirement, fetch: child.plan.fetch(), + preserve_input_order: preserve_input_order[idx], distribution_requirement: dists .get(idx) .cloned() @@ -98,6 +103,22 @@ pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { } } +/// Assigns sequence-preservation requirements to each child. +/// +/// A plan can introduce this requirement for one of its own inputs, or pass +/// down a requirement from its parent when it maintains that input's order. +fn assign_input_order_preservation( + sort_push_down: &mut SortPushDown, + preserve_output_order: bool, +) { + let own_requirements = sort_push_down.plan.requires_input_order_preservation(); + let maintains_input_order = sort_push_down.plan.maintains_input_order(); + for (idx, child) in sort_push_down.children.iter_mut().enumerate() { + child.data.preserve_input_order = own_requirements[idx] + || (preserve_output_order && maintains_input_order[idx]); + } +} + /// Tries to push down the sort requirements as far as possible, if decides a `SortExec` is unnecessary removes it. pub fn pushdown_sorts(sort_push_down: SortPushDown) -> Result { sort_push_down @@ -142,6 +163,7 @@ fn pushdown_sorts_helper( ) -> Result> { let plan = sort_push_down.plan; let parent_fetch = sort_push_down.data.fetch; + let preserve_input_order = sort_push_down.data.preserve_input_order; let parent_distribution = sort_push_down.data.distribution_requirement.clone(); let Some(parent_requirement) = sort_push_down.data.ordering_requirement.clone() @@ -152,8 +174,9 @@ fn pushdown_sorts_helper( let Some(sort_ordering) = plan.output_ordering().cloned() else { return internal_err!("SortExec should have output ordering"); }; - // The sort is unnecessary, just propagate the stricter fetch and - // ordering requirements. + // Propagate the Sort as a concrete ordering requirement. This can + // safely relocate the Sort through compatible operators without + // discarding an ordering consumed by an ancestor. let fetch = min_fetch(plan.fetch(), parent_fetch); sort_push_down = sort_push_down .children @@ -178,6 +201,7 @@ fn pushdown_sorts_helper( .cloned() .unwrap_or(Distribution::UnspecifiedDistribution); } + assign_input_order_preservation(&mut sort_push_down, preserve_input_order); return Ok(Transformed::no(sort_push_down)); }; @@ -196,6 +220,22 @@ fn pushdown_sorts_helper( sort_ordering.clone().into(), ); + // Replacing this Sort with an incompatible ordering would change the + // rows observed by the order-sensitive ancestor. Keep this Sort and + // satisfy the parent's independent requirement above it. + if preserve_input_order && !satisfy_parent && !parent_is_stricter { + sort_push_down.plan = plan; + sort_push_down = add_sort_above_with_distribution( + sort_push_down, + parent_requirement.into_single(), + parent_fetch, + &parent_distribution, + ); + assign_initial_requirements(&mut sort_push_down); + sort_push_down.children[0].data.preserve_input_order = true; + return Ok(Transformed::yes(sort_push_down)); + } + // Remove the current sort as we are either going to prove that it is // unnecessary, or replace it with a stricter sort. sort_push_down = sort_push_down @@ -217,6 +257,7 @@ fn pushdown_sorts_helper( sort_push_down.children[0].data = ParentRequirements { ordering_requirement: Some(OrderingRequirements::from(sort_ordering)), fetch: sort_fetch, + preserve_input_order: false, distribution_requirement: Distribution::UnspecifiedDistribution, }; return Ok(Transformed::yes(sort_push_down)); @@ -274,6 +315,7 @@ fn pushdown_sorts_helper( .unwrap_or(&Distribution::UnspecifiedDistribution), ); } + assign_input_order_preservation(&mut sort_push_down, preserve_input_order); } else if let Some(adjusted) = pushdown_requirement_to_children( &sort_push_down.plan, parent_requirement.clone(), @@ -305,6 +347,7 @@ fn pushdown_sorts_helper( .unwrap_or(&Distribution::UnspecifiedDistribution), ); } + assign_input_order_preservation(&mut sort_push_down, preserve_input_order); sort_push_down.data.ordering_requirement = None; } else { // Can not push down requirements, add new `SortExec` (distribution-aware): @@ -326,6 +369,26 @@ fn pushdown_requirement_to_children( parent_required: OrderingRequirements, parent_fetch: Option, ) -> Result>>> { + // A concrete ordering from above can cross an order-sensitive operator only + // when it is compatible with the operator's existing output ordering. A + // finer compatible ordering is safe, while an incompatible ordering could + // change which rows the operator produces. + if plan + .requires_input_order_preservation() + .into_iter() + .any(|preserve| preserve) + { + let Some(output_ordering) = plan.output_ordering() else { + return Ok(None); + }; + if !plan.equivalence_properties().requirements_compatible( + parent_required.first().clone(), + output_ordering.clone().into(), + ) { + return Ok(None); + } + } + // If there is a limit on the parent plan we cannot push it down through operators that change the cardinality. // E.g. consider if LIMIT 2 is applied below a FilteExec that filters out 1/2 of the rows we'll end up with 1 row instead of 2. // If the LIMIT is applied after the FilterExec and the FilterExec returns > 2 rows we'll end up with 2 rows (correct). diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs index 41a03bb031629..e5253d03cbc71 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs @@ -194,14 +194,15 @@ impl PhysicalOptimizerRule for EnsureRequirements { // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass) // For each node: distribution first, then sorting. - use super::enforce_distribution::{DistributionContext, ensure_distribution}; + use super::enforce_distribution::{ + DistributionContext, ensure_distribution_with_order_preservation, + }; use super::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting}; // Step 2a: Distribution enforcement (bottom-up) let dist_ctx = DistributionContext::new_default(plan); - let dist_ctx = dist_ctx - .transform_up(|ctx| ensure_distribution(ctx, config)) - .data()?; + let dist_ctx = + ensure_distribution_with_order_preservation(dist_ctx, config, false)?; // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan); diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 3b9d5d258a838..e6989a8bc32b1 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -202,6 +202,22 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { vec![None; self.children().len()] } + /// Specifies whether the identity or number of rows produced by this plan + /// depends on the order of each child input. + /// + /// Unlike [`Self::required_input_ordering`], this property does not require + /// a particular ordering that can be expressed using the child's schema. + /// It only requires that an existing input ordering is not discarded or + /// replaced with an incompatible ordering by physical optimizations. For + /// example, Limit must preserve an input ordering established by a Sort + /// even when an intervening Projection removes the sort columns from its + /// output schema. + /// + /// The default implementation returns `false` for every child. + fn requires_input_order_preservation(&self) -> Vec { + vec![false; self.children().len()] + } + /// Returns `false` if this `ExecutionPlan`'s implementation may reorder /// rows within or between partitions. /// @@ -1321,6 +1337,7 @@ pub fn check_default_invariants( check_len!(plan, maintains_input_order, children_len); check_len!(plan, required_input_ordering, children_len); + check_len!(plan, requires_input_order_preservation, children_len); check_len!(plan, benefits_from_input_partitioning, children_len); plan.input_distribution_requirements() .check_invariants(plan, check)?; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index a1f6074cb9ae6..4ea99c8301bb6 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -163,6 +163,10 @@ impl ExecutionPlan for GlobalLimitExec { vec![true] } + fn requires_input_order_preservation(&self) -> Vec { + vec![true] + } + fn benefits_from_input_partitioning(&self) -> Vec { vec![false] } @@ -398,6 +402,10 @@ impl ExecutionPlan for LocalLimitExec { vec![true] } + fn requires_input_order_preservation(&self) -> Vec { + vec![true] + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/sqllogictest/test_files/subquery_sort.slt b/datafusion/sqllogictest/test_files/subquery_sort.slt index 080fd57c274d8..b42066d2c3a52 100644 --- a/datafusion/sqllogictest/test_files/subquery_sort.slt +++ b/datafusion/sqllogictest/test_files/subquery_sort.slt @@ -191,3 +191,20 @@ b 5 c 4 d 1 e 1 + +# ORDER BY is semantically required by OFFSET, even when the subquery's +# consumer does not preserve that ordering. +query TI +SELECT grp, COUNT(*) AS n +FROM ( + SELECT grp, sort_key + FROM ( + VALUES ('physical_second', 2), ('sorted_first', 1) + ) AS t(grp, sort_key) + ORDER BY sort_key ASC NULLS LAST + OFFSET 1 +) q +GROUP BY grp +ORDER BY grp; +---- +physical_second 1