diff --git a/Cargo.lock b/Cargo.lock index 591eaae4e817..a1d6fe72fec1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3248,6 +3248,7 @@ dependencies = [ "polars-utils", "pyo3", "rayon", + "regex", "serde_json", "tokio", "version_check", diff --git a/crates/polars-lazy/Cargo.toml b/crates/polars-lazy/Cargo.toml index ed172a178cb1..d8e4b1cfe68b 100644 --- a/crates/polars-lazy/Cargo.toml +++ b/crates/polars-lazy/Cargo.toml @@ -29,6 +29,7 @@ memchr = { workspace = true } once_cell = { workspace = true } pyo3 = { workspace = true, optional = true } rayon = { workspace = true } +regex = { workspace = true } tokio = { workspace = true, optional = true } [dev-dependencies] @@ -226,6 +227,7 @@ dynamic_group_by = [ ewma = ["polars-plan/ewma"] ewma_by = ["polars-plan/ewma_by"] dot_diagram = ["polars-plan/dot_diagram"] +mermaid_diagram = ["polars-plan/mermaid_diagram"] diagonal_concat = [] unique_counts = ["polars-plan/unique_counts"] log = ["polars-plan/log"] @@ -382,6 +384,7 @@ features = [ "list_to_struct", "log", "merge_sorted", + "mermaid_diagram", "meta", "mode", "moment", diff --git a/crates/polars-lazy/src/lib.rs b/crates/polars-lazy/src/lib.rs index f3dff5710170..8bf61d4ed752 100644 --- a/crates/polars-lazy/src/lib.rs +++ b/crates/polars-lazy/src/lib.rs @@ -200,13 +200,16 @@ #![cfg_attr(docsrs, feature(doc_auto_cfg))] extern crate core; -#[cfg(feature = "dot_diagram")] +#[cfg(any(feature = "dot_diagram", feature = "mermaid_diagram"))] mod dot; pub mod dsl; pub mod frame; pub mod physical_plan; pub mod prelude; +#[cfg(feature = "mermaid_diagram")] +mod mermaid; + mod scan; #[cfg(test)] mod tests; diff --git a/crates/polars-lazy/src/mermaid.rs b/crates/polars-lazy/src/mermaid.rs new file mode 100644 index 000000000000..227aa40e7cd2 --- /dev/null +++ b/crates/polars-lazy/src/mermaid.rs @@ -0,0 +1,48 @@ +use polars_core::prelude::*; +use regex::Regex; + +use crate::prelude::*; + +impl LazyFrame { + pub fn to_mermaid(&self, optimized: bool) -> PolarsResult { + // Mermaid strings are very similar to dot strings, so + // we can reuse the dot implementation. + let dot = self.to_dot(optimized)?; + + let edge_regex = Regex::new(r"(?P\w+) -- (?P\w+)").unwrap(); + let node_regex = Regex::new(r#"(?P\w+)(\s+)?\[label="(?P