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::stage::{ExecutionTask, Stage};
-use crate::protobuf::{DistributedCodec, StageKey};
+use crate::protobuf::DistributedCodec;
 use crate::stage::DisplayCtx;
-use bytes::Bytes;
+use crate::stage::{ExecutionTask, Stage};
 use datafusion::common::exec_err;
 use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::error::DataFusionError;
@@ -27,7 +26,6 @@ use url::Url;
 #[derive(Debug, Clone)]
 pub struct DistributedExec {
     pub plan: Arc<dyn ExecutionPlan>,
-
     pub prepared_plan: Arc<Mutex<Option<Arc<dyn ExecutionPlan>>>>,
     pub display_ctx: Option<DisplayCtx>,
 }
@@ -41,25 +39,15 @@ 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 with_display_ctx(&self, display_ctx: DisplayCtx) -> Self {
         Self {
             display_ctx: Some(display_ctx),
             ..self.clone()
         }
     }
 
-    /// Returns the prepared plan which is lazily prepared on execute(). It is updated on every
-    /// call. Returns an error if .execute() has not been called.
+    /// Returns the plan which is lazily prepared on execute() and actually gets executed.
+    /// It is updated on every call to execute(). Returns an error if .execute() has not been called.
     pub(crate) fn pepared_plan(&self) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         self.prepared_plan
             .lock()

src/execution_plans/metrics.rs

@@ -11,7 +11,7 @@ use std::fmt::{Debug, Formatter};
 /// A transparent wrapper that delegates all execution to its child but returns custom metrics. This node is invisible during display.
 /// The structure of a plan tree is closely tied to the [TaskMetricsRewriter].
 pub struct MetricsWrapperExec {
-    inner: Arc<dyn ExecutionPlan>,
+    pub(crate) inner: Arc<dyn ExecutionPlan>,
     /// metrics for this plan node.
     metrics: MetricsSet,
     /// children is initially None. When used by the [TaskMetricsRewriter], the children will be updated
@@ -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

@@ -1,40 +1,20 @@
 use crate::display_plan_ascii;
 use crate::execution_plans::DistributedExec;
-use crate::metrics::MetricsCollectorResult;
-use crate::metrics::TaskMetricsCollector;
-use crate::metrics::proto::MetricsSetProto;
-use crate::metrics::proto::df_metrics_set_to_proto;
-use crate::stage::DisplayCtx;
+use crate::metrics::rewrite_distributed_plan_with_metrics;
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use std::sync::Arc;
 
 /// explain_analyze renders an [ExecutionPlan] with metrics.
 pub fn explain_analyze(executed: Arc<dyn ExecutionPlan>) -> Result<String, DataFusionError> {
-    // Check if the plan is distributed by looking for a root DistributedExec.
     match executed.as_any().downcast_ref::<DistributedExec>() {
         None => Ok(DisplayableExecutionPlan::with_metrics(executed.as_ref())
             .indent(true)
             .to_string()),
-        Some(dist_exec) => {
-            // If the plan was distributed, collect metrics from the coordinating stage exec.
-            // TODO: Should we move this into the DistributedExec itself or a new ExplainAnalyzeExec?
-            let MetricsCollectorResult {
-                task_metrics,
-                mut input_task_metrics,
-            } = TaskMetricsCollector::new().collect(dist_exec.pepared_plan()?)?;
-
-            input_task_metrics.insert(
-                dist_exec.to_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);
-            Ok(display_plan_ascii(&dist_exec.with_display_ctx(display_ctx)))
+        Some(_) => {
+            let executed = rewrite_distributed_plan_with_metrics(executed.clone())?;
+            Ok(display_plan_ascii(executed.as_ref()))
         }
     }
 }

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/lib.rs

@@ -29,6 +29,7 @@ pub use flight_service::{
     DistributedSessionBuilderContext, MappedDistributedSessionBuilder,
     MappedDistributedSessionBuilderExt,
 };
+pub use metrics::rewrite_distributed_plan_with_metrics;
 pub use stage::{
     DistributedTaskContext, ExecutionTask, Stage, display_plan_ascii, display_plan_graphviz,
 };

src/metrics/metrics_collecting_stream.rs

@@ -155,7 +155,10 @@ mod tests {
                     tasks: vec![TaskMetrics {
                         stage_key: Some(stage_key.clone()),
                         // use the task number to seed the test metrics set for convenience
-                        metrics: vec![make_test_metrics_set_proto_from_seed(stage_key.task_number)],
+                        metrics: vec![make_test_metrics_set_proto_from_seed(
+                            stage_key.task_number,
+                            4,
+                        )],
                     }],
                 })),
             })
@@ -200,7 +203,7 @@ mod tests {
             assert_eq!(collected_metrics.len(), 1);
             assert_eq!(
                 collected_metrics[0],
-                make_test_metrics_set_proto_from_seed(stage_key.task_number)
+                make_test_metrics_set_proto_from_seed(stage_key.task_number, 4)
             );
         }
     }
@@ -209,7 +212,7 @@ mod tests {
     async fn test_metrics_collecting_stream_error_missing_stage_key() {
         let task_metrics_with_no_stage_key = TaskMetrics {
             stage_key: None,
-            metrics: vec![make_test_metrics_set_proto_from_seed(1)],
+            metrics: vec![make_test_metrics_set_proto_from_seed(1, 4)],
         };
 
         let invalid_app_metadata = FlightAppMetadata {

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(crate) use task_metrics_rewriter::StageMetricsRewriter;
+pub 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 {