Skip to content

Commit 04a9650

Browse files
timsaucerclaude
timsaucer
and
claude
committed
feat: let extension bundles declare catalog providers
`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) <[email protected]>
1 parent fa22701 commit 04a9650

11 files changed

Lines changed: 371 additions & 36 deletions

File tree

‎.ai/skills/ffi-capsule-protocol/SKILL.md‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,25 @@ receiver had — the same session, and the same task-context provider, but not
7676
this call's codecs, not even your own. Read the host's codec chains in the
7777
planner hook, never in the extension hook.
7878

79+
**That is why a bundle declares unresolved components, not wrapped ones.** The
80+
components it returns split by what their getter asks for:
81+
82+
- Getters taking no argument — the three function kinds,
83+
`__datafusion_physical_optimizer_rule__` — have nothing session-scoped to
84+
bind, so a bundle may hand over either the raw exportable or an
85+
already-wrapped object.
86+
- Getters taking the session or a codec — `__datafusion_table_function__`,
87+
`__datafusion_table_provider__`, `__datafusion_catalog_provider__` — must be
88+
handed over **unwrapped**, with a name. Wrapping one inside the components
89+
hook would call its getter with the `ctx` that hook received, capturing a
90+
chain missing every library in the call. The host wraps these itself, against
91+
the handle carrying the final chains, which is the only place that chain
92+
exists.
93+
94+
`RecordingTableFunction` in `examples/datafusion-ffi-example/src/extension.rs`
95+
records the ids it was resolved against, so the difference is asserted rather
96+
than described.
97+
7998
A *codec* must always be handed over as an object implementing its getter, never
8099
as the bare capsule the getter returns; `with_extensions` refuses a capsule.
81100
A codec's wire id — the string a payload names on decode, which has to mean the
@@ -187,6 +206,18 @@ an instruction to derive one first. It is not: the factories are handed the
187206
receiver, and the returned handle shares its allocation. There is nothing to
188207
keep alive separately and nothing to garbage-collect out from under a provider.
189208

209+
It is also the one place where sharing an allocation has a cost, and the cost
210+
shapes how a component is added to it. Because the returned handle *is* the
211+
receiver's session, a failure part-way through has nothing to roll back to. So
212+
`with_extensions` does every fallible thing first — importing capsules,
213+
resolving names, running the planner hooks — and only then writes. **Adding a
214+
new kind of component means adding a resolve step, never a fallible commit
215+
step:** a `_resolve_extension_*` that returns an opaque carrier and an
216+
`_install_extension_*` that takes it and returns `()`. The one exception is
217+
table registration, whose insert goes through a `SchemaProvider` that a foreign
218+
library may implement; it is committed first so nothing else is written behind
219+
it. Do not add a second exception without the same justification.
220+
190221
`SessionContext.enable_url_table` is the one method that mints a second
191222
allocation for a session. Its result must not outlive the receiver, and it also
192223
forks the session's `SessionState` while keeping its id, so two handles report

‎crates/core/src/context.rs‎

Lines changed: 90 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -842,35 +842,9 @@ impl PySessionContext {
842842
pub fn register_catalog_provider(
843843
&self,
844844
name: &str,
845-
mut provider: Bound<'_, PyAny>,
845+
provider: Bound<'_, PyAny>,
846846
) -> PyDataFusionResult<()> {
847-
if provider.hasattr("__datafusion_catalog_provider__")? {
848-
let py = provider.py();
849-
let ffi = self.ffi_logical_codec();
850-
let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
851-
provider = call_capsule_getter(
852-
provider,
853-
"__datafusion_catalog_provider__",
854-
CapsuleGetterArg::LogicalCodec(&codec_capsule),
855-
)?;
856-
}
857-
858-
let provider = if let Ok(capsule) = provider.cast::<PyCapsule>() {
859-
let data: NonNull<FFI_CatalogProvider> = capsule
860-
.pointer_checked(Some(c"datafusion_catalog_provider"))?
861-
.cast();
862-
let provider = unsafe { data.as_ref() };
863-
let provider: Arc<dyn CatalogProvider> = provider.into();
864-
provider
865-
} else {
866-
match provider.extract::<PyCatalog>() {
867-
Ok(py_catalog) => py_catalog.catalog,
868-
Err(_) => Arc::new(RustWrappedPyCatalogProvider::new(
869-
provider.into(),
870-
self.ffi_logical_codec(),
871-
)) as Arc<dyn CatalogProvider>,
872-
}
873-
};
847+
let provider = self.resolve_catalog_provider(provider)?;
874848

875849
let _ = self.ctx.register_catalog(name, provider);
876850

@@ -1828,6 +1802,42 @@ impl PySessionContext {
18281802
Ok(())
18291803
}
18301804

1805+
/// Resolve the catalogs a `with_extensions` call declared.
1806+
///
1807+
/// The fallible half. Each provider is imported against `self` — the handle
1808+
/// carrying the completed codec chains, since
1809+
/// `__datafusion_catalog_provider__` is handed the logical codec it will
1810+
/// serialize through.
1811+
///
1812+
/// No name is refused here. `register_catalog` replaces rather than
1813+
/// rejects, and `datafusion` — the default catalog — always exists, so a
1814+
/// bundle replacing a catalog is ordinary rather than a mistake. Two
1815+
/// bundles claiming one name in the same call is refused on the Python
1816+
/// side, where both can be named.
1817+
///
1818+
/// **Writes nothing.**
1819+
pub fn _resolve_extension_catalogs<'py>(
1820+
&self,
1821+
catalogs: Vec<(String, Bound<'py, PyAny>)>,
1822+
) -> PyDataFusionResult<PyResolvedCatalogs> {
1823+
let catalogs = catalogs
1824+
.into_iter()
1825+
.map(|(name, provider)| Ok((name, self.resolve_catalog_provider(provider)?)))
1826+
.collect::<PyDataFusionResult<Vec<_>>>()?;
1827+
Ok(PyResolvedCatalogs { catalogs })
1828+
}
1829+
1830+
/// Commit the catalogs for a `with_extensions` call.
1831+
///
1832+
/// Nothing here can fail: the providers were imported by
1833+
/// [`Self::_resolve_extension_catalogs`], and `register_catalog` returns
1834+
/// whichever provider it displaced rather than refusing.
1835+
pub fn _install_extension_catalogs(&self, resolved: PyRef<'_, PyResolvedCatalogs>) {
1836+
for (name, provider) in &resolved.catalogs {
1837+
let _ = self.ctx.register_catalog(name, Arc::clone(provider));
1838+
}
1839+
}
1840+
18311841
/// Import the physical optimizer rules a `with_extensions` call declared.
18321842
///
18331843
/// The fallible half of installing them, run while the call can still fail
@@ -1900,6 +1910,14 @@ struct ResolvedTable {
19001910
provider: Arc<dyn TableProvider>,
19011911
}
19021912

1913+
/// Catalog providers imported for a `with_extensions` call.
1914+
///
1915+
/// Opaque to Python, like [`PyResolvedTables`] and [`PyPhysicalOptimizerRules`].
1916+
#[pyclass(name = "ResolvedCatalogs", module = "datafusion._internal")]
1917+
pub struct PyResolvedCatalogs {
1918+
catalogs: Vec<(String, Arc<dyn CatalogProvider>)>,
1919+
}
1920+
19031921
/// Physical optimizer rules imported for a `with_extensions` call.
19041922
///
19051923
/// Opaque to Python, and deliberately not added to the module: it exists only
@@ -1912,6 +1930,50 @@ pub struct PyPhysicalOptimizerRules {
19121930
}
19131931

19141932
impl PySessionContext {
1933+
/// Turn whatever a caller offered as a catalog provider into one.
1934+
///
1935+
/// The fallible half of registering a catalog, shared by
1936+
/// [`Self::register_catalog_provider`] and
1937+
/// [`Self::_resolve_extension_catalogs`] so both accept exactly the same
1938+
/// shapes: an object exposing `__datafusion_catalog_provider__`, a bare
1939+
/// capsule, a [`PyCatalog`], or a Python object implementing the provider
1940+
/// interface.
1941+
///
1942+
/// The getter is handed **this context's** logical codec, so which handle
1943+
/// this is called on decides what the provider will serialize through.
1944+
fn resolve_catalog_provider(
1945+
&self,
1946+
mut provider: Bound<'_, PyAny>,
1947+
) -> PyDataFusionResult<Arc<dyn CatalogProvider>> {
1948+
if provider.hasattr("__datafusion_catalog_provider__")? {
1949+
let py = provider.py();
1950+
let ffi = self.ffi_logical_codec();
1951+
let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?;
1952+
provider = call_capsule_getter(
1953+
provider,
1954+
"__datafusion_catalog_provider__",
1955+
CapsuleGetterArg::LogicalCodec(&codec_capsule),
1956+
)?;
1957+
}
1958+
1959+
Ok(if let Ok(capsule) = provider.cast::<PyCapsule>() {
1960+
let data: NonNull<FFI_CatalogProvider> = capsule
1961+
.pointer_checked(Some(c"datafusion_catalog_provider"))?
1962+
.cast();
1963+
let provider = unsafe { data.as_ref() };
1964+
let provider: Arc<dyn CatalogProvider> = provider.into();
1965+
provider
1966+
} else {
1967+
match provider.extract::<PyCatalog>() {
1968+
Ok(py_catalog) => py_catalog.catalog,
1969+
Err(_) => Arc::new(RustWrappedPyCatalogProvider::new(
1970+
provider.into(),
1971+
self.ffi_logical_codec(),
1972+
)) as Arc<dyn CatalogProvider>,
1973+
}
1974+
})
1975+
}
1976+
19151977
/// Write the session's query planner, in place.
19161978
///
19171979
/// Pass `Some(planner)` to install one, or `None` to rebuild whichever

‎docs/source/extension-guide/bundles.md‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -336,13 +336,13 @@ where the host can resolve them:
336336
- **Getters taking no argument** — the three function kinds and physical
337337
optimizer rules. Nothing is session-scoped, so a bundle may hand over either
338338
a wrapped object or the raw exportable.
339-
- **Getters taking the session or a codec** — table functions and table
340-
providers. These are resolved by the host against the *finished* handle, in
341-
step 3, which is why you hand over the unwrapped value and a name rather than
342-
a {py:class}`~datafusion.user_defined.TableFunction` you built yourself.
343-
Wrapping one inside your components hook binds it to the context that hook
344-
received, which has none of the call's codecs — so it would capture a chain
345-
missing every library in the call, including your own.
339+
- **Getters taking the session or a codec** — table functions, table providers,
340+
and catalog providers. These are resolved by the host against the *finished*
341+
handle, in step 3, which is why you hand over the unwrapped value and a name
342+
rather than a {py:class}`~datafusion.user_defined.TableFunction` you built
343+
yourself. Wrapping one inside your components hook binds it to the context
344+
that hook received, which has none of the call's codecs — so it would capture
345+
a chain missing every library in the call, including your own.
346346

347347
(extension_bundles_collisions)=
348348

‎docs/source/extension-guide/table-providers.md‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,20 @@ wrapped: `__datafusion_table_provider__` takes the session, and the one your
5151
components hook receives has none of the call's codecs yet. The host resolves
5252
it against the finished handle. See {ref}`extension_bundles_binding`.
5353

54+
Catalogs work the same way, as `catalog_providers`:
55+
56+
```python
57+
return SessionExtensionComponents(catalog_providers=(("engine", MyCatalog()),))
58+
```
59+
60+
with one difference worth knowing. A declared **table** name that is already
61+
registered is an error, because DataFusion refuses a duplicate table rather
62+
than replacing it. A **catalog** name is not: `register_catalog` returns
63+
whichever provider it displaced, and the default `datafusion` catalog always
64+
exists — so replacing one is the usual way a library backs a session with its
65+
own metadata. Only two bundles claiming the same catalog name in one call is
66+
refused.
67+
5468
Start with a table provider. Reach for the schema and catalog levels when your
5569
data source has its own namespace that should be browsable rather than
5670
registered table by table, and for the provider list only when your library is

‎examples/datafusion-ffi-example/python/tests/_test_session_extension.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import pytest
2424
from datafusion import SessionContext, SessionExtensionComponents
2525
from datafusion_ffi_example import (
26+
MyCatalogExtension,
2627
MyDataExtension,
2728
MyFunctionExtension,
2829
MyLogicalExtensionCodec,
@@ -237,6 +238,42 @@ def test_a_table_name_already_registered_is_refused():
237238
ctx.udf("my_custom_is_null")
238239

239240

241+
def test_a_declared_catalog_is_queryable():
242+
"""A catalog declared by a bundle is reachable by its qualified name."""
243+
ctx = SessionContext().with_extensions(MyCatalogExtension())
244+
245+
assert "declared_catalog" in ctx.catalog_names()
246+
result = ctx.sql("SELECT * FROM declared_catalog.my_schema.my_table").collect()
247+
assert result[0].num_rows > 0
248+
249+
250+
def test_four_libraries_install_in_one_call():
251+
"""The whole point, across a real FFI boundary.
252+
253+
Four independently declared bundles — functions, rules, a table and a table
254+
function, a catalog — in one call, and a single query that touches three of
255+
them while the fourth counts the planning it did.
256+
"""
257+
rules = MyRuleExtension()
258+
ctx = SessionContext().with_extensions(
259+
MyFunctionExtension(),
260+
rules,
261+
MyDataExtension(),
262+
MyCatalogExtension(),
263+
)
264+
265+
result = ctx.sql(
266+
'SELECT my_custom_is_null("A") AS n FROM declared_table '
267+
"UNION ALL "
268+
"SELECT my_custom_is_null(units) AS n "
269+
"FROM declared_catalog.my_schema.my_table"
270+
).collect()
271+
272+
assert sum(batch.num_rows for batch in result) > 0
273+
assert rules.first_calls() > 0
274+
assert rules.second_calls() > 0
275+
276+
240277
def test_the_hook_returns_the_components_type():
241278
"""The bundle builds a real dataclass, not a duck-typed stand-in.
242279

‎examples/datafusion-ffi-example/src/extension.rs‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use pyo3::prelude::*;
2121
use pyo3::types::{PyCapsule, PyDict};
2222

2323
use crate::aggregate_udf::MySumUDF;
24+
use crate::catalog_provider::MyCatalogProvider;
2425
use crate::physical_optimizer::MyPhysicalOptimizerRule;
2526
use crate::scalar_udf::IsNullUDF;
2627
use crate::table_function::MyTableFunction;
@@ -230,3 +231,43 @@ impl MyDataExtension {
230231
components.call((), Some(&kwargs))
231232
}
232233
}
234+
235+
/// A bundle contributing a catalog.
236+
///
237+
/// `__datafusion_catalog_provider__` takes the session and pulls the host's
238+
/// logical codec off it, so like a table provider it is handed over unresolved
239+
/// and the host binds it to the finished handle.
240+
#[pyclass(
241+
from_py_object,
242+
name = "MyCatalogExtension",
243+
module = "datafusion_ffi_example",
244+
subclass
245+
)]
246+
#[derive(Debug, Clone, Default)]
247+
pub(crate) struct MyCatalogExtension {}
248+
249+
#[pymethods]
250+
impl MyCatalogExtension {
251+
#[new]
252+
fn new() -> Self {
253+
Self {}
254+
}
255+
256+
fn __datafusion_session_components__<'py>(
257+
&self,
258+
py: Python<'py>,
259+
ctx: Bound<'py, PyAny>,
260+
) -> PyResult<Bound<'py, PyAny>> {
261+
let _ = ctx;
262+
263+
let components = py
264+
.import("datafusion")?
265+
.getattr("SessionExtensionComponents")?;
266+
let kwargs = PyDict::new(py);
267+
kwargs.set_item(
268+
"catalog_providers",
269+
(("declared_catalog", Py::new(py, MyCatalogProvider::new()?)?),),
270+
)?;
271+
components.call((), Some(&kwargs))
272+
}
273+
}

‎examples/datafusion-ffi-example/src/lib.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use pyo3::prelude::*;
2020
use crate::aggregate_udf::MySumUDF;
2121
use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList};
2222
use crate::config::MyConfig;
23-
use crate::extension::{MyDataExtension, MyFunctionExtension, MyRuleExtension};
23+
use crate::extension::{MyCatalogExtension, MyDataExtension, MyFunctionExtension, MyRuleExtension};
2424
use crate::logical_extension_codec::MyLogicalExtensionCodec;
2525
use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec};
2626
use crate::physical_extension_codec::MyPhysicalExtensionCodec;
@@ -68,5 +68,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
6868
m.add_class::<MyFunctionExtension>()?;
6969
m.add_class::<MyRuleExtension>()?;
7070
m.add_class::<MyDataExtension>()?;
71+
m.add_class::<MyCatalogExtension>()?;
7172
Ok(())
7273
}

‎python/datafusion/context.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,7 @@ def _collect_contributions(
231231
"udwfs": [],
232232
"udtfs": [],
233233
"table_providers": [],
234+
"catalog_providers": [],
234235
"physical_optimizer_rules": [],
235236
}
236237
for extension in extensions:
@@ -2198,6 +2199,18 @@ def with_extensions(
21982199
)
21992200
]
22002201
)
2202+
# Bound to `new` for the same reason: the catalog getter is handed the
2203+
# logical codec its provider will serialize through. Unlike a table a
2204+
# catalog may replace one the session holds, so only names claimed
2205+
# twice within the call are refused.
2206+
resolved_catalogs = new.ctx._resolve_extension_catalogs(
2207+
[
2208+
pair
2209+
for _, pair in _reject_repeated_names(
2210+
declared["catalog_providers"], "catalog"
2211+
)
2212+
]
2213+
)
22012214
# Rules accumulate, so there is no name to check and nothing to refuse
22022215
# -- only the capsules to import while failing is still free.
22032216
resolved_rules = new.ctx._resolve_extension_physical_optimizer_rules(
@@ -2250,6 +2263,7 @@ def with_extensions(
22502263
new.register_udwf(function)
22512264
for table_function in resolved_udtfs:
22522265
new.register_udtf(table_function)
2266+
new.ctx._install_extension_catalogs(resolved_catalogs)
22532267
new.ctx._install_extension_physical_optimizer_rules(resolved_rules)
22542268
return new
22552269

0 commit comments

Comments
 (0)