From dea1b0ef0e86492dedc9ab2ccde95cd8bdb70aa8 Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Wed, 24 Jun 2026 10:37:47 +0300 Subject: [PATCH 1/4] fix(physical-optimizer): capture ORDER BY under ScalarSubqueryExec root require_top_ordering_helper bailed on ScalarSubqueryExec because it has multiple children (main input + subquery plans), so the rule stamped an empty OutputRequirementExec (order_by=[], dist_by=Unspecified) at the root and never recorded the query's global ORDER BY requirement. Descend through child 0 (the order-transparent main input: maintains_input_order()[0] == true, no required input ordering) to find and protect the top SortExec, reattaching the subquery children unchanged. --- .../physical_optimizer/output_requirements.rs | 41 ++++++++++++++++++- .../src/output_requirements.rs | 21 ++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 846589104e4ca..531b348ee5181 100644 --- a/datafusion/core/tests/physical_optimizer/output_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -20,11 +20,12 @@ use std::sync::Arc; use crate::physical_optimizer::test_utils::{parquet_exec, schema, sort_exec, sort_expr}; 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::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::{ExecutionPlan, displayable, get_plan_string}; /// `OutputRequirements::new_add_mode()` must be idempotent: re-applying it to /// its own output must not stack additional `OutputRequirementExec` wrappers. @@ -63,3 +64,39 @@ fn assert_add_mode_idempotent(plan: Arc) { "second invocation of OutputRequirements::new_add_mode mutated the plan", ); } + +/// For a `ScalarSubqueryExec` root, `new_add_mode()` 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 add_mode_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()) + .unwrap(); + + 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 + "); +} diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index fb91ae46a2a08..6b2fef39245b8 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::{ @@ -359,6 +360,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 5f06f1be49febacac569cdbc3c4ec5b6fd003045 Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Sun, 12 Jul 2026 13:58:47 -0700 Subject: [PATCH 2/4] test(sqllogictest): show ORDER BY requirement captured under ScalarSubqueryExec root EXPLAIN VERBOSE pins the 'physical_plan after OutputRequirements' step: the OutputRequirementExec carrying the query's global ORDER BY is placed below the ScalarSubqueryExec root. Without the fix the rule stamps an empty order_by=[], dist_by=Unspecified requirement at the root and this snapshot fails. --- .../sqllogictest/test_files/subquery.slt | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index e38bd6001b43b..62def0bbfb202 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2023,6 +2023,111 @@ ORDER BY i; 9 10 +# The query's global ORDER BY must be captured as an output requirement even +# though the plan root is the multi-child `ScalarSubqueryExec`: the +# `OutputRequirements` rule descends through the order-transparent main input +# (child 0) and wraps the top `SortExec` with an `OutputRequirementExec` +# carrying the ordering (`physical_plan after OutputRequirements` below). +# Without this, the rule stamps an empty +# `OutputRequirementExec: order_by=[], dist_by=Unspecified` at the root and the +# global ORDER BY requirement is lost. +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN VERBOSE SELECT x FROM sq_main WHERE x > (SELECT min(v) FROM sq_values) ORDER BY x; +---- +initial_physical_plan +01)ScalarSubqueryExec: subqueries=1 +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----FilterExec: x@0 > scalar_subquery() +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)] +06)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)] +07)------DataSourceExec: partitions=1, partition_sizes=[1] +initial_physical_plan_with_stats +01)ScalarSubqueryExec: subqueries=1, statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] +03)----FilterExec: x@0 > scalar_subquery(), statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] +04)------DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(2), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] +05)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)], statistics=[Rows=Exact(1), Bytes=Inexact(38), [(Col[0]:)]] +06)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)], statistics=[Rows=Exact(1), Bytes=Inexact(38), [(Col[0]:)]] +07)------DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(3), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] +initial_physical_plan_with_schema +01)ScalarSubqueryExec: subqueries=1, schema=[x:Int32;N] +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], schema=[x:Int32;N] +03)----FilterExec: x@0 > scalar_subquery(), schema=[x:Int32;N] +04)------DataSourceExec: partitions=1, partition_sizes=[1], schema=[x:Int32;N] +05)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)], schema=[min(sq_values.v):Int32;N] +06)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)], schema=[min(sq_values.v)[value]:Int32;N] +07)------DataSourceExec: partitions=1, partition_sizes=[1], schema=[v:Int32;N] +physical_plan after OutputRequirements +01)ScalarSubqueryExec: subqueries=1 +02)--OutputRequirementExec: order_by=[(x@0, asc)], dist_by=SinglePartition +03)----SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------FilterExec: x@0 > scalar_subquery() +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)] +07)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)] +08)------DataSourceExec: partitions=1, partition_sizes=[1] +physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after join_selection SAME TEXT AS ABOVE +physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after FilterPushdown SAME TEXT AS ABOVE +physical_plan after EnsureRequirements SAME TEXT AS ABOVE +physical_plan after CombinePartialFinalAggregate +01)ScalarSubqueryExec: subqueries=1 +02)--OutputRequirementExec: order_by=[(x@0, asc)], dist_by=SinglePartition +03)----SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------FilterExec: x@0 > scalar_subquery() +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)] +07)----DataSourceExec: partitions=1, partition_sizes=[1] +physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE +physical_plan after WindowTopN SAME TEXT AS ABOVE +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after OutputRequirements +01)ScalarSubqueryExec: subqueries=1 +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----FilterExec: x@0 > scalar_subquery() +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)] +06)----DataSourceExec: partitions=1, partition_sizes=[1] +physical_plan after LimitAggregation SAME TEXT AS ABOVE +physical_plan after LimitPushPastWindows SAME TEXT AS ABOVE +physical_plan after HashJoinBuffering SAME TEXT AS ABOVE +physical_plan after LimitPushdown SAME TEXT AS ABOVE +physical_plan after TopKRepartition SAME TEXT AS ABOVE +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after PushdownSort SAME TEXT AS ABOVE +physical_plan after EnsureCooperative SAME TEXT AS ABOVE +physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE +physical_plan after SanityCheckPlan SAME TEXT AS ABOVE +physical_plan +01)ScalarSubqueryExec: subqueries=1 +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----FilterExec: x@0 > scalar_subquery() +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)] +06)----DataSourceExec: partitions=1, partition_sizes=[1] +physical_plan_with_stats +01)ScalarSubqueryExec: subqueries=1, statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] +03)----FilterExec: x@0 > scalar_subquery(), statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] +04)------DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(2), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] +05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)], statistics=[Rows=Exact(1), Bytes=Inexact(38), [(Col[0]:)]] +06)----DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(3), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] +physical_plan_with_schema +01)ScalarSubqueryExec: subqueries=1, schema=[x:Int32;N] +02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], schema=[x:Int32;N] +03)----FilterExec: x@0 > scalar_subquery(), schema=[x:Int32;N] +04)------DataSourceExec: partitions=1, partition_sizes=[1], schema=[x:Int32;N] +05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)], schema=[min(sq_values.v):Int32;N] +06)----DataSourceExec: partitions=1, partition_sizes=[1], schema=[v:Int32;N] + +statement ok +set datafusion.explain.physical_plan_only = false; + # Aggregate function argument containing uncorrelated scalar subquery query I SELECT sum(x + (SELECT max(v) FROM sq_values)) AS s FROM sq_main; From c1c5e4261c2fe7b799bf75d09ee230c6d96bd52a Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Sun, 12 Jul 2026 14:35:55 -0700 Subject: [PATCH 3/4] test: replace EXPLAIN VERBOSE slt with end-to-end wrong-results test Drop the pass-pinning EXPLAIN VERBOSE block from subquery.slt: SQL over built-in sources cannot produce wrong results from this defect (the top SortExec and the requirement-independent merge backstops in EnsureRequirements always restore the global order), so the slt could only ever assert on intermediate pass output while pinning the full pass list. Instead add full_pipeline_keeps_row_order_under_scalar_subquery_root: it runs the complete default physical optimizer pipeline over a ScalarSubqueryExec root whose ordering comes from a SortPreservingMergeExec over a two-partition ordered source (the shape federated planners produce, with no SortExec backstop) and executes the plan. Without the fix the empty root requirement lets the pipeline drop the merge and the rows come back partition-interleaved. --- .../physical_optimizer/output_requirements.rs | 98 +++++++++++++++- .../sqllogictest/test_files/subquery.slt | 105 ------------------ 2 files changed, 97 insertions(+), 106 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 531b348ee5181..1f8444509d39f 100644 --- a/datafusion/core/tests/physical_optimizer/output_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -19,13 +19,21 @@ use std::sync::Arc; use crate::physical_optimizer::test_utils::{parquet_exec, schema, sort_exec, sort_expr}; +use arrow::array::{Int32Array, cast::AsArray, types::Int32Type}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +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::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; -use datafusion_physical_plan::{ExecutionPlan, displayable, get_plan_string}; +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. @@ -100,3 +108,91 @@ fn add_mode_descends_through_scalar_subquery() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); } + +/// End-to-end wrong-results guard: run the full default physical optimizer +/// pipeline over a `ScalarSubqueryExec` root whose global ordering is provided +/// by a `SortPreservingMergeExec` over a multi-partition ordered source — the +/// shape federated/custom planners hand to the optimizer, with no `SortExec` +/// backstop. If `OutputRequirements` fails to record the ordering under the +/// multi-child root, the empty root requirement (no ordering, unspecified +/// distribution) lets the pipeline eliminate the merge, and the query returns +/// rows in partition-interleaved order — actually incorrect results, not +/// merely a different plan. +#[tokio::test] +async fn full_pipeline_keeps_row_order_under_scalar_subquery_root() { + // Two partitions, each sorted on `a`; only an order-preserving merge + // yields the global order. + let s = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions: Vec> = [vec![1, 3, 5, 7], vec![2, 4, 6, 8]] + .into_iter() + .map(|v| { + vec![ + RecordBatch::try_new(Arc::clone(&s), vec![Arc::new(Int32Array::from(v))]) + .unwrap(), + ] + }) + .collect(); + let ordering: LexOrdering = [sort_expr("a", &s)].into(); + let source = DataSourceExec::from_data_source( + MemorySourceConfig::try_new(&partitions, Arc::clone(&s), None) + .unwrap() + .try_with_sort_information(vec![ordering.clone()]) + .unwrap(), + ); + let merge = Arc::new(SortPreservingMergeExec::new(ordering, source)); + + // A scalar subquery child must produce exactly one column and one row. + let sq_schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let sq_batch = RecordBatch::try_new( + Arc::clone(&sq_schema), + vec![Arc::new(Int32Array::from(vec![42]))], + ) + .unwrap(); + let sq_plan = + MemorySourceConfig::try_new_exec(&[vec![sq_batch]], sq_schema, None).unwrap(); + + let plan = Arc::new(ScalarSubqueryExec::new( + merge, + vec![ScalarSubqueryLink { + plan: sq_plan, + index: SubqueryIndex::new(0), + }], + ScalarSubqueryResults::new(1), + )) as Arc; + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 4; + // Keep the plan focused on the merge-vs-coalesce decision rather than on + // added parallelism. + config.optimizer.enable_round_robin_repartition = false; + config.optimizer.prefer_existing_sort = true; + + let mut optimized = plan; + for rule in PhysicalOptimizer::new().rules { + optimized = rule.optimize(optimized, &config).unwrap(); + } + + // The order-preserving merge must survive the pipeline; the failure mode + // replaces it with a `CoalescePartitionsExec`. + let display = displayable(optimized.as_ref()).indent(true).to_string(); + assert!( + display.contains("SortPreservingMergeExec") + && !display.contains("CoalescePartitionsExec"), + "order-preserving merge was optimized away:\n{display}" + ); + + // And the rows must actually come back globally ordered. + let ctx = SessionContext::new(); + let batches = collect(optimized, ctx.task_ctx()).await.unwrap(); + 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/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 62def0bbfb202..e38bd6001b43b 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2023,111 +2023,6 @@ ORDER BY i; 9 10 -# The query's global ORDER BY must be captured as an output requirement even -# though the plan root is the multi-child `ScalarSubqueryExec`: the -# `OutputRequirements` rule descends through the order-transparent main input -# (child 0) and wraps the top `SortExec` with an `OutputRequirementExec` -# carrying the ordering (`physical_plan after OutputRequirements` below). -# Without this, the rule stamps an empty -# `OutputRequirementExec: order_by=[], dist_by=Unspecified` at the root and the -# global ORDER BY requirement is lost. -statement ok -set datafusion.explain.physical_plan_only = true; - -query TT -EXPLAIN VERBOSE SELECT x FROM sq_main WHERE x > (SELECT min(v) FROM sq_values) ORDER BY x; ----- -initial_physical_plan -01)ScalarSubqueryExec: subqueries=1 -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] -03)----FilterExec: x@0 > scalar_subquery() -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)] -06)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)] -07)------DataSourceExec: partitions=1, partition_sizes=[1] -initial_physical_plan_with_stats -01)ScalarSubqueryExec: subqueries=1, statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] -03)----FilterExec: x@0 > scalar_subquery(), statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] -04)------DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(2), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] -05)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)], statistics=[Rows=Exact(1), Bytes=Inexact(38), [(Col[0]:)]] -06)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)], statistics=[Rows=Exact(1), Bytes=Inexact(38), [(Col[0]:)]] -07)------DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(3), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] -initial_physical_plan_with_schema -01)ScalarSubqueryExec: subqueries=1, schema=[x:Int32;N] -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], schema=[x:Int32;N] -03)----FilterExec: x@0 > scalar_subquery(), schema=[x:Int32;N] -04)------DataSourceExec: partitions=1, partition_sizes=[1], schema=[x:Int32;N] -05)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)], schema=[min(sq_values.v):Int32;N] -06)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)], schema=[min(sq_values.v)[value]:Int32;N] -07)------DataSourceExec: partitions=1, partition_sizes=[1], schema=[v:Int32;N] -physical_plan after OutputRequirements -01)ScalarSubqueryExec: subqueries=1 -02)--OutputRequirementExec: order_by=[(x@0, asc)], dist_by=SinglePartition -03)----SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] -04)------FilterExec: x@0 > scalar_subquery() -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)--AggregateExec: mode=Final, gby=[], aggr=[min(sq_values.v)] -07)----AggregateExec: mode=Partial, gby=[], aggr=[min(sq_values.v)] -08)------DataSourceExec: partitions=1, partition_sizes=[1] -physical_plan after aggregate_statistics SAME TEXT AS ABOVE -physical_plan after join_selection SAME TEXT AS ABOVE -physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after FilterPushdown SAME TEXT AS ABOVE -physical_plan after EnsureRequirements SAME TEXT AS ABOVE -physical_plan after CombinePartialFinalAggregate -01)ScalarSubqueryExec: subqueries=1 -02)--OutputRequirementExec: order_by=[(x@0, asc)], dist_by=SinglePartition -03)----SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] -04)------FilterExec: x@0 > scalar_subquery() -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)] -07)----DataSourceExec: partitions=1, partition_sizes=[1] -physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE -physical_plan after WindowTopN SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE -physical_plan after OutputRequirements -01)ScalarSubqueryExec: subqueries=1 -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] -03)----FilterExec: x@0 > scalar_subquery() -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)] -06)----DataSourceExec: partitions=1, partition_sizes=[1] -physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after LimitPushPastWindows SAME TEXT AS ABOVE -physical_plan after HashJoinBuffering SAME TEXT AS ABOVE -physical_plan after LimitPushdown SAME TEXT AS ABOVE -physical_plan after TopKRepartition SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE -physical_plan after PushdownSort SAME TEXT AS ABOVE -physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE -physical_plan after SanityCheckPlan SAME TEXT AS ABOVE -physical_plan -01)ScalarSubqueryExec: subqueries=1 -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] -03)----FilterExec: x@0 > scalar_subquery() -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)] -06)----DataSourceExec: partitions=1, partition_sizes=[1] -physical_plan_with_stats -01)ScalarSubqueryExec: subqueries=1, statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] -03)----FilterExec: x@0 > scalar_subquery(), statistics=[Rows=Inexact(1), Bytes=Inexact(23), [(Col[0]: Null=Exact(0))]] -04)------DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(2), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] -05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)], statistics=[Rows=Exact(1), Bytes=Inexact(38), [(Col[0]:)]] -06)----DataSourceExec: partitions=1, partition_sizes=[1], statistics=[Rows=Exact(3), Bytes=Exact(112), [(Col[0]: Null=Exact(0))]] -physical_plan_with_schema -01)ScalarSubqueryExec: subqueries=1, schema=[x:Int32;N] -02)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false], schema=[x:Int32;N] -03)----FilterExec: x@0 > scalar_subquery(), schema=[x:Int32;N] -04)------DataSourceExec: partitions=1, partition_sizes=[1], schema=[x:Int32;N] -05)--AggregateExec: mode=Single, gby=[], aggr=[min(sq_values.v)], schema=[min(sq_values.v):Int32;N] -06)----DataSourceExec: partitions=1, partition_sizes=[1], schema=[v:Int32;N] - -statement ok -set datafusion.explain.physical_plan_only = false; - # Aggregate function argument containing uncorrelated scalar subquery query I SELECT sum(x + (SELECT max(v) FROM sq_values)) AS s FROM sq_main; From a2ce17cb86cb027582792fe00913dd4b0c30a04e Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Fri, 17 Jul 2026 10:30:23 -0700 Subject: [PATCH 4/4] test(output-requirements): clarify and simplify scalar-subquery ordering tests Rename the tests to name the unit each one exercises: the rule traversal (require_top_ordering_descends_through_scalar_subquery) and the end-to-end ordering guarantee (scalar_subquery_root_preserves_global_ordering_end_to_end). Rebuild the end-to-end test around record_batch!, drop the redundant plan-shape assertion and the unneeded optimizer config knobs (verified: the executed-row-order check alone still catches the regression), and describe the expected behavior rather than the failure mode. Replace unwrap() with descriptive expect(). --- .../physical_optimizer/output_requirements.rs | 121 ++++++++---------- 1 file changed, 54 insertions(+), 67 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs index 1f8444509d39f..4463986e42931 100644 --- a/datafusion/core/tests/physical_optimizer/output_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -19,9 +19,7 @@ use std::sync::Arc; use crate::physical_optimizer::test_utils::{parquet_exec, schema, sort_exec, sort_expr}; -use arrow::array::{Int32Array, cast::AsArray, types::Int32Type}; -use arrow::datatypes::{DataType, Field, Schema}; -use arrow::record_batch::RecordBatch; +use arrow::array::{cast::AsArray, record_batch, types::Int32Type}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::SessionContext; @@ -63,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), @@ -73,12 +75,13 @@ fn assert_add_mode_idempotent(plan: Arc) { ); } -/// For a `ScalarSubqueryExec` root, `new_add_mode()` 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. +/// 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 add_mode_descends_through_scalar_subquery() { +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))); @@ -96,7 +99,7 @@ fn add_mode_descends_through_scalar_subquery() { let optimized = OutputRequirements::new_add_mode() .optimize(plan, &ConfigOptions::new()) - .unwrap(); + .expect("add-mode optimize should succeed"); insta::assert_snapshot!( displayable(optimized.as_ref()).indent(true).to_string(), @@ -109,81 +112,65 @@ fn add_mode_descends_through_scalar_subquery() { "); } -/// End-to-end wrong-results guard: run the full default physical optimizer -/// pipeline over a `ScalarSubqueryExec` root whose global ordering is provided -/// by a `SortPreservingMergeExec` over a multi-partition ordered source — the -/// shape federated/custom planners hand to the optimizer, with no `SortExec` -/// backstop. If `OutputRequirements` fails to record the ordering under the -/// multi-child root, the empty root requirement (no ordering, unspecified -/// distribution) lets the pipeline eliminate the merge, and the query returns -/// rows in partition-interleaved order — actually incorrect results, not -/// merely a different plan. +/// 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 full_pipeline_keeps_row_order_under_scalar_subquery_root() { - // Two partitions, each sorted on `a`; only an order-preserving merge - // yields the global order. - let s = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let partitions: Vec> = [vec![1, 3, 5, 7], vec![2, 4, 6, 8]] - .into_iter() - .map(|v| { - vec![ - RecordBatch::try_new(Arc::clone(&s), vec![Arc::new(Int32Array::from(v))]) - .unwrap(), - ] - }) - .collect(); - let ordering: LexOrdering = [sort_expr("a", &s)].into(); +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(&partitions, Arc::clone(&s), None) - .unwrap() + MemorySourceConfig::try_new(&[vec![p1], vec![p2]], Arc::clone(&schema), None) + .expect("build memory source config") .try_with_sort_information(vec![ordering.clone()]) - .unwrap(), + .expect("attach sort information to source"), ); - let merge = Arc::new(SortPreservingMergeExec::new(ordering, source)); - - // A scalar subquery child must produce exactly one column and one row. - let sq_schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); - let sq_batch = RecordBatch::try_new( - Arc::clone(&sq_schema), - vec![Arc::new(Int32Array::from(vec![42]))], + // 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, ) - .unwrap(); - let sq_plan = - MemorySourceConfig::try_new_exec(&[vec![sq_batch]], sq_schema, None).unwrap(); + .expect("build subquery exec"); let plan = Arc::new(ScalarSubqueryExec::new( - merge, + main_input, vec![ScalarSubqueryLink { - plan: sq_plan, + 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; - // Keep the plan focused on the merge-vs-coalesce decision rather than on - // added parallelism. - config.optimizer.enable_round_robin_repartition = false; - config.optimizer.prefer_existing_sort = true; - let mut optimized = plan; for rule in PhysicalOptimizer::new().rules { - optimized = rule.optimize(optimized, &config).unwrap(); + optimized = rule + .optimize(optimized, &config) + .unwrap_or_else(|e| panic!("optimizer rule {} failed: {e}", rule.name())); } - // The order-preserving merge must survive the pipeline; the failure mode - // replaces it with a `CoalescePartitionsExec`. - let display = displayable(optimized.as_ref()).indent(true).to_string(); - assert!( - display.contains("SortPreservingMergeExec") - && !display.contains("CoalescePartitionsExec"), - "order-preserving merge was optimized away:\n{display}" - ); - - // And the rows must actually come back globally ordered. - let ctx = SessionContext::new(); - let batches = collect(optimized, ctx.task_ctx()).await.unwrap(); + // 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| {