Repository navigation
Extend SessionContext.with_extensions to cover additional extension points (functions, catalogs, object stores) #1676
Description
Activity
Now that #1679 has merged, the shape it settled on differs from what this issue was written against, so recording what went stale and what the plan is.
What is stale
This issue says Reality after #1679 __datafusion_session_extension__Renamed __datafusion_session_components__SessionExtensionExportableRenamed SessionComponentsExportable, with a new siblingSessionPlannerExportable"#1672 adds with_extensions"#1672 closed unmerged; the work landed as #1677 → #1678 → #1679 Components contain "optionally a query planner" Planners are a separate second-phase hook. SessionExtensionComponentssays so directly: "Query planners are not listed here.""matching the one-planner rule" The per-call one-planner refusal was removed in #1679 — planners nest now, so the analogy has no referent. The surviving precedent for erroring is the duplicate-codec-id ValueErrorinresolve_bundle_codec_id.Design decision 1, on the derived context sharing its catalog provider list Moot. _derive_for_extensionsno longer exists;_install_extension_codecsdoesArc::cloneon the source context, so there is oneArc<SessionContext>per session and no derivation point to isolate at. Deriving one is the hazard #1679 removed."the Rust _install_extensionshelper is private"It split into three private primitives ( _install_extension_codecs,_export_query_planner,_install_extension_planner) and the composition moved to Python. Adding the function fields is now Python-only, with no Rust at all."frozen dataclass with defaulted fields, so this can be added incrementally" Still true, but __post_init__normalizes fields by the_codecsname suffix, and its comment reserves non-suffixed fields for things that must not become tuples — pointed the wrong way for audfsfield. Needs replacing with field metadata.The core premise is not stale. #1679's own description renamed the hook "leaving room for the UDF and provider fields that will join the codec fields later."
Dropping
object_storesThere is no FFI object store type upstream —
datafusion-ffi55 has no such module, and this repository has no__datafusion_object_store__hook.register_object_storetakesStorageContexts, a closed enum over five built-in pyclasses, so the field could only ever carry datafusion-python's own objects. A third-party Rust cdylib cannot produce one at all; it would have to importdatafusion.object_storeand call back into the host. A library that wants to ship a configured store hands the user anAmazonS3and the user callsregister_object_store, which is already one line.I will file this separately as blocked on an upstream
FFI_ObjectStore. Everything else in the list stays in scope: functions, table providers, catalog providers, and physical optimizer rules.The constraint that shapes the rest
The issue's argument for declarative components — the host validates everything before mutating anything — is still the right one, but it needs restating now that there is no derived context.
with_extensionsis currently transactional for free because codec chains live on the returned Python handle rather than onSessionState, and_install_extension_planneris the single write. Every field proposed here writes into the sharedSessionState, and the returned handle is the same session as the source, so there is nothing to roll back to.The only design that keeps "nothing is written until every hook has returned and every capsule has been validated" literally true is to make the commit phase provably infallible:
1. collect every __datafusion_session_components__ fallible, writes nothing 2. chains _install_extension_codecs fallible, writes nothing 3. resolve declared components -> concrete objects fallible, writes nothing + every __datafusion_session_planner__ fallible, writes nothing 4. commit _install_extension_planner, then register_* must be infallibleRollback should not be attempted as an alternative.
deregister_udfremoves a name but does not restore a built-in the bundle shadowed, andregister_catalogreturns the displaced provider that the commit would have discarded — so a rollback that deletes a user's pre-existing registration is worse than a partial apply.Checking the commit halves against upstream:
SessionContext::register_udfreturns()and swallows errors, so it is infallible.register_catalogreturns the displaced provider, also infallible, butPySessionContext::register_catalog_providerdoes capsule import and insert in one function with no split point, so it needs a new private primitive.register_tablereturns aResultthroughschema_for_ref, so the schema lookup has to move into the resolve phase.add_physical_optimizer_ruledoes a fullSessionStateBuilder::new_from_existing(...).build()per rule, so N rules is N whole-state clones and a failure on rule 3 leaves 1 and 2 committed; it needs batching into one rebuild.Where resolution has to happen
The fields divide by what their capsule getter needs, and that decides the ordering more than anything else does:
Group A — the getter takes no argument.
udfs,udafs,udwfs,physical_optimizer_rules. Resolution is session-independent and, for the function kinds, already happens in the extension library's own frame before the bundle hands anything over.Group B — the getter receives the session or the logical codec.
udtfs(PyTableFunction::newpasses the session in),table_providers(PyTable::new(table, Some(session))),catalog_providers(builds a capsule from the session's logical codec and passes that in). These must resolve against the handle returned by_install_extension_codecs, not against thectxthe components hook received. A provider resolved against the pre-install context encodes through a chain missing every other bundle's codec, and the failure does not surface until a decode in a different process. This is the same hazard the guide already documents for phase one: "Reading them in phase one gets the chains from before the call, missing even the bundle's own codecs."Collisions
Duplicate names among components declared in the same
with_extensionscall are aValueErrornaming the name and both contributing bundles. Shadowing something the session already has stays legal —ctx.udfs()contains every DataFusion built-in, andenable_spark_functionsoverrides built-ins by design, so erroring on that would refuse legitimate overrides and would refuse"datafusion"for catalogs since the default catalog always exists. Namespaces are independent, so a UDF and a UDAF may share a name.Physical optimizer rules are exempt and never collide. The extension guide already says rules accumulate where planners nest, and erroring would contradict that.
Plan
Four stacked PRs, split on the Group A/B boundary plus whether new Rust is needed, since that is what changes the shape of the diff:
udfs,udafs,udwfs— Python only, no Rust. Smallest diff, largest conceptual load: it settles the field-metadata normalization mechanism, the resolve/commit structure, and the collision policy that the other three inherit.physical_optimizer_rules— one new private Rust primitive that imports every capsule and then applies them all to a singleSessionStaterebuild. Completes Group A.udtfsandtable_providers— Group B. No new Rust, sinceTableFunction(name, func, ctx)andTable(provider, ctx)are already constructible from Python; the work is placing resolution after the codec install and moving the schema lookup out of the commit.catalog_providers— a new private primitive splitting the capsule import out ofregister_catalog_provider. Closes this issue.
Each PR carries its own coverage across a real FFI boundary rather than deferring it, extending one bundle class in
examples/datafusion-ffi-exampleby one field as the stack proceeds — that crate already exports a fixture for every field in scope.No upgrade guide entry and no
api changelabel: every field is additive with a()default on a frozen dataclass, no hook signature changes, and no wire format changes.One honest caveat
For a library that ships a codec, worker-side function registration is already unnecessary — the codec alone rebuilds the function from the name in the plan, and the registry is tried first so the codec is the fallback rather than the path. The
udfsfield buys driver-side ergonomics, not worker parity. That is still worth shipping, it is just a smaller claim than the issue's framing makes.Worth stating in the guide too:
with_extensionscan never install everything a library provides. A config extension has to reachSessionConfigbefore the context exists, which no bundle hook can reach.- added 4 commits that reference this issue
on Sep 15, 2026
Is your feature request related to a problem or challenge?
#1672 adds
SessionContext.with_extensions(*extensions), which installs extension bundles atomically: each bundle implements__datafusion_session_extension__(ctx)and returns aSessionExtensionComponentscontaining logical extension codecs, physical extension codecs, and optionally a query planner. Those components carry weak task-context providers, so the atomic protocol exists to guarantee they bind to the exact context returned to the user.Extension libraries typically provide more than codecs and a planner. A distributed engine or data-source library may also ship scalar/aggregate/window/table functions, catalog providers, table providers, object stores, and physical optimizer rules. Today the user must install those with separate
register_*calls afterwith_extensions, which defeats the goal of a one-line installation of everything a library provides.Describe the solution you'd like
Extend
SessionExtensionComponentswith additional optional fields so a single bundle can declare every extension point its library offers:udfs,udafs,udwfs,udtfs(or a singlefunctionstuple with type dispatch) — accepting both FFI-capsule and Python-native definitions, flowing through the existingregister_*pathstable_providers: tuple[tuple[str, provider], ...]— named table registrationscatalog_providers: tuple[tuple[str, provider], ...]object_stores: tuple[tuple[str, store], ...]— keyed by schemephysical_optimizer_rulesSessionExtensionComponentsis a frozen dataclass with defaulted fields and the Rust_install_extensionshelper is private, so this is backward compatible and can be added incrementally. A distributed engine bundle could then install codecs, planner, its UDFs, and a scheduler-backed catalog in one call:Declarative components are preferred over having factories call
ctx.register_udf(...)on the context view directly: the host can validate every component before mutating anything, the configuration-only contract for factories stays honest, and installation order becomes explicit (codecs, then planner, then registrations, with fallible steps first).Design decisions to settle during implementation:
with_*method.Suggested starting scope:
udfsandcatalog_providers, with collision handling as an error.Describe alternatives you've considered
Factories can already register functions and catalogs imperatively on the context view passed to
__datafusion_session_extension__, since the view shares the destination context. This works but hides side effects inside the factory, cannot be validated up front, and leaves partial registrations behind on failure.Keeping registrations as separate user-facing
register_*calls afterwith_extensionsremains possible, but each library then documents its own multi-step setup recipe, which is the situationwith_extensionswas introduced to remove.Additional context
Follow-up to #1672. Relevant pieces there:
SessionExtensionComponents,SessionExtensionExportable, thewith_extensionstransaction incrates/core/src/context.rs, and theMyPlannerExtensionexample bundle inexamples/datafusion-ffi-query-planner-example.