Skip to content

Commit 922aa5e

Browse files
alambalamb
committed
Optionally display schema in explain plan
1 parent c80da91 commit 922aa5e

File tree

12 files changed

+193
-31
lines changed

12 files changed

+193
-31
lines changed

‎datafusion/common/src/config.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -613,6 +613,9 @@ config_namespace! {
613613

614614
/// When set to true, the explain statement will print the partition sizes
615615
pub show_sizes: bool, default = true
616+
617+
/// When set to true, the explain statement will print schema information
618+
pub show_schema: bool, default = false
616619
}
617620
}
618621

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,19 @@ pub enum PlanType {
4949
InitialPhysicalPlan,
5050
/// The initial physical plan with stats, prepared for execution
5151
InitialPhysicalPlanWithStats,
52+
/// The initial physical plan with schema, prepared for execution
53+
InitialPhysicalPlanWithSchema,
5254
/// The ExecutionPlan which results from applying an optimizer pass
5355
OptimizedPhysicalPlan {
5456
/// The name of the optimizer which produced this plan
5557
optimizer_name: String,
5658
},
5759
/// The final, fully optimized physical which would be executed
5860
FinalPhysicalPlan,
59-
/// The final with stats, fully optimized physical which would be executed
61+
/// The final with stats, fully optimized physical plan which would be executed
6062
FinalPhysicalPlanWithStats,
63+
/// The final with stats, fully optimized physical plan which would be executed
64+
FinalPhysicalPlanWithSchema,
6165
}
6266

6367
impl Display for PlanType {
@@ -76,11 +80,17 @@ impl Display for PlanType {
7680
PlanType::InitialPhysicalPlanWithStats => {
7781
write!(f, "initial_physical_plan_with_stats")
7882
}
83+
PlanType::InitialPhysicalPlanWithSchema => {
84+
write!(f, "initial_physical_plan_with_schema")
85+
}
7986
PlanType::OptimizedPhysicalPlan { optimizer_name } => {
8087
write!(f, "physical_plan after {optimizer_name}")
8188
}
8289
PlanType::FinalPhysicalPlan => write!(f, "physical_plan"),
8390
PlanType::FinalPhysicalPlanWithStats => write!(f, "physical_plan_with_stats"),
91+
PlanType::FinalPhysicalPlanWithSchema => {
92+
write!(f, "physical_plan_with_schema")
93+
}
8494
}
8595
}
8696
}

‎datafusion/core/src/physical_planner.rs

Lines changed: 50 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1983,23 +1983,37 @@ impl DefaultPhysicalPlanner {
19831983
.await
19841984
{
19851985
Ok(input) => {
1986-
// This plan will includes statistics if show_statistics is on
1986+
// Include statistics / schema if enabled
19871987
stringified_plans.push(
19881988
displayable(input.as_ref())
19891989
.set_show_statistics(config.show_statistics)
1990+
.set_show_schema(config.show_schema)
19901991
.to_stringified(e.verbose, InitialPhysicalPlan),
19911992
);
19921993

1993-
// If the show_statisitcs is off, add another line to show statsitics in the case of explain verbose
1994-
if e.verbose && !config.show_statistics {
1995-
stringified_plans.push(
1996-
displayable(input.as_ref())
1997-
.set_show_statistics(true)
1998-
.to_stringified(
1999-
e.verbose,
2000-
InitialPhysicalPlanWithStats,
2001-
),
2002-
);
1994+
// Show statistics + schema in verbose output even if not
1995+
// explicitly requested
1996+
if e.verbose {
1997+
if !config.show_statistics {
1998+
stringified_plans.push(
1999+
displayable(input.as_ref())
2000+
.set_show_statistics(true)
2001+
.to_stringified(
2002+
e.verbose,
2003+
InitialPhysicalPlanWithStats,
2004+
),
2005+
);
2006+
}
2007+
if !config.show_schema {
2008+
stringified_plans.push(
2009+
displayable(input.as_ref())
2010+
.set_show_schema(true)
2011+
.to_stringified(
2012+
e.verbose,
2013+
InitialPhysicalPlanWithSchema,
2014+
),
2015+
);
2016+
}
20032017
}
20042018

20052019
let optimized_plan = self.optimize_internal(
@@ -2011,6 +2025,7 @@ impl DefaultPhysicalPlanner {
20112025
stringified_plans.push(
20122026
displayable(plan)
20132027
.set_show_statistics(config.show_statistics)
2028+
.set_show_schema(config.show_schema)
20142029
.to_stringified(e.verbose, plan_type),
20152030
);
20162031
},
@@ -2021,19 +2036,33 @@ impl DefaultPhysicalPlanner {
20212036
stringified_plans.push(
20222037
displayable(input.as_ref())
20232038
.set_show_statistics(config.show_statistics)
2039+
.set_show_schema(config.show_schema)
20242040
.to_stringified(e.verbose, FinalPhysicalPlan),
20252041
);
20262042

2027-
// If the show_statisitcs is off, add another line to show statsitics in the case of explain verbose
2028-
if e.verbose && !config.show_statistics {
2029-
stringified_plans.push(
2030-
displayable(input.as_ref())
2031-
.set_show_statistics(true)
2032-
.to_stringified(
2033-
e.verbose,
2034-
FinalPhysicalPlanWithStats,
2035-
),
2036-
);
2043+
// Show statistics + schema in verbose output even if not
2044+
// explicitly requested
2045+
if e.verbose {
2046+
if !config.show_statistics {
2047+
stringified_plans.push(
2048+
displayable(input.as_ref())
2049+
.set_show_statistics(true)
2050+
.to_stringified(
2051+
e.verbose,
2052+
FinalPhysicalPlanWithStats,
2053+
),
2054+
);
2055+
}
2056+
if !config.show_schema {
2057+
stringified_plans.push(
2058+
displayable(input.as_ref())
2059+
.set_show_schema(true)
2060+
.to_stringified(
2061+
e.verbose,
2062+
FinalPhysicalPlanWithSchema,
2063+
),
2064+
);
2065+
}
20372066
}
20382067
}
20392068
Err(DataFusionError::Context(optimizer_name, e)) => {

‎datafusion/physical-plan/src/display.rs

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,14 @@
2121
use std::fmt;
2222
use std::fmt::Formatter;
2323

24-
use super::{accept, ExecutionPlan, ExecutionPlanVisitor};
25-
2624
use arrow_schema::SchemaRef;
25+
2726
use datafusion_common::display::{GraphvizBuilder, PlanType, StringifiedPlan};
27+
use datafusion_expr::display_schema;
2828
use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
2929

30+
use super::{accept, ExecutionPlan, ExecutionPlanVisitor};
31+
3032
/// Options for controlling how each [`ExecutionPlan`] should format itself
3133
#[derive(Debug, Clone, Copy)]
3234
pub enum DisplayFormatType {
@@ -37,12 +39,15 @@ pub enum DisplayFormatType {
3739
}
3840

3941
/// Wraps an `ExecutionPlan` with various ways to display this plan
42+
#[derive(Debug, Clone)]
4043
pub struct DisplayableExecutionPlan<'a> {
4144
inner: &'a dyn ExecutionPlan,
4245
/// How to show metrics
4346
show_metrics: ShowMetrics,
4447
/// If statistics should be displayed
4548
show_statistics: bool,
49+
/// If schema should be displayed. See [`Self::set_show_schema`]
50+
show_schema: bool,
4651
}
4752

4853
impl<'a> DisplayableExecutionPlan<'a> {
@@ -53,6 +58,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
5358
inner,
5459
show_metrics: ShowMetrics::None,
5560
show_statistics: false,
61+
show_schema: false,
5662
}
5763
}
5864

@@ -64,6 +70,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
6470
inner,
6571
show_metrics: ShowMetrics::Aggregated,
6672
show_statistics: false,
73+
show_schema: false,
6774
}
6875
}
6976

@@ -75,9 +82,19 @@ impl<'a> DisplayableExecutionPlan<'a> {
7582
inner,
7683
show_metrics: ShowMetrics::Full,
7784
show_statistics: false,
85+
show_schema: false,
7886
}
7987
}
8088

89+
/// Enable display of schema
90+
///
91+
/// If true, plans will be displayed with schema information at the end
92+
/// of each line. The format is `schema=[[a:Int32;N, b:Int32;N, c:Int32;N]]`
93+
pub fn set_show_schema(mut self, show_schema: bool) -> Self {
94+
self.show_schema = show_schema;
95+
self
96+
}
97+
8198
/// Enable display of statistics
8299
pub fn set_show_statistics(mut self, show_statistics: bool) -> Self {
83100
self.show_statistics = show_statistics;
@@ -105,6 +122,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
105122
plan: &'a dyn ExecutionPlan,
106123
show_metrics: ShowMetrics,
107124
show_statistics: bool,
125+
show_schema: bool,
108126
}
109127
impl<'a> fmt::Display for Wrapper<'a> {
110128
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -114,6 +132,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
114132
indent: 0,
115133
show_metrics: self.show_metrics,
116134
show_statistics: self.show_statistics,
135+
show_schema: self.show_schema,
117136
};
118137
accept(self.plan, &mut visitor)
119138
}
@@ -123,6 +142,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
123142
plan: self.inner,
124143
show_metrics: self.show_metrics,
125144
show_statistics: self.show_statistics,
145+
show_schema: self.show_schema,
126146
}
127147
}
128148

@@ -179,6 +199,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
179199
plan: &'a dyn ExecutionPlan,
180200
show_metrics: ShowMetrics,
181201
show_statistics: bool,
202+
show_schema: bool,
182203
}
183204

184205
impl<'a> fmt::Display for Wrapper<'a> {
@@ -189,6 +210,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
189210
indent: 0,
190211
show_metrics: self.show_metrics,
191212
show_statistics: self.show_statistics,
213+
show_schema: self.show_schema,
192214
};
193215
visitor.pre_visit(self.plan)?;
194216
Ok(())
@@ -199,6 +221,7 @@ impl<'a> DisplayableExecutionPlan<'a> {
199221
plan: self.inner,
200222
show_metrics: self.show_metrics,
201223
show_statistics: self.show_statistics,
224+
show_schema: self.show_schema,
202225
}
203226
}
204227

@@ -221,6 +244,14 @@ enum ShowMetrics {
221244
}
222245

223246
/// Formats plans with a single line per node.
247+
///
248+
/// # Example
249+
///
250+
/// ```text
251+
/// ProjectionExec: expr=[column1@0 + 2 as column1 + Int64(2)]
252+
/// FilterExec: column1@0 = 5
253+
/// ValuesExec
254+
/// ```
224255
struct IndentVisitor<'a, 'b> {
225256
/// How to format each node
226257
t: DisplayFormatType,
@@ -232,6 +263,8 @@ struct IndentVisitor<'a, 'b> {
232263
show_metrics: ShowMetrics,
233264
/// If statistics should be displayed
234265
show_statistics: bool,
266+
/// If schema should be displayed
267+
show_schema: bool,
235268
}
236269

237270
impl<'a, 'b> ExecutionPlanVisitor for IndentVisitor<'a, 'b> {
@@ -265,6 +298,13 @@ impl<'a, 'b> ExecutionPlanVisitor for IndentVisitor<'a, 'b> {
265298
let stats = plan.statistics().map_err(|_e| fmt::Error)?;
266299
write!(self.f, ", statistics=[{}]", stats)?;
267300
}
301+
if self.show_schema {
302+
write!(
303+
self.f,
304+
", schema={}",
305+
display_schema(plan.schema().as_ref())
306+
)?;
307+
}
268308
writeln!(self.f)?;
269309
self.indent += 1;
270310
Ok(true)
@@ -465,12 +505,13 @@ mod tests {
465505
use std::fmt::Write;
466506
use std::sync::Arc;
467507

468-
use super::DisplayableExecutionPlan;
469-
use crate::{DisplayAs, ExecutionPlan, PlanProperties};
470-
471508
use datafusion_common::{DataFusionError, Result, Statistics};
472509
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
473510

511+
use crate::{DisplayAs, ExecutionPlan, PlanProperties};
512+
513+
use super::DisplayableExecutionPlan;
514+
474515
#[derive(Debug, Clone, Copy)]
475516
enum TestStatsExecPlan {
476517
Panic,

‎datafusion/proto/proto/datafusion.proto

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -669,9 +669,11 @@ message PlanType {
669669
datafusion_common.EmptyMessage FinalLogicalPlan = 3;
670670
datafusion_common.EmptyMessage InitialPhysicalPlan = 4;
671671
datafusion_common.EmptyMessage InitialPhysicalPlanWithStats = 9;
672+
datafusion_common.EmptyMessage InitialPhysicalPlanWithSchema = 11;
672673
OptimizedPhysicalPlanType OptimizedPhysicalPlan = 5;
673674
datafusion_common.EmptyMessage FinalPhysicalPlan = 6;
674675
datafusion_common.EmptyMessage FinalPhysicalPlanWithStats = 10;
676+
datafusion_common.EmptyMessage FinalPhysicalPlanWithSchema = 12;
675677
}
676678
}
677679

0 commit comments

Comments
 (0)