TableProvider refactor & PyDataFrame integration

* Removed unused helpers (`extract_table_provider`, `_wrap`) and dead code to simplify maintenance.
* Consolidated and streamlined table-provider extraction and registration logic; improved error handling and replaced a hardcoded error message with `EXPECTED_PROVIDER_MSG`.
* Marked `from_view` as deprecated; updated deprecation message formatting and adjusted the warning `stacklevel` so it points to caller code.
* Removed the `Send` marker from TableProvider trait objects to increase type flexibility — review threading assumptions.
* Added type hints to `register_schema` and `deregister_table` methods.
* Adjusted tests and exceptions (e.g., changed one test to expect `RuntimeError`) and updated test coverage accordingly.
* Introduced a refactored `TableProvider` class and enhanced Python integration by adding support for extracting `PyDataFrame` in `PySchema`.

Notes:

* Consumers should migrate away from `TableProvider::from_view` to the new TableProvider integration.
* Audit any code relying on `Send` for trait objects passed across threads.
* Update downstream tests and documentation to reflect the changed exception types and deprecation.

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, EXPECTED_PROVIDER_MSG
 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,
@@ -76,6 +77,7 @@
     "DFSchema",
     "DataFrame",
     "Database",
+    "EXPECTED_PROVIDER_MSG",
     "ExecutionPlan",
     "Expr",
     "LogicalPlan",

python/datafusion/catalog.py

@@ -85,7 +85,11 @@ def database(self, name: str = "public") -> Schema:
         """Returns the database with the given ``name`` from this catalog."""
         return self.schema(name)
 
-    def register_schema(self, name, schema) -> Schema | None:
+    def register_schema(
+        self,
+        name: str,
+        schema: Schema | SchemaProvider | SchemaProviderExportable,
+    ) -> Schema | None:
         """Register a schema with this catalog."""
         if isinstance(schema, Schema):
             return self.catalog.register_schema(name, schema._raw_schema)
@@ -126,7 +130,7 @@ def table(self, name: str) -> Table:
         return Table(self._raw_schema.table(name))
 
     def register_table(
-        self, name, table: Table | TableProvider | TableProviderExportable
+        self, name: str, table: Table | TableProvider | TableProviderExportable
     ) -> None:
         """Register a table or table provider in this schema.
 
@@ -240,7 +244,7 @@ def register_table(  # noqa: B027
         and treated as :class:`TableProvider` instances.
         """
 
-    def deregister_table(self, name, cascade: bool) -> None:  # noqa: B027
+    def deregister_table(self, name: str, cascade: bool) -> None:  # noqa: B027
         """Remove a table from this schema.
 
         This method is optional. If your schema provides a fixed list of tables, you do

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,102 @@
+# 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
+
+import warnings
+from typing import Any
+
+import datafusion._internal as df_internal
+from datafusion._internal import EXPECTED_PROVIDER_MSG
+
+_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):
+            raise TypeError(EXPECTED_PROVIDER_MSG)
+
+        self._table_provider = table_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:
+        """Deprecated.
+
+        Use :meth:`DataFrame.into_view` or :meth:`TableProvider.from_dataframe`.
+        """
+        from datafusion.dataframe import DataFrame as DataFrameWrapper
+
+        if isinstance(df, DataFrameWrapper):
+            df = df.df
+
+        provider = _InternalTableProvider.from_view(df)
+        warnings.warn(
+            "TableProvider.from_view is deprecated; use DataFrame.into_view or "
+            "TableProvider.from_dataframe instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        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 __dir__(self) -> list[str]:
+        """Expose delegated attributes via :func:`dir`."""
+        return dir(self._table_provider) + super().__dir__()
+
+    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

@@ -23,6 +23,7 @@
 import pytest
 from datafusion import (
     DataFrame,
+    EXPECTED_PROVIDER_MSG,
     RuntimeEnvBuilder,
     SessionConfig,
     SessionContext,
@@ -350,7 +351,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()
@@ -374,19 +375,16 @@ def __datafusion_table_provider__(self):
 
 
 def test_table_provider_from_capsule_invalid():
-    with pytest.raises(Exception):  # noqa: B017
+    with pytest.raises(RuntimeError):
         TableProvider.from_capsule(object())
 
 
 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_PROVIDER_MSG
 
 
 def test_register_dataset(ctx):

src/catalog.rs

@@ -15,12 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::table::PyTableProvider;
 use crate::utils::{
-    table_provider_from_pycapsule, table_provider_send_to_table_provider, table_provider_to_send,
-    validate_pycapsule, wait_for_future,
+    table_provider_from_pycapsule, validate_pycapsule, wait_for_future, EXPECTED_PROVIDER_MSG,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -54,7 +54,7 @@ pub struct PySchema {
 #[pyclass(name = "RawTable", module = "datafusion.catalog", subclass)]
 #[derive(Clone)]
 pub struct PyTable {
-    pub table: Arc<dyn TableProvider + Send>,
+    pub table: Arc<dyn TableProvider>,
 }
 
 impl From<Arc<dyn CatalogProvider>> for PyCatalog {
@@ -70,11 +70,11 @@ impl From<Arc<dyn SchemaProvider>> for PySchema {
 }
 
 impl PyTable {
-    pub fn new(table: Arc<dyn TableProvider + Send>) -> Self {
+    pub fn new(table: Arc<dyn TableProvider>) -> Self {
         Self { table }
     }
 
-    pub fn table(&self) -> Arc<dyn TableProvider + Send> {
+    pub fn table(&self) -> Arc<dyn TableProvider> {
         self.table.clone()
     }
 }
@@ -203,12 +203,14 @@ impl PySchema {
             py_table.table
         } else if let Ok(py_provider) = table_provider.extract::<PyTableProvider>() {
             py_provider.into_inner()
+        } else if table_provider.extract::<PyDataFrame>().is_ok() {
+            return Err(PyDataFusionError::Common(EXPECTED_PROVIDER_MSG.to_string()).into());
         } else if let Some(provider) = table_provider_from_pycapsule(&table_provider)? {
             provider
         } else {
             let py = table_provider.py();
             let provider = Dataset::new(&table_provider, py)?;
-            Arc::new(provider) as Arc<dyn TableProvider + Send>
+            Arc::new(provider) as Arc<dyn TableProvider>
         };
 
         let _ = self
@@ -288,7 +290,7 @@ impl RustWrappedPySchemaProvider {
         }
     }
 
-    fn table_inner(&self, name: &str) -> PyResult<Option<Arc<dyn TableProvider + Send>>> {
+    fn table_inner(&self, name: &str) -> PyResult<Option<Arc<dyn TableProvider>>> {
         Python::with_gil(|py| {
             let provider = self.schema_provider.bind(py);
             let py_table_method = provider.getattr("table")?;
@@ -315,7 +317,7 @@ impl RustWrappedPySchemaProvider {
                     Ok(py_table) => Ok(Some(py_table.table)),
                     Err(_) => {
                         let ds = Dataset::new(&py_table, py).map_err(py_datafusion_err)?;
-                        Ok(Some(Arc::new(ds) as Arc<dyn TableProvider + Send>))
+                        Ok(Some(Arc::new(ds) as Arc<dyn TableProvider>))
                     }
                 }
             }
@@ -351,22 +353,15 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
         &self,
         name: &str,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
-        // Convert from our internal Send type to the trait expected type
-        match self.table_inner(name).map_err(to_datafusion_err)? {
-            Some(table) => Ok(Some(table_provider_send_to_table_provider(table))),
-            None => Ok(None),
-        }
+        self.table_inner(name).map_err(to_datafusion_err)
     }
 
     fn register_table(
         &self,
         name: String,
         table: Arc<dyn TableProvider>,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
-        // Convert from trait type to our internal Send type
-        let send_table = table_provider_to_send(table);
-
-        let py_table = PyTable::new(send_table);
+        let py_table = PyTable::new(table);
         Python::with_gil(|py| {
             let provider = self.schema_provider.bind(py);
             let _ = provider
@@ -395,10 +390,7 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
             // If we can turn this table provider into a `Dataset`, return it.
             // Otherwise, return None.
             let dataset = match Dataset::new(&table, py) {
-                Ok(dataset) => {
-                    let send_table = Arc::new(dataset) as Arc<dyn TableProvider + Send>;
-                    Some(table_provider_send_to_table_provider(send_table))
-                }
+                Ok(dataset) => Some(Arc::new(dataset) as Arc<dyn TableProvider>),
                 Err(_) => None,
             };
 

src/context.rs

@@ -48,7 +48,7 @@ use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
 use crate::utils::{
     get_global_ctx, get_tokio_runtime, table_provider_from_pycapsule, validate_pycapsule,
-    wait_for_future,
+    wait_for_future, EXPECTED_PROVIDER_MSG,
 };
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
@@ -618,7 +618,7 @@ impl PySessionContext {
             provider
         } else {
             return Err(crate::errors::PyDataFusionError::Common(
-                "Expected a Table or TableProvider. Convert DataFrames with \"DataFrame.into_view()\" or \"TableProvider.from_dataframe()\".".to_string(),
+                EXPECTED_PROVIDER_MSG.to_string(),
             ));
         };
 
@@ -852,7 +852,7 @@ impl PySessionContext {
         dataset: &Bound<'_, PyAny>,
         py: Python,
     ) -> PyDataFusionResult<()> {
-        let table: Arc<dyn TableProvider + Send> = Arc::new(Dataset::new(dataset, py)?);
+        let table: Arc<dyn TableProvider> = Arc::new(Dataset::new(dataset, py)?);
 
         self.ctx.register_table(name, table)?;
 

src/dataframe.rs

@@ -268,7 +268,7 @@ impl PyDataFrame {
         }
     }
 
-    pub(crate) fn to_view_provider(&self) -> Arc<dyn TableProvider + Send> {
+    pub(crate) fn to_view_provider(&self) -> Arc<dyn TableProvider> {
         self.df.as_ref().clone().into_view()
     }
 

src/lib.rs

@@ -81,6 +81,8 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     // Initialize logging
     pyo3_log::init();
 
+    m.add("EXPECTED_PROVIDER_MSG", crate::utils::EXPECTED_PROVIDER_MSG)?;
+
     // Register the python classes
     m.add_class::<context::PyRuntimeEnvBuilder>()?;
     m.add_class::<context::PySessionConfig>()?;