refactor: simplify table provider conversions by introducing utility functions

Author: kosiew

src/catalog.rs

@@ -18,7 +18,10 @@
 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, validate_pycapsule, wait_for_future};
+use crate::utils::{
+    table_provider_from_pycapsule, table_provider_send_to_table_provider, table_provider_to_send,
+    validate_pycapsule, wait_for_future,
+};
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion::common::DataFusionError;
@@ -354,13 +357,7 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
     ) -> 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) => {
-                // 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))
-            }
+            Some(table) => Ok(Some(table_provider_send_to_table_provider(table))),
             None => Ok(None),
         }
     }
@@ -371,11 +368,7 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
         table: Arc<dyn TableProvider>,
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
         // 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 send_table = table_provider_to_send(table);
 
         let py_table = PyTable::new(send_table);
         Python::with_gil(|py| {
@@ -407,12 +400,8 @@ impl SchemaProvider for RustWrappedPySchemaProvider {
             // Otherwise, return None.
             let dataset = match Dataset::new(&table, py) {
                 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)
+                    Some(table_provider_send_to_table_provider(send_table))
                 }
                 Err(_) => None,
             };

src/udtf.rs

@@ -21,7 +21,9 @@ 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, validate_pycapsule};
+use crate::utils::{
+    table_provider_from_pycapsule, table_provider_send_to_table_provider, validate_pycapsule,
+};
 use datafusion::catalog::{TableFunctionImpl, TableProvider};
 use datafusion::error::Result as DataFusionResult;
 use datafusion::logical_expr::Expr;
@@ -113,11 +115,7 @@ impl TableFunctionImpl for PyTableFunction {
             PyTableFunctionInner::FFIFunction(func) => func.call(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)
+                Ok(table_provider_send_to_table_provider(send_result))
             }
         }
     }

src/utils.rs

@@ -122,6 +122,40 @@ 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>>> {