Skip to content

Let extension bundles declare scalar, aggregate, and window functions - #1738

Open
timsaucer wants to merge 13 commits into
mainfrom
feat/bundle-functions
Open

timsaucer wants to merge 13 commits into
mainfrom
feat/bundle-functions

Conversation

@timsaucer

@timsaucer timsaucer commented Sep 15, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part 1 of 4 toward #1676.

  1. This PR -> scalar, aggregate, and window functions
  2. physical optimizer rules
  3. table and table functions
  4. catalog providers

Rationale for this change

#1679 added SessionContext.with_extensions so a library could ship its codecs and planner as one atomic bundle. This PR extends that idea so that a library can provide a set of user defined scalar, aggregate, and window functions. The allows users to add all of the functions from a library with one easy to use call.

What changes are included in this PR?

  • SessionExtensionComponents gains udfs, udafs, and udwfs.
  • Installation now separates what can fail from what cannot.
  • Name collisions within one call are a ValueError naming both extensions.
  • Improved documentation for the updates from this PR regarding what is fallible and what is not.

Are there any user-facing changes?

There are new functions now available on an extension bundle, but this is unreleased code and so not breaking and requires no upgrade guide.

…ctions

`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>
timsaucer and others added 12 commits September 15, 2026 11:57
The docs added alongside the new `udfs`/`udafs`/`udwfs` fields mixed three
readerships. Sorting them out:

The four-step Collect/Chains/Resolve/Commit list, and the rule it imposes on
whoever adds the next `SessionExtensionComponents` field, moves from the
extension guide to `contributor-guide/ffi-internals.md`, which already declares
itself the page you do *not* need to write an extension library. The extension
guide keeps only what an author acts on — declare, do not register — and links
across. The comment in `with_extensions` that states the same rule now points
at the new label rather than at the extension-facing one.

`Two bundles claiming one name` becomes a `##` and moves ahead of `Failure and
rollback`, which it had been splitting: the paragraphs closing that section
were rendering under the collision heading.

The user guide gains the collision as its own entry under what will bite you,
with the error text and the two-sessions workaround, since it is raised by a
call the user makes and cannot fix in their own code. Its section heading no
longer says "two kinds" over three, the functions paragraph moves above the
note that closes the section, and the discovery section covers `udfs()` and
friends rather than codec ids alone.

The bundle snippet in `functions.md` names `MyFunctionExtension` and this
crate's real function names, separates the author's class from the caller's
line, and says the cdylib can export the getter directly — it had implied a
Rust library ships a Python shim. The table's `—` for table functions now says
what to do instead.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Follow-up to the audience pass, all prose.

The paragraph after the two-hook example referred to the hooks as "the first"
and "the second", making the reader count back to the code block, and gave its
three cases three different shapes. It now names the hooks and groups by where
a component goes: codecs and functions in one hook, the planner in the other.

"Declare functions rather than registering them yourself inside the hook" named
no call, so the practice it warns against was never shown. It now says
`register_udf` on the `ctx` you were handed, and contrasts when each is
written. Its reason is also corrected: a registration made in the hook was said
to be "too early to see the other bundles' codecs", which is true of a table
provider and false of a function — the three function getters take no argument
at all, so nothing binds them to a codec chain. For functions the reason is the
transaction alone, which is what it now says.

Table providers are dropped from that paragraph rather than given the
codec-visibility caveat they deserve, since the follow-on PR covers them.

"Two bundles claiming one name" led with a gerund subject, carried its
rationale on a semicolon, and split two qualifications of equal weight across a
trailing clause and a paragraph. It now states the rule, shows the error,
explains the codec contrast on its own, and lists the two exceptions. One
example name runs through all three.

The error text quoted in both guides gains the `named` that the message
actually contains.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The collision check kept only the name, so a bundle declaring two
functions under one name reported "Two extensions declare ..." with the
same object printed on both sides, and advised installing them on
separate sessions — a remedy for a clash that was not happening.

Split the message on whether the first claimant is the same object.
Renaming is only available to a bundle colliding with itself, so it
moves to that branch; the cross-extension message loses it, matching
the user guide, which already says a caller cannot rename another
library's function.

Identity rather than equality: two objects in the argument list are two
installs even when the bundle is a dataclass that compares equal to its
twin.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`_collect_contributions` named the three function fields in a dict
literal and `with_extensions` called `_resolve_declared_functions` three
times with five positional arguments each, so a fourth kind meant
editing three places. `SessionExtensionComponents` had already moved the
other way: `__post_init__` is driven off field metadata precisely so a
field added later needs nobody to remember it.

Collapse the three call sites into `_FUNCTION_KINDS`, one row per field.
The rows name the `datafusion.user_defined` objects rather than holding
them, since that module imports this one; the lookups happen in
`with_extensions`, which also resolves the bound `register_*` method so
the infallible commit is nothing but calls.

The metadata and the table still answer different questions -- which
fields are collections to normalize, and which of those the installer
knows how to install -- so a field with metadata and no installer would
be accepted from a bundle and dropped in silence. Nothing observable
distinguishes that from a bundle declaring nothing, so
`test_every_component_field_has_an_installer` compares the two sets
directly.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`ffi-internals.md` records that functions register after the planner
hooks run, and `bundles.md` records the guarantee that ordering buys.
Neither said the consequence: `ctx.udfs()` inside
`__datafusion_session_planner__` does not list a function declared in
the same call, so a bundle resolving one at hook time gets a KeyError
and no clue why.

Also drop two stray blank lines left in the user guide by the collision
section.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The new bundle examples spelled volatility positionally while the
neighbouring `context.py` examples name it, which leaves a bare
"stable" sitting among three pyarrow arguments with nothing to say
what it is.

Two docstrings alongside: `_components` pointed a `:py:meth:` role at
`__post_init__`, which Sphinx has no target for and never renders
anyway from a private helper; and `_resolve_declared_functions` said
its `wrapper` argument was "passed through already", which says
nothing. It is the class a declaration may already be an instance of.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`_resolve_declared_functions` chose between its two collision messages
on object identity, so one bundle object passed twice -- an extension
list assembled from a plugin registry that names the same package
twice, which is the shape the FFI test already calls out -- read as a
bundle colliding with itself and was told to rename one of the two.
There is nothing to rename. Both claims come from the one declaration,
and the caller cannot rename another library's function anyway.

What the caller controls is the argument list, so key on position in
it. Two entries are two installs whether or not they are the same
object, and a repeat now gets the message aimed at a caller: install
them on separate sessions, or drop the duplicate. Only a single
argument declaring one name twice keeps the rename advice, which is
the one case where its author can act on it. `_collect_contributions`
carries the position alongside the extension to make the distinction
available, and the message names both positions so neither side has to
be guessed from two identical reprs.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`_FUNCTION_KINDS` holds the `user_defined` wrapper, factory, and
`register_*` method as strings, looked up during `with_extensions` so
the import stays out of this module's cycle. The cost is that a typo
in a row surfaces as an `AttributeError` part-way through an install
rather than at import, and `test_every_component_field_has_an_installer`
does not catch it -- that one compares field names, which a bad
`wrapper` or `factory` leaves untouched.

The three rows that exist today are each covered end to end by a
registration test, so this is for the fourth, which may well be added
before its own behaviour test is.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The rule that everything after the planner install must be infallible
was tacked onto the end of a comment arguing something else -- why the
rebind is guarded on a call that installs nothing. Two unrelated
arguments in one block, with the more important of the two reading as
a footnote to the other. Give it its own block.

Its pointer at the reasoning was a `:ref:` role, which renders nowhere
from a `#` comment and leaves a reader who follows it holding a label
with no way to resolve it. Name the file and the heading instead.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every other module in this crate names the pyo3 items it uses;
`extension.rs` arrived with a glob off the prelude. The crate's
rustfmt config asks for `imports_granularity = Module`, which would
have made the difference visible, but it is a nightly-only option and
CI checks formatting on stable, so nothing was going to flag it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The unusable-declaration error rendered "must be a AggregateUDF" for
two of the three kinds; dropping the article reads correctly for all
of them.

The collision branch carried a seven-line argument for keying on
position rather than object identity, duplicating what the extension
guide already argues. State the constraint and point at
`extension_bundles_collisions` instead.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The collision tests exercised only the scalar row of
`_FUNCTION_KINDS`, so a transposed label or field on the aggregate or
window rows would have passed. Parametrize the two-extension case over
all three kinds, with `_total`/`_first` factories the registration
test now shares.

The duplicate-extension test asserted only the message shape; add the
check that nothing reached the session, making it self-contained
rather than leaning on its siblings.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@timsaucer
timsaucer marked this pull request as ready for review September 15, 2026 20:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant