Use Tokio's task budget consistently

Author: pepijnve

datafusion/common/src/config.rs

@@ -722,15 +722,6 @@ config_namespace! {
         /// then the output will be coerced to a non-view.
         /// Coerces `Utf8View` to `LargeUtf8`, and `BinaryView` to `LargeBinary`.
         pub expand_views_at_output: bool, default = false
-
-        /// When DataFusion detects that a plan might not be promply cancellable
-        /// due to the presence of tight-looping operators, it will attempt to
-        /// mitigate this by inserting explicit yielding (in as few places as
-        /// possible to avoid performance degradation). This value represents the
-        /// yielding period (in batches) at such explicit yielding points. The
-        /// default value is 64. If set to 0, no DataFusion will not perform
-        /// any explicit yielding.
-        pub yield_period: usize, default = 64
     }
 }
 

datafusion/core/tests/execution/infinite_cancel.rs

@@ -38,7 +38,7 @@ use datafusion_expr_common::operator::Operator::Gt;
 use datafusion_physical_expr::expressions::{col, BinaryExpr, Column, Literal};
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
-use datafusion_physical_optimizer::insert_yield_exec::InsertYieldExec;
+use datafusion_physical_optimizer::ensure_coop::EnsureCooperative;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion_physical_plan::filter::FilterExec;
@@ -146,7 +146,7 @@ async fn test_infinite_agg_cancel(
 
     // 3) optimize the plan with InsertYieldExec to auto-insert Yield
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(aggr, &config)?;
+    let optimized = EnsureCooperative::new().optimize(aggr, &config)?;
 
     // 4) get the stream
     let mut stream = physical_plan::execute_stream(optimized, session_ctx.task_ctx())?;
@@ -199,7 +199,7 @@ async fn test_infinite_sort_cancel(
 
     // 4) optimize the plan with InsertYieldExec to auto-insert Yield
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(sort_exec, &config)?;
+    let optimized = EnsureCooperative::new().optimize(sort_exec, &config)?;
 
     // 5) get the stream
     let mut stream = physical_plan::execute_stream(optimized, session_ctx.task_ctx())?;
@@ -284,7 +284,7 @@ async fn test_infinite_interleave_cancel(
     // 6) Apply InsertYieldExec to insert YieldStreamExec under every leaf.
     //    Each InfiniteExec → FilterExec → CoalesceBatchesExec chain will yield periodically.
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(coalesced_top, &config)?;
+    let optimized = EnsureCooperative::new().optimize(coalesced_top, &config)?;
 
     // 7) Execute the optimized plan with a 1-second timeout.
     //    Because the top-level FilterExec always discards rows and the inputs are infinite,
@@ -385,7 +385,7 @@ async fn test_infinite_interleave_agg_cancel(
     //    That way, each InfiniteExec (through the FilterExec/CoalesceBatchesExec/RepartitionExec chain)
     //    yields to the runtime periodically instead of spinning CPU.
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(aggr, &config)?;
+    let optimized = EnsureCooperative::new().optimize(aggr, &config)?;
 
     // 6) Execute the stream. Because AggregateExec(mode=Single) only emits a final batch
     //    after all inputs finish—and those inputs are infinite—we expect no output
@@ -472,7 +472,7 @@ async fn test_infinite_join_cancel(
 
     // 3) Wrap yields under each infinite leaf
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(join, &config)?;
+    let optimized = EnsureCooperative::new().optimize(join, &config)?;
 
     // 4) Execute + 1 sec timeout
     let mut stream = physical_plan::execute_stream(optimized, session_ctx.task_ctx())?;
@@ -590,7 +590,7 @@ async fn test_infinite_join_agg_cancel(
 
     // 5) Wrap yields under each infinite leaf
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(aggr, &config)?;
+    let optimized = EnsureCooperative::new().optimize(aggr, &config)?;
 
     // 6) Execute + 1 sec timeout
     let mut stream = physical_plan::execute_stream(optimized, session_ctx.task_ctx())?;
@@ -644,7 +644,7 @@ async fn test_filter_reject_all_batches_cancel(
 
     // 3) InsertYieldExec to insert YieldExec—so that the InfiniteExec yields control between batches
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(coalesced, &config)?;
+    let optimized = EnsureCooperative::new().optimize(coalesced, &config)?;
 
     // 4) Execute with a 1-second timeout. Because Filter discards all 8192 rows each time
     //    without ever producing output, no batch will arrive within 1 second. And since
@@ -723,7 +723,7 @@ async fn test_infinite_hash_join_without_repartition_and_no_agg(
     //    because there is no aggregation so no wrapper is inserted. Here we simply do
     //    not call InsertYieldExec, ensuring the plan has neither aggregation nor repartition.
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(join, &config)?;
+    let optimized = EnsureCooperative::new().optimize(join, &config)?;
 
     // 4) Execute with a 1 second timeout
     let mut stream = physical_plan::execute_stream(optimized, session_ctx.task_ctx())?;
@@ -801,7 +801,7 @@ async fn test_infinite_sort_merge_join_without_repartition_and_no_agg(
 
     // 3) Do not apply InsertYieldExec (no aggregation, no repartition → no built-in yields).
     let config = ConfigOptions::new();
-    let optimized = InsertYieldExec::new().optimize(join, &config)?;
+    let optimized = EnsureCooperative::new().optimize(join, &config)?;
 
     // 4) Execute with a 1-second timeout. Because both sides are infinite and never match,
     //    the SortMergeJoin will never produce output within 1s.

datafusion/datasource/src/source.rs

@@ -35,10 +35,10 @@ use datafusion_common::{Constraints, Result, Statistics};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
+use datafusion_physical_plan::coop::cooperative_wrapper;
 use datafusion_physical_plan::filter_pushdown::{
     ChildPushdownResult, FilterPushdownPropagation,
 };
-use datafusion_physical_plan::yield_stream::wrap_yield_stream;
 
 /// A source of data, typically a list of files or memory
 ///
@@ -186,8 +186,6 @@ pub struct DataSourceExec {
     data_source: Arc<dyn DataSource>,
     /// Cached plan properties such as sort order
     cache: PlanProperties,
-    /// Indicates whether to enable cooperative yielding mode.
-    cooperative: bool,
 }
 
 impl DisplayAs for DataSourceExec {
@@ -261,11 +259,11 @@ impl ExecutionPlan for DataSourceExec {
     ) -> Result<SendableRecordBatchStream> {
         self.data_source
             .open(partition, Arc::clone(&context))
-            .map(|stream| wrap_yield_stream(stream, &context, self.cooperative))
+            .map(|stream| cooperative_wrapper(stream))
     }
 
     fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        self.cooperative.then_some(self)
+        Some(self)
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
@@ -298,11 +296,7 @@ impl ExecutionPlan for DataSourceExec {
         let data_source = self.data_source.with_fetch(limit)?;
         let cache = self.cache.clone();
 
-        Some(Arc::new(Self {
-            data_source,
-            cache,
-            cooperative: self.cooperative,
-        }))
+        Some(Arc::new(Self { data_source, cache }))
     }
 
     fn fetch(&self) -> Option<usize> {
@@ -353,11 +347,7 @@ impl DataSourceExec {
     // Default constructor for `DataSourceExec`, setting the `cooperative` flag to `true`.
     pub fn new(data_source: Arc<dyn DataSource>) -> Self {
         let cache = Self::compute_properties(Arc::clone(&data_source));
-        Self {
-            data_source,
-            cache,
-            cooperative: true,
-        }
+        Self { data_source, cache }
     }
 
     /// Return the source object
@@ -383,12 +373,6 @@ impl DataSourceExec {
         self
     }
 
-    /// Assign yielding mode
-    pub fn with_cooperative(mut self, cooperative: bool) -> Self {
-        self.cooperative = cooperative;
-        self
-    }
-
     fn compute_properties(data_source: Arc<dyn DataSource>) -> PlanProperties {
         PlanProperties::new(
             data_source.eq_properties(),

datafusion/physical-optimizer/src/insert_yield_exec.rs renamed to datafusion/physical-optimizer/src/ensure_coop.rs

@@ -15,12 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! 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.
+//! The [`EnsureCooperative`] optimizer rule inspects the physical plan to find all
+//! portions of the plan that will not yield cooperatively.
+//! It will insert `CooperativeExec` nodes where appropriate to ensure execution plans
+//! always yield cooperatively.
 
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
@@ -30,65 +28,59 @@ use crate::PhysicalOptimizerRule;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
 use datafusion_common::Result;
-use datafusion_physical_plan::yield_stream::YieldStreamExec;
+use datafusion_physical_plan::coop::CooperativeExec;
 use datafusion_physical_plan::ExecutionPlan;
 
-/// `InsertYieldExec` is a [`PhysicalOptimizerRule`] that finds every leaf node in
+/// `EnsureCooperative` is a [`PhysicalOptimizerRule`] that finds every leaf node in
 /// 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 {}
+/// [`ExecutionPlan::with_cooperative_yields`], it is wrapped in a [`CooperativeExec`] parent.
+pub struct EnsureCooperative {}
 
-impl InsertYieldExec {
+impl EnsureCooperative {
     pub fn new() -> Self {
         Self {}
     }
 }
 
-impl Default for InsertYieldExec {
+impl Default for EnsureCooperative {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl Debug for InsertYieldExec {
+impl Debug for EnsureCooperative {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("InsertYieldExec").finish()
     }
 }
 
-impl PhysicalOptimizerRule for InsertYieldExec {
+impl PhysicalOptimizerRule for EnsureCooperative {
     fn name(&self) -> &str {
         "insert_yield_exec"
     }
 
     fn optimize(
         &self,
         plan: Arc<dyn ExecutionPlan>,
-        config: &ConfigOptions,
+        _config: &ConfigOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        // Only activate if user has configured a non-zero yield frequency.
-        let yield_period = config.optimizer.yield_period;
-        if yield_period != 0 {
-            plan.transform_down(|plan| {
-                if !plan.children().is_empty() {
-                    // 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)
+        plan.transform_down(|plan| {
+            if !plan.children().is_empty() {
+                // 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(|| {
-                        // Only if no built-in variant exists, insert a `YieldStreamExec`.
-                        Arc::new(YieldStreamExec::new(plan, yield_period))
+                        // Only if no built-in variant exists, insert a `CooperativeExec`.
+                        Arc::new(CooperativeExec::new(plan))
                     });
-                Ok(Transformed::new(new_plan, true, TreeNodeRecursion::Jump))
-            })
-            .map(|t| t.data)
-        } else {
-            Ok(plan)
-        }
+            Ok(Transformed::new(new_plan, true, TreeNodeRecursion::Jump))
+        })
+        .map(|t| t.data)
     }
 
     fn schema_check(&self) -> bool {
@@ -105,10 +97,10 @@ mod tests {
     use insta::assert_snapshot;
 
     #[tokio::test]
-    async fn test_yield_stream_exec_for_custom_exec() {
+    async fn test_cooperative_exec_for_custom_exec() {
         let test_custom_exec = scan_partitioned(1);
         let config = ConfigOptions::new();
-        let optimized = InsertYieldExec::new()
+        let optimized = EnsureCooperative::new()
             .optimize(test_custom_exec, &config)
             .unwrap();
 

datafusion/physical-optimizer/src/lib.rs

@@ -29,8 +29,8 @@ pub mod coalesce_batches;
 pub mod combine_partial_final_agg;
 pub mod enforce_distribution;
 pub mod enforce_sorting;
+pub mod ensure_coop;
 pub mod filter_pushdown;
-pub mod insert_yield_exec;
 pub mod join_selection;
 pub mod limit_pushdown;
 pub mod limited_distinct_aggregation;

datafusion/physical-optimizer/src/optimizer.rs

@@ -25,8 +25,8 @@ use crate::coalesce_batches::CoalesceBatches;
 use crate::combine_partial_final_agg::CombinePartialFinalAggregate;
 use crate::enforce_distribution::EnforceDistribution;
 use crate::enforce_sorting::EnforceSorting;
+use crate::ensure_coop::EnsureCooperative;
 use crate::filter_pushdown::FilterPushdown;
-use crate::insert_yield_exec::InsertYieldExec;
 use crate::join_selection::JoinSelection;
 use crate::limit_pushdown::LimitPushdown;
 use crate::limited_distinct_aggregation::LimitedDistinctAggregation;
@@ -138,7 +138,7 @@ impl PhysicalOptimizer {
             // are not present, the load of executors such as join or union will be
             // reduced by narrowing their input tables.
             Arc::new(ProjectionPushdown::new()),
-            Arc::new(InsertYieldExec::new()),
+            Arc::new(EnsureCooperative::new()),
             // The SanityCheckPlan rule checks whether the order and
             // distribution requirements of each node in the plan
             // is satisfied. It will also reject non-runnable query