Getting results and explain analyze plan at the same time

[adriangb]

I would like to run a query once and get back both the result data and the explain analyze plan w/ metrics.
As far as I can tell I have to chose one or the other ATM, but it seems like it should be possible to get both in principle.
Is there an API I'm missing or is this not currently possible? I've looked at DataFrame.explain and LogicalPlanBuilder.explain.

[jonahgao]

I think this is not possible in DataFusion currently. Query results and explain results have different schemas and number of rows, and outputting both at the same time seems to conflict with the current design.

[adriangb]

We got it working, but it was not easy. I think it's possible and a cleaner solution can be arrived at.

use std::{collections::BTreeMap, fmt, fmt::Formatter, pin::Pin, sync::Arc, task::Poll};

use arrow::array::RecordBatch;
use arrow_schema::SchemaRef;
use datafusion::{
    error::Result as DataFusionResult,
    execution::{RecordBatchStream, SendableRecordBatchStream, SessionState},
    logical_expr::LogicalPlan,
    physical_plan::{
        accept, execute_stream, metrics::MetricsSet, DisplayFormatType, ExecutionPlan, ExecutionPlanVisitor,
    },
    prelude::DataFrame,
};
use futures::Stream;
use pin_project_lite::pin_project;
use serde::{Deserialize, Serialize};
use tracing::Span;

pin_project! {
    /// This emits a tracing event with the execution plan after completion.
    /// Note: requires the span has a `plan` field
    pub struct TracingRecordBatchStream {
        #[pin]
        inner: SendableRecordBatchStream,
        plan: Arc<dyn ExecutionPlan>,
        span: Span
    }

    impl PinnedDrop for TracingRecordBatchStream {

        /// We only record against the span on `drop`
        /// That way even if the request is cancelled we should still have *some* info
        fn drop(this: Pin<&mut Self>) {

            let mut json_plan = JsonExecutionPlanVisitor::new();

            if accept(&*this.plan, &mut json_plan).is_ok() {

                // We pop the stack, which should *alway* have a root node
                // Then we add it as a json string to the span
                if let Some(json_string) = json_plan.stack.pop().and_then(|val| serde_json::to_string(&val).ok()) {
                    this.span.record("plan", json_string);
                }

            }

        }
    }
}

impl RecordBatchStream for TracingRecordBatchStream {
    fn schema(&self) -> SchemaRef {
        self.inner.schema()
    }
}

impl TracingRecordBatchStream {
    pub async fn execute(span: Span, session_state: SessionState, plan: LogicalPlan) -> DataFusionResult<Self> {
        let df: DataFrame = DataFrame::new(session_state, plan);

        let task_ctx = Arc::new(df.task_ctx());
        let plan = df.create_physical_plan().await?;

        let inner = execute_stream(plan.clone(), task_ctx)?;

        Ok(Self { inner, plan, span })
    }
}

impl Stream for TracingRecordBatchStream {
    type Item = DataFusionResult<RecordBatch>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();

        this.inner.poll_next(cx)
    }
}

#[derive(Debug, Serialize, Deserialize)]
/// This is our JSON structure representing things about an Execution Plan
struct JsonExecutionPlan {
    /// The short name given to the Execution Plan, such as `FilterExec` or `ParquetExec` etc..
    name: String,
    /// The longer description of the Execution Plan
    /// Currently all of the native datafusion plans spit out a *lot* of information that doesn't fit into metrics/statistics
    /// I.e, `ParquetExec` will include all the parquet files it's going to read
    /// Unfortunately it's not structured, it is returned as a string
    desc: String,
    #[serde(skip_serializing_if = "BTreeMap::is_empty")]

    /// We always aggregate metrics, so there won't be any metric name collisions
    metrics: BTreeMap<String, String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]

    /// Any children of the Execution Plan
    children: Vec<JsonExecutionPlan>,
}

/// This is used to build up the `JsonExecutionPlan` nested hierarchy
/// Uses the visitor pattern to do so
struct JsonExecutionPlanVisitor {
    stack: Vec<JsonExecutionPlan>,
}

impl JsonExecutionPlanVisitor {
    pub fn new() -> Self {
        JsonExecutionPlanVisitor { stack: Vec::new() }
    }
}

/// This is used to pull out the bigger description.
/// Uses the `DisplayAs` trait but needs a `fmt::Formatter` which we can't construct directly.
/// Instead we impl `Display` to get the values out
struct PlanDescFormat<'a> {
    plan: &'a dyn ExecutionPlan,
}

impl<'a> fmt::Display for PlanDescFormat<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.plan.fmt_as(DisplayFormatType::Verbose, f)
    }
}

/// Converts a `MetricSet` to a BTreeMap (sorted alphabetically) for easier display
fn metrics_to_btree(metrics: &MetricsSet) -> BTreeMap<String, String> {
    let mut map = BTreeMap::new();

    let metrics = metrics.aggregate_by_name().sorted_for_display().timestamps_removed();

    for metric in metrics.iter() {
        map.insert(metric.value().name().into(), metric.value().to_string());
    }

    map
}

impl ExecutionPlanVisitor for JsonExecutionPlanVisitor {
    type Error = fmt::Error;

    fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> {
        let desc = PlanDescFormat { plan }.to_string();

        let json_plan = JsonExecutionPlan {
            name: plan.name().to_string(),
            desc,
            metrics: plan.metrics().as_ref().map(metrics_to_btree).unwrap_or_default(),
            children: vec![],
        };

        // Add this plan to our stack
        self.stack.push(json_plan);

        Ok(true)
    }

    fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> {
        // This is where we build the nested hierarchy

        // If we have a node, we decide where it goes
        if let Some(node) = self.stack.pop() {
            // If we're not at the root, then there will be a parent
            if let Some(parent) = self.stack.last_mut() {
                // Add this child to the current parent
                parent.children.push(node);
            } else {
                // Otherwise we're at the root, and we put the node back on the stack
                self.stack.push(node);
            }
        }

        Ok(true)
    }
}

(credit to @cetra3)

[alamb]

here is a previous discussion related to exporting this information as traces: #9415

[alamb]

It isn't quite the same as an EXPLAIN ANALYZE but it is close