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

## Which issue does this PR close?

Related to apache/datafusion-ballista#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) <noreply@anthropic.com>

Author: wirybeaver

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;

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",
+    );
+}

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)