diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index 4b6fa46d..3e6ddca5 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/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, + session_state: &SessionState, +) -> BoxFuture<'_, datafusion::common::Result>> { + Box::pin(async move { if original_plan.is::() { return Ok(original_plan); } + if let Some(analyze) = original_plan.downcast_ref::() { + // 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::() { + 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)] diff --git a/src/explain_analyze.rs b/src/explain_analyze.rs new file mode 100644 index 00000000..7e321244 --- /dev/null +++ b/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, + verbose: bool, + properties: Arc, +} + +impl DistributedAnalyzeExec { + pub(crate) fn new(analyze_exec: &AnalyzeExec, input: Arc) -> 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 { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: require_one_child(&children)?, + verbose: self.verbose, + properties: Arc::clone(&self.properties), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + 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, + schema: SchemaRef, +) -> Result { + 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) +} diff --git a/src/lib.rs b/src/lib.rs index b98856d1..ae113793 100644 --- a/src/lib.rs +++ b/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; diff --git a/tests/explain_analyze.rs b/tests/explain_analyze.rs new file mode 100644 index 00000000..5980f17e --- /dev/null +++ b/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> { + 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::>() + .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> { + 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::>() + .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::(); + let plans = batches[0].column(1).as_any().downcast_ref::(); + 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(()) + } +}