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 warnings
from typing import TYPE_CHECKING, Any, Literal

import datafusion._internal as df_internal

if TYPE_CHECKING:
    import datetime

    from datafusion.context import SessionContext

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


class LogicalPlan:  # noqa: PLW1641
    """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 formattable 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_bytes(ctx: SessionContext, data: bytes) -> LogicalPlan:
        """Create a LogicalPlan from serialized protobuf bytes.

        Decoding routes through the codecs installed on ``ctx`` with
        :py:meth:`~datafusion.SessionContext.with_logical_extension_codec`.
        Tables created in memory from record batches are currently not
        supported.

        Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and
        positional, and there is no fallback to a worker or global context.

        See Also:
            :py:meth:`to_bytes`, :py:meth:`ExecutionPlan.from_bytes`,
            :py:meth:`datafusion.Expr.from_bytes`.
        """
        return LogicalPlan(df_internal.LogicalPlan.from_bytes(ctx.ctx, data))

    def to_bytes(self, ctx: SessionContext | None = None) -> bytes:
        """Convert a LogicalPlan to serialized protobuf bytes.

        When ``ctx`` is supplied, encoding routes through the codecs
        installed on it with
        :py:meth:`~datafusion.SessionContext.with_logical_extension_codec`,
        so extension codecs see the encode path. With ``ctx=None`` a
        default codec is used. Tables created in memory from record
        batches are currently not supported.

        See Also:
            :py:meth:`from_bytes`, :py:meth:`ExecutionPlan.to_bytes`,
            :py:meth:`datafusion.Expr.to_bytes`.
        """
        ctx_arg = ctx.ctx if ctx is not None else None
        return self._raw_plan.to_bytes(ctx_arg)

    @staticmethod
    def from_proto(ctx: SessionContext, data: bytes) -> LogicalPlan:
        """Deprecated alias for :meth:`from_bytes`."""
        warnings.warn(
            "LogicalPlan.from_proto is deprecated; use from_bytes instead",
            DeprecationWarning,
            stacklevel=2,
        )
        return LogicalPlan.from_bytes(ctx, data)

    def to_proto(self) -> bytes:
        """Deprecated alias for :meth:`to_bytes`."""
        warnings.warn(
            "LogicalPlan.to_proto is deprecated; use to_bytes instead",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.to_bytes()

    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.

        A node contributed by an extension library prints as its FFI wrapper
        rather than as itself, which changes how it can be matched; see
        :ref:`extension_foreign_node_display`.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> df = ctx.from_pydict({"a": [1, 2, 3]})
            >>> df.execution_plan().display().strip()
            'DataSourceExec: partitions=1, partition_sizes=[1]'
        """
        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.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> df = ctx.from_pydict({"a": [1, 2, 3]})
            >>> df.execution_plan().partition_count
            1
        """
        return self._raw_plan.partition_count

    @property
    def output_partitioning(self) -> PhysicalPartitioning:
        """Returns how this plan's output rows are spread across its partitions.

        Where :py:attr:`partition_count` gives only the number of partitions,
        this also reports the scheme, so a caller executing partitions
        separately can tell whether they are hash-distributed on known keys or
        merely counted. A plan does not necessarily partition the way it was
        asked to; see :ref:`checking_partitioning`.

        Examples:
            >>> import pyarrow as pa
            >>> from datafusion import SessionConfig, SessionContext
            >>> ctx = SessionContext(SessionConfig().with_target_partitions(4))
            >>> ctx.register_record_batches("t", [
            ...     [pa.record_batch({"a": [1, 2, 3]})],
            ...     [pa.record_batch({"a": [4, 5, 6]})],
            ... ])
            >>> ctx.sql("select a from t").execution_plan().output_partitioning
            UnknownPartitioning(2)

            A group-by redistributes rows, so the plan reports the keys:

            >>> grouped = ctx.sql("select a, count(*) from t group by a")
            >>> partitioning = grouped.execution_plan().output_partitioning
            >>> partitioning.scheme
            'Hash'
            >>> partitioning.partition_count
            4
        """
        return PhysicalPartitioning(self._raw_plan.output_partitioning)

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

        Decoding routes through the codecs installed on ``ctx`` with
        :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`.

        Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and
        positional, and there is no fallback to a worker or global context.
        ``ctx`` need not share the encoding session's registered tables: a
        scan over a table registered from record batches decodes here, because
        the batches travel inside the encoded scan. It does need the same
        extension codecs, which are what decode any node an extension library
        defines.

        See Also:
            :py:meth:`to_bytes`, :py:meth:`LogicalPlan.from_bytes`.
        """
        return ExecutionPlan(df_internal.ExecutionPlan.from_bytes(ctx.ctx, data))

    def to_bytes(self, ctx: SessionContext | None = None) -> bytes:
        """Convert an ExecutionPlan into serialized protobuf bytes.

        When ``ctx`` is supplied, encoding routes through the codecs
        installed on it with
        :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`.
        Without it, no extension codec is consulted, so a node an extension
        library defines fails to encode; inline Python UDFs still travel,
        since they are not carried by an extension codec.

        Unlike :py:meth:`LogicalPlan.to_bytes`, a plan reading a table
        registered from record batches does round-trip: the batches travel
        inside the encoded scan.

        Round-tripping through this method and :py:meth:`from_bytes` is how
        an extension library checks that its own codec claimed its nodes,
        rather than a codec installed earlier in the chain — see
        :ref:`extension_codec_order`.

        See Also:
            :py:meth:`from_bytes`, :py:meth:`LogicalPlan.to_bytes`.
        """
        ctx_arg = ctx.ctx if ctx is not None else None
        return self._raw_plan.to_bytes(ctx_arg)

    @staticmethod
    def from_proto(ctx: SessionContext, data: bytes) -> ExecutionPlan:
        """Deprecated alias for :meth:`from_bytes`."""
        warnings.warn(
            "ExecutionPlan.from_proto is deprecated; use from_bytes instead",
            DeprecationWarning,
            stacklevel=2,
        )
        return ExecutionPlan.from_bytes(ctx, data)

    def to_proto(self) -> bytes:
        """Deprecated alias for :meth:`to_bytes`."""
        warnings.warn(
            "ExecutionPlan.to_proto is deprecated; use to_bytes instead",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.to_bytes()

    def metrics(self) -> MetricsSet | None:
        """Return metrics for this plan node, or None if this plan has no MetricsSet.

        Some operators (e.g. DataSourceExec) eagerly initialize a MetricsSet
        when the plan is created, so this may return a set even before
        execution.  Metric *values* (such as ``output_rows``) are only
        meaningful after the DataFrame has been executed.
        """
        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 PhysicalPartitioning:
    """How a physical plan's output rows are spread across its partitions.

    Returned by :py:attr:`ExecutionPlan.output_partitioning`. This is the
    partitioning a built plan *has*. Distinct from
    :py:class:`datafusion.expr.Partitioning`, the *logical* partitioning a
    ``Repartition`` node records and hands back from
    ``partitioning_scheme()`` — a request, which the plan need not honour. See
    :ref:`checking_partitioning`.
    """

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

    @property
    def scheme(
        self,
    ) -> Literal["RoundRobinBatch", "Hash", "Range", "UnknownPartitioning"]:
        """Which partitioning scheme this is.

        ``"UnknownPartitioning"`` means the plan knows how many partitions it
        has but nothing about how rows are distributed between them, which is
        the usual case for a file scan. ``"RoundRobinBatch"`` and ``"Hash"``
        come from a repartition the optimizer inserted. ``"Range"`` only
        appears on a plan this package did not build; see
        :ref:`checking_partitioning`.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> df = ctx.from_pydict({"a": [1, 2, 3]})
            >>> df.execution_plan().output_partitioning.scheme
            'UnknownPartitioning'
        """
        return self._raw_partitioning.scheme

    @property
    def partition_count(self) -> int:
        """The number of partitions.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> df = ctx.from_pydict({"a": [1, 2, 3]})
            >>> df.execution_plan().output_partitioning.partition_count
            1
        """
        return self._raw_partitioning.partition_count

    @property
    def hash_expressions(self) -> list[str] | None:
        """The expressions rows are hashed on, or ``None`` for other schemes.

        Physical expressions have no Python representation, so these are
        returned in their displayed form. ``"Range"`` reports ``None`` too,
        even though it does partition on expressions: its ordering and split
        points are not exposed here.

        Examples:
            >>> import pyarrow as pa
            >>> from datafusion import SessionConfig, SessionContext
            >>> ctx = SessionContext(SessionConfig().with_target_partitions(4))
            >>> ctx.register_record_batches("t", [
            ...     [pa.record_batch({"a": [1, 2, 3]})],
            ...     [pa.record_batch({"a": [4, 5, 6]})],
            ... ])
            >>> scan = ctx.sql("select a from t").execution_plan()
            >>> scan.output_partitioning.hash_expressions is None
            True
            >>> grouped = ctx.sql("select a, count(*) from t group by a")
            >>> grouped.execution_plan().output_partitioning.hash_expressions
            ['a@0']
        """
        return self._raw_partitioning.hash_expressions

    def __repr__(self) -> str:
        """Print a string representation of the partitioning.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> df = ctx.from_pydict({"a": [1, 2, 3]})
            >>> repr(df.execution_plan().output_partitioning)
            'UnknownPartitioning(1)'
        """
        return self._raw_partitioning.__repr__()

    def _key(self) -> tuple[str, int, tuple[str, ...] | None]:
        exprs = self.hash_expressions
        return (self.scheme, self.partition_count, tuple(exprs) if exprs else None)

    def __eq__(self, other: object) -> bool:
        """Compare two partitionings by scheme, count and hash expressions.

        Equality is structural, and does not mirror DataFusion's own
        comparison of the underlying type, under which two
        ``UnknownPartitioning`` values of the same width are unequal.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> a = ctx.from_pydict({"a": [1, 2, 3]}).execution_plan()
            >>> b = ctx.from_pydict({"b": [4, 5, 6]}).execution_plan()
            >>> a.output_partitioning == b.output_partitioning
            True
            >>> a.output_partitioning == "UnknownPartitioning(1)"
            False
        """
        if not isinstance(other, PhysicalPartitioning):
            return NotImplemented
        return self._key() == other._key()

    def __hash__(self) -> int:
        """Hash the partitioning, consistently with :py:meth:`__eq__`.

        Examples:
            >>> from datafusion import SessionContext
            >>> ctx = SessionContext()
            >>> a = ctx.from_pydict({"a": [1, 2, 3]}).execution_plan()
            >>> b = ctx.from_pydict({"b": [4, 5, 6]}).execution_plan()
            >>> len({a.output_partitioning, b.output_partitioning})
            1
        """
        return hash(self._key())


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 | datetime.datetime | None:
        """The value of this metric.

        Returns an ``int`` for counters, gauges, and time-based metrics
        (nanoseconds), a :py:class:`~datetime.datetime` (UTC) for
        ``start_timestamp`` / ``end_timestamp`` metrics, or ``None``
        when the value has not been set or is not representable.
        """
        return self._raw.value

    @property
    def value_as_datetime(self) -> datetime.datetime | None:
        """The value as a UTC :py:class:`~datetime.datetime` for timestamp metrics.

        Returns ``None`` for all non-timestamp metrics and for timestamp
        metrics whose value has not been set (e.g. before execution).
        """
        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)