implement distributed EXPLAIN ANALYZE

This change adds support for displaying a distributed EXPLAIN ANALYZE output. It updates the TPCH
validation tests to assert the EXPLAIN ANALYZE output for each query.

Implemenation notes:
- Adds `src/explain.rs` to stores the main entrypoint to rendering the output string
  - I left a TODO about pushing some of the work to `DistributedExec` or a new node type
- Adds a `Option<DisplayCtx>` field to `DistributedExec` to contain extra information for display purposes.
  - We use this to smuggle the information into `display_plan_ascii` because its only relevant in the
    distributed case
  - Then, at display-time, when displaying a task, we re-write each task plan to use the metrics
    from the `DislplayCtx`

Informs: #123

Remaning work:
- disable any metrics propagation if not running EXPLAIN ANALYZE as it adds extra overhead
- consider refactoring explain.rs
- graphviz
- add docs + exalidraw to explain the metrics protocol

Author: jayshrivastava

src/common/ttl_map.rs

@@ -289,7 +289,7 @@ where
 mod tests {
     use super::*;
     use std::sync::atomic::Ordering;
-    use tokio::time::{Duration, sleep};
+    use tokio::time::Duration;
 
     #[tokio::test]
     async fn test_basic_insert_and_get() {

src/execution_plans/distributed.rs

@@ -1,20 +1,24 @@
 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::protobuf::{DistributedCodec, StageKey};
+use crate::stage::DisplayCtx;
 use crate::{ExecutionTask, Stage};
+use bytes::Bytes;
 use datafusion::common::exec_err;
 use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use rand::Rng;
 use std::any::Any;
 use std::fmt::Formatter;
 use std::sync::Arc;
+use std::sync::Mutex;
 use url::Url;
 
-/// [ExecutionPlan] that executes the inner plan in distributed mode.
+/// [ExecutionPan] that executes the inner plan in distributed mode.
 /// Before executing it, two modifications are lazily performed on the plan:
 /// 1. Assigns worker URLs to all the stages. A random set of URLs are sampled from the
 ///    channel resolver and assigned to each task in each stage.
@@ -23,11 +27,47 @@ 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>,
 }
 
 impl DistributedExec {
     pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
-        Self { plan }
+        Self {
+            plan,
+            prepared_plan: Arc::new(Mutex::new(None)),
+            display_ctx: None,
+        }
+    }
+
+    /// 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.
+    pub(crate) fn pepared_plan(&self) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        self.prepared_plan
+            .lock()
+            .map_err(|e| DataFusionError::Internal(format!("Failed to lock prepared plan: {}", e)))?
+            .clone()
+            .ok_or(DataFusionError::Internal(
+                "No prepared plan found. Was execute() called?".to_string(),
+            ))
     }
 
     fn prepare_plan(
@@ -99,6 +139,8 @@ impl ExecutionPlan for DistributedExec {
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         Ok(Arc::new(DistributedExec {
             plan: require_one_child(&children)?,
+            prepared_plan: self.prepared_plan.clone(),
+            display_ctx: self.display_ctx.clone(),
         }))
     }
 
@@ -120,7 +162,14 @@ impl ExecutionPlan for DistributedExec {
         let channel_resolver = get_distributed_channel_resolver(context.session_config())?;
         let codec = DistributedCodec::new_combined_with_user(context.session_config());
 
-        let plan = self.prepare_plan(&channel_resolver.get_urls()?, &codec)?;
-        plan.execute(partition, context)
+        let prepared = self.prepare_plan(&channel_resolver.get_urls()?, &codec)?;
+        {
+            let mut guard = self.prepared_plan.lock().map_err(|e| {
+                DataFusionError::Internal(format!("Failed to lock prepared plan: {}", e))
+            })?;
+            *guard = Some(prepared.clone());
+        }
+
+        prepared.execute(partition, context)
     }
 }

src/explain.rs

@@ -0,0 +1,40 @@
+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 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)))
+        }
+    }
+}

src/flight_service/do_get.rs

@@ -143,8 +143,9 @@ impl ArrowFlightEndpoint {
         let stream = map_last_stream(stream, move |last| {
             if num_partitions_remaining.fetch_sub(1, Ordering::SeqCst) == 1 {
                 task_data_entries.remove(key.clone());
+                return last.and_then(|el| collect_and_create_metrics_flight_data(key, plan, el));
             }
-            last.and_then(|el| collect_and_create_metrics_flight_data(key, plan, el))
+            last
         });
 
         Ok(Response::new(Box::pin(stream.map_err(|err| match err {

src/lib.rs

@@ -6,6 +6,7 @@ mod config_extension_ext;
 mod distributed_ext;
 mod distributed_physical_optimizer_rule;
 mod execution_plans;
+mod explain;
 mod flight_service;
 mod metrics;
 mod stage;
@@ -20,6 +21,7 @@ pub use distributed_physical_optimizer_rule::{
     DistributedPhysicalOptimizerRule, NetworkBoundaryExt,
 };
 pub use execution_plans::{NetworkCoalesceExec, NetworkShuffleExec, PartitionIsolatorExec};
+pub use explain::explain_analyze;
 pub use flight_service::{
     ArrowFlightEndpoint, DefaultSessionBuilder, DistributedSessionBuilder,
     DistributedSessionBuilderContext, MappedDistributedSessionBuilder,

src/metrics/mod.rs

@@ -3,4 +3,5 @@ pub(crate) mod proto;
 mod task_metrics_collector;
 mod task_metrics_rewriter;
 pub(crate) use metrics_collecting_stream::MetricsCollectingStream;
-pub(crate) use task_metrics_collector::TaskMetricsCollector;
+pub(crate) use task_metrics_collector::{MetricsCollectorResult, TaskMetricsCollector};
+pub(crate) use task_metrics_rewriter::TaskMetricsRewriter;

src/metrics/task_metrics_rewriter.rs

@@ -32,14 +32,12 @@ pub struct TaskMetricsRewriter {
 
 impl TaskMetricsRewriter {
     /// Create a new TaskMetricsRewriter. The provided metrics will be used to enrich the plan.
-    #[allow(dead_code)]
     pub fn new(metrics: Vec<MetricsSetProto>) -> Self {
         Self { metrics, idx: 0 }
     }
 
     /// enrich_task_with_metrics rewrites the plan by wrapping nodes. If the length of the provided metrics set vec does not
     /// match the number of nodes in the plan, an error will be returned.
-    #[allow(dead_code)]
     pub fn enrich_task_with_metrics(
         mut self,
         plan: Arc<dyn ExecutionPlan>,