datafusion_distributed/distributed_planner/
network_boundary.rs

1use crate::execution_plans::SamplerExec;
2use crate::{
3    BroadcastExec, MaybeEncoded, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec,
4    Stage,
5};
6use datafusion::arrow::datatypes::SchemaRef;
7use datafusion::common::Result;
8use datafusion::execution::TaskContext;
9use datafusion::physical_expr::Partitioning;
10use datafusion::physical_plan::repartition::RepartitionExec;
11use datafusion::physical_plan::{
12    ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions,
13};
14use std::sync::Arc;
15
16/// This trait represents a node that introduces the necessity of a network boundary in the plan.
17/// The distributed planner, upon stepping into one of these, will break the plan and build a stage
18/// out of it.
19pub trait NetworkBoundary: ExecutionPlan {
20    /// Called when a [Stage] is correctly formed. The [NetworkBoundary] can use this
21    /// information to perform any internal transformations necessary for distributed execution.
22    ///
23    /// Typically, [NetworkBoundary]s will use this call for transitioning from "Pending" to "ready".
24    fn with_input_stage(&self, input_stage: Stage) -> Result<Arc<dyn NetworkBoundary>>;
25
26    /// Returns the assigned input [Stage], if any.
27    fn input_stage(&self) -> &Stage;
28
29    /// Defines what head node should the producer stage feeding this [NetworkBoundary]
30    /// implementation have. This information is used during planning an executing for ensuring
31    /// the head of a stage has the appropriate shape for consumption. Returns an error when that
32    /// shape cannot be constructed from the boundary's partitioning.
33    fn producer_head(&self, consumer_tasks: usize) -> Result<ProducerHead>;
34}
35
36/// Defines what shape should the head node of a stage have upon getting executed. Depending
37/// on the [NetworkBoundary] implementation, the stage below should have different head nodes.
38#[derive(Clone)]
39pub enum ProducerHead {
40    /// No specific head node is necessary.
41    None,
42    /// The head node should be a [BroadcastExec].
43    BroadcastExec { output_partitions: usize },
44    /// The head node should be a [RepartitionExec].
45    RepartitionExec {
46        partitioning: MaybeEncoded<Partitioning>,
47    },
48}
49
50/// Extension trait for downcasting dynamic types to [NetworkBoundary].
51pub trait NetworkBoundaryExt {
52    /// Downcasts self to a [NetworkBoundary] if possible.
53    fn as_network_boundary(&self) -> Option<&dyn NetworkBoundary>;
54    /// Returns whether self is a [NetworkBoundary] or not.
55    fn is_network_boundary(&self) -> bool {
56        self.as_network_boundary().is_some()
57    }
58}
59
60impl NetworkBoundaryExt for dyn ExecutionPlan {
61    fn as_network_boundary(&self) -> Option<&dyn NetworkBoundary> {
62        if let Some(node) = self.downcast_ref::<NetworkShuffleExec>() {
63            Some(node)
64        } else if let Some(node) = self.downcast_ref::<NetworkCoalesceExec>() {
65            Some(node)
66        } else if let Some(node) = self.downcast_ref::<NetworkBroadcastExec>() {
67            Some(node)
68        } else {
69            None
70        }
71    }
72}
73
74impl ProducerHead {
75    pub(crate) fn ensure_decoded(self, schema: SchemaRef, ctx: &TaskContext) -> Result<Self> {
76        Ok(match self {
77            Self::RepartitionExec { partitioning } => Self::RepartitionExec {
78                partitioning: MaybeEncoded::Decoded(partitioning.decode(schema, ctx)?),
79            },
80            v => v,
81        })
82    }
83
84    /// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This
85    /// can be called both during planning and lazily at runtime.
86    pub(crate) fn insert(self, input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
87        let input = if let Some(r_exec) = input.downcast_ref::<RepartitionExec>() {
88            Arc::clone(r_exec.input())
89        } else if let Some(b_exec) = input.downcast_ref::<BroadcastExec>() {
90            Arc::clone(b_exec.input())
91        } else {
92            input
93        };
94        let plan = match self {
95            ProducerHead::None => input,
96            ProducerHead::BroadcastExec { output_partitions } => {
97                let partitions = input.output_partitioning().partition_count();
98                Arc::new(BroadcastExec::new(input, output_partitions / partitions))
99            }
100            ProducerHead::RepartitionExec { partitioning } => Arc::new(RepartitionExec::try_new(
101                input,
102                partitioning.try_decoded()?,
103            )?),
104        };
105        Ok(plan)
106    }
107
108    /// Injects a [SamplerExec] right below a [RepartitionExec] or [BroadcastExec].
109    pub(crate) fn insert_sampler(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
110        if let Some(r_exec) = input.downcast_ref::<RepartitionExec>() {
111            let child = Arc::clone(r_exec.input());
112            input.replace_children(
113                vec![Arc::new(SamplerExec::new(child))],
114                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
115            )
116        } else if let Some(b_exec) = input.downcast_ref::<BroadcastExec>() {
117            let child = Arc::clone(b_exec.input());
118            input.replace_children(
119                vec![Arc::new(SamplerExec::new(child))],
120                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
121            )
122        } else {
123            Ok(input)
124        }
125    }
126}
datafusion_distributed/distributed_planner/network_boundary.rs

datafusion_distributed/distributed_planner/
network_boundary.rs
