fix: Capture global ORDER BY requirement under ScalarSubqueryExec root

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,12 +61,125 @@ 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),
         get_plan_string(&twice),
         "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]);
+}

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<dyn ExecutionPlan>,
 ) -> Result<(Arc<dyn ExecutionPlan>, 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::<ScalarSubqueryExec>().is_some() {
+        let (new_main, is_changed) =
+            require_top_ordering_helper(Arc::clone(children[0]))?;
+        if is_changed {
+            let mut new_children: Vec<Arc<dyn ExecutionPlan>> =
+                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))