timsaucer opened a new pull request, #1738: URL: https://github.com/apache/datafusion-python/pull/1738
# Which issue does this PR close? Part 1 of 4 toward #1676. The remaining parts will be stacked on this one: physical optimizer rules, then table providers and table functions, then catalog providers, which will carry the `Closes`. See https://github.com/apache/datafusion-python/issues/1676#issuecomment-5680959562 for the analysis this plan came out of, including what in that issue went stale after #1679 and why `object_stores` was split out to #1737 as blocked upstream. # Rationale for this change #1679 added `with_extensions` so a library could ship its codecs and planner as one atomic bundle, but `SessionExtensionComponents` carries codec fields only. A library contributing *functions* has nowhere to put them, so its users still follow a per-function setup recipe — which is the situation `with_extensions` was introduced to remove. The concrete case is in #1721: `dfx_udfs` deliberately exposes no bundle hook, and carries a test named `test_this_library_cannot_be_installed_as_a_bundle` asserting the gap rather than describing it. Its `build_session` needs two codec installs and three registrations in an order the caller has to get right alone, next to a sibling library that is one call. # What changes are included in this PR? **`SessionExtensionComponents` gains `udfs`, `udafs`, and `udwfs`.** Each accepts either the Python wrapper (`ScalarUDF`) or a raw object exposing the capsule getter; the host wraps the latter with `udf` / `udaf` / `udwf`. The registered name is read off the resolved wrapper, which on the FFI path is the name the capsule reports — not anything the bundle or the host invented. **Installation now separates what can fail from what cannot.** This is the part worth reviewing closely. `with_extensions` currently gets its transactionality for free: codec chains live on the returned Python handle rather than on `SessionState`, so `_install_extension_planner` is the single write and a failing hook leaves the session untouched. Registrations do not have that property — they write into the shared `SessionState`, the returned handle *is* the source session since #1679 removed the derived context, and there is nothing to roll back to. Rollback is also not available in principle: `deregister_udf` removes a name but does not restore a built-in the bundle shadowed, and `register_catalog` returns the displaced provider that a commit would have discarded, so an "undo" can delete a user's own registration. So the call is now four steps, and only the last one writes: 1. **Collect** — every `__datafusion_session_components__` runs. 2. **Chains** — `_install_extension_codecs` assembles the returned handle. Can fail; writes nothing. 3. **Resolve** — declared functions are wrapped, names are checked, and every `__datafusion_session_planner__` runs against the completed chains. Can fail; writes nothing. 4. **Commit** — the planner is bound, then the functions are registered. Cannot fail. `SessionContext::register_udf` is infallible upstream (it returns `()` and swallows through `.ok()`), which is what makes step 4 honest rather than merely hopeful. A comment at the commit boundary states the rule for the fields landing in the PRs above this one: a new component does its fallible work in step 3, never in step 4. Two tests pin the ordering — a components hook raising and a planner hook raising both leave zero functions registered on both handles. The second is the one that matters: it fails if the registrations are committed before phase two, which is a change that passes every other test here. **Name collisions within one call are a `ValueError` naming both extensions.** Codec ids can dispatch on decode, so a chain holds many and picks the right one; a function registry has no such fall-through and a second registration silently replaces the first. Names are compared per kind, so a scalar function and an aggregate may share one. Shadowing a name the session *already* holds stays legal — the registry contains every DataFusion built-in, and `enable_spark_functions` overrides built-ins by design, so refusing that would refuse a supported use rather than catch a mistake. **`__post_init__` normalizes by field metadata instead of the `_codecs` name suffix.** The suffix check was explicitly reserving non-suffixed fields for things that must *not* become tuples, which is the opposite of what `udfs` wants. A `_components(noun)` helper tags each field with what it holds, so the "you forgot the trailing comma" error names the right noun and the fields in the later PRs are covered without anyone remembering to list them. **Real-FFI coverage.** `MyFunctionExtension` in `datafusion-ffi-example` declares that crate's scalar, aggregate, and window functions from one hook, and `_test_session_extension.py` runs all three across the boundary, plus bundle reuse across sessions, the collision refusal, and the transaction. The three `#[new]` constructors became `pub(crate)` so the bundle can build them. **Docs.** `extension-guide/bundles.md` gains an `extension_bundles_transaction` label over a rewritten "Failure and rollback" that lays out the four steps and states the infallible-commit rule as a contract, plus an `extension_bundles_collisions` section — the existing text covered only duplicate codec ids. `functions.md` gains a "declared in a bundle as" column and a bundle example. `checklist.md` keeps the configuration-only prohibition and changes its reason to "declare it instead". `user-guide/extensions.md` no longer claims functions always register directly. # Are there any user-facing changes? Three new optional fields on `SessionExtensionComponents`, all defaulting to `()` on a frozen dataclass. No hook signatures change, no wire format changes, and no shipped extension library implements anything new, so there is no upgrade-guide entry and no `api change` label. One behaviour worth naming: functions declared by a bundle are registered on the session, which the source context shares — `with_extensions` returns a handle on the same session, and only the codec chains belong to the returned handle. A test pins that deliberately so it cannot change by accident. # Review notes **The `udfs` field buys driver-side ergonomics, not worker parity.** For a library that also ships a codec, worker-side registration is already unnecessary: `dfx_udfs`'s own tests show a worker with only the codec runs the plan, and note that "the registry is tried first, so the codec is the fallback, not the path". This is still worth shipping — it is what lets a user write the library's function in SQL after one call — it is just a smaller claim than #1676 makes. **`with_extensions` cannot install everything a library provides, and should not claim to.** A config extension has to reach `SessionConfig` before the context exists, which no bundle hook can reach. `dfx_engine` needs exactly that. **Nothing here touches `examples/distributed/`.** #1721 is open against `main` in parallel; whichever of the two lands second gets a small follow-up converting `dfx_udfs` to ship a bundle and flipping `test_this_library_cannot_be_installed_as_a_bundle`, whose premise this PR inverts. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
