Skip to content

Commit 6c95b56

Browse files
irenjjalamb
irenjj
and
alamb
authored
feat: Add tree / pretty explain mode (#14677)
* feat: pretty explain * fix compile errors * render_tree * fix compile error * fmt * add License * fix test * Update config docs * Add .slt test * Add test for MemoryExec * Use BTreeMap to ensure consistent order * credit & rename pretty to tree * update config md * remove collect_info interface * add comments * fix test * fix doc --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent b377725 commit 6c95b56

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

61 files changed

+1264
-52
lines changed

‎datafusion-examples/examples/planner_api.rs

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use datafusion::error::Result;
1919
use datafusion::logical_expr::{LogicalPlan, PlanType};
20-
use datafusion::physical_plan::displayable;
20+
use datafusion::physical_plan::{displayable, DisplayFormatType};
2121
use datafusion::physical_planner::DefaultPhysicalPlanner;
2222
use datafusion::prelude::*;
2323

@@ -78,7 +78,11 @@ async fn to_physical_plan_in_one_api_demo(
7878
println!(
7979
"Physical plan direct from logical plan:\n\n{}\n\n",
8080
displayable(physical_plan.as_ref())
81-
.to_stringified(false, PlanType::InitialPhysicalPlan)
81+
.to_stringified(
82+
false,
83+
PlanType::InitialPhysicalPlan,
84+
DisplayFormatType::Default
85+
)
8286
.plan
8387
);
8488

@@ -120,7 +124,11 @@ async fn to_physical_plan_step_by_step_demo(
120124
println!(
121125
"Final physical plan:\n\n{}\n\n",
122126
displayable(physical_plan.as_ref())
123-
.to_stringified(false, PlanType::InitialPhysicalPlan)
127+
.to_stringified(
128+
false,
129+
PlanType::InitialPhysicalPlan,
130+
DisplayFormatType::Default
131+
)
124132
.plan
125133
);
126134

@@ -135,7 +143,11 @@ async fn to_physical_plan_step_by_step_demo(
135143
println!(
136144
"Optimized physical plan:\n\n{}\n\n",
137145
displayable(physical_plan.as_ref())
138-
.to_stringified(false, PlanType::InitialPhysicalPlan)
146+
.to_stringified(
147+
false,
148+
PlanType::InitialPhysicalPlan,
149+
DisplayFormatType::Default
150+
)
139151
.plan
140152
);
141153

‎datafusion/common/src/config.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -711,6 +711,10 @@ config_namespace! {
711711

712712
/// When set to true, the explain statement will print schema information
713713
pub show_schema: bool, default = false
714+
715+
/// Display format of explain. Default is "indent".
716+
/// When set to "tree", it will print the plan in a tree-rendered format.
717+
pub format: String, default = "indent".to_string()
714718
}
715719
}
716720

‎datafusion/core/src/datasource/file_format/arrow.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,10 @@ impl DisplayAs for ArrowFileSink {
301301
FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?;
302302
write!(f, ")")
303303
}
304+
DisplayFormatType::TreeRender => {
305+
// TODO: collect info
306+
write!(f, "")
307+
}
304308
}
305309
}
306310
}

‎datafusion/core/src/datasource/file_format/csv.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
// under the License.
1717

1818
//! Re-exports the [`datafusion_datasource_csv::file_format`] module, and contains tests for it.
19-
2019
pub use datafusion_datasource_csv::file_format::*;
2120

2221
#[cfg(test)]

‎datafusion/core/src/datasource/file_format/json.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
// under the License.
1717

1818
//! Re-exports the [`datafusion_datasource_json::file_format`] module, and contains tests for it.
19-
2019
pub use datafusion_datasource_json::file_format::*;
2120

2221
#[cfg(test)]

‎datafusion/core/src/datasource/memory.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,10 @@ impl DisplayAs for MemSink {
315315
let partition_count = self.batches.len();
316316
write!(f, "MemoryTable (partitions={partition_count})")
317317
}
318+
DisplayFormatType::TreeRender => {
319+
// TODO: collect info
320+
write!(f, "")
321+
}
318322
}
319323
}
320324
}

‎datafusion/core/src/physical_planner.rs

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
2020
use std::borrow::Cow;
2121
use std::collections::HashMap;
22+
use std::str::FromStr;
2223
use std::sync::Arc;
2324

2425
use crate::datasource::file_format::file_type_to_format;
@@ -87,6 +88,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule;
8788
use datafusion_physical_plan::execution_plan::InvariantLevel;
8889
use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
8990
use datafusion_physical_plan::unnest::ListUnnest;
91+
use datafusion_physical_plan::DisplayFormatType;
9092

9193
use crate::schema_equivalence::schema_satisfied_by;
9294
use async_trait::async_trait;
@@ -1723,6 +1725,7 @@ impl DefaultPhysicalPlanner {
17231725
let mut stringified_plans = vec![];
17241726

17251727
let config = &session_state.config_options().explain;
1728+
let explain_format = DisplayFormatType::from_str(&config.format)?;
17261729

17271730
if !config.physical_plan_only {
17281731
stringified_plans.clone_from(&e.stringified_plans);
@@ -1742,7 +1745,11 @@ impl DefaultPhysicalPlanner {
17421745
displayable(input.as_ref())
17431746
.set_show_statistics(config.show_statistics)
17441747
.set_show_schema(config.show_schema)
1745-
.to_stringified(e.verbose, InitialPhysicalPlan),
1748+
.to_stringified(
1749+
e.verbose,
1750+
InitialPhysicalPlan,
1751+
explain_format,
1752+
),
17461753
);
17471754

17481755
// Show statistics + schema in verbose output even if not
@@ -1755,6 +1762,7 @@ impl DefaultPhysicalPlanner {
17551762
.to_stringified(
17561763
e.verbose,
17571764
InitialPhysicalPlanWithStats,
1765+
explain_format,
17581766
),
17591767
);
17601768
}
@@ -1765,6 +1773,7 @@ impl DefaultPhysicalPlanner {
17651773
.to_stringified(
17661774
e.verbose,
17671775
InitialPhysicalPlanWithSchema,
1776+
explain_format,
17681777
),
17691778
);
17701779
}
@@ -1780,7 +1789,11 @@ impl DefaultPhysicalPlanner {
17801789
displayable(plan)
17811790
.set_show_statistics(config.show_statistics)
17821791
.set_show_schema(config.show_schema)
1783-
.to_stringified(e.verbose, plan_type),
1792+
.to_stringified(
1793+
e.verbose,
1794+
plan_type,
1795+
explain_format,
1796+
),
17841797
);
17851798
},
17861799
);
@@ -1791,7 +1804,11 @@ impl DefaultPhysicalPlanner {
17911804
displayable(input.as_ref())
17921805
.set_show_statistics(config.show_statistics)
17931806
.set_show_schema(config.show_schema)
1794-
.to_stringified(e.verbose, FinalPhysicalPlan),
1807+
.to_stringified(
1808+
e.verbose,
1809+
FinalPhysicalPlan,
1810+
explain_format,
1811+
),
17951812
);
17961813

17971814
// Show statistics + schema in verbose output even if not
@@ -1804,6 +1821,7 @@ impl DefaultPhysicalPlanner {
18041821
.to_stringified(
18051822
e.verbose,
18061823
FinalPhysicalPlanWithStats,
1824+
explain_format,
18071825
),
18081826
);
18091827
}
@@ -1814,6 +1832,7 @@ impl DefaultPhysicalPlanner {
18141832
.to_stringified(
18151833
e.verbose,
18161834
FinalPhysicalPlanWithSchema,
1835+
explain_format,
18171836
),
18181837
);
18191838
}
@@ -2720,6 +2739,10 @@ mod tests {
27202739
DisplayFormatType::Default | DisplayFormatType::Verbose => {
27212740
write!(f, "NoOpExecutionPlan")
27222741
}
2742+
DisplayFormatType::TreeRender => {
2743+
// TODO: collect info
2744+
write!(f, "")
2745+
}
27232746
}
27242747
}
27252748
}

‎datafusion/core/tests/custom_sources_cases/mod.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,11 @@ impl DisplayAs for CustomExecutionPlan {
138138
DisplayFormatType::Default | DisplayFormatType::Verbose => {
139139
write!(f, "CustomExecutionPlan: projection={:#?}", self.projection)
140140
}
141+
142+
DisplayFormatType::TreeRender => {
143+
// TODO: collect info
144+
write!(f, "")
145+
}
141146
}
142147
}
143148
}

‎datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,10 @@ impl DisplayAs for CustomPlan {
9292
DisplayFormatType::Default | DisplayFormatType::Verbose => {
9393
write!(f, "CustomPlan: batch_size={}", self.batches.len(),)
9494
}
95+
DisplayFormatType::TreeRender => {
96+
// TODO: collect info
97+
write!(f, "")
98+
}
9599
}
96100
}
97101
}

‎datafusion/core/tests/custom_sources_cases/statistics.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,10 @@ impl DisplayAs for StatisticsValidation {
141141
self.stats.num_rows,
142142
)
143143
}
144+
DisplayFormatType::TreeRender => {
145+
// TODO: collect info
146+
write!(f, "")
147+
}
144148
}
145149
}
146150
}

‎datafusion/core/tests/physical_optimizer/enforce_distribution.rs

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,10 +99,18 @@ impl SortRequiredExec {
9999
impl DisplayAs for SortRequiredExec {
100100
fn fmt_as(
101101
&self,
102-
_t: DisplayFormatType,
102+
t: DisplayFormatType,
103103
f: &mut std::fmt::Formatter,
104104
) -> std::fmt::Result {
105-
write!(f, "SortRequiredExec: [{}]", self.expr)
105+
match t {
106+
DisplayFormatType::Default | DisplayFormatType::Verbose => {
107+
write!(f, "SortRequiredExec: [{}]", self.expr)
108+
}
109+
DisplayFormatType::TreeRender => {
110+
// TODO: collect info
111+
write!(f, "")
112+
}
113+
}
106114
}
107115
}
108116

‎datafusion/core/tests/physical_optimizer/join_selection.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -924,6 +924,10 @@ impl DisplayAs for UnboundedExec {
924924
self.batch_produce.is_none(),
925925
)
926926
}
927+
DisplayFormatType::TreeRender => {
928+
// TODO: collect info
929+
write!(f, "")
930+
}
927931
}
928932
}
929933
}
@@ -1019,6 +1023,11 @@ impl DisplayAs for StatisticsExec {
10191023
self.stats.num_rows,
10201024
)
10211025
}
1026+
1027+
DisplayFormatType::TreeRender => {
1028+
// TODO: collect info
1029+
write!(f, "")
1030+
}
10221031
}
10231032
}
10241033
}

‎datafusion/core/tests/physical_optimizer/test_utils.rs

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -339,8 +339,16 @@ impl RequirementsTestExec {
339339
}
340340

341341
impl DisplayAs for RequirementsTestExec {
342-
fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
343-
write!(f, "RequiredInputOrderingExec")
342+
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
343+
match t {
344+
DisplayFormatType::Default | DisplayFormatType::Verbose => {
345+
write!(f, "RequiredInputOrderingExec")
346+
}
347+
DisplayFormatType::TreeRender => {
348+
// TODO: collect info
349+
write!(f, "")
350+
}
351+
}
344352
}
345353
}
346354

‎datafusion/core/tests/user_defined/user_defined_plan.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,10 @@ impl DisplayAs for TopKExec {
700700
DisplayFormatType::Default | DisplayFormatType::Verbose => {
701701
write!(f, "TopKExec: k={}", self.k)
702702
}
703+
DisplayFormatType::TreeRender => {
704+
// TODO: collect info
705+
write!(f, "")
706+
}
703707
}
704708
}
705709
}

‎datafusion/datasource-csv/src/file_format.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -665,6 +665,10 @@ impl DisplayAs for CsvSink {
665665
FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?;
666666
write!(f, ")")
667667
}
668+
DisplayFormatType::TreeRender => {
669+
// TODO: collect info
670+
write!(f, "")
671+
}
668672
}
669673
}
670674
}

‎datafusion/datasource-json/src/file_format.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,10 @@ impl DisplayAs for JsonSink {
324324
FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?;
325325
write!(f, ")")
326326
}
327+
DisplayFormatType::TreeRender => {
328+
// TODO: collect info
329+
write!(f, "")
330+
}
327331
}
328332
}
329333
}

‎datafusion/datasource-parquet/src/file_format.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -871,6 +871,10 @@ impl DisplayAs for ParquetSink {
871871
FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?;
872872
write!(f, ")")
873873
}
874+
DisplayFormatType::TreeRender => {
875+
// TODO: collect info
876+
write!(f, "")
877+
}
874878
}
875879
}
876880
}

‎datafusion/datasource-parquet/src/source.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -554,7 +554,9 @@ impl FileSource for ParquetSource {
554554

555555
fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
556556
match t {
557-
DisplayFormatType::Default | DisplayFormatType::Verbose => {
557+
DisplayFormatType::Default
558+
| DisplayFormatType::Verbose
559+
| DisplayFormatType::TreeRender => {
558560
let predicate_string = self
559561
.predicate()
560562
.map(|p| format!(", predicate={p}"))

‎datafusion/datasource/src/display.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ impl DisplayAs for FileGroupsDisplay<'_> {
3636
let groups = if n_groups == 1 { "group" } else { "groups" };
3737
write!(f, "{{{n_groups} {groups}: [")?;
3838
match t {
39-
DisplayFormatType::Default => {
39+
DisplayFormatType::Default | DisplayFormatType::TreeRender => {
4040
// To avoid showing too many partitions
4141
let max_groups = 5;
4242
fmt_up_to_n_elements(self.0, max_groups, f, |group, f| {
@@ -66,7 +66,7 @@ impl DisplayAs for FileGroupDisplay<'_> {
6666
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult {
6767
write!(f, "[")?;
6868
match t {
69-
DisplayFormatType::Default => {
69+
DisplayFormatType::Default | DisplayFormatType::TreeRender => {
7070
// To avoid showing too many files
7171
let max_files = 5;
7272
fmt_up_to_n_elements(self.0, max_files, f, |pf, f| {

0 commit comments

Comments
 (0)