Let extension bundles declare scalar, aggregate, and window functions

crates/core/src/context.rs

@@ -1663,10 +1663,9 @@ impl PySessionContext {
     /// **Writes nothing.** The codec chains belong to the returned handle
     /// rather than to `SessionState`, so this phase is transactional for free:
     /// a codec that fails to import, or that collides with an installed id,
-    /// leaves the caller's context exactly as it was. Binding the planner is
-    /// the only step that touches the session, and it is deferred to
-    /// [`Self::_install_extension_planner`] so the planner hooks can run
-    /// against the final chains.
+    /// leaves the caller's context exactly as it was. Everything that touches
+    /// the session is deferred to [`Self::_commit_extensions`] so the planner
+    /// hooks can run against the final chains.
     ///
     /// Codecs must arrive as objects exposing the capsule getter, never as
     /// bare capsules — see [`resolve_bundle_codec_id`].
@@ -1717,49 +1716,68 @@ impl PySessionContext {
         })
     }
 
-    /// Re-export a planner a `__datafusion_session_planner__` hook returned as
-    /// a capsule, so the next hook in the chain receives one either way.
+    /// Run the planner hooks and commit a `with_extensions` call.
     ///
-    /// A hook may hand back an object exposing `__datafusion_query_planner__`
-    /// or a raw capsule; the next hook wraps whatever it is given and should
-    /// not have to branch on which. Importing here also surfaces a malformed
-    /// planner at the hook that produced it rather than at the final install.
-    /// Writes nothing.
-    pub fn _export_query_planner<'py>(
-        slf: &Bound<'py, Self>,
-        planner: Bound<'py, PyAny>,
-    ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {
-        let ffi = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?;
-        Ok(create_query_planner_capsule(slf.py(), &ffi)?)
-    }
-
-    /// Commit the query planner for a `with_extensions` call.
+    /// The second phase, run on the handle carrying the completed chains —
+    /// `session` is that same handle as the Python-level wrapper, which is
+    /// what each `__datafusion_session_planner__` hook receives. The hooks
+    /// run first, **in argument order**, each handed the planner built so
+    /// far as a capsule; a hook may hand back an object exposing
+    /// `__datafusion_query_planner__` or a raw capsule, and each return is
+    /// imported here so a malformed planner surfaces at the hook that
+    /// produced it rather than at the install. Returning `None` contributes
+    /// no planner. All of that writes nothing, so a hook that raises leaves
+    /// the session exactly as it was.
     ///
-    /// The second phase, run once every codec is installed and every planner
-    /// hook has returned, so the planner is bound against the final chains.
-    /// This is the one call in `with_extensions` that writes to the session,
-    /// and it goes through this context's own `state_ref()`, so providers
-    /// bound to it stay valid.
+    /// Binding the planner is the commit, and it cannot fail. Anything else a
+    /// call installs is committed by the Python caller once this returns,
+    /// through `register_udf`, `register_udaf` and `register_udwf`, which
+    /// cannot fail. Whatever can fail belongs before this call — see
+    /// docs/source/contributor-guide/ffi-internals.md, under
+    /// "Why `with_extensions` commits last".
     ///
-    /// `None` means no bundle supplied a planner. That still rebuilds
-    /// whichever planner the session already holds against the new chains,
-    /// exactly as `with_logical_extension_codec` does, and writes nothing at
-    /// all if the session has no FFI planner to rebuild.
-    ///
-    /// The caller skips this step entirely when the call installed no codec
-    /// and no planner, the same way [`Self::with_python_udf_inlining`] returns
-    /// early for a no-op toggle: there is nothing to rebind against, and the
-    /// rebuild would drag a planner sitting on another handle's codecs onto
-    /// this one's.
-    #[pyo3(signature = (planner=None))]
-    pub fn _install_extension_planner<'py>(
+    /// The planner is bound through this context's own `state_ref()`, so
+    /// providers bound to it stay valid. With no planner supplied the bind
+    /// still rebuilds whichever planner the session already holds against
+    /// the new chains, exactly as `with_logical_extension_codec` does —
+    /// unless `rebind_planner` is also false, meaning the call installed no
+    /// codec either. Then the bind is skipped entirely, the same way
+    /// [`Self::with_python_udf_inlining`] returns early for a no-op toggle:
+    /// there is nothing to rebind against, and the rebuild would drag a
+    /// planner sitting on another handle's codecs onto this one's.
+    pub fn _commit_extensions<'py>(
         slf: &Bound<'py, Self>,
-        planner: Option<Bound<'py, PyAny>>,
+        extensions: Vec<Bound<'py, PyAny>>,
+        session: Bound<'py, PyAny>,
+        rebind_planner: bool,
     ) -> PyDataFusionResult<()> {
-        let planner = planner
-            .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any())))
-            .transpose()?;
-        slf.borrow().set_session_query_planner(planner);
+        let py = slf.py();
+        // Nest the planners, outermost last. `planner` stays `None` when no
+        // bundle supplies one, which leaves an already-installed planner in
+        // place rather than wrapping the session's default in an FFI hop.
+        let mut planner: Option<FFI_QueryPlanner> = None;
+        for extension in &extensions {
+            if !extension.hasattr("__datafusion_session_planner__")? {
+                continue;
+            }
+            let fallback = match &planner {
+                Some(ffi) => create_query_planner_capsule(py, ffi)?,
+                None => slf.borrow().__datafusion_query_planner__(py, None)?,
+            };
+            let supplied =
+                extension.call_method1("__datafusion_session_planner__", (&session, fallback))?;
+            if supplied.is_none() {
+                continue;
+            }
+            planner = Some(ffi_query_planner_from_pycapsule(
+                &supplied,
+                Some(slf.as_any()),
+            )?);
+        }
+
+        if planner.is_some() || rebind_planner {
+            slf.borrow().set_session_query_planner(planner);
+        }
         Ok(())
     }
 }

docs/source/contributor-guide/ffi-internals.md

@@ -111,6 +111,45 @@ library would serialize, and would do it with the codecs it was imported with.
 The extension-facing consequence — install codecs before a layered planner, and
 prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`.
 
+(ffi_internals_commit_order)=
+
+## Why `with_extensions` commits last
+
+`with_extensions` promises that a bundle which raises leaves the session as it
+was. Keeping that promise is an ordering constraint on the implementation,
+because the components a bundle declares no longer all live on the returned
+handle — functions are registered on the shared `SessionState`, and the planner
+is bound there too.
+
+A call therefore splits into a part that may fail and a part that may not:
+
+1. **Collect.** Every `__datafusion_session_components__` runs.
+2. **Chains.** The codecs are assembled into the returned handle. Codec chains
+   live on that handle rather than on the session, so this step writes nothing
+   even though it can fail on a bad capsule or a duplicate id.
+3. **Resolve.** Every declared function is wrapped and every name is checked,
+   and every `__datafusion_session_planner__` runs against the completed
+   chains.
+4. **Commit.** The planner is bound, then the declared functions are
+   registered with `register_udf`, `register_udaf` and `register_udwf` — the
+   same methods a caller would use, and three that cannot fail.
+
+Only step 4 touches the session, and every step that can fail happens before
+it. This is a rule for the next field added to
+`SessionExtensionComponents`, not only a description of the current code: a new
+kind of component must do its fallible work — importing a capsule, resolving a
+name — in step 3, so that step 4 cannot raise part-way through.
+
+There is nothing to roll back to if it does. The returned handle shares one
+session with the receiver, so the damage is visible from every other handle;
+and undoing a registration is not the same as restoring what it displaced,
+because deregistering a function that shadowed a built-in removes the built-in
+too. The split is cheaper than an undo log that cannot be written correctly.
+
+The extension-facing statement of this is
+{ref}`extension_bundles_transaction`, which says only that declaring a
+component is safe where registering one during the hook is not.
+
 ## Two argument kinds for one convention
 
 `CapsuleGetterArg` in `crates/util/src/lib.rs` distinguishes three cases: no

docs/source/extension-guide/bundles.md

@@ -21,7 +21,7 @@
 
 # Extension bundles
 
-If your library ships codecs, or a query planner, or both, expose a **bundle**
+If your library ships codecs, functions, or a query planner, expose a **bundle**
 and let callers install it with
 {py:meth}`~datafusion.SessionContext.with_extensions`. This is the recommended
 way to package an extension, and the rest of this page explains what the
@@ -45,6 +45,7 @@ class MyEngineExtension:
         return SessionExtensionComponents(
             logical_extension_codecs=(self._make_logical_codec(ctx),),
             physical_extension_codecs=(self._make_physical_codec(ctx),),
+            udfs=(MyScalarUDF(),),
         )
 
     def __datafusion_session_planner__(self, ctx: SessionContext, fallback):
@@ -54,16 +55,20 @@ class MyEngineExtension:
         return self._make_planner(ctx, fallback=fallback)
 ```
 
-Implement whichever apply: a codec-only library defines the first, a library
-that ships only an optimizing planner defines the second. The caller then
-writes:
+Implement only the hooks you need. Codecs and functions both go in
+`__datafusion_session_components__`, with the fields you do not use left empty,
+so a codec-only library and a function-only library each define that one alone;
+a library shipping nothing but an optimizing planner defines only
+`__datafusion_session_planner__`. The caller then writes:
 
 ```python
 ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension())
 ctx.register_table("t", lib_a.TableProvider())
-ctx.register_udf(udf(lib_b.SomeUDF()))
 ```
 
+Return your functions rather than calling `register_udf` on the `ctx` you were
+handed — see {ref}`extension_bundles_transaction`.
+
 `MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete
 Rust implementation of the protocol, including taking the task-context provider
 off the supplied context, wrapping its codecs in `BundledLogicalCodec` /
@@ -281,14 +286,76 @@ for direct ones. The wrapper travels with the codec; the bundle does not.
 The query planner is exempt — it carries no wire id, so it may be an object or
 a capsule.
 
+(extension_bundles_collisions)=
+
+## Two bundles claiming one name
+
+One `with_extensions` call may not declare a function of the same kind under
+the same name twice. Doing so raises:
+
+```text
+ValueError: A scalar function named 'normalize' is declared twice: argument 0
+(<lib_a.Extension>) and argument 1 (<lib_b.Extension>). ...
+```
+
+Codecs get away with sharing a chain because a payload carries the id of the
+codec that wrote it, so decode routes to the right one. A function registry has
+no such fall-through — one name holds one function — so the second registration
+would quietly replace the first. The call refuses instead.
+
+The argument each claim came from is in the message because it is what picks
+the remedy, and the reader works it out from the two positions. Two different
+positions means two bundles, which only the caller can resolve — by dropping a
+repeat, or by installing them on separate sessions. One position named twice
+means a single bundle claimed the name twice, which is its author's own bug and
+the one case a rename fixes. Collisions are keyed on position rather than on
+object identity so that passing one extension twice reads as the caller's
+duplicate that it is, rather than as a bundle colliding with itself.
+
+Three cases this does *not* catch:
+
+- **Different kinds never collide.** Names are compared within a kind, so a
+  scalar function and an aggregate may both be called `normalize`.
+- **Shadowing a built-in is allowed.** The registry already holds every
+  DataFusion function, and replacing one by name is a supported thing to do —
+  `enable_spark_functions` works that way.
+- **One call at a time is checked.** The names compared are the ones declared
+  in a single `with_extensions` call. Two calls on the same session, or a call
+  following a hand-written `register_udf`, land in the registry one after the
+  other, and the later one silently replaces the earlier — that is the
+  shadowing rule above, applied to something you may not have meant to shadow.
+  Splitting colliding bundles across two calls therefore does not resolve the
+  collision; it hides it. Use two sessions.
+
+Your caller cannot rename your function, so stay out of the way: prefix the
+names with something tied to your library.
+
+(extension_bundles_transaction)=
+
 ## Failure and rollback
 
 Nothing is written to the session until every factory has returned and every
-capsule has been validated, so a factory that raises leaves the session exactly
-as it was. A factory that mutates the context it is handed — registering a
-table, say — is **not** rolled back, which is why bundle objects must be
-configuration-only: create fresh components on each call, never cache bound
-components, and do not retain the context passed in.
+component has been validated, so a factory that raises leaves the session
+exactly as it was. A factory that mutates the context it is handed —
+registering a table, say — is **not** rolled back, which is why bundle objects
+must be configuration-only: create fresh components on each call, never cache
+bound components, and do not retain the context passed in.
+
+Declaring a component is what buys you that guarantee, and it is the whole
+reason to prefer `udfs=(...)` over a `register_udf` call inside your hook.
+Anything you declare is resolved and checked while a failure still costs
+nothing, and is written only after every bundle in the call has succeeded.
+Anything you register yourself is written immediately, before the other bundles
+have even run. The ordering that makes this hold is recorded at
+{ref}`ffi_internals_commit_order`.
+
+The one thing that ordering costs you: functions are registered *after* the
+planner hooks run, so `ctx.udfs()` inside your
+`__datafusion_session_planner__` will not list a function declared in the same
+call — not yours, and not another bundle's. Look one up at plan time instead,
+where the registry is complete; a planner is called per query, long after the
+install has finished. If you need a function at hook time, you already have the
+object, because you are the one declaring it.
 
 Like every other derivation, the returned context is a handle on the *same*
 session as the receiver — see {ref}`extension_sessions`. Only the Python-side

docs/source/extension-guide/checklist.md

@@ -59,13 +59,13 @@ publish. Each links to the page that explains it.
 
 ## Bundles and planners
 
-- [ ] **You ship a bundle, not loose pieces**, if you have codecs or a planner.
+- [ ] **You ship a bundle, not loose pieces**, if you have codecs, functions,
+      or a planner.
       → {ref}`extension_bundles`
 - [ ] **Your bundle is configuration-only.** Fresh components on every call,
       no cached bound components, no retaining the context passed in, no
-      registering anything on it — a factory that mutates the context is not
-      rolled back if a later factory raises.
-      → {ref}`extension_bundles`
+      registering anything on it — declare what you contribute instead.
+      → {ref}`extension_bundles_transaction`
 - [ ] **Your codecs are objects exposing the getter, not bare capsules.**
       `with_extensions` refuses a capsule, because there would be nothing to
       name the codec by. → {ref}`extension_bundles_codecs_are_objects`

docs/source/extension-guide/functions.md

@@ -26,14 +26,18 @@ functions in pure Python — see
 {doc}`../user-guide/common-operations/udf-and-udfa` — and the two roads meet at
 the same registration methods.
 
-| Hook | Contributes | Wrapped by | Registered with |
-| --- | --- | --- | --- |
-| `__datafusion_scalar_udf__` | scalar function | {py:func}`datafusion.udf` | {py:meth}`~datafusion.SessionContext.register_udf` |
-| `__datafusion_aggregate_udf__` | aggregate function | {py:func}`datafusion.udaf` | {py:meth}`~datafusion.SessionContext.register_udaf` |
-| `__datafusion_window_udf__` | window function | {py:func}`datafusion.udwf` | {py:meth}`~datafusion.SessionContext.register_udwf` |
-| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` |
-
-All four are implemented in [`datafusion-ffi-example`], one per file.
+| Hook | Contributes | Wrapped by | Registered with | Declared in a bundle as |
+| --- | --- | --- | --- | --- |
+| `__datafusion_scalar_udf__` | scalar function | {py:func}`datafusion.udf` | {py:meth}`~datafusion.SessionContext.register_udf` | `udfs` |
+| `__datafusion_aggregate_udf__` | aggregate function | {py:func}`datafusion.udaf` | {py:meth}`~datafusion.SessionContext.register_udaf` | `udafs` |
+| `__datafusion_window_udf__` | window function | {py:func}`datafusion.udwf` | {py:meth}`~datafusion.SessionContext.register_udwf` | `udwfs` |
+| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` | — |
+
+All four are implemented in [`datafusion-ffi-example`], one per file. The last
+column is the {py:class}`~datafusion.SessionExtensionComponents` field a
+{ref}`bundle <extension_bundles>` declares the function in; table functions
+have no such field yet, so they are always registered by the caller with
+{py:meth}`~datafusion.SessionContext.register_udtf`.
 
 ## The three scalar-shaped hooks
 
@@ -67,6 +71,39 @@ from datafusion import udf
 ctx.register_udf(udf(my_library.MyScalarUDF()))
 ```
 
+If your library ships more than a function or two, do not make your users write
+that line once per function. Ship a {ref}`bundle <extension_bundles>` declaring
+them, so one call installs the lot:
+
+```python
+ctx = SessionContext().with_extensions(my_library.MyFunctionExtension())
+```
+
+The bundle is yours to write, and like the rest of the protocol it is an object
+exposing a getter — which your cdylib can export directly. That is what
+`MyFunctionExtension` in [`datafusion-ffi-example`] does for this crate's three
+functions; spelled in Python, it is:
+
+```python
+from datafusion import SessionExtensionComponents
+
+
+class MyFunctionExtension:
+    def __datafusion_session_components__(self, ctx):
+        return SessionExtensionComponents(
+            udfs=(IsNullUDF(),),
+            udafs=(MySumUDF(),),
+            udwfs=(MyRankUDF(),),
+        )
+```
+
+Declare either the raw exportable, as here, or an already-wrapped
+{py:class}`~datafusion.user_defined.ScalarUDF`; the registered name comes off
+the function either way. Declare rather than calling `register_udf` inside the
+hook — see {ref}`extension_bundles_transaction` for why — and pick names that
+will not collide with another library's
+({ref}`extension_bundles_collisions`).
+
 ## Table functions
 
 A table function takes literal `Expr` arguments and returns a table provider,