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

@@ -55,23 +55,23 @@ 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
 `__datafusion_session_planner__`. The caller then writes:
 
 ```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

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

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

@@ -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

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.