display plain output for explain queries in duckdb (#6439)

Light2Dark · web-flow · commit d7ca5cb2717c · 2025-09-19T10:25:51.000-04:00
## 📝 Summary  polars: <img width="526" height="344" alt="CleanShot 2025-09-19 at 18 12 33" src="https://github.com/user-attachments/assets/9dbd621e-db1e-464e-9143-a73ba4e0fe75" /> pandas: <img width="309" height="235" alt="CleanShot 2025-09-19 at 18 13 22" src="https://github.com/user-attachments/assets/4749035b-8857-45f6-b6dc-5ff34bc5ef1e" /> ## 🔍 Description of Changes  ## 📋 Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [ ] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [x] I have added tests for the changes made. - [x] I have run the code and verified that it works as expected.
diff --git a/marimo/_sql/sql.py b/marimo/_sql/sql.py
@@ -12,7 +12,11 @@
 from marimo._sql.engines.sqlalchemy import SQLAlchemyEngine
 from marimo._sql.engines.types import QueryEngine
 from marimo._sql.get_engines import SUPPORTED_ENGINES
-from marimo._sql.utils import raise_df_import_error
+from marimo._sql.utils import (
+    extract_explain_content,
+    is_explain_query,
+    raise_df_import_error,
+)
 from marimo._types.ids import VariableName
 from marimo._utils.narwhals_utils import can_narwhalify_lazyframe
 
@@ -103,9 +107,14 @@ def sql(
             raise_df_import_error("polars[pyarrow]")
 
     if output:
+        from marimo._plugins.stateless.plain_text import plain_text
         from marimo._plugins.ui._impl import table
 
-        if can_narwhalify_lazyframe(df):
+        if isinstance(sql_engine, DuckDBEngine) and is_explain_query(query):
+            # For EXPLAIN queries in DuckDB, display plain output to preserve box drawings
+            text_output = extract_explain_content(df)
+            t = plain_text(text_output)
+        elif can_narwhalify_lazyframe(df):
             # For pl.LazyFrame and DuckDBRelation, we only show the first few rows
             # to avoid loading all the data into memory.
             # Also preload the first page of data without user confirmation.
diff --git a/marimo/_sql/utils.py b/marimo/_sql/utils.py
@@ -176,3 +176,44 @@ def sql_type_to_data_type(type_str: str) -> DataType:
         return "string"
     else:
         return "string"
+
+
+def is_explain_query(query: str) -> bool:
+    """Check if a SQL query is an EXPLAIN query."""
+    return query.lstrip().lower().startswith("explain ")
+
+
+def extract_explain_content(df: Any) -> str:
+    """Extract all content from a DataFrame for EXPLAIN queries.
+
+    Args:
+        df: DataFrame (pandas or polars). If not pandas / polars, return repr(df).
+
+    Returns:
+        String containing content of dataframe
+    """
+    try:
+        if DependencyManager.polars.has():
+            import polars as pl
+
+            if isinstance(df, pl.LazyFrame):
+                df = df.collect()
+            if isinstance(df, pl.DataFrame):
+                # Display full strings without truncation
+                with pl.Config(fmt_str_lengths=1000):
+                    return str(df)
+
+        if DependencyManager.pandas.has():
+            import pandas as pd
+
+            if isinstance(df, pd.DataFrame):
+                # Preserve newlines in the data
+                all_values = df.values.flatten().tolist()
+                return "\n".join(str(val) for val in all_values)
+
+        # Fallback to repr for other types
+        return repr(df)
+
+    except Exception as e:
+        LOGGER.debug("Failed to extract explain content: %s", e)
+        return repr(df)
diff --git a/tests/_sql/test_sql.py b/tests/_sql/test_sql.py
@@ -13,6 +13,7 @@
 from marimo._sql.engines.ibis import IbisEngine
 from marimo._sql.engines.sqlalchemy import SQLAlchemyEngine
 from marimo._sql.sql import _query_includes_limit, sql
+from marimo._sql.utils import extract_explain_content, is_explain_query
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -360,3 +361,155 @@ def test_sql_with_ibis_expression_result():
     with patch.object(IbisEngine, "sql_output_format", return_value="native"):
         result = sql("SELECT * FROM test", engine=duckdb_backend)
         assert isinstance(result, Expr)
+
+
+class TestExplainQueries:
+    def test_is_explain_query(self):
+        """Test is_explain_query function."""
+        # Test valid EXPLAIN queries
+        assert is_explain_query("EXPLAIN SELECT 1")
+        assert is_explain_query("explain SELECT 1")
+        assert is_explain_query("  EXPLAIN SELECT 1")
+        assert is_explain_query("\tEXPLAIN SELECT 1")
+        assert is_explain_query("\nEXPLAIN SELECT 1")
+        assert is_explain_query("EXPLAIN (FORMAT JSON) SELECT 1")
+        assert is_explain_query("EXPLAIN ANALYZE SELECT 1")
+        assert is_explain_query("EXPLAIN QUERY PLAN SELECT 1")
+
+        # Test non-EXPLAIN queries
+        assert not is_explain_query("SELECT 1")
+        assert not is_explain_query("INSERT INTO t VALUES (1)")
+        assert not is_explain_query("UPDATE t SET col = 1")
+        assert not is_explain_query("DELETE FROM t")
+        assert not is_explain_query("CREATE TABLE t (id INT)")
+        assert not is_explain_query("")
+        assert not is_explain_query("   ")
+
+        # Test edge cases
+        assert not is_explain_query("EXPLAINED")  # Not exactly "explain"
+        assert not is_explain_query("EXPLAINING")  # Not exactly "explain"
+        assert not is_explain_query(
+            "-- EXPLAIN SELECT 1"
+        )  # Comment, not actual query
+
+    @pytest.fixture
+    def explain_df_data(self) -> dict[str, list[str]]:
+        return {
+            "explain_key": ["physical_plan", "logical_plan"],
+            "explain_value": [
+                "┌─────────────────────────────────────┐\n│              PROJECTION               │\n└─────────────────────────────────────┘",
+                "┌─────────────────────────────────────┐\n│              SELECTION                │\n└─────────────────────────────────────┘",
+            ],
+        }
+
+    @pytest.mark.skipif(not HAS_POLARS, reason="Polars not installed")
+    def test_extract_explain_content_polars(
+        self, explain_df_data: dict[str, list[str]]
+    ):
+        """Test extract_explain_content with polars DataFrames."""
+        import polars as pl
+
+        # Test with regular DataFrame
+        df = pl.DataFrame(explain_df_data)
+
+        expected_rendering = """shape: (2, 2)
+┌───────────────┬───────────────────────────────────────────┐
+│ explain_key   ┆ explain_value                             │
+│ ---           ┆ ---                                       │
+│ str           ┆ str                                       │
+╞═══════════════╪═══════════════════════════════════════════╡
+│ physical_plan ┆ ┌─────────────────────────────────────┐   │
+│               ┆ │              PROJECTION               │ │
+│               ┆ └─────────────────────────────────────┘   │
+│ logical_plan  ┆ ┌─────────────────────────────────────┐   │
+│               ┆ │              SELECTION                │ │
+│               ┆ └─────────────────────────────────────┘   │
+└───────────────┴───────────────────────────────────────────┘"""
+
+        result = extract_explain_content(df)
+        assert result == expected_rendering
+
+        # Test with LazyFrame
+        lazy_df = df.lazy()
+        result = extract_explain_content(lazy_df)
+        assert result == expected_rendering
+
+    @pytest.mark.skipif(not HAS_PANDAS, reason="Pandas not installed")
+    def test_extract_explain_content_pandas(
+        self, explain_df_data: dict[str, list[str]]
+    ):
+        """Test extract_explain_content with pandas DataFrames."""
+        import pandas as pd
+
+        df = pd.DataFrame(explain_df_data)
+
+        result = extract_explain_content(df)
+        assert (
+            result
+            == """physical_plan
+┌─────────────────────────────────────┐
+│              PROJECTION               │
+└─────────────────────────────────────┘
+logical_plan
+┌─────────────────────────────────────┐
+│              SELECTION                │
+└─────────────────────────────────────┘"""
+        )
+
+    def test_extract_explain_content_fallback(self):
+        """Test extract_explain_content fallback for non-DataFrame objects."""
+        # Test with non-DataFrame object
+        result = extract_explain_content("not a dataframe")
+        assert isinstance(result, str)
+        assert "not a dataframe" in result
+
+        # Test with None
+        result = extract_explain_content(None)
+        assert isinstance(result, str)
+        assert "None" in result
+
+    @pytest.mark.skipif(not HAS_POLARS, reason="Polars not installed")
+    def test_extract_explain_content_error_handling(self):
+        """Test extract_explain_content error handling."""
+
+        # Create a mock DataFrame that will raise an error
+        class MockDataFrame:
+            def __init__(self):
+                pass
+
+            def __str__(self):
+                raise RuntimeError("Test error")
+
+        mock_df = MockDataFrame()
+        result = extract_explain_content(mock_df)
+        assert isinstance(result, str)
+        assert "MockDataFrame" in result  # Should fallback to repr
+
+    @patch("marimo._sql.sql.replace")
+    @pytest.mark.skipif(
+        not HAS_POLARS or not HAS_DUCKDB, reason="polars and duckdb required"
+    )
+    def test_sql_explain_query_display(self, mock_replace):
+        """Test that EXPLAIN queries are displayed as plain text."""
+        import duckdb
+        import polars as pl
+
+        # Create a test table
+        duckdb.sql(
+            "CREATE OR REPLACE TABLE test_explain AS SELECT * FROM range(5)"
+        )
+
+        # Test EXPLAIN query
+        result = sql("EXPLAIN SELECT * FROM test_explain")
+        assert isinstance(result, pl.DataFrame)
+
+        # Should call replace with plain_text
+        mock_replace.assert_called_once()
+        call_args = mock_replace.call_args[0][0]
+
+        # The call should be a plain_text object
+        assert hasattr(call_args, "text")
+        assert isinstance(call_args.text, str)
+
+        # Clean up
+        duckdb.sql("DROP TABLE test_explain")
