Skip to content

Commit aca4d13

Browse files
feat: Add Protobuf support for Explain node (#21994)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #. ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> `EXPLAIN FORMAT TREE` is supported in logical plans, but protobuf serialization did not preserve the explain format. In Datafusion Ballista, we need the format field to generate corresponding distributed plan. apache/datafusion-ballista#1627 (comment) ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> - Add `ExplainFormat` to protobuf common definitions. - Add the `format` field to protobuf `ExplainNode`. - Regenerate protobuf code. ## Are these changes tested? Yes, we add a roundtrip test for `EXPLAIN FORMAT TREE`. <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? No <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: Kumar Ujjawal <ujjawalpathak6@gmail.com>
1 parent ea0da8b commit aca4d13

9 files changed

Lines changed: 222 additions & 1 deletion

File tree

datafusion/proto-common/proto/datafusion_common.proto

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -669,3 +669,10 @@ message ColumnStats {
669669
Precision distinct_count = 4;
670670
Precision byte_size = 6;
671671
}
672+
673+
enum ExplainFormat {
674+
EXPLAIN_FORMAT_INDENT = 0;
675+
EXPLAIN_FORMAT_TREE = 1;
676+
EXPLAIN_FORMAT_PGJSON = 2;
677+
EXPLAIN_FORMAT_GRAPHVIZ = 3;
678+
}

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4116,6 +4116,83 @@ impl<'de> serde::Deserialize<'de> for EmptyMessage {
41164116
deserializer.deserialize_struct("datafusion_common.EmptyMessage", FIELDS, GeneratedVisitor)
41174117
}
41184118
}
4119+
impl serde::Serialize for ExplainFormat {
4120+
#[allow(deprecated)]
4121+
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
4122+
where
4123+
S: serde::Serializer,
4124+
{
4125+
let variant = match self {
4126+
Self::Indent => "EXPLAIN_FORMAT_INDENT",
4127+
Self::Tree => "EXPLAIN_FORMAT_TREE",
4128+
Self::Pgjson => "EXPLAIN_FORMAT_PGJSON",
4129+
Self::Graphviz => "EXPLAIN_FORMAT_GRAPHVIZ",
4130+
};
4131+
serializer.serialize_str(variant)
4132+
}
4133+
}
4134+
impl<'de> serde::Deserialize<'de> for ExplainFormat {
4135+
#[allow(deprecated)]
4136+
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
4137+
where
4138+
D: serde::Deserializer<'de>,
4139+
{
4140+
const FIELDS: &[&str] = &[
4141+
"EXPLAIN_FORMAT_INDENT",
4142+
"EXPLAIN_FORMAT_TREE",
4143+
"EXPLAIN_FORMAT_PGJSON",
4144+
"EXPLAIN_FORMAT_GRAPHVIZ",
4145+
];
4146+
4147+
struct GeneratedVisitor;
4148+
4149+
impl serde::de::Visitor<'_> for GeneratedVisitor {
4150+
type Value = ExplainFormat;
4151+
4152+
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4153+
write!(formatter, "expected one of: {:?}", &FIELDS)
4154+
}
4155+
4156+
fn visit_i64<E>(self, v: i64) -> std::result::Result<Self::Value, E>
4157+
where
4158+
E: serde::de::Error,
4159+
{
4160+
i32::try_from(v)
4161+
.ok()
4162+
.and_then(|x| x.try_into().ok())
4163+
.ok_or_else(|| {
4164+
serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self)
4165+
})
4166+
}
4167+
4168+
fn visit_u64<E>(self, v: u64) -> std::result::Result<Self::Value, E>
4169+
where
4170+
E: serde::de::Error,
4171+
{
4172+
i32::try_from(v)
4173+
.ok()
4174+
.and_then(|x| x.try_into().ok())
4175+
.ok_or_else(|| {
4176+
serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self)
4177+
})
4178+
}
4179+
4180+
fn visit_str<E>(self, value: &str) -> std::result::Result<Self::Value, E>
4181+
where
4182+
E: serde::de::Error,
4183+
{
4184+
match value {
4185+
"EXPLAIN_FORMAT_INDENT" => Ok(ExplainFormat::Indent),
4186+
"EXPLAIN_FORMAT_TREE" => Ok(ExplainFormat::Tree),
4187+
"EXPLAIN_FORMAT_PGJSON" => Ok(ExplainFormat::Pgjson),
4188+
"EXPLAIN_FORMAT_GRAPHVIZ" => Ok(ExplainFormat::Graphviz),
4189+
_ => Err(serde::de::Error::unknown_variant(value, FIELDS)),
4190+
}
4191+
}
4192+
}
4193+
deserializer.deserialize_any(GeneratedVisitor)
4194+
}
4195+
}
41194196
impl serde::Serialize for Field {
41204197
#[allow(deprecated)]
41214198
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1313,3 +1313,35 @@ impl PrecisionInfo {
13131313
}
13141314
}
13151315
}
1316+
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
1317+
#[repr(i32)]
1318+
pub enum ExplainFormat {
1319+
Indent = 0,
1320+
Tree = 1,
1321+
Pgjson = 2,
1322+
Graphviz = 3,
1323+
}
1324+
impl ExplainFormat {
1325+
/// String value of the enum field names used in the ProtoBuf definition.
1326+
///
1327+
/// The values are not transformed in any way and thus are considered stable
1328+
/// (if the ProtoBuf definition does not change) and safe for programmatic use.
1329+
pub fn as_str_name(&self) -> &'static str {
1330+
match self {
1331+
Self::Indent => "EXPLAIN_FORMAT_INDENT",
1332+
Self::Tree => "EXPLAIN_FORMAT_TREE",
1333+
Self::Pgjson => "EXPLAIN_FORMAT_PGJSON",
1334+
Self::Graphviz => "EXPLAIN_FORMAT_GRAPHVIZ",
1335+
}
1336+
}
1337+
/// Creates an enum from field names used in the ProtoBuf definition.
1338+
pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
1339+
match value {
1340+
"EXPLAIN_FORMAT_INDENT" => Some(Self::Indent),
1341+
"EXPLAIN_FORMAT_TREE" => Some(Self::Tree),
1342+
"EXPLAIN_FORMAT_PGJSON" => Some(Self::Pgjson),
1343+
"EXPLAIN_FORMAT_GRAPHVIZ" => Some(Self::Graphviz),
1344+
_ => None,
1345+
}
1346+
}
1347+
}

datafusion/proto/proto/datafusion.proto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ message AnalyzeNode {
229229
message ExplainNode {
230230
LogicalPlanNode input = 1;
231231
bool verbose = 2;
232+
datafusion_common.ExplainFormat format = 3;
232233
}
233234

234235
message AggregateNode {

datafusion/proto/src/generated/datafusion_proto_common.rs

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1313,3 +1313,35 @@ impl PrecisionInfo {
13131313
}
13141314
}
13151315
}
1316+
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
1317+
#[repr(i32)]
1318+
pub enum ExplainFormat {
1319+
Indent = 0,
1320+
Tree = 1,
1321+
Pgjson = 2,
1322+
Graphviz = 3,
1323+
}
1324+
impl ExplainFormat {
1325+
/// String value of the enum field names used in the ProtoBuf definition.
1326+
///
1327+
/// The values are not transformed in any way and thus are considered stable
1328+
/// (if the ProtoBuf definition does not change) and safe for programmatic use.
1329+
pub fn as_str_name(&self) -> &'static str {
1330+
match self {
1331+
Self::Indent => "EXPLAIN_FORMAT_INDENT",
1332+
Self::Tree => "EXPLAIN_FORMAT_TREE",
1333+
Self::Pgjson => "EXPLAIN_FORMAT_PGJSON",
1334+
Self::Graphviz => "EXPLAIN_FORMAT_GRAPHVIZ",
1335+
}
1336+
}
1337+
/// Creates an enum from field names used in the ProtoBuf definition.
1338+
pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
1339+
match value {
1340+
"EXPLAIN_FORMAT_INDENT" => Some(Self::Indent),
1341+
"EXPLAIN_FORMAT_TREE" => Some(Self::Tree),
1342+
"EXPLAIN_FORMAT_PGJSON" => Some(Self::Pgjson),
1343+
"EXPLAIN_FORMAT_GRAPHVIZ" => Some(Self::Graphviz),
1344+
_ => None,
1345+
}
1346+
}
1347+
}

datafusion/proto/src/generated/pbjson.rs

Lines changed: 19 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

datafusion/proto/src/generated/prost.rs

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

datafusion/proto/src/logical_plan/mod.rs

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef};
3737
use datafusion_catalog::cte_worktable::CteWorkTable;
3838
use datafusion_catalog::empty::EmptyTable;
3939
use datafusion_common::file_options::file_type::FileType;
40+
use datafusion_common::format::ExplainFormat;
4041
use datafusion_common::{
4142
Result, TableReference, ToDFSchema, assert_or_internal_err, context,
4243
internal_datafusion_err, internal_err, not_impl_err, plan_err,
@@ -801,8 +802,25 @@ impl AsLogicalPlan for LogicalPlanNode {
801802
LogicalPlanType::Explain(explain) => {
802803
let input: LogicalPlan =
803804
into_logical_plan!(explain.input, ctx, extension_codec)?;
805+
let pb_format = protobuf::ExplainFormat::try_from(explain.format)
806+
.map_err(|_| {
807+
proto_error(format!(
808+
"Received an ExplainNode message with unknown ExplainFormat {}",
809+
explain.format
810+
))
811+
})?;
812+
let explain_format = match pb_format {
813+
protobuf::ExplainFormat::Indent => ExplainFormat::Indent,
814+
protobuf::ExplainFormat::Tree => ExplainFormat::Tree,
815+
protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON,
816+
protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz,
817+
};
818+
let explain_option =
819+
datafusion_expr::logical_plan::ExplainOption::default()
820+
.with_verbose(explain.verbose)
821+
.with_format(explain_format);
804822
LogicalPlanBuilder::from(input)
805-
.explain(explain.verbose, false)?
823+
.explain_option_format(explain_option)?
806824
.build()
807825
}
808826
LogicalPlanType::SubqueryAlias(aliased_relation) => {
@@ -1758,6 +1776,17 @@ impl AsLogicalPlan for LogicalPlanNode {
17581776
protobuf::ExplainNode {
17591777
input: Some(Box::new(input)),
17601778
verbose: a.verbose,
1779+
format: match &a.explain_format {
1780+
ExplainFormat::Indent => protobuf::ExplainFormat::Indent,
1781+
ExplainFormat::Tree => protobuf::ExplainFormat::Tree,
1782+
ExplainFormat::PostgresJSON => {
1783+
protobuf::ExplainFormat::Pgjson
1784+
}
1785+
ExplainFormat::Graphviz => {
1786+
protobuf::ExplainFormat::Graphviz
1787+
}
1788+
}
1789+
.into(),
17611790
},
17621791
))),
17631792
})

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ use datafusion::physical_expr::PhysicalExpr;
6868
use datafusion::prelude::*;
6969
use datafusion::test_util::{TestTableFactory, TestTableProvider};
7070
use datafusion_common::config::TableOptions;
71+
use datafusion_common::format::ExplainFormat;
7172
use datafusion_common::scalar::ScalarStructBuilder;
7273
use datafusion_common::{
7374
DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference,
@@ -277,6 +278,27 @@ async fn roundtrip_custom_memory_tables() -> Result<()> {
277278
Ok(())
278279
}
279280

281+
#[tokio::test]
282+
async fn roundtrip_explain_format_tree() -> Result<()> {
283+
let ctx = SessionContext::new();
284+
let plan = ctx
285+
.state()
286+
.create_logical_plan("EXPLAIN FORMAT TREE SELECT 1")
287+
.await?;
288+
289+
let bytes = logical_plan_to_bytes(&plan)?;
290+
let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?;
291+
292+
match logical_round_trip {
293+
LogicalPlan::Explain(explain) => {
294+
assert_eq!(explain.explain_format, ExplainFormat::Tree);
295+
}
296+
plan => panic!("expected Explain plan, got {plan:?}"),
297+
}
298+
299+
Ok(())
300+
}
301+
280302
#[tokio::test]
281303
async fn roundtrip_custom_listing_tables() -> Result<()> {
282304
let ctx = SessionContext::new();

0 commit comments

Comments
 (0)