feat: add metrics collection for execution plans

Adds ExecutionPlan.metrics() and ExecutionPlan.collect_metrics() for
accessing query execution metrics after running queries.

- New Metric and MetricsSet classes expose individual metrics and aggregations
- ExecutionPlan.execute_collect() populates metrics during execution
- ExecutionPlan.collect_metrics() walks plan tree to gather all metrics
- Proper None handling for timestamp conversions in Rust
- Concise docstrings and test assertions

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: ShreyeshArangath

python/datafusion/__init__.py

@@ -55,7 +55,7 @@
 from .expr import Expr, WindowFrame
 from .io import read_avro, read_csv, read_json, read_parquet
 from .options import CsvReadOptions
-from .plan import ExecutionPlan, LogicalPlan
+from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet
 from .record_batch import RecordBatch, RecordBatchStream
 from .user_defined import (
     Accumulator,
@@ -85,6 +85,8 @@
     "Expr",
     "InsertOp",
     "LogicalPlan",
+    "Metric",
+    "MetricsSet",
     "ParquetColumnOptions",
     "ParquetWriterOptions",
     "RecordBatch",

python/datafusion/plan.py

@@ -29,6 +29,8 @@
 __all__ = [
     "ExecutionPlan",
     "LogicalPlan",
+    "Metric",
+    "MetricsSet",
 ]
 
 
@@ -151,3 +153,111 @@ def to_proto(self) -> bytes:
         Tables created in memory from record batches are currently not supported.
         """
         return self._raw_plan.to_proto()
+
+    def execute_collect(self, ctx: SessionContext) -> list[Any]:
+        """Execute and collect results. Populates metrics on this plan."""
+        return self._raw_plan.execute_collect(ctx.ctx)
+
+    def metrics(self) -> MetricsSet | None:
+        """Return metrics for this plan node after execution, or None if unavailable."""
+        raw = self._raw_plan.metrics()
+        if raw is None:
+            return None
+        return MetricsSet(raw)
+
+    def collect_metrics(self) -> list[tuple[str, MetricsSet]]:
+        """Walk the plan tree and collect metrics from all operators.
+
+        Returns a list of (operator_name, MetricsSet) tuples.
+        """
+        result: list[tuple[str, MetricsSet]] = []
+
+        def _walk(node: ExecutionPlan) -> None:
+            ms = node.metrics()
+            if ms is not None:
+                result.append((node.display(), ms))
+            for child in node.children():
+                _walk(child)
+
+        _walk(self)
+        return result
+
+
+class MetricsSet:
+    """A set of metrics for a single execution plan operator.
+
+    Provides both individual metric access and convenience aggregations
+    across partitions.
+    """
+
+    def __init__(self, raw: df_internal.MetricsSet) -> None:
+        """This constructor should not be called by the end user."""
+        self._raw = raw
+
+    def metrics(self) -> list[Metric]:
+        """Return all individual metrics in this set."""
+        return [Metric(m) for m in self._raw.metrics()]
+
+    @property
+    def output_rows(self) -> int | None:
+        """Sum of output_rows across all partitions."""
+        return self._raw.output_rows()
+
+    @property
+    def elapsed_compute(self) -> int | None:
+        """Sum of elapsed_compute across all partitions, in nanoseconds."""
+        return self._raw.elapsed_compute()
+
+    @property
+    def spill_count(self) -> int | None:
+        """Sum of spill_count across all partitions."""
+        return self._raw.spill_count()
+
+    @property
+    def spilled_bytes(self) -> int | None:
+        """Sum of spilled_bytes across all partitions."""
+        return self._raw.spilled_bytes()
+
+    @property
+    def spilled_rows(self) -> int | None:
+        """Sum of spilled_rows across all partitions."""
+        return self._raw.spilled_rows()
+
+    def sum_by_name(self, name: str) -> int | None:
+        """Return the sum of metrics matching the given name."""
+        return self._raw.sum_by_name(name)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the metrics set."""
+        return repr(self._raw)
+
+
+class Metric:
+    """A single execution metric with name, value, partition, and labels."""
+
+    def __init__(self, raw: df_internal.Metric) -> None:
+        """This constructor should not be called by the end user."""
+        self._raw = raw
+
+    @property
+    def name(self) -> str:
+        """The name of this metric (e.g. ``output_rows``)."""
+        return self._raw.name
+
+    @property
+    def value(self) -> int | None:
+        """The numeric value of this metric, or None for non-numeric types."""
+        return self._raw.value
+
+    @property
+    def partition(self) -> int | None:
+        """The partition this metric applies to, or None if global."""
+        return self._raw.partition
+
+    def labels(self) -> dict[str, str]:
+        """Return the labels associated with this metric."""
+        return self._raw.labels()
+
+    def __repr__(self) -> str:
+        """Return a string representation of the metric."""
+        return repr(self._raw)

python/tests/test_plans.py

@@ -16,7 +16,13 @@
 # under the License.
 
 import pytest
-from datafusion import ExecutionPlan, LogicalPlan, SessionContext
+from datafusion import (
+    ExecutionPlan,
+    LogicalPlan,
+    Metric,
+    MetricsSet,
+    SessionContext,
+)
 
 
 # Note: We must use CSV because memory tables are currently not supported for
@@ -40,3 +46,87 @@ def test_logical_plan_to_proto(ctx, df) -> None:
     execution_plan = ExecutionPlan.from_proto(ctx, execution_plan_bytes)
 
     assert str(original_execution_plan) == str(execution_plan)
+
+
+def test_execution_plan_metrics() -> None:
+    ctx = SessionContext()
+    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
+
+    plan = df.execution_plan()
+    plan.execute_collect(ctx)
+
+    found_metrics = False
+
+    def _check(node):
+        nonlocal found_metrics
+        ms = node.metrics()
+        if ms is not None and ms.output_rows is not None and ms.output_rows > 0:
+            found_metrics = True
+        for child in node.children():
+            _check(child)
+
+    _check(plan)
+    assert found_metrics
+
+
+def test_metric_properties() -> None:
+    ctx = SessionContext()
+    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
+
+    plan = df.execution_plan()
+    plan.execute_collect(ctx)
+
+    for _, ms in plan.collect_metrics():
+        for metric in ms.metrics():
+            assert isinstance(metric, Metric)
+            assert isinstance(metric.name, str)
+            assert len(metric.name) > 0
+            assert metric.partition is None or isinstance(metric.partition, int)
+            assert isinstance(metric.labels(), dict)
+            return
+    pytest.skip("No metrics found")
+
+
+def test_metrics_tree_walk() -> None:
+    ctx = SessionContext()
+    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'a'), (4, 'b')")
+    df = ctx.sql("SELECT column2, COUNT(*) FROM t GROUP BY column2")
+
+    plan = df.execution_plan()
+    plan.execute_collect(ctx)
+
+    results = plan.collect_metrics()
+    assert len(results) >= 2
+    for name, ms in results:
+        assert isinstance(name, str)
+        assert isinstance(ms, MetricsSet)
+
+
+def test_no_metrics_before_execution() -> None:
+    ctx = SessionContext()
+    ctx.sql("CREATE TABLE t AS VALUES (1), (2), (3)")
+    df = ctx.sql("SELECT * FROM t")
+    plan = df.execution_plan()
+    ms = plan.metrics()
+    assert ms is None or ms.output_rows is None or ms.output_rows == 0
+
+
+def test_metrics_repr() -> None:
+    ctx = SessionContext()
+    ctx.sql("CREATE TABLE t AS VALUES (1), (2), (3)")
+    df = ctx.sql("SELECT * FROM t")
+
+    plan = df.execution_plan()
+    plan.execute_collect(ctx)
+
+    for _, ms in plan.collect_metrics():
+        r = repr(ms)
+        assert isinstance(r, str)
+        for metric in ms.metrics():
+            mr = repr(metric)
+            assert isinstance(mr, str)
+            assert len(mr) > 0
+        return
+    pytest.skip("No metrics found")

src/lib.rs

@@ -43,6 +43,7 @@ pub mod errors;
 pub mod expr;
 #[allow(clippy::borrow_deref_ref)]
 mod functions;
+pub mod metrics;
 mod options;
 pub mod physical_plan;
 mod pyarrow_filter_expression;
@@ -96,6 +97,8 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<udtf::PyTableFunction>()?;
     m.add_class::<config::PyConfig>()?;
     m.add_class::<sql::logical::PyLogicalPlan>()?;
+    m.add_class::<metrics::PyMetricsSet>()?;
+    m.add_class::<metrics::PyMetric>()?;
     m.add_class::<physical_plan::PyExecutionPlan>()?;
     m.add_class::<record_batch::PyRecordBatch>()?;
     m.add_class::<record_batch::PyRecordBatchStream>()?;