Add default in memory options for adding schema and catalogs

Author: timsaucer

python/datafusion/catalog.py

@@ -73,6 +73,11 @@ 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."""
         return self.catalog.register_schema(name, schema)

python/datafusion/context.py

@@ -755,6 +755,15 @@ def deregister_table(self, name: str) -> None:
         """Remove a table from the session."""
         self.ctx.deregister_table(name)
 
+    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
     ) -> None:

src/catalog.rs

@@ -19,6 +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::common::DataFusionError;
 use datafusion::{
     arrow::pyarrow::ToPyArrow,
@@ -105,6 +106,16 @@ 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

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;
+use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider};
 use datafusion::common::TableReference;
 use datafusion::common::{exec_err, ScalarValue};
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
@@ -612,6 +612,13 @@ 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,
@@ -870,6 +877,10 @@ impl PySessionContext {
         })
     }
 
+    pub fn catalog_names(&self) -> HashSet<String> {
+        self.ctx.catalog_names().into_iter().collect()
+    }
+
     pub fn tables(&self) -> HashSet<String> {
         self.ctx
             .catalog_names()