Adding implementation of python based catalog and schema providers

Author: timsaucer

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

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import pyarrow as pa
-
 from datafusion import SessionContext
 from datafusion_ffi_example import MyCatalogProvider
 

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,10 +104,13 @@ 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)]
+#[pyclass(
+    name = "MyCatalogProvider",
+    module = "datafusion_ffi_example",
+    subclass
+)]
 #[derive(Debug)]
 pub(crate) struct MyCatalogProvider {
     inner: MemoryCatalogProvider,
@@ -174,8 +171,9 @@ impl MyCatalogProvider {
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
         let name = cr"datafusion_catalog_provider".into();
-        let catalog_provider = FFI_CatalogProvider::new(Arc::new(MyCatalogProvider::default()), None);
+        let catalog_provider =
+            FFI_CatalogProvider::new(Arc::new(MyCatalogProvider::default()), None);
 
         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/__init__.py

@@ -32,7 +32,7 @@
 
 from datafusion.col import col, column
 
-from . import functions, object_store, substrait, unparser
+from . import catalog, functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
 from ._internal import Config
@@ -91,6 +91,7 @@
     "TableFunction",
     "WindowFrame",
     "WindowUDF",
+    "catalog",
     "col",
     "column",
     "common",

python/datafusion/catalog.py

@@ -26,11 +26,23 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
+try:
+    from warnings import deprecated  # Python 3.13+
+except ImportError:
+    from typing_extensions import deprecated  # Python 3.12
+
+
+__all__ = [
+    "Catalog",
+    "Schema",
+    "Table",
+]
+
 
 class Catalog:
     """DataFusion data catalog."""
 
-    def __init__(self, catalog: df_internal.Catalog) -> None:
+    def __init__(self, catalog: df_internal.catalog.RawCatalog) -> None:
         """This constructor is not typically called by the end user."""
         self.catalog = catalog
 
@@ -59,18 +71,74 @@ def __repr__(self) -> str:
         return self.db.__repr__()
 
     def names(self) -> set[str]:
-        """Returns the list of all tables in this database."""
-        return self.db.names()
+        """This is an alias for `schema_names`."""
+        return self.schema_names()
+
+    def schema_names(self) -> set[str]:
+        """Returns the list of schemas in this catalog."""
+        return self.catalog.schema_names()
+
+    def schema(self, name: str = "public") -> Schema:
+        """Returns the database with the given ``name`` from this catalog."""
+        schema = self.catalog.schema(name)
+
+        return (
+            Schema(schema)
+            if isinstance(schema, df_internal.catalog.RawSchema)
+            else schema
+        )
+
+    @deprecated("Use `schema` instead.")
+    def database(self, name: str = "public") -> Schema:
+        """Returns the database with the given ``name`` from this catalog."""
+        return self.schema(name)
+
+    def register_schema(self, name, schema) -> Schema | None:
+        """Register a schema with this catalog."""
+        return self.catalog.register_schema(name, schema)
+
+    def deregister_schema(self, name: str, cascade: bool = True) -> Schema | None:
+        """Deregister a schema from this catalog."""
+        return self.catalog.deregister_schema(name, cascade)
+
+
+class Schema:
+    """DataFusion Schema."""
+
+    def __init__(self, schema: df_internal.catalog.RawSchema) -> None:
+        """This constructor is not typically called by the end user."""
+        self._raw_schema = schema
+
+    def names(self) -> set[str]:
+        """This is an alias for `table_names`."""
+        return self.table_names()
+
+    def table_names(self) -> set[str]:
+        """Returns the list of all tables in this schema."""
+        return self._raw_schema.table_names
 
     def table(self, name: str) -> Table:
-        """Return the table with the given ``name`` from this database."""
-        return Table(self.db.table(name))
+        """Return the table with the given ``name`` from this schema."""
+        return Table(self._raw_schema.table(name))
+
+    def register_table(self, name, table) -> None:
+        """Register a table provider in this schema."""
+        return self._raw_schema.register_table(name, table)
+
+    def deregister_table(self, name: str) -> None:
+        """Deregister a table provider from this schema."""
+        return self._raw_schema.deregister_table(name)
+
+
+@deprecated("Use `Schema` instead.")
+class Database(Schema):
+    """See `Schema`."""
 
 
 class Table:
     """DataFusion table."""
 
-    def __init__(self, table: df_internal.Table) -> None:
+    def __init__(self, table: df_internal.catalog.RawTable) -> None:
         """This constructor is not typically called by the end user."""
         self.table = table
 

python/tests/test_catalog.py

@@ -15,8 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import datafusion as dfn
 import pyarrow as pa
+import pyarrow.dataset as ds
 import pytest
+from datafusion import SessionContext, Table
 
 
 # Note we take in `database` as a variable even though we don't use
@@ -27,7 +30,7 @@ def test_basic(ctx, database):
         ctx.catalog("non-existent")
 
     default = ctx.catalog()
-    assert default.names() == ["public"]
+    assert default.names() == {"public"}
 
     for db in [default.database("public"), default.database()]:
         assert db.names() == {"csv1", "csv", "csv2"}
@@ -41,3 +44,100 @@ def test_basic(ctx, database):
             pa.field("float", pa.float64(), nullable=True),
         ]
     )
+
+
+class CustomTableProvider:
+    def __init__(self):
+        pass
+
+
+def create_dataset() -> pa.dataset.Dataset:
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
+        names=["a", "b"],
+    )
+    return ds.dataset([batch])
+
+
+class CustomSchemaProvider:
+    def __init__(self):
+        self.tables = {"table1": create_dataset()}
+
+    def table_names(self) -> set[str]:
+        return set(self.tables.keys())
+
+    def register_table(self, name: str, table: Table):
+        self.tables[name] = table
+
+    def deregister_table(self, name, cascade: bool = True):
+        del self.tables[name]
+
+
+class CustomCatalogProvider:
+    def __init__(self):
+        self.schemas = {"my_schema": CustomSchemaProvider()}
+
+    def schema_names(self) -> set[str]:
+        return set(self.schemas.keys())
+
+    def schema(self, name: str):
+        return self.schemas[name]
+
+    def register_schema(self, name: str, schema: dfn.catalog.Schema):
+        self.schemas[name] = schema
+
+    def deregister_schema(self, name, cascade: bool):
+        del self.schemas[name]
+
+
+def test_python_catalog_provider(ctx: SessionContext):
+    ctx.register_catalog_provider("my_catalog", CustomCatalogProvider())
+
+    # Check the default catalog provider
+    assert ctx.catalog("datafusion").names() == {"public"}
+
+    my_catalog = ctx.catalog("my_catalog")
+    assert my_catalog.names() == {"my_schema"}
+
+    my_catalog.register_schema("second_schema", CustomSchemaProvider())
+    assert my_catalog.schema_names() == {"my_schema", "second_schema"}
+
+    my_catalog.deregister_schema("my_schema")
+    assert my_catalog.schema_names() == {"second_schema"}
+
+
+def test_python_schema_provider(ctx: SessionContext):
+    catalog = ctx.catalog()
+
+    catalog.deregister_schema("public")
+
+    catalog.register_schema("test_schema1", CustomSchemaProvider())
+    assert catalog.names() == {"test_schema1"}
+
+    catalog.register_schema("test_schema2", CustomSchemaProvider())
+    catalog.deregister_schema("test_schema1")
+    assert catalog.names() == {"test_schema2"}
+
+
+def test_python_table_provider(ctx: SessionContext):
+    catalog = ctx.catalog()
+
+    catalog.register_schema("custom_schema", CustomSchemaProvider())
+    schema = catalog.schema("custom_schema")
+
+    assert schema.table_names() == {"table1"}
+
+    schema.deregister_table("table1")
+    schema.register_table("table2", create_dataset())
+    assert schema.table_names() == {"table2"}
+
+    # Use the default schema instead of our custom schema
+
+    schema = catalog.schema()
+
+    schema.register_table("table3", create_dataset())
+    assert schema.table_names() == {"table3"}
+
+    schema.deregister_table("table3")
+    schema.register_table("table4", create_dataset())
+    assert schema.table_names() == {"table4"}