Add ProducerHead for ensuring that the appropriate node is placed at the head of stages

Author: gabotechs

src/distributed_planner/mod.rs

@@ -11,6 +11,7 @@ mod task_estimator;
 
 pub use distributed_config::DistributedConfig;
 pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt};
+pub(crate) use network_boundary::{ProducerHead, insert_producer_head};
 pub use session_state_builder_ext::SessionStateBuilderExt;
 pub(crate) use task_estimator::set_distributed_task_estimator;
 pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext};

src/distributed_planner/network_boundary.rs

@@ -1,6 +1,8 @@
-use crate::{NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage};
+use crate::{BroadcastExec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage};
 use datafusion::common::Result;
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_expr::Partitioning;
+use datafusion::physical_plan::repartition::RepartitionExec;
+use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use std::sync::Arc;
 
 /// This trait represents a node that introduces the necessity of a network boundary in the plan.
@@ -15,6 +17,22 @@ pub trait NetworkBoundary: ExecutionPlan {
 
     /// Returns the assigned input [Stage], if any.
     fn input_stage(&self) -> &Stage;
+
+    /// Defines what head node should the producer stage feeding this [NetworkBoundary]
+    /// implementation have. This information is used during planning an executing for ensuring
+    /// the head of a stage has the appropriate shape for consumption.
+    fn producer_head(&self, consumer_tasks: usize) -> ProducerHead;
+}
+
+/// Defines what shape should the head node of a stage have upon getting executed. Depending
+/// on the [NetworkBoundary] implementation, the stage below should have different head nodes.
+pub enum ProducerHead {
+    /// No specific head node is necessary.
+    None,
+    /// The head node should be a [BroadcastExec].
+    BroadcastExec { output_partitions: usize },
+    /// The head node should be a [RepartitionExec].
+    RepartitionExec { partitioning: Partitioning },
 }
 
 /// Extension trait for downcasting dynamic types to [NetworkBoundary].
@@ -41,38 +59,28 @@ impl NetworkBoundaryExt for dyn ExecutionPlan {
     }
 }
 
-/// Scales up the head node of the input stage of a network boundary. Different network boundaries
-/// have different needs for scaling up their input, like for example, scaling up a RepartitionExec
-/// during shuffles.
-pub(crate) fn network_boundary_scale_input(
+/// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This
+/// can be called both during planning and lazily at runtime.
+pub(crate) fn insert_producer_head(
     input: Arc<dyn ExecutionPlan>,
-    consumer_partitions: usize,
-    consumer_task_count: usize,
+    head: ProducerHead,
 ) -> Result<Arc<dyn ExecutionPlan>> {
-    let transformed = NetworkShuffleExec::scale_input(
-        Arc::clone(&input),
-        consumer_partitions,
-        consumer_task_count,
-    )?;
-    if transformed.transformed {
-        return Ok(transformed.data);
-    }
-    let transformed = NetworkBroadcastExec::scale_input(
-        Arc::clone(&input),
-        consumer_partitions,
-        consumer_task_count,
-    )?;
-    if transformed.transformed {
-        return Ok(transformed.data);
-    }
-    let transformed = NetworkCoalesceExec::scale_input(
-        Arc::clone(&input),
-        consumer_partitions,
-        consumer_task_count,
-    )?;
-    if transformed.transformed {
-        return Ok(transformed.data);
-    }
-
-    Ok(input)
+    let input = if let Some(r_exec) = input.downcast_ref::<RepartitionExec>() {
+        Arc::clone(r_exec.input())
+    } else if let Some(b_exec) = input.downcast_ref::<BroadcastExec>() {
+        Arc::clone(b_exec.input())
+    } else {
+        input
+    };
+    let plan = match head {
+        ProducerHead::None => input,
+        ProducerHead::BroadcastExec { output_partitions } => {
+            let partitions = input.output_partitioning().partition_count();
+            Arc::new(BroadcastExec::new(input, output_partitions / partitions))
+        }
+        ProducerHead::RepartitionExec { partitioning } => {
+            Arc::new(RepartitionExec::try_new(input, partitioning)?)
+        }
+    };
+    Ok(plan)
 }

src/distributed_planner/prepare_network_boundaries.rs

@@ -1,5 +1,5 @@
 use crate::common::TreeNodeExt;
-use crate::distributed_planner::network_boundary::network_boundary_scale_input;
+use crate::distributed_planner::network_boundary::insert_producer_head;
 use crate::stage::LocalStage;
 use crate::{NetworkBoundaryExt, Stage};
 use datafusion::common::Result;
@@ -35,11 +35,8 @@ pub(crate) fn prepare_network_boundaries(
 
         // 2) Scale up the head node of the input stage in order to account for the amount of partition
         //    and consumer count above it.
-        let plan = network_boundary_scale_input(
-            Arc::clone(&input_stage.plan),
-            nb.properties().partitioning.partition_count(),
-            task_count,
-        )?;
+        let plan =
+            insert_producer_head(Arc::clone(&input_stage.plan), nb.producer_head(task_count))?;
 
         // 3) Make sure the input stage can be uniquely identified with a stage index and query id.
         //    If there were already some `query_id` and `num` that's fine.

src/execution_plans/benchmarks/shuffle_bench.rs

@@ -223,7 +223,7 @@ impl ShuffleFixture {
             let shuffle = NetworkShuffleExec {
                 properties: Arc::new(PlanProperties::new(
                     EquivalenceProperties::new(Arc::clone(&self.schema)),
-                    Partitioning::UnknownPartitioning(self.bench.partitions),
+                    Partitioning::Hash(vec![Arc::new(Column::new("id", 0))], self.bench.partitions),
                     EmissionType::Incremental,
                     Boundedness::Bounded,
                 )),

src/execution_plans/benchmarks/transport_bench.rs

@@ -273,7 +273,7 @@ impl TransportFixture {
             let shuffle = NetworkShuffleExec {
                 properties: Arc::new(PlanProperties::new(
                     EquivalenceProperties::new(Arc::clone(&self.schema)),
-                    Partitioning::UnknownPartitioning(self.bench.partitions),
+                    Partitioning::RoundRobinBatch(self.bench.partitions),
                     EmissionType::Incremental,
                     Boundedness::Bounded,
                 )),

src/execution_plans/broadcast.rs

@@ -104,6 +104,10 @@ impl BroadcastExec {
     pub fn consumer_task_count(&self) -> usize {
         self.consumer_task_count
     }
+
+    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
+        &self.input
+    }
 }
 
 impl DisplayAs for BroadcastExec {

src/execution_plans/network_broadcast.rs

@@ -1,9 +1,8 @@
 use crate::common::require_one_child;
-use crate::distributed_planner::NetworkBoundary;
+use crate::distributed_planner::{NetworkBoundary, ProducerHead};
 use crate::stage::{LocalStage, Stage};
 use crate::worker::WorkerConnectionPool;
 use crate::{BroadcastExec, DistributedTaskContext};
-use datafusion::common::tree_node::Transformed;
 use datafusion::common::{Result, not_impl_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
@@ -123,21 +122,6 @@ pub struct NetworkBroadcastExec {
 }
 
 impl NetworkBroadcastExec {
-    pub(crate) fn scale_input(
-        plan: Arc<dyn ExecutionPlan>,
-        _consumer_partitions: usize,
-        consumer_task_count: usize,
-    ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
-        let Some(broadcast) = plan.downcast_ref::<BroadcastExec>() else {
-            return Ok(Transformed::no(plan));
-        };
-
-        Ok(Transformed::yes(Arc::new(BroadcastExec::new(
-            require_one_child(broadcast.children())?,
-            consumer_task_count,
-        ))))
-    }
-
     pub(crate) fn from_stage(input_stage: Stage, input_properties: Arc<PlanProperties>) -> Self {
         let input_partition_count = input_properties.partitioning.partition_count();
         let properties = Arc::new(
@@ -186,6 +170,13 @@ impl NetworkBoundary for NetworkBroadcastExec {
     fn input_stage(&self) -> &Stage {
         &self.input_stage
     }
+
+    fn producer_head(&self, consumer_task_count: usize) -> ProducerHead {
+        let partition_count = self.properties.output_partitioning().partition_count();
+        ProducerHead::BroadcastExec {
+            output_partitions: partition_count * consumer_task_count,
+        }
+    }
 }
 
 impl DisplayAs for NetworkBroadcastExec {
@@ -247,14 +238,16 @@ impl ExecutionPlan for NetworkBroadcastExec {
         };
 
         let task_context = DistributedTaskContext::from_ctx(&context);
-        let off = self.properties.partitioning.partition_count() * task_context.task_index;
+        let out_partitions = self.properties.partitioning.partition_count();
+        let off = out_partitions * task_context.task_index;
         let mut streams = Vec::with_capacity(self.input_stage.task_count());
 
         for input_task_index in 0..self.input_stage.task_count() {
             let worker_connection = self.worker_connections.get_or_init_worker_connection(
                 remote_stage,
                 off..(off + self.properties.partitioning.partition_count()),
                 input_task_index,
+                self.producer_head(task_context.task_count),
                 &context,
             )?;
 

src/execution_plans/network_coalesce.rs

@@ -1,10 +1,9 @@
 use crate::DistributedTaskContext;
 use crate::common::require_one_child;
-use crate::distributed_planner::NetworkBoundary;
+use crate::distributed_planner::{NetworkBoundary, ProducerHead};
 use crate::execution_plans::common::scale_partitioning_props;
 use crate::stage::{LocalStage, Stage};
 use crate::worker::WorkerConnectionPool;
-use datafusion::common::tree_node::Transformed;
 use datafusion::common::{exec_err, not_impl_err, plan_err};
 use datafusion::error::Result;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
@@ -81,16 +80,6 @@ pub struct NetworkCoalesceExec {
 }
 
 impl NetworkCoalesceExec {
-    /// Does nothing, but it's here for explicitly stating that this network boundary does not
-    /// need to mutate the input plan in other to account for more consumer tasks.
-    pub(crate) fn scale_input(
-        plan: Arc<dyn ExecutionPlan>,
-        _consumer_partitions: usize,
-        _consumer_task_count: usize,
-    ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
-        Ok(Transformed::no(plan))
-    }
-
     pub(crate) fn from_stage(
         input_stage: Stage,
         input_properties: Arc<PlanProperties>,
@@ -187,6 +176,10 @@ impl NetworkBoundary for NetworkCoalesceExec {
         self_clone.input_stage = input_stage;
         Ok(Arc::new(self_clone))
     }
+
+    fn producer_head(&self, _consumer_task_count: usize) -> ProducerHead {
+        ProducerHead::None
+    }
 }
 
 impl DisplayAs for NetworkCoalesceExec {
@@ -299,6 +292,7 @@ impl ExecutionPlan for NetworkCoalesceExec {
             remote_stage,
             0..partitions_per_task,
             target_task,
+            self.producer_head(task_context.task_count),
             &context,
         )?;
 

src/execution_plans/network_shuffle.rs

@@ -1,9 +1,9 @@
 use crate::common::require_one_child;
+use crate::distributed_planner::ProducerHead;
 use crate::execution_plans::common::scale_partitioning;
 use crate::stage::{LocalStage, Stage};
 use crate::worker::WorkerConnectionPool;
 use crate::{DistributedTaskContext, NetworkBoundary};
-use datafusion::common::tree_node::{Transformed, TreeNodeRecursion};
 use datafusion::common::{Result, not_impl_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
@@ -105,26 +105,6 @@ pub struct NetworkShuffleExec {
 }
 
 impl NetworkShuffleExec {
-    pub(crate) fn scale_input(
-        plan: Arc<dyn ExecutionPlan>,
-        consumer_partitions: usize,
-        consumer_task_count: usize,
-    ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
-        let Some(repartition_exec) = plan.downcast_ref::<RepartitionExec>() else {
-            return Ok(Transformed::no(plan));
-        };
-
-        let child = require_one_child(repartition_exec.children())?;
-        let partitioning = scale_partitioning(repartition_exec.partitioning(), |_| {
-            consumer_partitions * consumer_task_count
-        });
-
-        // Scale the input RepartitionExec to account for all the tasks to which it will
-        // need to fan data out.
-        let scaled = Arc::new(RepartitionExec::try_new(child, partitioning)?);
-        Ok(Transformed::new(scaled, true, TreeNodeRecursion::Stop))
-    }
-
     pub(crate) fn from_stage(input_stage: Stage, input_properties: Arc<PlanProperties>) -> Self {
         Self {
             properties: input_properties,
@@ -170,6 +150,14 @@ impl NetworkBoundary for NetworkShuffleExec {
         self_clone.input_stage = input_stage;
         Ok(Arc::new(self_clone))
     }
+
+    fn producer_head(&self, consumer_task_count: usize) -> ProducerHead {
+        ProducerHead::RepartitionExec {
+            partitioning: scale_partitioning(&self.properties.partitioning, |prev| {
+                prev * consumer_task_count
+            }),
+        }
+    }
 }
 
 impl DisplayAs for NetworkShuffleExec {
@@ -233,6 +221,7 @@ impl ExecutionPlan for NetworkShuffleExec {
                 remote_stage,
                 off..(off + self.properties.partitioning.partition_count()),
                 input_task_index,
+                self.producer_head(task_context.task_count),
                 &context,
             )?;
 

src/observability/service.rs

@@ -96,7 +96,7 @@ impl ObservabilityService for ObservabilityServiceImpl {
                 let total_partitions = task_data.total_partitions() as u64;
                 let remaining = task_data.num_partitions_remaining() as u64;
                 let completed_partitions = total_partitions.saturating_sub(remaining);
-                let output_rows = output_rows_from_plan(&task_data.plan);
+                let output_rows = output_rows_from_plan(&task_data.base_plan);
 
                 tasks.push(TaskProgress {
                     task_key: Some((*internal_key).clone()),