1919
2020from __future__ import annotations
2121
22+ import warnings
2223from abc import ABC , abstractmethod
23- from typing import TYPE_CHECKING , Protocol
24+ from typing import TYPE_CHECKING , Any , Protocol
2425
2526import datafusion ._internal as df_internal
27+ from datafusion ._internal import EXPECTED_PROVIDER_MSG
28+ from datafusion .utils import _normalize_table_provider
2629
2730if TYPE_CHECKING :
2831 import pyarrow as pa
2932
33+ from datafusion .context import TableProviderExportable
34+
3035try :
3136 from warnings import deprecated # Python 3.13+
3237except ImportError :
@@ -82,7 +87,11 @@ def database(self, name: str = "public") -> Schema:
8287 """Returns the database with the given ``name`` from this catalog."""
8388 return self .schema (name )
8489
85- def register_schema (self , name , schema ) -> Schema | None :
90+ def register_schema (
91+ self ,
92+ name : str ,
93+ schema : Schema | SchemaProvider | SchemaProviderExportable ,
94+ ) -> Schema | None :
8695 """Register a schema with this catalog."""
8796 if isinstance (schema , Schema ):
8897 return self .catalog .register_schema (name , schema ._raw_schema )
@@ -122,11 +131,16 @@ def table(self, name: str) -> Table:
122131 """Return the table with the given ``name`` from this schema."""
123132 return Table (self ._raw_schema .table (name ))
124133
125- def register_table (self , name , table ) -> None :
126- """Register a table provider in this schema."""
127- if isinstance (table , Table ):
128- return self ._raw_schema .register_table (name , table .table )
129- return self ._raw_schema .register_table (name , table )
134+ def register_table (
135+ self , name : str , table : Table | TableProviderExportable | Any
136+ ) -> None :
137+ """Register a table or table provider in this schema.
138+
139+ Objects implementing ``__datafusion_table_provider__`` are also supported
140+ and treated as table provider instances.
141+ """
142+ provider = _normalize_table_provider (table )
143+ return self ._raw_schema .register_table (name , provider )
130144
131145 def deregister_table (self , name : str ) -> None :
132146 """Deregister a table provider from this schema."""
@@ -138,31 +152,101 @@ class Database(Schema):
138152 """See `Schema`."""
139153
140154
155+ _InternalRawTable = df_internal .catalog .RawTable
156+ _InternalTableProvider = df_internal .TableProvider
157+
158+ # Keep in sync with ``datafusion._internal.TableProvider.from_view``.
159+ _FROM_VIEW_WARN_STACKLEVEL = 2
160+
161+
141162class Table :
142- """DataFusion table."""
163+ """DataFusion table or table provider wrapper ."""
143164
144- def __init__ (self , table : df_internal .catalog .RawTable ) -> None :
145- """This constructor is not typically called by the end user."""
146- self .table = table
165+ __slots__ = ("_table" ,)
166+
167+ def __init__ (
168+ self ,
169+ table : _InternalRawTable | _InternalTableProvider | Table ,
170+ ) -> None :
171+ """Wrap a low level table or table provider."""
172+ if isinstance (table , Table ):
173+ table = table .table
174+
175+ if not isinstance (table , (_InternalRawTable , _InternalTableProvider )):
176+ raise TypeError (EXPECTED_PROVIDER_MSG )
177+
178+ self ._table = table
179+
180+ def __getattribute__ (self , name : str ) -> Any :
181+ """Restrict provider-specific helpers to compatible tables."""
182+ if name == "__datafusion_table_provider__" :
183+ table = object .__getattribute__ (self , "_table" )
184+ if not hasattr (table , "__datafusion_table_provider__" ):
185+ raise AttributeError (name )
186+ return object .__getattribute__ (self , name )
147187
148188 def __repr__ (self ) -> str :
149189 """Print a string representation of the table."""
150- return self .table . __repr__ ( )
190+ return repr ( self ._table )
151191
152- @staticmethod
153- def from_dataset (dataset : pa .dataset .Dataset ) -> Table :
154- """Turn a pyarrow Dataset into a Table."""
155- return Table (df_internal .catalog .RawTable .from_dataset (dataset ))
192+ @property
193+ def table (self ) -> _InternalRawTable | _InternalTableProvider :
194+ """Return the wrapped low level table object."""
195+ return self ._table
196+
197+ @classmethod
198+ def from_dataset (cls , dataset : pa .dataset .Dataset ) -> Table :
199+ """Turn a :mod:`pyarrow.dataset` ``Dataset`` into a :class:`Table`."""
200+ return cls (_InternalRawTable .from_dataset (dataset ))
201+
202+ @classmethod
203+ def from_capsule (cls , capsule : Any ) -> Table :
204+ """Create a :class:`Table` from a PyCapsule exported provider."""
205+ provider = _InternalTableProvider .from_capsule (capsule )
206+ return cls (provider )
207+
208+ @classmethod
209+ def from_dataframe (cls , df : Any ) -> Table :
210+ """Create a :class:`Table` from tabular data."""
211+ from datafusion .dataframe import DataFrame as DataFrameWrapper
212+
213+ dataframe = df if isinstance (df , DataFrameWrapper ) else DataFrameWrapper (df )
214+ return dataframe .into_view ()
215+
216+ @classmethod
217+ def from_view (cls , df : Any ) -> Table :
218+ """Deprecated helper for constructing tables from views."""
219+ from datafusion .dataframe import DataFrame as DataFrameWrapper
220+
221+ if isinstance (df , DataFrameWrapper ):
222+ df = df .df
223+
224+ provider = _InternalTableProvider .from_view (df )
225+ warnings .warn (
226+ "Table.from_view is deprecated; use DataFrame.into_view or "
227+ "Table.from_dataframe instead." ,
228+ category = DeprecationWarning ,
229+ stacklevel = _FROM_VIEW_WARN_STACKLEVEL ,
230+ )
231+ return cls (provider )
156232
157233 @property
158234 def schema (self ) -> pa .Schema :
159235 """Returns the schema associated with this table."""
160- return self .table .schema
236+ return self ._table .schema
161237
162238 @property
163239 def kind (self ) -> str :
164240 """Returns the kind of table."""
165- return self .table .kind
241+ return self ._table .kind
242+
243+ def __datafusion_table_provider__ (self ) -> Any :
244+ """Expose the wrapped provider for FFI integrations."""
245+ exporter = getattr (self ._table , "__datafusion_table_provider__" , None )
246+ if exporter is None :
247+ msg = "Underlying object does not export __datafusion_table_provider__()"
248+ raise AttributeError (msg )
249+ return exporter ()
166250
167251
168252class CatalogProvider (ABC ):
@@ -219,14 +303,19 @@ def table(self, name: str) -> Table | None:
219303 """Retrieve a specific table from this schema."""
220304 ...
221305
222- def register_table (self , name : str , table : Table ) -> None : # noqa: B027
223- """Add a table from this schema.
306+ def register_table ( # noqa: B027
307+ self , name : str , table : Table | TableProviderExportable | Any
308+ ) -> None :
309+ """Add a table to this schema.
224310
225311 This method is optional. If your schema provides a fixed list of tables, you do
226312 not need to implement this method.
313+
314+ Objects implementing ``__datafusion_table_provider__`` are also supported
315+ and treated as table provider instances.
227316 """
228317
229- def deregister_table (self , name , cascade : bool ) -> None : # noqa: B027
318+ def deregister_table (self , name : str , cascade : bool ) -> None : # noqa: B027
230319 """Remove a table from this schema.
231320
232321 This method is optional. If your schema provides a fixed list of tables, you do
0 commit comments