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/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 791fdfc8e..c3f04d1e1 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): @@ -54,16 +55,24 @@ class MyEngineExtension: return self._make_planner(ctx, fallback=fallback) ``` -Implement whichever apply: a codec-only library defines the first, a library -that ships only an optimizing planner defines the second. The caller then -writes: +Implement only the hooks you need. Codecs and functions both go in +`__datafusion_session_components__`, with the fields you do not use left empty, +so a codec-only library and a function-only library each define that one alone; +a library shipping nothing but an optimizing planner defines only +`__datafusion_session_planner__`. The caller then writes: ```python ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) ctx.register_table("t", lib_a.TableProvider()) -ctx.register_udf(udf(lib_b.SomeUDF())) ``` +Return your functions rather than calling `register_udf` on the `ctx` you were +handed. Both put the function on the session, but a registration you make +inside the hook is written the moment it runs — before the other bundles have +been called, and not undone if one of them raises. What you declare is instead +resolved and checked while a failure still costs nothing, then written once +every bundle has succeeded. See {ref}`extension_bundles_transaction`. + `MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust implementation of the protocol, including taking the task-context provider off the supplied context, wrapping its codecs in `BundledLogicalCodec` / @@ -281,14 +290,69 @@ for direct ones. The wrapper travels with the codec; the bundle does not. The query planner is exempt — it carries no wire id, so it may be an object or a capsule. +(extension_bundles_collisions)= + +## Two bundles claiming one name + +Two extensions in one call may not declare a function of the same kind under +the same name. Doing so raises: + +```text +ValueError: Two extensions declare a scalar function named 'normalize': +argument 0 (...) and argument 1 (...). ... +``` + +Codecs get away with sharing a chain because a payload carries the id of the +codec that wrote it, so decode routes to the right one. A function registry has +no such fall-through — one name holds one function — so the second registration +would quietly replace the first. The call refuses instead. + +Which argument each claim came from is part of the message because it is what +picks the remedy. Two arguments colliding is the caller's to resolve, by +installing the two on separate sessions or by dropping a repeat; renaming is +not something a caller can do. One argument declaring a name twice is the +bundle author's own bug, and gets a different message saying so. Collisions are +keyed on position rather than on object identity, so passing one extension +twice reads as the caller's duplicate that it is, rather than as a bundle +colliding with itself. + +Two cases this does *not* catch: + +- **Different kinds never collide.** Names are compared within a kind, so a + scalar function and an aggregate may both be called `normalize`. +- **Shadowing a built-in is allowed.** The registry already holds every + DataFusion function, and replacing one by name is a supported thing to do — + `enable_spark_functions` works that way. + +Your caller cannot rename your function, so stay out of the way: prefix the +names with something tied to your library. + +(extension_bundles_transaction)= + ## Failure and rollback Nothing is written to the session until every factory has returned and every -capsule has been validated, so a factory that raises leaves the session exactly -as it was. A factory that mutates the context it is handed — registering a -table, say — is **not** rolled back, which is why bundle objects must be -configuration-only: create fresh components on each call, never cache bound -components, and do not retain the context passed in. +component has been validated, so a factory that raises leaves the session +exactly as it was. A factory that mutates the context it is handed — +registering a table, say — is **not** rolled back, which is why bundle objects +must be configuration-only: create fresh components on each call, never cache +bound components, and do not retain the context passed in. + +Declaring a component is what buys you that guarantee, and it is the whole +reason to prefer `udfs=(...)` over a `register_udf` call inside your hook. +Anything you declare is resolved and checked while a failure still costs +nothing, and is written only after every bundle in the call has succeeded. +Anything you register yourself is written immediately, before the other bundles +have even run. The ordering that makes this hold is recorded at +{ref}`ffi_internals_commit_order`. + +The one thing that ordering costs you: functions are registered *after* the +planner hooks run, so `ctx.udfs()` inside your +`__datafusion_session_planner__` will not list a function declared in the same +call — not yours, and not another bundle's. Look one up at plan time instead, +where the registry is complete; a planner is called per query, long after the +install has finished. If you need a function at hook time, you already have the +object, because you are the one declaring it. Like every other derivation, the returned context is a handle on the *same* session as the receiver — see {ref}`extension_sessions`. Only the Python-side 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..ee38090e2 100644 --- a/docs/source/extension-guide/functions.md +++ b/docs/source/extension-guide/functions.md @@ -26,14 +26,18 @@ functions in pure Python — see {doc}`../user-guide/common-operations/udf-and-udfa` — and the two roads meet at the same registration methods. -| Hook | Contributes | Wrapped by | Registered with | -| --- | --- | --- | --- | -| `__datafusion_scalar_udf__` | scalar function | {py:func}`datafusion.udf` | {py:meth}`~datafusion.SessionContext.register_udf` | -| `__datafusion_aggregate_udf__` | aggregate function | {py:func}`datafusion.udaf` | {py:meth}`~datafusion.SessionContext.register_udaf` | -| `__datafusion_window_udf__` | window function | {py:func}`datafusion.udwf` | {py:meth}`~datafusion.SessionContext.register_udwf` | -| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` | - -All four are implemented in [`datafusion-ffi-example`], one per file. +| Hook | Contributes | Wrapped by | Registered with | Declared in a bundle as | +| --- | --- | --- | --- | --- | +| `__datafusion_scalar_udf__` | scalar function | {py:func}`datafusion.udf` | {py:meth}`~datafusion.SessionContext.register_udf` | `udfs` | +| `__datafusion_aggregate_udf__` | aggregate function | {py:func}`datafusion.udaf` | {py:meth}`~datafusion.SessionContext.register_udaf` | `udafs` | +| `__datafusion_window_udf__` | window function | {py:func}`datafusion.udwf` | {py:meth}`~datafusion.SessionContext.register_udwf` | `udwfs` | +| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` | — | + +All four are implemented in [`datafusion-ffi-example`], one per file. The last +column is the {py:class}`~datafusion.SessionExtensionComponents` field a +{ref}`bundle ` declares the function in; table functions +have no such field yet, so they are always registered by the caller with +{py:meth}`~datafusion.SessionContext.register_udtf`. ## The three scalar-shaped hooks @@ -67,6 +71,39 @@ from datafusion import udf ctx.register_udf(udf(my_library.MyScalarUDF())) ``` +If your library ships more than a function or two, do not make your users write +that line once per function. Ship a {ref}`bundle ` declaring +them, so one call installs the lot: + +```python +ctx = SessionContext().with_extensions(my_library.MyFunctionExtension()) +``` + +The bundle is yours to write, and like the rest of the protocol it is an object +exposing a getter — which your cdylib can export directly. That is what +`MyFunctionExtension` in [`datafusion-ffi-example`] does for this crate's three +functions; spelled in Python, it is: + +```python +from datafusion import SessionExtensionComponents + + +class MyFunctionExtension: + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents( + udfs=(IsNullUDF(),), + udafs=(MySumUDF(),), + udwfs=(MyRankUDF(),), + ) +``` + +Declare either the raw exportable, as here, or an already-wrapped +{py:class}`~datafusion.user_defined.ScalarUDF`; the registered name comes off +the function either way. Declare rather than calling `register_udf` inside the +hook — see {ref}`extension_bundles_transaction` for why — and pick names that +will not collide with another library's +({ref}`extension_bundles_collisions`). + ## Table functions A table function takes literal `Expr` arguments and returns a table provider, diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index 8c86b4a67..0d02ad78b 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -32,13 +32,12 @@ which exposes Delta Lake tables to DataFusion, and the two worked examples in this repository under [`examples/`](https://github.com/apache/datafusion-python/tree/main/examples). -## Two kinds of extension +## How an extension reaches your session -Which one you have determines how much setup you do. +Which route your library takes determines how much setup you do. -**Tables and functions register directly.** If the library gives you a table -or a function, register it the same way you would register a CSV file. No -extra setup: +**Tables register directly.** If the library gives you a table, register it +the same way you would register a CSV file. No extra setup: ```python from datafusion import SessionContext @@ -64,6 +63,21 @@ ctx.register_table("events", my_engine.TableProvider("s3://bucket/events")) ctx.sql("SELECT count(*) FROM events").show() ``` +**Functions arrive by whichever route their library chose.** A library +offering one or two functions hands you the functions themselves, and you wrap +and register each: + +```python +from datafusion import udf + +ctx.register_udf(udf(my_library.MyScalarUDF())) +``` + +A library shipping a set of them packages them in its `Extension` object +instead, so `with_extensions` installs them all along with everything else it +provides, and there is nothing per-function for you to do. Its documentation +says which. + `with_extensions` returns a context; use the returned one. It shares everything else with the context you called it on, so tables you registered before the call are still there. @@ -85,7 +99,26 @@ rarely matters. When a library needs a particular position — usually "list me last" for something that wraps the others — it says so in its own documentation. -## Two things that will bite you +## Three things that will bite you + +**Two libraries can claim one function name.** If both ship a function of the +same kind under the same name, the call raises a `ValueError` naming both, +rather than letting one silently replace the other: + +```text +ValueError: Two extensions declare a scalar function named 'normalize': ... +``` + +You cannot rename another library's function from your own code, so the fix is +to use two sessions, one per library, and query each for what only it provides. +Worth reporting upstream too: the library whose names are the less specific +should be prefixing them. A function shadowing a *built-in* is not a collision +and raises nothing — that is a supported thing for a library to do. See +{ref}`extension_bundles_collisions`. + +Check the argument positions the message names before you go looking for a +second library. Passing one extension twice collides with itself, and an +extension list assembled from a plugin registry is the usual way that happens. **Keep your context alive.** A `DataFrame` or a plan does not keep its session alive on its own. If a context is garbage-collected while something built from @@ -133,6 +166,17 @@ ctx.logical_extension_codec_ids() An empty list means nothing extra is installed. +For functions, {py:meth}`~datafusion.SessionContext.udfs`, +{py:meth}`~datafusion.SessionContext.udafs` and +{py:meth}`~datafusion.SessionContext.udwfs` return the names a session knows. +Both the library's and every DataFusion built-in are in there, so look for the +name rather than reading the whole list: + +```python +"my_engine_normalize" in ctx.udfs() +# True +``` + ## Next steps - {ref}`distributed_query_engines` — running your queries across several 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..fe94a667f --- /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::types::{PyAnyMethods, PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods}; + +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..328840ded 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,203 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 +class _FunctionKind(NamedTuple): + """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 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 + 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.""" + + +_FUNCTION_KINDS = ( + _FunctionKind( + field="udfs", + wrapper="ScalarUDF", + getter="__datafusion_scalar_udf__", + factory="udf", + label="scalar function", + ), + _FunctionKind( + field="udafs", + wrapper="AggregateUDF", + getter="__datafusion_aggregate_udf__", + factory="udaf", + label="aggregate function", + ), + _FunctionKind( + field="udwfs", + wrapper="WindowUDF", + getter="__datafusion_window_udf__", + factory="udwf", + label="window function", + ), +) +"""Every kind of function a bundle can declare, in registration order.""" + + +def _collect_contributions( + extensions: tuple[object, ...], + ctx: SessionContext, +) -> 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 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. + 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 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 + 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[int, object, Any]]] = { + kind.field: [] for kind in _FUNCTION_KINDS + } + for position, extension in enumerate(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( + (position, extension, item) for item in getattr(components, field_name) + ) + return logical_codecs, physical_codecs, declared + + +def _resolve_declared_functions( + declared: list[tuple[int, object, Any]], + wrapper: type, + getter: str, + factory: Any, + label: 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: ``(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. + 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, 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 {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: + # 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 = ( + 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 {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] = (position, extension) + resolved.append(wrapped) + return resolved + + class SessionConfig: """Session configuration options.""" @@ -1915,10 +2112,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 +2137,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,70 +2177,54 @@ 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) - # 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__() + # 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 import user_defined as _user_defined # noqa: PLC0415 + + 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, ) - 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. - if planner is not None or logical_codecs or physical_codecs: - new.ctx._install_extension_planner(planner) + for kind in _FUNCTION_KINDS + } + + # 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/datafusion/extensions.py b/python/datafusion/extensions.py index 77ae92fc2..1bb413237 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 ``__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(), + ... volatility="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..e9b42b493 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -20,11 +20,14 @@ import gzip import pathlib import shutil +from dataclasses import fields 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 +38,11 @@ Table, column, literal, + udaf, udf, + udwf, ) +from datafusion.user_defined import WindowEvaluator def test_create_context_no_args(): @@ -1410,6 +1416,307 @@ 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(), + volatility="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) + + +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. + + 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.""" + 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" + + +@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. 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. + """ + name = make().name + + with pytest.raises(ValueError, match=rf"{label} named '{name}'"): + ctx.with_extensions( + _CodecOnlyExtension(), + _FunctionExtension(**{field: (make(),)}), + _FunctionExtension(**{field: (make(),)}), + ) + + with pytest.raises(KeyError): + getattr(ctx, lookup)(name) + 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) + with pytest.raises(KeyError): + 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. + + ``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_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_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 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]]) 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", } )