risingwave_frontend/utils/
stream_graph_formatter.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp::max;
16use std::collections::{BTreeMap, HashMap};
17
18use itertools::Itertools;
19use petgraph::Graph;
20use petgraph::dot::Dot;
21use petgraph::graph::NodeIndex;
22use pretty_xmlish::{Pretty, PrettyConfig};
23use risingwave_common::util::stream_graph_visitor;
24use risingwave_pb::catalog::Table;
25use risingwave_pb::stream_plan::stream_fragment_graph::StreamFragmentEdge;
26use risingwave_pb::stream_plan::{DispatcherType, StreamFragmentGraph, StreamNode, stream_node};
27
28use super::PrettySerde;
29use crate::TableCatalog;
30use crate::catalog::TableId;
31
32/// ice: in the future, we may allow configurable width, boundaries, etc.
33pub fn explain_stream_graph(
34    graph: &StreamFragmentGraph,
35    table: Option<Table>,
36    is_verbose: bool,
37) -> String {
38    let mut output = String::with_capacity(2048);
39    let mut config = PrettyConfig {
40        need_boundaries: false,
41        width: 80,
42        ..Default::default()
43    };
44    let mut fmt = StreamGraphFormatter::new(is_verbose);
45    if let Some(tb) = table {
46        fmt.add_table(&tb);
47    }
48    fmt.explain_graph(graph, &mut config, &mut output);
49    output
50}
51
52pub fn explain_stream_graph_as_dot(
53    sg: &StreamFragmentGraph,
54    table: Option<Table>,
55    is_verbose: bool,
56) -> String {
57    let mut fmt = StreamGraphFormatter::new(is_verbose);
58    if let Some(tb) = table {
59        fmt.add_table(&tb);
60    }
61    let graph = fmt.explain_graph_as_dot(sg);
62    let dot = Dot::new(&graph);
63    dot.to_string()
64}
65
66/// A formatter to display the final stream plan graph, used for `explain (distsql) create
67/// materialized view ...`
68struct StreamGraphFormatter {
69    /// exchange's `operator_id` -> edge
70    edges: HashMap<u64, StreamFragmentEdge>,
71    verbose: bool,
72    tables: BTreeMap<TableId, Table>,
73}
74
75impl StreamGraphFormatter {
76    fn new(verbose: bool) -> Self {
77        StreamGraphFormatter {
78            edges: HashMap::default(),
79            tables: BTreeMap::default(),
80            verbose,
81        }
82    }
83
84    /// collect the table catalog and return the table id
85    fn add_table(&mut self, tb: &Table) -> TableId {
86        self.tables.insert(tb.id, tb.clone());
87        tb.id
88    }
89
90    fn explain_graph(
91        &mut self,
92        graph: &StreamFragmentGraph,
93        config: &mut PrettyConfig,
94        output: &mut String,
95    ) {
96        self.edges.clear();
97        for edge in &graph.edges {
98            self.edges.insert(edge.link_id, edge.clone());
99        }
100        let mut max_width = 80;
101        for (_, fragment) in graph.fragments.iter().sorted_by_key(|(id, _)| **id) {
102            output.push_str("Fragment ");
103            output.push_str(&fragment.get_fragment_id().to_string());
104            output.push('\n');
105            let width = config.unicode(output, &self.explain_node(fragment.node.as_ref().unwrap()));
106            max_width = max(width, max_width);
107            config.width = max_width;
108            output.push_str("\n\n");
109        }
110        for tb in self.tables.values() {
111            let width = config.unicode(output, &self.explain_table(tb));
112            max_width = max(width, max_width);
113            config.width = max_width;
114            output.push_str("\n\n");
115        }
116    }
117
118    fn explain_graph_as_dot(&mut self, graph: &StreamFragmentGraph) -> Graph<String, String> {
119        self.edges.clear();
120        for edge in &graph.edges {
121            self.edges.insert(edge.link_id, edge.clone());
122        }
123
124        let mut g = Graph::<String, String>::new();
125        let mut nodes = HashMap::new();
126        for (_, fragment) in graph.fragments.iter().sorted_by_key(|(id, _)| **id) {
127            let mut label = String::new();
128            label.push_str("Fragment ");
129            label.push_str(&fragment.get_fragment_id().to_string());
130            label.push('\n');
131            nodes.insert(label.clone(), g.add_node(label.clone()));
132
133            build_graph_from_pretty(
134                &self.explain_node(fragment.node.as_ref().unwrap()),
135                &mut g,
136                &mut nodes,
137                Some(&label),
138            );
139        }
140        for tb in self.tables.values() {
141            build_graph_from_pretty(&self.explain_table(tb), &mut g, &mut nodes, None);
142        }
143        g
144    }
145
146    fn explain_table<'a>(&self, tb: &Table) -> Pretty<'a> {
147        let tb = TableCatalog::from(tb.clone());
148        let columns = tb
149            .columns
150            .iter()
151            .map(|c| {
152                let s = if self.verbose {
153                    format!("{}: {}", c.name(), c.data_type())
154                } else {
155                    c.name().to_owned()
156                };
157                Pretty::Text(s.into())
158            })
159            .collect();
160        let columns = Pretty::Array(columns);
161        let name = format!("Table {}", tb.id);
162        let mut fields = Vec::with_capacity(5);
163        fields.push(("columns", columns));
164        fields.push((
165            "primary key",
166            Pretty::Array(tb.pk.iter().map(Pretty::debug).collect()),
167        ));
168        fields.push((
169            "value indices",
170            Pretty::Array(tb.value_indices.iter().map(Pretty::debug).collect()),
171        ));
172        fields.push((
173            "distribution key",
174            Pretty::Array(tb.distribution_key.iter().map(Pretty::debug).collect()),
175        ));
176        fields.push((
177            "read pk prefix len hint",
178            Pretty::debug(&tb.read_prefix_len_hint),
179        ));
180        if let Some(vnode_col_idx) = tb.vnode_col_index {
181            fields.push(("vnode column idx", Pretty::debug(&vnode_col_idx)));
182        }
183        Pretty::childless_record(name, fields)
184    }
185
186    fn explain_node<'a>(&mut self, node: &StreamNode) -> Pretty<'a> {
187        let one_line_explain = match node.get_node_body().unwrap() {
188            stream_node::NodeBody::Exchange(_) => {
189                let edge = self.edges.get(&node.operator_id).unwrap();
190                let upstream_fragment_id = edge.upstream_id;
191                let dist = edge.dispatch_strategy.as_ref().unwrap();
192                format!(
193                    "StreamExchange {} from {}",
194                    match dist.r#type() {
195                        DispatcherType::Unspecified => unreachable!(),
196                        DispatcherType::Hash => format!("Hash({:?})", dist.dist_key_indices),
197                        DispatcherType::Broadcast => "Broadcast".to_owned(),
198                        DispatcherType::Simple => "Single".to_owned(),
199                        DispatcherType::NoShuffle => "NoShuffle".to_owned(),
200                    },
201                    upstream_fragment_id
202                )
203            }
204            _ => node.identity.clone(),
205        };
206
207        let mut tables: Vec<(String, TableId)> = Vec::with_capacity(7);
208        let mut node_copy = node.clone();
209
210        stream_graph_visitor::visit_stream_node_tables_inner(
211            &mut node_copy,
212            true,
213            false,
214            |table, table_name| {
215                tables.push((table_name.to_owned(), self.add_table(table)));
216            },
217        );
218
219        let mut fields = Vec::with_capacity(3);
220        if !tables.is_empty() {
221            fields.push((
222                "tables",
223                Pretty::Array(
224                    tables
225                        .into_iter()
226                        .map(|(name, id)| Pretty::Text(format!("{}: {}", name, id).into()))
227                        .collect(),
228                ),
229            ));
230        }
231        if self.verbose {
232            fields.push((
233                "output",
234                Pretty::Array(
235                    node.fields
236                        .iter()
237                        .map(|f| Pretty::display(f.get_name()))
238                        .collect(),
239                ),
240            ));
241            fields.push((
242                "stream key",
243                Pretty::Array(
244                    node.stream_key
245                        .iter()
246                        .map(|i| Pretty::display(node.fields[*i as usize].get_name()))
247                        .collect(),
248                ),
249            ));
250        }
251        let children = node
252            .input
253            .iter()
254            .map(|input| self.explain_node(input))
255            .collect();
256        Pretty::simple_record(one_line_explain, fields, children)
257    }
258}
259
260pub fn build_graph_from_pretty(
261    pretty: &Pretty<'_>,
262    graph: &mut Graph<String, String>,
263    nodes: &mut HashMap<String, NodeIndex>,
264    parent_label: Option<&str>,
265) {
266    if let Pretty::Record(r) = pretty {
267        let mut label = String::new();
268        label.push_str(&r.name);
269        for (k, v) in &r.fields {
270            label.push('\n');
271            label.push_str(k);
272            label.push_str(": ");
273            label.push_str(
274                &serde_json::to_string(&PrettySerde(v.clone(), false))
275                    .expect("failed to serialize plan to dot"),
276            );
277        }
278        // output alignment.
279        if !r.fields.is_empty() {
280            label.push('\n');
281        }
282
283        let current_node = *nodes
284            .entry(label.clone())
285            .or_insert_with(|| graph.add_node(label.clone()));
286
287        if let Some(parent_label) = parent_label
288            && let Some(&parent_node) = nodes.get(parent_label)
289        {
290            graph.add_edge(parent_node, current_node, "".to_owned());
291        }
292
293        for child in &r.children {
294            build_graph_from_pretty(child, graph, nodes, Some(&label));
295        }
296    }
297}
risingwave_frontend/utils/stream_graph_formatter.rs

risingwave_frontend/utils/
stream_graph_formatter.rs
