feat: enhance TableProvider to support Send trait across multiple modules

Author: kosiew

src/catalog.rs

@@ -52,7 +52,7 @@ pub struct PySchema {
 #[pyclass(name = "RawTable", module = "datafusion.catalog", subclass)]
 #[derive(Clone)]
 pub struct PyTable {
-    pub table: Arc<dyn TableProvider>,
+    pub table: Arc<dyn TableProvider + Send>,
 }
 
 impl From<Arc<dyn CatalogProvider>> for PyCatalog {
@@ -68,11 +68,11 @@ impl From<Arc<dyn SchemaProvider>> for PySchema {
 }
 
 impl PyTable {
-    pub fn new(table: Arc<dyn TableProvider>) -> Self {
+    pub fn new(table: Arc<dyn TableProvider + Send>) -> Self {
         Self { table }
     }
 
-    pub fn table(&self) -> Arc<dyn TableProvider> {
+    pub fn table(&self) -> Arc<dyn TableProvider + Send> {
         self.table.clone()
     }
 }
@@ -206,7 +206,7 @@ impl PySchema {
 
             let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
             let provider: ForeignTableProvider = provider.into();
-            Arc::new(provider) as Arc<dyn TableProvider>
+            Arc::new(provider) as Arc<dyn TableProvider + Send>
         } else {
             match table_provider.extract::<PyTable>() {
                 Ok(py_table) => py_table.table,
@@ -215,7 +215,7 @@ impl PySchema {
                     Err(_) => {
                         let py = table_provider.py();
                         let provider = Dataset::new(&table_provider, py)?;
-                        Arc::new(provider) as Arc<dyn TableProvider>
+                        Arc::new(provider) as Arc<dyn TableProvider + Send>
                     }
                 },
             }
@@ -298,7 +298,7 @@ impl RustWrappedPySchemaProvider {
         }
     }
 
-    fn table_inner(&self, name: &str) -> PyResult<Option<Arc<dyn TableProvider>>> {
+    fn table_inner(&self, name: &str) -> PyResult<Option<Arc<dyn TableProvider + Send>>> {
         Python::with_gil(|py| {
             let provider = self.schema_provider.bind(py);
             let py_table_method = provider.getattr("table")?;
@@ -316,7 +316,7 @@ impl RustWrappedPySchemaProvider {
                 let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
                 let provider: ForeignTableProvider = provider.into();
 
-                Ok(Some(Arc::new(provider) as Arc<dyn TableProvider>))
+                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>() {
@@ -332,7 +332,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>))
+                        Ok(Some(Arc::new(ds) as Arc<dyn TableProvider + Send>))
                     }
                 }
             }
@@ -368,15 +368,32 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
         &self,
         name: &str,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
-        self.table_inner(name).map_err(to_datafusion_err)
+        // Convert from our internal Send type to the trait expected type
+        match self.table_inner(name).map_err(to_datafusion_err)? {
+            Some(table) => {
+                // Safe conversion: we're widening the bounds (removing Send)
+                let raw = Arc::into_raw(table);
+                let wide: *const dyn TableProvider = raw as *const _;
+                let arc = unsafe { Arc::from_raw(wide) };
+                Ok(Some(arc))
+            }
+            None => Ok(None),
+        }
     }
 
     fn register_table(
         &self,
         name: String,
         table: Arc<dyn TableProvider>,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
-        let py_table = PyTable::new(table);
+        // Convert from trait type to our internal Send type
+        let send_table = {
+            let raw = Arc::into_raw(table);
+            let send: *const (dyn TableProvider + Send) = raw as *const _;
+            unsafe { Arc::from_raw(send) }
+        };
+
+        let py_table = PyTable::new(send_table);
         Python::with_gil(|py| {
             let provider = self.schema_provider.bind(py);
             let _ = provider
@@ -405,7 +422,14 @@ 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) => Some(Arc::new(dataset) as Arc<dyn TableProvider>),
+                Ok(dataset) => {
+                    // Convert from our internal Send type to trait expected type
+                    let send_table = Arc::new(dataset) as Arc<dyn TableProvider + Send>;
+                    let raw = Arc::into_raw(send_table);
+                    let wide: *const dyn TableProvider = raw as *const _;
+                    let arc = unsafe { Arc::from_raw(wide) };
+                    Some(arc)
+                }
                 Err(_) => None,
             };
 

src/context.rs

@@ -617,7 +617,7 @@ impl PySessionContext {
 
             let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
             let provider: ForeignTableProvider = provider.into();
-            Arc::new(provider) as Arc<dyn TableProvider>
+            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>() {
@@ -858,7 +858,7 @@ impl PySessionContext {
         dataset: &Bound<'_, PyAny>,
         py: Python,
     ) -> PyDataFusionResult<()> {
-        let table: Arc<dyn TableProvider> = Arc::new(Dataset::new(dataset, py)?);
+        let table: Arc<dyn TableProvider + Send> = Arc::new(Dataset::new(dataset, py)?);
 
         self.ctx.register_table(name, table)?;
 

src/dataframe.rs

@@ -48,7 +48,7 @@ use crate::expr::sort_expr::to_sort_expressions;
 use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
-use crate::table::PyTableProvider;
+pub use crate::table::PyTableProvider;
 use crate::utils::{
     get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, validate_pycapsule, wait_for_future,
 };
@@ -268,7 +268,7 @@ impl PyDataFrame {
         }
     }
 
-    pub(crate) fn into_view_provider(&self) -> Arc<dyn TableProvider> {
+    pub(crate) fn into_view_provider(&self) -> Arc<dyn TableProvider + Send> {
         self.df.as_ref().clone().into_view()
     }
 

src/table.rs

@@ -23,20 +23,28 @@ use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
 
+use crate::catalog::PyTable;
 use crate::dataframe::PyDataFrame;
 use crate::errors::{py_datafusion_err, PyDataFusionResult};
 use crate::utils::{get_tokio_runtime, validate_pycapsule};
 
 /// Represents a table provider that can be registered with DataFusion
 #[pyclass(name = "TableProvider", module = "datafusion")]
+#[derive(Clone)]
 pub struct PyTableProvider {
     pub(crate) provider: Arc<dyn TableProvider + Send>,
 }
 
 impl PyTableProvider {
-    pub(crate) fn new(provider: Arc<dyn TableProvider>) -> Self {
+    pub(crate) fn new(provider: Arc<dyn TableProvider + Send>) -> Self {
         Self { provider }
     }
+
+    /// Return a `PyTable` wrapper around this provider so callers can call
+    /// `as_table().table()` to get the underlying `Arc<dyn TableProvider + Send>`.
+    pub fn as_table(&self) -> PyTable {
+        PyTable::new(Arc::clone(&self.provider))
+    }
 }
 
 #[pymethods]

src/udtf.rs

@@ -87,7 +87,7 @@ impl PyTableFunction {
 fn call_python_table_function(
     func: &Arc<PyObject>,
     args: &[Expr],
-) -> DataFusionResult<Arc<dyn TableProvider>> {
+) -> DataFusionResult<Arc<dyn TableProvider + Send>> {
     let args = args
         .iter()
         .map(|arg| PyExpr::from(arg.clone()))
@@ -107,7 +107,7 @@ fn call_python_table_function(
             let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
             let provider: ForeignTableProvider = provider.into();
 
-            Ok(Arc::new(provider) as Arc<dyn TableProvider>)
+            Ok(Arc::new(provider) as Arc<dyn TableProvider + Send>)
         } else {
             Err(PyNotImplementedError::new_err(
                 "__datafusion_table_provider__ does not exist on Table Provider object.",
@@ -121,7 +121,14 @@ impl TableFunctionImpl for PyTableFunction {
     fn call(&self, args: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
         match &self.inner {
             PyTableFunctionInner::FFIFunction(func) => func.call(args),
-            PyTableFunctionInner::PythonFunction(obj) => call_python_table_function(obj, args),
+            PyTableFunctionInner::PythonFunction(obj) => {
+                let send_result = call_python_table_function(obj, args)?;
+                // Convert from our Send type to the trait expected type
+                let raw = Arc::into_raw(send_result);
+                let wide: *const dyn TableProvider = raw as *const _;
+                let arc = unsafe { Arc::from_raw(wide) };
+                Ok(arc)
+            }
         }
     }
 }