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 swap between the 3 PRs in the 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 in a single transaction. A bundle is a reusable configuration object — typically shipped by a compiled extension library — that implements one or both of two new protocol hooks. Nothing is written to the session until every hook has returned and every capsule has been validated, so a failing bundle leaves the session as it was.

Codecs and planners install in two phases, because they compose differently. A session chains many codecs and dispatches between them by id, so codecs accumulate and their order does not affect decoding. A session holds exactly one query planner, so planners cannot accumulate — they compose by nesting, each wrapping the one before it. Phase one calls every bundle's __datafusion_session_extension__(ctx), which returns its codecs as a SessionExtensionComponents, and installs all of them. Phase two then calls every bundle's __datafusion_session_planner__(ctx, fallback), in argument order, handing each the planner built so far; wrapping fallback nests that bundle outside the previous one, so the last bundle listed ends up outermost. A bundle implements whichever hooks apply, so a library shipping only an optimizing planner does not have to return empty components.

Running the planner hooks after every codec is installed is the point of the split. The rebuild that follows a later codec install reaches only the outermost planner layer (upstream apache/datafusion#24762), so a fallback captured against a partial chain would stay stale forever. Within a with_extensions call there is no "afterwards" for any layer. This is also what lets several planner-shipping libraries be installed together at all: collecting planners from one hook meant every factory ran before anything was installed, so no bundle could see another's planner to wrap it.

Codecs must be handed over as objects exposing the capsule getter, never as bare PyCapsule objects. A codec's wire id is read off the object it arrives as, and a capsule has no type to read one from; with_extensions takes no codec_id=, so the capsule is refused with a message naming the getter to implement. Deriving the id from the contributing bundle instead is not a fix — the bundle is whatever object the caller passed, so an application that packages a library inside a bundle of its own would silently re-tag that library's payloads and they would stop decoding in the process that reads them. The inner library cannot defend against that no matter what it declares, and the mismatch does not surface until a decode fails elsewhere. Wrapping keeps the identity with the codec, which is what __datafusion_codec_id__ is for.

Ownership and lifetime. 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. There is one Arc<SessionContext> per session, so every weak provider a bundle creates stays valid for as long as any handle on that session is alive. The context-outlives-DataFrame contract is documented and tested: a DataFrame outliving every handle on its session fails with a clean out-of-scope error rather than crashing.

Example library. MyPlannerExtension in datafusion-ffi-query-planner-example is a complete Rust implementation of both hooks, including taking the host's task-context provider off the supplied context and wrapping its codecs in BundledLogicalCodec / BundledPhysicalCodec so they carry declared ids. Its planner emits a DistributedExec — an execution plan node private to that library — and its physical codec claims that node by downcast and rebuilds it on decode. That pairing is the normal shape for a bundle: a planner that emits its own executor is only useful alongside the codec that can serialize it, which is why the two ship together and why every codec is installed before any planner is bound.

Documentation. 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 multi-library registration recipe. It also now covers what codec order does and does not control: decoding routes by id and is never order-dependent, while encoding stops at the first codec that claims a node, so a codec claiming a broad category can take nodes from a library installed after it. The query still succeeds — only the library that wrote the bytes changes, which breaks a plan that has to decode elsewhere. Where a bundle needs one position for its codec and another for its planner, the guide shows contributing each half separately with a small adapter rather than reordering, and treats the low-level sequence as the last resort it is. docs/source/user-guide/upgrade-guides.md points at with_extensions from the planner-install section. Rule 2 of .ai/skills/ffi-capsule-protocol/SKILL.md documents the planner hook's extra argument and the codecs-are-objects rule; Rule 6 calls out with_extensions as where "a session keeps one Arc<SessionContext> for life" 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 / SessionPlannerExportable protocols with their __datafusion_session_extension__ and __datafusion_session_planner__ hooks. QueryPlannerExportable moved from datafusion.context to datafusion.extensions and is now exported from the package root alongside the rest of the family. The context-outlives-DataFrame ownership contract is now documented.

No breaking changes to existing APIs. Both new hooks are new in this PR, so no shipped extension library implements them yet and no upgrade-guide migration entry is needed for them.

Review notes

Several things changed during review and the earlier revisions read differently. Recording them here so the discussion above stays legible.

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. test_with_extensions_shares_the_session_with_the_source 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 were briefly named after the contributing bundle, and are now refused instead. Naming them looked like it closed the gap that made the natural shape for a distributed engine — a Rust bundle handing over capsules — the one shape that could not produce portable plans. It did not: because the id came from whichever object the caller passed to with_extensions, wrapping one bundle inside another silently re-tagged the inner library's codecs, which is the ordinary way an application presents several libraries as one. Requiring an object instead puts the identity on the codec, where composition cannot move it, and deletes the resolution arm rather than documenting a footgun. test_with_extensions_codec_ids_survive_composition pins it.

with_extensions accepted at most one planner per call before the two-phase split. That was the right refusal for a single-hook protocol — every factory ran before anything was installed, so two bundles that both captured the session's planner would both capture the default and one would silently win — but it meant two libraries that each ship a planner could not be installed together, and splitting them across two calls discarded the first with no diagnostic. The planner hook removes the refusal by having the host thread each planner into the next rather than letting bundles capture one.

with_extensions() with no arguments no longer raises. Every sibling varargs method — DataFrame.select, filter, sort, drop, window — accepts zero arguments and returns a no-op result, and the two existing "at least one" guards in the codebase both cover cases with no meaningful identity element. Installing no extensions has an obvious answer, and a caller assembling the list from a plugin registry should not have to special-case it being empty.

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.