diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py index 055b954a979d..560542d6ea86 100644 --- a/py-polars/polars/lazyframe/frame.py +++ b/py-polars/polars/lazyframe/frame.py @@ -99,6 +99,7 @@ ClosedInterval, ColumnNameOrSelector, CsvQuoteStyle, + ExplainFormat, FillNullStrategy, FrameInitTypes, IntoExpr, @@ -881,6 +882,7 @@ def skip_minmax(dt: PolarsDataType) -> bool: def explain( self, *, + format: ExplainFormat = "plain", optimized: bool = True, type_coercion: bool = True, predicate_pushdown: bool = True, @@ -891,7 +893,7 @@ def explain( comm_subexpr_elim: bool = True, cluster_with_columns: bool = True, streaming: bool = False, - tree_format: bool = False, + tree_format: bool | None = None, ) -> str: """ Create a string representation of the query plan. @@ -900,6 +902,8 @@ def explain( Parameters ---------- + format : {'plain', 'tree'} + The format to use for displaying the logical plan. optimized Return an optimized query plan. Defaults to `True`. If this is set to `True` the subsequent @@ -924,7 +928,10 @@ def explain( streaming Run parts of the query in a streaming fashion (this is in an alpha state) tree_format - Format the output as a tree + Format the output as a tree. + + .. deprecated:: 0.20.30 + Use `format="tree"` instead. Examples -------- @@ -939,6 +946,15 @@ def explain( ... "a" ... ).explain() # doctest: +SKIP """ + if tree_format is not None: + issue_deprecation_warning( + "The `tree_format` parameter for `LazyFrame.explain` is deprecated" + " Use the `format` parameter instead.", + version="0.20.30", + ) + if tree_format: + format = "tree" + if optimized: ldf = self._ldf.optimization_toggle( type_coercion, @@ -952,13 +968,15 @@ def explain( streaming, _eager=False, ) - if tree_format: + if format == "tree": return ldf.describe_optimized_plan_tree() - return ldf.describe_optimized_plan() + else: + return ldf.describe_optimized_plan() - if tree_format: + if format == "tree": return self._ldf.describe_plan_tree() - return self._ldf.describe_plan() + else: + return self._ldf.describe_plan() def show_graph( self, diff --git a/py-polars/polars/type_aliases.py b/py-polars/polars/type_aliases.py index d65e36be40ea..924afb0f306e 100644 --- a/py-polars/polars/type_aliases.py +++ b/py-polars/polars/type_aliases.py @@ -169,6 +169,7 @@ TorchExportType: TypeAlias = Literal["tensor", "dataset", "dict"] TransferEncoding: TypeAlias = Literal["hex", "base64"] WindowMappingStrategy: TypeAlias = Literal["group_to_rows", "join", "explode"] +ExplainFormat: TypeAlias = Literal["plain", "tree"] # type signature for allowed frame init FrameInitTypes: TypeAlias = Union[ diff --git a/py-polars/tests/unit/lazyframe/test_explain.py b/py-polars/tests/unit/lazyframe/test_explain.py new file mode 100644 index 000000000000..9906d9db8e71 --- /dev/null +++ b/py-polars/tests/unit/lazyframe/test_explain.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import pytest + +import polars as pl + + +def test_lf_explain() -> None: + lf = pl.LazyFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) + plan = lf.select("a").select(pl.col("a").sum() + pl.len()) + + result = plan.explain() + + expected = """\ + SELECT [[(col("a").sum()) + (len().cast(Int64))]] FROM + DF ["a", "b"]; PROJECT 1/2 COLUMNS; SELECTION: None\ +""" + assert result == expected + + +def test_lf_explain_format_tree() -> None: + lf = pl.LazyFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) + plan = lf.select("a").select(pl.col("a").sum() + pl.len()) + + result = plan.explain(format="tree") + + expected = """\ + 0 1 + ┌───────────────────────────────────────────────── + │ + │ ╭────────╮ + 0 │ │ SELECT │ + │ ╰───┬┬───╯ + │ ││ + │ │╰───────────────────────╮ + │ │ │ + │ ╭─────────┴──────────╮ │ + │ │ expression: │ ╭──────────┴──────────╮ + │ │ [(col("a") │ │ FROM: │ + 1 │ │ .sum()) + (len() │ │ DF ["a", "b"] │ + │ │ .cast(Int64))] │ │ PROJECT 1/2 COLUMNS │ + │ ╰────────────────────╯ ╰─────────────────────╯ +\ +""" + assert result == expected + + +def test_lf_explain_tree_format_deprecated() -> None: + lf = pl.LazyFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) + + with pytest.deprecated_call(): + lf.explain(tree_format=True)