diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 87bcf3c14493c..5f4ba2b9acebd 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/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; +} \ No newline at end of file diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 6112c55793a2b..f6b5bbeaf3961 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/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(&self, serializer: S) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(self, v: i64) -> std::result::Result + 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(self, v: u64) -> std::result::Result + 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(self, value: &str) -> std::result::Result + 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(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index 4472ff0cde59b..f09af1db349e4 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/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 { + 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, + } + } +} diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 511e8eb1b012e..865887d41e111 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/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 { diff --git a/datafusion/proto/src/generated/datafusion_proto_common.rs b/datafusion/proto/src/generated/datafusion_proto_common.rs index 4472ff0cde59b..f09af1db349e4 100644 --- a/datafusion/proto/src/generated/datafusion_proto_common.rs +++ b/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 { + 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, + } + } +} diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index c05d3283eac8e..b8639afd04a89 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -6197,6 +6197,9 @@ impl serde::Serialize for ExplainNode { if self.verbose { len += 1; } + if self.format != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ExplainNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -6204,6 +6207,11 @@ impl serde::Serialize for ExplainNode { if self.verbose { struct_ser.serialize_field("verbose", &self.verbose)?; } + if self.format != 0 { + let v = super::datafusion_common::ExplainFormat::try_from(self.format) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.format)))?; + struct_ser.serialize_field("format", &v)?; + } struct_ser.end() } } @@ -6216,12 +6224,14 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { const FIELDS: &[&str] = &[ "input", "verbose", + "format", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Input, Verbose, + Format, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -6245,6 +6255,7 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { match value { "input" => Ok(GeneratedField::Input), "verbose" => Ok(GeneratedField::Verbose), + "format" => Ok(GeneratedField::Format), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -6266,6 +6277,7 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { { let mut input__ = None; let mut verbose__ = None; + let mut format__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -6280,11 +6292,18 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { } verbose__ = Some(map_.next_value()?); } + GeneratedField::Format => { + if format__.is_some() { + return Err(serde::de::Error::duplicate_field("format")); + } + format__ = Some(map_.next_value::()? as i32); + } } } Ok(ExplainNode { input: input__, verbose: verbose__.unwrap_or_default(), + format: format__.unwrap_or_default(), }) } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index af9b1404bb80a..b742e82ea24ec 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -351,6 +351,8 @@ pub struct ExplainNode { pub input: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(bool, tag = "2")] pub verbose: bool, + #[prost(enumeration = "super::datafusion_common::ExplainFormat", tag = "3")] + pub format: i32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct AggregateNode { diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 7ae5cbeed3e53..8228e8e6f2ff0 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/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(), }, ))), }) diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index dbc95536f0104..3e79ddab723eb 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/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();