feat: support AS (#131)

KipData · Feb 7, 2024 · b11a8c9 · b11a8c9
1 parent c6d0536
commit b11a8c9
Show file tree

Hide file tree

Showing 33 changed files with 285 additions and 416 deletions.
diff --git a/src/binder/aggregate.rs b/src/binder/aggregate.rs
@@ -160,7 +160,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
 
             if !group_raw_exprs.iter().contains(expr) {
                 return Err(DatabaseError::AggMiss(format!(
-                    "{:?} must appear in the GROUP BY clause or be used in an aggregate function",
+                    "`{}` must appear in the GROUP BY clause or be used in an aggregate function",
                     expr
                 )));
             }

diff --git a/src/binder/alter_table.rs b/src/binder/alter_table.rs
@@ -3,7 +3,7 @@ use sqlparser::ast::{AlterTableOperation, ObjectName};
 use std::sync::Arc;
 
 use super::{is_valid_identifier, Binder};
-use crate::binder::{lower_case_name, split_name};
+use crate::binder::lower_case_name;
 use crate::errors::DatabaseError;
 use crate::planner::operator::alter_table::add_column::AddColumnOperator;
 use crate::planner::operator::alter_table::drop_column::DropColumnOperator;
@@ -18,7 +18,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
         name: &ObjectName,
         operation: &AlterTableOperation,
     ) -> Result<LogicalPlan, DatabaseError> {
-        let table_name: Arc<String> = Arc::new(split_name(&lower_case_name(name))?.to_string());
+        let table_name: Arc<String> = Arc::new(lower_case_name(name)?);
 
         if let Some(table) = self.context.table(table_name.clone()) {
             let plan = match operation {

diff --git a/src/binder/analyze.rs b/src/binder/analyze.rs
@@ -1,4 +1,4 @@
-use crate::binder::{lower_case_name, split_name, Binder};
+use crate::binder::{lower_case_name, Binder};
 use crate::errors::DatabaseError;
 use crate::planner::operator::analyze::AnalyzeOperator;
 use crate::planner::operator::scan::ScanOperator;
@@ -11,9 +11,7 @@ use std::sync::Arc;
 
 impl<'a, T: Transaction> Binder<'a, T> {
     pub(crate) fn bind_analyze(&mut self, name: &ObjectName) -> Result<LogicalPlan, DatabaseError> {
-        let name = lower_case_name(name);
-        let name = split_name(&name)?;
-        let table_name = Arc::new(name.to_string());
+        let table_name = Arc::new(lower_case_name(name)?);
 
         let table_catalog = self.context.table_and_bind(table_name.clone(), None)?;
         let columns = table_catalog

diff --git a/src/binder/create_table.rs b/src/binder/create_table.rs
@@ -4,7 +4,7 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use super::{is_valid_identifier, Binder};
-use crate::binder::{lower_case_name, split_name};
+use crate::binder::lower_case_name;
 use crate::catalog::{ColumnCatalog, ColumnDesc};
 use crate::errors::DatabaseError;
 use crate::expression::ScalarExpression;
@@ -24,9 +24,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
         constraints: &[TableConstraint],
         if_not_exists: bool,
     ) -> Result<LogicalPlan, DatabaseError> {
-        let name = lower_case_name(name);
-        let name = split_name(&name)?;
-        let table_name = Arc::new(name.to_string());
+        let table_name = Arc::new(lower_case_name(name)?);
 
         if !is_valid_identifier(&table_name) {
             return Err(DatabaseError::InvalidTable(
@@ -132,7 +130,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
             }
         }
 
-        Ok(ColumnCatalog::new(column_name, nullable, column_desc, None))
+        Ok(ColumnCatalog::new(column_name, nullable, column_desc))
     }
 }
 

diff --git a/src/binder/delete.rs b/src/binder/delete.rs
@@ -1,4 +1,4 @@
-use crate::binder::{lower_case_name, split_name, Binder};
+use crate::binder::{lower_case_name, Binder};
 use crate::errors::DatabaseError;
 use crate::planner::operator::delete::DeleteOperator;
 use crate::planner::operator::scan::ScanOperator;
@@ -15,9 +15,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
         selection: &Option<Expr>,
     ) -> Result<LogicalPlan, DatabaseError> {
         if let TableFactor::Table { name, alias, .. } = &from.relation {
-            let name = lower_case_name(name);
-            let name = split_name(&name)?;
-            let table_name = Arc::new(name.to_string());
+            let table_name = Arc::new(lower_case_name(name)?);
 
             let table_catalog = self.context.table_and_bind(table_name.clone(), None)?;
             let primary_key_column = table_catalog
@@ -29,7 +27,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
 
             if let Some(alias) = alias {
                 self.context
-                    .add_table_alias(alias.to_string(), table_name.clone())?;
+                    .add_table_alias(alias.to_string(), table_name.clone());
             }
 
             if let Some(predicate) = selection {

diff --git a/src/binder/describe.rs b/src/binder/describe.rs
@@ -1,4 +1,4 @@
-use crate::binder::{lower_case_name, split_name, Binder};
+use crate::binder::{lower_case_name, Binder};
 use crate::errors::DatabaseError;
 use crate::planner::operator::describe::DescribeOperator;
 use crate::planner::operator::Operator;
@@ -12,9 +12,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
         &mut self,
         name: &ObjectName,
     ) -> Result<LogicalPlan, DatabaseError> {
-        let name = lower_case_name(&name);
-        let name = split_name(&name)?;
-        let table_name = Arc::new(name.to_string());
+        let table_name = Arc::new(lower_case_name(name)?);
 
         Ok(LogicalPlan {
             operator: Operator::Describe(DescribeOperator { table_name }),

diff --git a/src/binder/drop_table.rs b/src/binder/drop_table.rs
@@ -1,4 +1,4 @@
-use crate::binder::{lower_case_name, split_name, Binder};
+use crate::binder::{lower_case_name, Binder};
 use crate::errors::DatabaseError;
 use crate::planner::operator::drop_table::DropTableOperator;
 use crate::planner::operator::Operator;
@@ -13,9 +13,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
         name: &ObjectName,
         if_exists: &bool,
     ) -> Result<LogicalPlan, DatabaseError> {
-        let name = lower_case_name(name);
-        let name = split_name(&name)?;
-        let table_name = Arc::new(name.to_string());
+        let table_name = Arc::new(lower_case_name(name)?);
 
         let plan = LogicalPlan {
             operator: Operator::DropTable(DropTableOperator {

diff --git a/src/binder/expr.rs b/src/binder/expr.rs
@@ -8,12 +8,23 @@ use sqlparser::ast::{
 use std::slice;
 use std::sync::Arc;
 
-use super::Binder;
+use super::{lower_ident, Binder};
 use crate::expression::ScalarExpression;
 use crate::storage::Transaction;
 use crate::types::value::DataValue;
 use crate::types::LogicalType;
 
+macro_rules! try_alias {
+    ($context:expr, $column_name:expr) => {
+        if let Some(expr) = $context.expr_aliases.get(&$column_name) {
+            return Ok(ScalarExpression::Alias {
+                expr: Box::new(expr.clone()),
+                alias: $column_name,
+            });
+        }
+    };
+}
+
 impl<'a, T: Transaction> Binder<'a, T> {
     pub(crate) fn bind_expr(&mut self, expr: &Expr) -> Result<ScalarExpression, DatabaseError> {
         match expr {
@@ -76,16 +87,11 @@ impl<'a, T: Transaction> Binder<'a, T> {
     pub fn bind_column_ref_from_identifiers(
         &mut self,
         idents: &[Ident],
-        bind_table_name: Option<&String>,
+        bind_table_name: Option<String>,
     ) -> Result<ScalarExpression, DatabaseError> {
-        let idents = idents
-            .iter()
-            .map(|ident| Ident::new(ident.value.to_lowercase()))
-            .collect_vec();
-        let (_schema_name, table_name, column_name) = match idents.as_slice() {
-            [column] => (None, None, &column.value),
-            [table, column] => (None, Some(&table.value), &column.value),
-            [schema, table, column] => (Some(&schema.value), Some(&table.value), &column.value),
+        let (table_name, column_name) = match idents {
+            [column] => (None, lower_ident(column)),
+            [table, column] => (Some(lower_ident(table)), lower_ident(column)),
             _ => {
                 return Err(DatabaseError::InvalidColumn(
                     idents
@@ -96,38 +102,31 @@ impl<'a, T: Transaction> Binder<'a, T> {
                 ))
             }
         };
-
         if let Some(table) = table_name.or(bind_table_name) {
+            try_alias!(self.context, column_name);
             let table_catalog = self
                 .context
                 .table(Arc::new(table.clone()))
-                .ok_or_else(|| DatabaseError::InvalidTable(table.to_string()))?;
+                .ok_or_else(|| DatabaseError::TableNotFound)?;
 
             let column_catalog = table_catalog
-                .get_column_by_name(column_name)
-                .ok_or_else(|| DatabaseError::InvalidColumn(column_name.to_string()))?;
+                .get_column_by_name(&column_name)
+                .ok_or_else(|| DatabaseError::NotFound("column", column_name))?;
             Ok(ScalarExpression::ColumnRef(column_catalog.clone()))
         } else {
+            try_alias!(self.context, column_name);
             // handle col syntax
             let mut got_column = None;
             for (table_catalog, _) in self.context.bind_table.values() {
-                if let Some(column_catalog) = table_catalog.get_column_by_name(column_name) {
-                    if got_column.is_some() {
-                        return Err(DatabaseError::InvalidColumn(column_name.to_string()));
-                    }
+                if let Some(column_catalog) = table_catalog.get_column_by_name(&column_name) {
                     got_column = Some(column_catalog);
                 }
-            }
-            if got_column.is_none() {
-                if let Some(expr) = self.context.aliases.get(column_name) {
-                    return Ok(ScalarExpression::Alias {
-                        expr: Box::new(expr.clone()),
-                        alias: column_name.clone(),
-                    });
+                if got_column.is_some() {
+                    break;
                 }
             }
             let column_catalog =
-                got_column.ok_or_else(|| DatabaseError::InvalidColumn(column_name.to_string()))?;
+                got_column.ok_or_else(|| DatabaseError::NotFound("column", column_name))?;
             Ok(ScalarExpression::ColumnRef(column_catalog.clone()))
         }
     }

diff --git a/src/binder/insert.rs b/src/binder/insert.rs
@@ -1,4 +1,4 @@
-use crate::binder::{lower_case_name, split_name, Binder};
+use crate::binder::{lower_case_name, Binder};
 use crate::catalog::ColumnRef;
 use crate::errors::DatabaseError;
 use crate::expression::value_compute::unary_op;
@@ -16,14 +16,12 @@ use std::sync::Arc;
 impl<'a, T: Transaction> Binder<'a, T> {
     pub(crate) fn bind_insert(
         &mut self,
-        name: ObjectName,
+        name: &ObjectName,
         idents: &[Ident],
         expr_rows: &Vec<Vec<Expr>>,
         is_overwrite: bool,
     ) -> Result<LogicalPlan, DatabaseError> {
-        let name = lower_case_name(&name);
-        let name = split_name(&name)?;
-        let table_name = Arc::new(name.to_string());
+        let table_name = Arc::new(lower_case_name(name)?);
 
         if let Some(table) = self.context.table(table_name.clone()) {
             let mut columns = Vec::new();
@@ -35,11 +33,10 @@ impl<'a, T: Transaction> Binder<'a, T> {
                     return Err(DatabaseError::ValuesLenMismatch(columns.len(), values_len));
                 }
             } else {
-                let bind_table_name = Some(table_name.to_string());
                 for ident in idents {
                     match self.bind_column_ref_from_identifiers(
                         slice::from_ref(ident),
-                        bind_table_name.as_ref(),
+                        Some(table_name.to_string()),
                     )? {
                         ScalarExpression::ColumnRef(catalog) => columns.push(catalog),
                         _ => unreachable!(),