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,70 @@ 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 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.
+    /// `extensions` must already be narrowed to the planner hooks: every entry
+    /// is called, with no check of its own. The Python caller decides which
+    /// arguments are hooks.
     ///
-    /// `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 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 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>(
+    /// 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".
+    ///
+    /// 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 {
+            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

@@ -117,37 +117,35 @@ prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`.
 
 `with_extensions` promises that a bundle which raises leaves the session as it
 was, apart from anything a hook writes to the context it is handed
-({ref}`extension_bundles_transaction`). Keeping that promise is an ordering constraint on the implementation, not
-a property of any one step, because the planner is bound on the shared
-`SessionState` rather than on the returned handle.
+({ref}`extension_bundles_transaction`). Keeping that promise is an ordering
+constraint on the implementation, not a property of any one step, because the
+components a bundle declares do not 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 four steps, of which only the last writes:
 
-1. **Collect.** Every `__datafusion_session_components__` runs and its codecs
-   are gathered. Nothing is installed yet, so a hook that raises here has
-   touched nothing.
+1. **Collect.** Every `__datafusion_session_components__` runs and its
+   components are gathered. Nothing is installed yet, so a hook that raises here
+   has touched nothing.
 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
    to the session even though it can fail on a bad capsule or a duplicate id.
-3. **Resolve.** Every `__datafusion_session_planner__` runs, in argument order,
-   against the completed chains, and each supplied planner is exported to a
-   capsule. Every hook that can raise has run by the end of this step.
-4. **Commit.** The accumulated planner is re-imported from its capsule and
-   bound, in a single `SessionState` rebuild. The bind is skipped entirely when
-   the call installed nothing, so an empty call does not drag a planner sitting
-   on another handle's codecs onto this one's.
-
-Only step 4 touches the session, and it is not itself infallible:
-`_install_extension_planner` runs `ffi_query_planner_from_pycapsule` before it
-calls `set_session_query_planner`, which cannot fail. The property that keeps
-the promise is therefore about order, not about any step being incapable of
-raising — every fallible operation, including the ones inside the commit,
-completes before the first write.
-
-That is the 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 — before anything is
-written, so no failure can leave the session half-updated.
+3. **Resolve.** Every declared function is wrapped and every name is checked.
+   Then every `__datafusion_session_planner__` runs, in argument order, against
+   the completed chains, and each supplied planner is imported as it is
+   returned.
+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. The bind is
+   skipped entirely when the call installed no planner and no codec, so an
+   empty call does not drag a planner sitting on another handle's codecs onto
+   this one's.
+
+Only step 4 touches the session, and every fallible operation completes before
+its first write. 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 no failure can leave the session half-updated.
 
 There would be nothing to roll back to if one did. The returned handle shares one
 session with the receiver, so the damage is visible from every other handle;

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,22 +286,80 @@ 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.
+
+Four 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.
+- **Aliases are not compared.** DataFusion registers a function under each of
+  its aliases as well as its name, but only the names are checked. A function
+  whose alias matches another declared function's name replaces it, or is
+  replaced by it, with nothing raised.
+
+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.
-
-Declaring a component is what buys you that guarantee. Anything you return from
-your hook is validated 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`.
+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`