datafusion_physical_plan/
render_tree.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18// This code is based on the DuckDB’s implementation:
19// <https://github.com/duckdb/duckdb/blob/main/src/include/duckdb/common/render_tree.hpp>
20
21//! This module provides functionality for rendering an execution plan as a tree structure.
22//! It helps in visualizing how different operations in a query are connected and organized.
23
24use std::collections::HashMap;
25use std::fmt::Formatter;
26use std::sync::Arc;
27use std::{cmp, fmt};
28
29use crate::{DisplayFormatType, ExecutionPlan};
30
31// TODO: It's never used.
32/// Represents a 2D coordinate in the rendered tree.
33/// Used to track positions of nodes and their connections.
34#[allow(dead_code)]
35pub struct Coordinate {
36    /// Horizontal position in the tree
37    pub x: usize,
38    /// Vertical position in the tree
39    pub y: usize,
40}
41
42impl Coordinate {
43    pub fn new(x: usize, y: usize) -> Self {
44        Coordinate { x, y }
45    }
46}
47
48/// Represents a node in the render tree, containing information about an execution plan operator
49/// and its relationships to other operators.
50pub struct RenderTreeNode {
51    /// The name of physical `ExecutionPlan`.
52    pub name: String,
53    /// Execution info collected from `ExecutionPlan`.
54    pub extra_text: HashMap<String, String>,
55    /// Positions of child nodes in the rendered tree.
56    pub child_positions: Vec<Coordinate>,
57}
58
59impl RenderTreeNode {
60    pub fn new(name: String, extra_text: HashMap<String, String>) -> Self {
61        RenderTreeNode {
62            name,
63            extra_text,
64            child_positions: vec![],
65        }
66    }
67
68    fn add_child_position(&mut self, x: usize, y: usize) {
69        self.child_positions.push(Coordinate::new(x, y));
70    }
71}
72
73/// Main structure for rendering an execution plan as a tree.
74/// Manages a 2D grid of nodes and their layout information.
75pub struct RenderTree {
76    /// Storage for tree nodes in a flattened 2D grid
77    pub nodes: Vec<Option<Arc<RenderTreeNode>>>,
78    /// Total width of the rendered tree
79    pub width: usize,
80    /// Total height of the rendered tree
81    pub height: usize,
82}
83
84impl RenderTree {
85    /// Creates a new render tree from an execution plan.
86    pub fn create_tree(plan: &dyn ExecutionPlan) -> Self {
87        let (width, height) = get_tree_width_height(plan);
88
89        let mut result = Self::new(width, height);
90
91        create_tree_recursive(&mut result, plan, 0, 0);
92
93        result
94    }
95
96    fn new(width: usize, height: usize) -> Self {
97        RenderTree {
98            nodes: vec![None; (width + 1) * (height + 1)],
99            width,
100            height,
101        }
102    }
103
104    pub fn get_node(&self, x: usize, y: usize) -> Option<Arc<RenderTreeNode>> {
105        if x >= self.width || y >= self.height {
106            return None;
107        }
108
109        let pos = self.get_position(x, y);
110        self.nodes.get(pos).and_then(|node| node.clone())
111    }
112
113    pub fn set_node(&mut self, x: usize, y: usize, node: Arc<RenderTreeNode>) {
114        let pos = self.get_position(x, y);
115        if let Some(slot) = self.nodes.get_mut(pos) {
116            *slot = Some(node);
117        }
118    }
119
120    pub fn has_node(&self, x: usize, y: usize) -> bool {
121        if x >= self.width || y >= self.height {
122            return false;
123        }
124
125        let pos = self.get_position(x, y);
126        self.nodes.get(pos).is_some_and(|node| node.is_some())
127    }
128
129    fn get_position(&self, x: usize, y: usize) -> usize {
130        y * self.width + x
131    }
132}
133
134/// Calculates the required dimensions of the tree.
135/// This ensures we allocate enough space for the entire tree structure.
136///
137/// # Arguments
138/// * `plan` - The execution plan to measure
139///
140/// # Returns
141/// * A tuple of (width, height) representing the dimensions needed for the tree
142fn get_tree_width_height(plan: &dyn ExecutionPlan) -> (usize, usize) {
143    let children = plan.children();
144
145    // Leaf nodes take up 1x1 space
146    if children.is_empty() {
147        return (1, 1);
148    }
149
150    let mut width = 0;
151    let mut height = 0;
152
153    for child in children {
154        let (child_width, child_height) = get_tree_width_height(child.as_ref());
155        width += child_width;
156        height = cmp::max(height, child_height);
157    }
158
159    height += 1;
160
161    (width, height)
162}
163
164fn fmt_display(plan: &dyn ExecutionPlan) -> impl fmt::Display + '_ {
165    struct Wrapper<'a> {
166        plan: &'a dyn ExecutionPlan,
167    }
168
169    impl fmt::Display for Wrapper<'_> {
170        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
171            self.plan.fmt_as(DisplayFormatType::TreeRender, f)?;
172            Ok(())
173        }
174    }
175
176    Wrapper { plan }
177}
178
179/// Recursively builds the render tree structure.
180/// Traverses the execution plan and creates corresponding render nodes while
181/// maintaining proper positioning and parent-child relationships.
182///
183/// # Arguments
184/// * `result` - The render tree being constructed
185/// * `plan` - Current execution plan node being processed
186/// * `x` - Horizontal position in the tree
187/// * `y` - Vertical position in the tree
188///
189/// # Returns
190/// * The width of the subtree rooted at the current node
191fn create_tree_recursive(
192    result: &mut RenderTree,
193    plan: &dyn ExecutionPlan,
194    x: usize,
195    y: usize,
196) -> usize {
197    let display_info = fmt_display(plan).to_string();
198    let mut extra_info = HashMap::new();
199
200    // Parse the key-value pairs from the formatted string.
201    // See DisplayFormatType::TreeRender for details
202    for line in display_info.lines() {
203        if let Some((key, value)) = line.split_once('=') {
204            extra_info.insert(key.to_string(), value.to_string());
205        } else {
206            extra_info.insert(line.to_string(), "".to_string());
207        }
208    }
209
210    let mut node = RenderTreeNode::new(plan.name().to_string(), extra_info);
211
212    let children = plan.children();
213
214    if children.is_empty() {
215        result.set_node(x, y, Arc::new(node));
216        return 1;
217    }
218
219    let mut width = 0;
220    for child in children {
221        let child_x = x + width;
222        let child_y = y + 1;
223        node.add_child_position(child_x, child_y);
224        width += create_tree_recursive(result, child.as_ref(), child_x, child_y);
225    }
226
227    result.set_node(x, y, Arc::new(node));
228
229    width
230}
datafusion_physical_plan/render_tree.rs

datafusion_physical_plan/
render_tree.rs
