From e74f838f1b17d1487de142ec729a120469d4f480 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <kyle@goodricks.net>
Date: Mon, 13 Jan 2025 22:19:09 -0700
Subject: [PATCH 1/6] Add mermaid generation to rust

---
 crates/polars-lazy/Cargo.toml                 |  2 +
 crates/polars-lazy/src/lib.rs                 |  5 ++-
 crates/polars-lazy/src/mermaid.rs             | 40 +++++++++++++++++++
 crates/polars-plan/Cargo.toml                 |  2 +
 crates/polars-python/Cargo.toml               |  1 +
 crates/polars-python/src/lazyframe/general.rs |  5 +++
 crates/polars/Cargo.toml                      |  2 +
 crates/polars/src/lib.rs                      |  1 +
 docs/source/user-guide/installation.md        |  1 +
 9 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 crates/polars-lazy/src/mermaid.rs

diff --git a/crates/polars-lazy/Cargo.toml b/crates/polars-lazy/Cargo.toml
index ed172a178cb1..866cd61111c4 100644
--- a/crates/polars-lazy/Cargo.toml
+++ b/crates/polars-lazy/Cargo.toml
@@ -226,6 +226,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 +383,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..496213282d46
--- /dev/null
+++ b/crates/polars-lazy/src/mermaid.rs
@@ -0,0 +1,40 @@
+use polars_core::{export::regex, prelude::*};
+
+use crate::prelude::*;
+
+impl LazyFrame {
+    pub fn to_mermaid(&self, optimized: bool) -> PolarsResult<String> {
+        // 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::Regex::new(r"(?P<node1>\w+) -- (?P<node2>\w+)").unwrap();
+        let node_regex = regex::Regex::new(r#"(?P<node>\w+)(\s+)?\[label="(?P<label>.*)"]"#).unwrap();
+
+        let nodes = node_regex.captures_iter(&dot);
+        let edges = edge_regex.captures_iter(&dot);
+        
+        let node_lines = nodes.map(|node| {
+            format!(
+                "\t{}[\"{}\"]", 
+                node.name("node").unwrap().as_str(), 
+                node.name("label").unwrap().as_str()
+                .replace(r"\n", "\n") // replace escaped newlines
+                .replace(r#"\""#, "#quot;") // replace escaped quotes
+            )
+        }).collect::<Vec<_>>().join("\n");
+
+        let edge_lines = edges.map(|edge| {
+            format!(
+                "\t{} --- {}", 
+                edge.name("node1").unwrap().as_str(), 
+                edge.name("node2").unwrap().as_str()
+            )
+        }).collect::<Vec<_>>().join("\n");
+
+        let mermaid = format!("graph TD\n{node_lines}\n{edge_lines}");
+        
+        Ok(mermaid)
+
+    }
+}
\ No newline at end of file
diff --git a/crates/polars-plan/Cargo.toml b/crates/polars-plan/Cargo.toml
index 8f93faeebae8..c3709b19aada 100644
--- a/crates/polars-plan/Cargo.toml
+++ b/crates/polars-plan/Cargo.toml
@@ -148,6 +148,7 @@ dynamic_group_by = ["polars-core/dynamic_group_by"]
 ewma = ["polars-ops/ewma"]
 ewma_by = ["polars-ops/ewma_by"]
 dot_diagram = []
+mermaid_diagram = []
 unique_counts = ["polars-ops/unique_counts"]
 log = ["polars-ops/log"]
 chunked_ids = []
@@ -291,6 +292,7 @@ features = [
   "concat_str",
   "coalesce",
   "dot_diagram",
+  "mermaid_diagram",
   "trigonometry",
   "streaming",
   "true_div",
diff --git a/crates/polars-python/Cargo.toml b/crates/polars-python/Cargo.toml
index 14c918b44097..b0a94eec6d82 100644
--- a/crates/polars-python/Cargo.toml
+++ b/crates/polars-python/Cargo.toml
@@ -84,6 +84,7 @@ features = [
   "array_arithmetic",
   "array_to_struct",
   "log",
+  "mermaid_diagram",
   "mode",
   "moment",
   "ndarray",
diff --git a/crates/polars-python/src/lazyframe/general.rs b/crates/polars-python/src/lazyframe/general.rs
index ea6c39f8a2b9..67d31ef34e15 100644
--- a/crates/polars-python/src/lazyframe/general.rs
+++ b/crates/polars-python/src/lazyframe/general.rs
@@ -481,6 +481,11 @@ impl PyLazyFrame {
         Ok(result)
     }
 
+    fn to_mermaid(&self, optimized: bool) -> PyResult<String> {
+        let result = self.ldf.to_mermaid(optimized).map_err(PyPolarsErr::from)?;
+        Ok(result)
+    }
+
     fn optimization_toggle(
         &self,
         type_coercion: bool,
diff --git a/crates/polars/Cargo.toml b/crates/polars/Cargo.toml
index d256582503cf..6df9f0497b21 100644
--- a/crates/polars/Cargo.toml
+++ b/crates/polars/Cargo.toml
@@ -192,6 +192,7 @@ array_arithmetic = ["polars-core/array_arithmetic", "dtype-array"]
 array_to_struct = ["polars-ops/array_to_struct", "polars-lazy?/array_to_struct"]
 log = ["polars-ops/log", "polars-lazy?/log"]
 merge_sorted = ["polars-lazy?/merge_sorted"]
+mermaid_diagram = ["polars-lazy?/mermaid_diagram"]
 meta = ["polars-lazy?/meta"]
 mode = ["polars-ops/mode", "polars-lazy?/mode"]
 moment = ["polars-ops/moment", "polars-lazy?/moment"]
@@ -415,6 +416,7 @@ docs-selection = [
   "diagonal_concat",
   "abs",
   "dot_diagram",
+  "mermaid_diagram",
   "string_encoding",
   "product",
   "to_dummies",
diff --git a/crates/polars/src/lib.rs b/crates/polars/src/lib.rs
index 3ac56332cd01..cc8fb583cbaf 100644
--- a/crates/polars/src/lib.rs
+++ b/crates/polars/src/lib.rs
@@ -177,6 +177,7 @@
 //! * `lazy` - Lazy API
 //!     - `regex` - Use regexes in [column selection]
 //!     - `dot_diagram` - Create dot diagrams from lazy logical plans.
+//!     - `mermaid` - Create mermaid diagrams from lazy logical plans.
 //! * `sql` - Pass SQL queries to Polars.
 //! * `streaming` - Process datasets larger than RAM.
 //! * `random` - Generate arrays with randomly sampled values
diff --git a/docs/source/user-guide/installation.md b/docs/source/user-guide/installation.md
index 52882ee599ff..ac7219bab858 100644
--- a/docs/source/user-guide/installation.md
+++ b/docs/source/user-guide/installation.md
@@ -181,6 +181,7 @@ The opt-in features are:
 - `lazy` - Lazy API:
     - `regex` - Use regexes in column selection.
     - `dot_diagram` - Create dot diagrams from lazy logical plans.
+    - `mermaid_diagram` - Create mermaid diagrams from lazy logical plans.
 - `sql` - Pass SQL queries to Polars.
 - `streaming` - Be able to process datasets that are larger than RAM.
 - `random` - Generate arrays with randomly sampled values

From aba94064d8c5da4d975982a554db8d994077132a Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <kyle@goodricks.net>
Date: Mon, 13 Jan 2025 22:19:09 -0700
Subject: [PATCH 2/6] Add mermaid output

---
 Cargo.lock                          |  1 +
 crates/polars-lazy/Cargo.toml       |  1 +
 crates/polars-lazy/src/mermaid.rs   | 60 ++++++++++++++++-------------
 py-polars/polars/_typing.py         |  1 +
 py-polars/polars/_utils/various.py  | 12 ++++++
 py-polars/polars/lazyframe/frame.py | 45 +++++++++++++++++-----
 6 files changed, 84 insertions(+), 36 deletions(-)

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 866cd61111c4..98fd8f90a36e 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]
diff --git a/crates/polars-lazy/src/mermaid.rs b/crates/polars-lazy/src/mermaid.rs
index 496213282d46..227aa40e7cd2 100644
--- a/crates/polars-lazy/src/mermaid.rs
+++ b/crates/polars-lazy/src/mermaid.rs
@@ -1,40 +1,48 @@
-use polars_core::{export::regex, prelude::*};
+use polars_core::prelude::*;
+use regex::Regex;
 
 use crate::prelude::*;
 
 impl LazyFrame {
     pub fn to_mermaid(&self, optimized: bool) -> PolarsResult<String> {
-        // Mermaid strings are very similar to dot strings, so 
+        // 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::Regex::new(r"(?P<node1>\w+) -- (?P<node2>\w+)").unwrap();
-        let node_regex = regex::Regex::new(r#"(?P<node>\w+)(\s+)?\[label="(?P<label>.*)"]"#).unwrap();
+
+        let edge_regex = Regex::new(r"(?P<node1>\w+) -- (?P<node2>\w+)").unwrap();
+        let node_regex = Regex::new(r#"(?P<node>\w+)(\s+)?\[label="(?P<label>.*)"]"#).unwrap();
 
         let nodes = node_regex.captures_iter(&dot);
         let edges = edge_regex.captures_iter(&dot);
-        
-        let node_lines = nodes.map(|node| {
-            format!(
-                "\t{}[\"{}\"]", 
-                node.name("node").unwrap().as_str(), 
-                node.name("label").unwrap().as_str()
-                .replace(r"\n", "\n") // replace escaped newlines
-                .replace(r#"\""#, "#quot;") // replace escaped quotes
-            )
-        }).collect::<Vec<_>>().join("\n");
-
-        let edge_lines = edges.map(|edge| {
-            format!(
-                "\t{} --- {}", 
-                edge.name("node1").unwrap().as_str(), 
-                edge.name("node2").unwrap().as_str()
-            )
-        }).collect::<Vec<_>>().join("\n");
+
+        let node_lines = nodes
+            .map(|node| {
+                format!(
+                    "\t{}[\"{}\"]",
+                    node.name("node").unwrap().as_str(),
+                    node.name("label")
+                        .unwrap()
+                        .as_str()
+                        .replace(r"\n", "\n") // replace escaped newlines
+                        .replace(r#"\""#, "#quot;") // replace escaped quotes
+                )
+            })
+            .collect::<Vec<_>>()
+            .join("\n");
+
+        let edge_lines = edges
+            .map(|edge| {
+                format!(
+                    "\t{} --- {}",
+                    edge.name("node1").unwrap().as_str(),
+                    edge.name("node2").unwrap().as_str()
+                )
+            })
+            .collect::<Vec<_>>()
+            .join("\n");
 
         let mermaid = format!("graph TD\n{node_lines}\n{edge_lines}");
-        
-        Ok(mermaid)
 
+        Ok(mermaid)
     }
-}
\ No newline at end of file
+}
diff --git a/py-polars/polars/_typing.py b/py-polars/polars/_typing.py
index b362e2116502..1c0dbbd41a8e 100644
--- a/py-polars/polars/_typing.py
+++ b/py-polars/polars/_typing.py
@@ -188,6 +188,7 @@ def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
 TransferEncoding: TypeAlias = Literal["hex", "base64"]
 WindowMappingStrategy: TypeAlias = Literal["group_to_rows", "join", "explode"]
 ExplainFormat: TypeAlias = Literal["plain", "tree"]
+ShowGraphFormat: TypeAlias = Literal["dot", "mermaid"]
 
 # type signature for allowed frame init
 FrameInitTypes: TypeAlias = Union[
diff --git a/py-polars/polars/_utils/various.py b/py-polars/polars/_utils/various.py
index ef583250a254..05dc4e34bdbf 100644
--- a/py-polars/polars/_utils/various.py
+++ b/py-polars/polars/_utils/various.py
@@ -61,6 +61,8 @@
     P = ParamSpec("P")
     T = TypeVar("T")
 
+    from IPython.display import Markdown
+
 # note: reversed views don't match as instances of MappingView
 if sys.version_info >= (3, 11):
     _views: list[Reversible[Any]] = [{}.keys(), {}.values(), {}.items()]
@@ -702,3 +704,13 @@ def display_dot_graph(
         plt.imshow(img)
         plt.show()
         return None
+
+
+def display_mermaid_graph(mermaid: str) -> Markdown:
+    if not _in_notebook():
+        msg = "This function is only available in Jupyter notebooks."
+        raise OSError(msg)
+
+    from IPython.display import Markdown
+
+    return Markdown(f"```mermaid\n{mermaid}\n```")
diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
index 5453daa3995c..bbf960917d75 100644
--- a/py-polars/polars/lazyframe/frame.py
+++ b/py-polars/polars/lazyframe/frame.py
@@ -38,6 +38,7 @@
 from polars._utils.various import (
     _is_generator,
     display_dot_graph,
+    display_mermaid_graph,
     extend_bool,
     find_stacklevel,
     is_bool_sequence,
@@ -118,6 +119,7 @@
         SchemaDefinition,
         SchemaDict,
         SerializationFormat,
+        ShowGraphFormat,
         StartBy,
         UniqueKeepStrategy,
     )
@@ -1130,6 +1132,7 @@ def show_graph(
         show: bool = True,
         output_path: str | Path | None = None,
         raw_output: bool = False,
+        raw_output_format: ShowGraphFormat = "dot",
         figsize: tuple[float, float] = (16.0, 12.0),
         type_coercion: bool = True,
         _type_check: bool = True,
@@ -1147,8 +1150,9 @@ def show_graph(
         """
         Show a plot of the query plan.
 
-        Note that Graphviz must be installed to render the visualization (if not
-        already present, you can download it here: `<https://graphviz.org/download>`_).
+        Note that Graphviz must be installed to export the visualization
+        or show it outside of a notebook
+        (if not already present, you can download it here: `<https://graphviz.org/download>`_).
 
         Parameters
         ----------
@@ -1160,6 +1164,8 @@ def show_graph(
             Write the figure to disk.
         raw_output
             Return dot syntax. This cannot be combined with `show` and/or `output_path`.
+        raw_output_format
+            The format of the raw output.
         figsize
             Passed to matplotlib if `show == True`.
         type_coercion
@@ -1214,14 +1220,33 @@ def show_graph(
             new_streaming=False,
         )
 
-        dot = _ldf.to_dot(optimized)
-        return display_dot_graph(
-            dot=dot,
-            show=show,
-            output_path=output_path,
-            raw_output=raw_output,
-            figsize=figsize,
-        )
+        if raw_output:
+            if raw_output_format == "dot":
+                return _ldf.to_dot(optimized)
+            elif raw_output_format == "mermaid":
+                return _ldf.to_mermaid(optimized)
+            else:
+                msg = f'raw_output_format must be one of ["dot", "mermaid"] but got {raw_output_format}'
+                raise ValueError(msg)
+
+        try:
+            dot = _ldf.to_dot(optimized)
+            return display_dot_graph(
+                dot=dot,
+                show=show,
+                output_path=output_path,
+                figsize=figsize,
+            )
+        except (ImportError, FileNotFoundError) as e:
+            missing_graphviz_error = e
+
+        # if we reach this point, we failed to display the graph using Graphviz
+        # we should try as mermaid instead
+        try:
+            mermaid = _ldf.to_mermaid(optimized)
+            return display_mermaid_graph(mermaid)
+        except OSError:
+            raise missing_graphviz_error from None
 
     def inspect(self, fmt: str = "{}") -> LazyFrame:
         """

From 336d2b7ce061e1c1b2a02bf7fa5da473cc5b0b00 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <kyle@goodricks.net>
Date: Mon, 13 Jan 2025 22:19:09 -0700
Subject: [PATCH 3/6] formatting

---
 crates/polars-lazy/Cargo.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/polars-lazy/Cargo.toml b/crates/polars-lazy/Cargo.toml
index 98fd8f90a36e..d8e4b1cfe68b 100644
--- a/crates/polars-lazy/Cargo.toml
+++ b/crates/polars-lazy/Cargo.toml
@@ -29,7 +29,7 @@ memchr = { workspace = true }
 once_cell = { workspace = true }
 pyo3 = { workspace = true, optional = true }
 rayon = { workspace = true }
-regex = { workspace = true}
+regex = { workspace = true }
 tokio = { workspace = true, optional = true }
 
 [dev-dependencies]

From 1c9ad23295c61707c37b3bba260a80443df56f03 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <kyle@goodricks.net>
Date: Mon, 13 Jan 2025 22:19:09 -0700
Subject: [PATCH 4/6] Make mermaid and dot outputs more aligned.

mermaid `show_graph` now displays instead of returning the markdown object. It is now possible to export mermaid charts to file with `mmdc`. The user can now force polars to use graphviz or mermaid when showing the graph or outputting the files.
---
 py-polars/polars/_utils/various.py  | 43 ++++++++++++++++++++++---
 py-polars/polars/lazyframe/frame.py | 49 +++++++++++++++++------------
 2 files changed, 67 insertions(+), 25 deletions(-)

diff --git a/py-polars/polars/_utils/various.py b/py-polars/polars/_utils/various.py
index 05dc4e34bdbf..41c128502dc2 100644
--- a/py-polars/polars/_utils/various.py
+++ b/py-polars/polars/_utils/various.py
@@ -61,8 +61,6 @@
     P = ParamSpec("P")
     T = TypeVar("T")
 
-    from IPython.display import Markdown
-
 # note: reversed views don't match as instances of MappingView
 if sys.version_info >= (3, 11):
     _views: list[Reversible[Any]] = [{}.keys(), {}.values(), {}.items()]
@@ -706,11 +704,46 @@ def display_dot_graph(
         return None
 
 
-def display_mermaid_graph(mermaid: str) -> Markdown:
+def display_mermaid_graph(
+    *, mermaid: str, show: bool = True, output_path: str | Path | None = None
+) -> None:
+    # Make font monospace
+    mermaid += r'%%{init: {"fontFamily": "monospace"}}%%' + "\n"
+
+    if output_path:
+        output_path = Path(output_path)
+        try:
+            subprocess.check_output(
+                [
+                    "mmdc",
+                    "--output",
+                    output_path,
+                    "--outputFormat",
+                    output_path.suffix[1:],
+                    "--input",
+                    "-",
+                ],
+                input=f"{mermaid}".encode(),
+            )
+
+        except (ImportError, FileNotFoundError):
+            msg = (
+                "The mermaid-cli `mmdc` binary should be on your PATH. "
+                "(If not installed you can download here: https://github.com/mermaid-js/mermaid-cli)"
+            )
+
+            raise ImportError(msg) from None
+        except subprocess.CalledProcessError as e:
+            msg = "Error while calling `mmdc` ensure that your file extension is one of `mmdc`'s supported output types."
+            raise ValueError(msg) from e
+
+    if not show:
+        return None
+
     if not _in_notebook():
         msg = "This function is only available in Jupyter notebooks."
         raise OSError(msg)
 
-    from IPython.display import Markdown
+    from IPython.display import Markdown, display
 
-    return Markdown(f"```mermaid\n{mermaid}\n```")
+    display(Markdown(f"```mermaid\n{mermaid}\n```"))
diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
index bbf960917d75..5e1e7bf6129a 100644
--- a/py-polars/polars/lazyframe/frame.py
+++ b/py-polars/polars/lazyframe/frame.py
@@ -1132,7 +1132,7 @@ def show_graph(
         show: bool = True,
         output_path: str | Path | None = None,
         raw_output: bool = False,
-        raw_output_format: ShowGraphFormat = "dot",
+        format: ShowGraphFormat | None = None,
         figsize: tuple[float, float] = (16.0, 12.0),
         type_coercion: bool = True,
         _type_check: bool = True,
@@ -1150,9 +1150,11 @@ def show_graph(
         """
         Show a plot of the query plan.
 
-        Note that Graphviz must be installed to export the visualization
+        Note that Graphviz or mermaid must be installed to export the visualization
         or show it outside of a notebook
-        (if not already present, you can download it here: `<https://graphviz.org/download>`_).
+        (if not already present, you can download graphviz here: `<https://graphviz.org/download>`_
+        or mermaid here: `<https://github.com/mermaid-js/mermaid-cli>`_).
+
 
         Parameters
         ----------
@@ -1164,8 +1166,8 @@ def show_graph(
             Write the figure to disk.
         raw_output
             Return dot syntax. This cannot be combined with `show` and/or `output_path`.
-        raw_output_format
-            The format of the raw output.
+        format
+            The format of the graph output.
         figsize
             Passed to matplotlib if `show == True`.
         type_coercion
@@ -1221,30 +1223,37 @@ def show_graph(
         )
 
         if raw_output:
-            if raw_output_format == "dot":
+            if format is None or format == "dot":
                 return _ldf.to_dot(optimized)
-            elif raw_output_format == "mermaid":
+            elif format == "mermaid":
                 return _ldf.to_mermaid(optimized)
             else:
-                msg = f'raw_output_format must be one of ["dot", "mermaid"] but got {raw_output_format}'
+                msg = f'format must be one of ["dot", "mermaid"] but got {format}'
                 raise ValueError(msg)
 
-        try:
-            dot = _ldf.to_dot(optimized)
-            return display_dot_graph(
-                dot=dot,
-                show=show,
-                output_path=output_path,
-                figsize=figsize,
-            )
-        except (ImportError, FileNotFoundError) as e:
-            missing_graphviz_error = e
+        if format is None or format == "dot":
+            try:
+                dot = _ldf.to_dot(optimized)
+                return display_dot_graph(
+                    dot=dot,
+                    show=show,
+                    output_path=output_path,
+                    figsize=figsize,
+                )
+            except (ImportError, FileNotFoundError) as e:
+                missing_graphviz_error = e
+
+                # if the user specified "dot" raise immediately
+                if format == "dot":
+                    raise missing_graphviz_error from None
 
-        # if we reach this point, we failed to display the graph using Graphviz
+        # If we reach this point, we failed to display the graph using Graphviz
         # we should try as mermaid instead
         try:
             mermaid = _ldf.to_mermaid(optimized)
-            return display_mermaid_graph(mermaid)
+            return display_mermaid_graph(
+                mermaid=mermaid, show=show, output_path=output_path
+            )
         except OSError:
             raise missing_graphviz_error from None
 

From e24e1af03fe1cc7f8977cd3793507240a2884a09 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <kyle@goodricks.net>
Date: Mon, 13 Jan 2025 22:19:09 -0700
Subject: [PATCH 5/6] return display for mypy.

---
 py-polars/polars/_utils/various.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/py-polars/polars/_utils/various.py b/py-polars/polars/_utils/various.py
index 41c128502dc2..890087dc8dea 100644
--- a/py-polars/polars/_utils/various.py
+++ b/py-polars/polars/_utils/various.py
@@ -746,4 +746,4 @@ def display_mermaid_graph(
 
     from IPython.display import Markdown, display
 
-    display(Markdown(f"```mermaid\n{mermaid}\n```"))
+    return display(Markdown(f"```mermaid\n{mermaid}\n```"))

From 1b87720fff7f8bde51cff98df7513305e9c2a665 Mon Sep 17 00:00:00 2001
From: Kyle Goodrick <kyle@goodricks.net>
Date: Mon, 13 Jan 2025 22:19:09 -0700
Subject: [PATCH 6/6] Update `display_mermaid_graph` to match
 `display_dot_graph` including raw_output (mypy will not allow the function to
 return only None)

---
 py-polars/polars/_utils/various.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/py-polars/polars/_utils/various.py b/py-polars/polars/_utils/various.py
index 890087dc8dea..293f9153a436 100644
--- a/py-polars/polars/_utils/various.py
+++ b/py-polars/polars/_utils/various.py
@@ -705,8 +705,15 @@ def display_dot_graph(
 
 
 def display_mermaid_graph(
-    *, mermaid: str, show: bool = True, output_path: str | Path | None = None
-) -> None:
+    *,
+    mermaid: str,
+    show: bool = True,
+    output_path: str | Path | None = None,
+    raw_output: bool = False,
+) -> str | None:
+    if raw_output:
+        return mermaid
+
     # Make font monospace
     mermaid += r'%%{init: {"fontFamily": "monospace"}}%%' + "\n"