feat: let extension bundles declare tables and table functions

`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) <noreply@anthropic.com>

Author: timsaucer

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::{
@@ -1764,6 +1766,68 @@ 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<PyResolvedTables> {
+        let session = slf.clone().into_bound_py_any(slf.py())?;
+        let state = slf.borrow().ctx.state();
+
+        let mut resolved = Vec::with_capacity(tables.len());
+        for (name, obj) in tables {
+            let provider = PyTable::new(obj, Some(session.clone()))?.table;
+            let reference = TableReference::from(name.as_str());
+            let table_name = reference.table().to_owned();
+            let schema = state.schema_for_ref(reference)?;
+            // Checked against the schema rather than against this call's own
+            // list, so a name the session already holds is caught too. Both
+            // are the same error to a caller.
+            if schema.table_exist(&table_name) {
+                return Err(exec_datafusion_err!(
+                    "An extension declared a table named {name}, which is already registered"
+                )
+                .into());
+            }
+            resolved.push(ResolvedTable {
+                schema,
+                name: table_name,
+                provider,
+            });
+        }
+        Ok(PyResolvedTables { tables: resolved })
+    }
+
+    /// Commit the tables for a `with_extensions` call.
+    ///
+    /// Runs first among the commit steps. Every name was resolved and found
+    /// free by [`Self::_resolve_extension_tables`], so the only way an insert
+    /// still fails is a foreign `SchemaProvider` refusing a registration it
+    /// reported as available — the one place in `with_extensions` that can
+    /// leave a call part-applied. Going first is what keeps the blast radius
+    /// to the tables themselves: no planner is bound and no function is
+    /// registered behind it.
+    pub fn _install_extension_tables(
+        &self,
+        resolved: PyRef<'_, PyResolvedTables>,
+    ) -> PyDataFusionResult<()> {
+        for table in &resolved.tables {
+            table
+                .schema
+                .register_table(table.name.clone(), Arc::clone(&table.provider))?;
+        }
+        Ok(())
+    }
+
     /// Import the physical optimizer rules a `with_extensions` call declared.
     ///
     /// The fallible half of installing them, run while the call can still fail
@@ -1818,6 +1882,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<ResolvedTable>,
+}
+
+/// One entry of [`PyResolvedTables`]: where it goes, and what goes there.
+struct ResolvedTable {
+    schema: Arc<dyn SchemaProvider>,
+    name: String,
+    provider: Arc<dyn TableProvider>,
+}
+
 /// Physical optimizer rules imported for a `with_extensions` call.
 ///
 /// Opaque to Python, and deliberately not added to the module: it exists only

docs/source/extension-guide/bundles.md

@@ -62,15 +62,13 @@ caller then writes:
 
 ```python
 ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension())
-ctx.register_table("t", lib_a.TableProvider())
 ```
 
-Declare functions rather than registering them yourself inside the hook.
-Declared components are resolved before anything is written, and they are
-registered after every codec is installed; a registration you make during the
-hook happens too early to see the other bundles' codecs and is not undone if a
-later extension fails. Table providers are still registered by the caller, on
-the returned handle — see {ref}`extension_bundles_transaction`.
+Declare what you contribute rather than registering it yourself inside the
+hook. Declared components are resolved before anything is written, and they go
+on after every codec is installed; a registration you make during the hook
+happens too early to see the other bundles' codecs and is not undone if a later
+extension fails. 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
@@ -321,6 +319,31 @@ shares one session with the receiver, and undoing a registration is not the
 same as restoring what it displaced: deregistering a function that shadowed a
 built-in removes the built-in too.
 
+Tables are the single exception, and they are committed first because of it. A
+declared table has its provider imported and its destination schema resolved in
+step 3, and a name already taken is refused there — 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.
+
+(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, in
+  step 3, 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
@@ -340,6 +363,10 @@ Physical optimizer rules are exempt: 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.
+
 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

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.
 
@@ -107,6 +107,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

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 <datafusion.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

docs/source/user-guide/extensions.md

@@ -67,11 +67,10 @@ ctx.sql("SELECT count(*) FROM events").show()
 everything else with the context you called it on, so tables you registered
 before the call are still there.
 
-**Functions can arrive either way.** A single function is registered directly
-with {py:func}`~datafusion.udf` and
-{py:meth}`~datafusion.SessionContext.register_udf`. A library shipping a set of
-them usually packages them in the same `Extension` object instead, so
-`with_extensions` installs them along with everything else it provides. Follow
+**Tables and functions can arrive either way.** A single one is registered
+directly, as above. A library shipping a set of them usually packages them in
+the same `Extension` object instead, so `with_extensions` installs them along
+with everything else it provides — and the library picks the names. Follow
 whichever the library documents.
 
 ## Using more than one library

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():
@@ -184,6 +189,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.