fix: refactor foreign table provider handling for improved clarity and safety

Author: kosiew

src/catalog.rs

@@ -17,7 +17,9 @@
 
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
-use crate::utils::{get_tokio_runtime, validate_pycapsule, wait_for_future};
+use crate::utils::{
+    foreign_table_provider_from_capsule, get_tokio_runtime, validate_pycapsule, wait_for_future,
+};
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion::common::DataFusionError;
@@ -27,7 +29,7 @@ use datafusion::{
     datasource::{TableProvider, TableType},
 };
 use datafusion_ffi::schema_provider::{FFI_SchemaProvider, ForeignSchemaProvider};
-use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
+use datafusion_ffi::table_provider::FFI_TableProvider;
 use pyo3::exceptions::PyKeyError;
 use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
@@ -204,11 +206,9 @@ impl PySchema {
                 .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>
+            let provider = foreign_table_provider_from_capsule(capsule)?;
+            let provider: Arc<dyn TableProvider> = Arc::new(provider);
+            provider
         } else {
             let py = table_provider.py();
             let provider = Dataset::new(&table_provider, py)?;
@@ -323,12 +323,10 @@ impl RustWrappedPySchemaProvider {
             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();
+                let provider = foreign_table_provider_from_capsule(capsule)?;
+                let provider: Arc<dyn TableProvider> = Arc::new(provider);
 
-                Ok(Some(Arc::new(provider) as Arc<dyn TableProvider>))
+                Ok(Some(provider))
             } else {
                 if let Ok(inner_table) = py_table.getattr("table") {
                     if let Ok(inner_table) = inner_table.extract::<PyTable>() {

src/context.rs

@@ -45,7 +45,10 @@ use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
-use crate::utils::{get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future};
+use crate::utils::{
+    foreign_table_provider_from_capsule, get_global_ctx, get_tokio_runtime, validate_pycapsule,
+    wait_for_future,
+};
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
@@ -71,7 +74,6 @@ 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;
@@ -654,12 +656,10 @@ impl PySessionContext {
         if provider.hasattr("__datafusion_table_provider__")? {
             let capsule = 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();
+            let provider = foreign_table_provider_from_capsule(capsule)?;
+            let provider: Arc<dyn TableProvider> = Arc::new(provider);
 
-            let _ = self.ctx.register_table(name, Arc::new(provider))?;
+            let _ = self.ctx.register_table(name, provider)?;
 
             Ok(())
         } else {
@@ -1113,12 +1113,10 @@ impl PySessionContext {
         if table.hasattr("__datafusion_table_provider__")? {
             let capsule = 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();
+            let provider = foreign_table_provider_from_capsule(capsule)?;
+            let provider: Arc<dyn TableProvider> = Arc::new(provider);
 
-            let df = self.ctx.read_table(Arc::new(provider))?;
+            let df = self.ctx.read_table(provider)?;
             Ok(PyDataFrame::new(df))
         } else {
             Err(crate::errors::PyDataFusionError::Common(

src/udtf.rs

@@ -21,11 +21,10 @@ use std::sync::Arc;
 use crate::dataframe::PyTableProvider;
 use crate::errors::{py_datafusion_err, to_datafusion_err};
 use crate::expr::PyExpr;
-use crate::utils::validate_pycapsule;
+use crate::utils::{foreign_table_provider_from_capsule, validate_pycapsule};
 use datafusion::catalog::{TableFunctionImpl, TableProvider};
 use datafusion::error::Result as DataFusionResult;
 use datafusion::logical_expr::Expr;
-use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use datafusion_ffi::udtf::{FFI_TableFunction, ForeignTableFunction};
 use pyo3::exceptions::PyNotImplementedError;
 use pyo3::types::{PyCapsule, PyTuple};
@@ -102,12 +101,10 @@ fn call_python_table_function(
         if provider.hasattr("__datafusion_table_provider__")? {
             let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?;
             let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-            validate_pycapsule(capsule, "datafusion_table_provider")?;
+            let provider = foreign_table_provider_from_capsule(capsule)?;
+            let provider: Arc<dyn TableProvider> = Arc::new(provider);
 
-            let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
-            let provider: ForeignTableProvider = provider.into();
-
-            Ok(Arc::new(provider) as Arc<dyn TableProvider>)
+            Ok(provider)
         } else {
             Err(PyNotImplementedError::new_err(
                 "__datafusion_table_provider__ does not exist on Table Provider object.",

src/utils.rs

@@ -23,6 +23,7 @@ use crate::{
 use datafusion::{
     common::ScalarValue, execution::context::SessionContext, logical_expr::Volatility,
 };
+use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::prelude::*;
 use pyo3::{exceptions::PyValueError, types::PyCapsule};
 use std::{future::Future, sync::OnceLock, time::Duration};
@@ -116,6 +117,13 @@ pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyRe
     Ok(())
 }
 
+pub(crate) fn foreign_table_provider_from_capsule(
+    capsule: &Bound<PyCapsule>,
+) -> PyResult<ForeignTableProvider> {
+    validate_pycapsule(capsule, "datafusion_table_provider")?;
+    Ok(unsafe { capsule.reference::<FFI_TableProvider>() }.into())
+}
+
 pub(crate) fn py_obj_to_scalar_value(py: Python, obj: PyObject) -> PyResult<ScalarValue> {
     // convert Python object to PyScalarValue to ScalarValue