Rework ensure_coop to base itself on evaluation and scheduling properties

Author: pepijnve

datafusion/core/tests/user_defined/insert_operation.rs

@@ -26,6 +26,7 @@ use datafusion::{
 use datafusion_catalog::{Session, TableProvider};
 use datafusion_expr::{dml::InsertOp, Expr, TableType};
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
+use datafusion_physical_plan::execution_plan::SchedulingType;
 use datafusion_physical_plan::{
     execution_plan::{Boundedness, EmissionType},
     DisplayAs, ExecutionPlan, PlanProperties,
@@ -132,7 +133,8 @@ impl TestInsertExec {
                 Partitioning::UnknownPartitioning(1),
                 EmissionType::Incremental,
                 Boundedness::Bounded,
-            ),
+            )
+            .with_scheduling_type(SchedulingType::Cooperative),
         }
     }
 }
@@ -179,10 +181,6 @@ impl ExecutionPlan for TestInsertExec {
     ) -> Result<datafusion_execution::SendableRecordBatchStream> {
         unimplemented!("TestInsertExec is a stub for testing.")
     }
-
-    fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        Some(self)
-    }
 }
 
 fn make_count_schema() -> SchemaRef {

datafusion/datasource/src/sink.rs

@@ -36,6 +36,7 @@ use datafusion_physical_plan::{
 };
 
 use async_trait::async_trait;
+use datafusion_physical_plan::execution_plan::{EvaluationType, SchedulingType};
 use futures::StreamExt;
 
 /// `DataSink` implements writing streams of [`RecordBatch`]es to
@@ -141,6 +142,8 @@ impl DataSinkExec {
             input.pipeline_behavior(),
             input.boundedness(),
         )
+        .with_scheduling_type(SchedulingType::Cooperative)
+        .with_evaluation_type(EvaluationType::Eager)
     }
 }
 
@@ -246,10 +249,6 @@ impl ExecutionPlan for DataSinkExec {
     fn metrics(&self) -> Option<MetricsSet> {
         self.sink.metrics()
     }
-
-    fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        Some(self)
-    }
 }
 
 /// Create a output record batch with a count

datafusion/datasource/src/source.rs

@@ -22,7 +22,9 @@ use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
-use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion_physical_plan::execution_plan::{
+    Boundedness, EmissionType, SchedulingType,
+};
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::{
@@ -262,10 +264,6 @@ impl ExecutionPlan for DataSourceExec {
             .map(|stream| make_cooperative(stream))
     }
 
-    fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        Some(self)
-    }
-
     fn metrics(&self) -> Option<MetricsSet> {
         Some(self.data_source.metrics().clone_inner())
     }
@@ -380,6 +378,7 @@ impl DataSourceExec {
             EmissionType::Incremental,
             Boundedness::Bounded,
         )
+        .with_scheduling_type(SchedulingType::Cooperative)
     }
 
     /// Downcast the `DataSourceExec`'s `data_source` to a specific file source

datafusion/physical-optimizer/src/ensure_coop.rs

@@ -29,12 +29,14 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
 use datafusion_common::Result;
 use datafusion_physical_plan::coop::CooperativeExec;
+use datafusion_physical_plan::execution_plan::{EvaluationType, SchedulingType};
 use datafusion_physical_plan::ExecutionPlan;
 
-/// `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 [`CooperativeExec`] parent.
+/// `EnsureCooperative` is a [`PhysicalOptimizerRule`] that inspects the physical plan for
+/// sub plans that do not participate in cooperative scheduling. The plan is subdivided into sub
+/// plans on eager evaluation boundaries. Leaf nodes and eager evaluation roots are checked
+/// to see if they participate in cooperative scheduling. Those that do no are wrapped in
+/// a [`CooperativeExec`] parent.
 pub struct EnsureCooperative {}
 
 impl EnsureCooperative {
@@ -65,20 +67,22 @@ impl PhysicalOptimizerRule for EnsureCooperative {
         plan: Arc<dyn ExecutionPlan>,
         _config: &ConfigOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        plan.transform_down(|plan| {
-            if !plan.children().is_empty() {
-                // Not a leaf, keep recursing down.
-                return Ok(Transformed::no(plan));
+        plan.transform_up(|plan| {
+            let is_leaf = plan.children().is_empty();
+            let is_exchange = plan.properties().evaluation_type == EvaluationType::Eager;
+            if (is_leaf || is_exchange)
+                && plan.properties().scheduling_type != SchedulingType::Cooperative
+            {
+                // Wrap non-cooperative leaves or eager evaluation roots in a cooperative exec to
+                // ensure the plans they participate in are properly cooperative.
+                Ok(Transformed::new(
+                    Arc::new(CooperativeExec::new(Arc::clone(&plan))),
+                    true,
+                    TreeNodeRecursion::Continue,
+                ))
+            } else {
+                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 `CooperativeExec`.
-                        Arc::new(CooperativeExec::new(plan))
-                    });
-            Ok(Transformed::new(new_plan, true, TreeNodeRecursion::Jump))
         })
         .map(|t| t.data)
     }
@@ -107,7 +111,7 @@ mod tests {
         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
+            CooperativeExec
               DataSourceExec: partitions=1, partition_sizes=[1]
             "###);
     }

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -27,7 +27,7 @@ use super::{
     DisplayAs, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream,
     Statistics,
 };
-use crate::execution_plan::CardinalityEffect;
+use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
 use crate::projection::{make_with_child, ProjectionExec};
 use crate::{DisplayFormatType, ExecutionPlan, Partitioning};
 
@@ -72,6 +72,16 @@ impl CoalescePartitionsExec {
 
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
     fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
+        let input_partitions = input.output_partitioning().partition_count();
+        let (drive, scheduling) = if input_partitions > 1 {
+            (EvaluationType::Eager, SchedulingType::Cooperative)
+        } else {
+            (
+                input.properties().evaluation_type,
+                input.properties().scheduling_type,
+            )
+        };
+
         // Coalescing partitions loses existing orderings:
         let mut eq_properties = input.equivalence_properties().clone();
         eq_properties.clear_orderings();
@@ -82,6 +92,8 @@ impl CoalescePartitionsExec {
             input.pipeline_behavior(),
             input.boundedness(),
         )
+        .with_evaluation_type(drive)
+        .with_scheduling_type(scheduling)
     }
 }
 

datafusion/physical-plan/src/coop.rs

@@ -30,6 +30,7 @@ use arrow_schema::Schema;
 use datafusion_common::{internal_err, Result, Statistics};
 use datafusion_execution::TaskContext;
 
+use crate::execution_plan::SchedulingType;
 use crate::stream::RecordBatchStreamAdapter;
 use futures::{FutureExt, Stream};
 use pin_project_lite::pin_project;
@@ -102,13 +103,19 @@ pub struct CooperativeExec {
     /// The child execution plan that this operator "wraps" to make it
     /// cooperate with the runtime.
     input: Arc<dyn ExecutionPlan>,
+    properties: PlanProperties,
 }
 
 impl CooperativeExec {
     /// Creates a new `CooperativeExec` operator that wraps the given child
     /// execution plan.
     pub fn new(input: Arc<dyn ExecutionPlan>) -> Self {
-        Self { input }
+        let properties = input
+            .properties()
+            .clone()
+            .with_scheduling_type(SchedulingType::Cooperative);
+
+        Self { input, properties }
     }
 
     /// Returns the child execution plan this operator "wraps" to make it
@@ -142,7 +149,7 @@ impl ExecutionPlan for CooperativeExec {
     }
 
     fn properties(&self) -> &PlanProperties {
-        self.input.properties()
+        &self.properties
     }
 
     fn maintains_input_order(&self) -> Vec<bool> {

datafusion/physical-plan/src/empty.rs

@@ -33,6 +33,7 @@ use datafusion_common::{internal_err, Result};
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::EquivalenceProperties;
 
+use crate::execution_plan::SchedulingType;
 use log::trace;
 
 /// Execution plan for empty relation with produce_one_row=false
@@ -81,6 +82,7 @@ impl EmptyExec {
             EmissionType::Incremental,
             Boundedness::Bounded,
         )
+        .with_scheduling_type(SchedulingType::Cooperative)
     }
 }
 
@@ -173,10 +175,6 @@ impl ExecutionPlan for EmptyExec {
             None,
         ))
     }
-
-    fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        Some(self)
-    }
 }
 
 #[cfg(test)]

datafusion/physical-plan/src/execution_plan.rs

@@ -41,8 +41,6 @@ use crate::coalesce_partitions::CoalescePartitionsExec;
 use crate::display::DisplayableExecutionPlan;
 use crate::metrics::MetricsSet;
 use crate::projection::ProjectionExec;
-use crate::repartition::RepartitionExec;
-use crate::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use crate::stream::RecordBatchStreamAdapter;
 
 use arrow::array::{Array, RecordBatch};
@@ -559,16 +557,6 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
             child_pushdown_result,
         ))
     }
-
-    /// Returns a version of this plan that cooperates with the runtime via
-    /// built‐in yielding. If such a version doesn't exist, returns `None`.
-    /// You do not need to do provide such a version of a custom operator,
-    /// but DataFusion will utilize it while optimizing the plan if it exists.
-    fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        // Conservative default implementation assumes that a leaf does not
-        // cooperate with yielding.
-        None
-    }
 }
 
 /// [`ExecutionPlan`] Invariant Level
@@ -743,6 +731,26 @@ pub enum EmissionType {
     Both,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SchedulingType {
+    /// The stream generated by [`execute`](ExecutionPlan::execute) does not participate in cooperative scheduling
+    Blocking,
+    /// The stream generated by [`execute`](ExecutionPlan::execute) actively participates in cooperative scheduling
+    /// by consuming task budget
+    Cooperative,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum EvaluationType {
+    /// The stream generated by [`execute`](ExecutionPlan::execute) only generates `RecordBatch`
+    /// instances when it is demanded by invoking `Stream::poll_next`.
+    Lazy,
+    /// The stream generated by [`execute`](ExecutionPlan::execute) eagerly generates `RecordBatch`
+    /// in one or more spawned Tokio tasks. Eager evaluation is only started the first time
+    /// `Stream::poll_next` is called.
+    Eager,
+}
+
 /// Utility to determine an operator's boundedness based on its children's boundedness.
 ///
 /// Assumes boundedness can be inferred from child operators:
@@ -831,6 +839,8 @@ pub struct PlanProperties {
     pub emission_type: EmissionType,
     /// See [ExecutionPlanProperties::boundedness]
     pub boundedness: Boundedness,
+    pub evaluation_type: EvaluationType,
+    pub scheduling_type: SchedulingType,
     /// See [ExecutionPlanProperties::output_ordering]
     output_ordering: Option<LexOrdering>,
 }
@@ -850,6 +860,8 @@ impl PlanProperties {
             partitioning,
             emission_type,
             boundedness,
+            evaluation_type: EvaluationType::Lazy,
+            scheduling_type: SchedulingType::Blocking,
             output_ordering,
         }
     }
@@ -881,6 +893,16 @@ impl PlanProperties {
         self
     }
 
+    pub fn with_scheduling_type(mut self, scheduling_type: SchedulingType) -> Self {
+        self.scheduling_type = scheduling_type;
+        self
+    }
+
+    pub fn with_evaluation_type(mut self, drive_type: EvaluationType) -> Self {
+        self.evaluation_type = drive_type;
+        self
+    }
+
     /// Overwrite constraints with its new value.
     pub fn with_constraints(mut self, constraints: Constraints) -> Self {
         self.eq_properties = self.eq_properties.with_constraints(constraints);
@@ -912,25 +934,7 @@ impl PlanProperties {
 ///     2. CoalescePartitionsExec for collapsing all of the partitions into one without ordering guarantee
 ///     3. SortPreservingMergeExec for collapsing all of the sorted partitions into one with ordering guarantee
 pub fn need_data_exchange(plan: Arc<dyn ExecutionPlan>) -> bool {
-    if let Some(repartition) = plan.as_any().downcast_ref::<RepartitionExec>() {
-        !matches!(
-            repartition.properties().output_partitioning(),
-            Partitioning::RoundRobinBatch(_)
-        )
-    } else if let Some(coalesce) = plan.as_any().downcast_ref::<CoalescePartitionsExec>()
-    {
-        coalesce.input().output_partitioning().partition_count() > 1
-    } else if let Some(sort_preserving_merge) =
-        plan.as_any().downcast_ref::<SortPreservingMergeExec>()
-    {
-        sort_preserving_merge
-            .input()
-            .output_partitioning()
-            .partition_count()
-            > 1
-    } else {
-        false
-    }
+    plan.properties().evaluation_type == EvaluationType::Lazy
 }
 
 /// Returns a copy of this plan if we change any child according to the pointer comparison.

datafusion/physical-plan/src/memory.rs

@@ -23,7 +23,7 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use crate::coop::cooperative;
-use crate::execution_plan::{Boundedness, EmissionType};
+use crate::execution_plan::{Boundedness, EmissionType, SchedulingType};
 use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use crate::{
     DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
@@ -163,7 +163,9 @@ impl LazyMemoryExec {
             Partitioning::RoundRobinBatch(generators.len()),
             EmissionType::Incremental,
             Boundedness::Bounded,
-        );
+        )
+        .with_scheduling_type(SchedulingType::Cooperative);
+
         Ok(Self {
             schema,
             batch_generators: generators,
@@ -273,10 +275,6 @@ impl ExecutionPlan for LazyMemoryExec {
         Ok(Box::pin(cooperative(stream)))
     }
 
-    fn with_cooperative_yields(self: Arc<Self>) -> Option<Arc<dyn ExecutionPlan>> {
-        Some(self)
-    }
-
     fn metrics(&self) -> Option<MetricsSet> {
         Some(self.metrics.clone_inner())
     }