alamb commented on code in PR #16369:
URL: https://github.com/apache/datafusion/pull/16369#discussion_r2142595175
##########
datafusion/physical-optimizer/src/insert_yield_exec.rs:
##########
@@ -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
+/// `with_cooperative_yields`, it is wrapped in a [`YieldStreamExec`] parent to
Review Comment:
```suggestion
/// [`ExecutionPlan::with_cooperative_yields`], it is wrapped in a
[`YieldStreamExec`] parent to
```
##########
datafusion/physical-optimizer/src/insert_yield_exec.rs:
##########
@@ -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
Review Comment:
I think you can make this a link too
```suggestion
//! The [`InsertYieldExec`] optimizer rule inspects the physical plan to
find all leaf
```
##########
datafusion/physical-optimizer/src/insert_yield_exec.rs:
##########
@@ -92,3 +96,23 @@ impl PhysicalOptimizerRule for InsertYieldExec {
true
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion_common::assert_contains;
+ use datafusion_common::config::ConfigOptions;
+ use datafusion_physical_plan::{displayable, test::scan_partitioned};
+
+ #[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();
Review Comment:
Maybe it is worth using `insta` snapshots here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
