Let extension bundles declare scalar, aggregate, and window functions

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

@@ -111,6 +111,43 @@ 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 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 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,24 @@ 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. Both put the function on the session, but a registration you make
+inside the hook is written the moment it runs — before the other bundles have
+been called, and not undone if one of them raises. What you declare is instead
+resolved and checked while a failure still costs nothing, then written once
+every bundle has succeeded. 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 +290,69 @@ 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
+
+Two extensions in one call may not declare a function of the same kind under
+the same name. Doing so raises:
+
+```text
+ValueError: Two extensions declare a scalar function named 'normalize':
+argument 0 (...) and argument 1 (...). ...
+```
+
+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.
+
+Which argument each claim came from is part of the message because it is what
+picks the remedy. Two arguments colliding is the caller's to resolve, by
+installing the two on separate sessions or by dropping a repeat; renaming is
+not something a caller can do. One argument declaring a name twice is the
+bundle author's own bug, and gets a different message saying so. Collisions are
+keyed on position rather than on object identity, so passing one extension
+twice reads as the caller's duplicate that it is, rather than as a bundle
+colliding with itself.
+
+Two 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.
+
+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,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,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,

docs/source/user-guide/extensions.md

@@ -32,13 +32,12 @@ which exposes Delta Lake tables to DataFusion, and the two worked examples in
 this repository under
 [`examples/`](https://github.com/apache/datafusion-python/tree/main/examples).
 
-## Two kinds of extension
+## How an extension reaches your session
 
-Which one you have determines how much setup you do.
+Which route your library takes 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
@@ -64,6 +63,21 @@ ctx.register_table("events", my_engine.TableProvider("s3://bucket/events"))
 ctx.sql("SELECT count(*) FROM events").show()
 ```
 
+**Functions arrive by whichever route their library chose.** A library
+offering one or two functions hands you the functions themselves, and you wrap
+and register each:
+
+```python
+from datafusion import udf
+
+ctx.register_udf(udf(my_library.MyScalarUDF()))
+```
+
+A library shipping a set of them packages them in its `Extension` object
+instead, so `with_extensions` installs them all along with everything else it
+provides, and there is nothing per-function for you to do. Its documentation
+says which.
+
 `with_extensions` returns a context; use the returned one. It shares
 everything else with the context you called it on, so tables you registered
 before the call are still there.
@@ -85,7 +99,26 @@ rarely matters. When a library needs a particular position — usually "list me
 last" for something that wraps the others — it says so in its own
 documentation.
 
-## Two things that will bite you
+## Three things that will bite you
+
+**Two libraries can claim one function name.** If both ship a function of the
+same kind under the same name, the call raises a `ValueError` naming both,
+rather than letting one silently replace the other:
+
+```text
+ValueError: Two extensions declare a scalar function named 'normalize': ...
+```
+
+You cannot rename another library's function from your own code, so the fix is
+to use two sessions, one per library, and query each for what only it provides.
+Worth reporting upstream too: the library whose names are the less specific
+should be prefixing them. A function shadowing a *built-in* is not a collision
+and raises nothing — that is a supported thing for a library to do. See
+{ref}`extension_bundles_collisions`.
+
+Check the argument positions the message names before you go looking for a
+second library. Passing one extension twice collides with itself, and an
+extension list assembled from a plugin registry is the usual way that happens.
 
 **Keep your context alive.** A `DataFrame` or a plan does not keep its session
 alive on its own. If a context is garbage-collected while something built from
@@ -133,6 +166,17 @@ ctx.logical_extension_codec_ids()
 
 An empty list means nothing extra is installed.
 
+For functions, {py:meth}`~datafusion.SessionContext.udfs`,
+{py:meth}`~datafusion.SessionContext.udafs` and
+{py:meth}`~datafusion.SessionContext.udwfs` return the names a session knows.
+Both the library's and every DataFusion built-in are in there, so look for the
+name rather than reading the whole list:
+
+```python
+"my_engine_normalize" in ctx.udfs()
+# True
+```
+
 ## Next steps
 
 - {ref}`distributed_query_engines` — running your queries across several

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.