examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.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.

from __future__ import annotations

import doctest
import gc
import inspect
import io
import sys
import types

import pyarrow as pa
import pytest
from datafusion import (
    Expr,
    SessionConfig,
    SessionContext,
    SessionExtensionComponents,
    col,
    udf,
)
from datafusion_ffi_example import (
    IsNullUDF,
    MyCatalogProvider,
    MyLogicalExtensionCodec,
    MyPhysicalExtensionCodec,
    MyPhysicalOptimizerRule,
    MyTableProvider,
)
from datafusion_ffi_query_planner_example import (
    MyPlannerConfig,
    MyPlannerExtension,
    MyQueryPlanner,
)


def configured_context(max_rows: int):
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows))
    logical_codec = MyLogicalExtensionCodec()
    physical_codec = MyPhysicalExtensionCodec()
    ctx = SessionContext(config)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(physical_codec)
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    ctx.register_udf(udf(IsNullUDF()))
    return ctx, logical_codec, physical_codec


HOST_ONLY_UDF = "my_custom_is_null"
"""Scalar function registered on the host session and nowhere else."""

UNREGISTERED_UDF = "not_registered_anywhere"


def probe_context(
    *,
    logical_requires: str | None = None,
    physical_requires: str | None = None,
    max_rows: int = 3,
):
    """Three-library context whose codecs read the task context they are given.

    ``require_udf_on_decode`` makes each codec resolve a scalar function from
    the ``TaskContext`` handed to its FFI decode callback, which is otherwise
    unobservable: the example codecs restore objects from a token registry and
    never look at the registry they are passed.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows))
    logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=logical_requires)
    physical_codec = MyPhysicalExtensionCodec(require_udf_on_decode=physical_requires)
    ctx = SessionContext(config)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(physical_codec)
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    ctx.register_udf(udf(IsNullUDF()))
    ctx.set_query_planner(MyQueryPlanner())
    return ctx, logical_codec, physical_codec


def test_logical_codec_resolves_a_host_registered_udf():
    """``try_decode_table_provider`` sees the host session's registry.

    The codec takes its task context provider from the session it is installed
    on, so a function the host registered is resolvable inside a decode
    callback running in the other library.
    """
    ctx, logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF)

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert logical_codec.table_provider_decode_calls() > 0
    assert logical_codec.task_context_udf_resolutions() > 0


def test_physical_codec_resolves_a_host_registered_udf():
    """``try_decode`` sees the host session's registry, as above."""
    ctx, _logical_codec, physical_codec = probe_context(physical_requires=HOST_ONLY_UDF)

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert physical_codec.execution_plan_decode_calls() > 0
    assert physical_codec.task_context_udf_resolutions() > 0


def test_codec_still_reports_a_name_registered_nowhere():
    """Negative control: resolution really is a lookup, not an unconditional pass."""
    ctx, _logical_codec, _physical_codec = probe_context(
        logical_requires=UNREGISTERED_UDF
    )

    with pytest.raises(
        Exception, match=rf"could not resolve scalar function '{UNREGISTERED_UDF}'"
    ):
        ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()


def test_codec_sees_a_udf_registered_after_it_was_installed():
    """The provider is a live handle to the session, not a snapshot of it."""
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF)
    ctx = SessionContext(config)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    # Registered after the codec was installed and bound to this session.
    ctx.register_udf(udf(IsNullUDF()))
    ctx.set_query_planner(MyQueryPlanner())

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert logical_codec.task_context_udf_resolutions() > 0


def test_codec_sees_a_udf_registered_after_the_planner():
    """Installing a planner does not detach the codec from the session.

    The codec is bound to the session before the planner is installed and the
    function is registered afterwards. Installing writes through
    ``state_ref()`` rather than deriving a new ``SessionContext``, so the
    codec's task context provider still points at the one live session and
    sees the later registration.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF)
    ctx = SessionContext(config)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    ctx.set_query_planner(MyQueryPlanner())
    # Registered after the planner, on the same session the codec is bound to.
    ctx.register_udf(udf(IsNullUDF()))

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert logical_codec.table_provider_decode_calls() > 0
    assert logical_codec.task_context_udf_resolutions() > 0


@pytest.mark.parametrize("codecs_first", [True, False])
def test_registered_providers_survive_a_planner_install(codecs_first: bool):
    """Neither write path may orphan a previously registered provider.

    A foreign catalog provider is handed a codec carrying a *weak*
    ``FFI_TaskContextProvider`` pointing at the session it was registered on,
    and upgrades it on every ``supports_filters_pushdown`` and every ``scan``.
    That codec lives inside the registered ``FFI_CatalogProvider``, so nothing
    on the Python side can reach it to rebind it. Deriving a replacement
    ``SessionContext`` would drop the allocation those handles point at, and
    the next query would fail with ``TaskContextProvider went out of scope over
    FFI boundary``. Installing in place keeps the one allocation alive.

    Both parameters exercise that, because both write ``SessionState``:
    ``set_query_planner`` installs the planner, and installing a codec on a
    session that already has one rebuilds that planner against the new codec.

    The ``WHERE`` clause is load-bearing -- it forces filter pushdown, which
    upgrades the weak handle during logical optimization. It is also why both
    codecs have to be installed: without them the query fails at plan
    serialization with ``LogicalExtensionCodec is not provided``, which would
    mask a dangling handle rather than expose it.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=10))
    ctx = SessionContext(config)
    ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())

    def install_codecs(ctx):
        ctx = ctx.with_logical_extension_codec(MyLogicalExtensionCodec())
        return ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())

    if codecs_first:
        ctx = install_codecs(ctx)
        ctx.set_query_planner(MyQueryPlanner())
    else:
        ctx.set_query_planner(MyQueryPlanner())
        ctx = install_codecs(ctx)
    gc.collect()

    batches = ctx.sql(
        "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5"
    ).collect()
    assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [
        7,
        10,
        20,
        30,
    ]


def codec_context(max_rows: int = 3):
    """Context with both example codecs installed and a table to scan.

    Unlike :func:`probe_context` the codecs ask for no function, so the only
    thing they record is the session id of the task context they are handed.
    No planner yet -- the caller installs one.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows))
    logical_codec = MyLogicalExtensionCodec()
    physical_codec = MyPhysicalExtensionCodec()
    ctx = SessionContext(config)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(physical_codec)
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    return ctx, logical_codec, physical_codec


def physical_only_context(max_rows: int = 3):
    """Context with the physical codec installed but no logical codec.

    Encoding consults chained codecs in install order and the first to claim an
    object wins, so a codec installed later cannot be observed while an earlier
    one is already claiming table providers. Leaving the logical slot empty lets
    a test install exactly one logical codec -- through a handle it then throws
    away -- and watch its counters.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows))
    physical_codec = MyPhysicalExtensionCodec()
    ctx = SessionContext(config)
    ctx = ctx.with_physical_extension_codec(physical_codec)
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    return ctx, physical_codec


def test_installing_a_planner_keeps_the_session_id():
    """A session and its decode callbacks must agree on the session id.

    Installing a planner rebuilds ``SessionState`` through
    ``SessionStateBuilder``, which mints a fresh id unless handed one, while
    ``SessionContext`` caches its id in a field of its own. Dropping the id
    there leaves the session reporting one id from ``session_id()`` and a
    different one from every ``TaskContext`` it gives a foreign codec.

    Asserting on ``session_id()`` alone cannot catch that: it reads the cached
    copy, which stays correct either way. The codec-side id is the only
    observable that moves, which is what makes this worth a test rather than a
    one-line equality check.
    """
    ctx, logical_codec, physical_codec = codec_context()
    session_id = ctx.session_id()

    ctx.set_query_planner(MyQueryPlanner())
    assert ctx.session_id() == session_id

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert logical_codec.last_task_context_session_id() == session_id
    assert physical_codec.last_task_context_session_id() == session_id


def test_adding_a_physical_optimizer_rule_keeps_the_session_id():
    """The same guarantee for the other in-place ``SessionState`` rewrite.

    ``add_physical_optimizer_rule`` also rebuilds ``SessionState`` and writes
    it back, so a regenerated id would desync a context from itself.
    """
    ctx, logical_codec, physical_codec = codec_context()
    ctx.set_query_planner(MyQueryPlanner())
    session_id = ctx.session_id()

    ctx.add_physical_optimizer_rule(MyPhysicalOptimizerRule())
    assert ctx.session_id() == session_id

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert logical_codec.last_task_context_session_id() == session_id
    assert physical_codec.last_task_context_session_id() == session_id


def test_replacing_a_planner_keeps_the_session_id():
    """Installing repeatedly must not drift the id.

    Each install rebuilds ``SessionState``, so an id carried over only on the
    first write would still be lost by the second.
    """
    ctx, logical_codec, _physical_codec = codec_context()
    session_id = ctx.session_id()

    ctx.set_query_planner(MyQueryPlanner())
    ctx.set_query_planner(MyQueryPlanner())
    assert ctx.session_id() == session_id

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert logical_codec.last_task_context_session_id() == session_id


@pytest.mark.parametrize("raw_capsule", [False, True])
def test_three_library_query_planner(raw_capsule: bool):
    """Host, provider, and planner exchange a real non-empty plan over FFI."""
    ctx, logical_codec, physical_codec = configured_context(max_rows=3)
    planner = MyQueryPlanner()
    exported_planner = (
        planner.__datafusion_query_planner__(ctx) if raw_capsule else planner
    )
    ctx.set_query_planner(exported_planner)

    batches = ctx.sql(
        'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"'
    ).collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert batches[0].column(1).to_pylist() == [False, False, False]
    assert planner.last_max_rows() == 3

    ctx.sql("SET ffi_query_planner.max_rows = 2").collect()
    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner.last_max_rows() == 2

    assert planner.plan_calls() >= 2
    assert planner.foreign_session_observed()
    assert planner.foreign_provider_observed()
    assert planner.foreign_plan_observed()
    assert logical_codec.table_provider_encode_calls() > 0
    assert logical_codec.table_provider_decode_calls() > 0
    assert physical_codec.execution_plan_encode_calls() > 0
    assert physical_codec.execution_plan_decode_calls() > 0


def test_spawning_plan_across_three_libraries():
    """A plan that spawns Tokio tasks survives the full three-library round trip.

    ``target_partitions`` above one puts a ``RepartitionExec`` under the
    aggregate, and that operator spawns tasks while it runs. This exercises the
    codecs on a multi-node plan rather than the bare scan the other tests use.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=100))
    config = config.with_target_partitions(4)
    logical_codec = MyLogicalExtensionCodec()
    physical_codec = MyPhysicalExtensionCodec()
    ctx = SessionContext(config)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(physical_codec)
    ctx.register_table("numbers", MyTableProvider(1, 6, 3))

    planner = MyQueryPlanner()
    ctx.set_query_planner(planner)

    batches = ctx.sql(
        'SELECT "A" % 2 AS parity, count(*) AS n FROM numbers GROUP BY 1 ORDER BY 1'
    ).collect()
    counts = {
        row[0]: row[1]
        for batch in batches
        for row in zip(
            batch.column(0).to_pylist(), batch.column(1).to_pylist(), strict=True
        )
    }
    assert sum(counts.values()) == 6 + 7 + 8
    assert planner.plan_calls() > 0
    assert planner.foreign_provider_observed()


def test_planner_layers_on_the_session_planner():
    """A planner can wrap the one already installed and delegate to it.

    The capsule has to be captured before this planner is installed, because
    ``__datafusion_query_planner__`` exports whichever planner is installed when
    it is called. Capturing it afterwards would hand the planner a handle to
    itself, and planning would recurse.
    """
    ctx, logical_codec, physical_codec = configured_context(max_rows=3)
    fallback = ctx.__datafusion_query_planner__()
    planner = MyQueryPlanner(fallback=fallback)
    # The capsule's FFI codecs hold weak handles to this session. Installing in
    # place keeps that session alive, so the capsule stays usable; deriving a
    # replacement here would fail with "TaskContextProvider went out of scope
    # over FFI boundary".
    ctx.set_query_planner(planner)
    gc.collect()

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert planner.plan_calls() > 0
    assert planner.used_fallback()
    assert logical_codec.table_provider_decode_calls() > 0
    assert physical_codec.execution_plan_decode_calls() > 0


def test_a_planner_can_fall_back_to_another_planner_library():
    """A fallback may be another foreign planner, not only a session.

    The fallback is imported when this planner is installed rather than when
    it is constructed, so its own getter receives the session. Importing it at
    construction time would mean calling that getter with no session, which
    only a ``SessionContext`` or a raw capsule tolerates -- and layering on
    another planner is the case a distributed engine actually needs.
    """
    ctx, logical_codec, physical_codec = configured_context(max_rows=3)
    inner = MyQueryPlanner()
    outer = MyQueryPlanner(fallback=inner)
    ctx.set_query_planner(outer)

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert outer.plan_calls() > 0
    assert outer.used_fallback()
    # The delegation reached the inner planner rather than stopping at the
    # default physical planner.
    assert inner.plan_calls() > 0
    assert logical_codec.table_provider_decode_calls() > 0
    assert physical_codec.execution_plan_decode_calls() > 0


def test_a_session_fallback_delegates_to_its_installed_planner():
    """Passing a SessionContext delegates to whatever planner it holds."""
    ctx, _logical_codec, _physical_codec = configured_context(max_rows=3)
    first = MyQueryPlanner()
    ctx.set_query_planner(first)

    second = MyQueryPlanner(fallback=ctx)
    ctx.set_query_planner(second)

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert second.used_fallback()
    assert first.plan_calls() > 0


def test_observations_accumulate_across_queries():
    """A later plain query must not retract what an earlier query observed.

    The ``*_observed`` accessors answer "was this ever seen". They are written
    with ``fetch_or`` rather than ``store`` so a query that touches no foreign
    object cannot clear a flag an earlier one set. Written with ``store``,
    ``SELECT 1`` here clears ``foreign_provider_observed``, and every other
    test asserting these flags after more than one query is a coincidence away
    from failing.
    """
    ctx, _logical_codec, _physical_codec = configured_context(max_rows=3)
    planner = MyQueryPlanner()
    ctx.set_query_planner(planner)

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert planner.foreign_session_observed()
    assert planner.foreign_provider_observed()
    assert planner.foreign_plan_observed()

    # Touches no table, so this plan has no foreign provider of its own.
    ctx.sql("SELECT 1").collect()
    assert planner.foreign_session_observed()
    assert planner.foreign_provider_observed()
    assert planner.foreign_plan_observed()

    # `last_max_rows` is deliberately not cumulative; it reports the last plan.
    assert planner.last_max_rows() == 3
    assert planner.plan_calls() >= 2


def test_second_planner_replaces_the_first():
    """A session holds exactly one planner, so installing another replaces it."""
    ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)
    first = MyQueryPlanner()
    second = MyQueryPlanner()
    ctx.set_query_planner(first)
    ctx.set_query_planner(second)

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert second.plan_calls() > 0
    assert first.plan_calls() == 0


def test_the_planner_reaches_every_handle_on_the_session():
    """The planner is session state, so all handles on that session use it.

    ``set_query_planner`` writes through ``state_ref()``. A context returned by
    an earlier ``with_*`` call shares that session, so it plans through the new
    planner too -- there is one session and one planner, not a family of
    diverging copies.
    """
    ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)
    # Shares the session with `ctx`; predates the planner.
    sibling = ctx.with_python_udf_inlining(enabled=False)

    planner = MyQueryPlanner()
    ctx.set_query_planner(planner)

    batches = sibling.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner.plan_calls() > 0


def test_installed_codecs_outlive_python_exporters():
    ctx, logical_codec, physical_codec = configured_context(max_rows=2)
    del logical_codec, physical_codec
    gc.collect()

    ctx.set_query_planner(MyQueryPlanner())
    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]


def test_provider_codecs_can_be_installed_after_planner():
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    planner = MyQueryPlanner()
    logical_codec = MyLogicalExtensionCodec()
    physical_codec = MyPhysicalExtensionCodec()
    ctx = SessionContext(config)
    ctx.set_query_planner(planner)
    ctx = ctx.with_logical_extension_codec(logical_codec)
    ctx = ctx.with_physical_extension_codec(physical_codec)
    ctx.register_table("numbers", MyTableProvider(1, 4, 1))

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner.last_max_rows() == 2
    assert logical_codec.table_provider_decode_calls() > 0
    assert physical_codec.execution_plan_decode_calls() > 0


def test_a_discarded_derived_context_still_rebinds_the_planner():
    """Installing a codec rebinds the planner even if the handle is thrown away.

    `with_logical_extension_codec` returns a context sharing this session, and
    rebuilding the installed planner against the new codec happens on that
    shared session rather than on the returned handle. So the rebind outlives
    the handle, and the codec below takes effect on `ctx` despite `ctx`'s own
    codec field never changing.

    That is spooky enough to be worth pinning as a decision. It is also forced:
    `FFI_QueryPlanner` holds its codecs by value, so a planner cannot read the
    session's current codecs at plan time and the rebuild has to be eager.

    A fresh codec instance is what makes it observable -- it carries its own
    counters, and the planner encodes the outbound logical plan with whichever
    codec it is holding. The base context deliberately installs no logical
    codec, so this one is the only candidate; an already-installed codec would
    claim the provider first and hide the rebind.
    """
    ctx, _physical_codec = physical_only_context(max_rows=3)
    ctx.set_query_planner(MyQueryPlanner())

    later = MyLogicalExtensionCodec()
    # Deliberately discarded. The rebind still lands on the shared session.
    ctx.with_logical_extension_codec(later)
    gc.collect()

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert later.table_provider_encode_calls() > 0


def test_the_planner_and_the_handle_can_hold_different_codecs():
    """One session, two codecs in effect, depending on the path taken.

    The planner is session state and carries whichever codecs installed it
    last -- here, ones that arrived through a handle that was discarded.
    Everything else on a context uses that context's own codec field, which
    the discarded handle never touched. So `Expr.to_bytes(ctx)` and
    `ctx.sql(...)` encode with different codecs on the same `ctx`.

    Chaining ``ctx = ctx.with_...(...)`` keeps the two in step; this pins what
    happens when they are allowed to diverge.

    ``ctx`` installs no logical codec of its own, so its chain is empty and the
    planner's holds exactly one entry. That asymmetry is what makes the split
    visible: an entry on both chains would be claimed by the same codec either
    way, since encoding stops at the first codec to claim an object.

    Inlining has to be off for the assertion to say anything: with it on, a
    Python UDF is encoded inline by ``PythonLogicalCodec`` and never reaches
    the installed codec's ``try_encode_udf``.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    ctx = SessionContext(config).with_python_udf_inlining(enabled=False)
    ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    ctx.set_query_planner(MyQueryPlanner())

    planner_codec = MyLogicalExtensionCodec()
    # Discarded, but the planner keeps its codec.
    ctx.with_logical_extension_codec(planner_codec)
    gc.collect()

    identity = udf(
        lambda arr: arr,
        [pa.int64()],
        pa.int64(),
        volatility="immutable",
        name="identity_i64",
    )
    ctx.register_udf(identity)
    Expr.to_bytes(identity(col("A")), ctx)

    # Serializing through `ctx` uses `ctx`'s own codec field, which is empty --
    # the UDF goes out by name through the terminal codec.
    assert planner_codec.encode_udf_calls() == 0

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()

    # Planning through the same `ctx` uses the codec the planner was rebound to.
    assert planner_codec.table_provider_encode_calls() > 0


def test_an_unchanged_inlining_setting_leaves_the_planner_alone():
    """A no-op toggle must not rebind the session's planner.

    ``with_python_udf_inlining`` rebuilds the handle's codecs and rebinds the
    session's planner to them. Asking for the setting a context already has
    changes nothing, so it must not pay that side effect.

    Observable only once the planner is holding some *other* handle's codec:
    without the guard, a defensive no-op toggle on `ctx` drags the planner back
    onto `ctx`'s codecs and silently undoes the install below. `ctx` installs no
    logical codec, so being dragged back leaves the planner with an empty chain
    and the query fails outright rather than quietly using the wrong codec.
    """
    ctx, _physical_codec = physical_only_context()
    ctx.set_query_planner(MyQueryPlanner())

    planner_codec = MyLogicalExtensionCodec()
    ctx.with_logical_extension_codec(planner_codec)  # discarded; planner keeps it
    gc.collect()

    # The default is on, so this asks for what `ctx` already has.
    ctx.with_python_udf_inlining(enabled=True)
    gc.collect()

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert planner_codec.table_provider_encode_calls() > 0


def test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs():
    """A planner is built against the codecs of the handle it is installed from.

    The sequel to the test above, and the trap it sets up. Once a discarded
    derived handle has rebound the session's planner to its codec, installing
    the same planner again from the *original* handle rebuilds it against that
    handle's codec instead -- which never changed. The session's planner tracks
    whichever handle wrote it last, not the newest codec installed anywhere.

    So "re-install the planner after installing a codec" only repairs anything
    when it is done from the handle holding the new codec.
    """
    ctx, _physical_codec = physical_only_context()
    planner = MyQueryPlanner()
    ctx.set_query_planner(planner)

    later = MyLogicalExtensionCodec()
    # Deliberately discarded, exactly as in the test above.
    ctx.with_logical_extension_codec(later)
    gc.collect()

    ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert later.table_provider_encode_calls() > 0

    # `ctx`'s own codec field never changed -- it never had a logical codec --
    # so this rebuilds the planner against an empty chain and drops `later`.
    ctx.set_query_planner(planner)
    encodes_by_later = later.table_provider_encode_calls()

    with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"):
        ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert later.table_provider_encode_calls() == encodes_by_later


def test_query_planner_requires_provider_codec():
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    ctx = SessionContext(config)
    ctx.register_table("numbers", MyTableProvider(1, 3, 1))
    ctx.set_query_planner(MyQueryPlanner())

    with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"):
        ctx.sql('SELECT "A" FROM numbers').collect()


@pytest.mark.parametrize("max_rows", ["0", "oops"])
def test_query_planner_rejects_invalid_config(max_rows: str):
    ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)
    ctx.set_query_planner(MyQueryPlanner())

    with pytest.raises(Exception, match=r"max_rows|Invalid value"):
        ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect()


class ProviderCodecsExtension:
    """Bundles the provider library's codecs for ``with_extensions``.

    These codecs keep their own private task-context provider, so they only
    need to be created once; the bundle can hand out the same exporters on
    every call.
    """

    def __init__(self) -> None:
        self.logical_codec = MyLogicalExtensionCodec()
        self.physical_codec = MyPhysicalExtensionCodec()

    def __datafusion_session_extension__(
        self, ctx: SessionContext
    ) -> SessionExtensionComponents:
        return SessionExtensionComponents(
            logical_extension_codecs=(self.logical_codec,),
            physical_extension_codecs=(self.physical_codec,),
        )


class _NamedCodec:
    """Forwards a codec's capsule getters under a declared id.

    ``with_extensions`` takes no ``codec_id=``, so an extension that ships a
    codec class another extension also ships declares
    ``__datafusion_codec_id__`` on the object it hands over. Both getters are
    forwarded because one wrapper stands in for whichever kind it wraps.
    """

    def __init__(self, codec: object, codec_id: str) -> None:
        self._codec = codec
        self.__datafusion_codec_id__ = codec_id

    def __datafusion_logical_extension_codec__(self, session: object = None) -> object:
        return self._codec.__datafusion_logical_extension_codec__(session)

    def __datafusion_physical_extension_codec__(self, session: object = None) -> object:
        return self._codec.__datafusion_physical_extension_codec__(session)


class IdentifiedProviderCodecsExtension(ProviderCodecsExtension):
    """``ProviderCodecsExtension`` whose codecs carry ids of their own."""

    def __init__(self, prefix: str) -> None:
        super().__init__()
        self.logical = _NamedCodec(self.logical_codec, f"{prefix}.logical")
        self.physical = _NamedCodec(self.physical_codec, f"{prefix}.physical")

    def __datafusion_session_extension__(
        self, ctx: SessionContext
    ) -> SessionExtensionComponents:
        return SessionExtensionComponents(
            logical_extension_codecs=(self.logical,),
            physical_extension_codecs=(self.physical,),
        )


def test_with_extensions_three_library_query():
    """One with_extensions call installs provider codecs and a planner bundle,
    and a real non-empty plan flows across the three libraries."""
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    provider_ext = ProviderCodecsExtension()
    planner_ext = MyPlannerExtension()
    ctx = SessionContext(config).with_extensions(provider_ext, planner_ext)
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    ctx.register_udf(udf(IsNullUDF()))

    batches = ctx.sql(
        'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"'
    ).collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert batches[0].column(1).to_pylist() == [False, False, False]
    assert planner_ext.plan_calls() >= 1
    assert planner_ext.last_max_rows() == 3
    assert planner_ext.foreign_session_observed()
    assert planner_ext.foreign_provider_observed()
    assert planner_ext.foreign_plan_observed()
    assert provider_ext.logical_codec.table_provider_encode_calls() > 0
    assert provider_ext.logical_codec.table_provider_decode_calls() > 0
    assert provider_ext.physical_codec.execution_plan_encode_calls() > 0
    assert provider_ext.physical_codec.execution_plan_decode_calls() > 0


def test_with_extensions_names_a_rust_bundles_capsules_after_the_bundle():
    """A Rust bundle hands its codecs over as bare capsules, and they are
    named after the bundle's own import path.

    This is the identity that has to survive leaving the process: a plan a
    distributed engine writes here is decoded by its scheduler, which installs
    a codec under the same id. A session-private random id — what a bare
    capsule gets when installed directly — would make the plan undecodable
    there.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    ctx = SessionContext(config).with_extensions(
        ProviderCodecsExtension(), MyPlannerExtension()
    )

    bundle_id = "datafusion_ffi_query_planner_example.MyPlannerExtension"
    assert bundle_id in ctx.logical_extension_codec_ids()
    assert bundle_id in ctx.physical_extension_codec_ids()

    # The provider bundle hands over objects, so those keep their own class
    # names rather than picking up the bundle's.
    assert (
        "datafusion_ffi_example.MyLogicalExtensionCodec"
        in ctx.logical_extension_codec_ids()
    )
    assert not any(
        codec_id.startswith("anon:") for codec_id in ctx.logical_extension_codec_ids()
    )


def test_with_extensions_shares_the_session_with_the_source():
    """``with_extensions`` returns a handle on the source's session, and the
    bundle's task-context provider resolves against that one session.

    There is one ``Arc<SessionContext>`` per session, so a component bound
    during installation cannot be left pointing at a handle that is dropped
    later. A `SET` issued through the *source* after installation is therefore
    visible to the provider the bundle bound, which is what a codec's decode
    callback resolves through.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    source = SessionContext(config)
    planner_ext = MyPlannerExtension()
    result = source.with_extensions(ProviderCodecsExtension(), planner_ext)

    assert result.session_id() == source.session_id()

    # Registrations and config changes go through the source handle only.
    source.register_table("numbers", MyTableProvider(1, 6, 1))
    source.sql("SET ffi_query_planner.max_rows = 2").collect()

    batches = result.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner_ext.last_max_rows() == 2
    assert planner_ext.max_rows_through_provider() == 2

    # Symmetrically, the codec chains installed on the shared session are in
    # force for the source handle too.
    batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]


def test_with_extensions_survives_dropping_source_and_bundles():
    """The returned handle alone keeps the installed components alive.

    The context ``with_extensions`` was called on is a temporary here, and the
    bundle objects are dropped with it. Both share their allocation with the
    returned handle, so the components' task-context provider stays valid.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    ctx = SessionContext(config).with_extensions(
        ProviderCodecsExtension(), MyPlannerExtension()
    )
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    gc.collect()

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]


def test_with_extensions_sees_state_changes_after_install():
    """Tables, UDFs, and config changes made after installation are visible
    to the planner and to provider callbacks."""
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=4))
    planner_ext = MyPlannerExtension()
    ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), planner_ext)

    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    ctx.register_udf(udf(IsNullUDF()))
    ctx.sql("SET ffi_query_planner.max_rows = 2").collect()

    batches = ctx.sql(
        'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"'
    ).collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner_ext.last_max_rows() == 2
    assert planner_ext.max_rows_through_provider() == 2


def test_with_extensions_bundle_is_reusable():
    """Installing the same bundle into two contexts binds fresh components to
    each destination."""
    planner_ext = MyPlannerExtension()

    config_a = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    ctx_a = SessionContext(config_a).with_extensions(
        ProviderCodecsExtension(), planner_ext
    )
    ctx_a.register_table("numbers", MyTableProvider(1, 6, 1))

    config_b = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
    ctx_b = SessionContext(config_b).with_extensions(
        ProviderCodecsExtension(), planner_ext
    )
    ctx_b.register_table("numbers", MyTableProvider(1, 6, 1))

    batches = ctx_a.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner_ext.last_max_rows() == 2

    batches = ctx_b.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2]
    assert planner_ext.last_max_rows() == 3


def test_with_extensions_failure_leaves_source_usable():
    """A failing factory after a successful one leaves the source context
    fully functional."""

    class BoomExtension:
        def __datafusion_session_extension__(
            self, ctx: SessionContext
        ) -> SessionExtensionComponents:
            msg = "boom"
            raise RuntimeError(msg)

    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    source = SessionContext(config)
    source.register_table("numbers", MyTableProvider(1, 6, 1))

    with pytest.raises(RuntimeError, match="boom"):
        source.with_extensions(MyPlannerExtension(), BoomExtension())

    # No planner was installed, so the default planner runs unrestricted.
    batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5]


def test_with_extensions_rebinds_existing_planner():
    """Codec-only bundles installed on a context that already has an FFI
    planner rebind that planner to the new codec chains."""
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    planner = MyQueryPlanner()
    ctx = SessionContext(config)
    ctx.set_query_planner(planner)
    provider_ext = ProviderCodecsExtension()
    ctx = ctx.with_extensions(provider_ext)
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert planner.last_max_rows() == 2
    # The planner only sees these codecs if it was rebound to the chains
    # built during with_extensions.
    assert provider_ext.logical_codec.table_provider_decode_calls() > 0
    assert provider_ext.physical_codec.execution_plan_decode_calls() > 0


def test_with_extensions_rejects_two_bundles_of_the_same_codec_class():
    """Two bundles contributing the same codec class collide on id.

    Ids are derived from the exporting class, so two instances of one class
    claim the same id. A payload names its codec by id when it is decoded, so
    the ambiguity is refused at install time rather than resolved by position.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    with pytest.raises(ValueError, match="is already installed on this session"):
        SessionContext(config).with_extensions(
            ProviderCodecsExtension(), ProviderCodecsExtension(), MyPlannerExtension()
        )


def test_with_extensions_accepts_distinct_codec_ids():
    """Declaring ``__datafusion_codec_id__`` resolves the collision above.

    Both codec pairs then install, and the query still runs end to end: only
    the codec that wrote a payload is asked to decode it, so the second pair
    is simply never consulted.
    """
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    ext_a = ProviderCodecsExtension()
    ext_b = IdentifiedProviderCodecsExtension("second")
    ctx = SessionContext(config).with_extensions(ext_a, ext_b, MyPlannerExtension())
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]

    ids = ctx.logical_extension_codec_ids()
    assert "datafusion_ffi_example.MyLogicalExtensionCodec" in ids
    assert "second.logical" in ids

    # The first pair wrote the payloads, so decoding routes back to it alone.
    assert ext_a.logical_codec.table_provider_encode_calls() > 0
    assert ext_a.logical_codec.table_provider_decode_calls() > 0
    assert ext_b.logical_codec.table_provider_decode_calls() == 0


def test_dataframe_outliving_context_fails_cleanly():
    """A DataFrame does not keep its SessionContext alive. FFI components
    resolve the task context through a weak reference, so using the
    DataFrame after dropping the context raises a clean error instead of
    crashing. This locks in the documented ownership contract: the context
    must outlive DataFrames that depend on FFI codecs."""
    config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
    ctx = SessionContext(config).with_extensions(
        ProviderCodecsExtension(), MyPlannerExtension()
    )
    ctx.register_table("numbers", MyTableProvider(1, 6, 1))
    df = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"')
    del ctx
    gc.collect()

    with pytest.raises(Exception, match="went out of scope"):
        df.collect()


def test_composed_codecs_with_query_planner():
    """A second pair of codecs installed on top of the provider codecs
    composes with them instead of replacing them.

    The provider codecs are installed first, so encoding consults them
    first and they claim this library's tables and plans before the
    extra codecs (default-backed exports from a fresh session) get a
    turn; decoding goes straight to whichever codec wrote the bytes.
    The extra pair therefore changes nothing observable here, which is
    the assertion: under replace semantics it would have discarded the
    provider codecs and the planner-driven round trip would fail."""
    ctx, logical_codec, physical_codec = configured_context(max_rows=2)
    other = SessionContext()
    ctx = ctx.with_logical_extension_codec(
        other.__datafusion_logical_extension_codec__()
    )
    ctx = ctx.with_physical_extension_codec(
        other.__datafusion_physical_extension_codec__()
    )
    ctx.set_query_planner(MyQueryPlanner())

    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
    assert batches[0].column(0).to_pylist() == [0, 1]
    assert logical_codec.table_provider_encode_calls() > 0
    assert logical_codec.table_provider_decode_calls() > 0
    assert physical_codec.execution_plan_decode_calls() > 0


class _DocstringExampleExtension:
    """Stand-in for the ``my_extension`` bundle named in the docstring.

    The docstring shows a single engine bundle taking a scheduler address,
    which is what a real distributed engine ships: one object contributing a
    planner *and* the codecs that carry its plans. Here that is assembled from
    this repository's two example libraries. The address is accepted and
    ignored; everything else the example touches is the real API.
    """

    def __init__(self, endpoint: str) -> None:
        self.endpoint = endpoint
        self._codecs = ProviderCodecsExtension()
        self._planner = MyPlannerExtension()

    def __datafusion_session_extension__(
        self, ctx: SessionContext
    ) -> SessionExtensionComponents:
        codecs = self._codecs.__datafusion_session_extension__(ctx)
        planner = self._planner.__datafusion_session_extension__(ctx)
        return SessionExtensionComponents(
            logical_extension_codecs=(
                *codecs.logical_extension_codecs,
                *planner.logical_extension_codecs,
            ),
            physical_extension_codecs=(
                *codecs.physical_extension_codecs,
                *planner.physical_extension_codecs,
            ),
            query_planner=planner.query_planner,
        )


def test_with_extensions_docstring_example_still_runs():
    """Run the ``with_extensions`` docstring example verbatim.

    The example is marked ``+SKIP`` because the main suite has no built FFI
    extension to import, which is exactly how such an example rots. Here the
    statements are parsed out of the live docstring, the skip is dropped, and
    each one is executed and its output compared.

    Only names are redirected: ``my_extension`` resolves to the bundle above,
    and ``SessionContext`` supplies the config this library's planner reads.
    A renamed method, a changed signature, or a wrong expected output in the
    docstring fails here.
    """
    examples = doctest.DocTestParser().get_examples(
        inspect.getdoc(SessionContext.with_extensions)
    )
    assert examples, "with_extensions docstring has no examples to check"
    for example in examples:
        example.options.pop(doctest.SKIP, None)

    module = types.ModuleType("my_extension")
    module.DistributedEngineExtension = _DocstringExampleExtension

    def make_context(config: SessionConfig | None = None) -> SessionContext:
        # Accept a config so the example is free to pass one. Supplying it
        # positionally the way the real constructor does keeps a docstring
        # edit failing as a doctest diff rather than as a TypeError in here.
        config = SessionConfig() if config is None else config
        return SessionContext(config.with_extension(MyPlannerConfig(max_rows=3)))

    test = doctest.DocTest(
        examples,
        {"SessionContext": make_context},
        "SessionContext.with_extensions",
        None,
        None,
        None,
    )
    output = io.StringIO()
    sys.modules["my_extension"] = module
    try:
        results = doctest.DocTestRunner().run(test, out=output.write)
    finally:
        del sys.modules["my_extension"]
    assert results.failed == 0, output.getvalue()