Let extension bundles declare tables and table functions

crates/core/src/context.rs

@@ -27,8 +27,12 @@ use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory};
-use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_err};
+use datafusion::catalog::{
+    CatalogProvider, CatalogProviderList, SchemaProvider, TableProviderFactory,
+};
+use datafusion::common::{
+    DFSchema, ResolvedTableReference, ScalarValue, TableReference, exec_datafusion_err, exec_err,
+};
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::listing::{
@@ -1738,6 +1742,14 @@ impl PySessionContext {
     /// reasoning is in docs/source/contributor-guide/ffi-internals.md, under
     /// "Why `with_extensions` commits last".
     ///
+    /// The tables are the one exception and so go first. Every name was
+    /// resolved and found free by [`Self::_resolve_extension_tables`], so the
+    /// only way an insert still fails is a foreign `SchemaProvider` refusing
+    /// a registration it reported as available — the one place in
+    /// `with_extensions` that can leave a call part-applied. Going first is
+    /// what keeps the blast radius to the tables themselves: no planner is
+    /// bound and no function is registered behind it.
+    ///
     /// The planner is bound through this context's own `state_ref()`, so
     /// providers bound to it stay valid. With no planner supplied the bind
     /// still rebuilds whichever planner the session already holds against
@@ -1758,9 +1770,11 @@ impl PySessionContext {
         extensions: Vec<Bound<'py, PyAny>>,
         session: Bound<'py, PyAny>,
         rebind_planner: bool,
+        tables: PyRef<'_, PyResolvedTables>,
         udfs: Vec<PyScalarUDF>,
         udafs: Vec<PyAggregateUDF>,
         udwfs: Vec<PyWindowUDF>,
+        udtfs: Vec<PyTableFunction>,
         rules: PyRef<'_, PyPhysicalOptimizerRules>,
     ) -> PyDataFusionResult<()> {
         let py = slf.py();
@@ -1787,6 +1801,13 @@ impl PySessionContext {
             )?);
         }
 
+        // The first write. A foreign `SchemaProvider` refusing here leaves
+        // nothing else behind — see the tables exception above.
+        for table in &tables.tables {
+            table
+                .schema
+                .register_table(table.name.clone(), Arc::clone(&table.provider))?;
+        }
         if planner.is_some() || rebind_planner {
             slf.borrow().set_session_query_planner(planner);
         }
@@ -1800,6 +1821,9 @@ impl PySessionContext {
         for udwf in udwfs {
             this.ctx.register_udwf(udwf.function);
         }
+        for udtf in udtfs {
+            this.register_udtf(udtf);
+        }
         // Rules accumulate rather than replace, so unlike a planner there is
         // no composition order to get right and no collision to refuse. All
         // of them go on in **one** `SessionState` rebuild.
@@ -1824,6 +1848,81 @@ impl PySessionContext {
         Ok(())
     }
 
+    /// Resolve the tables a `with_extensions` call declared.
+    ///
+    /// The fallible half. Each provider is imported against `slf` — the handle
+    /// carrying the completed codec chains, not the context the components hook
+    /// was given — and each name is resolved to the schema that will hold it.
+    /// A name already taken is refused here rather than left to the insert.
+    /// Whether a duplicate replaces or refuses is the `SchemaProvider`'s own
+    /// call — the in-memory one refuses — and asking it would mean asking at
+    /// commit time, once refusing costs something. Deciding here buys one rule
+    /// for every destination and a refusal while a failure is still free.
+    ///
+    /// Two declarations landing on one destination are refused here too. Both
+    /// checks are against the *resolved* reference rather than the declared
+    /// spelling, which is the only thing that answers the question: a name is
+    /// lowercased when it is parsed and filled out from the session's default
+    /// catalog and schema, so `Events`, `events` and `public.events` are one
+    /// table under three spellings.
+    ///
+    /// **Writes nothing.**
+    pub fn _resolve_extension_tables<'py>(
+        slf: &Bound<'py, Self>,
+        tables: Vec<(String, Bound<'py, PyAny>)>,
+    ) -> PyDataFusionResult<PyResolvedTables> {
+        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<ResolvedTableReference, String> = HashMap::new();
+        let mut resolved = Vec::with_capacity(tables.len());
+        for (name, obj) in tables {
+            // The name is the culprit's identity: it is unique within the call,
+            // so an import failure that carries it points at one declaration.
+            // The importer's own message names neither the table nor the
+            // bundle, because it never knew them.
+            let provider = PyTable::new(obj, Some(session.clone()))
+                .map_err(|err| exec_datafusion_err!("Resolving the declared table {name}: {err}"))?
+                .table;
+            let reference = TableReference::from(name.as_str());
+            let table_name = reference.table().to_owned();
+            let destination = reference.clone().resolve(&default_catalog, &default_schema);
+            let schema = state.schema_for_ref(reference)?;
+            // Checked against the schema rather than against this call's own
+            // list, so a name the session already holds is caught too. Both
+            // are the same error to a caller.
+            if schema.table_exist(&table_name) {
+                return Err(exec_datafusion_err!(
+                    "An extension declared a table named {name}, which is already registered"
+                )
+                .into());
+            }
+            // The check above cannot catch a name this same call declared,
+            // because nothing is written until the commit step. Without this
+            // one, two spellings of a single table would both resolve and then
+            // collide during the commit with the first already inserted --
+            // exactly the part-applied outcome the two-phase split exists to
+            // rule out.
+            if let Some(claimed_as) = claimed.insert(destination.clone(), name.clone()) {
+                return Err(exec_datafusion_err!(
+                    "Two extensions declare the table {destination}: {claimed_as} and {name} \
+                     name one table, so one would have to replace the other. Rename one of \
+                     them, or install them on separate sessions"
+                )
+                .into());
+            }
+            resolved.push(ResolvedTable {
+                schema,
+                name: table_name,
+                provider,
+            });
+        }
+        Ok(PyResolvedTables { tables: resolved })
+    }
+
     /// Import the physical optimizer rules a `with_extensions` call declared.
     ///
     /// The fallible half of installing them, run while the call can still fail
@@ -1846,6 +1945,25 @@ impl PySessionContext {
     }
 }
 
+/// Tables resolved for a `with_extensions` call.
+///
+/// Opaque to Python, and deliberately not added to the module, like
+/// [`PyPhysicalOptimizerRules`]. Each entry is a provider that has already been
+/// imported and a schema that has already been looked up, so committing is an
+/// insert into a resolved destination rather than a fresh name resolution.
+/// `frozen` for the same reason as its sibling: the commit only reads.
+#[pyclass(frozen, name = "ResolvedTables", module = "datafusion._internal")]
+pub struct PyResolvedTables {
+    tables: Vec<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/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.

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,27 @@ Physical optimizer rules are exempt from all of this: they accumulate rather
 than replace, so two bundles contributing one each is the normal case and there
 is nothing to refuse. See {doc}`other-components`.
 
+Tables go the other way: a declared table name that is *already* on the session
+is an error too, so a table cannot shadow one the way a function can.
+
+Strictly, whether a duplicate registration replaces or refuses is the
+`SchemaProvider`'s own call, and the in-memory one a session starts with
+refuses. `with_extensions` does not ask. It settles the question during resolve
+and refuses whatever the destination would have done, which buys two things a
+per-provider answer could not: the same rule wherever your table lands, and the
+refusal arriving while a failure still costs nothing. A `SchemaProvider` that
+would have replaced is therefore stricter through a bundle than through
+{py:meth}`~datafusion.context.SessionContext.register_table`, which asks it
+directly. Deregister the old table first if replacing is what you meant.
+
+A table collides on where it lands rather than on how it was spelled. A name is
+lowercased when it is parsed and filled out from the session's default catalog
+and schema, so `Events`, `events` and `public.events` are one table, and two
+bundles declaring any two of them collide. Comparing the spellings would let the
+pair through resolution and leave the duplicate to surface from the insert, with
+the first table already written — the one thing {ref}`the commit order
+<ffi_internals_commit_order>` exists to prevent.
+
 Your caller cannot rename your function, so stay out of the way: prefix the
 names with something tied to your library.
 
@@ -358,6 +397,14 @@ where the registry is complete; a planner is called per query, long after the
 install has finished. If you need a function at hook time, you already have the
 object, because you are the one declaring it.
 
+Tables are the single exception to committing last, and they are committed
+first because of it. A declared table has its provider imported, its
+destination schema resolved, and a name already taken refused, all while a
+failure still costs nothing — but the insert itself goes through a
+`SchemaProvider`, and a foreign one can still refuse what it reported as free.
+Running that first means no planner is bound and no function is registered
+behind it when it does.
+
 Like every other derivation, the returned context is a handle on the *same*
 session as the receiver — see {ref}`extension_sessions`. Only the Python-side
 codec chains belong to the returned handle; the planner is installed on the

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,12 +75,16 @@ ctx.register_udf(udf(my_library.MyScalarUDF()))
 
 A library shipping a set of them packages them in its `Extension` object
 instead, so `with_extensions` installs them all along with everything else it
-provides, and there is nothing per-function for you to do. Its documentation
-says which.
+provides — and the library picks the names — leaving nothing per-item for you to
+do. Its documentation says which.
 
 `with_extensions` returns a context; use the returned one. It shares
 everything else with the context you called it on, so tables you registered
-before the call are still there.
+before the call are still there — and a library shipping a table under a name
+you have already used cannot replace it. The call fails and writes nothing, so
+drop yours with {py:meth}`~datafusion.context.SessionContext.deregister_table`
+first, or install onto a context that does not hold it. See
+{ref}`extension_bundles_collisions`.
 
 ## Using more than one library