Remove 2-pass based distributed planner in favor of just 1 pass

Author: gabotechs

src/distributed_planner/distribute_plan.rs

@@ -1,169 +1,82 @@
-use crate::common::require_one_child;
+use crate::NetworkBoundaryExt;
+use crate::distributed_planner::inject_network_boundaries::inject_network_boundaries;
 use crate::distributed_planner::insert_broadcast::insert_broadcast_execs;
 use crate::distributed_planner::partial_reduce_below_network_shuffles::partial_reduce_below_network_shuffles;
-use crate::distributed_planner::plan_annotator::{
-    AnnotatedPlan, PlanOrNetworkBoundary, annotate_plan,
-};
-use crate::{
-    DistributedConfig, NetworkBoundaryExt, NetworkBroadcastExec, NetworkCoalesceExec,
-    NetworkShuffleExec, TaskEstimator,
-};
-use datafusion::common::DataFusionError;
+use crate::distributed_planner::prepare_network_boundaries::prepare_network_boundaries;
 use datafusion::common::tree_node::TreeNode;
 use datafusion::config::ConfigOptions;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
-use std::ops::AddAssign;
 use std::sync::Arc;
-use uuid::Uuid;
 
-/// Inspects the plan, places the appropriate network boundaries, and breaks it down into stages
-/// that can be executed in a distributed manner.
+/// Transforms a single-node physical plan into a distributed plan by injecting network
+/// boundaries between stages.
 ///
-/// It performs the following operations:
+/// The pipeline runs four passes in order:
 ///
-/// 1. It prepares the plan for distribution, adding some extra single-node nodes like
-///    [BroadcastExec] or [CoalescePartitionsExec] that will signal the following steps to
-///    introduce network boundaries in the appropriate places.
+/// 1. **Pre-distribution shaping.** A `CoalescePartitionsExec` is wrapped on top of the plan
+///    when it has more than one output partition (so [inject_network_boundaries] later sees a
+///    partition-collecting parent and injects a `NetworkCoalesceExec` above its child). Then
+///    [insert_broadcast_execs] adds `BroadcastExec` nodes on the build side of `CollectLeft`
+///    hash joins so those build sides can later be wrapped in `NetworkBroadcastExec`.
 ///
-/// 2. Annotate the plan with [annotate_plan]: adds some annotations to each node about how
-///    many distributed tasks should be used in the stage containing them, and whether they
-///    need a network boundary below or not.
-///    For more information about this step, read [annotate_plan] docs.
+/// 2. **Boundary injection.** [inject_network_boundaries] walks the plan, computes a task count
+///    for each node, and inserts `NetworkShuffleExec` / `NetworkBroadcastExec` /
+///    `NetworkCoalesceExec` above the nodes that delimit a stage (hash `RepartitionExec`s,
+///    build-side `BroadcastExec`s, and any node sitting under a `CoalescePartitionsExec` /
+///    `SortPreservingMergeExec`).
 ///
-/// 3. Based on the [AnnotatedPlan] returned by [annotate_plan], place all the appropriate
-///    network boundaries ([NetworkShuffleExec] and [NetworkCoalesceExec]) with the task count
-///    assignation that the annotations required. After this, the plan is already a distributed
-///    executable plan.
+/// 3. **Boundary preparation.** [prepare_network_boundaries] readies each injected boundary
+///    for execution: elides ones that aren't actually needed and finalises the survivors. If
+///    no boundary survives, this function returns `None`.
 ///
-/// This function returns None if the plan was left undistributed.
+/// 4. **Shuffle-volume optimisation.** [partial_reduce_below_network_shuffles] inserts partial
+///    aggregation nodes underneath hash shuffles where it can, so less data crosses the network.
 pub(super) async fn distribute_plan(
     original: Arc<dyn ExecutionPlan>,
     cfg: &ConfigOptions,
 ) -> datafusion::common::Result<Option<Arc<dyn ExecutionPlan>>> {
-    // Keep this function idempotent.
+    // The plan already contains network boundaries set by the user. Just ensure they have nice
+    // unique identifiers for each stage, and move forward with it.
     if original.exists(|plan| Ok(plan.is_network_boundary()))? {
-        return Ok(None);
+        // Ensure the stages in the plan have nice unique identifiers.
+        let plan = prepare_network_boundaries(original)?;
+        if !plan.exists(|plan| Ok(plan.is_network_boundary()))? {
+            return Ok(None);
+        }
+        return Ok(Some(plan));
     }
 
     let mut plan = Arc::clone(&original);
 
-    // Add a CoalescePartitionsExec on top of the plan if necessary. The plan annotator will see
-    // this and will place a NetworkCoalesceExec below it.
+    // If the plan has multiple output partitions, wrap it in a `CoalescePartitionsExec` so
+    // `inject_network_boundaries` will recognise the partition-collecting parent and inject a
+    // `NetworkCoalesceExec` above the original root.
     if plan.output_partitioning().partition_count() > 1 {
         plan = Arc::new(CoalescePartitionsExec::new(plan));
     }
 
-    // Insert BroadcastExec nodes in collect left joins so that the plan annotator can inject
-    // broadcast network boundaries above.
+    // Insert `BroadcastExec` on the build side of CollectLeft hash joins so that
+    // `inject_network_boundaries` can wrap them in `NetworkBroadcastExec` afterwards.
     plan = insert_broadcast_execs(plan, cfg)?;
 
-    // Annotate the plan with network boundary and task count information.
-    let annotated = annotate_plan(plan, cfg).await?;
+    // Compute per-node task counts and inject `Network*Exec` nodes at the stage boundaries.
+    plan = inject_network_boundaries(plan, cfg).await?;
 
-    // Based on the annotations, place the actual network boundaries with the appropriate dimensions.
-    let mut stage_id = 1;
-    let plan = _distribute_plan(annotated, cfg, Uuid::new_v4(), &mut stage_id)?;
-    if stage_id == 1 {
+    // Run some preparations on the network boundaries, like optimizing them out if there is one
+    // task at both sides of the boundary or assigning it proper unique identifiers.
+    plan = prepare_network_boundaries(plan)?;
+    if !plan.exists(|plan| Ok(plan.is_network_boundary()))? {
         return Ok(None);
     }
 
-    // Insert PartialReduce aggregation nodes above hash repartitions to reduce shuffle data volume.
+    // Push partial aggregation below hash shuffles to shrink the volume of data sent over
+    // the network.
     let plan = partial_reduce_below_network_shuffles(plan, cfg)?;
 
     Ok(Some(plan))
 }
 
-/// Takes an [AnnotatedPlan] and returns a modified [ExecutionPlan] with all the network boundaries
-/// appropriately placed. This step performs the following modifications to the original
-/// [ExecutionPlan]:
-/// - The leaf nodes are scaled up in parallelism based on the number of distributed tasks in
-///   which they are going to run. This is configurable by the user via the [TaskEstimator] trait.
-/// - The appropriate network boundaries are placed in the plan depending on how it was annotated,
-///   so new nodes like [NetworkBroadcastExec], [NetworkCoalesceExec] and [NetworkShuffleExec] will be present.
-fn _distribute_plan(
-    annotated_plan: AnnotatedPlan,
-    cfg: &ConfigOptions,
-    query_id: Uuid,
-    stage_id: &mut usize,
-) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
-    let d_cfg = DistributedConfig::from_config_options(cfg)?;
-    let children = annotated_plan.children;
-    let task_count = annotated_plan.task_count.as_usize();
-    let max_child_task_count = children.iter().map(|v| v.task_count.as_usize()).max();
-    let new_children = children
-        .into_iter()
-        .map(|child| _distribute_plan(child, cfg, query_id, stage_id))
-        .collect::<Result<Vec<_>, _>>()?;
-    match annotated_plan.plan_or_nb {
-        // This is a leaf node. It needs to be scaled up in order to account for it running in
-        // multiple tasks.
-        PlanOrNetworkBoundary::Plan(plan) if plan.children().is_empty() => {
-            let scaled_up = d_cfg.__private_task_estimator.scale_up_leaf_node(
-                &plan,
-                annotated_plan.task_count.as_usize(),
-                cfg,
-            );
-            Ok(scaled_up.unwrap_or(plan))
-        }
-        // This is a normal intermediate plan, just pass it through with the mapped children.
-        PlanOrNetworkBoundary::Plan(plan) => plan.with_new_children(new_children),
-        // This is a shuffle, so inject a NetworkShuffleExec here in the plan.
-        PlanOrNetworkBoundary::Shuffle => {
-            // It would need a network boundary, but on both sides of the boundary there is just 1 task,
-            // so we are fine with not introducing any network boundary.
-            if task_count == 1 && max_child_task_count == Some(1) {
-                return require_one_child(new_children);
-            }
-            let node = Arc::new(NetworkShuffleExec::try_new(
-                require_one_child(new_children)?,
-                query_id,
-                *stage_id,
-                task_count,
-                max_child_task_count.unwrap_or(1),
-            )?);
-            stage_id.add_assign(1);
-            Ok(node)
-        }
-        // DataFusion is trying to coalesce multiple partitions into one, so we should do the
-        // same with tasks.
-        PlanOrNetworkBoundary::Coalesce => {
-            // It would need a network boundary, but on both sides of the boundary there is just 1 task,
-            // so we are fine with not introducing any network boundary.
-            if task_count == 1 && max_child_task_count == Some(1) {
-                return require_one_child(new_children);
-            }
-            let node = Arc::new(NetworkCoalesceExec::try_new(
-                require_one_child(new_children)?,
-                query_id,
-                *stage_id,
-                task_count,
-                max_child_task_count.unwrap_or(1),
-            )?);
-            stage_id.add_assign(1);
-            Ok(node)
-        }
-        // This is a CollectLeft HashJoinExec with the build side marked as being broadcast. we
-        // need to insert a NetworkBroadcastExec and scale up the BroadcastExec consumer_tasks.
-        PlanOrNetworkBoundary::Broadcast => {
-            // It would need a network boundary, but on both sides of the boundary there is just 1 task,
-            // so we are fine with not introducing any network boundary.
-            if task_count == 1 && max_child_task_count == Some(1) {
-                return require_one_child(new_children);
-            }
-            let node = Arc::new(NetworkBroadcastExec::try_new(
-                require_one_child(new_children)?,
-                query_id,
-                *stage_id,
-                task_count,
-                max_child_task_count.unwrap_or(1),
-            )?);
-            stage_id.add_assign(1);
-            Ok(node)
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver;