refactor: remove Send marker from TableProvider trait objects for improved type flexibility

Author: kosiew

src/catalog.rs

@@ -19,10 +19,7 @@ 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,
-};
+use crate::utils::{table_provider_from_pycapsule, validate_pycapsule, wait_for_future};
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion::common::DataFusionError;
@@ -55,7 +52,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 {
@@ -71,11 +68,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()
     }
 }
@@ -215,7 +212,7 @@ impl PySchema {
         } 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
@@ -295,7 +292,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")?;
@@ -322,7 +319,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>))
                     }
                 }
             }
@@ -358,22 +355,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
@@ -402,10 +392,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

@@ -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/table.rs

@@ -33,30 +33,30 @@ use crate::utils::{get_tokio_runtime, validate_pycapsule};
 #[pyclass(name = "TableProvider", module = "datafusion")]
 #[derive(Clone)]
 pub struct PyTableProvider {
-    pub(crate) provider: Arc<dyn TableProvider + Send>,
+    pub(crate) provider: Arc<dyn TableProvider>,
 }
 
 impl PyTableProvider {
-    pub(crate) fn new(provider: Arc<dyn TableProvider + Send>) -> Self {
+    pub(crate) fn new(provider: Arc<dyn TableProvider>) -> Self {
         Self { provider }
     }
 
     /// 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
+    /// underlying [`Arc<dyn TableProvider>`]. 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> {
+    pub fn as_arc(&self) -> Arc<dyn TableProvider> {
         Arc::clone(&self.provider)
     }
 
     /// Consume this wrapper and return the inner [`TableProvider`].
-    pub fn into_inner(self) -> Arc<dyn TableProvider + Send> {
+    pub fn into_inner(self) -> Arc<dyn TableProvider> {
         self.provider
     }
 }
@@ -110,7 +110,8 @@ impl PyTableProvider {
         let name = CString::new("datafusion_table_provider").unwrap();
 
         let runtime = get_tokio_runtime().0.handle().clone();
-        let provider = FFI_TableProvider::new(Arc::clone(&self.provider), false, Some(runtime));
+        let provider: Arc<dyn TableProvider + Send> = self.provider.clone();
+        let provider = FFI_TableProvider::new(provider, false, Some(runtime));
 
         PyCapsule::new(py, provider, Some(name.clone()))
     }

src/udtf.rs

@@ -21,9 +21,7 @@ use std::sync::Arc;
 use crate::errors::{py_datafusion_err, to_datafusion_err};
 use crate::expr::PyExpr;
 use crate::table::PyTableProvider;
-use crate::utils::{
-    table_provider_from_pycapsule, table_provider_send_to_table_provider, validate_pycapsule,
-};
+use crate::utils::{table_provider_from_pycapsule, validate_pycapsule};
 use datafusion::catalog::{TableFunctionImpl, TableProvider};
 use datafusion::error::Result as DataFusionResult;
 use datafusion::logical_expr::Expr;
@@ -88,7 +86,7 @@ impl PyTableFunction {
 fn call_python_table_function(
     func: &Arc<PyObject>,
     args: &[Expr],
-) -> DataFusionResult<Arc<dyn TableProvider + Send>> {
+) -> DataFusionResult<Arc<dyn TableProvider>> {
     let args = args
         .iter()
         .map(|arg| PyExpr::from(arg.clone()))
@@ -113,10 +111,7 @@ 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) => {
-                let send_result = call_python_table_function(obj, args)?;
-                Ok(table_provider_send_to_table_provider(send_result))
-            }
+            PyTableFunctionInner::PythonFunction(obj) => call_python_table_function(obj, args),
         }
     }
 }

src/utils.rs

@@ -122,51 +122,17 @@ pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyRe
     Ok(())
 }
 
-/// Convert a [`TableProvider`] wrapped in an [`Arc`] with a `Send` auto trait into one
-/// without the marker.
-///
-/// # Safety
-///
-/// Removing `Send` from a trait object only relaxes the bounds. The underlying vtable is
-/// unchanged, so it is safe to reuse the pointer produced by [`Arc::into_raw`].
-pub(crate) fn table_provider_send_to_table_provider(
-    table: Arc<dyn TableProvider + Send>,
-) -> Arc<dyn TableProvider> {
-    let raw: *const (dyn TableProvider + Send) = Arc::into_raw(table);
-    // SAFETY: `Send` is an auto trait with no associated data, so the trait object layout
-    // is identical and the pointer may be reinterpreted without changing the reference
-    // count.
-    unsafe { Arc::from_raw(raw as *const dyn TableProvider) }
-}
-
-/// Convert a [`TableProvider`] wrapped in an [`Arc`] into one that also carries the `Send`
-/// auto trait.
-///
-/// # Safety
-///
-/// DataFusion's `TableProvider` trait requires `Send`, so the underlying provider implements
-/// the marker. This allows us to reinterpret the pointer as a `TableProvider + Send` trait
-/// object.
-pub(crate) fn table_provider_to_send(
-    table: Arc<dyn TableProvider>,
-) -> Arc<dyn TableProvider + Send> {
-    let raw: *const dyn TableProvider = Arc::into_raw(table);
-    // SAFETY: The underlying type implements `Send`, so the pointer can be safely treated as
-    // a `TableProvider + Send` trait object.
-    unsafe { Arc::from_raw(raw as *const (dyn TableProvider + Send)) }
-}
-
 pub(crate) fn table_provider_from_pycapsule(
     obj: &Bound<PyAny>,
-) -> PyResult<Option<Arc<dyn TableProvider + Send>>> {
+) -> PyResult<Option<Arc<dyn TableProvider>>> {
     if obj.hasattr("__datafusion_table_provider__")? {
         let capsule = obj.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>))
+        Ok(Some(Arc::new(provider) as Arc<dyn TableProvider>))
     } else {
         Ok(None)
     }