datafusion_distributed/execution_plans/network_shuffle.rs
1use crate::common::require_one_child;
2use crate::distributed_planner::ProducerHead;
3use crate::execution_plans::common::scale_partitioning;
4use crate::stage::{LocalStage, Stage};
5use crate::worker::WorkerConnectionPool;
6use crate::{DistributedTaskContext, MaybeEncoded, NetworkBoundary};
7use datafusion::common::tree_node::TreeNodeRecursion;
8use datafusion::common::{Result, not_impl_err, plan_err};
9use datafusion::error::DataFusionError;
10use datafusion::execution::{SendableRecordBatchStream, TaskContext};
11use datafusion::physical_expr::{Partitioning, PhysicalExpr};
12use datafusion::physical_expr_common::metrics::MetricsSet;
13use datafusion::physical_plan::repartition::RepartitionExec;
14use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
15use datafusion::physical_plan::{
16 DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics, StatisticsArgs,
17};
18use std::fmt::Formatter;
19use std::sync::Arc;
20use uuid::Uuid;
21
22/// [ExecutionPlan] implementation that shuffles data across the network in a distributed context.
23///
24/// The easiest way of thinking about this node is as a plan [RepartitionExec] node that is
25/// capable of fanning out the different produced partitions to different tasks.
26/// This allows redistributing data across different tasks in different stages, so that different
27/// physical machines can make progress on different non-overlapping sets of data.
28///
29/// This node allows fanning out of data from N tasks to M tasks, with N and M being arbitrary non-zero
30/// positive numbers. Here are some examples of how data can be shuffled in different scenarios:
31///
32/// # 1 to many
33///
34/// ```text
35/// ┌───────────────────────────┐ ┌───────────────────────────┐ ┌───────────────────────────┐ ■
36/// │ NetworkShuffleExec │ │ NetworkShuffleExec │ │ NetworkShuffleExec │ │
37/// │ (task 1) │ │ (task 2) │ │ (task 3) │ │
38/// └┬─┬┬─┬┬─┬──────────────────┘ └─────────┬─┬┬─┬┬─┬─────────┘ └──────────────────┬─┬┬─┬┬─┬┘ Stage N+1
39/// │1││2││3│ │4││5││6│ │7││8││9│ │
40/// └─┘└─┘└─┘ └─┘└─┘└─┘ └─┘└─┘└─┘ │
41/// ▲ ▲ ▲ ▲ ▲ ▲ ▲ ▲ ▲ ■
42/// └──┴──┴────────────────────────┬──┬──┐ │ │ │ ┌──┬──┬───────────────────────┴──┴──┘
43/// │ │ │ │ │ │ │ │ │ ■
44/// ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ │
45/// │1││2││3││4││5││6││7││8││9│ │
46/// ┌┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┐ Stage N
47/// │ RepartitionExec │ │
48/// │ (task 1) │ │
49/// └───────────────────────────┘ ■
50/// ```
51///
52/// # many to 1
53///
54/// ```text
55/// ┌───────────────────────────┐ ■
56/// │ NetworkShuffleExec │ │
57/// │ (task 1) │ │
58/// └┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┘ Stage N+1
59/// │1││2││3││4││5││6││7││8││9│ │
60/// └─┘└─┘└─┘└─┘└─┘└─┘└─┘└─┘└─┘ │
61/// ▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲ ■
62/// ┌──┬──┬──┬──┬──┬──┬──┬──┬─────┴┼┴┴┼┴┴┼┴┴┼┴┴┼┴┴┼┴┴┼┴┴┼┴┴┼┴────┬──┬──┬──┬──┬──┬──┬──┬──┐
63/// │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ■
64/// ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ │
65/// │1││2││3││4││5││6││7││8││9│ │1││2││3││4││5││6││7││8││9│ │1││2││3││4││5││6││7││8││9│ │
66/// ┌┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┐ ┌┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┐ ┌┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┐ Stage N
67/// │ RepartitionExec │ │ RepartitionExec │ │ RepartitionExec │ │
68/// │ (task 1) │ │ (task 2) │ │ (task 3) │ │
69/// └───────────────────────────┘ └───────────────────────────┘ └───────────────────────────┘ ■
70/// ```
71///
72/// # many to many
73///
74/// ```text
75/// ┌───────────────────────────┐ ┌───────────────────────────┐ ■
76/// │ NetworkShuffleExec │ │ NetworkShuffleExec │ │
77/// │ (task 1) │ │ (task 2) │ │
78/// └┬─┬┬─┬┬─┬┬─┬───────────────┘ └───────────────┬─┬┬─┬┬─┬┬─┬┘ Stage N+1
79/// │1││2││3││4│ │5││6││7││8│ │
80/// └─┘└─┘└─┘└─┘ └─┘└─┘└─┘└─┘ │
81/// ▲▲▲▲▲▲▲▲▲▲▲▲ ▲▲▲▲▲▲▲▲▲▲▲▲ ■
82/// ┌──┬──┬──┬──┬──┬┴┴┼┴┴┼┴┴┴┴┴┴───┬──┬──┬──┬──┬──┬──┬──┬────────┬┴┴┼┴┴┼┴┴┼┴┴┼──┬──┬──┐
83/// │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ■
84/// ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ ┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐┌─┐ │
85/// │1││2││3││4││5││6││7││8│ │1││2││3││4││5││6││7││8│ │1││2││3││4││5││6││7││8│ │
86/// ┌──┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴─┐ ┌──┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴─┐ ┌──┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴┴─┴─┐ Stage N
87/// │ RepartitionExec │ │ RepartitionExec │ │ RepartitionExec │ │
88/// │ (task 1) │ │ (task 2) │ │ (task 3) │ │
89/// └───────────────────────────┘ └───────────────────────────┘ └───────────────────────────┘ ■
90/// ```
91///
92/// The communication between two stages across a [NetworkShuffleExec] has two implications:
93///
94/// - Each task in Stage N+1 gathers data from all tasks in Stage N
95/// - The total number of partitions across all tasks in Stage N+1 is equal to the
96/// number of partitions in a single task in Stage N. (e.g. (1,2,3,4)+(5,6,7,8) = (1,2,3,4,5,6,7,8) )
97///
98/// This node has two variants.
99/// 1. Pending: acts as a placeholder for the distributed optimization step to mark it as ready.
100/// 2. Ready: runs within a distributed stage and queries the next input stage over the network
101/// using Arrow Flight.
102#[derive(Debug, Clone)]
103pub struct NetworkShuffleExec {
104 /// the properties we advertise for this execution plan
105 pub(crate) properties: Arc<PlanProperties>,
106 pub(crate) input_stage: Stage,
107 pub(crate) worker_connections: WorkerConnectionPool,
108}
109
110impl NetworkShuffleExec {
111 pub(crate) fn from_stage(input_stage: Stage, input_properties: Arc<PlanProperties>) -> Self {
112 Self {
113 properties: input_properties,
114 worker_connections: WorkerConnectionPool::new(input_stage.task_count()),
115 input_stage,
116 }
117 }
118
119 /// Creates a new [NetworkShuffleExec] fed by the provided [RepartitionExec]. The input plan
120 /// will be executed in a remote worker in `producer_tasks` number of tasks.
121 pub fn try_new(input: Arc<dyn ExecutionPlan>, producer_tasks: usize) -> Result<Self> {
122 let Some(r_exec) = input.downcast_ref::<RepartitionExec>() else {
123 return plan_err!("The input of a NetworkShuffleExec can only be a RepartitionExec");
124 };
125 if !matches!(r_exec.partitioning(), Partitioning::Hash(_, _)) {
126 return plan_err!("The input of a NetworkShuffleExec must be hash partitioned");
127 }
128
129 let input_properties = Arc::clone(input.properties());
130 Ok(Self::from_stage(
131 Stage::Local(LocalStage {
132 // At this point, query_id and num are just placeholders that will be filled by
133 // prepare_network_boundaries.rs. Users are not expected to provide valid values for
134 // these two parameters.
135 query_id: Uuid::nil(),
136 num: 0,
137 plan: input,
138 tasks: producer_tasks,
139 metrics_set: Default::default(),
140 }),
141 input_properties,
142 ))
143 }
144}
145
146impl NetworkBoundary for NetworkShuffleExec {
147 fn input_stage(&self) -> &Stage {
148 &self.input_stage
149 }
150
151 fn with_input_stage(&self, input_stage: Stage) -> Result<Arc<dyn NetworkBoundary>> {
152 let mut self_clone = self.clone();
153 self_clone.worker_connections = WorkerConnectionPool::new(input_stage.task_count());
154 self_clone.input_stage = input_stage;
155 Ok(Arc::new(self_clone))
156 }
157
158 fn producer_head(&self, consumer_task_count: usize) -> Result<ProducerHead> {
159 Ok(ProducerHead::RepartitionExec {
160 partitioning: MaybeEncoded::Decoded(scale_partitioning(
161 &self.properties.partitioning,
162 |prev| prev * consumer_task_count,
163 )?),
164 })
165 }
166}
167
168impl DisplayAs for NetworkShuffleExec {
169 fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
170 let input_tasks = self.input_stage.task_count();
171 let partitions = self.properties.partitioning.partition_count();
172 let stage = self.input_stage.num();
173 write!(
174 f,
175 "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}",
176 )
177 }
178}
179
180impl ExecutionPlan for NetworkShuffleExec {
181 fn name(&self) -> &str {
182 "NetworkShuffleExec"
183 }
184
185 fn properties(&self) -> &Arc<PlanProperties> {
186 &self.properties
187 }
188
189 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
190 match &self.input_stage.local_plan() {
191 Some(v) => vec![v],
192 None => vec![],
193 }
194 }
195
196 fn apply_expressions(
197 &self,
198 _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
199 ) -> Result<TreeNodeRecursion> {
200 Ok(TreeNodeRecursion::Continue)
201 }
202
203 fn with_new_children(
204 self: Arc<Self>,
205 children: Vec<Arc<dyn ExecutionPlan>>,
206 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
207 let mut self_clone = self.as_ref().clone();
208 match &mut self_clone.input_stage {
209 Stage::Local(local) => {
210 local.plan = require_one_child(children)?;
211 }
212 Stage::Remote(_) => {
213 if !children.is_empty() {
214 not_impl_err!("NetworkBoundary cannot accept children")?
215 }
216 }
217 }
218 Ok(Arc::new(self_clone))
219 }
220
221 fn execute(
222 &self,
223 partition: usize,
224 context: Arc<TaskContext>,
225 ) -> Result<SendableRecordBatchStream, DataFusionError> {
226 let remote_stage = match &self.input_stage {
227 Stage::Local(local) => return local.execute(partition, context),
228 Stage::Remote(remote_stage) => remote_stage,
229 };
230
231 let task_context = DistributedTaskContext::from_ctx(&context);
232 let out_partitions = self.properties.partitioning.partition_count();
233 let off = out_partitions * task_context.task_index;
234
235 let mut streams = Vec::with_capacity(remote_stage.workers.len());
236 for input_task_index in 0..remote_stage.workers.len() {
237 streams.push(self.worker_connections.execute(
238 remote_stage,
239 off..(off + self.properties.partitioning.partition_count()),
240 input_task_index,
241 off + partition,
242 self.producer_head(task_context.task_count)?,
243 &context,
244 )?);
245 }
246
247 Ok(Box::pin(RecordBatchStreamAdapter::new(
248 self.schema(),
249 futures::stream::select_all(streams),
250 )))
251 }
252
253 fn metrics(&self) -> Option<MetricsSet> {
254 Some(self.worker_connections.metrics.clone_inner())
255 }
256
257 fn statistics_from_inputs(
258 &self,
259 _input_stats: &[Arc<Statistics>],
260 args: &StatisticsArgs,
261 ) -> Result<Arc<Statistics>> {
262 self.input_stage.partition_statistics(
263 args.partition(),
264 self.properties.output_partitioning().partition_count(),
265 self.schema(),
266 )
267 }
268}