Minor: add testing case for add YieldStreamExec and polish docs (#16369)

* Add more testing case for add YieldStreamExec

* fmt

* fix doc

* Update datafusion/physical-optimizer/src/insert_yield_exec.rs

Co-authored-by: Andrew Lamb <[email protected]>

* Update datafusion/physical-optimizer/src/insert_yield_exec.rs

Co-authored-by: Andrew Lamb <[email protected]>

* Address comments

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: zhuqi-lucas

Cargo.lock

@@ -54,3 +54,4 @@ recursive = { workspace = true, optional = true }
 datafusion-expr = { workspace = true }
 datafusion-functions-nested = { workspace = true }
 insta = { workspace = true }
+tokio = { workspace = true }

datafusion/physical-optimizer/Cargo.toml

@@ -15,10 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! The `InsertYieldExec` optimizer rule inspects the physical plan to look for
-//! tight-looping operators and inserts explicit yielding mechanisms (whether
-//! as a separate operator, or via a yielding variant) at leaf nodes to make
-//! the plan cancellation friendly.
+//! The [`InsertYieldExec`] optimizer rule inspects the physical plan to find all leaf
+//! nodes corresponding to tight-looping operators. It first attempts to replace
+//! each leaf with a cooperative-yielding variant via `with_cooperative_yields`,
+//! and only if no built-in variant exists does it wrap the node in a
+//! [`YieldStreamExec`] operator to enforce periodic yielding, ensuring the plan
+//! remains cancellation-friendly.
 
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
@@ -32,9 +34,10 @@ use datafusion_physical_plan::yield_stream::YieldStreamExec;
 use datafusion_physical_plan::ExecutionPlan;
 
 /// `InsertYieldExec` is a [`PhysicalOptimizerRule`] that finds every leaf node in
-/// the plan, and replaces it with a variant that cooperatively yields
-/// either using the its yielding variant given by `with_cooperative_yields`,
-/// or, if none exists, by inserting a [`YieldStreamExec`] operator as a parent.
+/// the plan and replaces it with a variant that yields cooperatively if supported.
+/// If the node does not provide a built-in yielding variant via
+/// [`ExecutionPlan::with_cooperative_yields`], it is wrapped in a [`YieldStreamExec`] parent to
+/// enforce a configured yield frequency.
 pub struct InsertYieldExec {}
 
 impl InsertYieldExec {
@@ -73,10 +76,11 @@ impl PhysicalOptimizerRule for InsertYieldExec {
                     // Not a leaf, keep recursing down.
                     return Ok(Transformed::no(plan));
                 }
+                // For leaf nodes, try to get a built-in cooperative-yielding variant.
                 let new_plan = Arc::clone(&plan)
                     .with_cooperative_yields()
                     .unwrap_or_else(|| {
-                        // Otherwise, insert a `YieldStreamExec` to enforce periodic yielding.
+                        // Only if no built-in variant exists, insert a `YieldStreamExec`.
                         Arc::new(YieldStreamExec::new(plan, yield_period))
                     });
                 Ok(Transformed::new(new_plan, true, TreeNodeRecursion::Jump))
@@ -92,3 +96,27 @@ impl PhysicalOptimizerRule for InsertYieldExec {
         true
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion_common::config::ConfigOptions;
+    use datafusion_physical_plan::{displayable, test::scan_partitioned};
+    use insta::assert_snapshot;
+
+    #[tokio::test]
+    async fn test_yield_stream_exec_for_custom_exec() {
+        let test_custom_exec = scan_partitioned(1);
+        let config = ConfigOptions::new();
+        let optimized = InsertYieldExec::new()
+            .optimize(test_custom_exec, &config)
+            .unwrap();
+
+        let display = displayable(optimized.as_ref()).indent(true).to_string();
+        // Use insta snapshot to ensure full plan structure
+        assert_snapshot!(display, @r###"
+            YieldStreamExec frequency=64
+              DataSourceExec: partitions=1, partition_sizes=[1]
+            "###);
+    }
+}

datafusion/physical-optimizer/src/insert_yield_exec.rs

@@ -92,6 +92,5 @@ pub mod udaf {
 }
 
 pub mod coalesce;
-#[cfg(any(test, feature = "bench"))]
 pub mod test;
 pub mod yield_stream;

datafusion/physical-plan/src/lib.rs

@@ -300,7 +300,7 @@ impl TestMemoryExec {
     }
 
     /// refer to `try_with_sort_information` at MemorySourceConfig for more information.
-    /// https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs
+    /// <https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs>
     pub fn try_with_sort_information(
         mut self,
         mut sort_information: Vec<LexOrdering>,