Skip to main content

vortex_layout/plan/physical/
display.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5
6pub use vortex_utils::tree::DepthContext as PlanTreeContext;
7pub use vortex_utils::tree::IndentedFormatter as PlanIndentedFormatter;
8use vortex_utils::tree::TreeDisplayAdapter;
9pub use vortex_utils::tree::TreeDisplayExtractor as PlanTreeExtractor;
10use vortex_utils::tree::write_indented_tree;
11
12use super::PlanRef;
13
14/// Adds the plan's display representation to a tree node's header.
15pub struct PlanSummaryExtractor;
16
17impl PlanSummaryExtractor {
18    /// Writes a plan directly to `formatter`.
19    pub fn write(plan: &PlanRef, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
20        write!(formatter, "{plan}")
21    }
22}
23
24impl PlanTreeExtractor<PlanRef, PlanTreeContext> for PlanSummaryExtractor {
25    fn write_header(
26        &self,
27        plan: &PlanRef,
28        _context: &PlanTreeContext,
29        formatter: &mut fmt::Formatter<'_>,
30    ) -> fmt::Result {
31        write!(formatter, " ")?;
32        Self::write(plan, formatter)
33    }
34}
35
36/// Composable display builder for a physical plan tree.
37///
38/// Call `plan.tree_display()` for the default extractors. Use `plan.tree_display_builder()` to
39/// start with only node and child names, then add extractors with [`Self::with`].
40pub struct PlanTreeDisplay<'a> {
41    plan: &'a PlanRef,
42    extractors: Vec<Box<dyn PlanTreeExtractor<PlanRef, PlanTreeContext>>>,
43}
44
45impl<'a> PlanTreeDisplay<'a> {
46    /// Creates a tree display for `plan` with no extractors.
47    pub fn new(plan: &'a PlanRef) -> Self {
48        Self {
49            plan,
50            extractors: Vec::new(),
51        }
52    }
53
54    /// Creates a tree display using each plan's display representation.
55    pub fn default_display(plan: &'a PlanRef) -> Self {
56        Self::new(plan).with(PlanSummaryExtractor)
57    }
58
59    /// Adds an extractor to the display pipeline.
60    pub fn with<E: PlanTreeExtractor<PlanRef, PlanTreeContext> + 'static>(
61        mut self,
62        extractor: E,
63    ) -> Self {
64        self.extractors.push(Box::new(extractor));
65        self
66    }
67
68    /// Adds a pre-boxed extractor to the display pipeline.
69    pub fn with_boxed(
70        mut self,
71        extractor: Box<dyn PlanTreeExtractor<PlanRef, PlanTreeContext>>,
72    ) -> Self {
73        self.extractors.push(extractor);
74        self
75    }
76}
77
78impl TreeDisplayAdapter for PlanTreeDisplay<'_> {
79    type Context = PlanTreeContext;
80    type Node = PlanRef;
81
82    fn write_node(
83        &self,
84        plan: &PlanRef,
85        context: &PlanTreeContext,
86        formatter: &mut fmt::Formatter<'_>,
87    ) -> fmt::Result {
88        for extractor in &self.extractors {
89            extractor.write_header(plan, context, formatter)?;
90        }
91        Ok(())
92    }
93
94    fn write_details(
95        &self,
96        plan: &PlanRef,
97        context: &PlanTreeContext,
98        formatter: &mut PlanIndentedFormatter<'_, '_>,
99    ) -> fmt::Result {
100        for extractor in &self.extractors {
101            extractor.write_details(plan, context, formatter)?;
102        }
103        Ok(())
104    }
105
106    fn visit_children(
107        &self,
108        plan: &PlanRef,
109        visit: &mut dyn FnMut(&str, &PlanRef, bool) -> fmt::Result,
110    ) -> fmt::Result {
111        let children = plan.children();
112        for index in 0..children.len() {
113            let child = plan.child_required(index).map_err(|_| fmt::Error)?;
114            let child_name = plan.child_name(index);
115            visit(child_name.as_ref(), &child, index + 1 == children.len())?;
116        }
117        Ok(())
118    }
119}
120
121impl fmt::Display for PlanTreeDisplay<'_> {
122    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
123        write_indented_tree(
124            self,
125            "root",
126            self.plan,
127            &mut PlanTreeContext::default(),
128            formatter,
129        )
130    }
131}