fix: extend register_table method to support TableProviderExportable

Author: kosiew

python/datafusion/catalog.py

@@ -28,6 +28,7 @@
     import pyarrow as pa
 
     from datafusion import TableProvider
+    from datafusion.context import TableProviderExportable
 
 try:
     from warnings import deprecated  # Python 3.13+
@@ -124,8 +125,14 @@ 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: Table | TableProvider) -> None:
-        """Register a table or table provider in this schema."""
+    def register_table(
+        self, name, table: Table | TableProvider | TableProviderExportable
+    ) -> None:
+        """Register a table or table provider in this schema.
+
+        Objects implementing ``__datafusion_table_provider__`` are also supported
+        and treated as :class:`TableProvider` instances.
+        """
         if isinstance(table, Table):
             return self._raw_schema.register_table(name, table.table)
         return self._raw_schema.register_table(name, table)
@@ -221,11 +228,16 @@ def table(self, name: str) -> Table | None:
         """Retrieve a specific table from this schema."""
         ...
 
-    def register_table(self, name: str, table: Table | TableProvider) -> None:  # noqa: B027
+    def register_table(  # noqa: B027
+        self, name: str, table: Table | TableProvider | TableProviderExportable
+    ) -> None:
         """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.
+
+        Objects implementing ``__datafusion_table_provider__`` are also supported
+        and treated as :class:`TableProvider` instances.
         """
 
     def deregister_table(self, name, cascade: bool) -> None:  # noqa: B027

python/datafusion/context.py

@@ -744,7 +744,9 @@ 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 | TableProvider) -> None:
+    def register_table(
+        self, name: str, table: Table | TableProvider | TableProviderExportable
+    ) -> None:
         """Register a :py:class:`~datafusion.catalog.Table` or ``TableProvider``.
 
         The registered table can be referenced from SQL statements executed against
@@ -754,10 +756,14 @@ def register_table(self, name: str, table: Table | TableProvider) -> None:
         convert them first with :meth:`datafusion.dataframe.DataFrame.into_view` or
         :meth:`datafusion.catalog.TableProvider.from_dataframe`.
 
+        Objects implementing ``__datafusion_table_provider__`` are also supported
+        and treated as :class:`~datafusion.catalog.TableProvider` instances.
+
         Args:
             name: Name of the resultant table.
-            table: DataFusion :class:`Table` or :class:`TableProvider` to add to the
-                session context.
+            table: DataFusion :class:`Table`, :class:`TableProvider`, or any object
+                implementing ``__datafusion_table_provider__`` to add to the session
+                context.
         """
         if isinstance(table, Table):
             self.ctx.register_table(name, table.table)
@@ -782,11 +788,14 @@ def register_catalog_provider(
             self.ctx.register_catalog_provider(name, provider)
 
     def register_table_provider(
-        self, name: str, provider: TableProviderExportable | TableProvider
+        self, name: str, provider: Table | TableProvider | TableProviderExportable
     ) -> None:
         """Register a table provider.
 
         Deprecated: use :meth:`register_table` instead.
+
+        Objects implementing ``__datafusion_table_provider__`` are also supported
+        and treated as :class:`~datafusion.catalog.TableProvider` instances.
         """
         warnings.warn(
             "register_table_provider is deprecated; use register_table",