diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 37be0bc9d..ab14dc732 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,18 @@ 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 an +`_install_extension_*` that takes it and returns `()`. 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 cc8d59079..fe7ea7623 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -27,8 +27,10 @@ use arrow::pyarrow::FromPyArrow; use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; -use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory}; -use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_err}; +use datafusion::catalog::{ + CatalogProvider, CatalogProviderList, SchemaProvider, TableProviderFactory, +}; +use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_datafusion_err, exec_err}; use datafusion::datasource::file_format::file_compression_type::FileCompressionType; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ @@ -840,35 +842,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); @@ -1764,6 +1740,104 @@ impl PySessionContext { Ok(()) } + /// Resolve the tables a `with_extensions` call declared. + /// + /// The fallible half. Each provider is imported against `slf` — the handle + /// carrying the completed codec chains, not the context the components hook + /// was given — and each name is resolved to the schema that will hold it. + /// A name already taken is refused here, because DataFusion refuses a + /// duplicate registration rather than replacing it, and a refusal is much + /// more useful before anything has been written. + /// + /// **Writes nothing.** + pub fn _resolve_extension_tables<'py>( + slf: &Bound<'py, Self>, + tables: Vec<(String, Bound<'py, PyAny>)>, + ) -> PyDataFusionResult { + let session = slf.clone().into_bound_py_any(slf.py())?; + let state = slf.borrow().ctx.state(); + + let mut resolved = Vec::with_capacity(tables.len()); + for (name, obj) in tables { + let provider = PyTable::new(obj, Some(session.clone()))?.table; + let reference = TableReference::from(name.as_str()); + let table_name = reference.table().to_owned(); + let schema = state.schema_for_ref(reference)?; + // Checked against the schema rather than against this call's own + // list, so a name the session already holds is caught too. Both + // are the same error to a caller. + if schema.table_exist(&table_name) { + return Err(exec_datafusion_err!( + "An extension declared a table named {name}, which is already registered" + ) + .into()); + } + resolved.push(ResolvedTable { + schema, + name: table_name, + provider, + }); + } + Ok(PyResolvedTables { tables: resolved }) + } + + /// Commit the tables for a `with_extensions` call. + /// + /// Runs first among the commit steps. Every name was resolved and found + /// free by [`Self::_resolve_extension_tables`], so the only way an insert + /// still fails is a foreign `SchemaProvider` refusing a registration it + /// reported as available — the one place in `with_extensions` that can + /// leave a call part-applied. Going first is what keeps the blast radius + /// to the tables themselves: no planner is bound and no function is + /// registered behind it. + pub fn _install_extension_tables( + &self, + resolved: PyRef<'_, PyResolvedTables>, + ) -> PyDataFusionResult<()> { + for table in &resolved.tables { + table + .schema + .register_table(table.name.clone(), Arc::clone(&table.provider))?; + } + Ok(()) + } + + /// 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 }) + } + + /// Commit the catalogs for a `with_extensions` call. + /// + /// Nothing here can fail: the providers were imported by + /// [`Self::_resolve_extension_catalogs`], and `register_catalog` returns + /// whichever provider it displaced rather than refusing. + pub fn _install_extension_catalogs(&self, resolved: PyRef<'_, PyResolvedCatalogs>) { + for (name, provider) in &resolved.catalogs { + let _ = self.ctx.register_catalog(name, Arc::clone(provider)); + } + } + /// Import the physical optimizer rules a `with_extensions` call declared. /// /// The fallible half of installing them, run while the call can still fail @@ -1818,6 +1892,32 @@ impl PySessionContext { } } +/// Tables resolved for a `with_extensions` call. +/// +/// Opaque to Python, and deliberately not added to the module, like +/// [`PyPhysicalOptimizerRules`]. Each entry is a provider that has already been +/// imported and a schema that has already been looked up, so committing is an +/// insert into a resolved destination rather than a fresh name resolution. +#[pyclass(name = "ResolvedTables", module = "datafusion._internal")] +pub struct PyResolvedTables { + tables: Vec, +} + +/// One entry of [`PyResolvedTables`]: where it goes, and what goes there. +struct ResolvedTable { + schema: Arc, + name: String, + 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 @@ -1830,6 +1930,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 57ee1da0b..8a97205dc 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -55,7 +55,7 @@ class MyEngineExtension: return self._make_planner(ctx, fallback=fallback) ``` -Implement only the hooks you need. Codecs and functions both go in +Implement only the hooks you need. Codecs, functions, and tables all 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 @@ -63,15 +63,15 @@ a library shipping nothing but an optimizing planner defines only ```python ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) -ctx.register_table("t", lib_a.TableProvider()) ``` -Return your functions rather than calling `register_udf` on the `ctx` you were -handed. Both put the function on the session, but a registration you make -inside the hook is written the moment it runs — before the other bundles have -been called, and not undone if one of them raises. What you declare is instead -resolved and checked while a failure still costs nothing, then written once -every bundle has succeeded. See {ref}`extension_bundles_transaction`. +Declare what you contribute rather than calling `register_udf` or +`register_table` on the `ctx` you were handed. Both put it on the session, but a +registration you make inside the hook is written the moment it runs — before the +other bundles have been called, too early to see their codecs, and not undone if +one of them raises. What you declare is instead resolved and checked while a +failure still costs nothing, then written once every bundle has succeeded. See +{ref}`extension_bundles_transaction`. `MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust implementation of the protocol, including taking the task-context provider @@ -290,6 +290,24 @@ 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_binding)= + +## What a component is resolved against + +Components split in two by what their capsule getter asks for, and it decides +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, 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. + (extension_bundles_collisions)= ## Two bundles claiming one name @@ -328,6 +346,12 @@ Physical optimizer rules are exempt from all of this: they accumulate rather than replace, so two bundles contributing one each is the normal case and there is nothing to refuse. See {doc}`other-components`. +Tables go the other way. DataFusion refuses a duplicate table registration +rather than replacing it, so a declared table name that is *already* on the +session is an error too — 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`. + Your caller cannot rename your function, so stay out of the way: prefix the names with something tied to your library. @@ -358,6 +382,14 @@ 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. +Tables are the single exception to committing last, and they are committed +first because of it. A declared table has its provider imported, its +destination schema resolved, and a name already taken refused, all while a +failure still costs nothing — but the insert itself goes through a +`SchemaProvider`, and a foreign one can still refuse what it reported as free. +Running that first means no planner is bound and no function is registered +behind it when it does. + 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 codec chains belong to the returned handle; the planner is installed on the diff --git a/docs/source/extension-guide/functions.md b/docs/source/extension-guide/functions.md index ee38090e2..35131dd1d 100644 --- a/docs/source/extension-guide/functions.md +++ b/docs/source/extension-guide/functions.md @@ -31,7 +31,7 @@ the same registration methods. | `__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` | — | +| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` | `udtfs`, as `(name, func)` | All four are implemented in [`datafusion-ffi-example`], one per file. The last column is the {py:class}`~datafusion.SessionExtensionComponents` field a @@ -128,6 +128,18 @@ Only literal expressions are supported as arguments. The Python side is described under {doc}`Table Functions <../user-guide/common-operations/udf-and-udfa>`. +Because that getter takes the session, a table function is declared on a bundle +as a `(name, func)` pair and the host wraps it — not as a +{py:class}`~datafusion.user_defined.TableFunction` you built, which would +capture the codec chain from before the call: + +```python +return SessionExtensionComponents(udtfs=(("expand", my_library.MyTableFunction()),)) +``` + +The name is given here rather than read off the capsule, which is the other way +this differs from the three above. See {ref}`extension_bundles_binding`. + ## Serializing functions A function that appears in a plan leaving the process has to be reconstructible diff --git a/docs/source/extension-guide/table-providers.md b/docs/source/extension-guide/table-providers.md index 8f5a0ee08..9e09e6a85 100644 --- a/docs/source/extension-guide/table-providers.md +++ b/docs/source/extension-guide/table-providers.md @@ -39,6 +39,32 @@ A schema provider is the one that does not register on the session: you reach a {py:meth}`SessionContext.catalog `, and register the schema on that. +If your library ships a table alongside anything else, declare it on your +bundle as `table_providers` and let one call install everything: + +```python +return SessionExtensionComponents(table_providers=(("events", MyProvider()),)) +``` + +Hand over the provider itself, not a {py:class}`~datafusion.catalog.Table` you +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/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index 0d02ad78b..e758bd525 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -63,9 +63,9 @@ 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: +**Tables and functions arrive by whichever route their library chose.** A +library offering one or two hands you the objects themselves, and you register +each — a table as above, a function after wrapping it: ```python from datafusion import udf @@ -75,8 +75,8 @@ 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. +provides — and the library picks the names — leaving nothing per-item 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 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 58a205497..bf1c36535 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py +++ b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py @@ -22,7 +22,13 @@ import pyarrow as pa import pytest from datafusion import SessionContext, SessionExtensionComponents -from datafusion_ffi_example import MyFunctionExtension, MyRuleExtension +from datafusion_ffi_example import ( + MyCatalogExtension, + MyDataExtension, + MyFunctionExtension, + MyLogicalExtensionCodec, + MyRuleExtension, +) def _session(): @@ -184,6 +190,90 @@ def __datafusion_session_planner__(self, ctx, fallback) -> None: assert rules.second_calls() == 0 +def test_declared_tables_and_table_functions_work(): + """Both arrive as ``(name, value)`` pairs and both are queryable.""" + ctx = SessionContext().with_extensions(MyDataExtension()) + + assert ctx.sql("SELECT * FROM declared_table").collect()[0].num_rows == 2 + assert ctx.sql("SELECT * FROM declared_function()").collect()[0].num_rows > 0 + + +class _CodecBundle: + """Contributes a codec, so a later bundle's chain is observably different.""" + + def __datafusion_session_components__(self, ctx) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(MyLogicalExtensionCodec(),) + ) + + +def test_a_declared_table_function_sees_the_finished_codec_chain(): + """The claim that makes declaring a table function worth doing. + + ``__datafusion_table_function__`` takes the session and pulls the host's + logical codec off it. A bundle wrapping one itself would hand it the + context the components hook received, which has none of the call's codecs — + so it would capture a chain missing every library in the call, including + the one contributed *after* it here. The host resolves it against the + finished handle instead, and this asserts the difference rather than + describing it. + """ + data = MyDataExtension() + ctx = SessionContext().with_extensions(data, _CodecBundle()) + + assert ctx.sql("SELECT * FROM declared_function()").collect()[0].num_rows > 0 + assert data.codec_ids_seen() == ctx.logical_extension_codec_ids() + assert data.codec_ids_seen() != [] + + +def test_a_table_name_already_registered_is_refused(): + """Tables cannot shadow, so a clash is caught before anything is written.""" + ctx = SessionContext() + ctx.from_pydict({"a": [1]}, name="declared_table") + + with pytest.raises(Exception, match=r"already registered"): + ctx.with_extensions(MyFunctionExtension(), MyDataExtension()) + + with pytest.raises(KeyError): + 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 4f6d114cc..b3b037aa5 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -15,12 +15,17 @@ // specific language governing permissions and limitations // under the License. -use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods}; +use std::sync::{Arc, Mutex}; + +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; +use crate::table_provider::MyTableProvider; use crate::window_udf::MyRankUDF; /// A bundle contributing this library's three functions in one install. @@ -121,3 +126,148 @@ impl MyRuleExtension { components.call((), Some(&kwargs)) } } + +/// A table function that records the codec chain it was handed. +/// +/// `__datafusion_table_function__` takes the session and pulls the host's +/// logical codec off it, so *which* session it is resolved against is +/// observable rather than a matter of taste. A bundle cannot wrap one itself: +/// the context its components hook receives has none of the call's codecs yet. +/// Recording the ids here is what lets a test assert the host resolved it +/// against the finished handle instead. +#[pyclass( + from_py_object, + name = "RecordingTableFunction", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone)] +pub(crate) struct RecordingTableFunction { + seen: Arc>>, + inner: MyTableFunction, +} + +#[pymethods] +impl RecordingTableFunction { + #[new] + fn new() -> Self { + Self { + seen: Arc::new(Mutex::new(Vec::new())), + inner: MyTableFunction::new(), + } + } + + /// The logical codec ids the session carried when this was resolved. + fn codec_ids_seen(&self) -> Vec { + self.seen.lock().map(|ids| ids.clone()).unwrap_or_default() + } + + fn __datafusion_table_function__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let ids: Vec = session + .call_method0("logical_extension_codec_ids")? + .extract()?; + if let Ok(mut seen) = self.seen.lock() { + *seen = ids; + } + self.inner.__datafusion_table_function__(py, session) + } +} + +/// A bundle contributing a table and a table function. +/// +/// Both are `(name, value)` pairs, because neither carries a name of its own +/// the way a scalar function's capsule does. +#[pyclass( + from_py_object, + name = "MyDataExtension", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone)] +pub(crate) struct MyDataExtension { + function: RecordingTableFunction, +} + +#[pymethods] +impl MyDataExtension { + #[new] + fn new() -> Self { + Self { + function: RecordingTableFunction::new(), + } + } + + /// The codec ids the declared table function was resolved against. + fn codec_ids_seen(&self) -> Vec { + self.function.codec_ids_seen() + } + + 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( + "table_providers", + (( + "declared_table", + Py::new(py, MyTableProvider::new(3, 2, 1))?, + ),), + )?; + kwargs.set_item( + "udtfs", + (("declared_function", Py::new(py, self.function.clone())?),), + )?; + 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 097c8ef87..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::{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; @@ -67,5 +67,7 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/table_function.rs b/examples/datafusion-ffi-example/src/table_function.rs index e653aeab1..846b3c72b 100644 --- a/examples/datafusion-ffi-example/src/table_function.rs +++ b/examples/datafusion-ffi-example/src/table_function.rs @@ -38,11 +38,11 @@ pub(crate) struct MyTableFunction {} #[pymethods] impl MyTableFunction { #[new] - fn new() -> Self { + pub(crate) fn new() -> Self { Self {} } - fn __datafusion_table_function__<'py>( + pub(crate) fn __datafusion_table_function__<'py>( &self, py: Python<'py>, session: Bound, diff --git a/python/datafusion/context.py b/python/datafusion/context.py index a29dd6fa2..1d0c3fcc0 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -216,6 +216,58 @@ class _FunctionKind(NamedTuple): """Every kind of function a bundle can declare, in registration order.""" +def _reject_repeated_names( + declared: list[tuple[int, object, Any]], + kind: str, +) -> list[tuple[int, object, tuple[str, Any]]]: + """Refuse a name two extensions both declared, for the pair-shaped fields. + + The sibling of the name check in :py:func:`_resolve_declared_functions`, + for components whose name is given alongside the value rather than read + off it. Only names declared *in this call* are compared; whether claiming + one the session already holds is allowed is the component's own business, + and for tables it is checked when the name is resolved. + + Args: + declared: ``(position, extension, (name, value))`` triples in + declaration order, where ``position`` indexes the + ``with_extensions`` argument list. + kind: What to call this sort of component in an error. + + Returns: + ``declared`` unchanged, so this reads as a step rather than a check. + + Raises: + ValueError: If two declarations claim the same name. + """ + claimed: dict[str, tuple[int, object]] = {} + for position, extension, pair in declared: + name = pair[0] + if name in claimed: + # Keyed on position for the same reason as in + # `_resolve_declared_functions`: the two cases have different + # remedies, and only one of them is the caller's. + claimed_at, claimed_by = claimed[name] + if claimed_at == position: + msg = ( + f"{extension!r} declares two {kind}s named {name!r}. " + "Registrations have no fall-through, so the second would " + "silently replace the first; rename one of them." + ) + else: + msg = ( + f"Two extensions declare a {kind} named {name!r}: " + f"argument {claimed_at} ({claimed_by!r}) and argument " + f"{position} ({extension!r}). Registrations have no " + "fall-through, so one would silently replace the other; " + "install them on separate sessions, or drop the repeat if " + "one extension was passed twice." + ) + raise ValueError(msg) + claimed[name] = (position, extension) + return declared + + def _collect_contributions( extensions: tuple[object, ...], ctx: SessionContext, @@ -236,9 +288,9 @@ def _collect_contributions( context derived from it. Returns: - The logical codecs, the physical codecs, and the declared functions and - optimizer rules as ``(position, extension, declaration)`` triples, keyed - by the ``SessionExtensionComponents`` field they arrived in. + The logical codecs, the physical codecs, and every other declared + component as ``(position, extension, declaration)`` triples, keyed by + the ``SessionExtensionComponents`` field they arrived in. Raises: TypeError: If an argument implements neither hook, or a hook returns @@ -260,9 +312,17 @@ def _collect_contributions( declared: dict[str, list[tuple[int, object, Any]]] = { kind.field: [] for kind in _FUNCTION_KINDS } - # Rules are not a function kind: they accumulate rather than replace, so - # they carry no collision rule and install through their own primitive. - declared["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): continue @@ -2111,9 +2171,10 @@ def with_extensions( Nothing is written to the session until every hook has returned and every component has been validated, so a hook that raises leaves the - session as it was. Declared functions and optimizer rules install after - the planner is bound, and are visible on every handle sharing this - session. A hook that *mutates* the context it is handed — registering a + session as it was. Declared tables install first and everything else + after the planner is bound; all of them are visible on every handle + sharing this session. A hook that *mutates* the context it is handed — + registering a table, say — is not rolled back, which is why bundle objects must be configuration-only. @@ -2208,6 +2269,38 @@ def with_extensions( ) for kind in _FUNCTION_KINDS ] + # Tables and table functions are not in that table: their getters *do* + # take the session, so each is wrapped against `new`, not `self` -- the + # handle the components hook saw is missing this call's codecs. + resolved_udtfs = [ + _user_defined.TableFunction(name, func, new) + for _, _, (name, func) in _reject_repeated_names( + declared["udtfs"], "table function" + ) + ] + # Imports each provider against `new` for the same reason, resolves + # each name to the schema that will hold it, and refuses a name that is + # already taken. + resolved_tables = new.ctx._resolve_extension_tables( + [ + pair + for _, _, pair in _reject_repeated_names( + declared["table_providers"], "table" + ) + ] + ) + # 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 = new.ctx._resolve_extension_physical_optimizer_rules( @@ -2248,11 +2341,19 @@ def with_extensions( # part-way through has nothing to roll back to. The reasoning is in # docs/source/contributor-guide/ffi-internals.md, under "Why # `with_extensions` commits last". + # + # Tables are the one exception and so go first: a foreign schema + # provider can still refuse an insert it reported as free, and running + # it here means nothing else has been written when it does. + new.ctx._install_extension_tables(resolved_tables) if planner is not None or logical_codecs or physical_codecs: new.ctx._install_extension_planner(planner) for register, functions in resolved: for function in functions: register(function) + for table_function in resolved_udtfs: + new.register_udtf(table_function) + new.ctx._install_extension_catalogs(resolved_catalogs) new.ctx._install_extension_physical_optimizer_rules(resolved_rules) return new diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 403526cc8..01af69828 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -293,6 +293,53 @@ class SessionExtensionComponents: :py:func:`~datafusion.udwf`. """ + udtfs: tuple[tuple[str, Any], ...] = _components("table function") + """Table functions to register, as ``(name, function)`` pairs. + + Unlike the other three function kinds the name is **not** read off the + capsule, so it is given here. The value is an object exposing + ``__datafusion_table_function__``, or a plain Python callable. + + Pass the unwrapped value, not a + :py:class:`~datafusion.user_defined.TableFunction`. Wrapping calls the + capsule getter with the session, and the context a bundle is handed has not + had this call's codecs installed yet — so a wrapper built inside the hook + would be bound to the wrong chains. The host wraps these against the + finished context instead. See :ref:`extension_bundles_two_phases`. + + Collides by name like :py:attr:`udfs`. + """ + + table_providers: tuple[tuple[str, Any], ...] = _components("table") + """Tables to register, as ``(name, provider)`` pairs. + + Anything :py:meth:`~datafusion.context.SessionContext.register_table` + accepts: an object exposing ``__datafusion_table_provider__``, a + :py:class:`~datafusion.catalog.Table`, a + :py:class:`~datafusion.dataframe.DataFrame`, or a PyArrow dataset. Names may + be qualified (``"cat.schema.events"``); an unqualified one lands in the + session's default schema. + + Bound to the finished context for the same reason as :py:attr:`udtfs`. + + A name that is **already registered** is an error, here and in + ``register_table`` alike — DataFusion refuses a duplicate table rather than + replacing it, so unlike a function a table cannot shadow one. + """ + + 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 5740df5c4..94001a3ec 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 @@ -1670,6 +1671,207 @@ def test_session_extension_components_rejects_a_single_function(field): SessionExtensionComponents(**{field: _doubler()}) +class _TableExtension: + """Contributes tables and table functions, as ``(name, value)`` pairs.""" + + def __init__(self, table_providers=(), udtfs=()): + self._table_providers = table_providers + self._udtfs = udtfs + + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents( + table_providers=self._table_providers, udtfs=self._udtfs + ) + + +def test_with_extensions_registers_a_declared_table(ctx): + """A declared table is queryable on the returned handle.""" + provider = ctx.from_pydict({"a": [1, 2, 3]}).into_view() + + result = ctx.with_extensions(_TableExtension(table_providers=(("t", provider),))) + + assert result.sql("SELECT sum(a) FROM t").collect()[0].column(0)[0].as_py() == 6 + + +def test_with_extensions_registers_a_declared_udtf(ctx): + """A declared table function is callable from SQL. + + Declared as a ``(name, callable)`` pair rather than a built + ``TableFunction``, because wrapping one hands the getter a session and the + bundle does not have the right one yet. + """ + table = ctx.from_pydict({"a": [1, 2, 3]}).into_view() + + result = ctx.with_extensions(_TableExtension(udtfs=(("always", lambda: table),))) + + assert result.sql("SELECT a FROM always()").collect()[0].num_rows == 3 + + +def test_with_extensions_rejects_a_table_name_two_extensions_claim(ctx): + """Two bundles claiming one table name is refused before anything lands.""" + provider = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(ValueError, match=r"table named 'events'"): + ctx.with_extensions( + _TableExtension(table_providers=(("events", provider),)), + _TableExtension(table_providers=(("events", provider),)), + ) + + assert not ctx.table_exist("events") + + +def test_with_extensions_rejects_a_table_name_the_session_holds(ctx): + """A table cannot shadow one, the way a function can. + + DataFusion refuses a duplicate registration rather than replacing it, so + this is its rule rather than a policy chosen here — and catching it during + resolution is what keeps the rest of the call from being written first. + """ + ctx.from_pydict({"a": [1]}, name="events") + provider = ctx.from_pydict({"a": [2]}).into_view() + + with pytest.raises(Exception, match=r"already registered"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + _TableExtension(table_providers=(("events", provider),)), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + + +def test_with_extensions_rejects_a_table_in_an_unknown_schema(ctx): + """Resolving the destination happens before anything is written too.""" + provider = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(Exception, match=r"nope"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + _TableExtension(table_providers=(("nope.public.t", provider),)), + ) + + with pytest.raises(KeyError): + 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( @@ -1726,7 +1928,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``, the codec pair, and the rules 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. """ @@ -1741,6 +1943,9 @@ def test_every_component_field_has_an_installer(): assert by_noun == { "codec": {"logical_extension_codecs", "physical_extension_codecs"}, "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 903fd4a83..2bd4ec941 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -43,6 +43,13 @@ # the call can still fail without leaving the session half-installed. "_resolve_extension_physical_optimizer_rules", "_install_extension_physical_optimizer_rules", + # Tables, split the same way: import the provider and resolve the + # destination schema during resolution, insert at commit. + "_resolve_extension_tables", + "_install_extension_tables", + # Catalogs, sharing their import half with register_catalog_provider. + "_resolve_extension_catalogs", + "_install_extension_catalogs", } )