Skip to content

Extend SessionContext.with_extensions to cover additional extension points (functions, catalogs, object stores) #1676

Description

@timsaucer

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 a SessionExtensionComponents containing 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 after with_extensions, which defeats the goal of a one-line installation of everything a library provides.

Describe the solution you'd like

Extend SessionExtensionComponents with additional optional fields so a single bundle can declare every extension point its library offers:

  • udfs, udafs, udwfs, udtfs (or a single functions tuple with type dispatch) — accepting both FFI-capsule and Python-native definitions, flowing through the existing register_* paths
  • table_providers: tuple[tuple[str, provider], ...] — named table registrations
  • catalog_providers: tuple[tuple[str, provider], ...]
  • object_stores: tuple[tuple[str, store], ...] — keyed by scheme
  • physical_optimizer_rules

SessionExtensionComponents is a frozen dataclass with defaulted fields and the Rust _install_extensions helper 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:

ctx = SessionContext(config).with_extensions(MyEngineExtension(address))

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:

  1. Catalog list sharing. The derived context shares its catalog provider list with the source context, so catalog registrations are visible in the source and are not rolled back if a later step fails. Registering catalogs last, after all fallible steps, closes most of the hole. Cloning the catalog provider list at derivation would give true isolation but diverges from the semantics of every other with_* method.
  2. Name collisions. Two bundles registering the same UDF name or catalog name should probably be an error (matching the one-planner rule) rather than last-wins, since registrations have no fall-through semantics like codec chains do.

Suggested starting scope: udfs and catalog_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 after with_extensions remains possible, but each library then documents its own multi-step setup recipe, which is the situation with_extensions was introduced to remove.

Additional context

Follow-up to #1672. Relevant pieces there: SessionExtensionComponents, SessionExtensionExportable, the with_extensions transaction in crates/core/src/context.rs, and the MyPlannerExtension example bundle in examples/datafusion-ffi-query-planner-example.

Activity

  1. timsaucer commented on Sep 15, 2026

    @timsaucer
    MemberAuthor

    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__
    SessionExtensionExportable Renamed SessionComponentsExportable, with a new sibling SessionPlannerExportable
    "#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. SessionExtensionComponents says 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 ValueError in resolve_bundle_codec_id.
    Design decision 1, on the derived context sharing its catalog provider list Moot. _derive_for_extensions no longer exists; _install_extension_codecs does Arc::clone on the source context, so there is one Arc<SessionContext> per session and no derivation point to isolate at. Deriving one is the hazard #1679 removed.
    "the Rust _install_extensions helper 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 _codecs name suffix, and its comment reserves non-suffixed fields for things that must not become tuples — pointed the wrong way for a udfs field. 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_stores

    There is no FFI object store type upstream — datafusion-ffi 55 has no such module, and this repository has no __datafusion_object_store__ hook. register_object_store takes StorageContexts, 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 import datafusion.object_store and call back into the host. A library that wants to ship a configured store hands the user an AmazonS3 and the user calls register_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_extensions is currently transactional for free because codec chains live on the returned Python handle rather than on SessionState, and _install_extension_planner is the single write. Every field proposed here writes into the shared SessionState, 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 infallible
    

    Rollback should not be attempted as an alternative. deregister_udf removes a name but does not restore a built-in the bundle shadowed, and register_catalog returns 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_udf returns () and swallows errors, so it is infallible. register_catalog returns the displaced provider, also infallible, but PySessionContext::register_catalog_provider does capsule import and insert in one function with no split point, so it needs a new private primitive. register_table returns a Result through schema_for_ref, so the schema lookup has to move into the resolve phase. add_physical_optimizer_rule does a full SessionStateBuilder::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::new passes 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 the ctx the 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_extensions call are a ValueError naming the name and both contributing bundles. Shadowing something the session already has stays legal — ctx.udfs() contains every DataFusion built-in, and enable_spark_functions overrides 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:

    1. 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.
    2. physical_optimizer_rules — one new private Rust primitive that imports every capsule and then applies them all to a single SessionState rebuild. Completes Group A.
    3. udtfs and table_providers — Group B. No new Rust, since TableFunction(name, func, ctx) and Table(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.
    4. catalog_providers — a new private primitive splitting the capsule import out of register_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-example by one field as the stack proceeds — that crate already exports a fixture for every field in scope.

    No upgrade guide entry and no api change label: 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 udfs field 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_extensions can never install everything a library provides. A config extension has to reach SessionConfig before the context exists, which no bundle hook can reach.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions