refactor: replace hardcoded error message with EXPECTED_PROVIDER_MSG constant

Author: kosiew

python/datafusion/__init__.py

@@ -33,7 +33,7 @@
 from . import functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
-from ._internal import Config
+from ._internal import Config, EXPECTED_PROVIDER_MSG
 from .catalog import Catalog, Database, Table
 from .col import col, column
 from .common import (
@@ -77,6 +77,7 @@
     "DFSchema",
     "DataFrame",
     "Database",
+    "EXPECTED_PROVIDER_MSG",
     "ExecutionPlan",
     "Expr",
     "LogicalPlan",

python/datafusion/table_provider.py

@@ -22,6 +22,7 @@
 from typing import Any
 
 import datafusion._internal as df_internal
+from datafusion._internal import EXPECTED_PROVIDER_MSG
 
 _InternalTableProvider = df_internal.TableProvider
 
@@ -37,8 +38,7 @@ def __init__(self, table_provider: _InternalTableProvider) -> None:
             table_provider = table_provider._table_provider
 
         if not isinstance(table_provider, _InternalTableProvider):
-            msg = "Expected a datafusion._internal.TableProvider instance."
-            raise TypeError(msg)
+            raise TypeError(EXPECTED_PROVIDER_MSG)
 
         self._table_provider = table_provider
 

python/tests/test_context.py

@@ -23,6 +23,7 @@
 import pytest
 from datafusion import (
     DataFrame,
+    EXPECTED_PROVIDER_MSG,
     RuntimeEnvBuilder,
     SessionConfig,
     SessionContext,
@@ -383,10 +384,7 @@ def test_register_table_with_dataframe_errors(ctx):
     with pytest.raises(Exception) as exc_info:
         ctx.register_table("bad", df)
 
-    assert str(exc_info.value) == (
-        "Expected a Table or TableProvider. Convert DataFrames with "
-        '"DataFrame.into_view()" or "TableProvider.from_dataframe()".'
-    )
+    assert str(exc_info.value) == EXPECTED_PROVIDER_MSG
 
 
 def test_register_dataset(ctx):

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::{
-    extract_table_provider, table_provider_from_pycapsule, validate_pycapsule, wait_for_future,
+    table_provider_from_pycapsule, validate_pycapsule, wait_for_future, EXPECTED_PROVIDER_MSG,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -198,7 +199,13 @@ impl PySchema {
     }
 
     fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> {
-        let provider = if let Some(provider) = extract_table_provider(&table_provider)? {
+        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();

src/context.rs

@@ -41,13 +41,14 @@ 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::{
-    extract_table_provider, get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future,
-    TABLE_OR_PROVIDER_EXPECTED_MESSAGE,
+    get_global_ctx, get_tokio_runtime, table_provider_from_pycapsule, validate_pycapsule,
+    wait_for_future, EXPECTED_PROVIDER_MSG,
 };
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
@@ -609,13 +610,17 @@ impl PySessionContext {
         name: &str,
         table_provider: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        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(),
-                )
-            })?;
+        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(),
+            ));
+        };
 
         self.ctx.register_table(name, provider)?;
         Ok(())

src/lib.rs

@@ -81,6 +81,8 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     // Initialize logging
     pyo3_log::init();
 
+    m.add("EXPECTED_PROVIDER_MSG", crate::utils::EXPECTED_PROVIDER_MSG)?;
+
     // Register the python classes
     m.add_class::<context::PyRuntimeEnvBuilder>()?;
     m.add_class::<context::PySessionConfig>()?;

src/utils.rs

@@ -37,7 +37,7 @@ use std::{
 };
 use tokio::{runtime::Runtime, time::sleep};
 
-pub(crate) const TABLE_OR_PROVIDER_EXPECTED_MESSAGE: &str =
+pub(crate) const EXPECTED_PROVIDER_MSG: &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]
@@ -156,9 +156,7 @@ pub(crate) fn extract_table_provider(
     }
 
     if py_obj.extract::<PyDataFrame>().is_ok() {
-        return Err(
-            PyDataFusionError::Common(TABLE_OR_PROVIDER_EXPECTED_MESSAGE.to_string()).into(),
-        );
+        return Err(PyDataFusionError::Common(EXPECTED_PROVIDER_MSG.to_string()).into());
     }
 
     if let Some(provider) = table_provider_from_pycapsule(py_obj)? {