feat: refactor table provider handling and update tests for improved access methods

Author: kosiew

src/catalog.rs

@@ -17,7 +17,7 @@
 
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
-use crate::table::pyany_to_table_provider;
+use crate::table::PyTableProvider;
 use crate::utils::{validate_pycapsule, wait_for_future};
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -28,6 +28,7 @@ use datafusion::{
     datasource::{TableProvider, TableType},
 };
 use datafusion_ffi::schema_provider::{FFI_SchemaProvider, ForeignSchemaProvider};
+use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::exceptions::PyKeyError;
 use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
@@ -196,7 +197,29 @@ impl PySchema {
     }
 
     fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> {
-        let provider = pyany_to_table_provider(&table_provider)?;
+        let provider = if table_provider.hasattr("__datafusion_table_provider__")? {
+            let capsule = table_provider
+                .getattr("__datafusion_table_provider__")?
+                .call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_table_provider")?;
+
+            let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+            let provider: ForeignTableProvider = provider.into();
+            Arc::new(provider) as Arc<dyn TableProvider + Send>
+        } else {
+            match table_provider.extract::<PyTable>() {
+                Ok(py_table) => py_table.table,
+                Err(_) => match table_provider.extract::<PyTableProvider>() {
+                    Ok(py_provider) => py_provider.into_inner(),
+                    Err(_) => {
+                        let py = table_provider.py();
+                        let provider = Dataset::new(&table_provider, py)?;
+                        Arc::new(provider) as Arc<dyn TableProvider + Send>
+                    }
+                },
+            }
+        };
 
         let _ = self
             .schema
@@ -285,7 +308,34 @@ impl RustWrappedPySchemaProvider {
                 return Ok(None);
             }
 
-            pyany_to_table_provider(&py_table).map(Some)
+            if py_table.hasattr("__datafusion_table_provider__")? {
+                let capsule = py_table.getattr("__datafusion_table_provider__")?.call0()?;
+                let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+                validate_pycapsule(capsule, "datafusion_table_provider")?;
+
+                let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+                let provider: ForeignTableProvider = provider.into();
+
+                Ok(Some(Arc::new(provider) as Arc<dyn TableProvider + Send>))
+            } else {
+                if let Ok(inner_table) = py_table.getattr("table") {
+                    if let Ok(inner_table) = inner_table.extract::<PyTable>() {
+                        return Ok(Some(inner_table.table));
+                    }
+                }
+
+                if let Ok(py_provider) = py_table.extract::<PyTableProvider>() {
+                    return Ok(Some(py_provider.into_inner()));
+                }
+
+                match py_table.extract::<PyTable>() {
+                    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>))
+                    }
+                }
+            }
         })
     }
 }

src/context.rs

@@ -34,14 +34,14 @@ use pyo3::prelude::*;
 use crate::catalog::{PyCatalog, PyTable, RustWrappedPyCatalogProvider};
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
-use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
+use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
 use crate::expr::sort_expr::PySortExpr;
 use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::exceptions::py_value_err;
 use crate::sql::logical::PyLogicalPlan;
 use crate::store::StorageContexts;
-use crate::table::pyany_to_table_provider;
+use crate::table::PyTableProvider;
 use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
@@ -72,6 +72,7 @@ use datafusion::prelude::{
     AvroReadOptions, CsvReadOptions, DataFrame, NdJsonReadOptions, ParquetReadOptions,
 };
 use datafusion_ffi::catalog_provider::{FFI_CatalogProvider, ForeignCatalogProvider};
+use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple, PyType};
 use pyo3::IntoPyObjectExt;
 use tokio::task::JoinHandle;
@@ -607,9 +608,26 @@ impl PySessionContext {
         name: &str,
         table_provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        let provider = pyany_to_table_provider(&table_provider).map_err(|_| {
-            PyDataFusionError::Common("Expected a Table or TableProvider.".to_string())
-        })?;
+        let provider = if table_provider.hasattr("__datafusion_table_provider__")? {
+            let capsule = table_provider
+                .getattr("__datafusion_table_provider__")?
+                .call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_table_provider")?;
+
+            let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+            let provider: ForeignTableProvider = provider.into();
+            Arc::new(provider) as Arc<dyn TableProvider + Send>
+        } else if let Ok(py_table) = table_provider.extract::<PyTable>() {
+            py_table.table()
+        } else if let Ok(py_provider) = table_provider.extract::<PyTableProvider>() {
+            py_provider.into_inner()
+        } else {
+            return Err(crate::errors::PyDataFusionError::Common(
+                "Expected a Table or TableProvider.".to_string(),
+            ));
+        };
+
         self.ctx.register_table(name, provider)?;
         Ok(())
     }

src/table.rs

@@ -42,11 +42,24 @@ impl PyTableProvider {
         Self { provider }
     }
 
-    /// Return a `PyTable` wrapper around this provider so callers can call
-    /// `as_table().table()` to get the underlying `Arc<dyn TableProvider + Send>`.
+    /// Return a `PyTable` wrapper around this provider.
+    ///
+    /// Historically callers chained `as_table().table()` to access the
+    /// underlying `Arc<dyn TableProvider + Send>`. Prefer [`as_arc`] or
+    /// [`into_inner`] for direct access instead.
     pub fn as_table(&self) -> PyTable {
         PyTable::new(Arc::clone(&self.provider))
     }
+
+    /// Return a clone of the inner [`TableProvider`].
+    pub fn as_arc(&self) -> Arc<dyn TableProvider + Send> {
+        Arc::clone(&self.provider)
+    }
+
+    /// Consume this wrapper and return the inner [`TableProvider`].
+    pub fn into_inner(self) -> Arc<dyn TableProvider + Send> {
+        self.provider
+    }
 }
 
 #[pymethods]

tests/dataframe_into_view.rs

@@ -26,8 +26,7 @@ fn dataframe_into_view_returns_table_provider() {
 
     // Register the view in a new context and ensure it can be queried.
     let ctx = SessionContext::new();
-    ctx.register_table("view", provider.as_table().table())
-        .unwrap();
+    ctx.register_table("view", provider.into_inner()).unwrap();
 
     let rt = tokio::runtime::Runtime::new().unwrap();
     let batches = rt.block_on(async {