Skip to content

Commit b035c97

Browse files
timsaucerclaude
andcommitted
feat: add show_statistics, analyze_level, and analyze_categories to explain
Expose the remaining upstream ExplainOption fields as keywords on DataFrame.explain. Each defaults to None, which falls back to the matching datafusion.explain.* session setting, so existing calls are unaffected. New ExplainAnalyzeLevel and ExplainMetricCategory enums sit beside ExplainFormat. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 24764bd commit b035c97

3 files changed

Lines changed: 133 additions & 4 deletions

File tree

‎crates/core/src/dataframe.rs‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -844,13 +844,24 @@ impl PyDataFrame {
844844
}
845845

846846
/// Print the query plan
847-
#[pyo3(signature = (verbose=false, analyze=false, format=None))]
847+
#[pyo3(signature = (
848+
verbose=false,
849+
analyze=false,
850+
format=None,
851+
show_statistics=None,
852+
analyze_level=None,
853+
analyze_categories=None
854+
))]
855+
#[allow(clippy::too_many_arguments)]
848856
fn explain(
849857
&self,
850858
py: Python,
851859
verbose: bool,
852860
analyze: bool,
853861
format: Option<&str>,
862+
show_statistics: Option<bool>,
863+
analyze_level: Option<&str>,
864+
analyze_categories: Option<Vec<String>>,
854865
) -> PyDataFusionResult<()> {
855866
let explain_format = match format {
856867
Some(f) => f
@@ -860,10 +871,24 @@ impl PyDataFrame {
860871
})?,
861872
None => datafusion::common::format::ExplainFormat::Indent,
862873
};
874+
let analyze_level = analyze_level
875+
.map(|l| l.parse::<datafusion::common::format::MetricType>())
876+
.transpose()?;
877+
let analyze_categories = analyze_categories
878+
.map(|cats| {
879+
cats.iter()
880+
.map(|c| c.parse::<datafusion::common::format::MetricCategory>())
881+
.collect::<datafusion::common::Result<Vec<_>>>()
882+
.map(datafusion::common::format::ExplainAnalyzeCategories::Only)
883+
})
884+
.transpose()?;
863885
let opts = datafusion::logical_expr::ExplainOption::default()
864886
.with_verbose(verbose)
865887
.with_analyze(analyze)
866-
.with_format(explain_format);
888+
.with_format(explain_format)
889+
.with_show_statistics(show_statistics)
890+
.with_analyze_level(analyze_level)
891+
.with_analyze_categories(analyze_categories);
867892
let df = self.df.as_ref().clone().explain_with_options(opts)?;
868893
print_dataframe(py, df)
869894
}

‎python/datafusion/dataframe.py‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,32 @@ class ExplainFormat(Enum):
108108
"""Graphviz DOT format for graph rendering."""
109109

110110

111+
class ExplainAnalyzeLevel(Enum):
112+
"""Which metrics :py:meth:`DataFrame.explain` reports when ``analyze=True``."""
113+
114+
SUMMARY = "summary"
115+
"""Common metrics for finding which operator is slow."""
116+
117+
DEV = "dev"
118+
"""All metrics, including those for deep operator-level introspection."""
119+
120+
121+
class ExplainMetricCategory(Enum):
122+
"""Category of metric reported by :py:meth:`DataFrame.explain` with ``analyze``."""
123+
124+
ROWS = "rows"
125+
"""Row counts, such as ``output_rows``."""
126+
127+
BYTES = "bytes"
128+
"""Byte sizes, such as ``output_bytes``."""
129+
130+
TIMING = "timing"
131+
"""Elapsed times, such as ``elapsed_compute``."""
132+
133+
UNCATEGORIZED = "uncategorized"
134+
"""Metrics that declare no category."""
135+
136+
111137
# excerpt from deltalake
112138
# https://github.com/apache/datafusion-python/pull/981#discussion_r1905619163
113139
class Compression(Enum):
@@ -1207,6 +1233,9 @@ def explain(
12071233
verbose: bool = False,
12081234
analyze: bool = False,
12091235
format: ExplainFormat | None = None,
1236+
show_statistics: bool | None = None,
1237+
analyze_level: ExplainAnalyzeLevel | None = None,
1238+
analyze_categories: Iterable[ExplainMetricCategory] | None = None,
12101239
) -> None:
12111240
"""Print an explanation of the DataFrame's plan so far.
12121241
@@ -1217,6 +1246,13 @@ def explain(
12171246
analyze: If ``True``, the plan will run and metrics reported.
12181247
format: Output format for the plan. Defaults to
12191248
:py:attr:`ExplainFormat.INDENT`.
1249+
show_statistics: If ``True``, include each operator's statistics.
1250+
``None`` uses the ``datafusion.explain.show_statistics`` setting.
1251+
analyze_level: Which metrics to report with ``analyze``. ``None``
1252+
uses the ``datafusion.explain.analyze_level`` setting.
1253+
analyze_categories: Report only metrics in these categories with
1254+
``analyze``; an empty iterable reports none. ``None`` uses the
1255+
``datafusion.explain.analyze_categories`` setting.
12201256
12211257
Examples:
12221258
Show the plan in tree format:
@@ -1229,9 +1265,22 @@ def explain(
12291265
Show plan with runtime metrics:
12301266
12311267
>>> df.explain(analyze=True) # doctest: +SKIP
1268+
1269+
Show only row-count metrics:
1270+
1271+
>>> from datafusion.dataframe import ExplainMetricCategory
1272+
>>> df.explain(
1273+
... analyze=True, analyze_categories=[ExplainMetricCategory.ROWS]
1274+
... ) # doctest: +SKIP
12321275
"""
12331276
fmt = format.value if format is not None else None
1234-
self.df.explain(verbose, analyze, fmt)
1277+
level = analyze_level.value if analyze_level is not None else None
1278+
categories = (
1279+
[c.value for c in analyze_categories]
1280+
if analyze_categories is not None
1281+
else None
1282+
)
1283+
self.df.explain(verbose, analyze, fmt, show_statistics, level, categories)
12351284

12361285
def logical_plan(self) -> LogicalPlan:
12371286
"""Return the unoptimized ``LogicalPlan``.

‎python/tests/test_dataframe.py‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,11 @@
4747
functions as f,
4848
)
4949
from datafusion.common import NullTreatment
50-
from datafusion.dataframe import DataFrameWriteOptions
50+
from datafusion.dataframe import (
51+
DataFrameWriteOptions,
52+
ExplainAnalyzeLevel,
53+
ExplainMetricCategory,
54+
)
5155
from datafusion.dataframe_formatter import (
5256
DataFrameHtmlFormatter,
5357
configure_formatter,
@@ -3871,6 +3875,57 @@ def test_explain_with_format(capsys, fmt, verbose, analyze, expected_substring):
38713875
assert expected_substring in captured.out
38723876

38733877

3878+
def _explain_output(capsys, **kwargs):
3879+
ctx = SessionContext()
3880+
df = ctx.from_pydict({"a": [1, 2]}).filter(column("a") > literal(1))
3881+
df.explain(**kwargs)
3882+
return capsys.readouterr().out
3883+
3884+
3885+
@pytest.mark.parametrize(
3886+
("kwargs", "present", "absent"),
3887+
[
3888+
pytest.param({}, [], ["statistics="], id="default_no_statistics"),
3889+
pytest.param(
3890+
{"show_statistics": True}, ["statistics=[Rows="], [], id="show_statistics"
3891+
),
3892+
pytest.param(
3893+
{"analyze": True, "analyze_level": ExplainAnalyzeLevel.DEV},
3894+
["output_rows=", "output_batches="],
3895+
[],
3896+
id="analyze_level_dev",
3897+
),
3898+
pytest.param(
3899+
{"analyze": True, "analyze_level": ExplainAnalyzeLevel.SUMMARY},
3900+
["output_rows="],
3901+
["output_batches="],
3902+
id="analyze_level_summary",
3903+
),
3904+
pytest.param(
3905+
{
3906+
"analyze": True,
3907+
"analyze_categories": [ExplainMetricCategory.ROWS],
3908+
},
3909+
["output_rows="],
3910+
["elapsed_compute=", "output_bytes="],
3911+
id="analyze_categories_rows",
3912+
),
3913+
pytest.param(
3914+
{"analyze": True, "analyze_categories": []},
3915+
["FilterExec: a@0 > 1, metrics=[]"],
3916+
["output_rows="],
3917+
id="analyze_categories_empty_suppresses_metrics",
3918+
),
3919+
],
3920+
)
3921+
def test_explain_options(capsys, kwargs, present, absent):
3922+
out = _explain_output(capsys, **kwargs)
3923+
for text in present:
3924+
assert text in out
3925+
for text in absent:
3926+
assert text not in out
3927+
3928+
38743929
@pytest.mark.parametrize(
38753930
("window_exprs", "expected_columns"),
38763931
[

0 commit comments

Comments
 (0)