feat: let extension bundles declare scalar, aggregate, and window functions

`SessionExtensionComponents` gains `udfs`, `udafs`, and `udwfs`, so a library
shipping functions can be installed with one `with_extensions` call instead of
documenting a per-function `register_*` recipe. Either the Python wrapper or a
raw capsule exportable is accepted; the registered name comes off the function.

Installation now splits into a fallible part and an infallible one. Collecting
hooks, building the codec chains, resolving the declared functions, and running
the planner hooks all write nothing; only the final step binds the planner and
registers. That keeps "nothing is written until every hook has returned" true
now that components reach the shared `SessionState`, where there is nothing to
roll back to. A new comment states the rule for whoever adds the next field.

Two extensions declaring one name in a single call is a `ValueError` naming
both, since a function registry has no fall-through the way a codec chain does.
Shadowing a name the session already has stays legal, which
`enable_spark_functions` relies on.

`__post_init__` now normalizes fields by metadata rather than by the `_codecs`
name suffix, so the new fields are covered and later ones will be too.

`MyFunctionExtension` in `datafusion-ffi-example` declares this crate's three
functions across a real FFI boundary.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

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):
@@ -55,15 +56,22 @@ class MyEngineExtension:
 ```
 
 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:
+that ships only an optimizing planner defines the second, and a library that
+ships only functions defines the first and leaves the codec fields empty. 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()))
 ```
 
+Declare functions rather than registering them yourself inside the hook.
+Declared components are resolved before anything is written, and they are
+registered after every codec is installed; a registration you make during the
+hook happens too early to see the other bundles' codecs and is not undone if a
+later extension fails. Table providers are still registered by the caller, on
+the returned handle — 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 +289,52 @@ 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_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.
+
+That guarantee is why the installation runs in the order it does. A call 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 and the functions are registered.
+
+Only step 4 touches the session, and every step that can fail happens before
+it. This is a rule for anyone extending `with_extensions`, not only a
+description: 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, and undoing a registration is not the
+same as restoring what it displaced: deregistering a function that shadowed a
+built-in removes the built-in too.
+
+(extension_bundles_collisions)=
+
+### Two bundles claiming one name
+
+Within a single call, two extensions declaring a function of the same kind
+under the same name is a `ValueError` naming both. Codec ids dispatch on
+decode, so a chain can hold many and pick the right one; a function registry
+has no such fall-through, and the second registration would silently replace
+the first. Names are compared per kind, so a scalar function and an aggregate
+may share one.
+
+Shadowing a name the session *already* has is allowed and is not a collision.
+The registry holds every DataFusion built-in, and overriding built-ins by name
+is a supported thing to do — `enable_spark_functions` is built on 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,15 @@ 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, so the
+      host can validate it before anything is written and install it after
+      every codec is in place.
+      → {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,12 +26,12 @@ 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` |
+| 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.
 
@@ -67,6 +67,22 @@ from datafusion import udf
 ctx.register_udf(udf(my_library.MyScalarUDF()))
 ```
 
+If your library ships more than a function or two, declare them on a bundle
+instead and let one call install everything:
+
+```python
+class MyLibraryExtension:
+    def __datafusion_session_components__(self, ctx):
+        return SessionExtensionComponents(udfs=(my_library.MyScalarUDF(),))
+
+
+ctx = SessionContext().with_extensions(MyLibraryExtension())
+```
+
+Either the raw exportable or an already-wrapped
+{py:class}`~datafusion.user_defined.ScalarUDF` is accepted; the name comes off
+the capsule either way. See {ref}`extension_bundles`.
+
 ## Table functions
 
 A table function takes literal `Expr` arguments and returns a table provider,

docs/source/user-guide/extensions.md

@@ -36,9 +36,8 @@ this repository under
 
 Which one you have determines how much setup you do.
 
-**Tables and functions register directly.** If the library gives you a table
-or a function, register it the same way you would register a CSV file. No
-extra setup:
+**Tables register directly.** If the library gives you a table, register it
+the same way you would register a CSV file. No extra setup:
 
 ```python
 from datafusion import SessionContext
@@ -68,6 +67,13 @@ ctx.sql("SELECT count(*) FROM events").show()
 everything else with the context you called it on, so tables you registered
 before the call are still there.
 
+**Functions can arrive either way.** A single function is registered directly
+with {py:func}`~datafusion.udf` and
+{py:meth}`~datafusion.SessionContext.register_udf`. A library shipping a set of
+them usually packages them in the same `Extension` object instead, so
+`with_extensions` installs them along with everything else it provides. Follow
+whichever the library documents.
+
 ## Using more than one library
 
 Pass them all to a single call:

examples/datafusion-ffi-example/README.md

@@ -29,6 +29,10 @@ The example intentionally uses separate `cdylib` crates for these roles:
 
 Separate shared libraries guarantee distinct DataFusion library markers. This catches type-identity mistakes that a planner and provider compiled into one shared library would hide.
 
+## Installing the functions as a bundle
+
+`MyFunctionExtension` implements `__datafusion_session_components__` and declares this crate's scalar, aggregate, and window functions, so a caller installs all three with one `SessionContext.with_extensions(MyFunctionExtension())` rather than wrapping and registering each in turn. It contributes no codecs and no planner, which is the shape a function-only library takes. `python/tests/_test_session_extension.py` covers it, including that a failure after the hook registers nothing.
+
 ## Codec behavior
 
 `MyLogicalExtensionCodec` serializes this example's in-memory table providers, and `MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them. Both use documented, process-local, one-shot token registries. The registries make ownership and callback routing visible without pretending to be a portable format. They assume trusted in-process payloads and consume each token during decoding. A production provider should instead encode durable metadata from which its provider and plans can be reconstructed.

examples/datafusion-ffi-example/python/tests/_test_session_extension.py

@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""What a function library gets for shipping a bundle instead of a recipe."""
+
+from __future__ import annotations
+
+import pyarrow as pa
+import pytest
+from datafusion import SessionContext, SessionExtensionComponents
+from datafusion_ffi_example import MyFunctionExtension
+
+
+def _session():
+    """A session with the library installed in one call.
+
+    The comparison this file exists to make: without a bundle this is three
+    ``register_*`` calls the caller has to know about, one per function.
+    """
+    ctx = SessionContext().with_extensions(MyFunctionExtension())
+    batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3, None])], names=["a"])
+    ctx.register_record_batches("test_table", [[batch]])
+    return ctx
+
+
+def test_one_call_installs_every_function():
+    """All three kinds arrive across the FFI boundary from one hook."""
+    ctx = _session()
+
+    scalar = ctx.sql("select my_custom_is_null(a) from test_table").collect()
+    assert [r.column(0) for r in scalar] == [
+        pa.array([False, False, False, True], type=pa.bool_())
+    ]
+
+    aggregate = ctx.sql("select my_custom_sum(a) from test_table").collect()
+    assert aggregate[0].column(0)[0].as_py() == 6
+
+    window = ctx.sql(
+        "select my_custom_rank() over (order by a) from test_table"
+    ).collect()
+    assert window[0].num_rows == 4
+
+
+def test_the_names_come_from_the_capsules():
+    """Not from anything the bundle or the host said.
+
+    The wrappers are built by the host during resolution, so a name it invented
+    would be the one a query had to use. These are the names the Rust
+    ``ScalarUDFImpl`` and friends report.
+    """
+    ctx = _session()
+
+    assert ctx.udf("my_custom_is_null").name == "my_custom_is_null"
+    assert ctx.udaf("my_custom_sum").name == "my_custom_sum"
+    assert ctx.udwf("my_custom_rank").name == "my_custom_rank"
+
+
+def test_the_bundle_is_reusable_across_sessions():
+    """One bundle object, two sessions: components are built per install."""
+    extension = MyFunctionExtension()
+    first = SessionContext().with_extensions(extension)
+    second = SessionContext().with_extensions(extension)
+
+    assert first.udf("my_custom_is_null").name == "my_custom_is_null"
+    assert second.udf("my_custom_is_null").name == "my_custom_is_null"
+    assert first.session_id() != second.session_id()
+
+
+def test_installing_the_library_twice_is_refused():
+    """The collision rule holds for functions arriving over FFI.
+
+    Two instances of one library is the shape this actually takes in the wild —
+    an application assembling its extension list from a plugin registry that
+    lists the same package twice.
+    """
+    ctx = SessionContext()
+
+    with pytest.raises(ValueError, match=r"scalar function named 'my_custom_is_null'"):
+        ctx.with_extensions(MyFunctionExtension(), MyFunctionExtension())
+
+    with pytest.raises(KeyError):
+        ctx.udf("my_custom_is_null")
+
+
+def test_a_failure_after_the_hook_registers_nothing():
+    """The transaction covers functions imported across the FFI boundary too."""
+    ctx = SessionContext()
+
+    class BoomPlanner:
+        def __datafusion_session_planner__(self, ctx, fallback) -> None:
+            msg = "boom"
+            raise RuntimeError(msg)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        ctx.with_extensions(MyFunctionExtension(), BoomPlanner())
+
+    with pytest.raises(KeyError):
+        ctx.udf("my_custom_is_null")
+
+
+def test_the_hook_returns_the_components_type():
+    """The bundle builds a real dataclass, not a duck-typed stand-in.
+
+    ``with_extensions`` rejects anything else, so a Rust bundle that imported
+    the wrong name would fail at install rather than silently contribute
+    nothing.
+    """
+    components = MyFunctionExtension().__datafusion_session_components__(
+        SessionContext()
+    )
+
+    assert isinstance(components, SessionExtensionComponents)
+    assert len(components.udfs) == 1
+    assert components.logical_extension_codecs == ()

examples/datafusion-ffi-example/src/aggregate_udf.rs

@@ -40,7 +40,7 @@ pub(crate) struct MySumUDF {
 #[pymethods]
 impl MySumUDF {
     #[new]
-    fn new() -> PyResult<Self> {
+    pub(crate) fn new() -> PyResult<Self> {
         Ok(Self {
             inner: Arc::new(Sum::new()),
         })