Initial implementation of unified table suggestion

Author: timsaucer

examples/datafusion-ffi-example/python/tests/_test_aggregate_udf.py

@@ -60,7 +60,7 @@ def test_ffi_aggregate_register():
 def test_ffi_aggregate_call_directly():
     ctx = setup_context_with_table()
     my_udaf = udaf(MySumUDF())
-    
+
     result = (
         ctx.table("test_table").aggregate([col("b")], [my_udaf(col("a"))]).collect()
     )

python/datafusion/__init__.py

@@ -34,7 +34,7 @@
 from . import functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
-from ._internal import EXPECTED_PROVIDER_MSG, Config
+from ._internal import Config
 from .catalog import Catalog, Database, Table
 from .col import col, column
 from .common import DFSchema
@@ -65,7 +65,6 @@
 __version__ = importlib_metadata.version(__name__)
 
 __all__ = [
-    "EXPECTED_PROVIDER_MSG",
     "Accumulator",
     "AggregateUDF",
     "Catalog",

python/datafusion/catalog.py

@@ -19,17 +19,15 @@
 
 from __future__ import annotations
 
-import warnings
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Protocol
 
 import datafusion._internal as df_internal
-from datafusion._internal import EXPECTED_PROVIDER_MSG
-from datafusion.utils import _normalize_table_provider
 
 if TYPE_CHECKING:
     import pyarrow as pa
 
+    from datafusion import DataFrame
     from datafusion.context import TableProviderExportable
 
 try:
@@ -139,8 +137,7 @@ def register_table(
         Objects implementing ``__datafusion_table_provider__`` are also supported
         and treated as table provider instances.
         """
-        provider = _normalize_table_provider(table)
-        return self._raw_schema.register_table(name, provider)
+        return self._raw_schema.register_table(name, table)
 
     def deregister_table(self, name: str) -> None:
         """Deregister a table provider from this schema."""
@@ -152,101 +149,37 @@ class Database(Schema):
     """See `Schema`."""
 
 
-_InternalRawTable = df_internal.catalog.RawTable
-_InternalTableProvider = df_internal.TableProvider
-
-# Keep in sync with ``datafusion._internal.TableProvider.from_view``.
-_FROM_VIEW_WARN_STACKLEVEL = 2
-
-
 class Table:
     """DataFusion table or table provider wrapper."""
 
-    __slots__ = ("_table",)
+    __slots__ = ("_inner",)
 
     def __init__(
         self,
-        table: _InternalRawTable | _InternalTableProvider | Table,
+        table: DataFrame | TableProviderExportable | pa.dataset.Dataset,
     ) -> None:
         """Wrap a low level table or table provider."""
-        if isinstance(table, Table):
-            table = table.table
-
-        if not isinstance(table, (_InternalRawTable, _InternalTableProvider)):
-            raise TypeError(EXPECTED_PROVIDER_MSG)
-
-        self._table = table
-
-    def __getattribute__(self, name: str) -> Any:
-        """Restrict provider-specific helpers to compatible tables."""
-        if name == "__datafusion_table_provider__":
-            table = object.__getattribute__(self, "_table")
-            if not hasattr(table, "__datafusion_table_provider__"):
-                raise AttributeError(name)
-        return object.__getattribute__(self, name)
+        self._inner = df_internal.catalog.RawTable(table)
 
     def __repr__(self) -> str:
         """Print a string representation of the table."""
-        return repr(self._table)
-
-    @property
-    def table(self) -> _InternalRawTable | _InternalTableProvider:
-        """Return the wrapped low level table object."""
-        return self._table
+        return repr(self._inner)
 
-    @classmethod
-    def from_dataset(cls, dataset: pa.dataset.Dataset) -> Table:
+    @deprecated("Use Table() constructor instead.")
+    @staticmethod
+    def from_dataset(dataset: pa.dataset.Dataset) -> Table:
         """Turn a :mod:`pyarrow.dataset` ``Dataset`` into a :class:`Table`."""
-        return cls(_InternalRawTable.from_dataset(dataset))
-
-    @classmethod
-    def from_capsule(cls, capsule: Any) -> Table:
-        """Create a :class:`Table` from a PyCapsule exported provider."""
-        provider = _InternalTableProvider.from_capsule(capsule)
-        return cls(provider)
-
-    @classmethod
-    def from_dataframe(cls, df: Any) -> Table:
-        """Create a :class:`Table` from tabular data."""
-        from datafusion.dataframe import DataFrame as DataFrameWrapper
-
-        dataframe = df if isinstance(df, DataFrameWrapper) else DataFrameWrapper(df)
-        return dataframe.into_view()
-
-    @classmethod
-    def from_view(cls, df: Any) -> Table:
-        """Deprecated helper for constructing tables from views."""
-        from datafusion.dataframe import DataFrame as DataFrameWrapper
-
-        if isinstance(df, DataFrameWrapper):
-            df = df.df
-
-        provider = _InternalTableProvider.from_view(df)
-        warnings.warn(
-            "Table.from_view is deprecated; use DataFrame.into_view or "
-            "Table.from_dataframe instead.",
-            category=DeprecationWarning,
-            stacklevel=_FROM_VIEW_WARN_STACKLEVEL,
-        )
-        return cls(provider)
+        return Table(dataset)
 
     @property
     def schema(self) -> pa.Schema:
         """Returns the schema associated with this table."""
-        return self._table.schema
+        return self._inner.schema
 
     @property
     def kind(self) -> str:
         """Returns the kind of table."""
-        return self._table.kind
-
-    def __datafusion_table_provider__(self) -> Any:
-        """Expose the wrapped provider for FFI integrations."""
-        exporter = getattr(self._table, "__datafusion_table_provider__", None)
-        if exporter is None:
-            msg = "Underlying object does not export __datafusion_table_provider__()"
-            raise AttributeError(msg)
-        return exporter()
+        return self._inner.kind
 
 
 class CatalogProvider(ABC):

python/datafusion/context.py

@@ -33,7 +33,6 @@
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
-from datafusion.utils import _normalize_table_provider
 
 from ._internal import RuntimeEnvBuilder as RuntimeEnvBuilderInternal
 from ._internal import SessionConfig as SessionConfigInternal
@@ -770,8 +769,7 @@ def register_table(
             table: DataFusion :class:`Table` or any object implementing
                 ``__datafusion_table_provider__`` to add to the session context.
         """
-        provider = _normalize_table_provider(table)
-        self.ctx.register_table(name, provider)
+        self.ctx.register_table(name, table)
 
     def deregister_table(self, name: str) -> None:
         """Remove a table from the session."""
@@ -1197,7 +1195,7 @@ def read_table(self, table: Table) -> DataFrame:
         :py:class:`~datafusion.catalog.ListingTable`, create a
         :py:class:`~datafusion.dataframe.DataFrame`.
         """
-        return DataFrame(self.ctx.read_table(table.table))
+        return DataFrame(self.ctx.read_table(table._inner))
 
     def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream:
         """Execute the ``plan`` and return the results."""

python/datafusion/utils.py

@@ -19,9 +19,7 @@
 from __future__ import annotations
 
 from importlib import import_module, util
-from typing import TYPE_CHECKING, Any
-
-from datafusion._internal import EXPECTED_PROVIDER_MSG
+from typing import Any
 
 _PYARROW_DATASET_TYPES: tuple[type[Any], ...]
 _dataset_spec = util.find_spec("pyarrow.dataset")
@@ -37,38 +35,3 @@
             if isinstance(value, type) and issubclass(value, dataset_base):
                 dataset_types.add(value)
     _PYARROW_DATASET_TYPES = tuple(dataset_types)
-
-if TYPE_CHECKING:  # pragma: no cover - imported for typing only
-    from datafusion.catalog import Table
-    from datafusion.context import TableProviderExportable
-
-
-def _normalize_table_provider(
-    table: Table | TableProviderExportable | Any,
-) -> Any:
-    """Return the underlying provider for supported table inputs.
-
-    Args:
-        table: A :class:`~datafusion.Table`, object exporting a DataFusion table
-            provider via ``__datafusion_table_provider__``, or compatible
-            :mod:`pyarrow.dataset` implementation.
-
-    Returns:
-        The object expected by the Rust bindings for table registration.
-
-    Raises:
-        TypeError: If ``table`` is not a supported table provider input.
-    """
-    from datafusion.catalog import Table as _Table
-
-    if isinstance(table, _Table):
-        return table.table
-
-    if _PYARROW_DATASET_TYPES and isinstance(table, _PYARROW_DATASET_TYPES):
-        return table
-
-    provider_factory = getattr(table, "__datafusion_table_provider__", None)
-    if callable(provider_factory):
-        return table
-
-    raise TypeError(EXPECTED_PROVIDER_MSG)

python/tests/test_catalog.py

@@ -20,7 +20,7 @@
 import pyarrow as pa
 import pyarrow.dataset as ds
 import pytest
-from datafusion import EXPECTED_PROVIDER_MSG, SessionContext, Table
+from datafusion import SessionContext, Table
 
 
 # Note we take in `database` as a variable even though we don't use
@@ -186,16 +186,6 @@ def test_schema_register_table_with_pyarrow_dataset(ctx: SessionContext):
         schema.deregister_table(table_name)
 
 
-def test_schema_register_table_with_dataframe_errors(ctx: SessionContext):
-    schema = ctx.catalog().schema()
-    df = ctx.from_pydict({"a": [1]})
-
-    with pytest.raises(Exception) as exc_info:
-        schema.register_table("bad", df)
-
-    assert str(exc_info.value) == EXPECTED_PROVIDER_MSG
-
-
 def test_in_end_to_end_python_providers(ctx: SessionContext):
     """Test registering all python providers and running a query against them."""
 

python/tests/test_context.py

@@ -22,7 +22,6 @@
 import pyarrow.dataset as ds
 import pytest
 from datafusion import (
-    EXPECTED_PROVIDER_MSG,
     DataFrame,
     RuntimeEnvBuilder,
     SessionConfig,
@@ -341,20 +340,9 @@ def test_register_table_from_dataframe_into_view(ctx):
     assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
 
 
-def test_table_from_capsule(ctx):
-    df = ctx.from_pydict({"a": [1, 2]})
-    table = df.into_view()
-    capsule = table.__datafusion_table_provider__()
-    table2 = Table.from_capsule(capsule)
-    assert isinstance(table2, Table)
-    ctx.register_table("capsule_tbl", table2)
-    result = ctx.sql("SELECT * FROM capsule_tbl").collect()
-    assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
-
-
 def test_table_from_dataframe(ctx):
     df = ctx.from_pydict({"a": [1, 2]})
-    table = Table.from_dataframe(df)
+    table = Table(df)
     assert isinstance(table, Table)
     ctx.register_table("from_dataframe_tbl", table)
     result = ctx.sql("SELECT * FROM from_dataframe_tbl").collect()
@@ -363,42 +351,13 @@ def test_table_from_dataframe(ctx):
 
 def test_table_from_dataframe_internal(ctx):
     df = ctx.from_pydict({"a": [1, 2]})
-    table = Table.from_dataframe(df.df)
+    table = Table(df.df)
     assert isinstance(table, Table)
     ctx.register_table("from_internal_dataframe_tbl", table)
     result = ctx.sql("SELECT * FROM from_internal_dataframe_tbl").collect()
     assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
 
 
-def test_register_table_capsule_direct(ctx):
-    df = ctx.from_pydict({"a": [1, 2]})
-    provider = df.into_view()
-
-    class CapsuleProvider:
-        def __init__(self, inner):
-            self._inner = inner
-
-        def __datafusion_table_provider__(self):
-            return self._inner.__datafusion_table_provider__()
-
-    ctx.register_table("capsule_direct_tbl", CapsuleProvider(provider))
-    result = ctx.sql("SELECT * FROM capsule_direct_tbl").collect()
-    assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
-
-
-def test_table_from_capsule_invalid():
-    with pytest.raises(RuntimeError):
-        Table.from_capsule(object())
-
-
-def test_register_table_with_dataframe_errors(ctx):
-    df = ctx.from_pydict({"a": [1]})
-    with pytest.raises(TypeError) as exc_info:
-        ctx.register_table("bad", df)
-
-    assert str(exc_info.value) == EXPECTED_PROVIDER_MSG
-
-
 def test_register_dataset(ctx):
     # create a RecordBatch and register it as a pyarrow.dataset.Dataset
     batch = pa.RecordBatch.from_arrays(