# 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: object) -> 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.

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 share nothing with the session that encoded the plan: a

scan over a table registered from record batches decodes here, because

the batches travel inside the encoded scan.

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`.

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. ``None`` covers ``"Range"`` as well,

whose ordering and split points this class does not expose.

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)