This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 9eae6da6fc fix: Capture global ORDER BY requirement under
ScalarSubqueryExec root (#23677)
9eae6da6fc is described below
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]
