feat: initial commit add support for metrics in pydatafusion

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, collect_metrics
 from .record_batch import RecordBatch, RecordBatchStream
 from .user_defined import (
     Accumulator,
@@ -85,6 +85,8 @@
     "Expr",
     "InsertOp",
     "LogicalPlan",
+    "Metric",
+    "MetricsSet",
     "ParquetColumnOptions",
     "ParquetWriterOptions",
     "RecordBatch",
@@ -100,6 +102,7 @@
     "WindowUDF",
     "catalog",
     "col",
+    "collect_metrics",
     "column",
     "common",
     "configure_formatter",

python/datafusion/plan.py

@@ -29,6 +29,9 @@
 __all__ = [
     "ExecutionPlan",
     "LogicalPlan",
+    "Metric",
+    "MetricsSet",
+    "collect_metrics",
 ]
 
 
@@ -151,3 +154,130 @@ 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 this plan and collect results as a list of RecordBatches.
+
+        After calling this method, :meth:`metrics` will return populated metrics
+        for this plan and its children.
+
+        Args:
+            ctx: The session context to use for execution.
+        """
+        return self._raw_plan.execute_collect(ctx.ctx)
+
+    def metrics(self) -> MetricsSet | None:
+        """Return the metrics for this plan node, or None if not available.
+
+        Metrics are only available after the plan has been executed via
+        :meth:`execute_collect`.
+        """
+        raw = self._raw_plan.metrics()
+        if raw is None:
+            return None
+        return MetricsSet(raw)
+
+
+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)
+
+    def __str__(self) -> str:
+        """Return a human-readable string of the metrics set."""
+        return str(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)
+
+
+def collect_metrics(
+    plan: ExecutionPlan,
+) -> list[tuple[str, MetricsSet]]:
+    """Walk an execution plan tree and collect metrics from all operators.
+
+    Returns a list of ``(operator_display_name, MetricsSet)`` tuples for every
+    operator that has metrics available.
+    """
+    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(plan)
+    return result

python/tests/test_plans.py

@@ -16,7 +16,14 @@
 # under the License.
 
 import pytest
-from datafusion import ExecutionPlan, LogicalPlan, SessionContext
+from datafusion import (
+    ExecutionPlan,
+    LogicalPlan,
+    Metric,
+    MetricsSet,
+    SessionContext,
+    collect_metrics,
+)
 
 
 # Note: We must use CSV because memory tables are currently not supported for
@@ -40,3 +47,97 @@ 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)
+
+    # Walk the tree to find a node with metrics
+    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, "Expected at least one operator with output_rows > 0"
+
+
+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 collect_metrics(plan):
+        for metric in ms.metrics():
+            assert isinstance(metric, Metric)
+            assert isinstance(metric.name, str)
+            assert len(metric.name) > 0
+            # partition is either None or an int
+            assert metric.partition is None or isinstance(metric.partition, int)
+            assert isinstance(metric.labels(), dict)
+            return  # only need to check one
+    pytest.skip("No metrics found on any operator")
+
+
+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 = collect_metrics(plan)
+    # An aggregation query should have multiple operators with metrics
+    assert len(results) >= 2, (
+        f"Expected multiple operators with metrics, got {len(results)}: "
+        f"{[name for name, _ in results]}"
+    )
+    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")
+    # Get execution plan WITHOUT executing
+    plan = df.execution_plan()
+    # Before execution, metrics should be None or have zero rows
+    ms = plan.metrics()
+    if ms is not None:
+        rows = ms.output_rows
+        assert rows is None or 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 collect_metrics(plan):
+        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 on any operator")

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>()?;