You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/21 21:22:24 UTC
(datafusion) branch gh-readonly-queue/main/pr-23677-eef101769d650b2b06ac2bcf298f71113a047c9d created (now 9eae6da6fc)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23677-eef101769d650b2b06ac2bcf298f71113a047c9d
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 9eae6da6fc fix: Capture global ORDER BY requirement under ScalarSubqueryExec root (#23677)
This branch includes the following new commits:
new 9eae6da6fc fix: Capture global ORDER BY requirement under ScalarSubqueryExec root (#23677)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: Capture global ORDER BY requirement under ScalarSubqueryExec root (#23677)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23677-eef101769d650b2b06ac2bcf298f71113a047c9d
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 9eae6da6fcddf7d1403b4fe5ccd4343d3b0386d7
Author: Sergei Grebnov <se...@gmail.com>
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: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org