feat: refactor TableProvider integration and add new TableProvider class

Author: kosiew

python/datafusion/__init__.py

@@ -33,7 +33,7 @@
 from . import functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
-from ._internal import Config, TableProvider
+from ._internal import Config
 from .catalog import Catalog, Database, Table
 from .col import col, column
 from .common import (
@@ -54,6 +54,7 @@
 from .io import read_avro, read_csv, read_json, read_parquet
 from .plan import ExecutionPlan, LogicalPlan
 from .record_batch import RecordBatch, RecordBatchStream
+from .table_provider import TableProvider
 from .user_defined import (
     Accumulator,
     AggregateUDF,

python/datafusion/dataframe.py

@@ -52,8 +52,8 @@
     import polars as pl
     import pyarrow as pa
 
-    from datafusion._internal import TableProvider
     from datafusion._internal import expr as expr_internal
+    from datafusion.table_provider import TableProvider
 
 from enum import Enum
 
@@ -316,7 +316,9 @@ def into_view(self) -> TableProvider:
         ``TableProvider.from_dataframe`` calls this method under the hood,
         and the older ``TableProvider.from_view`` helper is deprecated.
         """
-        return self.df.into_view()
+        from datafusion.table_provider import TableProvider as _TableProvider
+
+        return _TableProvider(self.df.into_view())
 
     def __getitem__(self, key: str | list[str]) -> DataFrame:
         """Return a new :py:class`DataFrame` with the specified column or columns.

python/datafusion/table_provider.py

@@ -0,0 +1,97 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Wrapper helpers for :mod:`datafusion._internal.TableProvider`."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import datafusion._internal as df_internal
+
+_InternalTableProvider = df_internal.TableProvider
+
+
+class TableProvider:
+    """High level wrapper around :mod:`datafusion._internal.TableProvider`."""
+
+    __slots__ = ("_table_provider",)
+
+    def __init__(self, table_provider: _InternalTableProvider) -> None:
+        """Wrap a low level :class:`~datafusion._internal.TableProvider`."""
+        if isinstance(table_provider, TableProvider):
+            table_provider = table_provider._table_provider
+
+        if not isinstance(table_provider, _InternalTableProvider):
+            msg = "Expected a datafusion._internal.TableProvider instance."
+            raise TypeError(msg)
+
+        self._table_provider = table_provider
+
+    # ------------------------------------------------------------------
+    # constructors
+    # ------------------------------------------------------------------
+    @classmethod
+    def _wrap(cls, provider: _InternalTableProvider | TableProvider) -> TableProvider:
+        if isinstance(provider, cls):
+            return provider
+        return cls(provider)
+
+    @classmethod
+    def from_capsule(cls, capsule: Any) -> TableProvider:
+        """Create a :class:`TableProvider` from a PyCapsule."""
+        provider = _InternalTableProvider.from_capsule(capsule)
+        return cls(provider)
+
+    @classmethod
+    def from_dataframe(cls, df: Any) -> TableProvider:
+        """Create a :class:`TableProvider` from a :class:`DataFrame`."""
+        from datafusion.dataframe import DataFrame as DataFrameWrapper
+
+        if isinstance(df, DataFrameWrapper):
+            df = df.df
+
+        provider = _InternalTableProvider.from_dataframe(df)
+        return cls(provider)
+
+    @classmethod
+    def from_view(cls, df: Any) -> TableProvider:
+        """Create a :class:`TableProvider` from a DataFrame view."""
+        from datafusion.dataframe import DataFrame as DataFrameWrapper
+
+        if isinstance(df, DataFrameWrapper):
+            df = df.df
+
+        provider = _InternalTableProvider.from_view(df)
+        return cls(provider)
+
+    # ------------------------------------------------------------------
+    # passthrough helpers
+    # ------------------------------------------------------------------
+    def __getattr__(self, name: str) -> Any:
+        """Delegate attribute lookup to the wrapped provider."""
+        return getattr(self._table_provider, name)
+
+    def __repr__(self) -> str:  # pragma: no cover - simple delegation
+        """Return a representation of the wrapped provider."""
+        return repr(self._table_provider)
+
+    def __datafusion_table_provider__(self) -> Any:
+        """Expose the wrapped provider for FFI integrations."""
+        return self._table_provider.__datafusion_table_provider__()
+
+
+__all__ = ["TableProvider"]

python/tests/test_context.py

@@ -350,7 +350,7 @@ def test_table_provider_from_capsule(ctx):
 
 
 def test_table_provider_from_dataframe(ctx):
-    df = ctx.from_pydict({"a": [1, 2]}).df
+    df = ctx.from_pydict({"a": [1, 2]})
     provider = TableProvider.from_dataframe(df)
     ctx.register_table("from_dataframe_tbl", provider)
     result = ctx.sql("SELECT * FROM from_dataframe_tbl").collect()
@@ -380,12 +380,12 @@ def test_table_provider_from_capsule_invalid():
 
 def test_register_table_with_dataframe_errors(ctx):
     df = ctx.from_pydict({"a": [1]})
-    with pytest.raises(Exception) as exc_info:  # noqa: B017
+    with pytest.raises(Exception) as exc_info:
         ctx.register_table("bad", df)
 
-    assert (
-        str(exc_info.value)
-        == 'Expected a Table or TableProvider. Convert DataFrames with "DataFrame.into_view()" or "TableProvider.from_dataframe()".'
+    assert str(exc_info.value) == (
+        "Expected a Table or TableProvider. Convert DataFrames with "
+        '"DataFrame.into_view()" or "TableProvider.from_dataframe()".'
     )