Adding capability to invoke python catalog and schema providers as well

Author: renato2099

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

@@ -19,16 +19,19 @@
 
 import pyarrow as pa
 
-from datafusion import SessionContext
+from datafusion import SessionContext, Table
 from datafusion_ffi_example import MyCatalogProvider
 
+from datafusion.context import PyCatalogProvider, PySchemaProvider
+
+
 def test_catalog_provider():
     ctx = SessionContext()
 
     my_catalog_name = "my_catalog"
     expected_schema_name = "my_schema"
     expected_table_name = "my_table"
-    expected_table_columns = ['units', 'price']
+    expected_table_columns = ["units", "price"]
 
     catalog_provider = MyCatalogProvider()
     ctx.register_catalog_provider(my_catalog_name, catalog_provider)
@@ -41,12 +44,9 @@ def test_catalog_provider():
     my_table = my_database.table(expected_table_name)
     assert expected_table_columns == my_table.schema.names
 
-    ctx.register_table(expected_table_name, my_table)
-    expected_df = ctx.sql(f"SELECT * FROM {expected_table_name}").to_pandas()
-    assert len(expected_df) == 5
-    assert expected_table_columns == expected_df.columns.tolist()
-
-    result = ctx.table(f"{my_catalog_name}.{expected_schema_name}.{expected_table_name}").collect()
+    result = ctx.table(
+        f"{my_catalog_name}.{expected_schema_name}.{expected_table_name}"
+    ).collect()
     assert len(result) == 2
 
     col0_result = [r.column(0) for r in result]
@@ -60,4 +60,58 @@ def test_catalog_provider():
         pa.array([1.5, 2.5], type=pa.float64()),
     ]
     assert col0_result == expected_col0
-    assert col1_result == expected_col1
+    assert col1_result == expected_col1
+
+
+class MyPyCatalogProvider(PyCatalogProvider):
+    my_schemas = ['my_schema']
+
+    def schema_names(self) -> list[str]:
+        return self.my_schemas
+
+    def schema(self, name: str) -> PySchemaProvider:
+        return MyPySchemaProvider()
+
+
+class MyPySchemaProvider(PySchemaProvider):
+    my_tables = ['table1', 'table2', 'table3']
+
+    def table_names(self) -> list[str]:
+        return self.my_tables
+
+    def table_exist(self, table_name: str) -> bool:
+        return table_name in self.my_tables
+
+    def table(self, table_name: str) -> Table:
+        raise RuntimeError(f"Can not get table: {table_name}")
+
+    def register_table(self, table: Table) -> None:
+        raise RuntimeError(f"Can not register {table} as table")
+
+    def deregister_table(self, table_name: str) -> None:
+        raise RuntimeError(f"Can not deregister table: {table_name}")
+
+
+def test_python_catalog_provider():
+    ctx = SessionContext()
+
+    my_catalog_name = "my_py_catalog"
+    expected_schema_name = "my_schema"
+    my_py_catalog_provider = MyPyCatalogProvider()
+    ctx.register_catalog_provider(my_catalog_name, my_py_catalog_provider)
+    my_py_catalog = ctx.catalog(my_catalog_name)
+    assert MyPyCatalogProvider.my_schemas == my_py_catalog.names()
+
+    my_database = my_py_catalog.database(expected_schema_name)
+    assert set(MyPySchemaProvider.my_tables) == my_database.names()
+
+    # asserting a non-compliant provider fails at the python level as expected
+    try:
+        ctx.register_catalog_provider(my_catalog_name, "non_compliant_provider")
+    except TypeError:
+        # expect a TypeError because we can not register a str as a catalog provider
+        pass
+
+
+if __name__ == "__main__":
+    test_python_catalog_provider()

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

@@ -15,15 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{any::Any, fmt::Debug, sync::Arc};
 use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use std::{any::Any, fmt::Debug, sync::Arc};
 
 use arrow::datatypes::Schema;
 use async_trait::async_trait;
 use datafusion::{
     catalog::{
-        CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider,
-        TableProvider,
+        CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider, TableProvider,
     },
     common::exec_err,
     datasource::MemTable,
@@ -46,12 +45,12 @@ pub fn my_table() -> Arc<dyn TableProvider + 'static> {
             ("units", Int32, vec![10, 20, 30]),
             ("price", Float64, vec![1.0, 2.0, 5.0])
         )
-            .unwrap(),
+        .unwrap(),
         record_batch!(
             ("units", Int32, vec![5, 7]),
             ("price", Float64, vec![1.5, 2.5])
         )
-            .unwrap(),
+        .unwrap(),
     ];
 
     Arc::new(MemTable::try_new(schema, vec![partitions]).unwrap())
@@ -68,9 +67,7 @@ impl Default for FixedSchemaProvider {
 
         let table = my_table();
 
-        let _ = inner
-            .register_table("my_table".to_string(), table)
-            .unwrap();
+        let _ = inner.register_table("my_table".to_string(), table).unwrap();
 
         Self { inner }
     }
@@ -86,10 +83,7 @@ impl SchemaProvider for FixedSchemaProvider {
         self.inner.table_names()
     }
 
-    async fn table(
-        &self,
-        name: &str,
-    ) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
+    async fn table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
         self.inner.table(name).await
     }
 
@@ -110,7 +104,6 @@ impl SchemaProvider for FixedSchemaProvider {
     }
 }
 
-
 /// This catalog provider is intended only for unit tests. It prepopulates with one
 /// schema and only allows for schemas named after four types of fruit.
 #[pyclass(name = "MyCatalogProvider", module = "datafusion_ffi_example", subclass)]
@@ -178,4 +171,4 @@ impl MyCatalogProvider {
 
         PyCapsule::new(py, catalog_provider, Some(name))
     }
-}
+}

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

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

python/datafusion/context.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 import pyarrow as pa
 
@@ -79,13 +79,40 @@ class TableProviderExportable(Protocol):
 
     def __datafusion_table_provider__(self) -> object: ...  # noqa: D105
 
-
+@runtime_checkable
 class CatalogProviderExportable(Protocol):
     """Type hint for object that has __datafusion_catalog_provider__ PyCapsule.
 
     https://docs.rs/datafusion/latest/datafusion/catalog/trait.CatalogProvider.html
     """
-    def __datafusion_catalog_provider__(self) -> object: ...
+
+    def __datafusion_catalog_provider__(self) -> object: ...  # noqa: D105
+
+@runtime_checkable
+class PySchemaProvider(Protocol):
+    def table_names(self) -> list[str]:
+        ...
+
+    def register_table(self, table: Table) -> None:
+        ...
+
+    def deregister_table(self, table_name: str) -> None:
+        ...
+
+    def table_exist(self, table_name: str) -> bool:
+        ...
+
+    def table(self, table_name: str) -> Table:
+        ...
+
+
+@runtime_checkable
+class PyCatalogProvider(Protocol):
+    def schema_names(self) -> list[str]:
+        ...
+
+    def schema(self, name: str) -> PySchemaProvider:
+        ...
 
 
 class SessionConfig:
@@ -758,9 +785,14 @@ def deregister_table(self, name: str) -> None:
         self.ctx.deregister_table(name)
 
     def register_catalog_provider(
-            self, name: str, provider: CatalogProviderExportable
+        self, name: str, provider: PyCatalogProvider | CatalogProviderExportable
     ) -> None:
         """Register a catalog provider."""
+        if not isinstance(provider, (PyCatalogProvider, CatalogProviderExportable)):
+            raise TypeError(
+                f"Expected provider to be CatalogProviderProtocol or rust version exposed through python, but got {type(provider)} instead."
+            )
+
         self.ctx.register_catalog_provider(name, provider)
 
     def register_table_provider(