Defining abstract methods for catalog and schema providers

Author: timsaucer

python/datafusion/catalog.py

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
 import datafusion._internal as df_internal
@@ -121,6 +122,8 @@ def table(self, name: str) -> Table:
 
     def register_table(self, name, table) -> None:
         """Register a table provider in this schema."""
+        if isinstance(table, Table):
+            return self._raw_schema.register_table(name, table.table)
         return self._raw_schema.register_table(name, table)
 
     def deregister_table(self, name: str) -> None:
@@ -144,6 +147,11 @@ def __repr__(self) -> str:
         """Print a string representation of the table."""
         return self.table.__repr__()
 
+    @staticmethod
+    def from_dataset(dataset: pa.dataset.Dataset) -> Table:
+        """Turn a pyarrow Dataset into a Table."""
+        return Table(df_internal.catalog.RawTable.from_dataset(dataset))
+
     @property
     def schema(self) -> pa.Schema:
         """Returns the schema associated with this table."""
@@ -153,3 +161,71 @@ def schema(self) -> pa.Schema:
     def kind(self) -> str:
         """Returns the kind of table."""
         return self.table.kind
+
+
+class CatalogProvider(ABC):
+    @abstractmethod
+    def schema_names(self) -> set[str]:
+        """Set of the names of all schemas in this catalog."""
+        ...
+
+    @abstractmethod
+    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
+        """Add a schema to this catalog.
+
+        This method is optional. If your catalog provides a fixed list of schemas, you
+        do not need to implement this method.
+        """
+
+    def deregister_schema(self, name: str, cascade: bool) -> None:  # noqa: B027
+        """Remove a schema from this catalog.
+
+        This method is optional. If your catalog provides a fixed list of schemas, you
+        do not need to implement this method.
+
+        Args:
+            name: The name of the schema to remove.
+            cascade: If true, deregister the tables within the schema.
+        """
+
+
+class SchemaProvider(ABC):
+    def owner_name(self) -> str | None:
+        """Returns the owner of the schema.
+
+        This is an optional method. The default return is None.
+        """
+        return None
+
+    @abstractmethod
+    def table_names(self) -> set[str]:
+        """Set of the names of all tables in this schema."""
+        ...
+
+    @abstractmethod
+    def table(self, name: str) -> Table | None:
+        """Retrieve a specific table from this schema."""
+        ...
+
+    def register_table(self, name: str, table: Table) -> None:  # noqa: B027
+        """Add a table from this schema.
+
+        This method is optional. If your schema provides a fixed list of tables, you do
+        not need to implement this method.
+        """
+
+    def deregister_table(self, name, cascade: bool) -> None:  # noqa: B027
+        """Remove a table from this schema.
+
+        This method is optional. If your schema provides a fixed list of tables, you do
+        not need to implement this method.
+        """
+
+    @abstractmethod
+    def table_exist(self, name: str) -> bool:
+        """Returns true if the table exists in this schema."""
+        ...

python/tests/test_catalog.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
 
 import datafusion as dfn
 import pyarrow as pa
@@ -46,20 +47,16 @@ def test_basic(ctx, database):
     )
 
 
-class CustomTableProvider:
-    def __init__(self):
-        pass
-
-
-def create_dataset() -> pa.dataset.Dataset:
+def create_dataset() -> Table:
     batch = pa.RecordBatch.from_arrays(
         [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
         names=["a", "b"],
     )
-    return ds.dataset([batch])
+    dataset = ds.dataset([batch])
+    return Table.from_dataset(dataset)
 
 
-class CustomSchemaProvider:
+class CustomSchemaProvider(dfn.catalog.SchemaProvider):
     def __init__(self):
         self.tables = {"table1": create_dataset()}
 
@@ -72,8 +69,14 @@ def register_table(self, name: str, table: Table):
     def deregister_table(self, name, cascade: bool = True):
         del self.tables[name]
 
+    def table(self, name: str) -> Table | None:
+        return self.tables[name]
+
+    def table_exist(self, name: str) -> bool:
+        return name in self.tables
+
 
-class CustomCatalogProvider:
+class CustomCatalogProvider(dfn.catalog.CatalogProvider):
     def __init__(self):
         self.schemas = {"my_schema": CustomSchemaProvider()}
 

src/catalog.rs

@@ -207,9 +207,14 @@ impl PySchema {
             let provider: ForeignTableProvider = provider.into();
             Arc::new(provider) as Arc<dyn TableProvider>
         } else {
-            let py = table_provider.py();
-            let provider = Dataset::new(&table_provider, py)?;
-            Arc::new(provider) as Arc<dyn TableProvider>
+            match table_provider.extract::<PyTable>() {
+                Ok(py_table) => py_table.table,
+                Err(_) => {
+                    let py = table_provider.py();
+                    let provider = Dataset::new(&table_provider, py)?;
+                    Arc::new(provider) as Arc<dyn TableProvider>
+                }
+            }
         };
 
         let _ = self
@@ -238,6 +243,14 @@ impl PyTable {
         self.table.schema().to_pyarrow(py)
     }
 
+    #[staticmethod]
+    fn from_dataset(py: Python<'_>, dataset: &Bound<'_, PyAny>) -> PyResult<Self> {
+        let ds = Arc::new(Dataset::new(dataset, py).map_err(py_datafusion_err)?)
+            as Arc<dyn TableProvider>;
+
+        Ok(Self::new(ds))
+    }
+
     /// Get the type of this table for metadata/catalog purposes.
     #[getter]
     fn kind(&self) -> &str {

src/context.rs

@@ -628,17 +628,10 @@ impl PySessionContext {
             let provider: ForeignCatalogProvider = provider.into();
             Arc::new(provider) as Arc<dyn CatalogProvider>
         } else {
-            println!("Provider has type {}", provider.get_type());
             match provider.extract::<PyCatalog>() {
-                Ok(py_catalog) => {
-                    println!("registering an existing PyCatalog");
-                    py_catalog.catalog
-                }
-                Err(_) => {
-                    println!("registering a rust wrapped catalog provider");
-                    Arc::new(RustWrappedPyCatalogProvider::new(provider.into()))
-                        as Arc<dyn CatalogProvider>
-                }
+                Ok(py_catalog) => py_catalog.catalog,
+                Err(_) => Arc::new(RustWrappedPyCatalogProvider::new(provider.into()))
+                    as Arc<dyn CatalogProvider>,
             }
         };