feat: implement extract_table_provider utility for streamlined table registration

Author: kosiew

src/catalog.rs

@@ -15,12 +15,11 @@
 // 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, EXPECTED_PROVIDER_MSG,
+    extract_table_provider, table_provider_from_pycapsule, validate_pycapsule, wait_for_future,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -199,18 +198,9 @@ 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_PROVIDER_MSG.to_string()).into());
-        } else if let Some(provider) = table_provider_from_pycapsule(&table_provider)? {
-            provider
-        } else {
-            let py = table_provider.py();
-            let provider = Dataset::new(&table_provider, py)?;
-            Arc::new(provider) as Arc<dyn TableProvider>
+        let provider = match extract_table_provider(&table_provider) {
+            Ok(provider) => provider,
+            Err(err) => return Err(err.into()),
         };
 
         let _ = self

src/context.rs

@@ -41,14 +41,12 @@ 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, EXPECTED_PROVIDER_MSG,
+    extract_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;
@@ -610,17 +608,7 @@ 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_PROVIDER_MSG.to_string(),
-            ));
-        };
+        let provider = extract_table_provider(&table_provider)?;
 
         self.ctx.register_table(name, provider)?;
         Ok(())

src/utils.rs

@@ -16,8 +16,12 @@
 // under the License.
 
 use crate::{
+    catalog::PyTable,
     common::data_type::PyScalarValue,
+    dataframe::PyDataFrame,
+    dataset::Dataset,
     errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult},
+    table::PyTableProvider,
     TokioRuntime,
 };
 use datafusion::{
@@ -141,6 +145,40 @@ pub(crate) fn table_provider_from_pycapsule(
     }
 }
 
+pub(crate) fn extract_table_provider(
+    table_like: &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()),
+    }
+}
+
 pub(crate) fn py_obj_to_scalar_value(py: Python, obj: PyObject) -> PyResult<ScalarValue> {
     // convert Python object to PyScalarValue to ScalarValue