feat: Add graphviz display format for execution plan. (#6726)

* Implement graphviz format for execution plan

* Update cargo.lock

* fix ci

* fix test

* Fix comment

* Resolve conflicts with main

datafusion/common/src/display/graphviz.rs

@@ -0,0 +1,105 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Logic related to creating DOT language graphs.
+
+use std::fmt;
+
+#[derive(Default)]
+pub struct GraphvizBuilder {
+    id_gen: usize,
+}
+
+impl GraphvizBuilder {
+    // Generate next id in graphviz.
+    pub fn next_id(&mut self) -> usize {
+        self.id_gen += 1;
+        self.id_gen
+    }
+
+    // Write out the start of whole graph.
+    pub fn start_graph(&mut self, f: &mut fmt::Formatter) -> fmt::Result {
+        writeln!(
+            f,
+            r#"
+// Begin DataFusion GraphViz Plan,
+// display it online here: https://dreampuf.github.io/GraphvizOnline
+"#
+        )?;
+        writeln!(f, "digraph {{")
+    }
+
+    pub fn end_graph(&mut self, f: &mut fmt::Formatter) -> fmt::Result {
+        writeln!(f, "}}")?;
+        writeln!(f, "// End DataFusion GraphViz Plan")
+    }
+
+    // write out the start of the subgraph cluster
+    pub fn start_cluster(&mut self, f: &mut fmt::Formatter, title: &str) -> fmt::Result {
+        writeln!(f, "  subgraph cluster_{}", self.next_id())?;
+        writeln!(f, "  {{")?;
+        writeln!(f, "    graph[label={}]", Self::quoted(title))
+    }
+
+    // write out the end of the subgraph cluster
+    pub fn end_cluster(&mut self, f: &mut fmt::Formatter) -> fmt::Result {
+        writeln!(f, "  }}")
+    }
+
+    /// makes a quoted string suitable for inclusion in a graphviz chart
+    pub fn quoted(label: &str) -> String {
+        let label = label.replace('"', "_");
+        format!("\"{label}\"")
+    }
+
+    pub fn add_node(
+        &self,
+        f: &mut fmt::Formatter,
+        id: usize,
+        label: &str,
+        tooltip: Option<&str>,
+    ) -> fmt::Result {
+        if let Some(tooltip) = tooltip {
+            writeln!(
+                f,
+                "    {}[shape=box label={}, tooltip={}]",
+                id,
+                GraphvizBuilder::quoted(label),
+                GraphvizBuilder::quoted(tooltip),
+            )
+        } else {
+            writeln!(
+                f,
+                "    {}[shape=box label={}]",
+                id,
+                GraphvizBuilder::quoted(label),
+            )
+        }
+    }
+
+    pub fn add_edge(
+        &self,
+        f: &mut fmt::Formatter,
+        from_id: usize,
+        to_id: usize,
+    ) -> fmt::Result {
+        writeln!(
+            f,
+            "    {from_id} -> {to_id} [arrowhead=none, arrowtail=normal, dir=back]"
+        )
+    }
+}

datafusion/common/src/display.rs → datafusion/common/src/display/mod.rs

@@ -17,6 +17,9 @@
 
 //! Types for plan display
 
+mod graphviz;
+pub use graphviz::*;
+
 use std::{
     fmt::{self, Display, Formatter},
     sync::Arc,

datafusion/core/src/physical_plan/display.rs

@@ -20,10 +20,12 @@
 //! format
 
 use std::fmt;
+use std::fmt::Formatter;
 
 use datafusion_common::display::StringifiedPlan;
 
 use super::{accept, ExecutionPlan, ExecutionPlanVisitor};
+use datafusion_common::display::GraphvizBuilder;
 
 /// Options for controlling how each [`ExecutionPlan`] should format itself
 #[derive(Debug, Clone, Copy)]
@@ -110,6 +112,49 @@ impl<'a> DisplayableExecutionPlan<'a> {
         }
     }
 
+    /// Returns a `format`able structure that produces graphviz format for execution plan, which can
+    /// be directly visualized [here](https://dreampuf.github.io/GraphvizOnline).
+    ///
+    /// An example is
+    /// ```dot
+    /// strict digraph dot_plan {
+    //     0[label="ProjectionExec: expr=[id@0 + 2 as employee.id + Int32(2)]",tooltip=""]
+    //     1[label="EmptyExec: produce_one_row=false",tooltip=""]
+    //     0 -> 1
+    // }
+    /// ```
+    pub fn graphviz(&self) -> impl fmt::Display + 'a {
+        struct Wrapper<'a> {
+            plan: &'a dyn ExecutionPlan,
+            show_metrics: ShowMetrics,
+        }
+        impl<'a> fmt::Display for Wrapper<'a> {
+            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+                let t = DisplayFormatType::Default;
+
+                let mut visitor = GraphvizVisitor {
+                    f,
+                    t,
+                    show_metrics: self.show_metrics,
+                    graphviz_builder: GraphvizBuilder::default(),
+                    parents: Vec::new(),
+                };
+
+                visitor.start_graph()?;
+
+                accept(self.plan, &mut visitor)?;
+
+                visitor.end_graph()?;
+                Ok(())
+            }
+        }
+
+        Wrapper {
+            plan: self.inner,
+            show_metrics: self.show_metrics,
+        }
+    }
+
     /// Return a single-line summary of the root of the plan
     /// Example: `ProjectionExec: expr=[a@0 as a]`.
     pub fn one_line(&self) -> impl fmt::Display + 'a {
@@ -209,6 +254,88 @@ impl<'a, 'b> ExecutionPlanVisitor for IndentVisitor<'a, 'b> {
     }
 }
 
+struct GraphvizVisitor<'a, 'b> {
+    f: &'a mut Formatter<'b>,
+    /// How to format each node
+    t: DisplayFormatType,
+    /// How to show metrics
+    show_metrics: ShowMetrics,
+    graphviz_builder: GraphvizBuilder,
+    /// Used to record parent node ids when visiting a plan.
+    parents: Vec<usize>,
+}
+
+impl GraphvizVisitor<'_, '_> {
+    fn start_graph(&mut self) -> fmt::Result {
+        self.graphviz_builder.start_graph(self.f)
+    }
+
+    fn end_graph(&mut self) -> fmt::Result {
+        self.graphviz_builder.end_graph(self.f)
+    }
+}
+
+impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> {
+    type Error = fmt::Error;
+
+    fn pre_visit(
+        &mut self,
+        plan: &dyn ExecutionPlan,
+    ) -> datafusion_common::Result<bool, Self::Error> {
+        let id = self.graphviz_builder.next_id();
+
+        struct Wrapper<'a>(&'a dyn ExecutionPlan, DisplayFormatType);
+
+        impl<'a> std::fmt::Display for Wrapper<'a> {
+            fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+                self.0.fmt_as(self.1, f)
+            }
+        }
+
+        let label = { format!("{}", Wrapper(plan, self.t)) };
+
+        let metrics = match self.show_metrics {
+            ShowMetrics::None => "".to_string(),
+            ShowMetrics::Aggregated => {
+                if let Some(metrics) = plan.metrics() {
+                    let metrics = metrics
+                        .aggregate_by_name()
+                        .sorted_for_display()
+                        .timestamps_removed();
+
+                    format!("metrics=[{metrics}]")
+                } else {
+                    "metrics=[]".to_string()
+                }
+            }
+            ShowMetrics::Full => {
+                if let Some(metrics) = plan.metrics() {
+                    format!("metrics=[{metrics}]")
+                } else {
+                    "metrics=[]".to_string()
+                }
+            }
+        };
+
+        self.graphviz_builder
+            .add_node(self.f, id, &label, Some(&metrics))?;
+
+        if let Some(parent_node_id) = self.parents.last() {
+            self.graphviz_builder
+                .add_edge(self.f, *parent_node_id, id)?;
+        }
+
+        self.parents.push(id);
+
+        Ok(true)
+    }
+
+    fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> {
+        self.parents.pop();
+        Ok(true)
+    }
+}
+
 /// Trait for types which could have additional details when formatted in `Verbose` mode
 pub trait DisplayAs {
     /// Format according to `DisplayFormatType`, used when verbose representation looks

datafusion/core/src/physical_plan/streaming.rs

@@ -28,12 +28,11 @@ use datafusion_common::{DataFusionError, Result, Statistics};
 use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
 use log::debug;
 
-use crate::datasource::physical_plan::{OutputOrderingDisplay, ProjectSchemaDisplay};
 use crate::physical_plan::stream::RecordBatchStreamAdapter;
 use crate::physical_plan::{ExecutionPlan, Partitioning, SendableRecordBatchStream};
 use datafusion_execution::TaskContext;
 
-use super::{DisplayAs, DisplayFormatType};
+use super::DisplayAs;
 
 /// A partition that can be converted into a [`SendableRecordBatchStream`]
 pub trait PartitionStream: Send + Sync {