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

import datetime

import pyarrow as pa
import pyarrow.parquet as pq
import pytest
from datafusion import (
    ExecutionPlan,
    LogicalPlan,
    Metric,
    MetricsSet,
    PhysicalPartitioning,
    SessionConfig,
    SessionContext,
    col,
    udf,
)
from datafusion.expr import Partitioning


# Note: CSV because a *logical* plan cannot carry a memory table. The physical
# layer can — see `test_execution_plan_over_memory_batches_round_trips`.
@pytest.fixture
def df():
    ctx = SessionContext()
    return ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select("c1")


def test_logical_plan_to_bytes_roundtrip(ctx, df) -> None:
    """Round-trip a LogicalPlan through the session's logical codec."""
    logical_plan_bytes = df.logical_plan().to_bytes()
    logical_plan = LogicalPlan.from_bytes(ctx, logical_plan_bytes)

    df_round_trip = ctx.create_dataframe_from_logical_plan(logical_plan)

    assert df.collect() == df_round_trip.collect()


def test_execution_plan_to_bytes_roundtrip(ctx, df) -> None:
    """Round-trip an ExecutionPlan through the session's physical codec."""
    original_execution_plan = df.execution_plan()
    execution_plan_bytes = original_execution_plan.to_bytes()
    execution_plan = ExecutionPlan.from_bytes(ctx, execution_plan_bytes)

    assert str(original_execution_plan) == str(execution_plan)


def test_logical_plan_to_proto_is_deprecated(ctx, df) -> None:
    """to_proto / from_proto still work but emit DeprecationWarning."""
    plan = df.logical_plan()

    with pytest.warns(DeprecationWarning, match="to_proto"):
        blob = plan.to_proto()
    with pytest.warns(DeprecationWarning, match="from_proto"):
        restored = LogicalPlan.from_proto(ctx, blob)

    df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
    assert df.collect() == df_round_trip.collect()


def test_execution_plan_to_proto_is_deprecated(ctx, df) -> None:
    plan = df.execution_plan()

    with pytest.warns(DeprecationWarning, match="to_proto"):
        blob = plan.to_proto()
    with pytest.warns(DeprecationWarning, match="from_proto"):
        restored = ExecutionPlan.from_proto(ctx, blob)

    assert str(plan) == str(restored)


def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None:
    """A session with a non-default logical codec still round-trips builtins.

    The codec slot is overridable via with_logical_extension_codec; the
    PythonLogicalCodec wrapper delegates unhandled cases to the inner
    codec, so plans without Python UDFs are unaffected by the swap.
    """
    # Default-routed session should round-trip via to_bytes.
    blob = df.logical_plan().to_bytes()
    restored = LogicalPlan.from_bytes(ctx, blob)
    df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
    assert df.collect() == df_round_trip.collect()


def test_execution_plan_over_memory_batches_round_trips() -> None:
    """A physical plan reading record batches decodes on an unrelated session.

    Only the *logical* layer cannot carry a memory table: its
    `try_encode_table_provider` has no arm for one. The physical scan inlines
    the batches, so it needs neither a shared session nor an extension codec —
    which is what lets a worker process execute a plan the driver encoded.
    """
    ctx = SessionContext()
    ctx.register_record_batches(
        "t",
        [[pa.record_batch({"a": [1, 2, 3]})], [pa.record_batch({"a": [4, 5, 6]})]],
    )
    plan_bytes = ctx.sql("select a from t").execution_plan().to_bytes(ctx)

    # A session that shares nothing with the encoder: no codecs, no tables.
    fresh = SessionContext()
    decoded = ExecutionPlan.from_bytes(fresh, plan_bytes)
    rows = sum(
        batch.to_pyarrow().num_rows
        for partition in range(decoded.partition_count)
        for batch in fresh.execute(decoded, partition)
    )
    assert rows == 6


def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None:
    """`output_partitioning` distinguishes hash-distributed output from counted."""
    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()
    scanned = scan.output_partitioning
    assert scanned.scheme == "UnknownPartitioning"
    assert scanned.hash_expressions is None
    # Agrees with the count-only accessor it supplements.
    assert scanned.partition_count == scan.partition_count

    grouped = ctx.sql("select a, count(*) from t group by a").execution_plan()
    partitioning = grouped.output_partitioning
    assert partitioning.scheme == "Hash"
    assert partitioning.hash_expressions == ["a@0"]
    assert partitioning.partition_count == 4
    assert repr(partitioning) == "Hash([a@0], 4)"


def test_a_requested_partitioning_and_the_resulting_one_disagree() -> None:
    """The logical request and the physical result are different things.

    `datafusion.expr.Partitioning` is what a `Repartition` node records — the
    request. `PhysicalPartitioning` is what the built plan does. Here the
    optimizer drops the repartition outright, because nothing above it needs
    the rows redistributed, so the two do not even agree on the scheme.
    """
    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]})]],
    )
    df = ctx.table("t").repartition_by_hash(col("a"), num=8)

    # The request survives on the logical plan, as an opaque object of the
    # other Partitioning type.
    requested = df.logical_plan().to_variant().partitioning_scheme()
    assert isinstance(requested, Partitioning)
    assert not isinstance(requested, PhysicalPartitioning)

    # The result honours neither the scheme nor the count that was asked for.
    resulting = df.execution_plan().output_partitioning
    assert isinstance(resulting, PhysicalPartitioning)
    assert resulting.scheme == "UnknownPartitioning"
    assert resulting.partition_count == 2


def test_output_partitioning_reports_round_robin(tmp_path) -> None:
    """A round-robin repartition reports `RoundRobinBatch`.

    The optimizer only inserts one above a source with fewer partitions than
    `target_partitions` and CPU work above it to parallelize, and it never
    survives at the root, so reach it by walking `children`.
    """
    path = tmp_path / "rr.parquet"
    pq.write_table(pa.table({"a": list(range(50)), "b": [1] * 50}), path)

    ctx = SessionContext(SessionConfig().with_target_partitions(8))
    ctx.register_parquet("t", str(path))
    plan = ctx.sql("select a, sum(b) from t where a > 5 group by a").execution_plan()

    schemes = set()
    stack = [plan]
    while stack:
        node = stack.pop()
        schemes.add(node.output_partitioning.scheme)
        stack.extend(node.children())

    # Membership, not equality: which other nodes the optimizer puts in this
    # tree is its business, and pinning the whole set here would make an
    # unrelated planner change look like a failure of this accessor. The other
    # schemes are asserted directly where they are the subject.
    assert "RoundRobinBatch" in schemes


def test_execute_rejects_an_out_of_range_partition() -> None:
    """An out-of-range partition index raises instead of panicking."""
    ctx = SessionContext()
    ctx.register_record_batches("t", [[pa.record_batch({"a": [1, 2, 3]})]])
    plan = ctx.sql("select a from t").execution_plan()
    assert plan.partition_count == 1

    with pytest.raises(ValueError, match="Partition index 5 is out of range"):
        ctx.execute(plan, 5)

    # The keyword is `partition`, as the upgrade guide says.
    with pytest.raises(ValueError, match="Partition index 5 is out of range"):
        ctx.execute(plan, partition=5)


def test_execute_rejects_a_negative_partition() -> None:
    """A negative index cannot reach the bounds check, so it overflows first.

    Documented on `execute` as `OverflowError` because that is what PyO3
    raises converting to `usize`, before any DataFusion code runs.
    """
    ctx = SessionContext()
    ctx.register_record_batches("t", [[pa.record_batch({"a": [1, 2, 3]})]])
    plan = ctx.sql("select a from t").execution_plan()

    with pytest.raises(OverflowError):
        ctx.execute(plan, -1)


def test_physical_partitioning_equality_is_structural() -> None:
    """Two partitionings are equal when scheme, count and keys agree.

    Not DataFusion's own comparison of the underlying type, which reports two
    `UnknownPartitioning` values of the same width as unequal. A reflexive
    `__eq__` is the Python expectation, and the count-only alternative would
    make `Hash` on different keys compare equal.
    """
    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]})]],
    )
    plan = ctx.sql("select a from t").execution_plan()
    other_scan = ctx.sql("select a as b from t").execution_plan().output_partitioning
    grouped = (
        ctx.sql("select a, count(*) from t group by a")
        .execution_plan()
        .output_partitioning
    )

    # The property builds a fresh wrapper per access, so these are two objects
    # over one partitioning. `UnknownPartitioning` is precisely the scheme
    # DataFusion's own comparison reports as unequal to itself.
    scan, scan_again = plan.output_partitioning, plan.output_partitioning
    assert scan is not scan_again
    assert scan == scan_again

    assert scan == other_scan
    assert scan != grouped
    assert scan != "UnknownPartitioning(2)"

    # Hashing agrees, so these collapse in a set the way equality implies.
    assert len({scan, other_scan, grouped}) == 2


def test_installing_a_physical_codec_preserves_strict_mode() -> None:
    """Installing a physical extension codec must not re-enable inlining.

    `with_physical_extension_codec` builds a replacement `PythonPhysicalCodec`
    around the imported one, and the constructor defaults inlining to on. A
    context that opted out via `with_python_udf_inlining(enabled=False)` has to
    keep its setting, or installing a codec silently starts embedding
    cloudpickled callables in serialized execution plans.

    The logical counterpart lives in `test_pickle_expr.py`; this one needs an
    `ExecutionPlan` because only the physical codec encodes it. `DFPYUDF` is
    the scalar Python-UDF family prefix, shared by both layers; see
    `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs.
    """
    identity = udf(
        lambda arr: arr,
        [pa.string()],
        pa.string(),
        volatility="immutable",
        name="identity_str",
    )

    def plan_bytes(ctx: SessionContext) -> bytes:
        df = ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select(
            identity(col("c1"))
        )
        return df.execution_plan().to_bytes(ctx)

    # The inlining default is on, so the strict blob is what has to differ.
    assert b"DFPYUDF" in plan_bytes(SessionContext())

    strict = SessionContext().with_python_udf_inlining(enabled=False)
    assert b"DFPYUDF" not in plan_bytes(strict)

    installed = strict.with_physical_extension_codec(
        strict.__datafusion_physical_extension_codec__()
    )
    assert b"DFPYUDF" not in plan_bytes(installed)


def test_session_codec_capsule_getters(ctx) -> None:
    """SessionContext exposes both logical and physical codec capsules."""
    logical = ctx.ctx.__datafusion_logical_extension_codec__()
    physical = ctx.ctx.__datafusion_physical_extension_codec__()
    assert logical is not None
    assert physical is not None


def test_metrics_tree_walk() -> 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")
    df.collect()
    plan = df.execution_plan()

    results = plan.collect_metrics()
    assert len(results) >= 1
    output_rows_by_op: dict[str, int] = {}
    for name, ms in results:
        assert isinstance(name, str)
        assert isinstance(ms, MetricsSet)
        if ms.output_rows is not None:
            output_rows_by_op[name] = ms.output_rows

    # The filter passes rows where column1 > 1, so exactly
    # 2 rows from (1,'a'),(2,'b'),(3,'c').
    # At least one operator must report exactly 2 output rows (the filter).
    assert 2 in output_rows_by_op.values(), (
        f"Expected an operator with output_rows=2, got {output_rows_by_op}"
    )


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")
    df.collect()
    plan = df.execution_plan()

    found_any_metric = False
    for _, ms in plan.collect_metrics():
        r = repr(ms)
        assert isinstance(r, str)
        for metric in ms.metrics():
            found_any_metric = True
            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 metric.value is None or isinstance(
                metric.value, int | datetime.datetime
            )
            assert isinstance(metric.labels(), dict)
            mr = repr(metric)
            assert isinstance(mr, str)
            assert len(mr) > 0
    assert found_any_metric, "Expected at least one metric after execution"


def test_no_meaningful_metrics_before_execution() -> 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_before = df.execution_plan()

    # Some plan nodes (e.g. DataSourceExec) eagerly initialize a MetricsSet,
    # so metrics() may return a set even before execution.  However, no rows
    # should have been processed yet — output_rows must be absent or zero.
    for _, ms in plan_before.collect_metrics():
        rows = ms.output_rows
        assert rows is None or rows == 0, (
            f"Expected 0 output_rows before execution, got {rows}"
        )

    # After execution, at least one operator must report rows processed.
    df.collect()
    plan_after = df.execution_plan()
    output_rows_after = [
        ms.output_rows
        for _, ms in plan_after.collect_metrics()
        if ms.output_rows is not None and ms.output_rows > 0
    ]
    assert len(output_rows_after) > 0, "Expected output_rows > 0 after execution"


def test_collect_partitioned_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")

    df.collect_partitioned()
    plan = df.execution_plan()

    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}"


def test_execute_stream_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")

    for _ in df.execute_stream():
        pass

    plan = df.execution_plan()
    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}"


def test_execute_stream_partitioned_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")

    for stream in df.execute_stream_partitioned():
        for _ in stream:
            pass

    plan = df.execution_plan()
    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}"


def test_value_as_datetime() -> 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")
    df.collect()
    plan = df.execution_plan()

    for _, ms in plan.collect_metrics():
        for metric in ms.metrics():
            if metric.name in ("start_timestamp", "end_timestamp"):
                dt = metric.value_as_datetime
                assert dt is None or isinstance(dt, datetime.datetime)
                if dt is not None:
                    assert dt.tzinfo is not None
            else:
                assert metric.value_as_datetime is None


def test_metric_names_and_labels() -> None:
    """Verify that known metric names appear and labels are well-formed."""
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
    df.collect()
    plan = df.execution_plan()

    all_metric_names: set[str] = set()
    for _, ms in plan.collect_metrics():
        for metric in ms.metrics():
            all_metric_names.add(metric.name)
            # Labels must be a dict of str->str
            labels = metric.labels()
            for k, v in labels.items():
                assert isinstance(k, str)
                assert isinstance(v, str)

    # After a filter query, we expect at minimum these standard metric names.
    assert "output_rows" in all_metric_names, (
        f"Expected 'output_rows' in {all_metric_names}"
    )
    assert "elapsed_compute" in all_metric_names, (
        f"Expected 'elapsed_compute' in {all_metric_names}"
    )


def test_collect_twice_has_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")

    df.collect()
    df.collect()

    plan = df.execution_plan()
    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert len(output_rows_values) > 0