Support EXPLAIN FORMAT <format>

Author: alamb

datafusion/core/src/execution/session_state.rs

@@ -577,6 +577,7 @@ impl SessionState {
                     return Ok(LogicalPlan::Explain(Explain {
                         verbose: e.verbose,
                         plan: Arc::clone(&e.plan),
+                        format: e.format.clone(),
                         stringified_plans,
                         schema: Arc::clone(&e.schema),
                         logical_optimization_succeeded: false,
@@ -612,6 +613,7 @@ impl SessionState {
 
             Ok(LogicalPlan::Explain(Explain {
                 verbose: e.verbose,
+                format: e.format.clone(),
                 plan,
                 stringified_plans,
                 schema: Arc::clone(&e.schema),

datafusion/core/src/physical_planner.rs

@@ -1211,6 +1211,7 @@ impl LogicalPlanBuilder {
             Ok(Self::new(LogicalPlan::Explain(Explain {
                 verbose,
                 plan: self.plan,
+                format: None,
                 stringified_plans,
                 schema,
                 logical_optimization_succeeded: false,

datafusion/expr/src/logical_plan/builder.rs

@@ -1084,6 +1084,7 @@ impl LogicalPlan {
                 Ok(LogicalPlan::Explain(Explain {
                     verbose: e.verbose,
                     plan: Arc::new(input),
+                    format: e.format.clone(),
                     stringified_plans: e.stringified_plans.clone(),
                     schema: Arc::clone(&e.schema),
                     logical_optimization_succeeded: e.logical_optimization_succeeded,
@@ -2930,6 +2931,9 @@ impl PartialOrd for DescribeTable {
 pub struct Explain {
     /// Should extra (detailed, intermediate plans) be included?
     pub verbose: bool,
+    /// Output format for explain, if specified.
+    /// If none, defaults to `text`
+    pub format: Option<String>,
     /// The logical plan that is being EXPLAIN'd
     pub plan: Arc<LogicalPlan>,
     /// Represent the various stages plans have gone through

datafusion/expr/src/logical_plan/plan.rs

@@ -202,13 +202,15 @@ impl TreeNode for LogicalPlan {
             .update_data(LogicalPlan::Distinct),
             LogicalPlan::Explain(Explain {
                 verbose,
+                format,
                 plan,
                 stringified_plans,
                 schema,
                 logical_optimization_succeeded,
             }) => plan.map_elements(f)?.update_data(|plan| {
                 LogicalPlan::Explain(Explain {
                     verbose,
+                    format,
                     plan,
                     stringified_plans,
                     schema,

datafusion/expr/src/logical_plan/tree_node.rs

@@ -19,13 +19,12 @@
 //! [`crate::displayable`] for examples of how to format
 
 use std::collections::{BTreeMap, HashMap};
+use std::fmt;
 use std::fmt::Formatter;
-use std::{fmt, str::FromStr};
 
 use arrow::datatypes::SchemaRef;
 
 use datafusion_common::display::{GraphvizBuilder, PlanType, StringifiedPlan};
-use datafusion_common::DataFusionError;
 use datafusion_expr::display_schema;
 use datafusion_physical_expr::LexOrdering;
 
@@ -39,7 +38,7 @@ pub enum DisplayFormatType {
     /// Default, compact format. Example: `FilterExec: c12 < 10.0`
     ///
     /// This format is designed to provide a detailed textual description
-    /// of all rele
+    /// of all parts of the plan.
     Default,
     /// Verbose, showing all available details.
     ///
@@ -79,21 +78,6 @@ pub enum DisplayFormatType {
     TreeRender,
 }
 
-impl FromStr for DisplayFormatType {
-    type Err = DataFusionError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s.to_lowercase().as_str() {
-            "indent" => Ok(Self::Default),
-            "tree" => Ok(Self::TreeRender),
-            _ => Err(DataFusionError::Configuration(format!(
-                "Invalid explain format: {}",
-                s
-            ))),
-        }
-    }
-}
-
 /// Wraps an `ExecutionPlan` with various methods for formatting
 ///
 ///

datafusion/physical-plan/src/display.rs

@@ -16,6 +16,9 @@
 // under the License.
 
 //! [`DFParser`]: DataFusion SQL Parser based on [`sqlparser`]
+//!
+//! This parser implements DataFusion specific statements such as
+//! `CREATE EXTERNAL TABLE`
 
 use std::collections::VecDeque;
 use std::fmt;
@@ -43,12 +46,23 @@ fn parse_file_type(s: &str) -> Result<String, ParserError> {
     Ok(s.to_uppercase())
 }
 
-/// DataFusion specific EXPLAIN (needed so we can EXPLAIN datafusion
-/// specific COPY and other statements)
+/// DataFusion specific `EXPLAIN`
+///
+/// Syntax:
+/// ```sql
+/// EXPLAIN <ANALYZE> <VERBOSE> [FORMAT format] statement
+///```
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ExplainStatement {
+    /// `EXPLAIN ANALYZE ..`
     pub analyze: bool,
+    /// `EXPLAIN .. VERBOSE ..`
     pub verbose: bool,
+    /// `EXPLAIN .. FORMAT `
+    pub format: Option<String>,
+    /// The statement to analyze. Note this is a DataFusion [`Statement`] (not a
+    /// [`sqlparser::ast::Statement`] so that we can `EXPLAIN`  `COPY` and other
+    /// DataFusion specific statements
     pub statement: Box<Statement>,
 }
 
@@ -57,6 +71,7 @@ impl fmt::Display for ExplainStatement {
         let Self {
             analyze,
             verbose,
+            format,
             statement,
         } = self;
 
@@ -67,6 +82,9 @@ impl fmt::Display for ExplainStatement {
         if *verbose {
             write!(f, "VERBOSE ")?;
         }
+        if let Some(format) = format.as_ref() {
+            write!(f, "FORMAT {format} ")?;
+        }
 
         write!(f, "{statement}")
     }
@@ -446,7 +464,6 @@ impl<'a> DFParser<'a> {
                         self.parse_copy()
                     }
                     Keyword::EXPLAIN => {
-                        // (TODO parse all supported statements)
                         self.parser.next_token(); // EXPLAIN
                         self.parse_explain()
                     }
@@ -620,15 +637,35 @@ impl<'a> DFParser<'a> {
     pub fn parse_explain(&mut self) -> Result<Statement, ParserError> {
         let analyze = self.parser.parse_keyword(Keyword::ANALYZE);
         let verbose = self.parser.parse_keyword(Keyword::VERBOSE);
+        let format = self.parse_explain_format()?;
+
         let statement = self.parse_statement()?;
 
         Ok(Statement::Explain(ExplainStatement {
             statement: Box::new(statement),
             analyze,
             verbose,
+            format,
         }))
     }
 
+    pub fn parse_explain_format(&mut self) -> Result<Option<String>, ParserError> {
+        if !self.parser.parse_keyword(Keyword::FORMAT) {
+            return Ok(None);
+        }
+
+        let next_token = self.parser.next_token();
+        let format = match next_token.token {
+            Token::Word(w) => Ok(w.value),
+            Token::SingleQuotedString(w) => Ok(w),
+            Token::DoubleQuotedString(w) => Ok(w),
+            _ => self
+                .parser
+                .expected("an explain format like JSON or TREE", next_token),
+        }?;
+        Ok(Some(format))
+    }
+
     /// Parse a SQL `CREATE` statement handling `CREATE EXTERNAL TABLE`
     pub fn parse_create(&mut self) -> Result<Statement, ParserError> {
         if self.parser.parse_keyword(Keyword::EXTERNAL) {
@@ -1543,6 +1580,7 @@ mod tests {
             let expected = Statement::Explain(ExplainStatement {
                 analyze,
                 verbose,
+                format: None,
                 statement: Box::new(expected_copy),
             });
             assert_eq!(verified_stmt(sql), expected);

datafusion/sql/src/parser.rs

@@ -177,8 +177,9 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             DFStatement::Explain(ExplainStatement {
                 verbose,
                 analyze,
+                format,
                 statement,
-            }) => self.explain_to_plan(verbose, analyze, *statement),
+            }) => self.explain_to_plan(verbose, analyze, format, *statement),
         }
     }
 
@@ -214,11 +215,13 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 verbose,
                 statement,
                 analyze,
-                format: _,
+                format,
                 describe_alias: _,
                 ..
             } => {
-                self.explain_to_plan(verbose, analyze, DFStatement::Statement(statement))
+                let format = format.map(|format| format.to_string());
+                let statement = DFStatement::Statement(statement);
+                self.explain_to_plan(verbose, analyze, format, statement)
             }
             Statement::Query(query) => self.query_to_plan(*query, planner_context),
             Statement::ShowVariable { variable } => self.show_variable_to_plan(&variable),
@@ -1564,17 +1567,26 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         &self,
         verbose: bool,
         analyze: bool,
+        format: Option<String>,
         statement: DFStatement,
     ) -> Result<LogicalPlan> {
         let plan = self.statement_to_plan(statement)?;
         if matches!(plan, LogicalPlan::Explain(_)) {
             return plan_err!("Nested EXPLAINs are not supported");
         }
+
         let plan = Arc::new(plan);
         let schema = LogicalPlan::explain_schema();
         let schema = schema.to_dfschema_ref()?;
 
+        if verbose && format.is_some() {
+            return plan_err!("EXPLAIN VERBOSE with FORMAT is not supported");
+        }
+
         if analyze {
+            if format.is_some() {
+                return plan_err!("EXPLAIN ANALYZE with FORMAT is not supported");
+            }
             Ok(LogicalPlan::Analyze(Analyze {
                 verbose,
                 input: plan,
@@ -1585,6 +1597,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 vec![plan.to_stringified(PlanType::InitialLogicalPlan)];
             Ok(LogicalPlan::Explain(Explain {
                 verbose,
+                format,
                 plan,
                 stringified_plans,
                 schema,