Working through issues between custom catalog and build in schema

Author: timsaucer

python/datafusion/catalog.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
 import datafusion._internal as df_internal
 
@@ -174,7 +174,9 @@ def schema(self, name: str) -> Schema | None:
         """Retrieve a specific schema from this catalog."""
         ...
 
-    def register_schema(self, name: str, schema: Schema) -> None:  # noqa: B027
+    def register_schema(  # noqa: B027
+        self, name: str, schema: SchemaProviderExportable | SchemaProvider | Schema
+    ) -> None:
         """Add a schema to this catalog.
 
         This method is optional. If your catalog provides a fixed list of schemas, you
@@ -229,3 +231,12 @@ def deregister_table(self, name, cascade: bool) -> None:  # noqa: B027
     def table_exist(self, name: str) -> bool:
         """Returns true if the table exists in this schema."""
         ...
+
+
+class SchemaProviderExportable(Protocol):
+    """Type hint for object that has __datafusion_schema_provider__ PyCapsule.
+
+    https://docs.rs/datafusion/latest/datafusion/catalog/trait.SchemaProvider.html
+    """
+
+    def __datafusion_schema_provider__(self) -> object: ...

python/datafusion/context.py

@@ -29,7 +29,7 @@
 except ImportError:
     from typing_extensions import deprecated  # Python 3.12
 
-from datafusion.catalog import Catalog, Table
+from datafusion.catalog import Catalog, CatalogProvider, Table
 from datafusion.dataframe import DataFrame
 from datafusion.expr import Expr, SortExpr, sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
@@ -763,7 +763,7 @@ def catalog_names(self) -> set[str]:
         return self.ctx.catalog_names()
 
     def register_catalog_provider(
-        self, name: str, provider: CatalogProviderExportable | Catalog
+        self, name: str, provider: CatalogProviderExportable | CatalogProvider | Catalog
     ) -> None:
         """Register a catalog provider."""
         if isinstance(provider, Catalog):

python/tests/test_catalog.py

@@ -162,3 +162,51 @@ def test_python_table_provider(ctx: SessionContext):
     schema.deregister_table("table3")
     schema.register_table("table4", create_dataset())
     assert schema.table_names() == {"table4"}
+
+
+def test_in_end_to_end_python_providers(ctx: SessionContext):
+    """Test registering all python providers and running a query against them."""
+
+    all_catalog_names = [
+        "datafusion",
+        "custom_catalog",
+        "in_mem_catalog",
+    ]
+
+    all_schema_names = [
+        "custom_schema",
+        "in_mem_schema",
+    ]
+
+    ctx.register_catalog_provider(all_catalog_names[1], CustomCatalogProvider())
+    ctx.register_catalog_provider(
+        all_catalog_names[2], dfn.catalog.Catalog.memory_catalog()
+    )
+
+    for catalog_name in all_catalog_names:
+        catalog = ctx.catalog(catalog_name)
+
+        # Clean out previous schemas if they exist so we can start clean
+        for schema_name in catalog.schema_names():
+            catalog.deregister_schema(schema_name, cascade=False)
+
+        catalog.register_schema(all_schema_names[0], CustomSchemaProvider())
+        catalog.register_schema(all_schema_names[1], dfn.catalog.Schema.memory_schema())
+
+        for schema_name in all_schema_names:
+            schema = catalog.schema(schema_name)
+
+            for table_name in schema.table_names():
+                schema.deregister_table(table_name)
+
+            schema.register_table("test_table", create_dataset())
+
+    for catalog_name in all_catalog_names:
+        for schema_name in all_schema_names:
+            table_full_name = f"{catalog_name}.{schema_name}.test_table"
+
+            batches = ctx.sql(f"select * from {table_full_name}").collect()
+
+            assert len(batches) == 1
+            assert batches[0].column(0) == pa.array([1, 2, 3])
+            assert batches[0].column(1) == pa.array([4, 5, 6])

src/catalog.rs

@@ -314,9 +314,19 @@ impl RustWrappedPySchemaProvider {
 
                 Ok(Some(Arc::new(provider) as Arc<dyn TableProvider>))
             } else {
-                let ds = Dataset::new(&py_table, py).map_err(py_datafusion_err)?;
+                if let Ok(inner_table) = py_table.getattr("table") {
+                    if let Ok(inner_table) = inner_table.extract::<PyTable>() {
+                        return Ok(Some(inner_table.table));
+                    }
+                }
 
-                Ok(Some(Arc::new(ds) as Arc<dyn TableProvider>))
+                match py_table.extract::<PyTable>() {
+                    Ok(py_table) => Ok(Some(py_table.table)),
+                    Err(_) => {
+                        let ds = Dataset::new(&py_table, py).map_err(py_datafusion_err)?;
+                        Ok(Some(Arc::new(ds) as Arc<dyn TableProvider>))
+                    }
+                }
             }
         })
     }