diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 554fc65d7..9a92c6055 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -27,8 +27,12 @@ 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, ResolvedTableReference, 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::{ @@ -1738,6 +1742,14 @@ impl PySessionContext { /// reasoning is in docs/source/contributor-guide/ffi-internals.md, under /// "Why `with_extensions` commits last". /// + /// The tables are the one exception and so go first. 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. + /// /// The planner is bound through this context's own `state_ref()`, so /// providers bound to it stay valid. With no planner supplied the bind /// still rebuilds whichever planner the session already holds against @@ -1758,9 +1770,11 @@ impl PySessionContext { extensions: Vec>, session: Bound<'py, PyAny>, rebind_planner: bool, + tables: PyRef<'_, PyResolvedTables>, udfs: Vec, udafs: Vec, udwfs: Vec, + udtfs: Vec, rules: PyRef<'_, PyPhysicalOptimizerRules>, ) -> PyDataFusionResult<()> { let py = slf.py(); @@ -1787,6 +1801,13 @@ impl PySessionContext { )?); } + // The first write. A foreign `SchemaProvider` refusing here leaves + // nothing else behind — see the tables exception above. + for table in &tables.tables { + table + .schema + .register_table(table.name.clone(), Arc::clone(&table.provider))?; + } if planner.is_some() || rebind_planner { slf.borrow().set_session_query_planner(planner); } @@ -1800,6 +1821,9 @@ impl PySessionContext { for udwf in udwfs { this.ctx.register_udwf(udwf.function); } + for udtf in udtfs { + this.register_udtf(udtf); + } // 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. @@ -1824,6 +1848,81 @@ 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 rather than left to the insert. + /// Whether a duplicate replaces or refuses is the `SchemaProvider`'s own + /// call — the in-memory one refuses — and asking it would mean asking at + /// commit time, once refusing costs something. Deciding here buys one rule + /// for every destination and a refusal while a failure is still free. + /// + /// Two declarations landing on one destination are refused here too. Both + /// checks are against the *resolved* reference rather than the declared + /// spelling, which is the only thing that answers the question: a name is + /// lowercased when it is parsed and filled out from the session's default + /// catalog and schema, so `Events`, `events` and `public.events` are one + /// table under three spellings. + /// + /// **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 catalog_options = &state.config_options().catalog; + let default_catalog = catalog_options.default_catalog.clone(); + let default_schema = catalog_options.default_schema.clone(); + + let mut claimed: HashMap = HashMap::new(); + let mut resolved = Vec::with_capacity(tables.len()); + for (name, obj) in tables { + // The name is the culprit's identity: it is unique within the call, + // so an import failure that carries it points at one declaration. + // The importer's own message names neither the table nor the + // bundle, because it never knew them. + let provider = PyTable::new(obj, Some(session.clone())) + .map_err(|err| exec_datafusion_err!("Resolving the declared table {name}: {err}"))? + .table; + let reference = TableReference::from(name.as_str()); + let table_name = reference.table().to_owned(); + let destination = reference.clone().resolve(&default_catalog, &default_schema); + 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()); + } + // The check above cannot catch a name this same call declared, + // because nothing is written until the commit step. Without this + // one, two spellings of a single table would both resolve and then + // collide during the commit with the first already inserted -- + // exactly the part-applied outcome the two-phase split exists to + // rule out. + if let Some(claimed_as) = claimed.insert(destination.clone(), name.clone()) { + return Err(exec_datafusion_err!( + "Two extensions declare the table {destination}: {claimed_as} and {name} \ + name one table, so one would have to replace the other. Rename one of \ + them, or install them on separate sessions" + ) + .into()); + } + resolved.push(ResolvedTable { + schema, + name: table_name, + provider, + }); + } + Ok(PyResolvedTables { tables: resolved }) + } + /// Import the physical optimizer rules a `with_extensions` call declared. /// /// The fallible half of installing them, run while the call can still fail @@ -1846,6 +1945,25 @@ 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. +/// `frozen` for the same reason as its sibling: the commit only reads. +#[pyclass(frozen, 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, +} + /// Physical optimizer rules imported for a `with_extensions` call. /// /// Opaque to Python, and deliberately not added to the module: it exists only diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index dfe6d6b17..a225f53f6 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -128,13 +128,19 @@ A call therefore splits into a part that may fail and a part that may not: live on that handle rather than on the session, so this step writes nothing 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 physical optimizer rule has its capsule imported, and every - `__datafusion_session_planner__` runs against the completed chains. -4. **Commit.** The planner is bound, the functions are registered, and the - optimizer rules are installed in a single `SessionState` rebuild. + 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. +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. Only step 4 touches the session, and every step that can fail happens before -it. This is a rule for the next field added to +it — with one honest exception: a table insert goes through a +`SchemaProvider`, and a foreign one can still refuse what it reported as free +during resolve. Tables commit first because of it, so nothing else has been +written when that happens. This is a rule for the next field added to `SessionExtensionComponents`, not only a description of the current code: a new kind of component must do its fallible work — importing a capsule, resolving a name — in step 3, so that step 4 cannot raise part-way through. diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 57ee1da0b..11dc278c7 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 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. + 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,27 @@ 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: 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. + +Strictly, whether a duplicate registration replaces or refuses is the +`SchemaProvider`'s own call, and the in-memory one a session starts with +refuses. `with_extensions` does not ask. It settles the question during resolve +and refuses whatever the destination would have done, which buys two things a +per-provider answer could not: the same rule wherever your table lands, and the +refusal arriving while a failure still costs nothing. A `SchemaProvider` that +would have replaced is therefore stricter through a bundle than through +{py:meth}`~datafusion.context.SessionContext.register_table`, which asks it +directly. Deregister the old table first if replacing is what you meant. + +A table collides on where it lands rather than on how it was spelled. A name is +lowercased when it is parsed and filled out from the session's default catalog +and schema, so `Events`, `events` and `public.events` are one table, and two +bundles declaring any two of them collide. Comparing the spellings would let the +pair through resolution and leave the duplicate to surface from the insert, with +the first table already written — the one thing {ref}`the commit order +` exists to prevent. + Your caller cannot rename your function, so stay out of the way: prefix the names with something tied to your library. @@ -358,6 +397,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..700ca1954 100644 --- a/docs/source/extension-guide/table-providers.md +++ b/docs/source/extension-guide/table-providers.md @@ -39,6 +39,18 @@ 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`. + 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..7b75b6490 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,12 +75,16 @@ 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 -before the call are still there. +before the call are still there — and a library shipping a table under a name +you have already used cannot replace it. The call fails and writes nothing, so +drop yours with {py:meth}`~datafusion.context.SessionContext.deregister_table` +first, or install onto a context that does not hold it. See +{ref}`extension_bundles_collisions`. ## Using more than one library 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 6d5a03bf0..c2e3a1dfa 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,12 @@ import pyarrow as pa import pytest from datafusion import SessionContext, SessionExtensionComponents -from datafusion_ffi_example import MyFunctionExtension, MyRuleExtension +from datafusion_ffi_example import ( + MyDataExtension, + MyFunctionExtension, + MyLogicalExtensionCodec, + MyRuleExtension, +) def _session(): @@ -203,6 +208,54 @@ 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_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 c0033b48c..7b5d668fb 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -17,12 +17,14 @@ use std::sync::{Arc, Mutex}; -use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods}; +use pyo3::types::{PyAnyMethods, PyCapsule, PyDict, PyDictMethods}; use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods}; use crate::aggregate_udf::MySumUDF; 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. @@ -137,3 +139,108 @@ 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)) + } +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 097c8ef87..25ebb2f3e 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::{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,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/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/examples/datafusion-ffi-example/uv.lock b/examples/datafusion-ffi-example/uv.lock new file mode 100644 index 000000000..8227259dd --- /dev/null +++ b/examples/datafusion-ffi-example/uv.lock @@ -0,0 +1,7 @@ +version = 1 +revision = 3 +requires-python = ">=3.10" + +[[package]] +name = "datafusion-ffi-example" +source = { editable = "." } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 3f2c0f811..f46deb11e 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -212,6 +212,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, @@ -232,9 +284,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 @@ -256,9 +308,11 @@ 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 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"): + declared[field_name] = [] for position, extension in enumerate(extensions): if not isinstance(extension, SessionComponentsExportable): continue @@ -2142,11 +2196,11 @@ 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 - table, say — is not rolled back, which is why bundle objects must be - configuration-only. + 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. A call that installs optimizer rules rebuilds the session state, which drops the session's prepared statements — see @@ -2175,14 +2229,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 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, 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. 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. Examples: The returned handle is a different object sharing one session, and @@ -2243,6 +2300,33 @@ 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 or that a second declaration resolves onto. + # + # The pass below is not redundant with that last check, and neither + # subsumes the other. Only the resolver knows that `Events` and + # `public.events` are one table, and only this side knows which + # argument each declaration came from -- the thing that picks the + # remedy. So the common case, two bundles writing the same name the + # same way, gets the message that names both of them. + resolved_tables = new.ctx._resolve_extension_tables( + [ + pair + for _, _, pair in _reject_repeated_names( + declared["table_providers"], "table" + ) + ] + ) # 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( @@ -2264,9 +2348,11 @@ def with_extensions( list(extensions), new, bool(logical_codecs or physical_codecs), + resolved_tables, [function._udf for function in resolved["udfs"]], [function._udaf for function in resolved["udafs"]], [function._udwf for function in resolved["udwfs"]], + [table_function._udtf for table_function in resolved_udtfs], resolved_rules, ) return new diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 5964f4c12..8e89c5c43 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -154,15 +154,58 @@ def _not_an_iterable(name: str, value: object, noun: str) -> str: ) -def _components(noun: str) -> Any: +def _not_a_pair(name: str, item: object, noun: str) -> str: + """Message for an item of a pair-shaped field that is not ``(name, value)``.""" + return ( + f"{name} must be an iterable of (name, {noun}) pairs, and {item!r} is " + f"not one. A {noun} is written with its name beside it — " + f'{name}=(("a_name", {noun}),) — and the inner parentheses are what ' + "make the two into one pair." + ) + + +def _pair_name_not_a_str(name: str, pair_name: object, noun: str) -> str: + """Message for a pair whose first element is not the name.""" + return ( + f"The name in a {name} pair must be a str, not a " + f"{type(pair_name).__name__}. The name comes first: " + f'{name}=(("a_name", {noun}),).' + ) + + +def _components(noun: str, *, pairs: bool = False) -> Any: """Declare a field holding a tuple of contributed components. ``noun`` names what the field holds, for the error a bundle sees when it hands over one component instead of a collection of them. Carrying it in the field metadata is what lets ``__post_init__`` normalize a field it was never told about by name. + + ``pairs`` marks a field whose items are ``(name, value)`` rather than bare + components, so ``__post_init__`` checks that shape too. """ - return field(default=(), metadata={"datafusion_component": noun}) + metadata: dict[str, Any] = {"datafusion_component": noun} + if pairs: + metadata["datafusion_component_pairs"] = True + return field(default=(), metadata=metadata) + + +def _as_pairs(name: str, components: tuple[Any, ...], noun: str) -> tuple[Any, ...]: + """Check every item of a pair-shaped field and normalize it to a tuple.""" + pairs = [] + for item in components: + # A str is iterable and unpacks into two characters, so a two-letter + # name would otherwise pass as a pair. + if isinstance(item, (str, bytes)): + raise TypeError(_not_a_pair(name, item, noun)) + try: + pair_name, value = item + except (TypeError, ValueError): + raise TypeError(_not_a_pair(name, item, noun)) from None + if not isinstance(pair_name, str): + raise TypeError(_pair_name_not_a_str(name, pair_name, noun)) + pairs.append((pair_name, value)) + return tuple(pairs) @dataclass(frozen=True) @@ -242,6 +285,14 @@ class SessionExtensionComponents: Traceback (most recent call last): ... TypeError: logical_extension_codecs must be an iterable of codec objects... + + Tables and table functions carry their name beside the value, so there + the same mistake is a missing *inner* pair of parentheses: + + >>> SessionExtensionComponents(udtfs=("expand", lambda: None)) + Traceback (most recent call last): + ... + TypeError: udtfs must be an iterable of (name, table function) pairs... """ logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = _components( @@ -293,6 +344,43 @@ class SessionExtensionComponents: :py:func:`~datafusion.udwf`. """ + udtfs: tuple[tuple[str, Any], ...] = _components("table function", pairs=True) + """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_binding`. + + Collides by name like :py:attr:`udfs`. + """ + + table_providers: tuple[tuple[str, Any], ...] = _components("table", pairs=True) + """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, so unlike a function a + table cannot shadow one. Two declarations collide when they resolve to one + table, not when they match as strings. Both rules hold wherever the table + lands, even where + :py:meth:`~datafusion.context.SessionContext.register_table` would have + replaced: see :ref:`extension_bundles_collisions`. + """ + physical_optimizer_rules: tuple[PhysicalOptimizerRuleExportable, ...] = _components( "optimizer rule" ) @@ -336,6 +424,14 @@ def __post_init__(self) -> None: components = tuple(value) except TypeError: raise TypeError(_not_an_iterable(name, value, noun)) from None + # The pair-shaped fields have a second version of the same mistake: + # `udtfs=("expand", func)` is one pair with its inner parentheses + # left off, and normalizes into two components rather than failing. + # Left unchecked it surfaces from `with_extensions` as + # `'function' object is not subscriptable`, which is the very thing + # this method exists to keep out of the caller's lap. + if spec.metadata.get("datafusion_component_pairs"): + components = _as_pairs(name, components, noun) object.__setattr__(self, name, components) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 83fea70a0..b160b5cae 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1670,6 +1670,179 @@ def test_session_extension_components_rejects_a_single_function(field): SessionExtensionComponents(**{field: _doubler()}) +@pytest.mark.parametrize( + ("field", "noun"), [("udtfs", "table function"), ("table_providers", "table")] +) +def test_session_extension_components_rejects_a_bare_pair(field, noun): + """One pair written without its inner parentheses is two components. + + The pair-shaped version of the lone-component mistake, and the one the + field's own shape invites: ``udtfs=("expand", func)`` is a two-element + tuple, so it normalizes without complaint and fails later inside + ``with_extensions`` under a name that says nothing about either. + """ + with pytest.raises( + TypeError, match=rf"{field} must be an iterable of \(name, {noun}\) pairs" + ): + SessionExtensionComponents(**{field: ("a_name", object())}) + + +@pytest.mark.parametrize("field", ["udtfs", "table_providers"]) +def test_session_extension_components_rejects_a_pair_of_the_wrong_length(field): + """Neither a bare value nor a triple is a ``(name, value)`` pair.""" + with pytest.raises(TypeError, match=r"is not one"): + SessionExtensionComponents(**{field: (object(),)}) + + with pytest.raises(TypeError, match=r"is not one"): + SessionExtensionComponents(**{field: (("a_name", object(), "extra"),)}) + + +@pytest.mark.parametrize("field", ["udtfs", "table_providers"]) +def test_session_extension_components_rejects_an_unnamed_pair(field): + """The name comes first, so a pair written the other way round is refused.""" + with pytest.raises(TypeError, match=r"name in a .* pair must be a str"): + SessionExtensionComponents(**{field: ((object(), "a_name"),)}) + + +@pytest.mark.parametrize("field", ["udtfs", "table_providers"]) +def test_session_extension_components_normalizes_a_pair_to_a_tuple(field): + """A pair given as a list is stored as a tuple, like the fields around it.""" + value = object() + + components = SessionExtensionComponents(**{field: [["a_name", value]]}) + + assert getattr(components, field) == (("a_name", value),) + + +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") + + +@pytest.mark.parametrize( + ("first", "second"), + [ + ("events", "public.events"), + ("events", "datafusion.public.events"), + ("Events", "events"), + ], +) +def test_with_extensions_rejects_two_spellings_of_one_table(ctx, first, second): + """Two names for one table is a collision, however differently they are written. + + A declared name is lowercased when it is parsed and filled out from the + session's default catalog and schema, so these pairs are one destination. + Comparing the spellings would not say so, and the duplicate would surface + from the insert with the first table already written. + """ + provider = ctx.from_pydict({"a": [1]}).into_view() + + with pytest.raises(Exception, match=r"Two extensions declare the table"): + ctx.with_extensions( + _TableExtension(table_providers=((first, provider),)), + _TableExtension(table_providers=((second, 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. + + The destination's own policy is not consulted: refusing during resolution + is what keeps the rest of the call from being written first, and asking a + ``SchemaProvider`` would mean asking at commit time, once refusing costs + something. + """ + 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") + + +def test_with_extensions_rejects_a_table_that_is_not_a_table_by_name(ctx): + """A declaration that is not a table at all is refused under its name. + + A table value can be any of four shapes, so unlike a rule the junk is only + discovered by the importer, after the bundle can be named. The declared + name is unique within the call — that is what identifies the culprit. + """ + + with pytest.raises(Exception, match=r"declared table junk"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + _TableExtension(table_providers=(("junk", object()),)), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + + def test_session_extension_components_rejects_a_single_optimizer_rule(): """The same for rules, naming what that field holds.""" with pytest.raises( @@ -1752,7 +1925,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 four 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. """ @@ -1767,6 +1940,8 @@ 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"}, "optimizer rule": {"physical_optimizer_rules"}, } diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 50a53a7f5..3a3eb3ca8 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -43,6 +43,9 @@ # half-installed; the commit itself is a `_commit_extensions` # parameter. "_resolve_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", } )