Add method on SessionContext to add all extensions from one library

[timsaucer]

Which issue does this PR close?

Part 3 of 3 in the split of #1672. These are enabled as a github stack so you should be able to swab between the 3 PRs in github interface (above, next to the "Open" oval).

Rationale for this change

Working on the Ballista integration showed that chaining the low-level with_* methods is easy to get wrong. FFI codecs and planners carry a weak task-context provider, and a query planner is built against whatever codec chains exist at the moment it is installed — so installing a codec afterwards leaves the planner encoding through a chain that is missing a library, and the caller is responsible for an ordering rule nothing enforces. SessionContext.with_extensions makes the whole installation one step so there is no "afterwards".

What changes are included in this PR?

• SessionContext.with_extensions(*extensions) installs one or more extension bundles atomically. Each bundle implements the new __datafusion_session_extension__(ctx) protocol: it receives the session it is being installed on, creates fresh components bound to that session, and returns them as a SessionExtensionComponents (new public dataclass; SessionExtensionExportable is the matching typing protocol).
• Internally the factories bind against the receiving context, all codecs are installed, and only then is a planner bound against the final chains. The chains belong to the returned handle rather than to SessionState, so a codec-only install onto a session with no FFI planner writes no state at all; when there is a planner to bind — one a bundle supplied, or one already installed that has to be rebuilt against the new chains — that is a single write through the session's own state_ref(). There is one Arc<SessionContext> per session, so every weak provider a factory creates stays valid for as long as any handle on that session is alive.
• Like the individual with_* methods, the returned context is a handle on the same session as the receiver: catalogs, tables, registered functions, and configuration are the one session, and the planner is installed on that shared session even if the returned handle is discarded. Only the Python-side codec chains belong to the returned handle.
• Nothing is written to the session until every factory has returned and every capsule has been validated, so a failing factory leaves the session as it was. The exception is a factory that mutates the context it is handed — a registration made during binding is not rolled back — which is why bundle objects must be configuration-only.
• If no bundle supplies a planner, an existing FFI planner is rebound to the new codec chains; at most one bundle may supply a planner.
• A codec handed over as a bare PyCapsule is named after the bundle that contributed it. A capsule carries no type of its own, so it previously fell through to anon:<uuid4> — an id private to the installing session, which makes the plans it writes undecodable anywhere else. A bundle is a plain object, so its module.QualName is library-owned and exactly as stable across processes as an exporting codec class's, which id resolution already trusts. This matters for a distributed engine, whose plans have to decode in another process: the bundle writes ballista.BallistaExtension and the scheduler installs a codec under that same id. Resolution goes through derive_codec_id itself, so a bundle can pin its id with __datafusion_codec_id__ against a later class rename, and an id declared on the handed-over object — or that object's own class — still wins.
• The ownership contract is documented and tested: the session owns the installed components' providers, and a DataFrame outliving every handle on it fails with a clean out-of-scope error rather than crashing.
• MyPlannerExtension in the example crate is a complete Rust implementation of the protocol, including extracting the host's task-context provider from the supplied context. Its codecs record the config value they resolve through the weak provider, letting tests prove the provider resolves against the session that runs the query.
• docs/source/contributor-guide/ffi.md documents with_extensions as the preferred API for extension bundles, keeps low-level chaining as advanced usage, and includes a full three-library registration recipe. docs/source/user-guide/upgrade-guides.md points at it from the planner-install section.
• Rule 6 of .ai/skills/ffi-capsule-protocol/SKILL.md — "a session keeps one Arc<SessionContext> for life" — now calls out with_extensions as where that rule is easiest to get wrong, since "bind the components to the context you are about to return" reads like an instruction to derive one first.

Are there any user-facing changes?

New public APIs: SessionContext.with_extensions, SessionExtensionComponents, and the SessionExtensionExportable / __datafusion_session_extension__ protocol. The context-outlives-DataFrame ownership contract is now documented. No breaking changes to existing APIs.

Review notes

Two things changed during review that are worth calling out, since the earlier revision reads differently.

The installation no longer derives a separate context. It used to — _derive_for_extensions, using SessionContext::new_with_state(self.ctx.state()) — and bound the factories against that. Because new_with_state carries the session id over while minting a fresh Arc<RwLock<SessionState>>, that produced two live sessions reporting one session_id() with independent state: configuration and the function registry diverged, catalogs stayed shared, and both handles reported the same __datafusion_codec_id__, which is session:<session_id> and exists precisely to tell codec chains apart.

The fork also bought nothing. It was meant to stop components binding to an intermediate context that could later be collected, but with one Arc<SessionContext> per session there is no such intermediate — deriving one is what creates the hazard, which is why Rule 6 of the capsule-protocol skill already said to mutate SessionState in place rather than derive a replacement. _derive_for_extensions is gone and the factories are handed the receiver. test_with_extensions_provider_targets_returned_context was vacuous once the session is shared and has been replaced by test_with_extensions_shares_the_session_with_the_source, which asserts matching session ids and that a SET issued through the source after installation is visible to the provider the bundle bound. Reintroducing the fork fails it.

Bare capsules from a bundle are named rather than randomized, as described above. This closes the gap that made the natural shape for a distributed engine — a Rust bundle handing over capsules, which is what MyPlannerExtension does — the one shape that could not produce portable plans. Two bare capsules of one kind from one bundle collide on that single id and are refused: numbering them by position would be exactly the kind of id codec.rs rejects for anon:, one another library can mint the same value from, and it would break stored plans the first time a bundle reordered what it returns. The collision message now names both routes to a distinct identity; it previously offered only codec_id=, which is unreachable from with_extensions.

I considered adding explicit id fields to SessionExtensionComponents instead. Once bundle-derived ids exist that would only serve the two-capsules-of-one-kind case, which an author already solves by wrapping the capsule in an object declaring __datafusion_codec_id__ — not worth a permanent field on a public dataclass.

Related follow-ups, neither of which this PR needs: enable_url_table is now once again the only method that mints a second Arc<SessionContext> for a session, and it forks state while keeping the session id (#1708); an agent skill for extension authors, as suggested in review, is #1707.

[ntjohnson1]

I don't have a fully coherent thought here. IIUC this is mostly to manage life times across the FFI boundary. I do wonder if there is a slightly cleaner way to mange this but this seems fine for now to provide a safer avenue.

I wonder if it makes sense to have a todo for some datafusion python extension skill/s. Being able to generate the 3 library example from the skill might be a nice smoke test to verify. I suspect after getting things setup for ballista/datafusion-distributed keeping it up to date shouldn't be too bad but I do suspect it will require some guidance to make sure they are doing it safely.

[milenkovicm]

LGTM, thanks @timsaucer
i guess the only open point is chaining of codecs raised in previous pr

[milenkovicm]

Following up on comment apache/datafusion-ballista#2252 (review) and follow up on this PR. FFILogicalCodec::encode|decode_file_format break our intent to have fully working distributed execution.

I might be wrong but encode|decode_file_format might not be an easy fix, at least not in short term and not in datafusion 55 timeframe, hence i have a proposal to make.

DistributedExec in ballista is nothing but a GRPC wrapper, it takes a logical plan, calls grpc endpoint and returns a stream of record batches. Would it make sense to create a CallbackPlanner (we might need a CallbackExec) in datafusion-py which would take a python closure LogicalPlan -> Stream<RecordBatches> (or LogicalPlanBlob -> Stream<RecordBatches>). Ballista would provide a closure which implements grcp logic in python code.

As CallbackPlanner have same library marker as df python FFILogicalCodec will not be triggered and we should have possibility to fully integrate distributed execution.

Basically we could implement query planner in python (limited but working)

wdyt @timsaucer and @ntjohnson1 ?

[timsaucer]

Thanks for chasing this down — the write_* failures are real, but I think the fix is much smaller than a new planner API, and it is not in this PR.

Where it actually breaks. FFI_LogicalExtensionCodec has no vtable entries for file formats at all. Both hooks are hardcoded stubs in datafusion-ffi (src/proto/logical_extension_codec.rs, try_decode_file_format / try_encode_file_format → not_impl_err!("FFI does not support ...")). So any codec that crosses the FFI boundary loses file-format support, and df.write_csv/write_parquet/write_json build a LogicalPlan::Copy that carries a FileFormatFactory. That is the whole failure.

This PR already survives it on the datafusion-python side. PythonLogicalCodec::try_{encode,decode}_file_format go through chain_encode/chain_decode, which collect a chained codec's error and fall through to the terminal DefaultLogicalExtensionCodec. Built-in formats encode and decode fine there. The direction that fails is the mirror image: Ballista holds an FFI_LogicalExtensionCodec wrapping the Python session's codec and calls it directly, with no Default fallback behind it.

Two fixes, both much cheaper than a new API:

1. Ballista side, today, no upstream change. Compose the FFI codec you get from Python with a DefaultLogicalExtensionCodec fallback on the two file-format hooks — the same shape PythonLogicalCodec uses here. One file, unblocks DF 55.

2. Upstream datafusion-ffi, ~8 lines. Make those two hooks delegate to a local DefaultLogicalExtensionCodec instead of returning not_impl_err!. No new vtable entries, so no ABI break, so it can ship in a 55.x patch. It works because built-in file formats are fully self-describing protobuf and nothing has to cross the boundary: DefaultLogicalExtensionCodec::try_encode_file_format in datafusion-proto covers csv/json/arrow/avro/parquet by downcasting the factory locally. What is left unsupported is a custom FileFormatFactory owned by the foreign library, and I agree that one needs a real FFI_FileFormatFactory and will not make the 55 window — but it is also not what is breaking write_*.

On CallbackPlanner. I do not think the concept holds up, and I want to be concrete about why rather than just deferring it.

A QueryPlanner returning a CallbackExec leaf replaces the entire plan with one opaque node: no children, no partitioning DataFusion can see, no pushdown, no repartition, no limit, no statistics. That is not planning, it is interception at the root. It is coherent for exactly one architecture — thin client, full delegation — which is what Ballista's DistributedExec already is; the local plan is degenerate because none of the work is local. But it has no answer for a hybrid plan, a local table joined against a remote one, because there is nothing left to split on. All or nothing.

It also buys very little over what already ships:

plan_bytes = df.logical_plan().to_bytes(ctx)
batches = my_grpc_client.execute(plan_bytes)   # your gRPC logic, in Python
result = ctx.from_arrow(batches)               # any __arrow_c_stream__ object

All three exist today, no new API. The only thing CallbackPlanner adds on top is transparency — the user keeps writing ctx.sql(...).collect() and never sees the interception. That is real, but it is a thin return for a permanent public API, and a second planner mechanism sitting next to the FFI query planner means two ways to do one thing.

If the goal is partial delegation rather than whole-plan delegation, the design that answers it already exists: a table provider exported from Python (__datafusion_table_provider__, see docs/source/user-guide/io/table_provider.md). Remote data appears as a table, DataFusion plans around it, filter and projection pushdown work, and local and remote sources mix in one query. That composes; a root-level callback does not.

The part that makes me most hesitant is the stated rationale: "As CallbackPlanner have same library marker as df python, FFILogicalCodec will not be triggered." That shapes a public API around which cdylib the code is compiled into, in order to route around an eight-line gap in datafusion-ffi. Fixing the gap is the smaller and more durable change.

One genuinely open question, independent of all of the above — and I think it is the same "chaining of codecs" point you flagged. Once Ballista's codec is installed on the Python session it becomes a chain entry, and chain entries write a framed payload (DFPYCHN + codec id) rather than bare bytes. A plain Rust BallistaCodec in the scheduler will not strip that envelope. So how does the scheduler decode — does it link PythonLogicalCodec, or does Ballista's codec stay off the chain? I am happy to add whatever hook makes the first option workable; I would rather solve that than work around it.

Proposal: keep #1679 as is, since the file-format gap predates it and is orthogonal, and I will open the upstream issue for (2). If you still want the callback route after the above, let us give it its own issue so the design can be argued on its own terms.