Update signatures for FFI passing of logical codec around

Author: timsaucer

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

@@ -19,19 +19,14 @@
 
 import pyarrow as pa
 from datafusion import SessionContext
-from datafusion_ffi_example import MyCatalogProvider
+from datafusion_ffi_example import FixedSchemaProvider, MyCatalogProvider
 
 
-def test_catalog_provider():
-    ctx = SessionContext()
-
-    my_catalog_name = "my_catalog"
+def common_checks(ctx: SessionContext, my_catalog_name: str) -> None:
     expected_schema_name = "my_schema"
     expected_table_name = "my_table"
     expected_table_columns = ["units", "price"]
 
-    catalog_provider = MyCatalogProvider(ctx)
-    ctx.register_catalog_provider(my_catalog_name, catalog_provider)
     my_catalog = ctx.catalog(my_catalog_name)
 
     my_catalog_schemas = my_catalog.names()
@@ -58,3 +53,23 @@ def test_catalog_provider():
     ]
     assert col0_result == expected_col0
     assert col1_result == expected_col1
+
+
+def test_catalog_provider():
+    ctx = SessionContext()
+
+    my_catalog_name = "my_catalog"
+
+    catalog_provider = MyCatalogProvider()
+    ctx.register_catalog_provider(my_catalog_name, catalog_provider)
+    common_checks(ctx, my_catalog_name)
+
+
+def test_schema_provider():
+    ctx = SessionContext()
+
+    my_schema_name = "my_schema"
+
+    schema_provider = FixedSchemaProvider()
+    ctx.catalog().register_schema(my_schema_name, schema_provider)
+    common_checks(ctx, "datafusion")

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

@@ -27,7 +27,7 @@ use datafusion_catalog::{
 };
 use datafusion_common::error::{DataFusionError, Result};
 use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
-use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use datafusion_ffi::schema_provider::FFI_SchemaProvider;
 use pyo3::types::PyCapsule;
 use pyo3::{pyclass, pymethods, Bound, PyAny, PyResult, Python};
 
@@ -58,14 +58,19 @@ pub fn my_table() -> Arc<dyn TableProvider + 'static> {
     Arc::new(MemTable::try_new(schema, vec![partitions]).unwrap())
 }
 
+#[pyclass(
+    name = "FixedSchemaProvider",
+    module = "datafusion_ffi_example",
+    subclass
+)]
 #[derive(Debug)]
 pub struct FixedSchemaProvider {
-    inner: MemorySchemaProvider,
+    inner: Arc<MemorySchemaProvider>,
 }
 
 impl Default for FixedSchemaProvider {
     fn default() -> Self {
-        let inner = MemorySchemaProvider::new();
+        let inner = Arc::new(MemorySchemaProvider::new());
 
         let table = my_table();
 
@@ -75,6 +80,29 @@ impl Default for FixedSchemaProvider {
     }
 }
 
+#[pymethods]
+impl FixedSchemaProvider {
+    #[new]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn __datafusion_schema_provider__<'py>(
+        &self,
+        py: Python<'py>,
+        session: Bound<PyAny>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_schema_provider".into();
+
+        let provider = Arc::clone(&self.inner) as Arc<dyn SchemaProvider + Send>;
+
+        let codec = ffi_logical_codec_from_pycapsule(session)?;
+        let provider = FFI_SchemaProvider::new_with_ffi_codec(provider, None, codec);
+
+        PyCapsule::new(py, provider, Some(name))
+    }
+}
+
 #[async_trait]
 impl SchemaProvider for FixedSchemaProvider {
     fn as_any(&self) -> &dyn Any {
@@ -116,7 +144,6 @@ impl SchemaProvider for FixedSchemaProvider {
 #[derive(Debug, Clone)]
 pub(crate) struct MyCatalogProvider {
     inner: Arc<MemoryCatalogProvider>,
-    logical_codec: FFI_LogicalExtensionCodec,
 }
 
 impl CatalogProvider for MyCatalogProvider {
@@ -152,28 +179,27 @@ impl CatalogProvider for MyCatalogProvider {
 #[pymethods]
 impl MyCatalogProvider {
     #[new]
-    pub fn new(session: &Bound<PyAny>) -> PyResult<Self> {
-        let logical_codec = ffi_logical_codec_from_pycapsule(session)?;
+    pub fn new() -> PyResult<Self> {
         let inner = Arc::new(MemoryCatalogProvider::new());
 
         let schema_name: &str = "my_schema";
         let _ = inner.register_schema(schema_name, Arc::new(FixedSchemaProvider::default()));
 
-        Ok(Self {
-            inner,
-            logical_codec,
-        })
+        Ok(Self { inner })
     }
 
     pub fn __datafusion_catalog_provider__<'py>(
         &self,
         py: Python<'py>,
+        session: Bound<PyAny>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
         let name = cr"datafusion_catalog_provider".into();
-        let codec = self.logical_codec.clone();
-        let catalog_provider =
-            FFI_CatalogProvider::new_with_ffi_codec(Arc::new(self.clone()), None, codec);
 
-        PyCapsule::new(py, catalog_provider, Some(name))
+        let provider = Arc::clone(&self.inner) as Arc<dyn CatalogProvider + Send>;
+
+        let codec = ffi_logical_codec_from_pycapsule(session)?;
+        let provider = FFI_CatalogProvider::new_with_ffi_codec(provider, None, codec);
+
+        PyCapsule::new(py, provider, Some(name))
     }
 }

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

@@ -18,7 +18,7 @@
 use pyo3::prelude::*;
 
 use crate::aggregate_udf::MySumUDF;
-use crate::catalog_provider::MyCatalogProvider;
+use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider};
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
@@ -37,6 +37,7 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyTableProvider>()?;
     m.add_class::<MyTableFunction>()?;
     m.add_class::<MyCatalogProvider>()?;
+    m.add_class::<FixedSchemaProvider>()?;
     m.add_class::<IsNullUDF>()?;
     m.add_class::<MySumUDF>()?;
     m.add_class::<MyRankUDF>()?;

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

@@ -41,7 +41,7 @@ impl MyTableFunction {
     fn __datafusion_table_function__<'py>(
         &self,
         py: Python<'py>,
-        session: &Bound<PyAny>,
+        session: Bound<PyAny>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
         let name = cr"datafusion_table_function".into();
 

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

@@ -93,13 +93,13 @@ impl MyTableProvider {
     pub fn __datafusion_table_provider__<'py>(
         &self,
         py: Python<'py>,
-        session: &Bound<PyAny>,
+        session: Bound<PyAny>,
     ) -> 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()))?;
+            .map_err(|e: DataFusionError| PyRuntimeError::new_err(e.to_string()))?;
 
         let codec = ffi_logical_codec_from_pycapsule(session)?;
         let provider =

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

@@ -5,23 +5,21 @@ use pyo3::types::PyCapsule;
 use pyo3::{Bound, PyAny, PyResult};
 
 pub(crate) fn ffi_logical_codec_from_pycapsule(
-    obj: &Bound<PyAny>,
+    obj: Bound<PyAny>,
 ) -> PyResult<FFI_LogicalExtensionCodec> {
     let attr_name = "__datafusion_logical_extension_codec__";
+    let capsule = if obj.hasattr(attr_name)? {
+        obj.getattr(attr_name)?.call0()?
+    } else {
+        obj
+    };
 
-    if obj.hasattr(attr_name)? {
-        let capsule = obj.getattr(attr_name)?.call0()?;
-        let capsule = capsule.downcast::<PyCapsule>()?;
-        validate_pycapsule(capsule, "datafusion_logical_extension_codec")?;
+    let capsule = capsule.downcast::<PyCapsule>()?;
+    validate_pycapsule(capsule, "datafusion_logical_extension_codec")?;
 
-        let provider = unsafe { capsule.reference::<FFI_LogicalExtensionCodec>() };
+    let codec = unsafe { capsule.reference::<FFI_LogicalExtensionCodec>() };
 
-        Ok(provider.clone())
-    } else {
-        Err(PyValueError::new_err(
-            "Expected PyCapsule object for FFI_LogicalExtensionCodec, but attribute does not exist",
-        ))
-    }
+    Ok(codec.clone())
 }
 
 pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyResult<()> {

python/datafusion/catalog.py

@@ -273,4 +273,4 @@ class SchemaProviderExportable(Protocol):
     https://docs.rs/datafusion/latest/datafusion/catalog/trait.SchemaProvider.html
     """
 
-    def __datafusion_schema_provider__(self) -> object: ...
+    def __datafusion_schema_provider__(self, session: Any) -> object: ...

python/datafusion/context.py

@@ -88,7 +88,7 @@ class TableProviderExportable(Protocol):
     https://datafusion.apache.org/python/user-guide/io/table_provider.html
     """
 
-    def __datafusion_table_provider__(self) -> object: ...  # noqa: D105
+    def __datafusion_table_provider__(self, session: Any) -> object: ...  # noqa: D105
 
 
 class CatalogProviderExportable(Protocol):
@@ -97,7 +97,7 @@ class CatalogProviderExportable(Protocol):
     https://docs.rs/datafusion/latest/datafusion/catalog/trait.CatalogProvider.html
     """
 
-    def __datafusion_catalog_provider__(self) -> object: ...  # noqa: D105
+    def __datafusion_catalog_provider__(self, session: Any) -> object: ...  # noqa: D105
 
 
 class SessionConfig:

src/catalog.rs

@@ -35,7 +35,10 @@ use pyo3::IntoPyObjectExt;
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::table::PyTable;
-use crate::utils::{extract_logical_extension_codec, validate_pycapsule, wait_for_future};
+use crate::utils::{
+    create_logical_extension_capsule, extract_logical_extension_codec, validate_pycapsule,
+    wait_for_future,
+};
 
 #[pyclass(frozen, name = "RawCatalog", module = "datafusion.catalog", subclass)]
 #[derive(Clone)]
@@ -111,23 +114,7 @@ impl PyCatalog {
     }
 
     pub fn register_schema(&self, name: &str, schema_provider: Bound<'_, PyAny>) -> PyResult<()> {
-        let provider = if schema_provider.hasattr("__datafusion_schema_provider__")? {
-            let capsule = schema_provider
-                .getattr("__datafusion_schema_provider__")?
-                .call0()?;
-            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-            validate_pycapsule(capsule, "datafusion_schema_provider")?;
-
-            let provider = unsafe { capsule.reference::<FFI_SchemaProvider>() };
-            let provider: Arc<dyn SchemaProvider + Send> = provider.into();
-            provider as Arc<dyn SchemaProvider>
-        } else {
-            match schema_provider.extract::<PySchema>() {
-                Ok(py_schema) => py_schema.schema,
-                Err(_) => Arc::new(RustWrappedPySchemaProvider::new(schema_provider.into()))
-                    as Arc<dyn SchemaProvider>,
-            }
-        };
+        let provider = extract_schema_provider_from_pyobj(schema_provider, self.codec.as_ref())?;
 
         let _ = self
             .catalog
@@ -195,14 +182,11 @@ impl PySchema {
         Ok(format!("Schema(table_names=[{}])", names.join(";")))
     }
 
-    fn register_table(&self, name: &str, table_provider: &Bound<'_, PyAny>) -> PyResult<()> {
+    fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> {
         let py = table_provider.py();
-        let codec_capsule = PyCapsule::new(
-            py,
-            self.codec.clone(),
-            Some(cr"datafusion_logical_extension_codec".into()),
-        )?
-        .into_bound_py_any(py)?;
+        let codec_capsule = create_logical_extension_capsule(py, self.codec.as_ref())?
+            .as_any()
+            .clone();
 
         let table = PyTable::new(table_provider, Some(codec_capsule))?;
 
@@ -256,7 +240,7 @@ impl RustWrappedPySchemaProvider {
                 return Ok(None);
             }
 
-            let table = PyTable::new(&py_table, None)?;
+            let table = PyTable::new(py_table, None)?;
 
             Ok(Some(table.table))
         })
@@ -370,32 +354,7 @@ impl RustWrappedPyCatalogProvider {
                 return Ok(None);
             }
 
-            if py_schema.hasattr("__datafusion_schema_provider__")? {
-                let capsule = provider
-                    .getattr("__datafusion_schema_provider__")?
-                    .call0()?;
-                let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
-                validate_pycapsule(capsule, "datafusion_schema_provider")?;
-
-                let provider = unsafe { capsule.reference::<FFI_SchemaProvider>() };
-                let provider: Arc<dyn SchemaProvider + Send> = provider.into();
-
-                Ok(Some(provider as Arc<dyn SchemaProvider>))
-            } else {
-                if let Ok(inner_schema) = py_schema.getattr("schema") {
-                    if let Ok(inner_schema) = inner_schema.extract::<PySchema>() {
-                        return Ok(Some(inner_schema.schema));
-                    }
-                }
-                match py_schema.extract::<PySchema>() {
-                    Ok(inner_schema) => Ok(Some(inner_schema.schema)),
-                    Err(_) => {
-                        let py_schema = RustWrappedPySchemaProvider::new(py_schema.into());
-
-                        Ok(Some(Arc::new(py_schema) as Arc<dyn SchemaProvider>))
-                    }
-                }
-            }
+            extract_schema_provider_from_pyobj(py_schema, self.codec.as_ref()).map(Some)
         })
     }
 }
@@ -479,6 +438,35 @@ impl CatalogProvider for RustWrappedPyCatalogProvider {
     }
 }
 
+fn extract_schema_provider_from_pyobj(
+    mut schema_provider: Bound<PyAny>,
+    codec: &FFI_LogicalExtensionCodec,
+) -> PyResult<Arc<dyn SchemaProvider>> {
+    if schema_provider.hasattr("__datafusion_schema_provider__")? {
+        let py = schema_provider.py();
+        let codec_capsule = create_logical_extension_capsule(py, codec)?;
+        schema_provider = schema_provider
+            .getattr("__datafusion_schema_provider__")?
+            .call1((codec_capsule,))?;
+    }
+
+    let provider = if let Ok(capsule) = schema_provider.downcast::<PyCapsule>() {
+        validate_pycapsule(capsule, "datafusion_schema_provider")?;
+
+        let provider = unsafe { capsule.reference::<FFI_SchemaProvider>() };
+        let provider: Arc<dyn SchemaProvider + Send> = provider.into();
+        provider as Arc<dyn SchemaProvider>
+    } else {
+        match schema_provider.extract::<PySchema>() {
+            Ok(py_schema) => py_schema.schema,
+            Err(_) => Arc::new(RustWrappedPySchemaProvider::new(schema_provider.into()))
+                as Arc<dyn SchemaProvider>,
+        }
+    };
+
+    Ok(provider)
+}
+
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyCatalog>()?;
     m.add_class::<PySchema>()?;