This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23677-eef101769d650b2b06ac2bcf298f71113a047c9d in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 9eae6da6fcddf7d1403b4fe5ccd4343d3b0386d7 Author: Sergei Grebnov <[email protected]> AuthorDate: Tue Jul 21 14:22:02 2026 -0700 fix: Capture global ORDER BY requirement under ScalarSubqueryExec root (#23677) ## Which issue does this PR close? After the DataFusion 53 → 54 upgrade, some queries with a `ScalarSubquery` stopped respecting the global `ORDER BY` and returned unordered results. ## Rationale for this change In add mode, `OutputRequirements` walks the plan via `require_top_ordering_helper` to capture the global `ORDER BY` as an `OutputRequirementExec`. The helper bails on any node with `children.len() != 1`. `ScalarSubqueryExec` has multiple children (child 0 = order-transparent main input, the rest = uncorrelated subquery plans), so when it is the root the rule stops immediately and stamps an empty `OutputRequirementExec(order_by=[], dist_by=Unspecified)` at the top — losing the ordering requirement. `ScalarSubqueryExec` is order-transparent on child 0 (`maintains_input_order()[0]== true`, no required input ordering), so the search should descend through it like any other single-child order-preserving operator. ## What changes are included in this PR? `require_top_ordering_helper` now special-cases `ScalarSubqueryExec`: it descends through child 0 to find and wrap the top `SortExec` / `SortPreservingMergeExec`, reattaching the subquery children unchanged. ## Are these changes tested? Yes — two tests in `datafusion/core/tests/physical_optimizer/output_requirements.rs`: 1. **Rule level** (`require_top_ordering_descends_through_scalar_subquery`): asserts the `OutputRequirementExec` carrying the ordering is placed *below* the `ScalarSubqueryExec`; without the fix it lands empty at the root. 2. **End to end** (`scalar_subquery_root_preserves_global_ordering_end_to_end`): runs the full default optimizer pipeline over a `ScalarSubqueryExec` root whose ordering comes from a `SortPreservingMergeExec` over a two-partition ordered source, then executes the plan and checks the rows. Verified by toggling the fix: | | executed output | |---|---| | **without fix** | `1, 3, 5, 7, 2, 4, 6, 8` — merge dropped, rows partition-interleaved | | **with fix** | `1, 2, 3, 4, 5, 6, 7, 8` — globally ordered | Existing `subquery.slt` / TPC-H snapshots are unchanged. I was unable to construct a SQL query that triggers the bug - over built-in sources the global `ORDER BY` is always preserved regardless of the missing requirement. It only surfaces when a custom physical planner supplies a plan shaped like the end-to-end test: a `SortPreservingMergeExec` over already-sorted partitions, with no `SortExec` above it. ## Are there any user-facing changes? No. --- .../physical_optimizer/output_requirements.rs | 152 ++++++++++++++++++++- .../physical-optimizer/src/output_requirements.rs | 80 +++++++---- 2 files changed, 199 insertions(+), 33 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 846589104e..79b47dc441 100644 --- a/datafusion/core/tests/physical_optimizer/output_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -19,12 +19,19 @@ use std::sync::Arc; use crate::physical_optimizer::test_utils::{parquet_exec, schema, sort_exec, sort_expr}; +use arrow::array::{cast::AsArray, record_batch, types::Int32Type}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::prelude::SessionContext; use datafusion_common::config::ConfigOptions; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::optimizer::PhysicalOptimizer; use datafusion_physical_optimizer::output_requirements::OutputRequirements; -use datafusion_physical_plan::ExecutionPlan; -use datafusion_physical_plan::get_plan_string; +use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ExecutionPlan, collect, displayable, get_plan_string}; /// `OutputRequirements::new_add_mode()` must be idempotent: re-applying it to /// its own output must not stack additional `OutputRequirementExec` wrappers. @@ -50,12 +57,40 @@ fn add_mode_is_idempotent_on_sorted_plan() { assert_add_mode_idempotent(plan); } +#[test] +fn add_mode_is_idempotent_on_scalar_subquery() { + // Exercises the below-root case: the wrapper carrying the ordering lands + // under the `ScalarSubqueryExec`, so the root guard in `require_top_ordering` + // does not fire on the second pass. Without treating the existing wrapper as + // already-handled, the second pass would stamp a redundant empty wrapper on + // top of the subquery. + let s = schema(); + let ordering: LexOrdering = [sort_expr("a", &s)].into(); + let sort = sort_exec(ordering, parquet_exec(Arc::clone(&s))); + + let subqueries = vec![ScalarSubqueryLink { + plan: parquet_exec(Arc::clone(&s)), + index: SubqueryIndex::new(0), + }]; + let plan = Arc::new(ScalarSubqueryExec::new( + sort, + subqueries, + ScalarSubqueryResults::new(1), + )) as Arc<dyn ExecutionPlan>; + + assert_add_mode_idempotent(plan); +} + fn assert_add_mode_idempotent(plan: Arc<dyn ExecutionPlan>) { let config = ConfigOptions::new(); let rule = OutputRequirements::new_add_mode(); - let once = rule.optimize(plan, &config).unwrap(); - let twice = rule.optimize(Arc::clone(&once), &config).unwrap(); + let once = rule + .optimize(plan, &config) + .expect("first add-mode optimize pass should succeed"); + let twice = rule + .optimize(Arc::clone(&once), &config) + .expect("second add-mode optimize pass should succeed"); assert_eq!( get_plan_string(&once), @@ -63,3 +98,112 @@ fn assert_add_mode_idempotent(plan: Arc<dyn ExecutionPlan>) { "second invocation of OutputRequirements::new_add_mode mutated the plan", ); } + +/// For a `ScalarSubqueryExec` root, `require_top_ordering_helper` descends +/// through the main input (child 0) and wraps the global `SortExec` with an +/// `OutputRequirementExec` carrying its ordering, leaving the subquery child +/// untouched. Without this, the multi-child root is skipped and the query's +/// global ORDER BY requirement is lost. +#[test] +fn require_top_ordering_descends_through_scalar_subquery() { + let s = schema(); + let ordering: LexOrdering = [sort_expr("a", &s)].into(); + let sort = sort_exec(ordering, parquet_exec(Arc::clone(&s))); + + // A subquery child makes `children.len() == 2`, exercising the multi-child path. + let subqueries = vec![ScalarSubqueryLink { + plan: parquet_exec(Arc::clone(&s)), + index: SubqueryIndex::new(0), + }]; + let plan = Arc::new(ScalarSubqueryExec::new( + sort, + subqueries, + ScalarSubqueryResults::new(1), + )) as Arc<dyn ExecutionPlan>; + + let optimized = OutputRequirements::new_add_mode() + .optimize(plan, &ConfigOptions::new()) + .expect("add-mode optimize should succeed"); + + insta::assert_snapshot!( + displayable(optimized.as_ref()).indent(true).to_string(), + @r" + ScalarSubqueryExec: subqueries=1 + OutputRequirementExec: order_by=[(a@0, asc)], dist_by=SinglePartition + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); +} + +/// A `ScalarSubqueryExec` plan root must preserve its main input's global +/// ordering end to end. +/// +/// The main input is a `SortPreservingMergeExec` over a two-partition ordered +/// source — the shape federated/custom planners hand to the optimizer: an +/// order-preserving merge with no `SortExec` above it. `OutputRequirements` +/// records the global ORDER BY under the multi-child subquery root, the rest of +/// the pipeline keeps the merge, and executing the optimized plan returns the +/// rows in global order regardless of how the source is partitioned. +#[tokio::test] +async fn scalar_subquery_root_preserves_global_ordering_end_to_end() { + // Two partitions, each already sorted on `a`. Global order requires a sort-preserving merge; + // a plain concatenation would interleave them as 1, 3, 5, 7, 2, 4, 6, 8. + let p1 = record_batch!(("a", Int32, [1, 3, 5, 7])).expect("build partition 1 batch"); + let p2 = record_batch!(("a", Int32, [2, 4, 6, 8])).expect("build partition 2 batch"); + let schema = p1.schema(); + let ordering: LexOrdering = [sort_expr("a", &schema)].into(); + let source = DataSourceExec::from_data_source( + MemorySourceConfig::try_new(&[vec![p1], vec![p2]], Arc::clone(&schema), None) + .expect("build memory source config") + .try_with_sort_information(vec![ordering.clone()]) + .expect("attach sort information to source"), + ); + // The main plan establishes the query's global ordering via an `SortPreservingMergeExec` over the two sorted partitions. + let main_input = Arc::new(SortPreservingMergeExec::new(ordering, source)); + + // Dummy subquery that returns a single row + let sq_batch = record_batch!(("v", Int32, [42])).expect("build subquery batch"); + let subquery = MemorySourceConfig::try_new_exec( + &[vec![sq_batch.clone()]], + sq_batch.schema(), + None, + ) + .expect("build subquery exec"); + + let plan = Arc::new(ScalarSubqueryExec::new( + main_input, + vec![ScalarSubqueryLink { + plan: subquery, + index: SubqueryIndex::new(0), + }], + ScalarSubqueryResults::new(1), + )) as Arc<dyn ExecutionPlan>; + + // Run the full default physical optimizer pipeline. + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 4; + let mut optimized = plan; + for rule in PhysicalOptimizer::new().rules { + optimized = rule + .optimize(optimized, &config) + .unwrap_or_else(|e| panic!("optimizer rule {} failed: {e}", rule.name())); + } + + // The executed rows come back in global order: the two sorted partitions + // are merged into 1, 2, 3, 4, 5, 6, 7, 8. + let batches = collect(optimized, SessionContext::new().task_ctx()) + .await + .expect("execute optimized plan"); + let values: Vec<i32> = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_primitive::<Int32Type>() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![1, 2, 3, 4, 5, 6, 7, 8]); +} diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 40ac5643de..b9d0d06da1 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -36,6 +36,7 @@ use datafusion_physical_plan::execution_plan::Boundedness; use datafusion_physical_plan::projection::{ ProjectionExec, make_with_child, update_expr, update_ordering_requirement, }; +use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ @@ -62,8 +63,8 @@ impl OutputRequirements { /// top-level [`OutputRequirementExec`] into the physical plan to keep track /// of global ordering and distribution requirements if there are any. /// Note that this rule should run at the beginning. It is idempotent: when - /// invoked on a plan that is already topped by an `OutputRequirementExec`, - /// it returns the plan unchanged. + /// invoked on a plan that already contains an `OutputRequirementExec` (at + /// the root or below it), it returns the plan unchanged. pub fn new_add_mode() -> Self { Self { mode: RuleMode::Add, @@ -357,10 +358,10 @@ impl PhysicalOptimizerRule for OutputRequirements { /// This functions adds ancillary `OutputRequirementExec` to the physical plan, so that /// global requirements are not lost during optimization. /// -/// Idempotent: if the plan is already topped by an `OutputRequirementExec`, it -/// is returned unchanged so that re-running this rule (as adaptive execution -/// in datafusion-ballista AQE does after every completed stage, see -/// datafusion-ballista#1359) does not stack wrappers. +/// Idempotent: re-running this rule (as adaptive execution in datafusion-ballista +/// AQE does after every completed stage, see datafusion-ballista#1359) does not +/// stack wrappers, whether the previously-added `OutputRequirementExec` sits at +/// the root (handled here) or below it (handled in `require_top_ordering_helper`). fn require_top_ordering(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> { if plan.downcast_ref::<OutputRequirementExec>().is_some() { return Ok(plan); @@ -380,17 +381,36 @@ fn require_top_ordering(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn Executio } } +/// Which child (if any) `require_top_ordering_helper` should descend into when +/// searching for the operator that establishes the global ordering. +fn output_requirement_child(plan: &dyn ExecutionPlan) -> Option<usize> { + if plan.children().len() == 1 { + Some(0) + } else if plan.downcast_ref::<ScalarSubqueryExec>().is_some() { + // `ScalarSubqueryExec` is multi-child but order-transparent on child 0 + // (the main input); its other children are subquery plans that don't + // affect output ordering, so descend into child 0. Without this the + // search stops here and loses the query's global ORDER BY. + Some(0) + } else { + None + } +} + /// Helper function that adds an ancillary `OutputRequirementExec` to the given plan. /// First entry in the tuple is resulting plan, second entry indicates whether any /// `OutputRequirementExec` is added to the plan. fn require_top_ordering_helper( plan: Arc<dyn ExecutionPlan>, ) -> Result<(Arc<dyn ExecutionPlan>, bool)> { - let mut children = plan.children(); + // A previous run of this rule already captured the ordering requirement at + // this node. Report it as already handled. + if plan.downcast_ref::<OutputRequirementExec>().is_some() { + return Ok((plan, true)); + } + // Global ordering defines desired ordering in the final result. - if children.len() != 1 { - Ok((plan, false)) - } else if let Some(sort_exec) = plan.downcast_ref::<SortExec>() { + if let Some(sort_exec) = plan.downcast_ref::<SortExec>() { // In case of constant columns, output ordering of the `SortExec` would // be an empty set. Therefore; we check the sort expression field to // assign the requirements. @@ -423,25 +443,27 @@ fn require_top_ordering_helper( )) as _, true, )) - } else if plan.maintains_input_order()[0] - && (plan.required_input_ordering()[0] - .as_ref() - .is_none_or(|o| matches!(o, OrderingRequirements::Soft(_)))) - { - // Keep searching for a `SortExec` as long as ordering is maintained, - // and on-the-way operators do not themselves require an ordering. - // When an operator requires an ordering, any `SortExec` below can not - // be responsible for (i.e. the originator of) the global ordering. - let (new_child, is_changed) = - require_top_ordering_helper(Arc::clone(children.swap_remove(0)))?; - - let plan = if is_changed { - plan.with_new_children(vec![new_child])? - } else { - plan - }; - - Ok((plan, is_changed)) + } else if let Some(idx) = output_requirement_child(plan.as_ref()) { + // Keep searching for a `SortExec` / `SortPreservingMergeExec` as long as + // ordering is maintained, and on-the-way operators do not themselves + // require an ordering. When an operator requires an ordering, any + // `SortExec` below can not be responsible for (i.e. the originator of) + // the global ordering. + if plan.maintains_input_order()[idx] + && plan.required_input_ordering()[idx] + .as_ref() + .is_none_or(|o| matches!(o, OrderingRequirements::Soft(_))) + { + let mut children: Vec<Arc<dyn ExecutionPlan>> = + plan.children().into_iter().map(Arc::clone).collect(); + let (new_child, is_changed) = + require_top_ordering_helper(Arc::clone(&children[idx]))?; + if is_changed { + children[idx] = new_child; + return Ok((plan.with_new_children(children)?, true)); + } + } + Ok((plan, false)) } else { // Stop searching, there is no global ordering desired for the query. Ok((plan, false)) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
