fix: standardize TABLE_PROVIDER_CAPSULE_NAME usage across modules for consistency

Author: kosiew

examples/datafusion-ffi-example/Cargo.toml

@@ -23,6 +23,7 @@ edition = "2021"
 [dependencies]
 datafusion = { version = "49.0.2" }
 datafusion-ffi = { version = "49.0.2" }
+datafusion-python = { path = "../../" }
 pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py39"] }
 arrow = { version = "55.0.0" }
 arrow-array = { version = "55.0.0" }

examples/datafusion-ffi-example/src/table_provider.rs

@@ -19,6 +19,7 @@ use arrow_array::{ArrayRef, RecordBatch};
 use arrow_schema::{DataType, Field, Schema};
 use datafusion::catalog::MemTable;
 use datafusion::error::{DataFusionError, Result as DataFusionResult};
+use datafusion_python::utils::TABLE_PROVIDER_CAPSULE_NAME;
 use datafusion_ffi::table_provider::FFI_TableProvider;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::types::PyCapsule;
@@ -91,13 +92,11 @@ impl MyTableProvider {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = cr"datafusion_table_provider".into();
-
         let provider = self
             .create_table()
             .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
         let provider = FFI_TableProvider::new(Arc::new(provider), false, None);
 
-        PyCapsule::new(py, provider, Some(name))
+        PyCapsule::new(py, provider, Some(TABLE_PROVIDER_CAPSULE_NAME.to_owned()))
     }
 }

src/catalog.rs

@@ -19,6 +19,7 @@ use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::utils::{
     get_tokio_runtime, try_table_provider_from_object, validate_pycapsule, wait_for_future,
+    TABLE_PROVIDER_CAPSULE_NAME,
 };
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
@@ -36,7 +37,6 @@ use pyo3::types::PyCapsule;
 use pyo3::IntoPyObjectExt;
 use std::any::Any;
 use std::collections::HashSet;
-use std::ffi::CString;
 use std::sync::Arc;
 
 #[pyclass(name = "RawCatalog", module = "datafusion.catalog", subclass)]
@@ -257,14 +257,13 @@ impl PyTable {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = CString::new("datafusion_table_provider").unwrap();
         let runtime = get_tokio_runtime().0.handle().clone();
 
         let provider = Arc::clone(&self.table);
         let provider: Arc<dyn TableProvider + Send> = provider;
         let provider = FFI_TableProvider::new(provider, false, Some(runtime));
 
-        PyCapsule::new(py, provider, Some(name.clone()))
+        PyCapsule::new(py, provider, Some(TABLE_PROVIDER_CAPSULE_NAME.to_owned()))
     }
 
     fn __repr__(&self) -> PyResult<String> {

src/dataframe.rs

@@ -52,6 +52,7 @@ use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{
     get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, validate_pycapsule, wait_for_future,
+    TABLE_PROVIDER_CAPSULE_NAME,
 };
 use crate::{
     errors::PyDataFusionResult,
@@ -83,12 +84,10 @@ impl PyTableProvider {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = CString::new("datafusion_table_provider").unwrap();
-
         let runtime = get_tokio_runtime().0.handle().clone();
         let provider = FFI_TableProvider::new(Arc::clone(&self.provider), false, Some(runtime));
 
-        PyCapsule::new(py, provider, Some(name.clone()))
+        PyCapsule::new(py, provider, Some(TABLE_PROVIDER_CAPSULE_NAME.to_owned()))
     }
 }
 

src/utils.rs

@@ -28,11 +28,17 @@ use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::prelude::*;
 use pyo3::{exceptions::PyValueError, types::PyCapsule};
 use std::{
+    ffi::CStr,
     future::Future,
     sync::{Arc, OnceLock},
     time::Duration,
 };
 use tokio::{runtime::Runtime, time::sleep};
+
+pub const TABLE_PROVIDER_CAPSULE_NAME: &CStr =
+    unsafe { CStr::from_bytes_with_nul_unchecked(b"datafusion_table_provider\0") };
+
+pub const TABLE_PROVIDER_CAPSULE_NAME_STR: &str = "datafusion_table_provider";
 /// Utility to get the Tokio Runtime from Python
 #[inline]
 pub(crate) fn get_tokio_runtime() -> &'static TokioRuntime {
@@ -125,7 +131,7 @@ pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyRe
 pub(crate) fn foreign_table_provider_from_capsule(
     capsule: &Bound<PyCapsule>,
 ) -> PyResult<ForeignTableProvider> {
-    validate_pycapsule(capsule, "datafusion_table_provider")?;
+    validate_pycapsule(capsule, TABLE_PROVIDER_CAPSULE_NAME_STR)?;
     Ok(unsafe { capsule.reference::<FFI_TableProvider>() }.into())
 }