From f057361e2e0fb010f662d84bcb9136f24fde4800 Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Fri, 17 Jul 2026 10:55:21 -0700 Subject: [PATCH 1/4] fix: Capture global ORDER BY requirement under ScalarSubqueryExec root `OutputRequirements` (add mode) walks the plan via `require_top_ordering_helper` to capture the query's global ORDER BY as an `OutputRequirementExec`. The helper bails on any node with `children.len() != 1`. `ScalarSubqueryExec` is such a node (child 0 is the order-transparent main input, the rest are uncorrelated subquery plans), so when it is the plan root the rule stamps an empty `OutputRequirementExec(order_by=[], dist_by=Unspecified)` at the top and never records the global ordering requirement. `ScalarSubqueryExec` is order-transparent on child 0 (`maintains_input_order()[0] == true`, no required input ordering), so `require_top_ordering_helper` now descends through child 0 to find and wrap the top `SortExec` / `SortPreservingMergeExec`, reattaching the subquery children unchanged. Tested at two levels in `datafusion/core/tests/physical_optimizer/output_requirements.rs`: - `require_top_ordering_descends_through_scalar_subquery` (rule level) - `scalar_subquery_root_preserves_global_ordering_end_to_end` (end to end) --- .../physical_optimizer/output_requirements.rs | 128 +++++++++++++++++- .../src/output_requirements.rs | 21 +++ 2 files changed, 145 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 846589104e4ca..4463986e42931 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::execution_props::{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. @@ -54,8 +61,12 @@ fn assert_add_mode_idempotent(plan: Arc) { 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 +74,112 @@ fn assert_add_mode_idempotent(plan: Arc) { "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; + + 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; + + // 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 = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_primitive::() + .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 c6f5f87622bea..f0808c8809b19 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -37,6 +37,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::{ @@ -380,6 +381,26 @@ fn require_top_ordering_helper( plan: Arc, ) -> Result<(Arc, bool)> { let mut children = plan.children(); + + // `ScalarSubqueryExec` is a multi-child but order-transparent root: child 0 is + // the main input (it copies that child's `PlanProperties` and reports + // `maintains_input_order()[0] == true` with no required input ordering), while + // the remaining children are subquery plans that don't contribute to output + // ordering. The generic `children.len() != 1` guard below would stop the search + // at this node and lose the query's global ORDER BY (the top `SortExec` lives + // below the main input), so descend through child 0 and reattach the rest. + if plan.downcast_ref::().is_some() { + let (new_main, is_changed) = + require_top_ordering_helper(Arc::clone(children[0]))?; + if is_changed { + let mut new_children: Vec> = + children.iter().map(|&c| Arc::clone(c)).collect(); + new_children[0] = new_main; + return Ok((plan.with_new_children(new_children)?, true)); + } + return Ok((plan, false)); + } + // Global ordering defines desired ordering in the final result. if children.len() != 1 { Ok((plan, false)) From 2cf98ce60d3394200f7a99043d9a9b42ed91d74b Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Mon, 20 Jul 2026 15:11:10 -0700 Subject: [PATCH 2/4] refactor: extract output_requirement_child helper for require_top_ordering Address review feedback: isolate the ScalarSubqueryExec descent decision in a single output_requirement_child() helper and route both single-child operators and ScalarSubqueryExec through one index-based descent path in require_top_ordering_helper, removing the dedicated subquery block and the children.len() != 1 guard. Behavior is unchanged. Also generalize the search comment to mention SortPreservingMergeExec alongside SortExec. --- .../src/output_requirements.rs | 81 +++++++++---------- 1 file changed, 38 insertions(+), 43 deletions(-) diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index f0808c8809b19..da0930e421994 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -374,37 +374,30 @@ fn require_top_ordering(plan: Arc) -> Result Option { + if plan.children().len() == 1 { + Some(0) + } else if plan.downcast_ref::().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, ) -> Result<(Arc, bool)> { - let mut children = plan.children(); - - // `ScalarSubqueryExec` is a multi-child but order-transparent root: child 0 is - // the main input (it copies that child's `PlanProperties` and reports - // `maintains_input_order()[0] == true` with no required input ordering), while - // the remaining children are subquery plans that don't contribute to output - // ordering. The generic `children.len() != 1` guard below would stop the search - // at this node and lose the query's global ORDER BY (the top `SortExec` lives - // below the main input), so descend through child 0 and reattach the rest. - if plan.downcast_ref::().is_some() { - let (new_main, is_changed) = - require_top_ordering_helper(Arc::clone(children[0]))?; - if is_changed { - let mut new_children: Vec> = - children.iter().map(|&c| Arc::clone(c)).collect(); - new_children[0] = new_main; - return Ok((plan.with_new_children(new_children)?, true)); - } - return Ok((plan, false)); - } - // 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::() { + if let Some(sort_exec) = plan.downcast_ref::() { // 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. @@ -437,25 +430,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> = + 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)) From 681c346dcc3d5be2577260e7f8b62c74f2bec10c Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Mon, 20 Jul 2026 15:19:59 -0700 Subject: [PATCH 3/4] fix: make OutputRequirements(Add) idempotent for below-root wrappers The root-only guard in require_top_ordering only detects a previously-added OutputRequirementExec when it sits at the plan root. When the wrapper lands deeper (e.g. below a ScalarSubqueryExec, as this rule now produces), a second Add pass descended past it, returned is_changed = false, and stamped a redundant empty OutputRequirementExec on top - so the rule was not idempotent under AQE replans (datafusion-ballista#1359). require_top_ordering_helper now treats an existing OutputRequirementExec as already-handled (is_changed = true), stopping the search without adding a redundant wrapper. Add add_mode_is_idempotent_on_scalar_subquery covering the below-root case; the pre-existing idempotency tests only exercised root wrappers. --- .../physical_optimizer/output_requirements.rs | 24 +++++++++++++++++++ .../src/output_requirements.rs | 14 +++++++---- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 4463986e42931..5030580f881a2 100644 --- a/datafusion/core/tests/physical_optimizer/output_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -57,6 +57,30 @@ 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; + + assert_add_mode_idempotent(plan); +} + fn assert_add_mode_idempotent(plan: Arc) { let config = ConfigOptions::new(); let rule = OutputRequirements::new_add_mode(); diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index da0930e421994..fd654c6317e1e 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -351,10 +351,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) -> Result> { if plan.downcast_ref::().is_some() { return Ok(plan); @@ -396,6 +396,12 @@ fn output_requirement_child(plan: &dyn ExecutionPlan) -> Option { fn require_top_ordering_helper( plan: Arc, ) -> Result<(Arc, bool)> { + // A previous run of this rule already captured the ordering requirement at + // this node. Report it as already handled. + if plan.downcast_ref::().is_some() { + return Ok((plan, true)); + } + // Global ordering defines desired ordering in the final result. if let Some(sort_exec) = plan.downcast_ref::() { // In case of constant columns, output ordering of the `SortExec` would From 3ff2f2f495aca91a2ed7bb6b9ed0a105290e9872 Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Mon, 20 Jul 2026 16:53:31 -0700 Subject: [PATCH 4/4] fix: update ScalarSubqueryResults/SubqueryIndex import to physical_planning_context The merge from main relocated these types from execution_props to the new physical_planning_context module, but the import in the output_requirements test was missed, breaking the core_integration test build. --- datafusion/core/tests/physical_optimizer/output_requirements.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 5030580f881a2..79b47dc4418a7 100644 --- a/datafusion/core/tests/physical_optimizer/output_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -24,7 +24,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::SessionContext; use datafusion_common::config::ConfigOptions; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +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;