alamb commented on a change in pull request #9762: URL: https://github.com/apache/arrow/pull/9762#discussion_r599027565
########## File path: 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>; Review comment: What do you think about returning something more like `Vec<&str>` to prevent requiring a copy? Ideally it would be nice if we could do something like ``` fn schema_names(&self) -> impl Iterator<Item=&str> ``` As all uses of the results here will need to iterate over the names I suspect ########## File path: 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 { Review comment: In some other PR perhaps we can put the concrete implementations into their own modules. I don't think this one is big enough to warrant that yet, however; I just wanted to point it out ########## File path: rust/datafusion/src/catalog/schema.rs ########## @@ -0,0 +1,104 @@ +// 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 schemas, +//! representing collections of named tables. + +use crate::datasource::TableProvider; +use crate::error::{DataFusionError, Result}; +use std::any::Any; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Represents a schema, comprising a number of named tables. +pub trait SchemaProvider: Sync + Send { + /// Returns the schema 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 table names in this schema. + fn table_names(&self) -> Vec<String>; + + /// Retrieves a specific table from the schema by name, provided it exists. + fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>>; + + /// If supported by the implementation, adds a new table to this schema. + /// If a table of the same name existed before, it is replaced in the schema and returned. + #[allow(unused_variables)] + fn register_table( + &self, + name: String, + table: Arc<dyn TableProvider>, + ) -> Result<Option<Arc<dyn TableProvider>>> { + Err(DataFusionError::Execution( + "schema provider does not support registering tables".to_owned(), + )) + } + + /// If supported by the implementation, removes an existing table from this schema and returns it. + /// If no table of that name exists, returns Ok(None). + #[allow(unused_variables)] + fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> { Review comment: ```suggestion fn deregister_table(&self, _name: &str) -> Result<Option<Arc<dyn TableProvider>>> { ``` ########## File path: rust/datafusion/src/catalog/schema.rs ########## @@ -0,0 +1,104 @@ +// 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 schemas, +//! representing collections of named tables. + +use crate::datasource::TableProvider; +use crate::error::{DataFusionError, Result}; +use std::any::Any; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Represents a schema, comprising a number of named tables. +pub trait SchemaProvider: Sync + Send { + /// Returns the schema 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 table names in this schema. + fn table_names(&self) -> Vec<String>; Review comment: same suggestion here to try and avoid copying strings if we can avoid it ########## File path: rust/datafusion/tests/sql.rs ########## @@ -2504,3 +2507,36 @@ async fn inner_join_qualified_names() -> Result<()> { } Ok(()) } + +#[tokio::test] +async fn qualified_table_references() -> Result<()> { + let mut ctx = ExecutionContext::new(); + register_aggregate_csv(&mut ctx)?; + + for table_ref in &[ + "aggregate_test_100", + "public.aggregate_test_100", + "datafusion.public.aggregate_test_100", + ] { + let sql = format!("SELECT COUNT(*) FROM {}", table_ref); + let results = execute(&mut ctx, &sql).await; + assert_eq!(results, vec![vec!["100"]]); + } + Ok(()) +} + +#[tokio::test] +async fn invalid_qualified_table_references() -> Result<()> { + let mut ctx = ExecutionContext::new(); + register_aggregate_csv(&mut ctx)?; + + for table_ref in &[ + "nonexistentschema.aggregate_test_100", + "nonexistentcatalog.public.aggregate_test_100", + "way.too.many.namespaces.as.ident.prefixes.aggregate_test_100", Review comment: 😆 ########## File path: rust/datafusion/src/execution/context.rs ########## @@ -564,13 +634,30 @@ impl ExecutionConfig { self.optimizers.push(optimizer_rule); self } + + /// Selects a name for the default catalog and schema + pub fn with_default_catalog_and_schema( + mut self, + catalog: impl Into<String>, + schema: impl Into<String>, + ) -> Self { + self.default_catalog = catalog.into(); + self.default_schema = schema.into(); + self + } + + /// Controls whether the default catalog and schema will be automatically created + pub fn create_default_catalog_and_schema(mut self, create: bool) -> Self { + self.create_default_catalog_and_schema = create; + self + } } /// Execution context for registering data sources and executing queries #[derive(Clone)] pub struct ExecutionContextState { - /// Data sources that are registered with the context - pub datasources: HashMap<String, Arc<dyn TableProvider + Send + Sync>>, + /// Collection of catalogs containing schemas and ultimately tables Review comment: ```suggestion /// Collection of catalogs containing schemas and ultimately TableProviders ``` ########## File path: 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> { Review comment: 👍 ########## File path: rust/datafusion/src/execution/context.rs ########## @@ -512,6 +573,12 @@ pub struct ExecutionConfig { optimizers: Vec<Arc<dyn OptimizerRule + Send + Sync>>, /// Responsible for planning `LogicalPlan`s, and `ExecutionPlan` query_planner: Arc<dyn QueryPlanner + Send + Sync>, + /// Default catalog name for table resolution + default_catalog: String, Review comment: Would we ever need to distinguish between `""` and `None`? If so, perhaps we can make this field `default_catalog: Option<String>` instead ########## File path: 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, Review comment: 👍 ########## File path: rust/datafusion/src/execution/context.rs ########## @@ -1922,6 +2035,114 @@ mod tests { Ok(()) } Review comment: these are excellent tests . 👍 ########## File path: rust/datafusion/src/execution/context.rs ########## @@ -1922,6 +2035,114 @@ mod tests { Ok(()) } + fn table_with_sequence( + seq_start: i32, + seq_end: i32, + ) -> Result<Arc<dyn TableProvider>> { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let arr = Arc::new(Int32Array::from((seq_start..=seq_end).collect::<Vec<_>>())); + let partitions = vec![vec![RecordBatch::try_new( + schema.clone(), + vec![arr as ArrayRef], + )?]]; + Ok(Arc::new(MemTable::try_new(schema, partitions)?)) + } + + #[tokio::test] + async fn disabled_default_catalog_and_schema() -> Result<()> { + let mut ctx = ExecutionContext::with_config( + ExecutionConfig::new().create_default_catalog_and_schema(false), + ); + + assert!(matches!( + ctx.register_table("test", table_with_sequence(1, 1)?), + Err(DataFusionError::Plan(_)) + )); + + assert!(matches!( + ctx.sql("select * from datafusion.public.test"), + Err(DataFusionError::Plan(_)) + )); + + Ok(()) + } + + #[tokio::test] + async fn custom_catalog_and_schema() -> Result<()> { + let mut ctx = ExecutionContext::with_config( + ExecutionConfig::new() + .create_default_catalog_and_schema(false) + .with_default_catalog_and_schema("my_catalog", "my_schema"), + ); + + let catalog = MemoryCatalogProvider::new(); + let schema = MemorySchemaProvider::new(); + schema.register_table("test".to_owned(), table_with_sequence(1, 1)?)?; + catalog.register_schema("my_schema", Arc::new(schema)); + ctx.register_catalog("my_catalog", Arc::new(catalog)); + + for table_ref in &["my_catalog.my_schema.test", "my_schema.test", "test"] { + let result = plan_and_collect( + &mut ctx, + &format!("SELECT COUNT(*) AS count FROM {}", table_ref), + ) + .await?; + + let expected = vec![ + "+-------+", + "| count |", + "+-------+", + "| 1 |", + "+-------+", + ]; + assert_batches_eq!(expected, &result); + } + + Ok(()) + } + + #[tokio::test] + async fn cross_catalog_access() -> Result<()> { + let mut ctx = ExecutionContext::new(); + + let catalog_a = MemoryCatalogProvider::new(); + let schema_a = MemorySchemaProvider::new(); + schema_a.register_table("table_a".to_owned(), table_with_sequence(1, 1)?)?; + catalog_a.register_schema("schema_a", Arc::new(schema_a)); + ctx.register_catalog("catalog_a", Arc::new(catalog_a)); + + let catalog_b = MemoryCatalogProvider::new(); + let schema_b = MemorySchemaProvider::new(); + schema_b.register_table("table_b".to_owned(), table_with_sequence(1, 2)?)?; + catalog_b.register_schema("schema_b", Arc::new(schema_b)); + ctx.register_catalog("catalog_b", Arc::new(catalog_b)); + + let result = plan_and_collect( + &mut ctx, + "SELECT cat, SUM(i) AS total FROM ( + SELECT i, 'a' AS cat FROM catalog_a.schema_a.table_a + UNION ALL Review comment: @Dandandan `UNION` is already being used! ########## File path: 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))?; Review comment: The change to make `register_table` fallible is a breaking change, but a very reasonable one I think. ########## File path: rust/datafusion/src/catalog/schema.rs ########## @@ -0,0 +1,104 @@ +// 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 schemas, +//! representing collections of named tables. + +use crate::datasource::TableProvider; +use crate::error::{DataFusionError, Result}; +use std::any::Any; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Represents a schema, comprising a number of named tables. +pub trait SchemaProvider: Sync + Send { + /// Returns the schema 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 table names in this schema. + fn table_names(&self) -> Vec<String>; + + /// Retrieves a specific table from the schema by name, provided it exists. + fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>>; + + /// If supported by the implementation, adds a new table to this schema. + /// If a table of the same name existed before, it is replaced in the schema and returned. + #[allow(unused_variables)] + fn register_table( + &self, + name: String, + table: Arc<dyn TableProvider>, Review comment: ```suggestion fn register_table( &self, _name: String, _table: Arc<dyn TableProvider>, ``` ########## File path: rust/benchmarks/src/bin/tpch.rs ########## @@ -1105,7 +1105,7 @@ fn get_table( table: &str, table_format: &str, max_concurrency: usize, -) -> Result<Arc<dyn TableProvider + Send + Sync>> { Review comment: FWIW I was confused about this change initially until I saw that you added `Send + Sync` to the `TableProvider` trait itself. I think that is a good change (so we don't have to remember to add `Send + Sync` everywhere) ########## File path: rust/datafusion/src/execution/context.rs ########## @@ -564,13 +634,30 @@ impl ExecutionConfig { self.optimizers.push(optimizer_rule); self } + + /// Selects a name for the default catalog and schema + pub fn with_default_catalog_and_schema( Review comment: I like this API. 👍 I wonder if we need an API for getting `default_catalog` and `default_schema`? ########## File path: rust/datafusion/src/catalog/schema.rs ########## @@ -0,0 +1,104 @@ +// 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 schemas, +//! representing collections of named tables. + +use crate::datasource::TableProvider; +use crate::error::{DataFusionError, Result}; +use std::any::Any; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Represents a schema, comprising a number of named tables. +pub trait SchemaProvider: Sync + Send { + /// Returns the schema 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 table names in this schema. + fn table_names(&self) -> Vec<String>; + + /// Retrieves a specific table from the schema by name, provided it exists. + fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>>; + + /// If supported by the implementation, adds a new table to this schema. + /// If a table of the same name existed before, it is replaced in the schema and returned. + #[allow(unused_variables)] + fn register_table( + &self, + name: String, + table: Arc<dyn TableProvider>, Review comment: I think that pattern is more "idiomatic rust" but the unused variables thing works too -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org
