This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new e1e39da10 feat: add DistributedQueryExec::explain_executed_plan and
show logical plan in display (#2042)
e1e39da10 is described below
commit e1e39da103ad0d3bb08e58e5dd42f14d98f2f212
Author: Andy Grove <[email protected]>
AuthorDate: Sat Aug 8 00:42:42 2026 -0600
feat: add DistributedQueryExec::explain_executed_plan and show logical plan
in display (#2042)
---
ballista/client/tests/context_checks.rs | 49 +++++++++
.../execution_plans/distributed_explain_analyze.rs | 117 ++++++++++++++++-----
.../core/src/execution_plans/distributed_query.rs | 91 +++++++++++++++-
3 files changed, 230 insertions(+), 27 deletions(-)
diff --git a/ballista/client/tests/context_checks.rs
b/ballista/client/tests/context_checks.rs
index 5d9cf83ea..fc3dc2243 100644
--- a/ballista/client/tests/context_checks.rs
+++ b/ballista/client/tests/context_checks.rs
@@ -24,6 +24,8 @@ mod supported {
standalone_context_with_state,
};
use ballista_core::config::BallistaConfig;
+ use ballista_core::execution_plans::DistributedQueryExec;
+ use datafusion_proto::protobuf::LogicalPlanNode;
use datafusion::arrow::array::StringArray;
use datafusion::arrow::record_batch::RecordBatch;
@@ -1204,6 +1206,53 @@ mod supported {
Ok(())
}
+ #[rstest]
+ #[case::standalone(standalone_context())]
+ #[case::remote(remote_context())]
+ #[tokio::test]
+ async fn should_explain_executed_plan(
+ #[future(awt)]
+ #[case]
+ ctx: SessionContext,
+ ) -> datafusion::error::Result<()> {
+ let df = ctx
+ .sql(
+ "select count(*), id from (select unnest([1,2,3,4,5]) as id)
group by id",
+ )
+ .await?;
+
+ let plan = df.create_physical_plan().await?;
+
+ // Execute so the scheduler assigns a job id and completes the stages.
+ let _ = collect(plan.clone(), ctx.task_ctx()).await?;
+
+ let dqe = plan
+ .downcast_ref::<DistributedQueryExec<LogicalPlanNode>>()
+ .expect("top operator should be a
DistributedQueryExec<LogicalPlanNode>");
+
+ let config = ctx.copied_config();
+
+ // Plan without metrics.
+ let plan_text = dqe.explain_executed_plan(&config, false).await?;
+ assert!(
+ plan_text.contains("SuccessfulStage"),
+ "expected per-stage plan, got: {plan_text}"
+ );
+ assert!(
+ !plan_text.contains("metrics=["),
+ "expected no metrics, got: {plan_text}"
+ );
+
+ // Same plan with metrics.
+ let plan_with_metrics = dqe.explain_executed_plan(&config,
true).await?;
+ assert!(
+ plan_with_metrics.contains("metrics=["),
+ "expected metrics, got: {plan_with_metrics}"
+ );
+
+ Ok(())
+ }
+
#[rstest]
#[case::standalone(standalone_context())]
#[case::remote(remote_context())]
diff --git a/ballista/core/src/execution_plans/distributed_explain_analyze.rs
b/ballista/core/src/execution_plans/distributed_explain_analyze.rs
index c562df885..116165280 100644
--- a/ballista/core/src/execution_plans/distributed_explain_analyze.rs
+++ b/ballista/core/src/execution_plans/distributed_explain_analyze.rs
@@ -196,7 +196,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for
DistributedExplainAnalyzeExec
}
}
-async fn fetch_job_metrics(
+pub(crate) async fn fetch_job_metrics(
scheduler_url: &str,
job_id: &JobId,
session_config: datafusion::prelude::SessionConfig,
@@ -240,13 +240,15 @@ async fn fetch_job_metrics(
.map_err(|e| DataFusionError::Execution(format!("{e:?}")))
}
-fn format_metrics_as_record_batch(
+/// Render the per-stage executed plan carried in `job_metrics` as a single
+/// string. When `with_metrics` is true, each operator line is suffixed with
its
+/// aggregated `metrics=[...]`; when false, only the operator description is
+/// shown.
+pub(crate) fn format_job_plan(
job_metrics: &GetJobMetricsResult,
- _job_id: &JobId,
- schema: SchemaRef,
- _verbose: bool,
-) -> Result<RecordBatch> {
- let plan = job_metrics
+ with_metrics: bool,
+) -> Result<String> {
+ job_metrics
.stages
.iter()
.map(|stage| {
@@ -257,28 +259,41 @@ fn format_metrics_as_record_batch(
));
for operator in &stage.operators {
- let metrics_set:
datafusion::physical_plan::metrics::MetricsSet =
- protobuf::OperatorMetricsSet {
- metrics: operator.metrics.clone(),
- }
- .try_into()
- .map_err(|e|
DataFusionError::Execution(format!("{e:?}")))?;
- let metrics = metrics_set
- .aggregate_by_name()
- .sorted_for_display()
- .timestamps_removed()
- .to_string();
let indent = " ".repeat(operator.depth as usize);
- stage_plan.push(format!(
- "{indent}{}, metrics=[{metrics}]",
- operator.operator_desc
- ));
+ if with_metrics {
+ let metrics_set:
datafusion::physical_plan::metrics::MetricsSet =
+ protobuf::OperatorMetricsSet {
+ metrics: operator.metrics.clone(),
+ }
+ .try_into()
+ .map_err(|e|
DataFusionError::Execution(format!("{e:?}")))?;
+ let metrics = metrics_set
+ .aggregate_by_name()
+ .sorted_for_display()
+ .timestamps_removed()
+ .to_string();
+ stage_plan.push(format!(
+ "{indent}{}, metrics=[{metrics}]",
+ operator.operator_desc
+ ));
+ } else {
+ stage_plan.push(format!("{indent}{}",
operator.operator_desc));
+ }
}
Ok(stage_plan.join("\n"))
})
- .collect::<Result<Vec<_>>>()?
- .join("\n\n");
+ .collect::<Result<Vec<_>>>()
+ .map(|stages| stages.join("\n\n"))
+}
+
+fn format_metrics_as_record_batch(
+ job_metrics: &GetJobMetricsResult,
+ _job_id: &JobId,
+ schema: SchemaRef,
+ _verbose: bool,
+) -> Result<RecordBatch> {
+ let plan = format_job_plan(job_metrics, true)?;
let mut type_builder = StringBuilder::with_capacity(1, "Plan with
Metrics".len());
let mut plan_builder = StringBuilder::with_capacity(1, plan.len());
@@ -298,7 +313,7 @@ fn format_metrics_as_record_batch(
#[cfg(test)]
mod tests {
- use super::format_metrics_as_record_batch;
+ use super::{format_job_plan, format_metrics_as_record_batch};
use datafusion::arrow::array::StringArray;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
@@ -308,6 +323,58 @@ mod tests {
operator_metric,
};
+ #[test]
+ fn test_format_job_plan_without_metrics() {
+ let response = GetJobMetricsResult {
+ stages: vec![
+ JobStageMetrics {
+ stage_id: 0,
+ partitions: 1,
+ operators: vec![OperatorWithMetrics {
+ depth: 0,
+ operator_type: "ProjectionExec".to_string(),
+ operator_desc: "ProjectionExec: expr=[a]".to_string(),
+ metrics: vec![
+ OperatorMetric {
+ metric:
Some(operator_metric::Metric::OutputRows(4)),
+ },
+ OperatorMetric {
+ metric:
Some(operator_metric::Metric::ElapseTime(
+ 15_000_000,
+ )),
+ },
+ ],
+ }],
+ },
+ JobStageMetrics {
+ stage_id: 1,
+ partitions: 2,
+ operators: vec![OperatorWithMetrics {
+ depth: 0,
+ operator_type: "FilterExec".to_string(),
+ operator_desc: "FilterExec: a@0 > 1".to_string(),
+ metrics: vec![OperatorMetric {
+ metric:
Some(operator_metric::Metric::OutputRows(2)),
+ }],
+ }],
+ },
+ ],
+ };
+
+ let plan = format_job_plan(&response, false).unwrap();
+
+ let expected = [
+ "=========SuccessfulStage[stage_id=0, partitions=1]=========",
+ "ProjectionExec: expr=[a]",
+ "",
+ "=========SuccessfulStage[stage_id=1, partitions=2]=========",
+ "FilterExec: a@0 > 1",
+ ]
+ .join("\n");
+ assert_eq!(plan, expected);
+ assert!(!plan.contains("metrics=["));
+ }
+
#[test]
fn test_format_metrics_as_record_batch() {
let schema = Arc::new(Schema::new(vec![
diff --git a/ballista/core/src/execution_plans/distributed_query.rs
b/ballista/core/src/execution_plans/distributed_query.rs
index 20f3234cf..a27abd0a1 100644
--- a/ballista/core/src/execution_plans/distributed_query.rs
+++ b/ballista/core/src/execution_plans/distributed_query.rs
@@ -141,6 +141,45 @@ impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
self.job_id.lock().clone()
}
+ /// Fetch and render the physical plan that executed on the cluster for
this
+ /// query, one section per query stage.
+ ///
+ /// This must be called after the query has run (i.e. after
+ /// [`ExecutionPlan::execute`]/`collect`), because it relies on the
+ /// scheduler-assigned job id and the plan is only available once the
stages
+ /// have completed. Returns an error if the query has not been executed
yet.
+ ///
+ /// Only stages that completed successfully are included, inherited from
the
+ /// scheduler's `get_job_metrics`.
+ ///
+ /// When `with_metrics` is true the output matches `EXPLAIN ANALYZE`; when
+ /// false the `metrics=[...]` suffixes are omitted.
+ pub async fn explain_executed_plan(
+ &self,
+ session_config: &SessionConfig,
+ with_metrics: bool,
+ ) -> Result<String> {
+ let job_id = self.job_id().ok_or_else(|| {
+ DataFusionError::Execution(
+ "Cannot explain executed plan: query has not been executed yet"
+ .to_string(),
+ )
+ })?;
+
+ let job_metrics =
+
crate::execution_plans::distributed_explain_analyze::fetch_job_metrics(
+ &self.scheduler_url,
+ &job_id,
+ session_config.clone(),
+ )
+ .await?;
+
+ crate::execution_plans::distributed_explain_analyze::format_job_plan(
+ &job_metrics,
+ with_metrics,
+ )
+ }
+
fn compute_properties(schema: SchemaRef) -> PlanProperties {
PlanProperties::new(
EquivalenceProperties::new(schema),
@@ -159,11 +198,12 @@ impl<T: 'static + AsLogicalPlan> DisplayAs for
DistributedQueryExec<T> {
) -> std::fmt::Result {
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
- write!(
+ writeln!(
f,
"DistributedQueryExec: scheduler_url={}",
self.scheduler_url
- )
+ )?;
+ write!(f, "logical_plan:\n{}", self.plan.display_indent())
}
DisplayFormatType::TreeRender => {
writeln!(f, "scheduler_url={}", self.scheduler_url)
@@ -855,6 +895,8 @@ mod test {
use crate::serde::protobuf::get_job_status_result::FlightProxy;
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::physical_plan::displayable;
+ use datafusion::prelude::SessionConfig;
use datafusion_proto::protobuf::LogicalPlanNode;
use std::sync::Arc;
@@ -930,4 +972,49 @@ mod test {
assert_eq!(new_exec.job_id(), Some(JobId::new("job-123")));
}
+
+ #[test]
+ fn test_display_includes_logical_plan() {
+ let exec = Arc::new(DistributedQueryExec::<LogicalPlanNode>::new(
+ "http://scheduler:50050".to_string(),
+ BallistaConfig::default(),
+ LogicalPlan::default(),
+ "session".to_string(),
+ ));
+
+ let rendered = displayable(exec.as_ref()).indent(false).to_string();
+
+ assert!(
+ rendered.contains("scheduler_url=http://scheduler:50050";),
+ "missing scheduler_url line: {rendered}"
+ );
+ assert!(
+ rendered.contains("logical_plan:"),
+ "missing logical_plan section: {rendered}"
+ );
+ // LogicalPlan::default() renders as EmptyRelation
+ assert!(
+ rendered.contains("EmptyRelation"),
+ "missing rendered logical plan: {rendered}"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_explain_executed_plan_errors_before_execution() {
+ let exec = DistributedQueryExec::<LogicalPlanNode>::new(
+ "http://scheduler:50050".to_string(),
+ BallistaConfig::default(),
+ LogicalPlan::default(),
+ "session".to_string(),
+ );
+ // job_id is None because the query has not been executed
+ let err = exec
+ .explain_executed_plan(&SessionConfig::new(), false)
+ .await
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("has not been executed"),
+ "unexpected error: {err}"
+ );
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
