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))