ARROW-12037: [Rust] [DataFusion] Support catalogs and schemas for table namespacing

rust/benchmarks/src/bin/tpch.rs

@@ -157,9 +157,9 @@ async fn benchmark(opt: BenchmarkOpt) -> Result<Vec<arrow::record_batch::RecordB
                 table,
                 start.elapsed().as_millis()
             );
-            ctx.register_table(table, Arc::new(memtable));
+            ctx.register_table(*table, Arc::new(memtable))?;
         } else {
-            ctx.register_table(table, table_provider);
+            ctx.register_table(*table, table_provider)?;
         }
     }
 
@@ -1105,7 +1105,7 @@ fn get_table(
     table: &str,
     table_format: &str,
     max_concurrency: usize,
-) -> Result<Arc<dyn TableProvider + Send + Sync>> {
+) -> Result<Arc<dyn TableProvider>> {
     match table_format {
         // dbgen creates .tbl ('|' delimited) files without header
         "tbl" => {
@@ -1614,7 +1614,7 @@ mod tests {
 
             let provider = MemTable::try_new(Arc::new(schema), vec![vec![batch]])?;
 
-            ctx.register_table(table, Arc::new(provider));
+            ctx.register_table(table, Arc::new(provider))?;
         }
 
         let plan = create_logical_plan(&mut ctx, n)?;

rust/datafusion/benches/aggregate_query_sql.rs

@@ -150,7 +150,7 @@ fn create_context(
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     let provider = MemTable::try_new(schema, partitions)?;
-    ctx.register_table("t", Arc::new(provider));
+    ctx.register_table("t", Arc::new(provider))?;
 
     Ok(Arc::new(Mutex::new(ctx)))
 }

rust/datafusion/benches/filter_query_sql.rs

@@ -62,7 +62,7 @@ fn create_context(array_len: usize, batch_size: usize) -> Result<ExecutionContex
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     let provider = MemTable::try_new(schema, vec![batches])?;
-    ctx.register_table("t", Arc::new(provider));
+    ctx.register_table("t", Arc::new(provider))?;
 
     Ok(ctx)
 }

rust/datafusion/benches/math_query_sql.rs

@@ -72,7 +72,7 @@ fn create_context(
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     let provider = MemTable::try_new(schema, vec![batches])?;
-    ctx.register_table("t", Arc::new(provider));
+    ctx.register_table("t", Arc::new(provider))?;
 
     Ok(Arc::new(Mutex::new(ctx)))
 }

rust/datafusion/benches/sort_limit_query_sql.rs

@@ -81,7 +81,8 @@ fn create_context() -> Arc<Mutex<ExecutionContext>> {
         // create local execution context
         let mut ctx = ExecutionContext::new();
         ctx.state.lock().unwrap().config.concurrency = 1;
-        ctx.register_table("aggregate_test_100", Arc::new(mem_table));
+        ctx.register_table("aggregate_test_100", Arc::new(mem_table))
+            .unwrap();
         ctx_holder.lock().unwrap().push(Arc::new(Mutex::new(ctx)))
     });
 

rust/datafusion/examples/dataframe_in_memory.rs

@@ -49,7 +49,7 @@ async fn main() -> Result<()> {
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     let provider = MemTable::try_new(schema, vec![vec![batch]])?;
-    ctx.register_table("t", Arc::new(provider));
+    ctx.register_table("t", Arc::new(provider))?;
     let df = ctx.table("t")?;
 
     // construct an expression corresponding to "SELECT a, b FROM t WHERE b = 10" in SQL

rust/datafusion/examples/simple_udaf.rs

@@ -48,7 +48,7 @@ fn create_context() -> Result<ExecutionContext> {
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     let provider = MemTable::try_new(schema, vec![vec![batch1], vec![batch2]])?;
-    ctx.register_table("t", Arc::new(provider));
+    ctx.register_table("t", Arc::new(provider))?;
     Ok(ctx)
 }
 

rust/datafusion/examples/simple_udf.rs

@@ -50,7 +50,7 @@ fn create_context() -> Result<ExecutionContext> {
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     let provider = MemTable::try_new(schema, vec![vec![batch]])?;
-    ctx.register_table("t", Arc::new(provider));
+    ctx.register_table("t", Arc::new(provider))?;
     Ok(ctx)
 }
 

rust/datafusion/src/catalog/catalog.rs

@@ -0,0 +1,78 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Describes the interface and built-in implementations of catalogs,
+//! representing collections of named schemas.
+
+use crate::catalog::schema::SchemaProvider;
+use std::any::Any;
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+/// Represents a catalog, comprising a number of named schemas.
+pub trait CatalogProvider: Sync + Send {
+    /// Returns the catalog provider as [`Any`](std::any::Any)
+    /// so that it can be downcast to a specific implementation.
+    fn as_any(&self) -> &dyn Any;
+
+    /// Retrieves the list of available schema names in this catalog.
+    fn schema_names(&self) -> Vec<String>;
+
+    /// Retrieves a specific schema from the catalog by name, provided it exists.
+    fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>>;
+}
+
+/// Simple in-memory implementation of a catalog.
+pub struct MemoryCatalogProvider {
+    schemas: RwLock<HashMap<String, Arc<dyn SchemaProvider>>>,
+}
+
+impl MemoryCatalogProvider {
+    /// Instantiates a new MemoryCatalogProvider with an empty collection of schemas.
+    pub fn new() -> Self {
+        Self {
+            schemas: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Adds a new schema to this catalog.
+    /// If a schema of the same name existed before, it is replaced in the catalog and returned.
+    pub fn register_schema(
+        &self,
+        name: impl Into<String>,
+        schema: Arc<dyn SchemaProvider>,
+    ) -> Option<Arc<dyn SchemaProvider>> {
+        let mut schemas = self.schemas.write().unwrap();
+        schemas.insert(name.into(), schema)
+    }
+}
+
+impl CatalogProvider for MemoryCatalogProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema_names(&self) -> Vec<String> {
+        let schemas = self.schemas.read().unwrap();
+        schemas.keys().cloned().collect()
+    }
+
+    fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
+        let schemas = self.schemas.read().unwrap();
+        schemas.get(name).cloned()
+    }
+}

rust/datafusion/src/catalog/mod.rs

@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This module contains interfaces and default implementations
+//! of table namespacing concepts, including catalogs and schemas.
+
+pub mod catalog;
+pub mod schema;
+
+use crate::error::DataFusionError;
+use std::convert::TryFrom;
+
+/// Represents a resolved path to a table of the form "catalog.schema.table"
+#[derive(Clone, Copy)]
+pub struct ResolvedTableReference<'a> {
+    /// The catalog (aka database) containing the table
+    pub catalog: &'a str,
+    /// The schema containing the table
+    pub schema: &'a str,
+    /// The table name
+    pub table: &'a str,
+}
+
+/// Represents a path to a table that may require further resolution
+#[derive(Clone, Copy)]
+pub enum TableReference<'a> {
+    /// An unqualified table reference, e.g. "table"
+    Bare {
+        /// The table name
+        table: &'a str,
+    },
+    /// A partially resolved table reference, e.g. "schema.table"
+    Partial {
+        /// The schema containing the table
+        schema: &'a str,
+        /// The table name
+        table: &'a str,
+    },
+    /// A fully resolved table reference, e.g. "catalog.schema.table"
+    Full {
+        /// The catalog (aka database) containing the table
+        catalog: &'a str,
+        /// The schema containing the table
+        schema: &'a str,
+        /// The table name
+        table: &'a str,
+    },
+}
+
+impl<'a> TableReference<'a> {
+    /// Retrieve the actual table name, regardless of qualification
+    pub fn table(&self) -> &str {
+        match self {
+            Self::Full { table, .. }
+            | Self::Partial { table, .. }
+            | Self::Bare { table } => table,
+        }
+    }
+
+    /// Given a default catalog and schema, ensure this table reference is fully resolved
+    pub fn resolve(
+        self,
+        default_catalog: &'a str,
+        default_schema: &'a str,
+    ) -> ResolvedTableReference<'a> {
+        match self {
+            Self::Full {
+                catalog,
+                schema,
+                table,
+            } => ResolvedTableReference {
+                catalog,
+                schema,
+                table,
+            },
+            Self::Partial { schema, table } => ResolvedTableReference {
+                catalog: default_catalog,
+                schema,
+                table,
+            },
+            Self::Bare { table } => ResolvedTableReference {
+                catalog: default_catalog,
+                schema: default_schema,
+                table,
+            },
+        }
+    }
+}
+
+impl<'a> From<&'a str> for TableReference<'a> {
+    fn from(s: &'a str) -> Self {
+        Self::Bare { table: s }
+    }
+}
+
+impl<'a> From<ResolvedTableReference<'a>> for TableReference<'a> {
+    fn from(resolved: ResolvedTableReference<'a>) -> Self {
+        Self::Full {
+            catalog: resolved.catalog,
+            schema: resolved.schema,
+            table: resolved.table,
+        }
+    }
+}
+
+impl<'a> TryFrom<&'a sqlparser::ast::ObjectName> for TableReference<'a> {
+    type Error = DataFusionError;
+
+    fn try_from(value: &'a sqlparser::ast::ObjectName) -> Result<Self, Self::Error> {
+        let idents = &value.0;
+
+        match idents.len() {
+            1 => Ok(Self::Bare {
+                table: &idents[0].value,
+            }),
+            2 => Ok(Self::Partial {
+                schema: &idents[0].value,
+                table: &idents[1].value,
+            }),
+            3 => Ok(Self::Full {
+                catalog: &idents[0].value,
+                schema: &idents[1].value,
+                table: &idents[2].value,
+            }),
+            _ => Err(DataFusionError::Plan(format!(
+                "invalid table reference: {}",
+                value
+            ))),
+        }
+    }
+}