From 0af2a50e886f95e6820b3f533cb26c2fc2609a13 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 09:50:30 -0400 Subject: [PATCH 01/14] 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) --- docs/source/extension-guide/bundles.md | 64 ++++- docs/source/extension-guide/checklist.md | 10 +- docs/source/extension-guide/functions.md | 28 ++- docs/source/user-guide/extensions.md | 12 +- examples/datafusion-ffi-example/README.md | 4 + .../python/tests/_test_session_extension.py | 128 ++++++++++ .../src/aggregate_udf.rs | 2 +- .../datafusion-ffi-example/src/extension.rs | 64 +++++ examples/datafusion-ffi-example/src/lib.rs | 3 + .../datafusion-ffi-example/src/scalar_udf.rs | 2 +- .../datafusion-ffi-example/src/window_udf.rs | 2 +- python/datafusion/context.py | 218 +++++++++++++++--- python/datafusion/extensions.py | 125 ++++++++-- python/tests/test_context.py | 181 +++++++++++++++ 14 files changed, 758 insertions(+), 85 deletions(-) create mode 100644 examples/datafusion-ffi-example/python/tests/_test_session_extension.py create mode 100644 examples/datafusion-ffi-example/src/extension.rs diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 791fdfc8e..afeaeb4cb 100644 --- a/docs/source/extension-guide/bundles.md +++ b/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 diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md index 6e204b3f5..f5da351a1 100644 --- a/docs/source/extension-guide/checklist.md +++ b/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` diff --git a/docs/source/extension-guide/functions.md b/docs/source/extension-guide/functions.md index fcfef9f1a..50f9f2ded 100644 --- a/docs/source/extension-guide/functions.md +++ b/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, diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index 8c86b4a67..00ad5b46d 100644 --- a/docs/source/user-guide/extensions.md +++ b/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: diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index aadea909f..b96c798a6 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/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. diff --git a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py new file mode 100644 index 000000000..fadc1fd70 --- /dev/null +++ b/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 == () diff --git a/examples/datafusion-ffi-example/src/aggregate_udf.rs b/examples/datafusion-ffi-example/src/aggregate_udf.rs index ea1518365..b0c7790b3 100644 --- a/examples/datafusion-ffi-example/src/aggregate_udf.rs +++ b/examples/datafusion-ffi-example/src/aggregate_udf.rs @@ -40,7 +40,7 @@ pub(crate) struct MySumUDF { #[pymethods] impl MySumUDF { #[new] - fn new() -> PyResult { + pub(crate) fn new() -> PyResult { Ok(Self { inner: Arc::new(Sum::new()), }) diff --git a/examples/datafusion-ffi-example/src/extension.rs b/examples/datafusion-ffi-example/src/extension.rs new file mode 100644 index 000000000..777db0cb5 --- /dev/null +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -0,0 +1,64 @@ +// 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. + +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use crate::aggregate_udf::MySumUDF; +use crate::scalar_udf::IsNullUDF; +use crate::window_udf::MyRankUDF; + +/// A bundle contributing this library's three functions in one install. +/// +/// The shape a function library takes: no codecs and no planner, so the whole +/// of its installation is what it declares here. The three function getters +/// take no argument, so unlike a provider these need nothing from `ctx` and +/// the objects are handed over unresolved for the host to wrap. +#[pyclass( + from_py_object, + name = "MyFunctionExtension", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone, Default)] +pub(crate) struct MyFunctionExtension {} + +#[pymethods] +impl MyFunctionExtension { + #[new] + fn new() -> Self { + Self {} + } + + /// `ctx` is unused: nothing declared here is bound to a session. + fn __datafusion_session_components__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + let _ = ctx; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("udfs", (Py::new(py, IsNullUDF::new())?,))?; + kwargs.set_item("udafs", (Py::new(py, MySumUDF::new()?)?,))?; + kwargs.set_item("udwfs", (Py::new(py, MyRankUDF::new()?)?,))?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 92fccb1e2..b680c84de 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -20,6 +20,7 @@ use pyo3::prelude::*; use crate::aggregate_udf::MySumUDF; use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList}; use crate::config::MyConfig; +use crate::extension::MyFunctionExtension; use crate::logical_extension_codec::MyLogicalExtensionCodec; use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec}; use crate::physical_extension_codec::MyPhysicalExtensionCodec; @@ -33,6 +34,7 @@ use crate::window_udf::MyRankUDF; pub(crate) mod aggregate_udf; pub(crate) mod catalog_provider; pub(crate) mod config; +pub(crate) mod extension; pub(crate) mod logical_extension_codec; pub(crate) mod name_only_codec; pub(crate) mod physical_extension_codec; @@ -63,5 +65,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/scalar_udf.rs b/examples/datafusion-ffi-example/src/scalar_udf.rs index 85b884ec6..7c25d4ade 100644 --- a/examples/datafusion-ffi-example/src/scalar_udf.rs +++ b/examples/datafusion-ffi-example/src/scalar_udf.rs @@ -43,7 +43,7 @@ pub(crate) struct IsNullUDF { #[pymethods] impl IsNullUDF { #[new] - fn new() -> Self { + pub(crate) fn new() -> Self { Self { signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable), } diff --git a/examples/datafusion-ffi-example/src/window_udf.rs b/examples/datafusion-ffi-example/src/window_udf.rs index 2956ad64c..e15ddc3a0 100644 --- a/examples/datafusion-ffi-example/src/window_udf.rs +++ b/examples/datafusion-ffi-example/src/window_udf.rs @@ -40,7 +40,7 @@ pub(crate) struct MyRankUDF { #[pymethods] impl MyRankUDF { #[new] - fn new() -> PyResult { + pub(crate) fn new() -> PyResult { Ok(Self { inner: rank_udwf() }) } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 63cdfd487..1bee0804f 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -161,6 +161,122 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 +def _collect_contributions( + extensions: tuple[object, ...], + ctx: SessionContext, +) -> tuple[list[Any], list[Any], dict[str, list[tuple[object, Any]]]]: + """Run every components hook and gather what the extensions contribute. + + Validates the whole argument list before calling anything, so an argument + that implements neither hook is refused before a well-formed extension + ahead of it has done any work. Writes nothing to the session. + + Functions are kept paired with the extension that declared them, so a name + claimed twice can name both sides. + + Args: + extensions: The arguments ``with_extensions`` was given. + ctx: The context the components are bound against — the receiver, not a + context derived from it. + + Returns: + The logical codecs, the physical codecs, and the declared functions by + field name. + + Raises: + TypeError: If an argument implements neither hook, or a hook returns + something other than ``SessionExtensionComponents``. + """ + for extension in extensions: + if not isinstance( + extension, (SessionComponentsExportable, SessionPlannerExportable) + ): + msg = ( + "Extension implements neither " + "__datafusion_session_components__ nor " + f"__datafusion_session_planner__: {extension!r}" + ) + raise TypeError(msg) + + logical_codecs: list[LogicalExtensionCodecExportable] = [] + physical_codecs: list[PhysicalExtensionCodecExportable] = [] + declared: dict[str, list[tuple[object, Any]]] = { + "udfs": [], + "udafs": [], + "udwfs": [], + } + for extension in extensions: + if not isinstance(extension, SessionComponentsExportable): + continue + components = extension.__datafusion_session_components__(ctx) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_components__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + logical_codecs.extend(components.logical_extension_codecs) + physical_codecs.extend(components.physical_extension_codecs) + for field_name, declarations in declared.items(): + declarations.extend( + (extension, item) for item in getattr(components, field_name) + ) + return logical_codecs, physical_codecs, declared + + +def _resolve_declared_functions( + declared: list[tuple[object, Any]], + wrapper: type, + getter: str, + factory: Any, + kind: str, +) -> list[Any]: + """Wrap every function an extension declared, refusing a name claimed twice. + + Runs before anything is written to the session, so a bundle that hands over + something unusable — or two bundles that both claim a name — fail with + nothing registered. The name is read off the wrapper rather than off the + declaration: on the FFI path the capsule reports it, and the constructor + argument is ignored. + + Args: + declared: ``(extension, function)`` pairs in declaration order. + wrapper: The Python wrapper class for this kind, passed through already. + getter: The capsule getter an unwrapped declaration must expose. + factory: The helper that wraps a declaration — ``udf`` and friends. + kind: What to call this sort of function in an error. + + Returns: + The wrappers to register, in declaration order. + """ + resolved = [] + claimed: dict[str, object] = {} + for extension, function in declared: + if isinstance(function, wrapper): + wrapped = function + elif hasattr(function, getter): + wrapped = factory(function) + else: + msg = ( + f"A declared {kind} must be a {wrapper.__name__} or expose " + f"{getter}, got {function!r} from {extension!r}" + ) + raise TypeError(msg) + name = wrapped.name + if name in claimed: + msg = ( + f"Two extensions declare a {kind} named {name!r}: " + f"{claimed[name]!r} and {extension!r}. Registrations have no " + "fall-through, so one would silently replace the other; " + "install them on separate sessions, or rename one." + ) + raise ValueError(msg) + claimed[name] = extension + resolved.append(wrapped) + return resolved + + class SessionConfig: """Session configuration options.""" @@ -1915,10 +2031,12 @@ def with_extensions( :py:class:`~datafusion.extensions.SessionPlannerExportable`. Nothing is written to the session until every hook has returned and - every capsule has been validated, so a hook that raises leaves the - session as it was. A hook that *mutates* the context it is handed — - registering a table, say — is not rolled back, which is why bundle - objects must be configuration-only. + every component has been validated, so a hook that raises leaves the + session as it was. Declared functions register after the planner is + bound, and are visible on every handle sharing this session. A hook + that *mutates* the context it is handed — registering a table, say — is + not rolled back, which is why bundle objects must be + configuration-only. Shares its session with this context — see :py:class:`SessionContext`. @@ -1938,10 +2056,13 @@ def with_extensions( Raises: TypeError: If an argument implements neither hook, if a hook - returns the wrong type, or if a codec is contributed as a bare - ``PyCapsule`` rather than an object exposing the getter. - ValueError: If two codecs claim the same id, or a getter returns a - capsule of the wrong kind. See + returns the wrong type, if a codec is contributed as a bare + ``PyCapsule`` rather than an object exposing the getter, or if + a declared function is neither a wrapper nor exposes its + capsule getter. + ValueError: If two codecs claim the same id, if two extensions + declare a function of one kind under the same name, or if a + getter returns a capsule of the wrong kind. See :py:meth:`with_logical_extension_codec` for how ids are assigned. @@ -1975,42 +2096,58 @@ def with_extensions( >>> batches[0].column(0).to_pylist() # doctest: +SKIP [1] """ - for extension in extensions: - if not isinstance( - extension, (SessionComponentsExportable, SessionPlannerExportable) - ): - msg = ( - "Extension implements neither " - "__datafusion_session_components__ nor " - f"__datafusion_session_planner__: {extension!r}" - ) - raise TypeError(msg) - - # Phase one: collect every bundle's codecs. Components are bound + # Phase one: collect every bundle's components. Components are bound # against this context, not a context derived from it. There is one # `Arc` per session, so a component bound here holds a # task-context provider that the returned handle keeps alive. - logical_codecs: list[LogicalExtensionCodecExportable] = [] - physical_codecs: list[PhysicalExtensionCodecExportable] = [] - for extension in extensions: - if not isinstance(extension, SessionComponentsExportable): - continue - components = extension.__datafusion_session_components__(self) - if not isinstance(components, SessionExtensionComponents): - msg = ( - "__datafusion_session_components__ must return " - "SessionExtensionComponents, got " - f"{type(components).__name__} from {extension!r}" - ) - raise TypeError(msg) - logical_codecs.extend(components.logical_extension_codecs) - physical_codecs.extend(components.physical_extension_codecs) + logical_codecs, physical_codecs, declared = _collect_contributions( + extensions, self + ) # Writes nothing: the chains belong to the new handle, so a failure # above or below leaves this context as it was. new = SessionContext.__new__(SessionContext) new.ctx = self.ctx._install_extension_codecs(logical_codecs, physical_codecs) + # Resolve every declared function to the wrapper that registers it, and + # settle name collisions, while a failure still costs nothing. None of + # these getters take an argument, so unlike a provider they do not care + # which handle they are resolved against. + from datafusion.user_defined import ( # noqa: PLC0415 + AggregateUDF as _AggregateUDF, + ) + from datafusion.user_defined import ( # noqa: PLC0415 + ScalarUDF as _ScalarUDF, + ) + from datafusion.user_defined import ( # noqa: PLC0415 + WindowUDF as _WindowUDF, + ) + from datafusion.user_defined import udaf as _udaf # noqa: PLC0415 + from datafusion.user_defined import udf as _udf # noqa: PLC0415 + from datafusion.user_defined import udwf as _udwf # noqa: PLC0415 + + resolved_udfs = _resolve_declared_functions( + declared["udfs"], + _ScalarUDF, + "__datafusion_scalar_udf__", + _udf, + "scalar function", + ) + resolved_udafs = _resolve_declared_functions( + declared["udafs"], + _AggregateUDF, + "__datafusion_aggregate_udf__", + _udaf, + "aggregate function", + ) + resolved_udwfs = _resolve_declared_functions( + declared["udwfs"], + _WindowUDF, + "__datafusion_window_udf__", + _udwf, + "window function", + ) + # Phase two: nest the planners, outermost last. Each hook runs against # `new`, which carries the final chains, so a planner captured here # never sees a partial codec set. `planner` stays None when no bundle @@ -2037,8 +2174,19 @@ def with_extensions( # it already holds, so the rebuild is unobservable except in the one case # where it does harm: a planner sitting on some *other* handle's codecs # gets dragged onto this handle's, silently undoing that install. + # Everything below this line must be infallible. A registration whose + # commit can fail belongs above, split into an import step that returns + # a resolved object and an insert step that cannot raise -- there is one + # session here, shared with the receiver, so a failure part-way through + # has nothing to roll back to. See :ref:`extension_bundles_transaction`. if planner is not None or logical_codecs or physical_codecs: new.ctx._install_extension_planner(planner) + for function in resolved_udfs: + new.register_udf(function) + for function in resolved_udafs: + new.register_udaf(function) + for function in resolved_udwfs: + new.register_udwf(function) return new def table_provider(self, name: str) -> Table: diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 77ae92fc2..5988edc3b 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -44,7 +44,7 @@ from __future__ import annotations -from dataclasses import dataclass, fields +from dataclasses import dataclass, field, fields from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -52,8 +52,14 @@ from datafusion.context import SessionContext from datafusion.user_defined import ( + AggregateUDF, + AggregateUDFExportable, LogicalExtensionCodecExportable, PhysicalExtensionCodecExportable, + ScalarUDF, + ScalarUDFExportable, + WindowUDF, + WindowUDFExportable, ) __all__ = [ @@ -101,15 +107,26 @@ class QueryPlannerExportable(Protocol): def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 -def _not_a_codec_iterable(field: str, value: object) -> str: - """Message for a codec field that cannot be read as a collection.""" +def _not_an_iterable(name: str, value: object, noun: str) -> str: + """Message for a component field that cannot be read as a collection.""" return ( - f"{field} must be an iterable of codec objects, not a single " - f"{type(value).__name__}. A lone codec is written as a one-element " - f"tuple — {field}=(codec,) — and the trailing comma is what makes it one." + f"{name} must be an iterable of {noun} objects, not a single " + f"{type(value).__name__}. A lone {noun} is written as a one-element " + f"tuple — {name}=({noun},) — and the trailing comma is what makes it one." ) +def _components(noun: str) -> Any: + """Declare a field holding a tuple of contributed components. + + ``noun`` names what the field holds, for the error a bundle sees when it + hands over one component instead of a collection of them. Carrying it in + the field metadata is what lets :py:meth:`SessionExtensionComponents.__post_init__` + normalize a field it was never told about by name. + """ + return field(default=(), metadata={"datafusion_component": noun}) + + @dataclass(frozen=True) class SessionExtensionComponents: """Components an extension contributes to a session context. @@ -124,6 +141,12 @@ class SessionExtensionComponents: Query planners are not listed here. They install in a second phase so each can wrap the one before it — see :py:class:`SessionPlannerExportable`. + Codecs are held by the returned handle; everything else is registered on + the session the handle shares. Declaring a component is not the same as + registering it yourself during the hook: declared components are resolved + before anything is written, so a bundle that fails leaves nothing behind. + See :ref:`extension_bundles_transaction`. + Examples: A bundle that contributes no codecs is valid — a planner-only library returns this, or omits the hook entirely: @@ -158,7 +181,23 @@ class SessionExtensionComponents: >>> components.physical_extension_codecs () - A single codec is not an iterable of codecs, and forgetting the + Functions are declared the same way, and register under the name the + function itself reports: + + >>> import pyarrow as pa + >>> from datafusion import udf + >>> double = udf( + ... lambda arr: pa.array([v.as_py() * 2 for v in arr]), + ... [pa.int64()], + ... pa.int64(), + ... "stable", + ... name="double", + ... ) + >>> components = SessionExtensionComponents(udfs=(double,)) + >>> [fn.name for fn in components.udfs] + ['double'] + + A single component is not an iterable of them, and forgetting the trailing comma is the easy way to write one by accident: >>> SessionExtensionComponents(logical_extension_codecs=NamedCodec(ctx)) @@ -167,7 +206,9 @@ class SessionExtensionComponents: TypeError: logical_extension_codecs must be an iterable of codec objects... """ - logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = () + logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = _components( + "codec" + ) """Logical codecs to add to the session's codec chain, in declaration order. Objects exposing ``__datafusion_logical_extension_codec__``, never bare @@ -176,15 +217,46 @@ class SessionExtensionComponents: :ref:`extension_bundles_codecs_are_objects`. """ - physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () + physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = ( + _components("codec") + ) """Physical codecs to add to the session's codec chain, in declaration order. As :py:attr:`logical_extension_codecs`, for ``__datafusion_physical_extension_codec__``. """ + udfs: tuple[ScalarUDF | ScalarUDFExportable, ...] = _components("function") + """Scalar functions to register on the session. + + Either a :py:class:`~datafusion.user_defined.ScalarUDF` or an object + exposing ``__datafusion_scalar_udf__``, which is wrapped with + :py:func:`~datafusion.udf` on the way in. The registered name comes from + the function itself, not from this field. + + Two extensions in one + :py:meth:`~datafusion.context.SessionContext.with_extensions` call may not + declare the same name; shadowing a function the session already has is + allowed. See :ref:`extension_bundles_collisions`. + """ + + udafs: tuple[AggregateUDF | AggregateUDFExportable, ...] = _components("function") + """Aggregate functions to register on the session. + + As :py:attr:`udfs`, for ``__datafusion_aggregate_udf__`` and + :py:func:`~datafusion.udaf`. Names are compared within their own kind, so + an aggregate may share a name with a scalar function. + """ + + udwfs: tuple[WindowUDF | WindowUDFExportable, ...] = _components("function") + """Window functions to register on the session. + + As :py:attr:`udfs`, for ``__datafusion_window_udf__`` and + :py:func:`~datafusion.udwf`. + """ + def __post_init__(self) -> None: - """Normalize each codec field to a tuple, rejecting what cannot become one.""" + """Normalize each component field, rejecting what cannot become a tuple.""" # A bundle that writes `logical_extension_codecs=codec` instead of # `(codec,)` is contributing one codec, not an iterable of them. # Without this, the mistake surfaces inside `with_extensions` as @@ -198,24 +270,25 @@ def __post_init__(self) -> None: # the first read. # # Driven off `dataclasses.fields` rather than a written-out list, so a - # codec field added later is normalized without anyone remembering to - # name it here. The `_codecs` suffix is what marks a field as one of - # them, leaving room for a future field that is not a codec collection - # and must not be turned into a tuple. - for field in fields(self): - name = field.name - if not name.endswith("_codecs"): + # component field added later is normalized without anyone remembering + # to name it here. `_components` metadata is what marks a field as one + # of them, leaving room for a future field that is not a collection and + # must not be turned into a tuple. + for spec in fields(self): + noun = spec.metadata.get("datafusion_component") + if noun is None: continue + name = spec.name value = getattr(self, name) # A str is iterable, so it would otherwise normalize into a tuple - # of characters and fail much later as that many bogus codecs. + # of characters and fail much later as that many bogus components. if isinstance(value, (str, bytes)): - raise TypeError(_not_a_codec_iterable(name, value)) + raise TypeError(_not_an_iterable(name, value, noun)) try: - codecs = tuple(value) + components = tuple(value) except TypeError: - raise TypeError(_not_a_codec_iterable(name, value)) from None - object.__setattr__(self, name, codecs) + raise TypeError(_not_an_iterable(name, value, noun)) from None + object.__setattr__(self, name, components) @runtime_checkable @@ -231,9 +304,11 @@ class SessionComponentsExportable(Protocol): components on every call using the context supplied by :py:meth:`~datafusion.context.SessionContext.with_extensions`, and must not retain that context or cache the components they bound to it, since the - next call may install onto a different session. They should also avoid - mutating the context they are handed — a registration made during binding - is not rolled back if a later extension fails. + next call may install onto a different session. Declare what you contribute + rather than registering it on the context you are handed: a registration + made during the hook is not rolled back if a later extension fails, and it + binds to the codec chains from before the call. See + :ref:`extension_bundles_transaction`. A bundle that also contributes a query planner implements :py:class:`SessionPlannerExportable` alongside this protocol. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 9fbc4744f..50842c690 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -22,9 +22,11 @@ import shutil import pyarrow as pa +import pyarrow.compute as pc import pyarrow.dataset as ds import pytest from datafusion import ( + Accumulator, CsvReadOptions, DataFrame, RuntimeEnvBuilder, @@ -35,8 +37,11 @@ Table, column, literal, + udaf, udf, + udwf, ) +from datafusion.user_defined import WindowEvaluator def test_create_context_no_args(): @@ -1410,6 +1415,182 @@ def __datafusion_session_components__(self, ctx): assert batches[0].column(0) == pa.array([1]) +def _doubler(name="double"): + """A scalar function under a name the caller picks.""" + return udf( + lambda arr: pa.array([v.as_py() * 2 for v in arr]), + [pa.int64()], + pa.int64(), + "stable", + name=name, + ) + + +class _Total(Accumulator): + """The smallest accumulator that survives a partial/final split.""" + + def __init__(self): + self._sum = 0 + + def state(self) -> list[pa.Scalar]: + return [pa.scalar(self._sum)] + + def update(self, values: pa.Array) -> None: + self._sum += pc.sum(values).as_py() or 0 + + def merge(self, states: list[pa.Array]) -> None: + self._sum += pc.sum(states[0]).as_py() or 0 + + def evaluate(self) -> pa.Scalar: + return pa.scalar(self._sum) + + +class _First(WindowEvaluator): + """Repeats the first value of the partition across every row.""" + + def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array: + first = values[0][0].as_py() + return pa.array([first] * num_rows) + + +class _FunctionExtension: + """Contributes functions and nothing else. + + The shape a library shipping only functions has: no codecs, no planner, + so the whole of its installation is what it declares here. + """ + + def __init__(self, udfs=(), udafs=(), udwfs=()): + self._udfs = udfs + self._udafs = udafs + self._udwfs = udwfs + + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents( + udfs=self._udfs, udafs=self._udafs, udwfs=self._udwfs + ) + + +def test_with_extensions_registers_a_declared_udf(ctx): + """A declared scalar function is callable from SQL on the returned handle.""" + result = ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),))) + result.from_pydict({"a": [1, 2, 3]}, name="nums") + + batches = result.sql("SELECT double(a) AS doubled FROM nums").collect() + assert batches[0].column(0) == pa.array([2, 4, 6]) + + +def test_with_extensions_registers_udafs_and_udwfs(ctx): + """The other two function kinds install the same way.""" + total = udaf(_Total, pa.int64(), pa.int64(), [pa.int64()], "stable", name="total") + first = udwf(_First, pa.int64(), pa.int64(), "immutable", name="first_value_of") + + result = ctx.with_extensions(_FunctionExtension(udafs=(total,), udwfs=(first,))) + result.from_pydict({"a": [1, 2, 3]}, name="nums") + + assert result.sql("SELECT total(a) FROM nums").collect()[0].column(0) == pa.array( + [6] + ) + batches = result.sql("SELECT first_value_of(a) OVER () FROM nums").collect() + assert batches[0].column(0) == pa.array([1, 1, 1]) + + +def test_with_extensions_registers_on_the_shared_session(ctx): + """Registrations land on the session, which the source context also holds. + + Only the codec chains belong to the returned handle. Pinned deliberately: + a future change that made registrations private to the handle would be a + behaviour change, not a fix. + """ + ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),))) + + assert ctx.udf("double").name == "double" + + +def test_with_extensions_rejects_a_name_two_extensions_claim(ctx): + """Registrations have no fall-through, so a clash cannot be resolved by order. + + Unlike codecs, which dispatch by id, a second function under one name would + silently replace the first. + """ + with pytest.raises(ValueError, match=r"scalar function named 'double'"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + _FunctionExtension(udfs=(_doubler(),)), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + + +def test_with_extensions_allows_shadowing_an_existing_function(ctx): + """Claiming a name the session already has is legal. + + ``ctx.udfs()`` holds every built-in, and ``enable_spark_functions`` + overrides built-ins by design, so refusing this would refuse a supported + use rather than catch a mistake. + """ + result = ctx.with_extensions(_FunctionExtension(udfs=(_doubler(name="abs"),))) + result.from_pydict({"a": [1, 2, 3]}, name="nums") + + batches = result.sql("SELECT abs(a) AS shadowed FROM nums").collect() + assert batches[0].column(0) == pa.array([2, 4, 6]) + + +def test_with_extensions_registers_nothing_when_a_components_hook_raises(ctx): + """A failure in phase one leaves the first extension's functions uninstalled.""" + + class BoomExtension: + def __datafusion_session_components__(self, ctx): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + BoomExtension(), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + + +def test_with_extensions_registers_nothing_when_a_planner_hook_raises(ctx): + """A failure in phase two does too, which is what pins the ordering. + + By the time the planner hooks run, the functions have been resolved and + their names checked. Committing them at that point rather than after would + pass every other test here and still leave this one registered. + """ + + class BoomPlanner: + def __datafusion_session_planner__(self, ctx, fallback): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + BoomPlanner(), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + + +def test_with_extensions_rejects_an_unusable_declaration(ctx): + """Something that is neither a wrapper nor an exportable names both sides.""" + with pytest.raises(TypeError, match=r"__datafusion_scalar_udf__"): + ctx.with_extensions(_FunctionExtension(udfs=(object(),))) + + +@pytest.mark.parametrize("field", ["udfs", "udafs", "udwfs"]) +def test_session_extension_components_rejects_a_single_function(field): + """A lone function is not an iterable of them, as for codecs.""" + with pytest.raises(TypeError, match=r"must be an iterable of function objects"): + SessionExtensionComponents(**{field: _doubler()}) + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) From d47055e1f205ade2081b112d43a2a1988e10a01e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 11:57:05 -0400 Subject: [PATCH 02/14] docs: put each bundle-functions claim in front of its own audience MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../source/contributor-guide/ffi-internals.md | 37 ++++++++++++ docs/source/extension-guide/bundles.md | 57 ++++++++----------- docs/source/extension-guide/functions.md | 41 +++++++++---- docs/source/user-guide/extensions.md | 56 ++++++++++++++---- python/datafusion/context.py | 2 +- 5 files changed, 139 insertions(+), 54 deletions(-) diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index 75c32444c..f25e14091 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/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 diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index afeaeb4cb..fb5bd3e89 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -289,41 +289,9 @@ 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 -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 +## 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 @@ -336,6 +304,29 @@ 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. +Since the caller cannot repair a collision from their own code, name your +functions so this does not arise: a prefix tying them to your library is the +usual answer. + +(extension_bundles_transaction)= + +## Failure and rollback + +Nothing is written to the session until every factory has returned and every +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`. + 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 codec chains belong to the returned handle; the planner is installed on the diff --git a/docs/source/extension-guide/functions.md b/docs/source/extension-guide/functions.md index 50f9f2ded..ee38090e2 100644 --- a/docs/source/extension-guide/functions.md +++ b/docs/source/extension-guide/functions.md @@ -33,7 +33,11 @@ the same registration methods. | `__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. +All four are implemented in [`datafusion-ffi-example`], one per file. The last +column is the {py:class}`~datafusion.SessionExtensionComponents` field a +{ref}`bundle ` 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,21 +71,38 @@ 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: +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 ` declaring +them, so one call installs the lot: ```python -class MyLibraryExtension: - def __datafusion_session_components__(self, ctx): - return SessionExtensionComponents(udfs=(my_library.MyScalarUDF(),)) +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 -ctx = SessionContext().with_extensions(MyLibraryExtension()) +class MyFunctionExtension: + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents( + udfs=(IsNullUDF(),), + udafs=(MySumUDF(),), + udwfs=(MyRankUDF(),), + ) ``` -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`. +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 diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index 00ad5b46d..197e660a8 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -32,9 +32,9 @@ 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 register directly.** If the library gives you a table, register it the same way you would register a CSV file. No extra setup: @@ -63,17 +63,25 @@ 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. -**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: @@ -91,7 +99,24 @@ 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 '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`. + + **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 @@ -139,6 +164,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 diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 1bee0804f..5b78ccc54 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2178,7 +2178,7 @@ def with_extensions( # commit can fail belongs above, split into an import step that returns # a resolved object and an insert step that cannot raise -- there is one # session here, shared with the receiver, so a failure part-way through - # has nothing to roll back to. See :ref:`extension_bundles_transaction`. + # has nothing to roll back to. See :ref:`ffi_internals_commit_order`. if planner is not None or logical_codecs or physical_codecs: new.ctx._install_extension_planner(planner) for function in resolved_udfs: From 7ca32c68482d8ca69a70ddaabf626ff15528c8f2 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 12:25:42 -0400 Subject: [PATCH 03/14] docs: unpack the dense passages in the bundle guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- docs/source/extension-guide/bundles.md | 53 +++++++++++++++----------- docs/source/user-guide/extensions.md | 2 +- 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index fb5bd3e89..487fdf618 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -55,22 +55,23 @@ 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, and a library that -ships only functions defines the first and leaves the codec fields empty. 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()) ``` -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`. +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 @@ -293,20 +294,28 @@ a capsule. ## 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. +Two extensions in one call may not declare a function of the same kind under +the same name. Doing so raises: -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. +```text +ValueError: Two extensions declare a scalar function named 'normalize': ... +``` + +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. + +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. -Since the caller cannot repair a collision from their own code, name your -functions so this does not arise: a prefix tying them to your library is the -usual answer. +Your caller cannot rename your function, so stay out of the way: prefix the +names with something tied to your library. (extension_bundles_transaction)= diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index 197e660a8..e41747512 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -106,7 +106,7 @@ 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 'normalize': ... +ValueError: Two extensions declare a scalar function named 'normalize': ... ``` You cannot rename another library's function from your own code, so the fix is From 30d1c7ab0d2bde88725d95e2aa052dd8f87c9a5f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 14:30:02 -0400 Subject: [PATCH 04/14] fix: name the right culprit when one bundle claims a name twice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- python/datafusion/context.py | 22 ++++++++++++++++------ python/tests/test_context.py | 16 ++++++++++++++++ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 5b78ccc54..3e02a8d12 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -265,12 +265,22 @@ def _resolve_declared_functions( raise TypeError(msg) name = wrapped.name if name in claimed: - msg = ( - f"Two extensions declare a {kind} named {name!r}: " - f"{claimed[name]!r} and {extension!r}. Registrations have no " - "fall-through, so one would silently replace the other; " - "install them on separate sessions, or rename one." - ) + # Identity, not equality: two objects in the argument list are two + # installs even when the bundle is a dataclass that compares equal + # to its twin. Only a bundle colliding with *itself* can rename. + if claimed[name] is extension: + msg = ( + f"{extension!r} declares two {kind}s named {name!r}. " + "Registrations have no fall-through, so the second would " + "silently replace the first; rename one of them." + ) + else: + msg = ( + f"Two extensions declare a {kind} named {name!r}: " + f"{claimed[name]!r} and {extension!r}. Registrations have " + "no fall-through, so one would silently replace the other; " + "install them on separate sessions." + ) raise ValueError(msg) claimed[name] = extension resolved.append(wrapped) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 50842c690..0db4e9a97 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1523,6 +1523,22 @@ def test_with_extensions_rejects_a_name_two_extensions_claim(ctx): ctx.udf("double") +def test_with_extensions_rejects_a_name_one_extension_claims_twice(ctx): + """A bundle colliding with itself is its own bug, not a clash of libraries. + + Separated from the two-extension case because the remedy differs: a bundle + author can rename their own function, and a caller cannot rename someone + else's. + """ + with pytest.raises( + ValueError, match=r"declares two scalar functions named 'double'" + ): + ctx.with_extensions(_FunctionExtension(udfs=(_doubler(), _doubler()))) + + with pytest.raises(KeyError): + ctx.udf("double") + + def test_with_extensions_allows_shadowing_an_existing_function(ctx): """Claiming a name the session already has is legal. From fc26394ce0aebf089cdd3d4895f6d6cc3b7820c1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 14:33:19 -0400 Subject: [PATCH 05/14] refactor: drive declared-function install off one table `_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) --- python/datafusion/context.py | 135 +++++++++++++++++++++++------------ python/tests/test_context.py | 29 ++++++++ 2 files changed, 117 insertions(+), 47 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 3e02a8d12..3076fedf7 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -46,7 +46,7 @@ import uuid import warnings -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, NamedTuple, Protocol try: from warnings import deprecated # Python 3.13+ @@ -161,6 +161,70 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 +class _FunctionKind(NamedTuple): + """How one kind of declared function is resolved and registered. + + One row per function field on + :py:class:`~datafusion.extensions.SessionExtensionComponents`, so adding a + kind is adding a row rather than editing three places. The dataclass + metadata says which fields are collections to normalize; this table says + which of them are functions and what to do with one. + ``test_every_component_field_has_an_installer`` pins the two together. + + The members naming a ``datafusion.user_defined`` object hold its name + rather than the object: that module imports this one, so the lookups are + deferred to :py:meth:`SessionContext.with_extensions`, which runs with the + cycle long settled. + """ + + field: str + """The ``SessionExtensionComponents`` field a bundle declares these in.""" + + wrapper: str + """Wrapper class a declaration may already be an instance of.""" + + getter: str + """Capsule getter an unwrapped declaration must expose instead.""" + + factory: str + """Helper that turns an unwrapped declaration into a wrapper.""" + + label: str + """What to call this sort of function in an error message.""" + + register: str + """:py:class:`SessionContext` method that commits one to the session.""" + + +_FUNCTION_KINDS = ( + _FunctionKind( + field="udfs", + wrapper="ScalarUDF", + getter="__datafusion_scalar_udf__", + factory="udf", + label="scalar function", + register="register_udf", + ), + _FunctionKind( + field="udafs", + wrapper="AggregateUDF", + getter="__datafusion_aggregate_udf__", + factory="udaf", + label="aggregate function", + register="register_udaf", + ), + _FunctionKind( + field="udwfs", + wrapper="WindowUDF", + getter="__datafusion_window_udf__", + factory="udwf", + label="window function", + register="register_udwf", + ), +) +"""Every kind of function a bundle can declare, in registration order.""" + + def _collect_contributions( extensions: tuple[object, ...], ctx: SessionContext, @@ -180,8 +244,8 @@ def _collect_contributions( context derived from it. Returns: - The logical codecs, the physical codecs, and the declared functions by - field name. + The logical codecs, the physical codecs, and the declared functions + keyed by the :py:data:`_FUNCTION_KINDS` field they arrived in. Raises: TypeError: If an argument implements neither hook, or a hook returns @@ -201,9 +265,7 @@ def _collect_contributions( logical_codecs: list[LogicalExtensionCodecExportable] = [] physical_codecs: list[PhysicalExtensionCodecExportable] = [] declared: dict[str, list[tuple[object, Any]]] = { - "udfs": [], - "udafs": [], - "udwfs": [], + kind.field: [] for kind in _FUNCTION_KINDS } for extension in extensions: if not isinstance(extension, SessionComponentsExportable): @@ -2122,41 +2184,23 @@ def with_extensions( # Resolve every declared function to the wrapper that registers it, and # settle name collisions, while a failure still costs nothing. None of # these getters take an argument, so unlike a provider they do not care - # which handle they are resolved against. - from datafusion.user_defined import ( # noqa: PLC0415 - AggregateUDF as _AggregateUDF, - ) - from datafusion.user_defined import ( # noqa: PLC0415 - ScalarUDF as _ScalarUDF, - ) - from datafusion.user_defined import ( # noqa: PLC0415 - WindowUDF as _WindowUDF, - ) - from datafusion.user_defined import udaf as _udaf # noqa: PLC0415 - from datafusion.user_defined import udf as _udf # noqa: PLC0415 - from datafusion.user_defined import udwf as _udwf # noqa: PLC0415 - - resolved_udfs = _resolve_declared_functions( - declared["udfs"], - _ScalarUDF, - "__datafusion_scalar_udf__", - _udf, - "scalar function", - ) - resolved_udafs = _resolve_declared_functions( - declared["udafs"], - _AggregateUDF, - "__datafusion_aggregate_udf__", - _udaf, - "aggregate function", - ) - resolved_udwfs = _resolve_declared_functions( - declared["udwfs"], - _WindowUDF, - "__datafusion_window_udf__", - _udwf, - "window function", - ) + # which handle they are resolved against. The bound `register_*` method + # is looked up here too, leaving the commit below nothing but calls. + from datafusion import user_defined as _user_defined # noqa: PLC0415 + + resolved: list[tuple[Any, list[Any]]] = [ + ( + getattr(new, kind.register), + _resolve_declared_functions( + declared[kind.field], + getattr(_user_defined, kind.wrapper), + kind.getter, + getattr(_user_defined, kind.factory), + kind.label, + ), + ) + for kind in _FUNCTION_KINDS + ] # Phase two: nest the planners, outermost last. Each hook runs against # `new`, which carries the final chains, so a planner captured here @@ -2191,12 +2235,9 @@ def with_extensions( # has nothing to roll back to. See :ref:`ffi_internals_commit_order`. if planner is not None or logical_codecs or physical_codecs: new.ctx._install_extension_planner(planner) - for function in resolved_udfs: - new.register_udf(function) - for function in resolved_udafs: - new.register_udaf(function) - for function in resolved_udwfs: - new.register_udwf(function) + for register, functions in resolved: + for function in functions: + register(function) return new def table_provider(self, name: str) -> Table: diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 0db4e9a97..d59083d1c 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -20,6 +20,7 @@ import gzip import pathlib import shutil +from dataclasses import fields import pyarrow as pa import pyarrow.compute as pc @@ -1607,6 +1608,34 @@ def test_session_extension_components_rejects_a_single_function(field): SessionExtensionComponents(**{field: _doubler()}) +def test_every_component_field_has_an_installer(): + """A field added to the components dataclass must be wired into the install. + + ``SessionExtensionComponents`` normalizes any field carrying the component + metadata, so one added without an installer would be accepted from a + bundle and then quietly dropped — the failure this pins is a contributed + component going nowhere, with no error to say so. + + Reaching into private names on purpose: the two sides answer different + questions. The metadata says which fields are collections to normalize; + ``_FUNCTION_KINDS`` and the codec pair say which of them ``with_extensions`` + knows how to install. Nothing observable from outside can tell you they + have drifted, because the symptom is silence. + """ + from datafusion.context import _FUNCTION_KINDS + + by_noun: dict[str, set[str]] = {} + for spec in fields(SessionExtensionComponents): + noun = spec.metadata.get("datafusion_component") + if noun is not None: + by_noun.setdefault(noun, set()).add(spec.name) + + assert by_noun == { + "codec": {"logical_extension_codecs", "physical_extension_codecs"}, + "function": {kind.field for kind in _FUNCTION_KINDS}, + } + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) From 2dc115949e592b586963483e22037fe21467cf0c Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 14:36:16 -0400 Subject: [PATCH 06/14] docs: say what the commit order costs a bundle author `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) --- docs/source/extension-guide/bundles.md | 8 ++++++++ docs/source/user-guide/extensions.md | 2 -- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 487fdf618..4d327f17b 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -336,6 +336,14 @@ 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 codec chains belong to the returned handle; the planner is installed on the diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index e41747512..72c839654 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -116,8 +116,6 @@ 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`. - - **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 it is still in use, the next query fails with: From 6216bf10bf492b5e22f78ece82ef187f04ef965b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 14:36:16 -0400 Subject: [PATCH 07/14] docs: pass volatility by keyword, drop a dead cross-reference 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) --- python/datafusion/context.py | 3 ++- python/datafusion/extensions.py | 6 +++--- python/tests/test_context.py | 19 ++++++++++++++++--- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 3076fedf7..445df72b1 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -304,7 +304,8 @@ def _resolve_declared_functions( Args: declared: ``(extension, function)`` pairs in declaration order. - wrapper: The Python wrapper class for this kind, passed through already. + wrapper: The wrapper class a declaration may already be an instance of, + in which case it is taken as-is. getter: The capsule getter an unwrapped declaration must expose. factory: The helper that wraps a declaration — ``udf`` and friends. kind: What to call this sort of function in an error. diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 5988edc3b..1bb413237 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -121,8 +121,8 @@ def _components(noun: str) -> Any: ``noun`` names what the field holds, for the error a bundle sees when it hands over one component instead of a collection of them. Carrying it in - the field metadata is what lets :py:meth:`SessionExtensionComponents.__post_init__` - normalize a field it was never told about by name. + the field metadata is what lets ``__post_init__`` normalize a field it was + never told about by name. """ return field(default=(), metadata={"datafusion_component": noun}) @@ -190,7 +190,7 @@ class SessionExtensionComponents: ... lambda arr: pa.array([v.as_py() * 2 for v in arr]), ... [pa.int64()], ... pa.int64(), - ... "stable", + ... volatility="stable", ... name="double", ... ) >>> components = SessionExtensionComponents(udfs=(double,)) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index d59083d1c..3a289b3f5 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1422,7 +1422,7 @@ def _doubler(name="double"): lambda arr: pa.array([v.as_py() * 2 for v in arr]), [pa.int64()], pa.int64(), - "stable", + volatility="stable", name=name, ) @@ -1483,8 +1483,21 @@ def test_with_extensions_registers_a_declared_udf(ctx): def test_with_extensions_registers_udafs_and_udwfs(ctx): """The other two function kinds install the same way.""" - total = udaf(_Total, pa.int64(), pa.int64(), [pa.int64()], "stable", name="total") - first = udwf(_First, pa.int64(), pa.int64(), "immutable", name="first_value_of") + total = udaf( + _Total, + pa.int64(), + pa.int64(), + [pa.int64()], + volatility="stable", + name="total", + ) + first = udwf( + _First, + pa.int64(), + pa.int64(), + volatility="immutable", + name="first_value_of", + ) result = ctx.with_extensions(_FunctionExtension(udafs=(total,), udwfs=(first,))) result.from_pydict({"a": [1, 2, 3]}, name="nums") From e6889482e5e4f2b7128d6d867217a13a22cc8187 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 15:00:27 -0400 Subject: [PATCH 08/14] fix: pick the collision remedy by argument position `_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) --- docs/source/extension-guide/bundles.md | 12 ++++- docs/source/user-guide/extensions.md | 4 ++ python/datafusion/context.py | 66 ++++++++++++++++---------- python/tests/test_context.py | 25 ++++++++++ 4 files changed, 81 insertions(+), 26 deletions(-) diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 4d327f17b..c3f04d1e1 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -298,7 +298,8 @@ 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': ... +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 @@ -306,6 +307,15 @@ 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 diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index 72c839654..0d02ad78b 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -116,6 +116,10 @@ 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 it is still in use, the next query fails with: diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 445df72b1..988645da4 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -228,15 +228,16 @@ class _FunctionKind(NamedTuple): def _collect_contributions( extensions: tuple[object, ...], ctx: SessionContext, -) -> tuple[list[Any], list[Any], dict[str, list[tuple[object, Any]]]]: +) -> tuple[list[Any], list[Any], dict[str, list[tuple[int, object, Any]]]]: """Run every components hook and gather what the extensions contribute. Validates the whole argument list before calling anything, so an argument that implements neither hook is refused before a well-formed extension ahead of it has done any work. Writes nothing to the session. - Functions are kept paired with the extension that declared them, so a name - claimed twice can name both sides. + Functions are kept paired with the extension that declared them and with + that extension's position in the argument list, so a name claimed twice can + name both sides and tell which remedy applies. Args: extensions: The arguments ``with_extensions`` was given. @@ -244,8 +245,9 @@ def _collect_contributions( context derived from it. Returns: - The logical codecs, the physical codecs, and the declared functions - keyed by the :py:data:`_FUNCTION_KINDS` field they arrived in. + The logical codecs, the physical codecs, and the declared functions as + ``(position, extension, function)`` triples, keyed by the + :py:data:`_FUNCTION_KINDS` field they arrived in. Raises: TypeError: If an argument implements neither hook, or a hook returns @@ -264,10 +266,10 @@ def _collect_contributions( logical_codecs: list[LogicalExtensionCodecExportable] = [] physical_codecs: list[PhysicalExtensionCodecExportable] = [] - declared: dict[str, list[tuple[object, Any]]] = { + declared: dict[str, list[tuple[int, object, Any]]] = { kind.field: [] for kind in _FUNCTION_KINDS } - for extension in extensions: + for position, extension in enumerate(extensions): if not isinstance(extension, SessionComponentsExportable): continue components = extension.__datafusion_session_components__(ctx) @@ -282,17 +284,17 @@ def _collect_contributions( physical_codecs.extend(components.physical_extension_codecs) for field_name, declarations in declared.items(): declarations.extend( - (extension, item) for item in getattr(components, field_name) + (position, extension, item) for item in getattr(components, field_name) ) return logical_codecs, physical_codecs, declared def _resolve_declared_functions( - declared: list[tuple[object, Any]], + declared: list[tuple[int, object, Any]], wrapper: type, getter: str, factory: Any, - kind: str, + label: str, ) -> list[Any]: """Wrap every function an extension declared, refusing a name claimed twice. @@ -303,49 +305,63 @@ def _resolve_declared_functions( argument is ignored. Args: - declared: ``(extension, function)`` pairs in declaration order. + declared: ``(position, extension, function)`` triples in declaration + order, where ``position`` indexes the ``with_extensions`` argument + list. wrapper: The wrapper class a declaration may already be an instance of, in which case it is taken as-is. getter: The capsule getter an unwrapped declaration must expose. factory: The helper that wraps a declaration — ``udf`` and friends. - kind: What to call this sort of function in an error. + label: What to call this sort of function in an error. Returns: The wrappers to register, in declaration order. + + Raises: + TypeError: If a declaration is neither an instance of ``wrapper`` nor + an object exposing ``getter``. + ValueError: If two declarations resolve to the same name. """ resolved = [] - claimed: dict[str, object] = {} - for extension, function in declared: + claimed: dict[str, tuple[int, object]] = {} + for position, extension, function in declared: if isinstance(function, wrapper): wrapped = function elif hasattr(function, getter): wrapped = factory(function) else: msg = ( - f"A declared {kind} must be a {wrapper.__name__} or expose " + f"A declared {label} must be a {wrapper.__name__} or expose " f"{getter}, got {function!r} from {extension!r}" ) raise TypeError(msg) name = wrapped.name if name in claimed: - # Identity, not equality: two objects in the argument list are two - # installs even when the bundle is a dataclass that compares equal - # to its twin. Only a bundle colliding with *itself* can rename. - if claimed[name] is extension: + # Position, not object identity: what the caller controls is the + # argument list, and two entries in it are two installs whether or + # not they are the same object. The distinction the message needs + # is which remedy exists. One argument colliding with itself is a + # bundle author's own bug, and only they can rename a function; + # two arguments colliding is the caller's to resolve, and renaming + # is not among the things a caller can do. + claimed_at, claimed_by = claimed[name] + if claimed_at == position: msg = ( - f"{extension!r} declares two {kind}s named {name!r}. " + f"{extension!r} declares two {label}s named {name!r}. " "Registrations have no fall-through, so the second would " "silently replace the first; rename one of them." ) else: msg = ( - f"Two extensions declare a {kind} named {name!r}: " - f"{claimed[name]!r} and {extension!r}. Registrations have " - "no fall-through, so one would silently replace the other; " - "install them on separate sessions." + f"Two extensions declare a {label} named {name!r}: " + f"argument {claimed_at} ({claimed_by!r}) and argument " + f"{position} ({extension!r}). Registrations have no " + "fall-through, so one would silently replace the other; " + "install them on separate sessions, or drop the repeat if " + "one extension was passed twice." ) raise ValueError(msg) - claimed[name] = extension + claimed[name] = (position, extension) resolved.append(wrapped) return resolved diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3a289b3f5..288b4c197 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1526,15 +1526,40 @@ def test_with_extensions_rejects_a_name_two_extensions_claim(ctx): Unlike codecs, which dispatch by id, a second function under one name would silently replace the first. + + A codec-carrying bundle rides along to pin the other half of the + transaction: resolution runs *after* the codec chains are built, so this + failure lands between the two steps, and the chains must not reach the + session either. """ with pytest.raises(ValueError, match=r"scalar function named 'double'"): ctx.with_extensions( + _CodecOnlyExtension(), _FunctionExtension(udfs=(_doubler(),)), _FunctionExtension(udfs=(_doubler(),)), ) with pytest.raises(KeyError): ctx.udf("double") + assert ctx.logical_extension_codec_ids() == [] + assert ctx.physical_extension_codec_ids() == [] + + +def test_with_extensions_rejects_one_extension_passed_twice(ctx): + """A bundle object listed twice is the caller's duplicate, not a naming bug. + + The remedy has to match the mistake, and nothing the bundle author renames + helps here — both claims come from the one declaration. Collisions are + therefore keyed on argument position rather than on object identity, which + would read a repeat as a bundle colliding with itself and offer a rename + that cannot be made. + """ + extension = _FunctionExtension(udfs=(_doubler(),)) + + with pytest.raises(ValueError, match=r"argument 0 .* and argument 1 ") as excinfo: + ctx.with_extensions(extension, extension) + + assert "rename" not in str(excinfo.value) def test_with_extensions_rejects_a_name_one_extension_claims_twice(ctx): From 659cda73df197c1b8bcd7fc61242971213c00577 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 15:01:41 -0400 Subject: [PATCH 09/14] test: pin that every function-kind row names something real `_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) --- python/tests/test_context.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 288b4c197..427e52f6e 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1674,6 +1674,26 @@ def test_every_component_field_has_an_installer(): } +def test_every_function_kind_names_something_real(): + """The names on a ``_FunctionKind`` row resolve to what it says they do. + + They are held as strings and looked up during ``with_extensions``, to keep + the ``user_defined`` import out of this module's import cycle. The cost is + that a typo in a row surfaces as an ``AttributeError`` part-way through an + install rather than at import. Every row that exists today is covered by a + behaviour test above; this is what covers the next one, which may be added + before its own test is. + """ + from datafusion import user_defined + from datafusion.context import _FUNCTION_KINDS + + for kind in _FUNCTION_KINDS: + assert isinstance(getattr(user_defined, kind.wrapper), type) + assert callable(getattr(user_defined, kind.factory)) + assert callable(getattr(SessionContext, kind.register)) + assert kind.field in {spec.name for spec in fields(SessionExtensionComponents)} + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) From a9df02b10b7942eb7c5d80e6b7460844ce5bd758 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 15:01:48 -0400 Subject: [PATCH 10/14] refactor: split the commit rule from the planner-rebinding note 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) --- python/datafusion/context.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 988645da4..ca5e0c45c 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2245,11 +2245,14 @@ def with_extensions( # it already holds, so the rebuild is unobservable except in the one case # where it does harm: a planner sitting on some *other* handle's codecs # gets dragged onto this handle's, silently undoing that install. - # Everything below this line must be infallible. A registration whose - # commit can fail belongs above, split into an import step that returns - # a resolved object and an insert step that cannot raise -- there is one - # session here, shared with the receiver, so a failure part-way through - # has nothing to roll back to. See :ref:`ffi_internals_commit_order`. + + # Commit. Everything below this line must be infallible. A registration + # whose commit can fail belongs above, split into an import step that + # returns a resolved object and an insert step that cannot raise -- + # there is one session here, shared with the receiver, so a failure + # part-way through has nothing to roll back to. The reasoning is in + # docs/source/contributor-guide/ffi-internals.md, under "Why + # `with_extensions` commits last". if planner is not None or logical_codecs or physical_codecs: new.ctx._install_extension_planner(planner) for register, functions in resolved: From f05b80858737f5fd6501bd1d2581d0c5d38b46cf Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 15:02:00 -0400 Subject: [PATCH 11/14] style: spell out the pyo3 imports in the bundle example 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) --- examples/datafusion-ffi-example/src/extension.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/datafusion-ffi-example/src/extension.rs b/examples/datafusion-ffi-example/src/extension.rs index 777db0cb5..fe94a667f 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods}; use crate::aggregate_udf::MySumUDF; use crate::scalar_udf::IsNullUDF; From 6bf48e77ea1dbe6b6acc534926f7898a72440e76 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 15:12:49 -0400 Subject: [PATCH 12/14] style: tidy wording in _resolve_declared_functions 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 --- python/datafusion/context.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index ca5e0c45c..75d8e66ab 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -331,19 +331,15 @@ def _resolve_declared_functions( wrapped = factory(function) else: msg = ( - f"A declared {label} must be a {wrapper.__name__} or expose " + f"A declared {label} must be {wrapper.__name__} or expose " f"{getter}, got {function!r} from {extension!r}" ) raise TypeError(msg) name = wrapped.name if name in claimed: - # Position, not object identity: what the caller controls is the - # argument list, and two entries in it are two installs whether or - # not they are the same object. The distinction the message needs - # is which remedy exists. One argument colliding with itself is a - # bundle author's own bug, and only they can rename a function; - # two arguments colliding is the caller's to resolve, and renaming - # is not among the things a caller can do. + # Keyed on position, not object identity, so each message names + # the remedy its reader actually has — see + # `extension_bundles_collisions` in the extension guide. claimed_at, claimed_by = claimed[name] if claimed_at == position: msg = ( From 5c2b6c7bd7ce5c36189cb222855139bf003165f2 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 15:12:58 -0400 Subject: [PATCH 13/14] test: cover function-name collisions for every kind 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 --- python/tests/test_context.py | 68 ++++++++++++++++++++++++------------ 1 file changed, 46 insertions(+), 22 deletions(-) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 427e52f6e..a0276063a 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1454,6 +1454,29 @@ def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array: return pa.array([first] * num_rows) +def _total(name="total"): + """An aggregate function under a name the caller picks.""" + return udaf( + _Total, + pa.int64(), + pa.int64(), + [pa.int64()], + volatility="stable", + name=name, + ) + + +def _first(name="first_value_of"): + """A window function under a name the caller picks.""" + return udwf( + _First, + pa.int64(), + pa.int64(), + volatility="immutable", + name=name, + ) + + class _FunctionExtension: """Contributes functions and nothing else. @@ -1483,23 +1506,9 @@ def test_with_extensions_registers_a_declared_udf(ctx): def test_with_extensions_registers_udafs_and_udwfs(ctx): """The other two function kinds install the same way.""" - total = udaf( - _Total, - pa.int64(), - pa.int64(), - [pa.int64()], - volatility="stable", - name="total", - ) - first = udwf( - _First, - pa.int64(), - pa.int64(), - volatility="immutable", - name="first_value_of", + result = ctx.with_extensions( + _FunctionExtension(udafs=(_total(),), udwfs=(_first(),)) ) - - result = ctx.with_extensions(_FunctionExtension(udafs=(total,), udwfs=(first,))) result.from_pydict({"a": [1, 2, 3]}, name="nums") assert result.sql("SELECT total(a) FROM nums").collect()[0].column(0) == pa.array( @@ -1521,26 +1530,39 @@ def test_with_extensions_registers_on_the_shared_session(ctx): assert ctx.udf("double").name == "double" -def test_with_extensions_rejects_a_name_two_extensions_claim(ctx): +@pytest.mark.parametrize( + ("field", "make", "label", "lookup"), + [ + ("udfs", _doubler, "scalar function", "udf"), + ("udafs", _total, "aggregate function", "udaf"), + ("udwfs", _first, "window function", "udwf"), + ], +) +def test_with_extensions_rejects_a_name_two_extensions_claim( + ctx, field, make, label, lookup +): """Registrations have no fall-through, so a clash cannot be resolved by order. Unlike codecs, which dispatch by id, a second function under one name would - silently replace the first. + silently replace the first. Parametrized over the kinds to pin each + ``_FUNCTION_KINDS`` row's field and label wiring, not just the machinery. A codec-carrying bundle rides along to pin the other half of the transaction: resolution runs *after* the codec chains are built, so this failure lands between the two steps, and the chains must not reach the session either. """ - with pytest.raises(ValueError, match=r"scalar function named 'double'"): + name = make().name + + with pytest.raises(ValueError, match=rf"{label} named '{name}'"): ctx.with_extensions( _CodecOnlyExtension(), - _FunctionExtension(udfs=(_doubler(),)), - _FunctionExtension(udfs=(_doubler(),)), + _FunctionExtension(**{field: (make(),)}), + _FunctionExtension(**{field: (make(),)}), ) with pytest.raises(KeyError): - ctx.udf("double") + getattr(ctx, lookup)(name) assert ctx.logical_extension_codec_ids() == [] assert ctx.physical_extension_codec_ids() == [] @@ -1560,6 +1582,8 @@ def test_with_extensions_rejects_one_extension_passed_twice(ctx): ctx.with_extensions(extension, extension) assert "rename" not in str(excinfo.value) + with pytest.raises(KeyError): + ctx.udf("double") def test_with_extensions_rejects_a_name_one_extension_claims_twice(ctx): From ebb894364d16e6277473716944ddc9169e50df65 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 13:44:46 -0400 Subject: [PATCH 14/14] refactor: fold the planner hooks and the commit into one Rust call The tail of `with_extensions` was three private pymethods and two Python loops: nest the planner hooks, re-export each return as a capsule, bind the planner, then walk a table of bound `register_*` methods. Replace all of it with one `_commit_extensions` primitive that runs the hooks and commits everything the bundles declared, so the ordering contract lives in a single function next to the reasoning it answers to, and the boundary drops `_export_query_planner` and `_install_extension_planner`. The Python side keeps everything that reads better in Python: protocol dispatch, collision messages naming argument positions, and the resolve step. Each function kind is now its own `_commit_extensions` parameter, so `_FunctionKind` loses its `register` member -- a kind that resolves but never commits cannot be written, because the call's arity refuses it. The hook loop dispatches on the presence of `__datafusion_session_planner__`, which is the same question the runtime-checkable protocol asked. Behaviour is pinned unchanged: no test assertion moves beyond the private-method allowlist and the dropped `register` line in the meta-test. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 119 +++++++++++++++++--------- python/datafusion/context.py | 93 +++++++------------- python/tests/test_context.py | 1 - python/tests/test_wrapper_coverage.py | 9 +- 4 files changed, 114 insertions(+), 108 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index fe75668fb..a3be7dda0 100644 --- a/crates/core/src/context.rs +++ b/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,87 @@ 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> { - 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. + /// Everything after the hooks is the commit, and none of it can fail: a + /// registration whose commit can fail belongs in the resolve step, split + /// into an import that returns a resolved object and an insert that + /// cannot raise. There is one session here, shared with the receiver, so + /// a failure part-way through would have nothing to roll back to. The + /// reasoning is in 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 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. /// - /// 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 functions are registered *after* the planner hooks have run, so a + /// hook never sees this call's functions in the registry — the + /// registrations have no fall-through, and a name is free to shadow one + /// the session already had. + pub fn _commit_extensions<'py>( slf: &Bound<'py, Self>, - planner: Option>, + extensions: Vec>, + session: Bound<'py, PyAny>, + rebind_planner: bool, + udfs: Vec, + udafs: Vec, + udwfs: Vec, ) -> 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 = 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); + } + let this = slf.borrow(); + for udf in udfs { + this.ctx.register_udf(udf.function); + } + for udaf in udafs { + this.ctx.register_udaf(udaf.function); + } + for udwf in udwfs { + this.ctx.register_udwf(udwf.function); + } Ok(()) } } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 75d8e66ab..328840ded 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -162,13 +162,15 @@ def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 class _FunctionKind(NamedTuple): - """How one kind of declared function is resolved and registered. + """How one kind of declared function is resolved. One row per function field on :py:class:`~datafusion.extensions.SessionExtensionComponents`, so adding a kind is adding a row rather than editing three places. The dataclass metadata says which fields are collections to normalize; this table says - which of them are functions and what to do with one. + which of them are functions and how to wrap one. Committing is not here: + each kind is its own ``_commit_extensions`` parameter, so a kind that + resolves but never commits cannot be written. ``test_every_component_field_has_an_installer`` pins the two together. The members naming a ``datafusion.user_defined`` object hold its name @@ -192,9 +194,6 @@ class _FunctionKind(NamedTuple): label: str """What to call this sort of function in an error message.""" - register: str - """:py:class:`SessionContext` method that commits one to the session.""" - _FUNCTION_KINDS = ( _FunctionKind( @@ -203,7 +202,6 @@ class _FunctionKind(NamedTuple): getter="__datafusion_scalar_udf__", factory="udf", label="scalar function", - register="register_udf", ), _FunctionKind( field="udafs", @@ -211,7 +209,6 @@ class _FunctionKind(NamedTuple): getter="__datafusion_aggregate_udf__", factory="udaf", label="aggregate function", - register="register_udaf", ), _FunctionKind( field="udwfs", @@ -219,7 +216,6 @@ class _FunctionKind(NamedTuple): getter="__datafusion_window_udf__", factory="udwf", label="window function", - register="register_udwf", ), ) """Every kind of function a bundle can declare, in registration order.""" @@ -2197,63 +2193,38 @@ def with_extensions( # Resolve every declared function to the wrapper that registers it, and # settle name collisions, while a failure still costs nothing. None of # these getters take an argument, so unlike a provider they do not care - # which handle they are resolved against. The bound `register_*` method - # is looked up here too, leaving the commit below nothing but calls. + # which handle they are resolved against. from datafusion import user_defined as _user_defined # noqa: PLC0415 - resolved: list[tuple[Any, list[Any]]] = [ - ( - getattr(new, kind.register), - _resolve_declared_functions( - declared[kind.field], - getattr(_user_defined, kind.wrapper), - kind.getter, - getattr(_user_defined, kind.factory), - kind.label, - ), + resolved: dict[str, list[Any]] = { + kind.field: _resolve_declared_functions( + declared[kind.field], + getattr(_user_defined, kind.wrapper), + kind.getter, + getattr(_user_defined, kind.factory), + kind.label, ) for kind in _FUNCTION_KINDS - ] - - # Phase two: nest the planners, outermost last. Each hook runs against - # `new`, which carries the final chains, so a planner captured here - # never sees a partial codec set. `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. - planner: _PyCapsule | None = None - for extension in extensions: - if not isinstance(extension, SessionPlannerExportable): - continue - fallback = ( - planner - if planner is not None - else new.ctx.__datafusion_query_planner__() - ) - supplied = extension.__datafusion_session_planner__(new, fallback) - if supplied is None: - continue - planner = new.ctx._export_query_planner(supplied) - - # Rebinding the session's planner is a side effect on state shared with - # every other handle, so do not pay it for a call that installs nothing - # -- the same guard `with_python_udf_inlining` carries. With no codec - # installed the chains the planner would be rebuilt against are the ones - # it already holds, so the rebuild is unobservable except in the one case - # where it does harm: a planner sitting on some *other* handle's codecs - # gets dragged onto this handle's, silently undoing that install. - - # Commit. Everything below this line must be infallible. A registration - # whose commit can fail belongs above, split into an import step that - # returns a resolved object and an insert step that cannot raise -- - # there is one session here, shared with the receiver, so a failure - # part-way through has nothing to roll back to. The reasoning is in - # docs/source/contributor-guide/ffi-internals.md, under "Why - # `with_extensions` commits last". - if planner is not None or logical_codecs or physical_codecs: - new.ctx._install_extension_planner(planner) - for register, functions in resolved: - for function in functions: - register(function) + } + + # Phase two: run the planner hooks and commit, in one call. Each hook + # runs against `new`, which carries the final chains, so a planner + # captured there never sees a partial codec set. The hook loop, the + # ordering of the commit, and the guard that skips the planner rebind + # for a call that installs nothing all live on the Rust side -- see + # `_commit_extensions` and docs/source/contributor-guide/ + # ffi-internals.md, under "Why `with_extensions` commits last". A new + # component field must be resolved above and given its own + # `_commit_extensions` parameter; the call's arity is what keeps a + # declared component from being quietly dropped. + new.ctx._commit_extensions( + list(extensions), + new, + bool(logical_codecs or physical_codecs), + [function._udf for function in resolved["udfs"]], + [function._udaf for function in resolved["udafs"]], + [function._udwf for function in resolved["udwfs"]], + ) return new def table_provider(self, name: str) -> Table: diff --git a/python/tests/test_context.py b/python/tests/test_context.py index a0276063a..e9b42b493 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1714,7 +1714,6 @@ def test_every_function_kind_names_something_real(): for kind in _FUNCTION_KINDS: assert isinstance(getattr(user_defined, kind.wrapper), type) assert callable(getattr(user_defined, kind.factory)) - assert callable(getattr(SessionContext, kind.register)) assert kind.field in {spec.name for spec in fields(SessionExtensionComponents)} diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 6927632b9..00bbe920b 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -33,12 +33,11 @@ # gap in coverage. PRIVATE_SUPPORT_METHODS = frozenset( { - # The three steps of SessionContext.with_extensions: install the - # codecs, re-export each planner hook's return value as a capsule, - # commit the planner. + # The two steps of SessionContext.with_extensions: install the + # codecs, then run the planner hooks and commit everything the + # bundles declared. "_install_extension_codecs", - "_export_query_planner", - "_install_extension_planner", + "_commit_extensions", } )