chore: Use SchemaReference from datafusion instead of our own. (#1327)

Signed-off-by: Vaibhav <[email protected]>
GlareDB · Jul 17, 2023 · 6d2eb40 · 6d2eb40
1 parent 1b3e939
commit 6d2eb40
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 25 deletions.
diff --git a/crates/sqlexec/src/context.rs b/crates/sqlexec/src/context.rs
@@ -10,7 +10,7 @@ use crate::planner::session_planner::SessionPlanner;
 use crate::vars::SessionVars;
 use datafusion::arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::common::Column as DfColumn;
+use datafusion::common::{Column as DfColumn, SchemaReference};
 use datafusion::config::{CatalogOptions, ConfigOptions, OptimizerOptions};
 use datafusion::datasource::MemTable;
 use datafusion::execution::context::{SessionConfig, SessionState, TaskContext};
@@ -752,10 +752,12 @@ impl SessionContext {
     }
 
     /// Resolve schema reference.
-    fn resolve_schema_ref(r: SchemaReference) -> (String, String) {
+    fn resolve_schema_ref(r: SchemaReference<'_>) -> (String, String) {
         match r {
-            SchemaReference::Bare { schema } => (DEFAULT_CATALOG.to_owned(), schema),
-            SchemaReference::Full { schema, catalog } => (catalog, schema),
+            SchemaReference::Bare { schema } => (DEFAULT_CATALOG.to_owned(), schema.into_owned()),
+            SchemaReference::Full { schema, catalog } => {
+                (catalog.into_owned(), schema.into_owned())
+            }
         }
     }
 

diff --git a/crates/sqlexec/src/planner/logical_plan.rs b/crates/sqlexec/src/planner/logical_plan.rs
@@ -1,6 +1,6 @@
 use crate::errors::{internal, Result};
 use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
-use datafusion::common::OwnedTableReference;
+use datafusion::common::{OwnedSchemaReference, OwnedTableReference};
 use datafusion::datasource::TableProvider;
 use datafusion::logical_expr::{Explain, LogicalPlan as DfLogicalPlan};
 use datafusion::scalar::ScalarValue;
@@ -200,7 +200,7 @@ pub struct CreateCredentials {
 
 #[derive(Clone, Debug)]
 pub struct CreateSchema {
-    pub schema_name: SchemaReference,
+    pub schema_name: OwnedSchemaReference,
     pub if_not_exists: bool,
 }
 
@@ -260,7 +260,7 @@ pub struct DropViews {
 
 #[derive(Clone, Debug)]
 pub struct DropSchemas {
-    pub names: Vec<SchemaReference>,
+    pub names: Vec<OwnedSchemaReference>,
     pub if_exists: bool,
     pub cascade: bool,
 }
@@ -355,11 +355,3 @@ impl SetVariable {
 pub struct ShowVariable {
     pub variable: String,
 }
-
-/// [`SchemaReference`] represents a multi-part schema that may require further
-/// resolution.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum SchemaReference {
-    Bare { schema: String },
-    Full { schema: String, catalog: String },
-}
diff --git a/crates/sqlexec/src/planner/session_planner.rs b/crates/sqlexec/src/planner/session_planner.rs
@@ -12,7 +12,7 @@ use crate::planner::preprocess::{preprocess, CastRegclassReplacer, EscapedString
 use datafusion::arrow::datatypes::{
     DataType, Field, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE,
 };
-use datafusion::common::OwnedTableReference;
+use datafusion::common::{OwnedSchemaReference, OwnedTableReference};
 use datafusion::sql::planner::{object_name_to_table_reference, IdentNormalizer};
 use datafusion::sql::sqlparser::ast::AlterTableOperation;
 use datafusion::sql::sqlparser::ast::{self, Ident, ObjectName, ObjectType};
@@ -647,8 +647,8 @@ impl<'a> SessionPlanner<'a> {
                     }
                     ast::SchemaName::UnnamedAuthorization(ident) => {
                         validate_ident(&ident)?;
-                        SchemaReference::Bare {
-                            schema: normalize_ident(ident),
+                        OwnedSchemaReference::Bare {
+                            schema: normalize_ident(ident).into(),
                         }
                     }
                     ast::SchemaName::NamedAuthorization(name, ident) => {
@@ -1264,15 +1264,13 @@ fn quoted_table_ref(table_ref: TableReference<'_>) -> String {
     }
 }
 
-fn object_name_to_schema_ref(name: ObjectName) -> Result<SchemaReference> {
+fn object_name_to_schema_ref(name: ObjectName) -> Result<OwnedSchemaReference> {
     let r = match object_name_to_table_ref(name)? {
         // Table becomes the schema and schema becomes the catalog.
-        TableReference::Bare { table } => SchemaReference::Bare {
-            schema: table.into_owned(),
-        },
-        TableReference::Partial { schema, table } => SchemaReference::Full {
-            schema: table.into_owned(),
-            catalog: schema.into_owned(),
+        OwnedTableReference::Bare { table } => OwnedSchemaReference::Bare { schema: table },
+        OwnedTableReference::Partial { schema, table } => OwnedSchemaReference::Full {
+            schema: table,
+            catalog: schema,
         },
         tr => return Err(internal!("invalid schema object: {tr}")),
     };