feat: add examples BlockingTableProvider with scan functionality and corresponding tests

Author: kosiew

examples/datafusion-ffi-example/python/tests/test_blocking_provider.py

@@ -0,0 +1,10 @@
+from datafusion import SessionContext
+from datafusion_ffi_example import BlockingTableProvider
+
+
+def test_runtime_available_in_scan() -> None:
+    ctx = SessionContext()
+    provider = BlockingTableProvider()
+    ctx.register_table_provider("t", provider)
+    result = ctx.table("t").collect()
+    assert [b.column(0).to_pylist() for b in result] == [[1, 2, 3]]

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

@@ -0,0 +1,67 @@
+use arrow_array::{ArrayRef, Int32Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use datafusion_catalog::{MemTable, Session, TableProvider, TableType};
+use datafusion_common::Result as DataFusionResult;
+use datafusion_expr::Expr;
+use datafusion_ffi::table_provider::FFI_TableProvider;
+use datafusion_physical_plan::ExecutionPlan;
+use pyo3::types::PyCapsule;
+use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use std::any::Any;
+use std::sync::Arc;
+
+#[pyclass(
+    name = "BlockingTableProvider",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Clone)]
+pub(crate) struct BlockingTableProvider;
+
+#[pymethods]
+impl BlockingTableProvider {
+    #[new]
+    fn new() -> Self {
+        Self
+    }
+
+    fn __datafusion_table_provider__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_table_provider".into();
+        let provider = FFI_TableProvider::new(Arc::new(self.clone()), false, None);
+        PyCapsule::new(py, provider, Some(name))
+    }
+}
+
+#[async_trait::async_trait]
+impl TableProvider for BlockingTableProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
+        tokio::spawn_blocking(|| ()).await.unwrap();
+        let batch = RecordBatch::try_new(
+            self.schema(),
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef],
+        )?;
+        let table = MemTable::try_new(batch.schema(), vec![vec![batch]])?;
+        table.scan(state, projection, &[], None).await
+    }
+}

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

@@ -15,16 +15,19 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::blocking_provider::BlockingTableProvider;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
 use pyo3::prelude::*;
 
+pub(crate) mod blocking_provider;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
 
 #[pymodule]
 fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyTableProvider>()?;
+    m.add_class::<BlockingTableProvider>()?;
     m.add_class::<MyTableFunction>()?;
     Ok(())
 }