wip

Author: jayshrivastava

src/execution_plans/distributed.rs

@@ -1,10 +1,9 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
 use crate::distributed_physical_optimizer_rule::NetworkBoundaryExt;
 use crate::execution_plans::common::require_one_child;
+use crate::protobuf::DistributedCodec;
+use crate::stage::{DisplayCtx, MaybeEncodedPlan};
 use crate::stage::{ExecutionTask, Stage};
-use crate::protobuf::{DistributedCodec, StageKey};
-use crate::stage::DisplayCtx;
-use bytes::Bytes;
 use datafusion::common::exec_err;
 use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::error::DataFusionError;
@@ -43,11 +42,12 @@ impl DistributedExec {
 
     /// Returns a special stage key used to identify the root "stage" of the distributed plan.
     /// TODO: reconcile this with display_plan_graphviz
-    pub(crate) fn to_stage_key(&self) -> StageKey {
-        StageKey {
-            query_id: Bytes::new(),
-            stage_id: 0_u64,
-            task_number: 0,
+    pub(crate) fn synthetic_stage(&self) -> Stage {
+        Stage {
+            query_id: uuid::Uuid::nil(),
+            num: 0,
+            plan: MaybeEncodedPlan::Decoded(self.plan.clone()),
+            tasks: vec![ExecutionTask { url: None }],
         }
     }
 

src/execution_plans/metrics.rs

@@ -53,7 +53,7 @@ impl ExecutionPlan for MetricsWrapperExec {
     }
 
     fn properties(&self) -> &PlanProperties {
-        unimplemented!("MetricsWrapperExec does not implement properties")
+        self.inner.properties()
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {

src/explain.rs

@@ -4,7 +4,10 @@ use crate::metrics::MetricsCollectorResult;
 use crate::metrics::TaskMetricsCollector;
 use crate::metrics::proto::MetricsSetProto;
 use crate::metrics::proto::df_metrics_set_to_proto;
+use crate::metrics::rewrite_distributed_plan_with_metrics;
+use crate::protobuf::StageKey;
 use crate::stage::DisplayCtx;
+use bytes::Bytes;
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
@@ -25,15 +28,28 @@ pub fn explain_analyze(executed: Arc<dyn ExecutionPlan>) -> Result<String, DataF
                 mut input_task_metrics,
             } = TaskMetricsCollector::new().collect(dist_exec.pepared_plan()?)?;
 
+            // let input_stage = dist_exec.pepared_plan()?
+            let fake_stage = dist_exec.synthetic_stage();
+            let stage_key = StageKey {
+                query_id: Bytes::from(fake_stage.query_id.as_bytes().to_vec()),
+                stage_id: fake_stage.num as u64,
+                task_number: 0,
+            };
             input_task_metrics.insert(
-                dist_exec.to_stage_key(),
+                stage_key,
                 task_metrics
                     .into_iter()
                     .map(|metrics| df_metrics_set_to_proto(&metrics))
                     .collect::<Result<Vec<MetricsSetProto>, DataFusionError>>()?,
             );
 
-            let display_ctx = DisplayCtx::new(input_task_metrics);
+            let executed = rewrite_distributed_plan_with_metrics(
+                executed.clone(),
+                Arc::new(input_task_metrics),
+            )?;
+            let dist_exec = executed.as_any().downcast_ref::<DistributedExec>().unwrap();
+
+            let display_ctx = DisplayCtx::new().with_metrics();
             Ok(display_plan_ascii(&dist_exec.with_display_ctx(display_ctx)))
         }
     }

src/flight_service/do_get.rs

@@ -165,7 +165,7 @@ fn collect_and_create_metrics_flight_data(
     plan: Arc<dyn ExecutionPlan>,
     incoming: FlightData,
 ) -> Result<FlightData, FlightError> {
-    // Get the metrics for the task executed on this worker. Separately, collect metrics for child tasks.
+    // Get the metrics for the task executed on this worker + child tasks.
     let mut result = TaskMetricsCollector::new()
         .collect(plan)
         .map_err(|err| FlightError::ProtocolError(err.to_string()))?;

src/metrics/mod.rs

@@ -4,4 +4,6 @@ mod task_metrics_collector;
 mod task_metrics_rewriter;
 pub(crate) use metrics_collecting_stream::MetricsCollectingStream;
 pub(crate) use task_metrics_collector::{MetricsCollectorResult, TaskMetricsCollector};
-pub(crate) use task_metrics_rewriter::TaskMetricsRewriter;
+#[allow(unused_imports)]
+pub use task_metrics_rewriter::StageMetricsRewriter;
+pub(crate) use task_metrics_rewriter::rewrite_distributed_plan_with_metrics;

src/metrics/proto.rs

@@ -26,6 +26,18 @@ pub struct MetricsSetProto {
     pub metrics: Vec<MetricProto>,
 }
 
+impl MetricsSetProto {
+    pub fn new() -> Self {
+        Self {
+            metrics: Vec::new(),
+        }
+    }
+
+    pub fn push(&mut self, metric: MetricProto) {
+        self.metrics.push(metric)
+    }
+}
+
 /// MetricValueProto is a protobuf mirror of the [datafusion::physical_plan::metrics::MetricValue] enum.
 #[derive(Clone, PartialEq, Eq, ::prost::Oneof)]
 pub enum MetricValueProto {