(datafusion) 01/01: fix(physical-optimizer): make OutputRequirements idempotent (#22522)

[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-22522-fa724c1ad527c0cd171e14c856299a1e1ad30e01
in repository https://gitbox.apache.org/repos/asf/datafusion.git

commit d9ea38b95123159161c017840d3e6256e41988dd
Author: Xuanyi Li <sh...@outlook.com>
AuthorDate: Sun May 31 01:59:10 2026 -0700

    fix(physical-optimizer): make OutputRequirements idempotent (#22522)
    
    ## Which issue does this PR close?
    
    Related to https://github.com/apache/datafusion-ballista/issues/1359
    
    ## Rationale
    
    Ballista's Adaptive Query Execution (AQE) planner re-invokes
    DataFusion's full `PhysicalOptimizer` chain after every completed stage
    (`AdaptivePlanner::replan_stages`). Rules that are not idempotent
    (`rule(rule(x)) != rule(x)`) stack execution-plan nodes on each pass.
    
    `OutputRequirements::new_add_mode()` wraps the plan root with
    `OutputRequirementExec` to preserve global ordering/distribution
    requirements. On a second pass the wrapper's `maintains_input_order() ==
    [true]` and `required_input_ordering() == [None]` cause
    `require_top_ordering_helper` to recurse through it and produce a
    *second* wrapper, yielding
    `OutputRequirementExec(OutputRequirementExec(...))`. Each AQE replan
    adds another layer.
    
    ## What changes are included in this PR?
    
    - **Guard in `require_top_ordering()`**: if the plan root is already an
    `OutputRequirementExec`, return it unchanged. This makes the rule
    idempotent with zero overhead for single-pass use.
    - **Doc-comment update** on `new_add_mode()` and
    `require_top_ordering()` documenting the idempotence guarantee.
    - **Two tests** in `tests/physical_optimizer/output_requirements.rs`:
    - `add_mode_is_idempotent_on_bare_scan` — bare `ParquetExec` (exercises
    `is_changed = false` path).
    - `add_mode_is_idempotent_on_sorted_plan` — `SortExec → ParquetExec`
    (exercises `is_changed = true` path).
    
    ## Are these changes tested?
    
    Yes. Two new tests run the rule twice on distinct fixtures and assert
    structural equality via `get_plan_string`. Both fail without the fix
    (double-wrapped `OutputRequirementExec`) and pass with it.
    
    ## Are there any user-facing changes?
    
    No. `OutputRequirementExec` is an internal ancillary node stripped
    before execution; the idempotence guard only affects re-optimization
    scenarios (AQE).
    
    Co-authored-by: Claude Opus 4.7 (1M context) <no...@anthropic.com>
---
 datafusion/core/tests/physical_optimizer/mod.rs    |  1 +
 .../physical_optimizer/output_requirements.rs      | 65 ++++++++++++++++++++++
 .../physical-optimizer/src/output_requirements.rs  | 12 +++-
 3 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs
index 801c2f30f9..f3b2884dab 100644
--- a/datafusion/core/tests/physical_optimizer/mod.rs
+++ b/datafusion/core/tests/physical_optimizer/mod.rs
@@ -30,6 +30,7 @@ mod join_selection;
 #[expect(clippy::needless_pass_by_value)]
 mod limit_pushdown;
 mod limited_distinct_aggregation;
+mod output_requirements;
 mod partition_statistics;
 mod projection_pushdown;
 mod pushdown_sort;
diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs
new file mode 100644
index 0000000000..846589104e
--- /dev/null
+++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs
@@ -0,0 +1,65 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use crate::physical_optimizer::test_utils::{parquet_exec, schema, sort_exec, sort_expr};
+
+use datafusion_common::config::ConfigOptions;
+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;
+
+/// `OutputRequirements::new_add_mode()` must be idempotent: re-applying it to
+/// its own output must not stack additional `OutputRequirementExec` wrappers.
+///
+/// AQE (datafusion-ballista#1359) re-runs the optimizer chain after every
+/// completed stage; without this guarantee, every replan adds another wrapper.
+#[test]
+fn add_mode_is_idempotent_on_bare_scan() {
+    // Exercises the path where `require_top_ordering_helper` returns
+    // `is_changed = false` and the rule adds a default (empty-requirement)
+    // wrapper.
+    assert_add_mode_idempotent(parquet_exec(schema()));
+}
+
+#[test]
+fn add_mode_is_idempotent_on_sorted_plan() {
+    // Exercises the path where the helper recognizes a top-level `SortExec`
+    // and produces a wrapper carrying that ordering requirement
+    // (`is_changed = true` branch).
+    let s = schema();
+    let ordering: LexOrdering = [sort_expr("a", &s)].into();
+    let plan = sort_exec(ordering, parquet_exec(Arc::clone(&s)));
+    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();
+
+    assert_eq!(
+        get_plan_string(&once),
+        get_plan_string(&twice),
+        "second invocation of OutputRequirements::new_add_mode mutated the plan",
+    );
+}
diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs
index 24eb3af5f5..899abcc88b 100644
--- a/datafusion/physical-optimizer/src/output_requirements.rs
+++ b/datafusion/physical-optimizer/src/output_requirements.rs
@@ -61,7 +61,9 @@ impl OutputRequirements {
     /// Create a new rule which works in `Add` mode; i.e. it simply adds a
     /// 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.
+    /// 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.
     pub fn new_add_mode() -> Self {
         Self {
             mode: RuleMode::Add,
@@ -325,7 +327,15 @@ 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.
 fn require_top_ordering(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
+    if plan.downcast_ref::<OutputRequirementExec>().is_some() {
+        return Ok(plan);
+    }
     let (new_plan, is_changed) = require_top_ordering_helper(plan)?;
     if is_changed {
         Ok(new_plan)


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org