From cc391dfaecbb42f80bf2b161a547e01395fad453 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 10:25:44 -0400 Subject: [PATCH 1/4] feat: let extension bundles declare catalog providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `SessionExtensionComponents.catalog_providers` completes the set #1676 asked for, minus object stores, which have no FFI type upstream and are tracked separately in #1737. `register_catalog_provider` had its import and its insert in one function with no split point, so the import half moves to a shared `resolve_catalog_provider` that both it and the new `_resolve_extension_catalogs` call. Both therefore accept exactly the same shapes, and the bundle path gets the same treatment as tables: imported against the handle carrying the finished codec chains, since the getter is handed the logical codec its provider will serialize through. Catalogs replace rather than collide. `register_catalog` returns whichever provider it displaced, and the default `datafusion` catalog always exists, so a library backing a session with its own metadata has to be able to replace one. Only two bundles claiming a name in the same call is refused. That is the opposite of tables, where a duplicate is an error, and both now say so where a reader meets them. `_install_extension_catalogs` returns `()`: nothing is left that can fail once the providers are imported. Rules 2 and 6 of the capsule-protocol skill now carry the two conventions the stack established — which components a bundle hands over unwrapped and why, and that a new component means a new resolve step rather than a fallible commit step. Closes #1676. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 32 +++++ crates/core/src/context.rs | 114 ++++++++++++---- docs/source/extension-guide/bundles.md | 10 +- .../source/extension-guide/table-providers.md | 14 ++ .../python/tests/_test_session_extension.py | 37 ++++++ .../datafusion-ffi-example/src/extension.rs | 41 ++++++ examples/datafusion-ffi-example/src/lib.rs | 3 +- python/datafusion/context.py | 27 +++- python/datafusion/extensions.py | 13 ++ python/tests/test_context.py | 122 +++++++++++++++++- python/tests/test_wrapper_coverage.py | 2 + 11 files changed, 377 insertions(+), 38 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 37be0bc9d..0cfbd3c9b 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -76,6 +76,25 @@ receiver had — the same session, and the same task-context provider, but not this call's codecs, not even your own. Read the host's codec chains in the planner hook, never in the extension hook. +**That is why a bundle declares unresolved components, not wrapped ones.** The +components it returns split by what their getter asks for: + +- Getters taking no argument — the three function kinds, + `__datafusion_physical_optimizer_rule__` — have nothing session-scoped to + bind, so a bundle may hand over either the raw exportable or an + already-wrapped object. +- Getters taking the session or a codec — `__datafusion_table_function__`, + `__datafusion_table_provider__`, `__datafusion_catalog_provider__` — must be + handed over **unwrapped**, with a name. Wrapping one inside the components + hook would call its getter with the `ctx` that hook received, capturing a + chain missing every library in the call. The host wraps these itself, against + the handle carrying the final chains, which is the only place that chain + exists. + +`RecordingTableFunction` in `examples/datafusion-ffi-example/src/extension.rs` +records the ids it was resolved against, so the difference is asserted rather +than described. + A *codec* must always be handed over as an object implementing its getter, never as the bare capsule the getter returns; `with_extensions` refuses a capsule. A codec's wire id — the string a payload names on decode, which has to mean the @@ -187,6 +206,19 @@ an instruction to derive one first. It is not: the factories are handed the receiver, and the returned handle shares its allocation. There is nothing to keep alive separately and nothing to garbage-collect out from under a provider. +It is also the one place where sharing an allocation has a cost, and the cost +shapes how a component is added to it. Because the returned handle *is* the +receiver's session, a failure part-way through has nothing to roll back to. So +`with_extensions` does every fallible thing first — importing capsules, +resolving names, running the planner hooks — and only then writes. **Adding a +new kind of component means adding a resolve step, never a fallible commit +step:** a `_resolve_extension_*` that returns an opaque carrier, and a new +parameter on `_commit_extensions` whose commit cannot fail. The one exception +is table registration, whose insert goes through a `SchemaProvider` that a +foreign library may implement; it is committed first so nothing else is +written behind it. Do not add a second exception without the same +justification. + `SessionContext.enable_url_table` is the one method that mints a second allocation for a session. Its result must not outlive the receiver, and it also forks the session's `SessionState` while keeping its id, so two handles report diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 9a92c6055..3e8d89587 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -844,35 +844,9 @@ impl PySessionContext { pub fn register_catalog_provider( &self, name: &str, - mut provider: Bound<'_, PyAny>, + provider: Bound<'_, PyAny>, ) -> PyDataFusionResult<()> { - if provider.hasattr("__datafusion_catalog_provider__")? { - let py = provider.py(); - let ffi = self.ffi_logical_codec(); - let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; - provider = call_capsule_getter( - provider, - "__datafusion_catalog_provider__", - CapsuleGetterArg::LogicalCodec(&codec_capsule), - )?; - } - - let provider = if let Ok(capsule) = provider.cast::() { - let data: NonNull = capsule - .pointer_checked(Some(c"datafusion_catalog_provider"))? - .cast(); - let provider = unsafe { data.as_ref() }; - let provider: Arc = provider.into(); - provider - } else { - match provider.extract::() { - Ok(py_catalog) => py_catalog.catalog, - Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( - provider.into(), - self.ffi_logical_codec(), - )) as Arc, - } - }; + let provider = self.resolve_catalog_provider(provider)?; let _ = self.ctx.register_catalog(name, provider); @@ -1775,6 +1749,7 @@ impl PySessionContext { udafs: Vec, udwfs: Vec, udtfs: Vec, + catalogs: PyRef<'_, PyResolvedCatalogs>, rules: PyRef<'_, PyPhysicalOptimizerRules>, ) -> PyDataFusionResult<()> { let py = slf.py(); @@ -1824,6 +1799,12 @@ impl PySessionContext { for udtf in udtfs { this.register_udtf(udtf); } + // Nothing here can fail: the providers were imported by + // [`Self::_resolve_extension_catalogs`], and `register_catalog` + // returns whichever provider it displaced rather than refusing. + for (name, provider) in &catalogs.catalogs { + let _ = this.ctx.register_catalog(name, Arc::clone(provider)); + } // Rules accumulate rather than replace, so unlike a planner there is // no composition order to get right and no collision to refuse. All // of them go on in **one** `SessionState` rebuild. @@ -1923,6 +1904,31 @@ impl PySessionContext { Ok(PyResolvedTables { tables: resolved }) } + /// Resolve the catalogs a `with_extensions` call declared. + /// + /// The fallible half. Each provider is imported against `self` — the handle + /// carrying the completed codec chains, since + /// `__datafusion_catalog_provider__` is handed the logical codec it will + /// serialize through. + /// + /// No name is refused here. `register_catalog` replaces rather than + /// rejects, and `datafusion` — the default catalog — always exists, so a + /// bundle replacing a catalog is ordinary rather than a mistake. Two + /// bundles claiming one name in the same call is refused on the Python + /// side, where both can be named. + /// + /// **Writes nothing.** + pub fn _resolve_extension_catalogs<'py>( + &self, + catalogs: Vec<(String, Bound<'py, PyAny>)>, + ) -> PyDataFusionResult { + let catalogs = catalogs + .into_iter() + .map(|(name, provider)| Ok((name, self.resolve_catalog_provider(provider)?))) + .collect::>>()?; + Ok(PyResolvedCatalogs { catalogs }) + } + /// Import the physical optimizer rules a `with_extensions` call declared. /// /// The fallible half of installing them, run while the call can still fail @@ -1964,6 +1970,14 @@ struct ResolvedTable { provider: Arc, } +/// Catalog providers imported for a `with_extensions` call. +/// +/// Opaque to Python, like [`PyResolvedTables`] and [`PyPhysicalOptimizerRules`]. +#[pyclass(name = "ResolvedCatalogs", module = "datafusion._internal")] +pub struct PyResolvedCatalogs { + catalogs: Vec<(String, Arc)>, +} + /// Physical optimizer rules imported for a `with_extensions` call. /// /// Opaque to Python, and deliberately not added to the module: it exists only @@ -1984,6 +1998,50 @@ pub struct PyPhysicalOptimizerRules { } impl PySessionContext { + /// Turn whatever a caller offered as a catalog provider into one. + /// + /// The fallible half of registering a catalog, shared by + /// [`Self::register_catalog_provider`] and + /// [`Self::_resolve_extension_catalogs`] so both accept exactly the same + /// shapes: an object exposing `__datafusion_catalog_provider__`, a bare + /// capsule, a [`PyCatalog`], or a Python object implementing the provider + /// interface. + /// + /// The getter is handed **this context's** logical codec, so which handle + /// this is called on decides what the provider will serialize through. + fn resolve_catalog_provider( + &self, + mut provider: Bound<'_, PyAny>, + ) -> PyDataFusionResult> { + if provider.hasattr("__datafusion_catalog_provider__")? { + let py = provider.py(); + let ffi = self.ffi_logical_codec(); + let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; + provider = call_capsule_getter( + provider, + "__datafusion_catalog_provider__", + CapsuleGetterArg::LogicalCodec(&codec_capsule), + )?; + } + + Ok(if let Ok(capsule) = provider.cast::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_catalog_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider + } else { + match provider.extract::() { + Ok(py_catalog) => py_catalog.catalog, + Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( + provider.into(), + self.ffi_logical_codec(), + )) as Arc, + } + }) + } + /// Write the session's query planner, in place. /// /// Pass `Some(planner)` to install one, or `None` to rebuild whichever diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 11dc278c7..33d3c685f 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -300,10 +300,10 @@ where the host can resolve them: - **Getters taking no argument** — the three function kinds and physical optimizer rules. Nothing is session-scoped, so a bundle may hand over either a wrapped object or the raw exportable. -- **Getters taking the session or a codec** — table functions and table - providers. These are resolved by the host against the *finished* handle, - which is why you hand over the unwrapped value and a name rather than a - {py:class}`~datafusion.user_defined.TableFunction` you built yourself. +- **Getters taking the session or a codec** — table functions, table providers, + and catalog providers. These are resolved by the host against the *finished* + handle, which is why you hand over the unwrapped value and a name rather than + a {py:class}`~datafusion.user_defined.TableFunction` you built yourself. Wrapping one inside your components hook binds it to the context that hook received, which has none of the call's codecs — so it would capture a chain missing every library in the call, including your own. @@ -348,6 +348,8 @@ is nothing to refuse. See {doc}`other-components`. Tables go the other way: a declared table name that is *already* on the session is an error too, so a table cannot shadow one the way a function can. +Catalogs are back on the function side of that line, and for a reason worth +reading before you declare one: see {doc}`table-providers`. Strictly, whether a duplicate registration replaces or refuses is the `SchemaProvider`'s own call, and the in-memory one a session starts with diff --git a/docs/source/extension-guide/table-providers.md b/docs/source/extension-guide/table-providers.md index 700ca1954..9e09e6a85 100644 --- a/docs/source/extension-guide/table-providers.md +++ b/docs/source/extension-guide/table-providers.md @@ -51,6 +51,20 @@ wrapped: `__datafusion_table_provider__` takes the session, and the one your components hook receives has none of the call's codecs yet. The host resolves it against the finished handle. See {ref}`extension_bundles_binding`. +Catalogs work the same way, as `catalog_providers`: + +```python +return SessionExtensionComponents(catalog_providers=(("engine", MyCatalog()),)) +``` + +with one difference worth knowing. A declared **table** name that is already +registered is an error, because DataFusion refuses a duplicate table rather +than replacing it. A **catalog** name is not: `register_catalog` returns +whichever provider it displaced, and the default `datafusion` catalog always +exists — so replacing one is the usual way a library backs a session with its +own metadata. Only two bundles claiming the same catalog name in one call is +refused. + Start with a table provider. Reach for the schema and catalog levels when your data source has its own namespace that should be browsable rather than registered table by table, and for the provider list only when your library is diff --git a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py index c2e3a1dfa..6cd7cf27e 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py +++ b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py @@ -23,6 +23,7 @@ import pytest from datafusion import SessionContext, SessionExtensionComponents from datafusion_ffi_example import ( + MyCatalogExtension, MyDataExtension, MyFunctionExtension, MyLogicalExtensionCodec, @@ -256,6 +257,42 @@ def test_a_table_name_already_registered_is_refused(): ctx.udf("my_custom_is_null") +def test_a_declared_catalog_is_queryable(): + """A catalog declared by a bundle is reachable by its qualified name.""" + ctx = SessionContext().with_extensions(MyCatalogExtension()) + + assert "declared_catalog" in ctx.catalog_names() + result = ctx.sql("SELECT * FROM declared_catalog.my_schema.my_table").collect() + assert result[0].num_rows > 0 + + +def test_four_libraries_install_in_one_call(): + """The whole point, across a real FFI boundary. + + Four independently declared bundles — functions, rules, a table and a table + function, a catalog — in one call, and a single query that touches three of + them while the fourth counts the planning it did. + """ + rules = MyRuleExtension() + ctx = SessionContext().with_extensions( + MyFunctionExtension(), + rules, + MyDataExtension(), + MyCatalogExtension(), + ) + + result = ctx.sql( + 'SELECT my_custom_is_null("A") AS n FROM declared_table ' + "UNION ALL " + "SELECT my_custom_is_null(units) AS n " + "FROM declared_catalog.my_schema.my_table" + ).collect() + + assert sum(batch.num_rows for batch in result) > 0 + assert rules.first_calls() > 0 + assert rules.second_calls() > 0 + + def test_the_hook_returns_the_components_type(): """The bundle builds a real dataclass, not a duck-typed stand-in. diff --git a/examples/datafusion-ffi-example/src/extension.rs b/examples/datafusion-ffi-example/src/extension.rs index 7b5d668fb..148439738 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -21,6 +21,7 @@ use pyo3::types::{PyAnyMethods, PyCapsule, PyDict, PyDictMethods}; use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods}; use crate::aggregate_udf::MySumUDF; +use crate::catalog_provider::MyCatalogProvider; use crate::physical_optimizer::MyPhysicalOptimizerRule; use crate::scalar_udf::IsNullUDF; use crate::table_function::MyTableFunction; @@ -244,3 +245,43 @@ impl MyDataExtension { components.call((), Some(&kwargs)) } } + +/// A bundle contributing a catalog. +/// +/// `__datafusion_catalog_provider__` takes the session and pulls the host's +/// logical codec off it, so like a table provider it is handed over unresolved +/// and the host binds it to the finished handle. +#[pyclass( + from_py_object, + name = "MyCatalogExtension", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone, Default)] +pub(crate) struct MyCatalogExtension {} + +#[pymethods] +impl MyCatalogExtension { + #[new] + fn new() -> Self { + Self {} + } + + 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( + "catalog_providers", + (("declared_catalog", Py::new(py, MyCatalogProvider::new()?)?),), + )?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 25ebb2f3e..a7dceb3e9 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -20,7 +20,7 @@ use pyo3::prelude::*; use crate::aggregate_udf::MySumUDF; use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList}; use crate::config::MyConfig; -use crate::extension::{MyDataExtension, MyFunctionExtension, MyRuleExtension}; +use crate::extension::{MyCatalogExtension, MyDataExtension, MyFunctionExtension, MyRuleExtension}; use crate::logical_extension_codec::MyLogicalExtensionCodec; use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec}; use crate::physical_extension_codec::MyPhysicalExtensionCodec; @@ -68,5 +68,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/python/datafusion/context.py b/python/datafusion/context.py index f46deb11e..4ea518220 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -308,10 +308,16 @@ def _collect_contributions( declared: dict[str, list[tuple[int, object, Any]]] = { kind.field: [] for kind in _FUNCTION_KINDS } - # The rest are not function kinds: tables and table functions are named by - # the bundle rather than by the value, and rules carry no name at all, so - # each has its own collision rule -- or none -- and its own installer. - for field_name in ("udtfs", "table_providers", "physical_optimizer_rules"): + # The rest are not function kinds: tables, table functions, and catalogs are + # named by the bundle rather than by the value, and rules carry no name at + # all, so each has its own collision rule -- or none -- and its own + # installer. + for field_name in ( + "udtfs", + "table_providers", + "catalog_providers", + "physical_optimizer_rules", + ): declared[field_name] = [] for position, extension in enumerate(extensions): if not isinstance(extension, SessionComponentsExportable): @@ -2327,6 +2333,18 @@ def with_extensions( ) ] ) + # Bound to `new` for the same reason: the catalog getter is handed the + # logical codec its provider will serialize through. Unlike a table a + # catalog may replace one the session holds, so only names claimed + # twice within the call are refused. + resolved_catalogs = new.ctx._resolve_extension_catalogs( + [ + pair + for _, _, pair in _reject_repeated_names( + declared["catalog_providers"], "catalog" + ) + ] + ) # Rules accumulate, so there is no name to check and nothing to refuse # -- only the capsules to import while failing is still free. resolved_rules = _resolve_declared_rules( @@ -2353,6 +2371,7 @@ def with_extensions( [function._udaf for function in resolved["udafs"]], [function._udwf for function in resolved["udwfs"]], [table_function._udtf for table_function in resolved_udtfs], + resolved_catalogs, resolved_rules, ) return new diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 8e89c5c43..020d36eef 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -381,6 +381,19 @@ class SessionExtensionComponents: replaced: see :ref:`extension_bundles_collisions`. """ + catalog_providers: tuple[tuple[str, Any], ...] = _components("catalog") + """Catalogs to register, as ``(name, provider)`` pairs. + + Anything + :py:meth:`~datafusion.context.SessionContext.register_catalog_provider` + accepts. Bound to the finished context like :py:attr:`table_providers`. + + Two extensions claiming one name in the same call is refused. Replacing a + catalog the session already has is not — the default ``datafusion`` catalog + always exists, and swapping it is the usual way a library backs a session + with its own metadata. + """ + physical_optimizer_rules: tuple[PhysicalOptimizerRuleExportable, ...] = _components( "optimizer rule" ) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index b160b5cae..4c6d24b63 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -22,6 +22,7 @@ import shutil from dataclasses import fields +import datafusion.catalog import pyarrow as pa import pyarrow.compute as pc import pyarrow.dataset as ds @@ -1843,6 +1844,124 @@ def test_with_extensions_rejects_a_table_that_is_not_a_table_by_name(ctx): ctx.udf("double") +class _Schema(datafusion.catalog.SchemaProvider): + """One table, enough to prove a declared catalog is reachable from SQL.""" + + def __init__(self, table): + self.tables = {"t": table} + + def table_names(self) -> set[str]: + return set(self.tables) + + def register_table(self, name, table): + self.tables[name] = table + + def deregister_table(self, name, cascade: bool = True): + del self.tables[name] + + def table(self, name): + return self.tables.get(name) + + def table_exist(self, name) -> bool: + return name in self.tables + + +class _Catalog(datafusion.catalog.CatalogProvider): + def __init__(self, table): + self.schemas = {"s": _Schema(table)} + + def schema_names(self) -> set[str]: + return set(self.schemas) + + def schema(self, name): + return self.schemas.get(name) + + def register_schema(self, name, schema): + self.schemas[name] = schema + + def deregister_schema(self, name, cascade: bool): + del self.schemas[name] + + +class _CatalogExtension: + """Contributes catalogs as ``(name, provider)`` pairs.""" + + def __init__(self, catalog_providers=()): + self._catalog_providers = catalog_providers + + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents(catalog_providers=self._catalog_providers) + + +def test_with_extensions_registers_a_declared_catalog(ctx): + """A declared catalog is queryable through its qualified name.""" + table = ctx.from_pydict({"a": [1, 2, 3]}).into_view() + + result = ctx.with_extensions( + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)) + ) + + assert "engine" in result.catalog_names() + assert ( + result.sql("SELECT sum(a) FROM engine.s.t").collect()[0].column(0)[0].as_py() + == 6 + ) + + +def test_with_extensions_rejects_a_catalog_two_extensions_claim(ctx): + """One name, two bundles: refused with both named.""" + table = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(ValueError, match=r"catalog named 'engine'"): + ctx.with_extensions( + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)), + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)), + ) + + assert "engine" not in ctx.catalog_names() + + +def test_with_extensions_allows_replacing_an_existing_catalog(ctx): + """Replacing a catalog the session holds is ordinary, unlike a table. + + ``register_catalog`` returns whichever provider it displaced rather than + refusing, and the default ``datafusion`` catalog always exists — so a + library backing a session with its own metadata has to be able to do this. + """ + table = ctx.from_pydict({"a": [1, 2, 3]}).into_view() + + result = ctx.with_extensions( + _CatalogExtension(catalog_providers=(("datafusion", _Catalog(table)),)) + ) + + assert ( + result.sql("SELECT sum(a) FROM datafusion.s.t") + .collect()[0] + .column(0)[0] + .as_py() + == 6 + ) + + +def test_with_extensions_registers_no_catalog_when_a_later_hook_raises(ctx): + """Catalogs share the transaction, even though they write to a shared list.""" + + class BoomPlanner: + def __datafusion_session_planner__(self, ctx, fallback): + msg = "boom" + raise RuntimeError(msg) + + table = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions( + _CatalogExtension(catalog_providers=(("engine", _Catalog(table)),)), + BoomPlanner(), + ) + + assert "engine" not in ctx.catalog_names() + + def test_session_extension_components_rejects_a_single_optimizer_rule(): """The same for rules, naming what that field holds.""" with pytest.raises( @@ -1925,7 +2044,7 @@ def test_every_component_field_has_an_installer(): 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 four fields named here say which of them + ``_FUNCTION_KINDS`` and the fields named here 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. """ @@ -1942,6 +2061,7 @@ def test_every_component_field_has_an_installer(): "function": {kind.field for kind in _FUNCTION_KINDS}, "table function": {"udtfs"}, "table": {"table_providers"}, + "catalog": {"catalog_providers"}, "optimizer rule": {"physical_optimizer_rules"}, } diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 3a3eb3ca8..f0c86a401 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -46,6 +46,8 @@ # Tables, split the same way: import the provider and resolve the # destination schema during resolution, insert at commit. "_resolve_extension_tables", + # Catalogs, sharing their import half with register_catalog_provider. + "_resolve_extension_catalogs", } ) From 004b95a3bd77ed755e6f8e119f967b2bc7b41000 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 12:04:52 -0400 Subject: [PATCH 2/4] fix: name the declared catalog when its capsule will not import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The same repair the tables got: a resolve failure surfaced the importer's message, which names neither the catalog nor the bundle. Wrapping it with the declared name — unique within the call — points at one declaration. For a catalog the failure that remains is a capsule of the wrong kind: a non-capsule return falls through to the duck-typed wrap, the same fall-through every getter in the family has. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 10 +++++++++- python/datafusion/context.py | 14 +++++++------- python/tests/test_context.py | 25 +++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 8 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 3e8d89587..6dfa61655 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1924,7 +1924,15 @@ impl PySessionContext { ) -> PyDataFusionResult { let catalogs = catalogs .into_iter() - .map(|(name, provider)| Ok((name, self.resolve_catalog_provider(provider)?))) + .map(|(name, provider)| { + // The name identifies the culprit declaration, as for tables: + // a getter that returns junk fails in the importer, which + // knows neither the catalog nor the bundle. + let provider = self.resolve_catalog_provider(provider).map_err(|err| { + exec_datafusion_err!("Resolving the declared catalog {name}: {err}") + })?; + Ok((name, provider)) + }) .collect::>>()?; Ok(PyResolvedCatalogs { catalogs }) } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 4ea518220..e8e792eee 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2235,17 +2235,17 @@ def with_extensions( a declared function or optimizer rule does not expose its capsule getter and is not already a wrapper. ValueError: If two codecs claim the same id, if two extensions - declare a function, table, or table 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. + declare a function, table, table function, or catalog 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. RuntimeError: If a getter is present but returns something that is not a ``PyCapsule`` at all. The message comes from the importer and does not name the bundle, because by then the declaration has already been accepted as the right shape. - Exception: If a declared table cannot be resolved — the name is - already registered, two declarations resolve to one table, the - schema is unknown, or the value is not a table. + Exception: If a declared table or catalog cannot be resolved — + name taken, two declarations resolving to one table, unknown + schema, value not a table, or a capsule of the wrong kind. Examples: The returned handle is a different object sharing one session, and diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 4c6d24b63..ef397fbf1 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1962,6 +1962,31 @@ def __datafusion_session_planner__(self, ctx, fallback): assert "engine" not in ctx.catalog_names() +def test_with_extensions_rejects_a_catalog_capsule_of_the_wrong_kind(ctx): + """A catalog whose getter returns the wrong capsule is refused by name. + + A non-capsule return is not refusable — any object can be a catalog + provider, so it falls through to the duck-typed wrap, like a table getter's + falls through to the dataset one. A capsule of the wrong *kind* is the + resolve failure that remains, and the declared name — unique within the + call — is what points at the culprit declaration, as for tables. + """ + + class WrongCapsuleCatalog: + def __datafusion_catalog_provider__(self, codec): + return ctx.__datafusion_query_planner__() + + with pytest.raises(Exception, match=r"declared catalog engine"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + _CatalogExtension(catalog_providers=(("engine", WrongCapsuleCatalog()),)), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + assert "engine" not in ctx.catalog_names() + + def test_session_extension_components_rejects_a_single_optimizer_rule(): """The same for rules, naming what that field holds.""" with pytest.raises( From ce9454a69967ec6a90d8709c92c493fb416affce Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 12:05:15 -0400 Subject: [PATCH 3/4] style: freeze the resolved-catalogs carrier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing mutates it between resolve and commit, so there is no reason to pay for the runtime borrow flag a mutable pyclass carries — the same reasoning that froze PhysicalOptimizerRules and ResolvedTables. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 6dfa61655..0268b523e 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1980,8 +1980,9 @@ struct ResolvedTable { /// Catalog providers imported for a `with_extensions` call. /// -/// Opaque to Python, like [`PyResolvedTables`] and [`PyPhysicalOptimizerRules`]. -#[pyclass(name = "ResolvedCatalogs", module = "datafusion._internal")] +/// Opaque to Python, like [`PyResolvedTables`] and [`PyPhysicalOptimizerRules`], +/// and `frozen` for the same reason as both: the commit only reads. +#[pyclass(frozen, name = "ResolvedCatalogs", module = "datafusion._internal")] pub struct PyResolvedCatalogs { catalogs: Vec<(String, Arc)>, } From 78332d3f6efcf9c655c4751a155f3393deaaf14c Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 12:05:24 -0400 Subject: [PATCH 4/4] docs: add catalogs to the commit-order contract Catalogs resolve alongside everything else and commit with the rest; register_catalog replaces rather than refuses, so their commit step cannot fail and needs no exception of its own. Co-Authored-By: Claude Fable 5 --- docs/source/contributor-guide/ffi-internals.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index a225f53f6..da003facc 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -129,12 +129,13 @@ A call therefore splits into a part that may fail and a part that may not: 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, every declared table has its provider imported and its destination schema - resolved, every declared physical optimizer rule has its capsule imported, - and every `__datafusion_session_planner__` runs against the completed - chains. + resolved, every declared catalog has its provider imported, every declared + physical optimizer rule has its capsule imported, and every + `__datafusion_session_planner__` runs against the completed chains. 4. **Commit.** The tables are inserted, the planner is bound, the functions - are registered, and the optimizer rules are installed in a single - `SessionState` rebuild. + are registered, the catalogs are registered — `register_catalog` replaces + rather than refuses, so it cannot fail — and the optimizer rules are + installed in a single `SessionState` rebuild. Only step 4 touches the session, and every step that can fail happens before it — with one honest exception: a table insert goes through a