Lazily set the producer head at execution time (#478)

This is one PR from the following stack of PRs:
- #477
- #463
- #464
- #478
<- you are here
- #479
- #486
- #432

---

Introduces the `ProducerHead` type:

```rust
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 },
}
```

Which is passed over the network while remotely executing tasks in order
to set the appropriate node at the head of a stage.

Today, this is a noop because the right head node in stages is ensured
statically at planning time, but in follow up PRs, network boundaries
can get swamped and reorganized arbitrarily.

One example that happens in AQE:

1. A JOIN is planned as a CollectLeft

```js
HashJoinExec: mode=CollectLeft
  CoalescePartitionsExec:
    [Stage 1] => NetworkBroadcastExec
      BroadcastExec
        DistributedLeafExec: unknown size
  DistributedLeafExec: unknown size
```

2. While collecting runtime statistics, it happens that `Stage 1` is
huge, and during AQE the JOINs are swapped

```js
HashJoinExec: mode=CollectLeft
  DistributedLeafExec: small size
  CoalescePartitionsExec:
    [Stage 1] => NetworkBroadcastExec
      BroadcastExec
        DistributedLeafExec: big size
```

3. The `Stage 1` is now on the probe side, so it needs to be rewritten
to a `NetworkShuffleExec`, otherwise duplicate data will be returned:

```js
HashJoinExec: mode=CollectLeft
  DistributedLeafExec: small size
  CoalescePartitionsExec:
    [Stage 1] => NetworkShuffleExec
      RepartitionExec // <- dynamically swapped at runtime based on the passed `ProducerHead`
        DistributedLeafExec: big size
```

Passing a `ProducerHead` at execution time unlocks two things:
1. dynamically set the fanout width accounting for a dynamically scaled
upper stage
2. dynamically set the correct operator `BroadcastExec` or
`RepartitionExec` based on the network boundary above, which might have
changed because of AQE

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()),