feat: Add Protobuf support for Explain node

datafusion/proto-common/proto/datafusion_common.proto

@@ -669,3 +669,10 @@ message ColumnStats {
   Precision distinct_count = 4;
   Precision byte_size = 6;
 }
+
+enum ExplainFormat {
+  EXPLAIN_FORMAT_INDENT = 0;
+  EXPLAIN_FORMAT_TREE = 1;
+  EXPLAIN_FORMAT_PGJSON = 2;
+  EXPLAIN_FORMAT_GRAPHVIZ = 3;
+}

datafusion/proto-common/src/generated/pbjson.rs

@@ -4116,6 +4116,83 @@ impl<'de> serde::Deserialize<'de> for EmptyMessage {
         deserializer.deserialize_struct("datafusion_common.EmptyMessage", FIELDS, GeneratedVisitor)
     }
 }
+impl serde::Serialize for ExplainFormat {
+    #[allow(deprecated)]
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        let variant = match self {
+            Self::Indent => "EXPLAIN_FORMAT_INDENT",
+            Self::Tree => "EXPLAIN_FORMAT_TREE",
+            Self::Pgjson => "EXPLAIN_FORMAT_PGJSON",
+            Self::Graphviz => "EXPLAIN_FORMAT_GRAPHVIZ",
+        };
+        serializer.serialize_str(variant)
+    }
+}
+impl<'de> serde::Deserialize<'de> for ExplainFormat {
+    #[allow(deprecated)]
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        const FIELDS: &[&str] = &[
+            "EXPLAIN_FORMAT_INDENT",
+            "EXPLAIN_FORMAT_TREE",
+            "EXPLAIN_FORMAT_PGJSON",
+            "EXPLAIN_FORMAT_GRAPHVIZ",
+        ];
+
+        struct GeneratedVisitor;
+
+        impl serde::de::Visitor<'_> for GeneratedVisitor {
+            type Value = ExplainFormat;
+
+            fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                write!(formatter, "expected one of: {:?}", &FIELDS)
+            }
+
+            fn visit_i64<E>(self, v: i64) -> std::result::Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                i32::try_from(v)
+                    .ok()
+                    .and_then(|x| x.try_into().ok())
+                    .ok_or_else(|| {
+                        serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self)
+                    })
+            }
+
+            fn visit_u64<E>(self, v: u64) -> std::result::Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                i32::try_from(v)
+                    .ok()
+                    .and_then(|x| x.try_into().ok())
+                    .ok_or_else(|| {
+                        serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self)
+                    })
+            }
+
+            fn visit_str<E>(self, value: &str) -> std::result::Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                match value {
+                    "EXPLAIN_FORMAT_INDENT" => Ok(ExplainFormat::Indent),
+                    "EXPLAIN_FORMAT_TREE" => Ok(ExplainFormat::Tree),
+                    "EXPLAIN_FORMAT_PGJSON" => Ok(ExplainFormat::Pgjson),
+                    "EXPLAIN_FORMAT_GRAPHVIZ" => Ok(ExplainFormat::Graphviz),
+                    _ => Err(serde::de::Error::unknown_variant(value, FIELDS)),
+                }
+            }
+        }
+        deserializer.deserialize_any(GeneratedVisitor)
+    }
+}
 impl serde::Serialize for Field {
     #[allow(deprecated)]
     fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>

datafusion/proto-common/src/generated/prost.rs

@@ -1313,3 +1313,35 @@ impl PrecisionInfo {
         }
     }
 }
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+pub enum ExplainFormat {
+    Indent = 0,
+    Tree = 1,
+    Pgjson = 2,
+    Graphviz = 3,
+}
+impl ExplainFormat {
+    /// String value of the enum field names used in the ProtoBuf definition.
+    ///
+    /// The values are not transformed in any way and thus are considered stable
+    /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+    pub fn as_str_name(&self) -> &'static str {
+        match self {
+            Self::Indent => "EXPLAIN_FORMAT_INDENT",
+            Self::Tree => "EXPLAIN_FORMAT_TREE",
+            Self::Pgjson => "EXPLAIN_FORMAT_PGJSON",
+            Self::Graphviz => "EXPLAIN_FORMAT_GRAPHVIZ",
+        }
+    }
+    /// Creates an enum from field names used in the ProtoBuf definition.
+    pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+        match value {
+            "EXPLAIN_FORMAT_INDENT" => Some(Self::Indent),
+            "EXPLAIN_FORMAT_TREE" => Some(Self::Tree),
+            "EXPLAIN_FORMAT_PGJSON" => Some(Self::Pgjson),
+            "EXPLAIN_FORMAT_GRAPHVIZ" => Some(Self::Graphviz),
+            _ => None,
+        }
+    }
+}

datafusion/proto/proto/datafusion.proto

@@ -229,6 +229,7 @@ message AnalyzeNode {
 message ExplainNode {
   LogicalPlanNode input = 1;
   bool verbose = 2;
+  datafusion_common.ExplainFormat format = 3;
 }
 
 message AggregateNode {

datafusion/proto/src/generated/datafusion_proto_common.rs

@@ -1313,3 +1313,35 @@ impl PrecisionInfo {
         }
     }
 }
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+pub enum ExplainFormat {
+    Indent = 0,
+    Tree = 1,
+    Pgjson = 2,
+    Graphviz = 3,
+}
+impl ExplainFormat {
+    /// String value of the enum field names used in the ProtoBuf definition.
+    ///
+    /// The values are not transformed in any way and thus are considered stable
+    /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+    pub fn as_str_name(&self) -> &'static str {
+        match self {
+            Self::Indent => "EXPLAIN_FORMAT_INDENT",
+            Self::Tree => "EXPLAIN_FORMAT_TREE",
+            Self::Pgjson => "EXPLAIN_FORMAT_PGJSON",
+            Self::Graphviz => "EXPLAIN_FORMAT_GRAPHVIZ",
+        }
+    }
+    /// Creates an enum from field names used in the ProtoBuf definition.
+    pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+        match value {
+            "EXPLAIN_FORMAT_INDENT" => Some(Self::Indent),
+            "EXPLAIN_FORMAT_TREE" => Some(Self::Tree),
+            "EXPLAIN_FORMAT_PGJSON" => Some(Self::Pgjson),
+            "EXPLAIN_FORMAT_GRAPHVIZ" => Some(Self::Graphviz),
+            _ => None,
+        }
+    }
+}

datafusion/proto/src/logical_plan/mod.rs

@@ -37,6 +37,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef};
 use datafusion_catalog::cte_worktable::CteWorkTable;
 use datafusion_catalog::empty::EmptyTable;
 use datafusion_common::file_options::file_type::FileType;
+use datafusion_common::format::ExplainFormat;
 use datafusion_common::{
     Result, TableReference, ToDFSchema, assert_or_internal_err, context,
     internal_datafusion_err, internal_err, not_impl_err, plan_err,
@@ -801,8 +802,25 @@ impl AsLogicalPlan for LogicalPlanNode {
             LogicalPlanType::Explain(explain) => {
                 let input: LogicalPlan =
                     into_logical_plan!(explain.input, ctx, extension_codec)?;
+                let pb_format = protobuf::ExplainFormat::try_from(explain.format)
+                    .map_err(|_| {
+                        proto_error(format!(
+                            "Received an ExplainNode message with unknown ExplainFormat {}",
+                            explain.format
+                        ))
+                    })?;
+                let explain_format = match pb_format {
+                    protobuf::ExplainFormat::Indent => ExplainFormat::Indent,
+                    protobuf::ExplainFormat::Tree => ExplainFormat::Tree,
+                    protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON,
+                    protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz,
+                };
+                let explain_option =
+                    datafusion_expr::logical_plan::ExplainOption::default()
+                        .with_verbose(explain.verbose)
+                        .with_format(explain_format);
                 LogicalPlanBuilder::from(input)
-                    .explain(explain.verbose, false)?
+                    .explain_option_format(explain_option)?
                     .build()
             }
             LogicalPlanType::SubqueryAlias(aliased_relation) => {
@@ -1758,6 +1776,17 @@ impl AsLogicalPlan for LogicalPlanNode {
                         protobuf::ExplainNode {
                             input: Some(Box::new(input)),
                             verbose: a.verbose,
+                            format: match &a.explain_format {
+                                ExplainFormat::Indent => protobuf::ExplainFormat::Indent,
+                                ExplainFormat::Tree => protobuf::ExplainFormat::Tree,
+                                ExplainFormat::PostgresJSON => {
+                                    protobuf::ExplainFormat::Pgjson
+                                }
+                                ExplainFormat::Graphviz => {
+                                    protobuf::ExplainFormat::Graphviz
+                                }
+                            }
+                            .into(),
                         },
                     ))),
                 })

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -68,6 +68,7 @@ use datafusion::physical_expr::PhysicalExpr;
 use datafusion::prelude::*;
 use datafusion::test_util::{TestTableFactory, TestTableProvider};
 use datafusion_common::config::TableOptions;
+use datafusion_common::format::ExplainFormat;
 use datafusion_common::scalar::ScalarStructBuilder;
 use datafusion_common::{
     DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference,
@@ -277,6 +278,27 @@ async fn roundtrip_custom_memory_tables() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn roundtrip_explain_format_tree() -> Result<()> {
+    let ctx = SessionContext::new();
+    let plan = ctx
+        .state()
+        .create_logical_plan("EXPLAIN FORMAT TREE SELECT 1")
+        .await?;
+
+    let bytes = logical_plan_to_bytes(&plan)?;
+    let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?;
+
+    match logical_round_trip {
+        LogicalPlan::Explain(explain) => {
+            assert_eq!(explain.explain_format, ExplainFormat::Tree);
+        }
+        plan => panic!("expected Explain plan, got {plan:?}"),
+    }
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn roundtrip_custom_listing_tables() -> Result<()> {
     let ctx = SessionContext::new();