fix: replace foreign_table_provider_from_capsule with try_table_provider_from_object for improved clarity and error handling

Author: kosiew

src/catalog.rs

@@ -18,7 +18,7 @@
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::utils::{
-    foreign_table_provider_from_capsule, get_tokio_runtime, validate_pycapsule, wait_for_future,
+    get_tokio_runtime, try_table_provider_from_object, validate_pycapsule, wait_for_future,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -201,13 +201,7 @@ impl PySchema {
     fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> {
         let provider = if let Ok(py_table) = table_provider.extract::<PyTable>() {
             py_table.table
-        } else 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)?;
-            let provider = foreign_table_provider_from_capsule(capsule)?;
-            let provider: Arc<dyn TableProvider> = Arc::new(provider);
+        } else if let Some(provider) = try_table_provider_from_object(&table_provider)? {
             provider
         } else {
             let py = table_provider.py();
@@ -320,12 +314,7 @@ impl RustWrappedPySchemaProvider {
                 return Ok(Some(inner_table.table));
             }
 
-            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)?;
-                let provider = foreign_table_provider_from_capsule(capsule)?;
-                let provider: Arc<dyn TableProvider> = Arc::new(provider);
-
+            if let Some(provider) = try_table_provider_from_object(&py_table)? {
                 Ok(Some(provider))
             } else {
                 if let Ok(inner_table) = py_table.getattr("table") {

src/context.rs

@@ -46,7 +46,7 @@ use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
 use crate::utils::{
-    foreign_table_provider_from_capsule, get_global_ctx, get_tokio_runtime, validate_pycapsule,
+    get_global_ctx, get_tokio_runtime, try_table_provider_from_object, validate_pycapsule,
     wait_for_future,
 };
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
@@ -653,12 +653,7 @@ impl PySessionContext {
         name: &str,
         provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        if provider.hasattr("__datafusion_table_provider__")? {
-            let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?;
-            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-            let provider = foreign_table_provider_from_capsule(capsule)?;
-            let provider: Arc<dyn TableProvider> = Arc::new(provider);
-
+        if let Some(provider) = try_table_provider_from_object(&provider)? {
             let _ = self.ctx.register_table(name, provider)?;
 
             Ok(())
@@ -1110,12 +1105,7 @@ impl PySessionContext {
             return Ok(PyDataFrame::new(df));
         }
 
-        if table.hasattr("__datafusion_table_provider__")? {
-            let capsule = table.getattr("__datafusion_table_provider__")?.call0()?;
-            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-            let provider = foreign_table_provider_from_capsule(capsule)?;
-            let provider: Arc<dyn TableProvider> = Arc::new(provider);
-
+        if let Some(provider) = try_table_provider_from_object(&table)? {
             let df = self.ctx.read_table(provider)?;
             Ok(PyDataFrame::new(df))
         } else {

src/udtf.rs

@@ -21,7 +21,7 @@ use std::sync::Arc;
 use crate::dataframe::PyTableProvider;
 use crate::errors::{py_datafusion_err, to_datafusion_err};
 use crate::expr::PyExpr;
-use crate::utils::{foreign_table_provider_from_capsule, validate_pycapsule};
+use crate::utils::{try_table_provider_from_object, validate_pycapsule};
 use datafusion::catalog::{TableFunctionImpl, TableProvider};
 use datafusion::error::Result as DataFusionResult;
 use datafusion::logical_expr::Expr;
@@ -98,12 +98,7 @@ fn call_python_table_function(
         let provider_obj = func.call1(py, py_args)?;
         let provider = provider_obj.bind(py);
 
-        if provider.hasattr("__datafusion_table_provider__")? {
-            let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?;
-            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-            let provider = foreign_table_provider_from_capsule(capsule)?;
-            let provider: Arc<dyn TableProvider> = Arc::new(provider);
-
+        if let Some(provider) = try_table_provider_from_object(provider)? {
             Ok(provider)
         } else {
             Err(PyNotImplementedError::new_err(

src/utils.rs

@@ -17,16 +17,21 @@
 
 use crate::{
     common::data_type::PyScalarValue,
-    errors::{PyDataFusionError, PyDataFusionResult},
+    errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult},
     TokioRuntime,
 };
 use datafusion::{
-    common::ScalarValue, execution::context::SessionContext, logical_expr::Volatility,
+    common::ScalarValue, datasource::TableProvider, 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};
+use std::{
+    future::Future,
+    sync::{Arc, OnceLock},
+    time::Duration,
+};
 use tokio::{runtime::Runtime, time::sleep};
 /// Utility to get the Tokio Runtime from Python
 #[inline]
@@ -124,6 +129,21 @@ pub(crate) fn foreign_table_provider_from_capsule(
     Ok(unsafe { capsule.reference::<FFI_TableProvider>() }.into())
 }
 
+pub(crate) fn try_table_provider_from_object(
+    provider: &Bound<'_, PyAny>,
+) -> PyResult<Option<Arc<dyn TableProvider>>> {
+    if !provider.hasattr("__datafusion_table_provider__")? {
+        return Ok(None);
+    }
+
+    let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?;
+    let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+    let provider = foreign_table_provider_from_capsule(capsule)?;
+    let provider: Arc<dyn TableProvider> = Arc::new(provider);
+
+    Ok(Some(provider))
+}
+
 pub(crate) fn py_obj_to_scalar_value(py: Python, obj: PyObject) -> PyResult<ScalarValue> {
     // convert Python object to PyScalarValue to ScalarValue