Support distributed EXPLAIN ANALYZE

src/distributed_planner/distributed_query_planner.rs

@@ -7,15 +7,18 @@ use crate::distributed_planner::insert_broadcast::insert_broadcast_execs;
 use crate::distributed_planner::partial_reduce_below_network_shuffles::partial_reduce_below_network_shuffles;
 use crate::distributed_planner::prepare_network_boundaries::prepare_network_boundaries;
 use crate::distributed_planner::push_fetch_into_network_coalesce::push_fetch_into_network_coalesce;
+use crate::explain_analyze::DistributedAnalyzeExec;
 use crate::{DistributedConfig, DistributedExec, NetworkBoundaryExt, TaskEstimator};
 use async_trait::async_trait;
 use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::execution::SessionState;
 use datafusion::execution::context::QueryPlanner;
 use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::analyze::AnalyzeExec;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
+use futures::future::BoxFuture;
 use std::sync::Arc;
 
 /// Transforms a single-node physical plan into a distributed plan by injecting network
@@ -67,10 +70,34 @@ impl QueryPlanner for DistributedQueryPlanner {
             }
         };
 
+        create_distributed_plan(original_plan, session_state).await
+    }
+}
+
+/// Applies distributed planning to an already-created physical plan.
+///
+/// Inputs to [`AnalyzeExec`] are planned recursively. When one becomes distributed, the native
+/// node is replaced by [`DistributedAnalyzeExec`]. The future is boxed because recursive `async
+/// fn`s would otherwise have an infinitely-sized return type.
+fn create_distributed_plan(
+    original_plan: Arc<dyn ExecutionPlan>,
+    session_state: &SessionState,
+) -> BoxFuture<'_, datafusion::common::Result<Arc<dyn ExecutionPlan>>> {
+    Box::pin(async move {
         if original_plan.is::<DistributedExec>() {
             return Ok(original_plan);
         }
 
+        if let Some(analyze) = original_plan.downcast_ref::<AnalyzeExec>() {
+            // Recursively distribute the query being analyzed. Replace `AnalyzeExec` only when its
+            // input became distributed; otherwise retain DataFusion's native implementation.
+            let input = create_distributed_plan(Arc::clone(analyze.input()), session_state).await?;
+            return match input.is::<DistributedExec>() {
+                true => Ok(Arc::new(DistributedAnalyzeExec::new(analyze, input))),
+                false => Arc::new(analyze.clone()).with_new_children(vec![input]),
+            };
+        }
+
         let session_cfg = session_state.config();
         let cfg = session_cfg.options();
         let d_cfg = DistributedConfig::from_config_options(cfg)?;
@@ -132,7 +159,7 @@ impl QueryPlanner for DistributedQueryPlanner {
         Ok(Arc::new(
             DistributedExec::new(plan).with_metrics_collection(d_cfg.collect_metrics),
         ))
-    }
+    })
 }
 
 #[cfg(test)]

src/explain_analyze.rs

@@ -0,0 +1,144 @@
+use crate::DistributedMetricsFormat;
+use crate::common::require_one_child;
+use crate::stage::explain_analyze;
+use datafusion::arrow::array::StringBuilder;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::instant::Instant;
+use datafusion::common::{DataFusionError, Result, assert_eq_or_internal_err};
+use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_plan::analyze::AnalyzeExec;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties,
+};
+use futures::{StreamExt, stream};
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+/// Distributed counterpart of DataFusion's [`AnalyzeExec`].
+///
+/// It drains a [`crate::DistributedExec`], waits for metrics from every worker, and returns the
+/// distributed plan with those metrics as the `EXPLAIN ANALYZE` result.
+#[derive(Debug)]
+pub(crate) struct DistributedAnalyzeExec {
+    input: Arc<dyn ExecutionPlan>,
+    verbose: bool,
+    properties: Arc<PlanProperties>,
+}
+
+impl DistributedAnalyzeExec {
+    pub(crate) fn new(analyze_exec: &AnalyzeExec, input: Arc<dyn ExecutionPlan>) -> Self {
+        Self {
+            input,
+            verbose: analyze_exec.verbose(),
+            properties: Arc::clone(analyze_exec.properties()),
+        }
+    }
+}
+
+impl DisplayAs for DistributedAnalyzeExec {
+    fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "DistributedAnalyzeExec verbose={}", self.verbose)
+    }
+}
+
+impl ExecutionPlan for DistributedAnalyzeExec {
+    fn name(&self) -> &str {
+        "DistributedAnalyzeExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.properties
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn required_input_distribution(&self) -> Vec<Distribution> {
+        vec![Distribution::UnspecifiedDistribution]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(Arc::new(Self {
+            input: require_one_child(&children)?,
+            verbose: self.verbose,
+            properties: Arc::clone(&self.properties),
+        }))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        assert_eq_or_internal_err!(
+            partition,
+            0,
+            "DistributedAnalyzeExec invalid partition. Expected 0, got {partition}"
+        );
+
+        let mut input_stream = self.input.execute(0, context)?;
+        let input = Arc::clone(&self.input);
+        let schema = self.schema();
+        let verbose = self.verbose;
+        let output = async move {
+            let start = Instant::now();
+            let mut total_rows = 0;
+            while let Some(batch) = input_stream.next().await.transpose()? {
+                total_rows += batch.num_rows();
+            }
+
+            let plan =
+                explain_analyze(Arc::clone(&input), DistributedMetricsFormat::Aggregated).await?;
+            let full_plan = match verbose {
+                true => Some(explain_analyze(input, DistributedMetricsFormat::PerTask).await?),
+                false => None,
+            };
+            create_output_batch(total_rows, start.elapsed(), plan, full_plan, schema)
+        };
+
+        Ok(Box::pin(RecordBatchStreamAdapter::new(
+            self.schema(),
+            stream::once(output),
+        )))
+    }
+}
+
+fn create_output_batch(
+    total_rows: usize,
+    duration: std::time::Duration,
+    plan: String,
+    full_plan: Option<String>,
+    schema: SchemaRef,
+) -> Result<RecordBatch> {
+    let mut type_builder = StringBuilder::with_capacity(1, 1024);
+    let mut plan_builder = StringBuilder::with_capacity(1, 1024);
+
+    type_builder.append_value("Plan with Metrics");
+    plan_builder.append_value(&plan);
+
+    if let Some(full_plan) = full_plan {
+        type_builder.append_value("Plan with Full Metrics");
+        plan_builder.append_value(full_plan);
+
+        type_builder.append_value("Output Rows");
+        plan_builder.append_value(total_rows.to_string());
+
+        type_builder.append_value("Duration");
+        plan_builder.append_value(format!("{duration:?}"));
+    }
+
+    RecordBatch::try_new(
+        schema,
+        vec![
+            Arc::new(type_builder.finish()),
+            Arc::new(plan_builder.finish()),
+        ],
+    )
+    .map_err(DataFusionError::from)
+}

src/lib.rs

@@ -7,6 +7,7 @@ mod coordinator;
 mod distributed_ext;
 mod distributed_planner;
 mod execution_plans;
+mod explain_analyze;
 mod metrics;
 mod passthrough_headers;
 mod protocol;

tests/explain_analyze.rs

@@ -0,0 +1,101 @@
+#[cfg(all(feature = "integration", test))]
+mod tests {
+    use datafusion::arrow::array::{Array, StringArray};
+    use datafusion::arrow::util::pretty::pretty_format_batches;
+    use datafusion::common::{Result, assert_contains, assert_not_contains};
+    use datafusion::physical_plan::execute_stream;
+    use datafusion_distributed::test_utils::localhost::start_localhost_context;
+    use datafusion_distributed::test_utils::parquet::register_parquet_tables;
+    use datafusion_distributed::{DefaultSessionBuilder, DistributedExt};
+    use futures::TryStreamExt;
+    use test_case::test_case;
+
+    #[test_case(false ; "static_task_count")]
+    #[test_case(true ; "dynamic_task_count")]
+    #[tokio::test]
+    async fn explain_analyze_displays_distributed_metrics(
+        dynamic_task_count: bool,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let (mut ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await;
+        ctx.set_distributed_dynamic_task_count(dynamic_task_count)?;
+        register_parquet_tables(&ctx).await?;
+
+        let plan = ctx
+            .sql(
+                r#"EXPLAIN ANALYZE
+                   SELECT count(*), "RainToday"
+                   FROM weather
+                   GROUP BY "RainToday"
+                   ORDER BY count(*)"#,
+            )
+            .await?
+            .create_physical_plan()
+            .await?;
+        assert_eq!(plan.name(), "DistributedAnalyzeExec");
+        assert_eq!(plan.children()[0].name(), "DistributedExec");
+
+        let batches = execute_stream(plan, ctx.task_ctx())?
+            .try_collect::<Vec<_>>()
+            .await?;
+        let formatted = pretty_format_batches(&batches)?.to_string();
+        println!("{formatted}");
+
+        assert_contains!(&formatted, "Plan with Metrics");
+        assert_contains!(&formatted, "DistributedExec");
+        assert_contains!(&formatted, "NetworkShuffleExec");
+        assert_contains!(&formatted, "metrics=[output_rows=");
+        assert_not_contains!(&formatted, "metrics=[output_rows={");
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn explain_analyze_verbose_displays_additional_metrics()
+    -> Result<(), Box<dyn std::error::Error>> {
+        let (ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await;
+        register_parquet_tables(&ctx).await?;
+
+        let plan = ctx
+            .sql(
+                r#"EXPLAIN ANALYZE VERBOSE
+                   SELECT count(*), "RainToday"
+                   FROM weather
+                   GROUP BY "RainToday"
+                   ORDER BY count(*)"#,
+            )
+            .await?
+            .create_physical_plan()
+            .await?;
+        assert_eq!(plan.name(), "DistributedAnalyzeExec");
+        assert_eq!(plan.children()[0].name(), "DistributedExec");
+
+        let batches = execute_stream(plan, ctx.task_ctx())?
+            .try_collect::<Vec<_>>()
+            .await?;
+        let formatted = pretty_format_batches(&batches)?.to_string();
+        println!("{formatted}");
+
+        assert_contains!(&formatted, "Plan with Full Metrics");
+        assert_contains!(&formatted, "Output Rows");
+        assert_contains!(&formatted, "Duration");
+        assert_contains!(&formatted, "metrics=[output_rows=");
+
+        let plan_types = batches[0].column(0).as_any().downcast_ref::<StringArray>();
+        let plans = batches[0].column(1).as_any().downcast_ref::<StringArray>();
+        let (plan_types, plans) = match (plan_types, plans) {
+            (Some(plan_types), Some(plans)) => (plan_types, plans),
+            _ => panic!("EXPLAIN ANALYZE columns should be strings"),
+        };
+        let plan_for = |plan_type| {
+            let index = (0..plan_types.len())
+                .find(|index| plan_types.value(*index) == plan_type)
+                .unwrap_or_else(|| panic!("verbose output should contain a {plan_type} entry"));
+            plans.value(index)
+        };
+        assert_not_contains!(plan_for("Plan with Metrics"), "metrics=[output_rows={");
+        assert_contains!(plan_for("Plan with Full Metrics"), "metrics=[output_rows={");
+        assert_eq!(plan_for("Output Rows"), "2");
+
+        Ok(())
+    }
+}