1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

use std::fmt::Write;

use polars_core::prelude::*;
use polars_plan::dot::*;
use polars_plan::prelude::*;

use crate::prelude::*;

impl LazyFrame {
    /// Get a dot language representation of the LogicalPlan.
    pub fn to_dot(&self, optimized: bool) -> PolarsResult<String> {
        let mut s = String::with_capacity(512);

        let mut logical_plan = self.clone().get_plan_builder().build();
        if optimized {
            // initialize arena's
            let mut expr_arena = Arena::with_capacity(64);
            let mut lp_arena = Arena::with_capacity(32);

            let lp_top = self.clone().optimize_with_scratch(
                &mut lp_arena,
                &mut expr_arena,
                &mut vec![],
                true,
            )?;
            logical_plan = node_to_lp(lp_top, &expr_arena, &mut lp_arena);
        }

        let prev_node = DotNode {
            branch: 0,
            id: 0,
            fmt: "",
        };

        // maps graphviz id to label
        // we use this to create this graph
        // first we create nodes including ids to make sure they are unique
        // A [id] -- B [id]
        // B [id] -- C [id]
        //
        // then later we hide the [id] by adding this to the graph
        // A [id] [label="A"]
        // B [id] [label="B"]
        // C [id] [label="C"]

        let mut id_map = PlHashMap::with_capacity(8);
        logical_plan
            .dot(&mut s, (0, 0), prev_node, &mut id_map)
            .expect("io error");
        s.push('\n');

        for (id, label) in id_map {
            // the label is wrapped in double quotes
            // the id already is wrapped in double quotes
            writeln!(s, "{id}[label=\"{label}\"]").unwrap();
        }
        s.push_str("\n}");
        Ok(s)
    }
}