refactor: streamline table provider extraction logic and improve error handling

Author: kosiew

src/catalog.rs

@@ -15,11 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-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, validate_pycapsule, wait_for_future};
+use crate::utils::{
+    extract_table_provider, table_provider_from_pycapsule, validate_pycapsule, wait_for_future,
+};
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion::common::DataFusionError;
@@ -197,17 +198,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 let Ok(py_provider) = table_provider.extract::<PyTableProvider>() {
-            py_provider.into_inner()
-        } else if table_provider.extract::<PyDataFrame>().is_ok() {
-            return Err(PyDataFusionError::Common(
-                "Expected a Table or TableProvider. Convert DataFrames with \"DataFrame.into_view()\" or \"TableProvider.from_dataframe()\"."
-                    .to_string(),
-            )
-            .into());
-        } else if let Some(provider) = table_provider_from_pycapsule(&table_provider)? {
+        let provider = if let Some(provider) = extract_table_provider(&table_provider)? {
             provider
         } else {
             let py = table_provider.py();

src/context.rs

@@ -41,14 +41,13 @@ use crate::record_batch::PyRecordBatchStream;
 use crate::sql::exceptions::py_value_err;
 use crate::sql::logical::PyLogicalPlan;
 use crate::store::StorageContexts;
-use crate::table::PyTableProvider;
 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, table_provider_from_pycapsule, validate_pycapsule,
-    wait_for_future,
+    extract_table_provider, get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future,
+    TABLE_OR_PROVIDER_EXPECTED_MESSAGE,
 };
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
@@ -610,17 +609,13 @@ impl PySessionContext {
         name: &str,
         table_provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        let provider = if let Ok(py_table) = table_provider.extract::<PyTable>() {
-            py_table.table()
-        } else if let Ok(py_provider) = table_provider.extract::<PyTableProvider>() {
-            py_provider.into_inner()
-        } else if let Some(provider) = table_provider_from_pycapsule(&table_provider)? {
-            provider
-        } else {
-            return Err(crate::errors::PyDataFusionError::Common(
-                "Expected a Table or TableProvider. Convert DataFrames with \"DataFrame.into_view()\" or \"TableProvider.from_dataframe()\".".to_string(),
-            ));
-        };
+        let provider = extract_table_provider(&table_provider)
+            .map_err(crate::errors::PyDataFusionError::from)?
+            .ok_or_else(|| {
+                crate::errors::PyDataFusionError::Common(
+                    TABLE_OR_PROVIDER_EXPECTED_MESSAGE.to_string(),
+                )
+            })?;
 
         self.ctx.register_table(name, provider)?;
         Ok(())

src/utils.rs

@@ -16,8 +16,11 @@
 // under the License.
 
 use crate::{
+    catalog::PyTable,
     common::data_type::PyScalarValue,
+    dataframe::PyDataFrame,
     errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult},
+    table::PyTableProvider,
     TokioRuntime,
 };
 use datafusion::{
@@ -33,6 +36,9 @@ use std::{
     time::Duration,
 };
 use tokio::{runtime::Runtime, time::sleep};
+
+pub(crate) const TABLE_OR_PROVIDER_EXPECTED_MESSAGE: &str =
+    "Expected a Table or TableProvider. Convert DataFrames with \"DataFrame.into_view()\" or \"TableProvider.from_dataframe()\".";
 /// Utility to get the Tokio Runtime from Python
 #[inline]
 pub(crate) fn get_tokio_runtime() -> &'static TokioRuntime {
@@ -138,6 +144,30 @@ pub(crate) fn table_provider_from_pycapsule(
     }
 }
 
+pub(crate) fn extract_table_provider(
+    py_obj: &Bound<PyAny>,
+) -> PyResult<Option<Arc<dyn TableProvider>>> {
+    if let Ok(py_table) = py_obj.extract::<PyTable>() {
+        return Ok(Some(py_table.table()));
+    }
+
+    if let Ok(py_provider) = py_obj.extract::<PyTableProvider>() {
+        return Ok(Some(py_provider.into_inner()));
+    }
+
+    if py_obj.extract::<PyDataFrame>().is_ok() {
+        return Err(
+            PyDataFusionError::Common(TABLE_OR_PROVIDER_EXPECTED_MESSAGE.to_string()).into(),
+        );
+    }
+
+    if let Some(provider) = table_provider_from_pycapsule(py_obj)? {
+        return Ok(Some(provider));
+    }
+
+    Ok(None)
+}
+
 pub(crate) fn py_obj_to_scalar_value(py: Python, obj: PyObject) -> PyResult<ScalarValue> {
     // convert Python object to PyScalarValue to ScalarValue