Add table provider capsule to PyTableProvider

Author: timsaucer

src/dataframe.rs

@@ -34,6 +34,7 @@ use datafusion::error::DataFusionError;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::prelude::*;
+use datafusion_ffi::table_provider::FFI_TableProvider;
 use futures::{StreamExt, TryStreamExt};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
@@ -58,7 +59,7 @@ use crate::{
 // this is an interim implementation
 #[pyclass(name = "TableProvider", module = "datafusion")]
 pub struct PyTableProvider {
-    provider: Arc<dyn TableProvider>,
+    provider: Arc<dyn TableProvider + Send>,
 }
 
 impl PyTableProvider {
@@ -72,6 +73,21 @@ impl PyTableProvider {
     }
 }
 
+#[pymethods]
+impl PyTableProvider {
+    fn __datafusion_table_provider__<'py>(
+        &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()))
+    }
+}
+
 /// Configuration for DataFrame display formatting
 #[derive(Debug, Clone)]
 pub struct FormatterConfig {