feat: enhance table registration to support TableProvider in catalog and context

Author: kosiew

python/datafusion/catalog.py

@@ -27,6 +27,8 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
+    from datafusion import TableProvider
+
 try:
     from warnings import deprecated  # Python 3.13+
 except ImportError:
@@ -122,8 +124,8 @@ def table(self, name: str) -> Table:
         """Return the table with the given ``name`` from this schema."""
         return Table(self._raw_schema.table(name))
 
-    def register_table(self, name, table) -> None:
-        """Register a table provider in this schema."""
+    def register_table(self, name, table: Table | TableProvider) -> None:
+        """Register a table or table provider in this schema."""
         if isinstance(table, Table):
             return self._raw_schema.register_table(name, table.table)
         return self._raw_schema.register_table(name, table)
@@ -219,8 +221,8 @@ def table(self, name: str) -> Table | None:
         """Retrieve a specific table from this schema."""
         ...
 
-    def register_table(self, name: str, table: Table) -> None:  # noqa: B027
-        """Add a table from this schema.
+    def register_table(self, name: str, table: Table | TableProvider) -> None:  # noqa: B027
+        """Add a table to this schema.
 
         This method is optional. If your schema provides a fixed list of tables, you do
         not need to implement this method.

python/datafusion/context.py

@@ -46,6 +46,7 @@
     import pandas as pd
     import polars as pl
 
+    from datafusion import TableProvider
     from datafusion.plan import ExecutionPlan, LogicalPlan
 
 
@@ -743,16 +744,21 @@ def register_view(self, name: str, df: DataFrame) -> None:
         view = df.into_view()
         self.ctx.register_table(name, view)
 
-    def register_table(self, name: str, table: Table) -> None:
-        """Register a :py:class: `~datafusion.catalog.Table` as a table.
+    def register_table(self, name: str, table: Table | TableProvider) -> None:
+        """Register a :py:class:`~datafusion.catalog.Table` or ``TableProvider``.
 
-        The registered table can be referenced from SQL statement executed against.
+        The registered table can be referenced from SQL statements executed against
+        this context.
 
         Args:
             name: Name of the resultant table.
-            table: DataFusion table to add to the session context.
+            table: DataFusion :class:`Table` or :class:`TableProvider` to add to the
+                session context.
         """
-        self.ctx.register_table(name, table.table)
+        if isinstance(table, Table):
+            self.ctx.register_table(name, table.table)
+        else:
+            self.ctx.register_table(name, table)
 
     def deregister_table(self, name: str) -> None:
         """Remove a table from the session."""
@@ -772,14 +778,18 @@ def register_catalog_provider(
             self.ctx.register_catalog_provider(name, provider)
 
     def register_table_provider(
-        self, name: str, provider: TableProviderExportable
+        self, name: str, provider: TableProviderExportable | TableProvider
     ) -> None:
         """Register a table provider.
 
-        This table provider must have a method called ``__datafusion_table_provider__``
-        which returns a PyCapsule that exposes a ``FFI_TableProvider``.
+        Deprecated: use :meth:`register_table` instead.
         """
-        self.ctx.register_table_provider(name, provider)
+        warnings.warn(
+            "register_table_provider is deprecated; use register_table",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self.register_table(name, provider)
 
     def register_udtf(self, func: TableFunction) -> None:
         """Register a user defined table function."""

src/catalog.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::dataframe::PyTableProvider;
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::utils::{validate_pycapsule, wait_for_future};
@@ -209,11 +210,14 @@ impl PySchema {
         } else {
             match table_provider.extract::<PyTable>() {
                 Ok(py_table) => py_table.table,
-                Err(_) => {
-                    let py = table_provider.py();
-                    let provider = Dataset::new(&table_provider, py)?;
-                    Arc::new(provider) as Arc<dyn TableProvider>
-                }
+                Err(_) => match table_provider.extract::<PyTableProvider>() {
+                    Ok(py_provider) => py_provider.as_table().table(),
+                    Err(_) => {
+                        let py = table_provider.py();
+                        let provider = Dataset::new(&table_provider, py)?;
+                        Arc::new(provider) as Arc<dyn TableProvider>
+                    }
+                },
             }
         };
 
@@ -305,7 +309,7 @@ impl RustWrappedPySchemaProvider {
             }
 
             if py_table.hasattr("__datafusion_table_provider__")? {
-                let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?;
+                let capsule = py_table.getattr("__datafusion_table_provider__")?.call0()?;
                 let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
                 validate_pycapsule(capsule, "datafusion_table_provider")?;
 
@@ -320,6 +324,10 @@ impl RustWrappedPySchemaProvider {
                     }
                 }
 
+                if let Ok(py_provider) = py_table.extract::<PyTableProvider>() {
+                    return Ok(Some(py_provider.as_table().table()));
+                }
+
                 match py_table.extract::<PyTable>() {
                     Ok(py_table) => Ok(Some(py_table.table)),
                     Err(_) => {

src/context.rs

@@ -33,6 +33,7 @@ use pyo3::prelude::*;
 
 use crate::catalog::{PyCatalog, PyTable, RustWrappedPyCatalogProvider};
 use crate::dataframe::PyDataFrame;
+use crate::dataframe::PyTableProvider;
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
 use crate::expr::sort_expr::PySortExpr;
@@ -417,12 +418,7 @@ impl PySessionContext {
             .with_listing_options(options)
             .with_schema(resolved_schema);
         let table = ListingTable::try_new(config)?;
-        self.register_table(
-            name,
-            &PyTable {
-                table: Arc::new(table),
-            },
-        )?;
+        self.ctx.register_table(name, Arc::new(table))?;
         Ok(())
     }
 
@@ -607,8 +603,32 @@ impl PySessionContext {
         Ok(df)
     }
 
-    pub fn register_table(&mut self, name: &str, table: &PyTable) -> PyDataFusionResult<()> {
-        self.ctx.register_table(name, table.table())?;
+    pub fn register_table(
+        &mut self,
+        name: &str,
+        table_provider: Bound<'_, PyAny>,
+    ) -> PyDataFusionResult<()> {
+        let provider = if table_provider.hasattr("__datafusion_table_provider__")? {
+            let capsule = table_provider
+                .getattr("__datafusion_table_provider__")?
+                .call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_table_provider")?;
+
+            let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+            let provider: ForeignTableProvider = provider.into();
+            Arc::new(provider) as Arc<dyn TableProvider>
+        } else if let Ok(py_table) = table_provider.extract::<PyTable>() {
+            py_table.table()
+        } else if let Ok(py_provider) = table_provider.extract::<PyTableProvider>() {
+            py_provider.as_table().table()
+        } else {
+            return Err(crate::errors::PyDataFusionError::Common(
+                "Expected a Table or TableProvider.".to_string(),
+            ));
+        };
+
+        self.ctx.register_table(name, provider)?;
         Ok(())
     }
 
@@ -651,23 +671,8 @@ impl PySessionContext {
         name: &str,
         provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        if provider.hasattr("__datafusion_table_provider__")? {
-            let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?;
-            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-            validate_pycapsule(capsule, "datafusion_table_provider")?;
-
-            let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
-            let provider: ForeignTableProvider = provider.into();
-
-            let _ = self.ctx.register_table(name, Arc::new(provider))?;
-
-            Ok(())
-        } else {
-            Err(crate::errors::PyDataFusionError::Common(
-                "__datafusion_table_provider__ does not exist on Table Provider object."
-                    .to_string(),
-            ))
-        }
+        // Deprecated: use `register_table` instead
+        self.register_table(name, provider)
     }
 
     pub fn register_record_batches(

src/lib.rs

@@ -89,6 +89,7 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<dataframe::PyDataFrame>()?;
     m.add_class::<dataframe::PyParquetColumnOptions>()?;
     m.add_class::<dataframe::PyParquetWriterOptions>()?;
+    m.add_class::<dataframe::PyTableProvider>()?;
     m.add_class::<udf::PyScalarUDF>()?;
     m.add_class::<udaf::PyAggregateUDF>()?;
     m.add_class::<udwf::PyWindowUDF>()?;