Add support for creating in memory catalog and schema

Author: timsaucer

python/datafusion/catalog.py

@@ -58,6 +58,12 @@ def schema_names(self) -> set[str]:
         """Returns the list of schemas in this catalog."""
         return self.catalog.schema_names()
 
+    @staticmethod
+    def memory_catalog() -> Catalog:
+        """Create an in-memory catalog provider."""
+        catalog = df_internal.catalog.RawCatalog.memory_catalog()
+        return Catalog(catalog)
+
     def schema(self, name: str = "public") -> Schema:
         """Returns the database with the given ``name`` from this catalog."""
         schema = self.catalog.schema(name)
@@ -73,13 +79,10 @@ def database(self, name: str = "public") -> Schema:
         """Returns the database with the given ``name`` from this catalog."""
         return self.schema(name)
 
-    def new_in_memory_schema(self, name: str) -> Schema:
-        """Create a new schema in this catalog using an in-memory provider."""
-        self.catalog.new_in_memory_schema(name)
-        return self.schema(name)
-
     def register_schema(self, name, schema) -> Schema | None:
         """Register a schema with this catalog."""
+        if isinstance(schema, Schema):
+            return self.catalog.register_schema(name, schema._raw_schema)
         return self.catalog.register_schema(name, schema)
 
     def deregister_schema(self, name: str, cascade: bool = True) -> Schema | None:
@@ -98,6 +101,12 @@ def __repr__(self) -> str:
         """Print a string representation of the schema."""
         return self._raw_schema.__repr__()
 
+    @staticmethod
+    def memory_schema() -> Schema:
+        """Create an in-memory schema provider."""
+        schema = df_internal.catalog.RawSchema.memory_schema()
+        return Schema(schema)
+
     def names(self) -> set[str]:
         """This is an alias for `table_names`."""
         return self.table_names()

python/datafusion/context.py

@@ -759,16 +759,14 @@ def catalog_names(self) -> set[str]:
         """Returns the list of catalogs in this context."""
         return self.ctx.catalog_names()
 
-    def new_in_memory_catalog(self, name: str) -> Catalog:
-        """Create a new catalog in this context using an in-memory provider."""
-        self.ctx.new_in_memory_catalog(name)
-        return self.catalog(name)
-
     def register_catalog_provider(
-        self, name: str, provider: CatalogProviderExportable
+        self, name: str, provider: CatalogProviderExportable | Catalog
     ) -> None:
         """Register a catalog provider."""
-        self.ctx.register_catalog_provider(name, provider)
+        if isinstance(provider, Catalog):
+            self.ctx.register_catalog_provider(name, provider.catalog)
+        else:
+            self.ctx.register_catalog_provider(name, provider)
 
     def register_table_provider(
         self, name: str, provider: TableProviderExportable

python/tests/test_catalog.py

@@ -106,6 +106,24 @@ def test_python_catalog_provider(ctx: SessionContext):
     assert my_catalog.schema_names() == {"second_schema"}
 
 
+def test_in_memory_providers(ctx: SessionContext):
+    catalog = dfn.catalog.Catalog.memory_catalog()
+    ctx.register_catalog_provider("in_mem_catalog", catalog)
+
+    assert ctx.catalog_names() == {"datafusion", "in_mem_catalog"}
+
+    schema = dfn.catalog.Schema.memory_schema()
+    catalog.register_schema("in_mem_schema", schema)
+
+    schema.register_table("my_table", create_dataset())
+
+    batches = ctx.sql("select * from in_mem_catalog.in_mem_schema.my_table").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])
+
+
 def test_python_schema_provider(ctx: SessionContext):
     catalog = ctx.catalog()
 

src/catalog.rs

@@ -19,7 +19,7 @@ use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
 use crate::utils::{validate_pycapsule, wait_for_future};
 use async_trait::async_trait;
-use datafusion::catalog::MemorySchemaProvider;
+use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion::common::DataFusionError;
 use datafusion::{
     arrow::pyarrow::ToPyArrow,
@@ -37,16 +37,19 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 #[pyclass(name = "RawCatalog", module = "datafusion.catalog", subclass)]
+#[derive(Clone)]
 pub struct PyCatalog {
     pub catalog: Arc<dyn CatalogProvider>,
 }
 
 #[pyclass(name = "RawSchema", module = "datafusion.catalog", subclass)]
+#[derive(Clone)]
 pub struct PySchema {
     pub schema: Arc<dyn SchemaProvider>,
 }
 
 #[pyclass(name = "RawTable", module = "datafusion.catalog", subclass)]
+#[derive(Clone)]
 pub struct PyTable {
     pub table: Arc<dyn TableProvider>,
 }
@@ -82,6 +85,13 @@ impl PyCatalog {
         catalog_provider.into()
     }
 
+    #[staticmethod]
+    fn memory_catalog() -> Self {
+        let catalog_provider =
+            Arc::new(MemoryCatalogProvider::default()) as Arc<dyn CatalogProvider>;
+        catalog_provider.into()
+    }
+
     fn schema_names(&self) -> HashSet<String> {
         self.catalog.schema_names().into_iter().collect()
     }
@@ -106,16 +116,6 @@ impl PyCatalog {
         })
     }
 
-    fn new_in_memory_schema(&mut self, name: &str) -> PyResult<()> {
-        let schema = Arc::new(MemorySchemaProvider::new()) as Arc<dyn SchemaProvider>;
-        let _ = self
-            .catalog
-            .register_schema(name, schema)
-            .map_err(py_datafusion_err)?;
-
-        Ok(())
-    }
-
     fn register_schema(&self, name: &str, schema_provider: Bound<'_, PyAny>) -> PyResult<()> {
         let provider = if schema_provider.hasattr("__datafusion_schema_provider__")? {
             let capsule = schema_provider
@@ -128,8 +128,11 @@ impl PyCatalog {
             let provider: ForeignSchemaProvider = provider.into();
             Arc::new(provider) as Arc<dyn SchemaProvider>
         } else {
-            let provider = RustWrappedPySchemaProvider::new(schema_provider.into());
-            Arc::new(provider) as Arc<dyn SchemaProvider>
+            match schema_provider.extract::<PySchema>() {
+                Ok(py_schema) => py_schema.schema,
+                Err(_) => Arc::new(RustWrappedPySchemaProvider::new(schema_provider.into()))
+                    as Arc<dyn SchemaProvider>,
+            }
         };
 
         let _ = self
@@ -165,6 +168,12 @@ impl PySchema {
         schema_provider.into()
     }
 
+    #[staticmethod]
+    fn memory_schema() -> Self {
+        let schema_provider = Arc::new(MemorySchemaProvider::default()) as Arc<dyn SchemaProvider>;
+        schema_provider.into()
+    }
+
     #[getter]
     fn table_names(&self) -> HashSet<String> {
         self.schema.table_names().into_iter().collect()

src/context.rs

@@ -49,7 +49,7 @@ use crate::utils::{get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_f
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider};
+use datafusion::catalog::CatalogProvider;
 use datafusion::common::TableReference;
 use datafusion::common::{exec_err, ScalarValue};
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
@@ -612,13 +612,6 @@ impl PySessionContext {
         Ok(())
     }
 
-    pub fn new_in_memory_catalog(&mut self, name: &str) -> PyResult<()> {
-        let catalog = Arc::new(MemoryCatalogProvider::new()) as Arc<dyn CatalogProvider>;
-        let _ = self.ctx.register_catalog(name, catalog);
-
-        Ok(())
-    }
-
     pub fn register_catalog_provider(
         &mut self,
         name: &str,
@@ -635,8 +628,18 @@ impl PySessionContext {
             let provider: ForeignCatalogProvider = provider.into();
             Arc::new(provider) as Arc<dyn CatalogProvider>
         } else {
-            let provider = RustWrappedPyCatalogProvider::new(provider.into());
-            Arc::new(provider) as Arc<dyn CatalogProvider>
+            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>
+                }
+            }
         };
 
         let _ = self.ctx.register_catalog(name, provider);