asdf

Author: jayshrivastava

src/execution_plans/distributed.rs

@@ -1,20 +1,26 @@
 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::f32::consts::E;
+use std::fmt::{format, Formatter};
 use std::sync::Arc;
+use std::sync::Mutex;
 use url::Url;
+use std::sync::OnceLock;
 
-/// [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 +29,34 @@ 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 as u64,
+            task_number: 0,
+        }
+    }
+
+    pub(crate) fn with_display_ctx(&self, display_ctx: DisplayCtx) -> Self {
+        Self { display_ctx: Some(display_ctx), ..self.clone() }
+    }
+
+    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 +128,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 +151,17 @@ 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)?;
+        let plan = {
+            let mut guard = self.prepared_plan.lock().map_err(|e| DataFusionError::Internal(format!("Failed to lock prepared plan: {}", e)))?;
+            match guard.clone() {
+                Some(plan) => plan,
+                None => {
+                    let prepared = self.prepare_plan(&channel_resolver.get_urls()?, &codec)?;
+                    *guard = Some(prepared.clone());
+                    prepared
+                }
+            }
+        };
         plan.execute(partition, context)
     }
 }

src/execution_plans/metrics.rs

@@ -27,6 +27,10 @@ impl MetricsWrapperExec {
             children: None,
         }
     }
+
+    pub fn inner(&self) -> &Arc<dyn ExecutionPlan> {
+        &self.inner
+    }
 }
 
 /// MetricsWrapperExec is invisible during display.

src/explain.rs

@@ -0,0 +1,47 @@
+use crate::stage::{DisplayCtx};
+use crate::metrics::MetricsCollectorResult;
+use crate::metrics::TaskMetricsCollector;
+use crate::metrics::proto::MetricsSetProto;
+use crate::metrics::proto::df_metrics_set_to_proto;
+use crate::{display_plan_ascii, NetworkBoundaryExt};
+use datafusion::common::tree_node::Transformed;
+use datafusion::common::tree_node::{TreeNode, TreeNodeRewriter};
+use datafusion::error::DataFusionError;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::display::DisplayableExecutionPlan;
+use std::sync::Arc;
+use crate::execution_plans::DistributedExec;
+use bytes::Bytes;
+use crate::protobuf::StageKey;
+
+
+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(display_plan_ascii(executed.as_ref())),
+        Some(dist_exec) => {
+            // If the plan was distributed, collect metrics from the coordinating stage exec.
+            let MetricsCollectorResult {
+                task_metrics,
+                mut input_task_metrics,
+            } = TaskMetricsCollector::new().collect(
+                dist_exec.pepared_plan()?
+            )?;
+            println!("task_metrics: {}", task_metrics.len());
+            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 keys = input_task_metrics.iter().map(|(k, v)| (k, v.len())).collect::<Vec<_>>();
+            println!("items: {:#?}", keys);
+        
+
+            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

@@ -9,6 +9,7 @@ mod execution_plans;
 mod flight_service;
 mod metrics;
 mod stage;
+mod explain;
 
 mod protobuf;
 #[cfg(any(feature = "integration", test))]
@@ -28,3 +29,4 @@ pub use flight_service::{
 pub use stage::{
     DistributedTaskContext, ExecutionTask, Stage, display_plan_ascii, display_plan_graphviz,
 };
+pub use explain::explain_analyze;

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;