Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: move all describe, describe_tree and dot-viz code to IR instead of DslPlan #16237

Merged
merged 11 commits into from
May 15, 2024
Merged

refactor: move all describe, describe_tree and dot-viz code to IR instead of DslPlan #16237

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
88b90ce
refactor: allow for display of arbitrary ExprIR
coastalwhite May 14, 2024
3e7d2e1
feat: add recursive macro to alp formatter
coastalwhite May 14, 2024
c99ea42
refactor: base implementation for IR tree_format
coastalwhite May 14, 2024
0a8b37b
refactor: replace old LazyFrame::describe
coastalwhite May 14, 2024
5d1dbe2
refactor: completely refactor dot-viz output
coastalwhite May 14, 2024
032ffe9
feat: properly print simple projects
coastalwhite May 15, 2024
6f4c109
refactor: remove old DSL describe code
coastalwhite May 15, 2024
2138ab9
fix: readd Expr debug
coastalwhite May 15, 2024
eb6cf70
fix: readd AExpr printing in tree_format
coastalwhite May 15, 2024
a465bf3
fix: clippy
coastalwhite May 15, 2024
ca70478
fix: ruff format
coastalwhite May 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
refactor: completely refactor dot-viz output
  • Loading branch information
@coastalwhite
coastalwhite committed May 15, 2024
commit 5d1dbe2d652f91ef885f34bafc6dbcc4b6fb0b13
56 changes: 6 additions & 50 deletions crates/polars-lazy/src/dot.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,60 +1,16 @@
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');
let lp = if optimized {
self.clone().to_alp_optimized()
} else {
self.clone().to_alp()
}?;

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)
Ok(lp.display_dot().to_string())
}
}
Loading