plan.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""This module supports physical and logical plans in DataFusion."""

from __future__ import annotations

import datetime

from typing import TYPE_CHECKING, Any

import datafusion._internal as df_internal

if TYPE_CHECKING:
    from datafusion.context import SessionContext

__all__ = [
    "ExecutionPlan",
    "LogicalPlan",
    "Metric",
    "MetricsSet",
]


class LogicalPlan:
    """Logical Plan.

    A `LogicalPlan` is a node in a tree of relational operators (such as
    Projection or Filter).

    Represents transforming an input relation (table) to an output relation
    (table) with a potentially different schema. Plans form a dataflow tree
    where data flows from leaves up to the root to produce the query result.

    A `LogicalPlan` can be created by the SQL query planner, the DataFrame API,
    or programmatically (for example custom query languages).
    """

    def __init__(self, plan: df_internal.LogicalPlan) -> None:
        """This constructor should not be called by the end user."""
        self._raw_plan = plan

    def to_variant(self) -> Any:
        """Convert the logical plan into its specific variant."""
        return self._raw_plan.to_variant()

    def inputs(self) -> list[LogicalPlan]:
        """Returns the list of inputs to the logical plan."""
        return [LogicalPlan(p) for p in self._raw_plan.inputs()]

    def __repr__(self) -> str:
        """Generate a printable representation of the plan."""
        return self._raw_plan.__repr__()

    def display(self) -> str:
        """Print the logical plan."""
        return self._raw_plan.display()

    def display_indent(self) -> str:
        """Print an indented form of the logical plan."""
        return self._raw_plan.display_indent()

    def display_indent_schema(self) -> str:
        """Print an indented form of the schema for the logical plan."""
        return self._raw_plan.display_indent_schema()

    def display_graphviz(self) -> str:
        """Print the graph visualization of the logical plan.

        Returns a `format`able structure that produces lines meant for graphical display
        using the `DOT` language. This format can be visualized using software from
        [`graphviz`](https://graphviz.org/)
        """
        return self._raw_plan.display_graphviz()

    @staticmethod
    def from_proto(ctx: SessionContext, data: bytes) -> LogicalPlan:
        """Create a LogicalPlan from protobuf bytes.

        Tables created in memory from record batches are currently not supported.
        """
        return LogicalPlan(df_internal.LogicalPlan.from_proto(ctx.ctx, data))

    def to_proto(self) -> bytes:
        """Convert a LogicalPlan to protobuf bytes.

        Tables created in memory from record batches are currently not supported.
        """
        return self._raw_plan.to_proto()

    def __eq__(self, other: LogicalPlan) -> bool:
        """Test equality."""
        if not isinstance(other, LogicalPlan):
            return False
        return self._raw_plan.__eq__(other._raw_plan)


class ExecutionPlan:
    """Represent nodes in the DataFusion Physical Plan."""

    def __init__(self, plan: df_internal.ExecutionPlan) -> None:
        """This constructor should not be called by the end user."""
        self._raw_plan = plan

    def children(self) -> list[ExecutionPlan]:
        """Get a list of children `ExecutionPlan` that act as inputs to this plan.

        The returned list will be empty for leaf nodes such as scans, will contain a
        single value for unary nodes, or two values for binary nodes (such as joins).
        """
        return [ExecutionPlan(e) for e in self._raw_plan.children()]

    def display(self) -> str:
        """Print the physical plan."""
        return self._raw_plan.display()

    def display_indent(self) -> str:
        """Print an indented form of the physical plan."""
        return self._raw_plan.display_indent()

    def __repr__(self) -> str:
        """Print a string representation of the physical plan."""
        return self._raw_plan.__repr__()

    @property
    def partition_count(self) -> int:
        """Returns the number of partitions in the physical plan."""
        return self._raw_plan.partition_count

    @staticmethod
    def from_proto(ctx: SessionContext, data: bytes) -> ExecutionPlan:
        """Create an ExecutionPlan from protobuf bytes.

        Tables created in memory from record batches are currently not supported.
        """
        return ExecutionPlan(df_internal.ExecutionPlan.from_proto(ctx.ctx, data))

    def to_proto(self) -> bytes:
        """Convert an ExecutionPlan into protobuf bytes.

        Tables created in memory from record batches are currently not supported.
        """
        return self._raw_plan.to_proto()

    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]]:
        """Return runtime statistics for each step of the query execution.

        DataFusion executes a query as a pipeline of operators — for example a
        data source scan, followed by a filter, followed by a projection. After
        the DataFrame has been executed (via
        :py:meth:`~datafusion.DataFrame.collect`,
        :py:meth:`~datafusion.DataFrame.execute_stream`, etc.), each operator
        records statistics such as how many rows it produced and how much CPU
        time it consumed.

        Each entry in the returned list corresponds to one operator that
        recorded metrics. The first element of the tuple is the operator's
        description string — the same text shown by
        :py:meth:`display_indent` — which identifies both the operator type
        and its key parameters, for example ``"FilterExec: column1@0 > 1"``
        or ``"DataSourceExec: partitions=1"``.

        Returns:
            A list of ``(description, MetricsSet)`` tuples ordered from the
            outermost operator (top of the execution tree) down to the
            data-source leaves. Only operators that recorded at least one
            metric are included. Returns an empty list if called before the
            DataFrame has been executed.
        """
        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.

    A physical plan operator runs independently across one or more partitions.
    :py:meth:`metrics` returns the raw per-partition :py:class:`Metric` objects.
    The convenience properties (:py:attr:`output_rows`, :py:attr:`elapsed_compute`,
    etc.) automatically sum the named metric across *all* partitions, giving a
    single aggregate value for the operator as a whole.
    """

    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:
        """Total CPU time (in nanoseconds) spent inside this operator's execute loop.

        Summed across all partitions. Returns ``None`` if no ``elapsed_compute``
        metric was recorded.
        """
        return self._raw.elapsed_compute()

    @property
    def spill_count(self) -> int | None:
        """Number of times this operator spilled data to disk due to memory pressure.

        This is a count of spill events, not a byte count. Summed across all
        partitions. Returns ``None`` if no ``spill_count`` metric was recorded.
        """
        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:
        """Sum the named metric across all partitions.

        Useful for accessing any metric not exposed as a first-class property.
        Returns ``None`` if no metric with the given name was recorded.

        Args:
            name: The metric name, e.g. ``"output_rows"`` or ``"elapsed_compute"``.
        """
        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`` when not representable.

        ``None`` is returned for metric types whose value has not yet been set
        (e.g. ``StartTimestamp`` / ``EndTimestamp`` before the operator runs)
        and for any metric variant whose value cannot be expressed as an integer.
        Timestamp metrics, when available, are returned as nanoseconds since the
        Unix epoch.
        """
        return self._raw.value

    @property
    def value_as_datetime(self) -> datetime.datetime | None:
        """The value as a UTC datetime for timestamp metrics, or ``None``."""
        return self._raw.value_as_datetime()

    @property
    def partition(self) -> int | None:
        """The 0-based partition index this metric applies to.

        Returns ``None`` for metrics that are not partition-specific (i.e. they
        apply globally across all partitions of the operator).
        """
        return self._raw.partition

    def labels(self) -> dict[str, str]:
        """Return the labels associated with this metric.

        Labels provide additional context for a metric.  For example::

            >>> metric.labels()
            {'output_type': 'final'}
        """
        return self._raw.labels()

    def __repr__(self) -> str:
        """Return a string representation of the metric."""
        return repr(self._raw)