refactor: replace extract_table_provider with coerce_table_provider for improved clarity and error handling

Author: kosiew

python/tests/test_catalog.py

@@ -20,7 +20,7 @@
 import pyarrow as pa
 import pyarrow.dataset as ds
 import pytest
-from datafusion import SessionContext, Table
+from datafusion import EXPECTED_PROVIDER_MSG, SessionContext, Table
 
 
 # Note we take in `database` as a variable even though we don't use
@@ -164,6 +164,16 @@ def test_python_table_provider(ctx: SessionContext):
     assert schema.table_names() == {"table4"}
 
 
+def test_schema_register_table_with_dataframe_errors(ctx: SessionContext):
+    schema = ctx.catalog().schema()
+    df = ctx.from_pydict({"a": [1]})
+
+    with pytest.raises(Exception) as exc_info:
+        schema.register_table("bad", df)
+
+    assert str(exc_info.value) == EXPECTED_PROVIDER_MSG
+
+
 def test_in_end_to_end_python_providers(ctx: SessionContext):
     """Test registering all python providers and running a query against them."""
 

src/catalog.rs

@@ -19,7 +19,7 @@ use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::table::PyTableProvider;
 use crate::utils::{
-    extract_table_provider, table_provider_from_pycapsule, validate_pycapsule, wait_for_future,
+    coerce_table_provider, table_provider_from_pycapsule, validate_pycapsule, wait_for_future,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -198,10 +198,7 @@ impl PySchema {
     }
 
     fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> {
-        let provider = match extract_table_provider(&table_provider) {
-            Ok(provider) => provider,
-            Err(err) => return Err(err.into()),
-        };
+        let provider = coerce_table_provider(&table_provider).map_err(PyErr::from)?;
 
         let _ = self
             .schema

src/context.rs

@@ -46,7 +46,7 @@ use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
 use crate::utils::{
-    extract_table_provider, get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future,
+    coerce_table_provider, get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future,
 };
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
@@ -608,7 +608,7 @@ impl PySessionContext {
         name: &str,
         table_provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        let provider = extract_table_provider(&table_provider)?;
+        let provider = coerce_table_provider(&table_provider)?;
 
         self.ctx.register_table(name, provider)?;
         Ok(())

src/utils.rs

@@ -20,7 +20,7 @@ use crate::{
     common::data_type::PyScalarValue,
     dataframe::PyDataFrame,
     dataset::Dataset,
-    errors::{PyDataFusionError, PyDataFusionResult},
+    errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult},
     table::PyTableProvider,
     TokioRuntime,
 };
@@ -140,37 +140,25 @@ pub(crate) fn table_provider_from_pycapsule(
     }
 }
 
-pub(crate) fn extract_table_provider(
-    table_like: &Bound<PyAny>,
+pub(crate) fn coerce_table_provider(
+    obj: &Bound<PyAny>,
 ) -> PyDataFusionResult<Arc<dyn TableProvider>> {
-    if let Ok(py_table) = table_like.extract::<PyTable>() {
-        return Ok(py_table.table());
-    }
-
-    if let Ok(py_provider) = table_like.extract::<PyTableProvider>() {
-        return Ok(py_provider.into_inner());
-    }
-
-    if table_like.extract::<PyDataFrame>().is_ok() {
-        return Err(PyDataFusionError::Common(EXPECTED_PROVIDER_MSG.to_string()));
-    }
-
-    match table_provider_from_pycapsule(table_like) {
-        Ok(Some(provider)) => Ok(provider),
-        Ok(None) => {
-            let py = table_like.py();
-            match Dataset::new(table_like, py) {
-                Ok(dataset) => Ok(Arc::new(dataset) as Arc<dyn TableProvider>),
-                Err(err) => {
-                    if err.is_instance_of::<PyValueError>(py) {
-                        Err(PyDataFusionError::Common(EXPECTED_PROVIDER_MSG.to_string()))
-                    } else {
-                        Err(err.into())
-                    }
-                }
-            }
-        }
-        Err(err) => Err(err.into()),
+    if let Ok(py_table) = obj.extract::<PyTable>() {
+        Ok(py_table.table())
+    } else if let Ok(py_provider) = obj.extract::<PyTableProvider>() {
+        Ok(py_provider.into_inner())
+    } else if obj.is_instance_of::<PyDataFrame>()
+        || obj
+            .getattr("df")
+            .is_ok_and(|inner| inner.is_instance_of::<PyDataFrame>())
+    {
+        Err(PyDataFusionError::Common(EXPECTED_PROVIDER_MSG.to_string()))
+    } else if let Some(provider) = table_provider_from_pycapsule(obj)? {
+        Ok(provider)
+    } else {
+        let py = obj.py();
+        let provider = Dataset::new(obj, py)?;
+        Ok(Arc::new(provider) as Arc<dyn TableProvider>)
     }
 }