Repository: marimo-team/marimo

Issue:
Polars LazyFrame `show_graph` has poor graphic quality in Marimo
### Description

I am submitting a PR for this suggestion.

When using polars LazyFrames it is often desirable to display the query plan prior to collecting the query. There are multiple ways to do this in Marimo. 

1. Leave the LazyFrame object as the last line in a cell.
   - This works well and displays a nice SVG image in the cell (as long as graphviz is installed and on the path). ![Image](https://github.com/user-attachments/assets/0756de45-05ba-4565-8f65-214855821d01)
3. Manually call the `show_graph` function    
   - This is necessary if one wants to show the optimized version of the query plan
   - Unfortunately, this produces a very poor output if matplotlib is installed or an error if it is not![Image](https://github.com/user-attachments/assets/2e39572c-ab10-48f6-a372-5cb237b628fd)
3. Use `explain` to output a text-based query plan
   - This works well and Marimo even formats the text (whereas Jupyter notebooks by default show the raw string) 
   - However, it can be harder to follow than the graph for more complicated queries. ![Image](https://github.com/user-attachments/assets/dd21baf3-6fa7-4228-af85-9ef3c6b89712)

It would be nice if Marimo could display high quality graphics for the query plan even when the user wants to view the optimized plan and/or doesn't have matplotlib / graphviz installed. It would also be nice if the graph followed the theme of the notebook (dark / light) and could nicely display plans that have many terms (e.g. a query that transforms many columns in one step).


### Suggested solution

My suggestion is to register a polars extension that adds a marimo `mo` namespace to LazyFrames and allows users to display high quality query plan graphs in all situations with only a slight change to their code. This added polars extension code will be maintained within the marimo repository so that it will not require adding any code to the polars repository. 

Because marimo already has support for displaying mermaid graphs and polars can return the raw text defining the graph (in dot notation) it makes sense to parse this and convert to mermaid.

Once added this approach would allow for results that:

1. Are natively high quality for all means of showing the graph
2. Wrap the text in long columns
3. Do not require graphviz or matplotlib to be installed
4. Adopt to the user theme by default

| Case | Current | Proposed | Notes |
|---|---|---|---|
|Show Graph|![Image](https://github.com/user-attachments/assets/e1f18cff-33b5-45d7-8d54-98b496010bb2)| ![Image](https://github.com/user-attachments/assets/f9eff464-f5ec-4e3c-85d6-27e20d301c42) |
| Light Mode Example | ![Image](https://github.com/user-attachments/assets/9e3a9710-1c68-40b7-b02b-5cad1ca77ff0)|![Image](https://github.com/user-attachments/assets/4fdd82d4-77a5-45ab-8b3c-22373c9c1ee7)|
| LazyFrame last line | ![Image](https://github.com/user-attachments/assets/6ddf5a97-b8a9-4254-a8d7-1d0e710d0a82)|![Image](https://github.com/user-attachments/assets/4436b0ed-46fa-4486-b075-9e2e60c5a231)| No Change here, would likely require a change to polars unless marimo has a way to change the behavior of objects as the last line.|
| Join Graph | ![Image](https://github.com/user-attachments/assets/4299679c-7d7a-4bea-b139-76036a79b046)|![Image](https://github.com/user-attachments/assets/b9933924-15e6-40ed-8216-7404e68bac30)| Current behavior sometimes does not fit on the screen if the width is not "right"| 
| Wide example |![Image](https://github.com/user-attachments/assets/94ae0c3f-fac9-4840-b649-e7ec5d6aa9d3)| ![Image](https://github.com/user-attachments/assets/66880367-2daa-45c3-b694-d538ca077b64)| Proposed shown with output expanded|
 

### Alternative

Instead of outputting meramaid code it might be possible to recreate polars `_repr_html_` in order to display a high-quality image of the graph. However, this would require the user to have dot installed, would not change color with the user theme, and would not give the line wrapping behavior for wide graphs.  

### Additional context

I think ultimately the best solution would be to have `show_graph` and `_repr_html_` in polars recognize that they are in marimo and change their behavior. This would require their cooperation though, which is of unknown likelihood to me. I have written the implementation for the PR such that it would be easy for them to call marimo's functions if they detect a marimo environment (similar to what they already do for notebooks).

This snippet could be added in to the polars `display_dot_graph` function after the `raw_output` check to get the same behavior when using polars' `show_graph` function.

```python
try:
    from marimo import running_in_notebook
    from marimo._polars.lazyframe import Marimo 

    if running_in_notebook():
        return Marimo._dot_to_mermaid_html(dot)
except ImportError:
    pass
```

Discussion:
@mscolnick: Hey @kjgoodrick this seems like a great idea/addition. You can actually easily add that logic to https://github.com/marimo-team/marimo/blob/main/marimo/_output/formatters/df_formatters.py

You can see what we do for dataframes (just call table()) and just call mermaid() for the lazyframe 

Solution patch:
diff --git a/marimo/_output/formatters/df_formatters.py b/marimo/_output/formatters/df_formatters.py
index bc0dc69e..16a03c28 100644
--- a/marimo/_output/formatters/df_formatters.py
+++ b/marimo/_output/formatters/df_formatters.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import os
+import re
 
 import narwhals.stable.v1 as nw
 
@@ -11,7 +12,7 @@ from marimo._output.formatters.formatter_factory import (
     FormatterFactory,
     Unregister,
 )
-from marimo._output.md import md
+from marimo._plugins.stateless.mermaid import mermaid
 from marimo._plugins.ui._impl import tabs
 from marimo._plugins.ui._impl.table import get_default_table_page_size, table
 from marimo._runtime.patches import patch_polars_write_json
@@ -34,6 +35,43 @@ def include_opinionated() -> bool:
     return True
 
 
+def polars_dot_to_mermaid(dot: str) -> str:
+    """Converts polars DOT query plan renderings to mermaid.
+
+    Adapted from https://github.com/pola-rs/polars/pull/20607
+    Note: Not comprehensive, only handles components of the dot language used by polars.
+    """
+
+    edge_regex = r"(?P<node1>\w+) -- (?P<node2>\w+)"
+    node_regex = r"(?P<node>\w+)(\s+)?\[label=\"(?P<label>.*)\"]"
+
+    nodes = re.finditer(node_regex, dot)
+    edges = re.finditer(edge_regex, dot)
+
+    mermaid_str = "\n".join(
+        [
+            "graph TD",
+            *[f'\t{n["node"]}["{n["label"]}"]' for n in nodes],
+            *[f"\t{e['node1']} --- {e['node2']}" for e in edges],
+        ]
+    )
+
+    # replace [https://...] with <a> tags to avoid Mermaid interpreting it as markdown
+    mermaid_str = re.sub(
+        r"\[(https?://[^\]]+)\]",
+        lambda m: f"[<a href='{m.group(1)}'>{m.group(1)}</a>]",
+        mermaid_str,
+    )
+
+    # replace escaped newlines
+    mermaid_str = mermaid_str.replace(r"\n", "\n")
+
+    # replace escaped quotes
+    mermaid_str = mermaid_str.replace(r"\"", "#quot;")
+
+    return mermaid_str
+
+
 class PolarsFormatter(FormatterFactory):
     @staticmethod
     def package_name() -> str:
@@ -80,8 +118,14 @@ class PolarsFormatter(FormatterFactory):
         ) -> tuple[KnownMimeType, str]:
             return tabs.tabs(
                 {
+                    # NB(Trevor): Use `optimized=True` to match other methods' defaults (`show_graph`, `collect`).
+                    # The _repr_html_ uses `optimized=False`, but "cost" is probably minimal and this is more
+                    # accurate/opinionated default.
+                    # See: https://github.com/pola-rs/polars/blob/911352/py-polars/polars/lazyframe/frame.py#L773-L790
+                    "Query plan": mermaid(
+                        polars_dot_to_mermaid(df._ldf.to_dot(optimized=True))
+                    ),
                     "Table": table.lazy(df),
-                    "Query plan": md(df._repr_html_()),
                 }
             )._mime_()
 


Test patch:
diff --git a/tests/_output/formatters/test_formatters.py b/tests/_output/formatters/test_formatters.py
index 9efec01c..8fc665f4 100644
--- a/tests/_output/formatters/test_formatters.py
+++ b/tests/_output/formatters/test_formatters.py
@@ -10,6 +10,7 @@ import pytest
 
 from marimo._dependencies.dependencies import DependencyManager
 from marimo._messaging.mimetypes import KnownMimeType
+from marimo._output.formatters.df_formatters import polars_dot_to_mermaid
 from marimo._output.formatters.formatters import register_formatters
 from marimo._output.formatting import (
     Plain,
@@ -457,3 +458,65 @@ def test_display_protocol_takes_precedence() -> None:
     mime, content = formatter(obj)
     assert mime == "text/html"
     assert content == "<span>foo</span>"
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
+def test_lazyframe_renders_mermaid_html() -> None:
+    register_formatters()
+
+    import polars as pl
+
+    ldf = pl.DataFrame({"A": [1, 2, 3], "B": ["a", "a", "a"]}).lazy()
+
+    formatter = get_formatter(ldf)
+    assert formatter is not None
+    mimetype, contents = formatter(ldf)
+    assert mimetype == "text/html"
+    assert "<marimo-mermaid" in contents
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
+def test_polars_dot_to_mermaid() -> None:
+    import polars as pl
+
+    ldf = pl.DataFrame({"A": [1, 2, 3], "B": ["a", "a", "a"]}).lazy()
+
+    result = polars_dot_to_mermaid(ldf._ldf.to_dot(optimized=True))
+    assert result == 'graph TD\n\tp1["TABLE\nπ */2"]'
+
+
+@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
+def test_polars_dot_to_mermaid_complex() -> None:
+    import polars as pl
+
+    ldf = (
+        pl.DataFrame({"A": [1, 2, 3], "B": ["a", "b", "c"]})
+        .lazy()
+        .filter(pl.col("A") > 1)
+        .select(pl.col("A"), pl.col("B").alias("B_renamed"))
+        .join(pl.DataFrame({"A": [1]}).lazy(), on="A")
+    )
+
+    assert (
+        polars_dot_to_mermaid(ldf._ldf.to_dot(optimized=True))
+        == """graph TD
+\tp4["TABLE\nπ 2/2"]
+\tp3["FILTER BY [(col(#quot;A#quot;)) > (1)]"]
+\tp2["π 2/2"]
+\tp5["TABLE\nπ */1"]
+\tp1["JOIN INNER\nleft: [col(#quot;A#quot;)];\nright: [col(#quot;A#quot;)]"]
+\tp1 --- p2
+\tp2 --- p3
+\tp3 --- p4
+\tp1 --- p5"""
+    )
+
+
+def test_polars_dot_to_mermaid_handles_urls() -> None:
+    assert (
+        polars_dot_to_mermaid("""graph polars_query {
+  p1[label="Check [https://example.com] for more"]
+}""")
+        == """graph TD
+\tp1[\"Check [<a href='https://example.com'>https://example.com</a>] for more\"]"""
+    )