datafusion_distributed/worker/
task_data.rs

1use crate::common::OnceLockResult;
2use crate::common::now_ns;
3use crate::{MaxLatencyMetric, ProducerHead, TaskMetrics};
4use datafusion::common::{DataFusionError, Result};
5use datafusion::execution::TaskContext;
6use datafusion::physical_plan::ExecutionPlan;
7use datafusion::physical_plan::metrics::{Metric, MetricValue, MetricsSet};
8use std::borrow::Cow;
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::sync::oneshot;
12
13#[derive(Clone, Debug)]
14/// TaskData stores state for a single task being executed by this Endpoint. It may be shared
15/// by concurrent requests for the same task which execute separate partitions.
16pub struct TaskData {
17    /// Task context suitable for execute different partitions from the same task.
18    pub(crate) task_ctx: Arc<TaskContext>,
19    pub(crate) base_plan: Arc<dyn ExecutionPlan>,
20    pub(crate) final_plan: Arc<OnceLockResult<Arc<dyn ExecutionPlan>>>,
21    /// Sender half of the metrics channel. `impl_coordinator_channel` takes this (via
22    /// `Option::take`) when the coordinator channel reaches EOS, sending the collected metrics
23    /// back to the coordinator through the `CoordinatorChannel` side channel.
24    pub(super) metrics_tx: Arc<std::sync::Mutex<Option<oneshot::Sender<TaskMetrics>>>>,
25    /// Metrics related to the execution of a task within a stage. This metrics, instead of being
26    /// associated to a specific node, they are global to the task, like the time at which the plan
27    /// was fed by the coordinator to the worker.
28    pub(super) task_data_metrics: Arc<TaskDataMetrics>,
29}
30
31pub(crate) const PLAN_ADDED_AT_METRIC: &str = "plan_added_at";
32pub(crate) const PLAN_EXECUTED_AT_METRIC: &str = "plan_executed_at";
33pub(crate) const PLAN_FINISHED_AT_METRIC: &str = "plan_finished_at";
34
35#[derive(Debug)]
36pub(super) struct TaskDataMetrics {
37    pub(super) query_start_time_ns: usize,
38    /// When the plan was set by the coordinator.
39    pub(super) plan_added_at: MaxLatencyMetric,
40    /// When the plan execution was triggered by the parent worker.
41    pub(super) plan_executed_at: MaxLatencyMetric,
42    /// When the execution stream finished.
43    pub(super) plan_finished_at: MaxLatencyMetric,
44}
45
46impl TaskDataMetrics {
47    pub(super) fn new(query_start_time_ns: usize) -> Self {
48        let plan_added_at = MaxLatencyMetric::default();
49        plan_added_at.add_duration(Duration::from_nanos(
50            now_ns::<u64>().saturating_sub(query_start_time_ns as u64),
51        ));
52        Self {
53            query_start_time_ns,
54            plan_added_at,
55            plan_finished_at: MaxLatencyMetric::default(),
56            plan_executed_at: MaxLatencyMetric::default(),
57        }
58    }
59
60    pub(super) fn mark_execution_started_once(&self) {
61        if self.plan_executed_at.value() == 0 {
62            self.plan_executed_at.add_duration(Duration::from_nanos(
63                now_ns::<u64>().saturating_sub(self.query_start_time_ns as u64),
64            ))
65        }
66    }
67
68    pub(super) fn mark_execution_finished(&self) {
69        self.plan_finished_at.add_duration(Duration::from_nanos(
70            now_ns::<u64>().saturating_sub(self.query_start_time_ns as u64),
71        ))
72    }
73
74    pub(super) fn to_metrics_set(&self) -> MetricsSet {
75        let mut metrics_set = MetricsSet::new();
76        metrics_set.push(max_latency_metric(
77            PLAN_ADDED_AT_METRIC,
78            &self.plan_added_at,
79        ));
80        metrics_set.push(max_latency_metric(
81            PLAN_EXECUTED_AT_METRIC,
82            &self.plan_executed_at,
83        ));
84        metrics_set.push(max_latency_metric(
85            PLAN_FINISHED_AT_METRIC,
86            &self.plan_finished_at,
87        ));
88
89        metrics_set
90    }
91}
92
93fn max_latency_metric(name: &'static str, value: &MaxLatencyMetric) -> Arc<Metric> {
94    Arc::new(Metric::new(
95        MetricValue::Custom {
96            name: Cow::Borrowed(name),
97            value: Arc::new(MaxLatencyMetric::from_nanos(value.value())),
98        },
99        None,
100    ))
101}
102
103impl TaskData {
104    pub(crate) fn plan(&self, producer_head: ProducerHead) -> Result<Arc<dyn ExecutionPlan>> {
105        let result = self.final_plan.get_or_init(|| {
106            let producer_head =
107                producer_head.ensure_decoded(self.base_plan.schema(), &self.task_ctx)?;
108
109            Ok(producer_head.insert(Arc::clone(&self.base_plan))?)
110        });
111        match result {
112            Ok(plan) => Ok(Arc::clone(plan)),
113            Err(err) => Err(DataFusionError::Shared(Arc::clone(err))),
114        }
115    }
116}
datafusion_distributed/worker/task_data.rs

datafusion_distributed/worker/
task_data.rs
