datafusion_distributed/execution_plans/network_broadcast.rs
1use crate::common::require_one_child;
2use crate::distributed_planner::{NetworkBoundary, ProducerHead};
3use crate::stage::{LocalStage, Stage};
4use crate::worker::WorkerConnectionPool;
5use crate::{BroadcastExec, DistributedTaskContext};
6use datafusion::common::tree_node::TreeNodeRecursion;
7use datafusion::common::{Result, not_impl_err, plan_err};
8use datafusion::error::DataFusionError;
9use datafusion::execution::{SendableRecordBatchStream, TaskContext};
10use datafusion::physical_expr::PhysicalExpr;
11use datafusion::physical_expr_common::metrics::MetricsSet;
12use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
13use datafusion::physical_plan::{
14 DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics,
15 StatisticsArgs,
16};
17use std::fmt::Formatter;
18use std::sync::Arc;
19use uuid::Uuid;
20
21/// Network boundary for broadcasting data to all consumer tasks.
22///
23/// This operator works with [BroadcastExec] which scales up partitions so each
24/// consumer task fetches a unique set of partition numbers. Each partition request
25/// is sent to all stage tasks because each task's leaf node is specialized to serve
26/// a different slice of the data for the same logical partition number.
27///
28/// Here are some examples of how [NetworkBroadcastExec] distributes data:
29///
30/// # 1 to many
31///
32/// ```text
33/// ┌────────────────────────┐ ┌────────────────────────┐ ■
34/// │ NetworkBroadcastExec │ │ NetworkBroadcastExec │ │
35/// │ (task 1) │ ... │ (task M) │ │
36/// │ │ │ │ Stage N
37/// │ Populates Caches │ │ Populates Caches │ │
38/// └────────┬─┬┬─┬┬─┬───────┘ └────────┬─┬┬─┬┬─┬───────┘ │
39/// │0││1││2│ │0││1││2│ │
40/// └▲┘└▲┘└▲┘ └▲┘└▲┘└▲┘ ■
41/// │ │ │ │ │ │
42/// │ │ │ │ │ │
43/// │ │ │ │ │ │
44/// │ │ └─────────────┐ ┌──────────────────┘ │ │
45/// │ └─────────────┐ │ │ ┌───────────────┘ │
46/// └─────────────┐ │ │ │ │ ┌─────────────┘
47/// │ │ │ │ │ │
48/// ┌┴┐┌┴┐┌┴┐ ... ┌───┴┐┌───┴┐┌──┴─┐
49/// │1││2││3│ │NM-3││NM-2││NM-1│ ■
50/// ┌┴─┴┴─┴┴─┴─────┴────┴┴────┴┴────┴─┐ │
51/// │ BroadcastExec │ │
52/// │ ┌───────────────┐ │ Stage N-1
53/// │ │ Batch Cache │ │ │
54/// │ │ ┌─┐ ┌─┐ ┌─┐ │ │ │
55/// │ │ │0│ │1│ │2│ │ │ │
56/// │ │ └─┘ └─┘ └─┘ │ │ │
57/// │ └───────────────┘ │ │
58/// └───────────┬─┬─┬─┬─┬─┬───────────┘ │
59/// │0│ │1│ │2│ │
60/// └▲┘ └▲┘ └▲┘ ■
61/// │ │ │
62/// │ │ │
63/// │ │ │
64/// ┌┴┐ ┌┴┐ ┌┴┐ ■
65/// │0│ │1│ │2│ │
66/// ┌──────┴─┴─┴─┴─┴─┴──────┐ Stage N-2
67/// │Arc<dyn ExecutionPlan> │ │
68/// │ (task 1) │ │
69/// └───────────────────────┘ ■
70/// ```
71///
72/// # Many to many
73///
74/// ```text
75/// ┌────────────────────────┐ ┌────────────────────────┐ ■
76/// │ NetworkBroadcastExec │ │ NetworkBroadcastExec │ │
77/// │ (task 1) │ │ (task M) │ │
78/// │ │ ... │ │ Stage N
79/// │ Populates Caches │ │ Cache Hits │ │
80/// └────────┬─┬┬─┬┬─┬───────┘ └────────┬─┬┬─┬┬─┬───────┘ │
81/// │0││1││2│ │0││1││2│ │
82/// └▲┘└▲┘└▲┘ └▲┘└▲┘└▲┘ ■
83/// │ │ │ │ │ │
84/// ┌──────────┴──┼──┼────────────────────────────────┐ │ │ │
85/// │ ┌──────────┴──┼────────────────────────────────┼──┐ │ │ │
86/// │ │ ┌──────────┴────────────────────────────────┼──┼──┐ │ │ │
87/// │ │ │ │ │ │ │ │ │
88/// │ │ │ ┌─────────────────────────────────┼──┼──┼────┴──┼─┐│
89/// │ │ │ │ ┌───────────────────────────┼──┼──┼───────┴─┼┼─────┐
90/// │ │ │ │ │ ┌─────────────────────┼──┼──┼─────────┼┴─────┼────┐
91/// │ │ │ │ │ │ │ │ │ │ │ │
92/// ┌┴┐┌┴┐┌┴┐ ... ┌──┴─┐┌──┴─┐┌──┴─┐ ┌┴┐┌┴┐┌┴┐ ... ┌──┴─┐┌───┴┐┌──┴─┐ ■
93/// │0││1││2│ │3M-3││3M-2││3M-1│ │0││1││2│ │3M-3││3M-2││3M-1│ │
94/// ┌┴─┴┴─┴┴─┴─────┴────┴┴────┴┴────┴┐ ┌┴─┴┴─┴┴─┴─────┴────┴┴────┴┴────┴┐ │
95/// │ BroadcastExec │ │ BroadcastExec │ │
96/// │ ┌───────────────┐ │ │ ┌───────────────┐ │ │
97/// │ │ Batch Cache │ │ │ │ Batch Cache │ │ │
98/// │ │ ┌─┐ ┌─┐ ┌─┐ │ │ ... │ │ ┌─┐ ┌─┐ ┌─┐ │ │ Stage N-1
99/// │ │ │0│ │1│ │2│ │ │ │ │ │0│ │1│ │2│ │ │ │
100/// │ │ └─┘ └─┘ └─┘ │ │ │ │ └─┘ └─┘ └─┘ │ │ │
101/// │ └───────────────┘ │ │ └───────────────┘ │ │
102/// └───────────┬─┬─┬─┬─┬─┬──────────┘ └───────────┬─┬─┬─┬─┬─┬──────────┘ │
103/// │0│ │1│ │2│ │0│ │1│ │2│ │
104/// └▲┘ └▲┘ └▲┘ └▲┘ └▲┘ └▲┘ ■
105/// │ │ │ │ │ │
106/// │ │ │ │ │ │
107/// │ │ │ │ │ │
108/// ┌┴┐ ┌┴┐ ┌┴┐ ┌┴┐ ┌┴┐ ┌┴┐ ■
109/// │0│ │1│ │2│ │0│ │1│ │2│ │
110/// ┌──────┴─┴─┴─┴─┴─┴──────┐ ┌──────┴─┴─┴─┴─┴─┴──────┐ Stage N-2
111/// │Arc<dyn ExecutionPlan> │ ... │Arc<dyn ExecutionPlan> │ │
112/// │ (task 1) │ │ (task N) │ │
113/// └───────────────────────┘ └───────────────────────┘ ■
114/// ```
115///
116/// Notice in this diagram that each [NetworkBroadcastExec] sends a request to fetch data from each
117/// [BroadcastExec] in the stage below per partition. This is because each [BroadcastExec] has its
118/// own cache which contains partial results for the partition. It is the [NetworkBroadcastExec]'s
119/// job to merge these partial partitions to then broadcast complete data to the consumers.
120#[derive(Debug, Clone)]
121pub struct NetworkBroadcastExec {
122 pub(crate) properties: Arc<PlanProperties>,
123 pub(crate) input_stage: Stage,
124 pub(crate) worker_connections: WorkerConnectionPool,
125}
126
127impl NetworkBroadcastExec {
128 pub(crate) fn from_stage(input_stage: Stage, input_properties: Arc<PlanProperties>) -> Self {
129 let input_partition_count = input_properties.partitioning.partition_count();
130 let properties = Arc::new(
131 PlanProperties::clone(&input_properties)
132 .with_partitioning(Partitioning::UnknownPartitioning(input_partition_count)),
133 );
134
135 Self {
136 properties,
137 worker_connections: WorkerConnectionPool::new(input_stage.task_count()),
138 input_stage,
139 }
140 }
141
142 /// Creates a new [NetworkBroadcastExec] fed by the provided [BroadcastExec]. The input plan
143 /// will be executed in a remote worker in `producer_tasks` number of tasks.
144 pub fn try_new(input: Arc<dyn ExecutionPlan>, producer_tasks: usize) -> Result<Self> {
145 if !input.is::<BroadcastExec>() {
146 return plan_err!("The input of a NetworkBroadcastExec can only be a BroadcastExec");
147 }
148
149 let input_properties = Arc::clone(input.properties());
150 Ok(Self::from_stage(
151 Stage::Local(LocalStage {
152 // At this point, query_id and num are just placeholders that will be filled by
153 // prepare_network_boundaries.rs. Users are not expected to provide valid values for
154 // these two parameters.
155 query_id: Uuid::nil(),
156 num: 0,
157 plan: input,
158 tasks: producer_tasks,
159 metrics_set: Default::default(),
160 }),
161 input_properties,
162 ))
163 }
164}
165
166impl NetworkBoundary for NetworkBroadcastExec {
167 fn with_input_stage(&self, input_stage: Stage) -> Result<Arc<dyn NetworkBoundary>> {
168 let mut self_clone = self.clone();
169 self_clone.worker_connections = WorkerConnectionPool::new(input_stage.task_count());
170 self_clone.input_stage = input_stage;
171 Ok(Arc::new(self_clone))
172 }
173
174 fn input_stage(&self) -> &Stage {
175 &self.input_stage
176 }
177
178 fn producer_head(&self, consumer_task_count: usize) -> Result<ProducerHead> {
179 let partition_count = self.properties.output_partitioning().partition_count();
180 Ok(ProducerHead::BroadcastExec {
181 output_partitions: partition_count * consumer_task_count,
182 })
183 }
184}
185
186impl DisplayAs for NetworkBroadcastExec {
187 fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
188 let input_tasks = self.input_stage.task_count();
189 let stage = self.input_stage.num();
190 let consumer_partitions = self.properties.partitioning.partition_count();
191 let stage_partitions = self
192 .input_stage
193 .local_plan()
194 .as_ref()
195 .map(|p| p.properties().partitioning.partition_count())
196 .unwrap_or(0);
197 write!(
198 f,
199 "[Stage {stage}] => NetworkBroadcastExec: partitions_per_consumer={consumer_partitions}, stage_partitions={stage_partitions}, input_tasks={input_tasks}",
200 )
201 }
202}
203
204impl ExecutionPlan for NetworkBroadcastExec {
205 fn name(&self) -> &str {
206 "NetworkBroadcastExec"
207 }
208
209 fn properties(&self) -> &Arc<PlanProperties> {
210 &self.properties
211 }
212
213 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
214 match &self.input_stage.local_plan() {
215 Some(plan) => vec![plan],
216 None => vec![],
217 }
218 }
219
220 fn apply_expressions(
221 &self,
222 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
223 ) -> Result<TreeNodeRecursion> {
224 Ok(TreeNodeRecursion::Continue)
225 }
226
227 fn with_new_children(
228 self: Arc<Self>,
229 children: Vec<Arc<dyn ExecutionPlan>>,
230 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
231 let mut self_clone = self.as_ref().clone();
232 match &mut self_clone.input_stage {
233 Stage::Local(local) => {
234 local.plan = require_one_child(children)?;
235 }
236 Stage::Remote(_) => {
237 if !children.is_empty() {
238 not_impl_err!("NetworkBoundary cannot accept children")?
239 }
240 }
241 }
242 Ok(Arc::new(self_clone))
243 }
244
245 fn execute(
246 &self,
247 partition: usize,
248 context: Arc<TaskContext>,
249 ) -> Result<SendableRecordBatchStream, DataFusionError> {
250 let remote_stage = match &self.input_stage {
251 Stage::Local(local) => return local.execute(partition, context),
252 Stage::Remote(remote_stage) => remote_stage,
253 };
254
255 let task_context = DistributedTaskContext::from_ctx(&context);
256 let out_partitions = self.properties.partitioning.partition_count();
257 let off = out_partitions * task_context.task_index;
258 let mut streams = Vec::with_capacity(self.input_stage.task_count());
259
260 for input_task_index in 0..self.input_stage.task_count() {
261 streams.push(self.worker_connections.execute(
262 remote_stage,
263 off..(off + self.properties.partitioning.partition_count()),
264 input_task_index,
265 off + partition,
266 self.producer_head(task_context.task_count)?,
267 &context,
268 )?);
269 }
270
271 Ok(Box::pin(RecordBatchStreamAdapter::new(
272 self.schema(),
273 futures::stream::select_all(streams),
274 )))
275 }
276
277 fn metrics(&self) -> Option<MetricsSet> {
278 Some(self.worker_connections.metrics.clone_inner())
279 }
280
281 fn statistics_from_inputs(
282 &self,
283 _input_stats: &[Arc<Statistics>],
284 args: &StatisticsArgs,
285 ) -> Result<Arc<Statistics>> {
286 self.input_stage.partition_statistics(
287 args.partition(),
288 self.properties.output_partitioning().partition_count(),
289 self.schema(),
290 )
291 }
292}