datafusion_distributed/execution_plans/
network_coalesce.rs

1use crate::DistributedTaskContext;
2use crate::common::require_one_child;
3use crate::distributed_planner::{NetworkBoundary, ProducerHead};
4use crate::execution_plans::common::scale_partitioning_props;
5use crate::stage::{LocalStage, Stage};
6use crate::worker::WorkerConnectionPool;
7use datafusion::common::tree_node::TreeNodeRecursion;
8use datafusion::common::{exec_err, not_impl_err, plan_err};
9use datafusion::error::Result;
10use datafusion::execution::{SendableRecordBatchStream, TaskContext};
11use datafusion::physical_expr::PhysicalExpr;
12use datafusion::physical_expr_common::metrics::MetricsSet;
13use datafusion::physical_plan::limit::LocalLimitExec;
14use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
15use datafusion::physical_plan::{
16    DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, PlanProperties,
17    Statistics, StatisticsArgs, internal_err,
18};
19use std::fmt::{Debug, Formatter};
20use std::sync::Arc;
21use uuid::Uuid;
22
23/// [ExecutionPlan] that coalesces partitions from multiple tasks into a one or more task without
24/// performing any repartition, and maintaining the same partitioning scheme.
25///
26/// This is the equivalent of a [CoalescePartitionsExec] but coalescing tasks across the network
27/// between distributed stages.
28///
29/// ```text
30///                                ┌───────────────────────────┐                                   ■
31///                                │    NetworkCoalesceExec    │                                   │
32///                                │         (task 1)          │                                   │
33///                                └┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┬─┬┘                                Stage N+1
34///                                 │1││2││3││4││5││6││7││8││9│                                    │
35///                                 └─┘└─┘└─┘└─┘└─┘└─┘└─┘└─┘└─┘                                    │
36///                                 ▲  ▲  ▲   ▲  ▲  ▲   ▲  ▲  ▲                                    ■
37///   ┌──┬──┬───────────────────────┴──┴──┘   │  │  │   └──┴──┴──────────────────────┬──┬──┐
38///   │  │  │                                 │  │  │                                │  │  │       ■
39///  ┌─┐┌─┐┌─┐                               ┌─┐┌─┐┌─┐                              ┌─┐┌─┐┌─┐      │
40///  │1││2││3│                               │4││5││6│                              │7││8││9│      │
41/// ┌┴─┴┴─┴┴─┴──────────────────┐  ┌─────────┴─┴┴─┴┴─┴─────────┐ ┌──────────────────┴─┴┴─┴┴─┴┐  Stage N
42/// │  Arc<dyn ExecutionPlan>   │  │  Arc<dyn ExecutionPlan>   │ │  Arc<dyn ExecutionPlan>   │     │
43/// │         (task 1)          │  │         (task 2)          │ │         (task 3)          │     │
44/// └───────────────────────────┘  └───────────────────────────┘ └───────────────────────────┘     ■
45/// ```
46///
47/// The communication between two stages across a [NetworkCoalesceExec] has two implications:
48///
49/// - Stage N+1 may have one or more tasks. Each consumer task reads a contiguous group of upstream
50///   tasks from Stage N.
51/// - Output partitioning for Stage N+1 is sized based on the maximum upstream-group size. When
52///   groups are uneven, consumer tasks with smaller groups return empty streams for the “extra”
53///   partitions.
54/// ```text
55///                    ┌───────────────────────────┐        ┌───────────────────────────┐          ■
56///                    │    NetworkCoalesceExec    │        │    NetworkCoalesceExec    │          │
57///                    │         (task 1)          │        │         (task 2)          │          │
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│      │
66/// ┌┴─┴┴─┴┴─┴──────────────────┐  ┌─────────┴─┴┴─┴┴─┴─────────┐ ┌──────────────────┴─┴┴─┴┴─┴┐  Stage N
67/// │  Arc<dyn ExecutionPlan>   │  │  Arc<dyn ExecutionPlan>   │ │  Arc<dyn ExecutionPlan>   │     │
68/// │         (task 1)          │  │         (task 2)          │ │         (task 3)          │     │
69/// └───────────────────────────┘  └───────────────────────────┘ └───────────────────────────┘     ■
70/// ```
71///
72/// This node has two variants.
73/// 1. Pending: acts as a placeholder for the distributed optimization step to mark it as ready.
74/// 2. Ready: runs within a distributed stage and queries the next input stage over the network
75///    using Arrow Flight.
76#[derive(Debug, Clone)]
77pub struct NetworkCoalesceExec {
78    /// the properties we advertise for this execution plan
79    pub(crate) properties: Arc<PlanProperties>,
80    pub(crate) input_stage: Stage,
81    pub(crate) worker_connections: WorkerConnectionPool,
82}
83
84impl NetworkCoalesceExec {
85    pub(crate) fn try_from_stage(
86        input_stage: Stage,
87        input_properties: Arc<PlanProperties>,
88        consumer_tasks: usize,
89    ) -> Result<Self> {
90        // Each output task coalesces a group of input tasks. We size the output partition count
91        // per output task based on the maximum group size, returning empty streams for tasks with
92        // smaller groups.
93        let max_input_task_count = input_stage.task_count().div_ceil(consumer_tasks).max(1);
94        let props = scale_partitioning_props(&input_properties, |p| p * max_input_task_count)?;
95
96        Ok(Self {
97            properties: props,
98            worker_connections: WorkerConnectionPool::new(input_stage.task_count()),
99            input_stage,
100        })
101    }
102
103    /// Creates a new [NetworkCoalesceExec] fed by the provided `input` plan.
104    ///
105    /// The `input` plan will be remotely executed in `producer_tasks` tasks, while the
106    /// [NetworkCoalesceExec] will be executed in `consumer_tasks` tasks in the stage above.
107    ///
108    /// Typically, this node should be placed right after nodes that coalesce all the input
109    /// partitions into one, for example:
110    /// - [CoalescePartitionsExec]
111    /// - [SortPreservingMergeExec]
112    ///
113    /// ## Warning
114    ///
115    /// The caller must ensure that the provided `consumer_tasks` count matches the `producer_tasks`
116    /// of the network boundary immediately above.
117    pub fn try_new(
118        input: Arc<dyn ExecutionPlan>,
119        producer_tasks: usize,
120        consumer_tasks: usize,
121    ) -> Result<Self> {
122        if consumer_tasks == 0 {
123            return plan_err!("The `consumer_tasks` input of a NetworkCoalesceExec must not be 0");
124        }
125
126        let input_properties = Arc::clone(input.properties());
127        Self::try_from_stage(
128            Stage::Local(LocalStage {
129                // At this point, query_id and num are just placeholders that will be filled by
130                // prepare_network_boundaries.rs. Users are not expected to provide valid values for
131                // these two parameters.
132                query_id: Uuid::nil(),
133                num: 0,
134                plan: input,
135                tasks: producer_tasks,
136                metrics_set: Default::default(),
137            }),
138            input_properties,
139            consumer_tasks,
140        )
141    }
142
143    pub(crate) fn with_fetch_on_input_stage(&self, fetch: usize) -> Result<Arc<dyn ExecutionPlan>> {
144        let Stage::Local(local) = &self.input_stage else {
145            return Ok(Arc::new(self.clone()));
146        };
147
148        let input_with_fetch = if local.plan.fetch().is_some_and(|existing| existing <= fetch) {
149            Arc::clone(&local.plan)
150        } else {
151            local
152                .plan
153                .with_fetch(Some(fetch))
154                .unwrap_or_else(|| Arc::new(LocalLimitExec::new(Arc::clone(&local.plan), fetch)))
155        };
156
157        let mut self_clone = self.clone();
158        self_clone.input_stage = Stage::Local(LocalStage {
159            query_id: local.query_id,
160            num: local.num,
161            plan: input_with_fetch,
162            tasks: local.tasks,
163            metrics_set: Default::default(),
164        });
165        Ok(Arc::new(self_clone))
166    }
167}
168
169impl NetworkBoundary for NetworkCoalesceExec {
170    fn input_stage(&self) -> &Stage {
171        &self.input_stage
172    }
173
174    fn with_input_stage(&self, input_stage: Stage) -> Result<Arc<dyn NetworkBoundary>> {
175        let mut self_clone = self.clone();
176        self_clone.properties = scale_partitioning_props(self_clone.properties(), |p| {
177            p * input_stage.task_count() / self_clone.input_stage.task_count().max(1)
178        })?;
179        self_clone.worker_connections = WorkerConnectionPool::new(input_stage.task_count());
180        self_clone.input_stage = input_stage;
181        Ok(Arc::new(self_clone))
182    }
183
184    fn producer_head(&self, _consumer_task_count: usize) -> Result<ProducerHead> {
185        Ok(ProducerHead::None)
186    }
187}
188
189impl DisplayAs for NetworkCoalesceExec {
190    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
191        let input_tasks = self.input_stage.task_count();
192        let partitions = self.properties.partitioning.partition_count();
193        let stage = self.input_stage.num();
194        write!(
195            f,
196            "[Stage {stage}] => NetworkCoalesceExec: output_partitions={partitions}, input_tasks={input_tasks}",
197        )
198    }
199}
200
201impl ExecutionPlan for NetworkCoalesceExec {
202    fn name(&self) -> &str {
203        "NetworkCoalesceExec"
204    }
205
206    fn properties(&self) -> &Arc<PlanProperties> {
207        &self.properties
208    }
209
210    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
211        match &self.input_stage.local_plan() {
212            Some(v) => vec![v],
213            None => vec![],
214        }
215    }
216
217    fn apply_expressions(
218        &self,
219        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
220    ) -> Result<TreeNodeRecursion> {
221        Ok(TreeNodeRecursion::Continue)
222    }
223
224    fn with_new_children(
225        self: Arc<Self>,
226        children: Vec<Arc<dyn ExecutionPlan>>,
227    ) -> Result<Arc<dyn ExecutionPlan>> {
228        let mut self_clone = self.as_ref().clone();
229        match &mut self_clone.input_stage {
230            Stage::Local(local) => {
231                local.plan = require_one_child(children)?;
232            }
233            Stage::Remote(_) => {
234                if !children.is_empty() {
235                    not_impl_err!("NetworkBoundary cannot accept children")?
236                }
237            }
238        }
239        Ok(Arc::new(self_clone))
240    }
241
242    fn execute(
243        &self,
244        partition: usize,
245        context: Arc<TaskContext>,
246    ) -> Result<SendableRecordBatchStream> {
247        let remote_stage = match &self.input_stage {
248            Stage::Local(local) => return local.execute(partition, context),
249            Stage::Remote(remote_stage) => remote_stage,
250        };
251
252        let task_context = DistributedTaskContext::from_ctx(&context);
253        if task_context.task_index >= task_context.task_count {
254            return exec_err!(
255                "NetworkCoalesceExec invalid task context: task_index={} >= task_count={}",
256                task_context.task_index,
257                task_context.task_count
258            );
259        }
260
261        let out_partitions = self.properties().partitioning.partition_count();
262        let partitions_per_task = out_partitions
263            .checked_div(
264                self.input_stage
265                    .task_count()
266                    .div_ceil(task_context.task_count)
267                    .max(1),
268            )
269            .unwrap_or(0);
270        if partitions_per_task == 0 {
271            return exec_err!("NetworkCoalesceExec has 0 partitions per input task");
272        }
273
274        let input_task_count = self.input_stage.task_count();
275        let group = task_group(
276            input_task_count,
277            task_context.task_index,
278            task_context.task_count,
279        );
280
281        let input_task_offset = partition / partitions_per_task;
282        let target_partition = partition % partitions_per_task;
283
284        // Some consumer tasks are assigned fewer upstream tasks when
285        // `input_task_count % task_count != 0` (uneven grouping).
286        // We still size partitions based on the maximum group size, so partitions that
287        // would map to a missing upstream task slot are treated as padding and return
288        // an empty stream (no network call).
289        if input_task_offset >= group.len {
290            return Ok(Box::pin(EmptyRecordBatchStream::new(self.schema())));
291        }
292
293        // This should never happen.
294        if input_task_offset >= group.max_len {
295            return internal_err!(
296                "NetworkCoalesceExec input_task_offset={} >= group.max_len={}",
297                input_task_offset,
298                group.max_len
299            );
300        }
301
302        let target_task = group.start_task + input_task_offset;
303
304        let stream = self.worker_connections.execute(
305            remote_stage,
306            0..partitions_per_task,
307            target_task,
308            target_partition,
309            self.producer_head(task_context.task_count)?,
310            &context,
311        )?;
312
313        Ok(Box::pin(RecordBatchStreamAdapter::new(
314            self.schema(),
315            stream,
316        )))
317    }
318
319    fn metrics(&self) -> Option<MetricsSet> {
320        Some(self.worker_connections.metrics.clone_inner())
321    }
322
323    fn statistics_from_inputs(
324        &self,
325        _input_stats: &[Arc<Statistics>],
326        args: &StatisticsArgs,
327    ) -> Result<Arc<Statistics>> {
328        self.input_stage.partition_statistics(
329            args.partition(),
330            self.properties.output_partitioning().partition_count(),
331            self.schema(),
332        )
333    }
334}
335
336#[derive(Debug, Clone, Copy)]
337struct TaskGroup {
338    /// The first input task index in this group.
339    start_task: usize,
340    /// The number of input tasks in this group.
341    len: usize,
342    /// The maximum possible group size across all groups.
343    ///
344    /// When groups are uneven (input_tasks % task_count != 0), some groups are shorter. We still
345    /// size the output partitioning based on this max and return empty streams for the extra
346    /// partitions in smaller groups.
347    max_len: usize,
348}
349
350/// Returns the contiguous group of input tasks assigned to DistributedTaskContext::task_index.
351fn task_group(input_task_count: usize, task_index: usize, task_count: usize) -> TaskGroup {
352    if task_count == 0 {
353        return TaskGroup {
354            start_task: 0,
355            len: 0,
356            max_len: 0,
357        };
358    }
359
360    // Split `input_task_count` into `task_count` contiguous groups.
361    // - base_tasks_per_group: floor(input_task_count / task_count)
362    // - groups_with_extra_task: first N groups that get one extra task (remainder)
363    let base_tasks_per_group = input_task_count / task_count;
364    let groups_with_extra_task = input_task_count % task_count;
365
366    let len = base_tasks_per_group + usize::from(task_index < groups_with_extra_task);
367    let start_task = (task_index * base_tasks_per_group) + task_index.min(groups_with_extra_task);
368    let max_len = base_tasks_per_group + usize::from(groups_with_extra_task > 0);
369
370    TaskGroup {
371        start_task,
372        len,
373        max_len,
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380    use datafusion::arrow::datatypes::Schema;
381    use datafusion::physical_plan::empty::EmptyExec;
382
383    #[derive(Clone, Copy)]
384    struct Case {
385        name: &'static str,
386        input_tasks: usize,
387        consumer_tasks: usize,
388    }
389
390    fn expected_groups(input_tasks: usize, consumer_tasks: usize) -> Vec<(usize, usize)> {
391        assert!(consumer_tasks > 0, "consumer_tasks must be non-zero");
392
393        let base_tasks_per_group = input_tasks / consumer_tasks;
394        let groups_with_extra_task = input_tasks % consumer_tasks;
395        let mut groups = Vec::with_capacity(consumer_tasks);
396        let mut start_task = 0;
397
398        for task_index in 0..consumer_tasks {
399            let len = base_tasks_per_group + usize::from(task_index < groups_with_extra_task);
400            groups.push((start_task, len));
401            start_task += len;
402        }
403
404        groups
405    }
406
407    fn assert_case(case: Case) -> Result<()> {
408        // Child plan used only for properties/schema (we won't reach network codepaths).
409        let child: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
410        let child_partitions = child.properties().partitioning.partition_count();
411
412        let exec = NetworkCoalesceExec::try_new(
413            Arc::clone(&child),
414            case.input_tasks,
415            case.consumer_tasks,
416        )?;
417
418        // Output partitions are sized by the maximum group size.
419        let max_group_size = case.input_tasks.div_ceil(case.consumer_tasks).max(1);
420        assert_eq!(
421            exec.properties().partitioning.partition_count(),
422            child_partitions * max_group_size
423        );
424
425        let groups = expected_groups(case.input_tasks, case.consumer_tasks);
426        assert_eq!(groups.len(), case.consumer_tasks);
427
428        let mut seen = vec![false; case.input_tasks];
429        let mut expected_start = 0;
430        let mut padding_slots = 0;
431
432        for (index, (start, len)) in groups.into_iter().enumerate() {
433            assert_eq!(
434                start, expected_start,
435                "case {} group {} should be contiguous",
436                case.name, index
437            );
438            assert!(
439                start + len <= case.input_tasks,
440                "case {} group {} exceeds input task count",
441                case.name,
442                index
443            );
444
445            for (offset, seen_task) in seen.iter_mut().skip(start).take(len).enumerate() {
446                let task = start + offset;
447                assert!(
448                    !*seen_task,
449                    "case {} input task {} appears twice",
450                    case.name, task
451                );
452                *seen_task = true;
453            }
454
455            expected_start = start + len;
456            padding_slots += max_group_size - len;
457        }
458
459        assert_eq!(
460            expected_start, case.input_tasks,
461            "case {} groups should cover all input tasks",
462            case.name
463        );
464        assert!(
465            seen.iter().all(|v| *v),
466            "case {} missing at least one input task",
467            case.name
468        );
469
470        let total_slots = case.consumer_tasks * max_group_size;
471        let total_padding = total_slots - case.input_tasks;
472        assert_eq!(
473            padding_slots, total_padding,
474            "case {} padding slots mismatch",
475            case.name
476        );
477
478        Ok(())
479    }
480
481    const ONE_TO_MANY_INPUT: usize = 1;
482    const ONE_TO_MANY_OUTPUT: usize = 3;
483    const MANY_TO_ONE_INPUT: usize = 4;
484    const MANY_TO_ONE_OUTPUT: usize = 1;
485    const MANY_TO_FEWER_INPUT: usize = 5;
486    const MANY_TO_FEWER_OUTPUT: usize = 2;
487    const FEWER_TO_MANY_INPUT: usize = 2;
488    const FEWER_TO_MANY_OUTPUT: usize = 5;
489
490    #[test]
491    fn validates_partition_coverage_one_to_many() -> Result<()> {
492        assert_case(Case {
493            name: "1_to_n",
494            input_tasks: ONE_TO_MANY_INPUT,
495            consumer_tasks: ONE_TO_MANY_OUTPUT,
496        })
497    }
498
499    #[test]
500    fn validates_partition_coverage_many_to_one() -> Result<()> {
501        assert_case(Case {
502            name: "n_to_1",
503            input_tasks: MANY_TO_ONE_INPUT,
504            consumer_tasks: MANY_TO_ONE_OUTPUT,
505        })
506    }
507
508    #[test]
509    fn validates_partition_coverage_many_to_fewer() -> Result<()> {
510        assert_case(Case {
511            name: "n_to_m_n_gt_m",
512            input_tasks: MANY_TO_FEWER_INPUT,
513            consumer_tasks: MANY_TO_FEWER_OUTPUT,
514        })
515    }
516
517    #[test]
518    fn validates_partition_coverage_fewer_to_many() -> Result<()> {
519        assert_case(Case {
520            name: "m_to_n_n_gt_m",
521            input_tasks: FEWER_TO_MANY_INPUT,
522            consumer_tasks: FEWER_TO_MANY_OUTPUT,
523        })
524    }
525}
datafusion_distributed/execution_plans/network_coalesce.rs

datafusion_distributed/execution_plans/
network_coalesce.rs
