From 4585f0c2f7b07017f69687fe2dafb7040d5865df Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 10:13:00 -0400 Subject: [PATCH 1/7] feat: let extension bundles declare tables and table functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `SessionExtensionComponents` gains `udtfs` and `table_providers`, both as `(name, value)` pairs. Neither carries a name of its own the way a scalar function's capsule does, and both getters take the session — which is what makes them different from everything the stack has added so far. Because they take the session, the host resolves them against the handle carrying the completed codec chains rather than against the context the components hook received. A bundle wrapping one itself would bind it to a chain missing every library in the call, including its own, and the failure would not surface until a decode somewhere else. So a bundle hands over the unwrapped value and lets the host wrap it. `RecordingTableFunction` in the example crate records the codec ids it was handed, which turns that claim into an assertion rather than a paragraph. Tables do not shadow. DataFusion refuses a duplicate table registration rather than replacing it, so a declared name already on the session is an error too, not just one two bundles both claim. Both are caught while resolving, alongside resolving the destination schema, so a bad name costs nothing. `_resolve_extension_tables` and `_install_extension_tables` keep the same split as the rules, with one honest exception: the insert goes through a `SchemaProvider`, and a foreign one can still refuse what it reported as free. Tables are therefore committed first, so nothing else has been written when that happens. The guide says so rather than claiming a guarantee that does not hold. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 85 +++++++++++++- docs/source/extension-guide/bundles.md | 46 ++++++-- docs/source/extension-guide/functions.md | 14 ++- .../source/extension-guide/table-providers.md | 12 ++ docs/source/user-guide/extensions.md | 10 +- .../python/tests/_test_session_extension.py | 55 ++++++++- .../datafusion-ffi-example/src/extension.rs | 109 +++++++++++++++++- examples/datafusion-ffi-example/src/lib.rs | 3 +- .../src/table_function.rs | 4 +- python/datafusion/context.py | 95 +++++++++++++-- python/datafusion/extensions.py | 34 ++++++ python/tests/test_context.py | 87 +++++++++++++- python/tests/test_wrapper_coverage.py | 3 + 13 files changed, 526 insertions(+), 31 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 554fc65d7..7c1c5c3cb 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::{ @@ -1738,6 +1740,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 +1768,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 +1799,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 +1819,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 +1846,47 @@ 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 }) + } + /// 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 +1909,24 @@ 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, +} + /// 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/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 57ee1da0b..d2c6ee77d 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,10 @@ 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. + Your caller cannot rename your function, so stay out of the way: prefix the names with something tied to your library. @@ -358,6 +380,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..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 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/python/datafusion/context.py b/python/datafusion/context.py index 3f2c0f811..5b2aecda7 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,9 +2196,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. @@ -2243,6 +2298,26 @@ 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" + ) + ] + ) # 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 +2339,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..a3cb05632 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -293,6 +293,40 @@ 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. + """ + physical_optimizer_rules: tuple[PhysicalOptimizerRuleExportable, ...] = _components( "optimizer rule" ) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 83fea70a0..ac10b6b8a 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1670,6 +1670,89 @@ 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") + + def test_session_extension_components_rejects_a_single_optimizer_rule(): """The same for rules, naming what that field holds.""" with pytest.raises( @@ -1752,7 +1835,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 +1850,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", } ) From 37fcabefc419de1685dff2962593ba4d673bda17 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 11:20:13 -0400 Subject: [PATCH 2/7] fix: name the declared table when its value will not import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A junk table value fell through to the pyarrow Dataset fallback, whose error names neither the table nor the bundle. The declared name is unique within the call, so wrapping the import error with it points at one declaration — the same repair the rules got, adapted to a component whose value has four legal shapes and so cannot be pre-checked in Python. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 8 +++++++- python/datafusion/context.py | 10 ++++++---- python/tests/test_context.py | 18 ++++++++++++++++++ 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7c1c5c3cb..9eb271770 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1865,7 +1865,13 @@ impl PySessionContext { let mut resolved = Vec::with_capacity(tables.len()); for (name, obj) in tables { - let provider = PyTable::new(obj, Some(session.clone()))?.table; + // 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 schema = state.schema_for_ref(reference)?; diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 5b2aecda7..c3fc40d6d 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2230,14 +2230,16 @@ 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 + taken, the schema unknown, or the value not a table. Examples: The returned handle is a different object sharing one session, and diff --git a/python/tests/test_context.py b/python/tests/test_context.py index ac10b6b8a..2e4957160 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1753,6 +1753,24 @@ def test_with_extensions_rejects_a_table_in_an_unknown_schema(ctx): 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( From da4d1b814cd457b69ac3149b89207676d9064d1b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 11:20:36 -0400 Subject: [PATCH 3/7] style: freeze the resolved-tables carrier Same reasoning as PhysicalOptimizerRules gaining frozen on the base branch: nothing mutates it between resolve and commit, so there is no reason to pay for the runtime borrow flag a mutable pyclass carries. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 9eb271770..03ad46ec0 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1921,7 +1921,8 @@ impl PySessionContext { /// [`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")] +/// `frozen` for the same reason as its sibling: the commit only reads. +#[pyclass(frozen, name = "ResolvedTables", module = "datafusion._internal")] pub struct PyResolvedTables { tables: Vec, } From bdad9b813a720ddba6675857a4ebcb639da100c6 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 11:20:45 -0400 Subject: [PATCH 4/7] docs: add tables to the commit-order contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The resolve and commit steps in ffi-internals.md predate declared tables. Tables resolve alongside everything else, but their insert goes through a SchemaProvider, and a foreign one can still refuse what it reported as free — the one honest exception to "step 4 cannot raise", already stated in the bundle guide, now stated where the rule for the next field lives. Co-Authored-By: Claude Fable 5 --- docs/source/contributor-guide/ffi-internals.md | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) 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. From e340b539ce20d41ba88e6b3c9a77d868d58b71e4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 12:18:46 -0400 Subject: [PATCH 5/7] fix: refuse two declared spellings of one table `_reject_repeated_names` compares the declared strings, and the `table_exist` check in `_resolve_extension_tables` runs against a schema nothing has been written to yet. Two declarations that resolve to one table but differ as strings passed both: `TableReference::from` lowercases a name and splits it, and the default catalog and schema fill in the rest, so `Events`, `events` and `public.events` are one destination under three spellings. Both then resolved, and the duplicate surfaced from the insert during the commit with the first table already registered -- the part-applied outcome the two-phase split exists to rule out, reachable through the default in-memory schema provider rather than only through a foreign one. Resolve each name to a `ResolvedTableReference` and key the check on that. The Python pass stays: only the resolver knows two spellings are one table, and only Python knows which argument each declaration came from, which is what picks the remedy. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 30 +++++++++++++++++++++++++- docs/source/extension-guide/bundles.md | 8 +++++++ python/datafusion/context.py | 9 +++++++- python/datafusion/extensions.py | 4 +++- python/tests/test_context.py | 27 +++++++++++++++++++++++ 5 files changed, 75 insertions(+), 3 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 03ad46ec0..5065a6575 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -30,7 +30,9 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::catalog::{ CatalogProvider, CatalogProviderList, SchemaProvider, TableProviderFactory, }; -use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_datafusion_err, exec_err}; +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::{ @@ -1855,6 +1857,13 @@ impl PySessionContext { /// duplicate registration rather than replacing it, and a refusal is much /// more useful before anything has been written. /// + /// 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>, @@ -1862,7 +1871,11 @@ impl PySessionContext { ) -> 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, @@ -1874,6 +1887,7 @@ impl PySessionContext { .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 @@ -1884,6 +1898,20 @@ impl PySessionContext { ) .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, diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index d2c6ee77d..6922acb2a 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -350,6 +350,14 @@ 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. +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. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index c3fc40d6d..e7c446ed2 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2311,7 +2311,14 @@ def with_extensions( ] # 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. + # 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 diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index a3cb05632..452bc901b 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -324,7 +324,9 @@ class SessionExtensionComponents: 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. + replacing it, so unlike a function a table cannot shadow one. Two + declarations collide when they resolve to one table, not when they match as + strings: see :ref:`extension_bundles_collisions`. """ physical_optimizer_rules: tuple[PhysicalOptimizerRuleExportable, ...] = _components( diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 2e4957160..8eb2b25ed 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1719,6 +1719,33 @@ def test_with_extensions_rejects_a_table_name_two_extensions_claim(ctx): 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. From fa647c03cb47e442e1183ac2f34716e61b3aa138 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 12:23:12 -0400 Subject: [PATCH 6/7] fix: refuse a declared pair written without its inner parentheses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `__post_init__` catches `udfs=fn` written for `udfs=(fn,)` so the error lands in the extension library's own frame rather than deep inside `with_extensions`. The pair-shaped fields have the same mistake one level in — `udtfs=("expand", fn)` is a two-element tuple, so it normalized without complaint and surfaced later as `'function' object is not subscriptable`, naming neither the field nor the bundle. Mark those fields in the field metadata and check each item is a `(str, value)` pair, normalizing it to a tuple like everything else here. A str item is rejected before unpacking: a two-letter name would otherwise unpack into two characters and pass. Co-Authored-By: Claude Opus 5 (1M context) --- examples/datafusion-ffi-example/uv.lock | 7 +++ python/datafusion/extensions.py | 67 +++++++++++++++++++++++-- python/tests/test_context.py | 44 ++++++++++++++++ 3 files changed, 114 insertions(+), 4 deletions(-) create mode 100644 examples/datafusion-ffi-example/uv.lock 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/extensions.py b/python/datafusion/extensions.py index 452bc901b..0181d0d10 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,7 +344,7 @@ class SessionExtensionComponents: :py:func:`~datafusion.udwf`. """ - udtfs: tuple[tuple[str, Any], ...] = _components("table function") + 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 @@ -310,7 +361,7 @@ class SessionExtensionComponents: Collides by name like :py:attr:`udfs`. """ - table_providers: tuple[tuple[str, Any], ...] = _components("table") + 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` @@ -372,6 +423,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 8eb2b25ed..99e6902af 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1670,6 +1670,50 @@ 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.""" From 300e321403524a7c61b7387a572f2c8b132e5c1e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 12:29:23 -0400 Subject: [PATCH 7/7] docs: correct the table duplicate rule and point at the binding section MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four review nits, all documentation. The no-shadowing rule was credited to DataFusion. It is not DataFusion's: a `SchemaProvider` decides for itself whether a duplicate replaces or refuses, and only the in-memory one a session starts with refuses. `with_extensions` does not ask — it settles the question during resolve, which is what buys one rule for every destination and a refusal while a failure is still free, and which makes a bundle stricter than `register_table` against a provider that would have replaced. Said once in the collisions section and pointed at from the field docstring, the Rust resolver and the test that pins it. `udtfs` pointed at the two-phase section for a claim whose canonical home is the binding section this PR added; `functions.md` and `table-providers.md` already point there. The `Raises:` entry for a declared table did not mention two declarations resolving to one, and a rewrap had left "registering a" alone on a line. The user guide now says what a caller meets: a bundle cannot take a table name they already used. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 8 +++++--- docs/source/extension-guide/bundles.md | 15 ++++++++++++--- docs/source/user-guide/extensions.md | 6 +++++- python/datafusion/context.py | 8 ++++---- python/datafusion/extensions.py | 13 +++++++------ python/tests/test_context.py | 7 ++++--- 6 files changed, 37 insertions(+), 20 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 5065a6575..9a92c6055 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1853,9 +1853,11 @@ impl PySessionContext { /// 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. + /// 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 diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 6922acb2a..11dc278c7 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -346,9 +346,18 @@ 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. +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 diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md index e758bd525..7b75b6490 100644 --- a/docs/source/user-guide/extensions.md +++ b/docs/source/user-guide/extensions.md @@ -80,7 +80,11 @@ 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/python/datafusion/context.py b/python/datafusion/context.py index e7c446ed2..f46deb11e 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2199,9 +2199,8 @@ def with_extensions( 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. + 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 @@ -2239,7 +2238,8 @@ def with_extensions( 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 - taken, the schema unknown, or the value not a table. + 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 diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 0181d0d10..8e89c5c43 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -356,7 +356,7 @@ class SessionExtensionComponents: 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`. + finished context instead. See :ref:`extension_bundles_binding`. Collides by name like :py:attr:`udfs`. """ @@ -373,11 +373,12 @@ class SessionExtensionComponents: 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. Two - declarations collide when they resolve to one table, not when they match as - strings: see :ref:`extension_bundles_collisions`. + 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( diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 99e6902af..b160b5cae 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1793,9 +1793,10 @@ def test_with_extensions_rejects_two_spellings_of_one_table(ctx, first, second): 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. + 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()