diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index fe75668fb..63b7a6c82 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,68 @@ 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. + /// Binding the planner is the commit, and it cannot fail. Anything else a + /// call installs is committed by the Python caller once this returns, + /// through `register_udf`, `register_udaf` and `register_udwf`, which + /// cannot fail. Whatever can fail belongs before this call — see + /// 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 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 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. + pub fn _commit_extensions<'py>( slf: &Bound<'py, Self>, - planner: Option>, + extensions: Vec>, + session: Bound<'py, PyAny>, + rebind_planner: bool, ) -> 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); + } Ok(()) } } diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index bbf6b65e7..b00ab11d8 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -117,37 +117,35 @@ prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`. `with_extensions` promises that a bundle which raises leaves the session as it was, apart from anything a hook writes to the context it is handed -({ref}`extension_bundles_transaction`). Keeping that promise is an ordering constraint on the implementation, not -a property of any one step, because the planner is bound on the shared -`SessionState` rather than on the returned handle. +({ref}`extension_bundles_transaction`). Keeping that promise is an ordering +constraint on the implementation, not a property of any one step, because the +components a bundle declares do not 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 four steps, of which only the last writes: -1. **Collect.** Every `__datafusion_session_components__` runs and its codecs - are gathered. Nothing is installed yet, so a hook that raises here has - touched nothing. +1. **Collect.** Every `__datafusion_session_components__` runs and its + components are gathered. Nothing is installed yet, so a hook that raises here + has touched nothing. 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 to the session even though it can fail on a bad capsule or a duplicate id. -3. **Resolve.** Every `__datafusion_session_planner__` runs, in argument order, - against the completed chains, and each supplied planner is exported to a - capsule. Every hook that can raise has run by the end of this step. -4. **Commit.** The accumulated planner is re-imported from its capsule and - bound, in a single `SessionState` rebuild. The bind is skipped entirely when - the call installed nothing, so an empty call does not drag a planner sitting - on another handle's codecs onto this one's. - -Only step 4 touches the session, and it is not itself infallible: -`_install_extension_planner` runs `ffi_query_planner_from_pycapsule` before it -calls `set_session_query_planner`, which cannot fail. The property that keeps -the promise is therefore about order, not about any step being incapable of -raising — every fallible operation, including the ones inside the commit, -completes before the first write. - -That is the 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 — before anything is -written, so no failure can leave the session half-updated. +3. **Resolve.** Every declared function is wrapped and every name is checked. + Then every `__datafusion_session_planner__` runs, in argument order, against + the completed chains, and each supplied planner is imported as it is + returned. +4. **Commit.** The planner is bound, then the declared functions are + registered with `register_udf`, `register_udaf` and `register_udwf` — the + same methods a caller would use, and three that cannot fail. The bind is + skipped entirely when the call installed no planner and no codec, so an + empty call does not drag a planner sitting on another handle's codecs onto + this one's. + +Only step 4 touches the session, and every fallible operation completes before +its first write. 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 no failure can leave the session half-updated. There would be nothing to roll back to if one did. The returned handle shares one session with the receiver, so the damage is visible from every other handle; diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 437deb165..cedbe0f4c 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,20 @@ 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 — 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,22 +286,76 @@ 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 + +One `with_extensions` call may not declare a function of the same kind under +the same name twice. Doing so raises: + +```text +ValueError: A scalar function named 'normalize' is declared twice: 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. + +The argument each claim came from is in the message because it is what picks +the remedy, and the reader works it out from the two positions. Two different +positions means two bundles, which only the caller can resolve — by dropping a +repeat, or by installing them on separate sessions. One position named twice +means a single bundle claimed the name twice, which is its author's own bug and +the one case a rename fixes. Collisions are keyed on position rather than on +object identity so that passing one extension twice reads as the caller's +duplicate that it is, rather than as a bundle colliding with itself. + +Three 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. +- **One call at a time is checked.** The names compared are the ones declared + in a single `with_extensions` call. Two calls on the same session, or a call + following a hand-written `register_udf`, land in the registry one after the + other, and the later one silently replaces the earlier — that is the + shadowing rule above, applied to something you may not have meant to shadow. + Splitting colliding bundles across two calls therefore does not resolve the + collision; it hides it. Use two sessions. + +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. - -Declaring a component is what buys you that guarantee. Anything you return from -your hook is validated 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`. +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..f2d816cb6 100644 --- a/docs/source/extension-guide/checklist.md +++ b/docs/source/extension-guide/checklist.md @@ -59,13 +59,13 @@ 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. + → {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..62225446f 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,45 @@ 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 +## Four 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: A scalar function named 'normalize' is declared twice: argument 0 +() and argument 1 (). ... +``` + +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. +Installing the two in separate `with_extensions` calls on one session is not a +fix: only the names within a single call are compared, so the second library's +function quietly replaces the first's. 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. +If both positions are the *same* number, only one library is involved and it +declared the name twice — nothing on your side fixes that, so report it. + +**Functions outlive the handle you installed them on.** `with_extensions` +returns a new context, and its codecs belong to that context alone — but +functions are registered on the *session*, which every handle shares. So this +changes the context you called it on: + +```python +ctx.with_extensions(my_library.Extension()) # return value dropped +ctx.udf("my_library_normalize") # ...and it is there anyway +``` + +Use the returned context regardless — you need it for the codecs, and it is +what the next section's checks read. But do not count on dropping it to undo +an install. **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 +185,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 c663dea2b..c3b3b6f0f 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+ @@ -163,6 +163,164 @@ class TableProviderExportable(Protocol): def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 +class _Contributions(NamedTuple): + """Everything a call's extensions contributed, before any of it is resolved. + + One member per component field on + :py:class:`~datafusion.extensions.SessionExtensionComponents`. A field added + there needs a member here and a resolve step in + :py:meth:`SessionContext.with_extensions`; + ``test_every_component_field_has_an_installer`` is what fails if it gets + neither. + + 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. + """ + + logical_codecs: list[LogicalExtensionCodecExportable] + """Logical codecs, in declaration order.""" + + physical_codecs: list[PhysicalExtensionCodecExportable] + """Physical codecs, in declaration order.""" + + udfs: list[tuple[int, object, Any]] + """Declared scalar functions as ``(position, extension, function)``.""" + + udafs: list[tuple[int, object, Any]] + """Declared aggregate functions, as :py:attr:`udfs`.""" + + udwfs: list[tuple[int, object, Any]] + """Declared window functions, as :py:attr:`udfs`.""" + + +def _collect_contributions( + extensions: tuple[object, ...], + ctx: SessionContext, +) -> _Contributions: + """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. + + 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: + What every bundle declared, gathered by kind. + + Raises: + TypeError: If an argument implements neither hook, or a hook returns + something other than ``SessionExtensionComponents``. + """ + for extension in extensions: + if not isinstance( + extension, + ( + _extensions.SessionComponentsExportable, + _extensions.SessionPlannerExportable, + ), + ): + msg = ( + "Extension implements neither " + "__datafusion_session_components__ nor " + f"__datafusion_session_planner__: {extension!r}" + ) + raise TypeError(msg) + + contributed = _Contributions([], [], [], [], []) + for position, extension in enumerate(extensions): + if not isinstance(extension, _extensions.SessionComponentsExportable): + continue + components = extension.__datafusion_session_components__(ctx) + if not isinstance(components, _extensions.SessionExtensionComponents): + msg = ( + "__datafusion_session_components__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + contributed.logical_codecs.extend(components.logical_extension_codecs) + contributed.physical_codecs.extend(components.physical_extension_codecs) + for declared, functions in ( + (contributed.udfs, components.udfs), + (contributed.udafs, components.udafs), + (contributed.udwfs, components.udwfs), + ): + declared.extend((position, extension, item) for item in functions) + return contributed + + +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: one argument listed + # twice is what tells the reader a single bundle claimed the name + # twice, which is the only case a rename can fix. See + # `extension_bundles_collisions` in the extension guide. + claimed_at, claimed_by = claimed[name] + msg = ( + f"A {label} named {name!r} is declared twice: argument " + f"{claimed_at} ({claimed_by!r}) and argument {position} " + f"({extension!r}). Registrations have no fall-through, so one " + "would silently replace the other. Rename one if both came " + "from an extension you own, drop the repeat if one extension " + "was passed twice, or install them on separate sessions." + ) + raise ValueError(msg) + claimed[name] = (position, extension) + resolved.append(wrapped) + return resolved + + class SessionConfig: """Session configuration options.""" @@ -1918,10 +2076,10 @@ 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. See + 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 is not rolled back — see :ref:`extension_bundles_transaction`. Shares its session with this context — see :py:class:`SessionContext`. @@ -1942,10 +2100,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 one call declares a + function of one kind under the same name twice, or if a + getter returns a capsule of the wrong kind. See :py:meth:`with_logical_extension_codec` for how ids are assigned. @@ -1979,82 +2140,76 @@ def with_extensions( >>> batches[0].column(0).to_pylist() # doctest: +SKIP [1] """ - for extension in extensions: - if not isinstance( - extension, - ( - _extensions.SessionComponentsExportable, - _extensions.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, _extensions.SessionComponentsExportable): - continue - components = extension.__datafusion_session_components__(self) - if not isinstance(components, _extensions.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) + contributed = _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, _extensions.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) - - # The commit step, and the only one that writes to the session. It can - # still raise -- `_install_extension_planner` re-imports the capsule - # before binding it -- so what keeps the promise is the order, not any - # step being incapable of failing: every fallible operation finishes - # before the first write. See docs/source/contributor-guide/ - # ffi-internals.md, "Why `with_extensions` commits last", for what a new - # component kind has to do to keep that true. - # - # 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) + new.ctx = self.ctx._install_extension_codecs( + contributed.logical_codecs, contributed.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. `user_defined` imports this + # module, so the import waits until here, with the cycle long settled. + from datafusion import user_defined as _ud # noqa: PLC0415 + + udfs = _resolve_declared_functions( + contributed.udfs, + _ud.ScalarUDF, + "__datafusion_scalar_udf__", + _ud.udf, + "scalar function", + ) + udafs = _resolve_declared_functions( + contributed.udafs, + _ud.AggregateUDF, + "__datafusion_aggregate_udf__", + _ud.udaf, + "aggregate function", + ) + udwfs = _resolve_declared_functions( + contributed.udwfs, + _ud.WindowUDF, + "__datafusion_window_udf__", + _ud.udwf, + "window function", + ) + + # Phase two: run the planner hooks, then bind the planner. Each hook + # runs against `new`, which carries the final chains, so a planner + # captured there never sees a partial codec set. The hook loop and the + # guard that skips the rebind for a call that installs nothing live on + # the Rust side -- see `_commit_extensions`. The list is narrowed with + # the same check the validation above uses, so one predicate decides + # both what is admitted and what is called. + new.ctx._commit_extensions( + [ + e + for e in extensions + if isinstance(e, _extensions.SessionPlannerExportable) + ], + new, + bool(contributed.logical_codecs or contributed.physical_codecs), + ) + + # The rest of the commit. These three registrations cannot fail, which + # is what lets them run after the point of no return -- see + # docs/source/contributor-guide/ffi-internals.md, under "Why + # `with_extensions` commits last". It runs after the planner hooks, so + # a hook never sees this call's functions in the registry. + for function in udfs: + new.register_udf(function) + for function in udafs: + new.register_udaf(function) + for function in 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 cbcc1a9ac..61b590888 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -45,7 +45,7 @@ from __future__ import annotations -from dataclasses import dataclass, fields +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -58,8 +58,14 @@ from datafusion.context import SessionContext from datafusion.user_defined import ( + AggregateUDF, + AggregateUDFExportable, LogicalExtensionCodecExportable, PhysicalExtensionCodecExportable, + ScalarUDF, + ScalarUDFExportable, + WindowUDF, + WindowUDFExportable, ) __all__ = [ @@ -144,15 +150,31 @@ 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." ) +_COMPONENT_NOUNS = { + "logical_extension_codecs": "codec", + "physical_extension_codecs": "codec", + "udfs": "function", + "udafs": "function", + "udwfs": "function", +} +"""Fields of :py:class:`SessionExtensionComponents` holding a tuple of +components, and what to call one of them in an error. + +Listing them here rather than marking each field is what leaves room for a +future field that is not a collection and must not be normalized into a tuple. +``test_every_component_field_has_an_installer`` pins the contents. +""" + + @dataclass(frozen=True, kw_only=True) class SessionExtensionComponents: """Components an extension contributes to a session context. @@ -170,6 +192,10 @@ 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. Declare a component here rather than + registering it during the hook — 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: @@ -204,7 +230,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)) @@ -229,8 +271,37 @@ class SessionExtensionComponents: ``__datafusion_physical_extension_codec__``. """ + udfs: tuple[ScalarUDF | ScalarUDFExportable, ...] = () + """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. + + One :py:meth:`~datafusion.context.SessionContext.with_extensions` call may + not declare a name twice, whether the two claims come from one extension or + two; shadowing a function the session already has is allowed. See + :ref:`extension_bundles_collisions`. + """ + + udafs: tuple[AggregateUDF | AggregateUDFExportable, ...] = () + """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, ...] = () + """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 @@ -243,25 +314,17 @@ def __post_init__(self) -> None: # member of an immutable value, and a generator would be exhausted by # 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"): - continue + for name, noun in _COMPONENT_NOUNS.items(): 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 @@ -277,9 +340,8 @@ 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. See + next call may install onto a different session. Declare what you contribute + rather than registering it on the context you are handed — see :ref:`extension_bundles_transaction`. A bundle that also contributes a query planner implements diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 274f2a9b1..f95757852 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(): @@ -1200,6 +1206,32 @@ def __datafusion_session_planner__(self, ctx, fallback): assert batches[0].column(0) == pa.array([1]) +def test_with_extensions_ignores_a_planner_attribute_set_to_none(ctx): + """Regression test: ``__datafusion_session_planner__ = None`` is not a hook. + + A base class spelling out that its subclasses contribute no planner binds + the attribute to ``None`` rather than omitting it. ``isinstance`` against + :py:class:`~datafusion.extensions.SessionPlannerExportable` rejects that, + but a ``hasattr`` does not, so a host that validated with one and dispatched + with the other called the ``None``: ``'NoneType' object is not callable``, + naming neither the extension nor the hook. + + ``_commit_extensions`` still carries a Rust-side ``hasattr``, and the + Python caller now hands it a list already narrowed by ``isinstance``. Two + checks means one of them decides, and this pins which: a refactor that + passes the unfiltered arguments through reintroduces the failure, and + nothing else in this file notices. + """ + + class NoPlannerExtension(_CodecOnlyExtension): + __datafusion_session_planner__ = None + + result = ctx.with_extensions(NoPlannerExtension()) + + assert result.logical_extension_codec_ids() == ["my_library.logical"] + assert result.physical_extension_codec_ids() == ["my_library.physical"] + + def test_with_extensions_rejects_bad_codec_capsule(ctx): """A correctly shaped object still has to return the right capsule.""" @@ -1425,6 +1457,291 @@ 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 one'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. + + Two argument positions is what says so, and it is the reason collisions are + keyed on position rather than on object identity: identity would read a + repeat as one bundle colliding with itself, pointing the reader at a rename + of someone else's function that they cannot make. + """ + extension = _FunctionExtension(udfs=(_doubler(),)) + + with pytest.raises(ValueError, match=r"argument 0 .* and argument 1 "): + ctx.with_extensions(extension, extension) + + 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 caught by the same check. + + One argument position named twice is what distinguishes it, and it is the + case the message's rename remedy is for — the only reader who can rename a + function is the one who declared both of them. + """ + with pytest.raises(ValueError, match=r"argument 0 .* and argument 0 ") as excinfo: + ctx.with_extensions(_FunctionExtension(udfs=(_doubler(), _doubler()))) + + assert "scalar function named 'double'" in str(excinfo.value) + 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 every field named in + ``_COMPONENT_NOUNS``, 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. Nothing + observable from outside can catch that, because the symptom is silence. + + If this fails because you added a field: give it a member on + ``datafusion.context._Contributions``, collect it in + ``_collect_contributions``, and resolve and install it in + ``with_extensions``. Then add it below. + + Keep the comparison an equality. A subset check is satisfied by any new + field on its own, which is the one case this exists to catch. A field that + is deliberately not a component collection belongs in ``not_components``. + """ + from datafusion.extensions import _COMPONENT_NOUNS + + assert set(_COMPONENT_NOUNS) == { + "logical_extension_codecs", + "physical_extension_codecs", + "udfs", + "udafs", + "udwfs", + } + # Fields that are not component collections: normalized by nothing, + # installed by nothing, and exempt on purpose. Empty today. + not_components: set[str] = set() + declared = {spec.name for spec in fields(SessionExtensionComponents)} + assert declared - not_components == set(_COMPONENT_NOUNS) + + 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..1f108c376 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -33,12 +33,12 @@ # 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 Rust steps of SessionContext.with_extensions: install the + # codecs, then run the planner hooks and bind the planner. Declared + # functions are registered from Python with the public `register_*` + # methods once this returns. "_install_extension_codecs", - "_export_query_planner", - "_install_extension_planner", + "_commit_extensions", } )