feat: support drop column (#109)

* feat: add examples

* code fmt

* version up

* feat: add_column support unique index

* feat: support drop column

---------

Co-authored-by: Kould <[email protected]>

Cargo.toml

@@ -2,7 +2,7 @@
 
 [package]
 name = "kip-sql"
-version = "0.0.1-alpha.5"
+version = "0.0.1-alpha.6"
 edition = "2021"
 authors = ["Kould <[email protected]>", "Xwg <[email protected]>"]
 description = "build the SQL layer of KipDB database"

src/binder/alter_table.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 use super::Binder;
 use crate::binder::{lower_case_name, split_name, BindError};
 use crate::planner::operator::alter_table::add_column::AddColumnOperator;
+use crate::planner::operator::alter_table::drop_column::DropColumnOperator;
 use crate::planner::operator::scan::ScanOperator;
 use crate::planner::operator::Operator;
 use crate::planner::LogicalPlan;
@@ -18,13 +19,13 @@ impl<'a, T: Transaction> Binder<'a, T> {
     ) -> Result<LogicalPlan, BindError> {
         let table_name: Arc<String> = Arc::new(split_name(&lower_case_name(name))?.1.to_string());
 
-        let plan = match operation {
-            AlterTableOperation::AddColumn {
-                column_keyword: _,
-                if_not_exists,
-                column_def,
-            } => {
-                if let Some(table) = self.context.table(table_name.clone()) {
+        if let Some(table) = self.context.table(table_name.clone()) {
+            let plan = match operation {
+                AlterTableOperation::AddColumn {
+                    column_keyword: _,
+                    if_not_exists,
+                    column_def,
+                } => {
                     let plan = ScanOperator::build(table_name.clone(), table);
 
                     LogicalPlan {
@@ -35,37 +36,49 @@ impl<'a, T: Transaction> Binder<'a, T> {
                         }),
                         childrens: vec![plan],
                     }
-                } else {
-                    return Err(BindError::InvalidTable(format!(
-                        "not found table {}",
-                        table_name
-                    )));
                 }
-            }
-            AlterTableOperation::DropColumn {
-                column_name: _,
-                if_exists: _,
-                cascade: _,
-            } => todo!(),
-            AlterTableOperation::DropPrimaryKey => todo!(),
-            AlterTableOperation::RenameColumn {
-                old_column_name: _,
-                new_column_name: _,
-            } => todo!(),
-            AlterTableOperation::RenameTable { table_name: _ } => todo!(),
-            AlterTableOperation::ChangeColumn {
-                old_name: _,
-                new_name: _,
-                data_type: _,
-                options: _,
-            } => todo!(),
-            AlterTableOperation::AlterColumn {
-                column_name: _,
-                op: _,
-            } => todo!(),
-            _ => todo!(),
-        };
+                AlterTableOperation::DropColumn {
+                    column_name,
+                    if_exists,
+                    ..
+                } => {
+                    let plan = ScanOperator::build(table_name.clone(), table);
+                    let column_name = column_name.value.clone();
+
+                    LogicalPlan {
+                        operator: Operator::DropColumn(DropColumnOperator {
+                            table_name,
+                            if_exists: *if_exists,
+                            column_name,
+                        }),
+                        childrens: vec![plan],
+                    }
+                }
+                AlterTableOperation::DropPrimaryKey => todo!(),
+                AlterTableOperation::RenameColumn {
+                    old_column_name: _,
+                    new_column_name: _,
+                } => todo!(),
+                AlterTableOperation::RenameTable { table_name: _ } => todo!(),
+                AlterTableOperation::ChangeColumn {
+                    old_name: _,
+                    new_name: _,
+                    data_type: _,
+                    options: _,
+                } => todo!(),
+                AlterTableOperation::AlterColumn {
+                    column_name: _,
+                    op: _,
+                } => todo!(),
+                _ => todo!(),
+            };
 
-        Ok(plan)
+            Ok(plan)
+        } else {
+            Err(BindError::InvalidTable(format!(
+                "not found table {}",
+                table_name
+            )))
+        }
     }
 }

src/catalog/table.rs

@@ -29,11 +29,11 @@ impl TableCatalog {
     }
 
     #[allow(dead_code)]
-    pub(crate) fn get_column_id_by_name(&self, name: &String) -> Option<ColumnId> {
+    pub(crate) fn get_column_id_by_name(&self, name: &str) -> Option<ColumnId> {
         self.column_idxs.get(name).cloned()
     }
 
-    pub(crate) fn get_column_by_name(&self, name: &String) -> Option<&ColumnRef> {
+    pub(crate) fn get_column_by_name(&self, name: &str) -> Option<&ColumnRef> {
         let id = self.column_idxs.get(name)?;
         self.columns.get(id)
     }
@@ -65,10 +65,22 @@ impl TableCatalog {
         Ok(col_id)
     }
 
-    pub(crate) fn add_index_meta(&mut self, mut index: IndexMeta) -> &IndexMeta {
+    pub(crate) fn add_index_meta(
+        &mut self,
+        name: String,
+        column_ids: Vec<ColumnId>,
+        is_unique: bool,
+        is_primary: bool,
+    ) -> &IndexMeta {
         let index_id = self.indexes.len();
 
-        index.id = index_id as u32;
+        let index = IndexMeta {
+            id: index_id as u32,
+            column_ids,
+            name,
+            is_unique,
+            is_primary,
+        };
         self.indexes.push(Arc::new(index));
 
         &self.indexes[index_id]

src/execution/executor/ddl/alter_table/add_column.rs

@@ -6,6 +6,7 @@ use futures_async_stream::try_stream;
 use std::cell::RefCell;
 use std::sync::Arc;
 
+use crate::types::index::Index;
 use crate::{
     execution::executor::Executor, planner::operator::alter_table::add_column::AddColumnOperator,
     storage::Transaction,
@@ -23,35 +24,54 @@ impl From<(AddColumnOperator, BoxedExecutor)> for AddColumn {
 }
 
 impl<T: Transaction> Executor<T> for AddColumn {
-    fn execute(self, transaction: &RefCell<T>) -> crate::execution::executor::BoxedExecutor {
+    fn execute(self, transaction: &RefCell<T>) -> BoxedExecutor {
         unsafe { self._execute(transaction.as_ptr().as_mut().unwrap()) }
     }
 }
 
 impl AddColumn {
     #[try_stream(boxed, ok = Tuple, error = ExecutorError)]
     async fn _execute<T: Transaction>(self, transaction: &mut T) {
-        let _ = transaction.add_column(&self.op)?;
-
         let AddColumnOperator {
-            table_name, column, ..
+            table_name,
+            column,
+            if_not_exists,
         } = &self.op;
+        let mut unique_values = column.desc().is_unique.then(|| Vec::new());
 
         #[for_await]
         for tuple in self.input {
             let mut tuple: Tuple = tuple?;
-            let is_overwrite = true;
 
             tuple.columns.push(Arc::new(column.clone()));
             if let Some(value) = column.default_value() {
+                if let Some(unique_values) = &mut unique_values {
+                    unique_values.push((tuple.id.clone().unwrap(), value.clone()));
+                }
                 tuple.values.push(value);
             } else {
                 tuple.values.push(Arc::new(DataValue::Null));
             }
+            transaction.append(table_name, tuple, true)?;
+        }
+        let col_id = transaction.add_column(table_name, column, *if_not_exists)?;
 
-            transaction.append(table_name, tuple, is_overwrite)?;
+        // Unique Index
+        if let (Some(unique_values), Some(unique_meta)) = (
+            unique_values,
+            transaction
+                .table(table_name.clone())
+                .and_then(|table| table.get_unique_index(&col_id))
+                .cloned(),
+        ) {
+            for (tuple_id, value) in unique_values {
+                let index = Index {
+                    id: unique_meta.id,
+                    column_values: vec![value],
+                };
+                transaction.add_index(&table_name, index, vec![tuple_id], true)?;
+            }
         }
-        transaction.remove_cache(&table_name)?;
 
         let tuple_builder = TupleBuilder::new_result();
         let tuple = tuple_builder.push_result("ALTER TABLE SUCCESS", "1")?;

src/execution/executor/ddl/alter_table/drop_column.rs

@@ -0,0 +1,76 @@
+use crate::binder::BindError;
+use crate::execution::executor::{BoxedExecutor, Executor};
+use crate::execution::ExecutorError;
+use crate::planner::operator::alter_table::drop_column::DropColumnOperator;
+use crate::storage::Transaction;
+use crate::types::tuple::Tuple;
+use crate::types::tuple_builder::TupleBuilder;
+use futures_async_stream::try_stream;
+use std::cell::RefCell;
+
+pub struct DropColumn {
+    op: DropColumnOperator,
+    input: BoxedExecutor,
+}
+
+impl From<(DropColumnOperator, BoxedExecutor)> for DropColumn {
+    fn from((op, input): (DropColumnOperator, BoxedExecutor)) -> Self {
+        Self { op, input }
+    }
+}
+
+impl<T: Transaction> Executor<T> for DropColumn {
+    fn execute(self, transaction: &RefCell<T>) -> BoxedExecutor {
+        unsafe { self._execute(transaction.as_ptr().as_mut().unwrap()) }
+    }
+}
+
+impl DropColumn {
+    #[try_stream(boxed, ok = Tuple, error = ExecutorError)]
+    async fn _execute<T: Transaction>(self, transaction: &mut T) {
+        let DropColumnOperator {
+            table_name,
+            column_name,
+            if_exists,
+        } = &self.op;
+        let mut option_column_index = None;
+
+        #[for_await]
+        for tuple in self.input {
+            let mut tuple: Tuple = tuple?;
+
+            if option_column_index.is_none() {
+                if let Some((column_index, is_primary)) = tuple
+                    .columns
+                    .iter()
+                    .enumerate()
+                    .find(|(_, column)| column.name() == column_name)
+                    .map(|(i, column)| (i, column.desc.is_primary))
+                {
+                    if is_primary {
+                        Err(BindError::InvalidColumn(
+                            "drop of primary key column is not allowed.".to_owned(),
+                        ))?;
+                    }
+                    option_column_index = Some(column_index);
+                }
+            }
+            if option_column_index.is_none() && *if_exists {
+                return Ok(());
+            }
+            let column_index = option_column_index
+                .ok_or_else(|| BindError::InvalidColumn("not found column".to_string()))?;
+
+            let _ = tuple.columns.remove(column_index);
+            let _ = tuple.values.remove(column_index);
+
+            transaction.append(table_name, tuple, true)?;
+        }
+        transaction.drop_column(table_name, column_name, *if_exists)?;
+
+        let tuple_builder = TupleBuilder::new_result();
+        let tuple = tuple_builder.push_result("ALTER TABLE SUCCESS", "1")?;
+
+        yield tuple;
+    }
+}

src/execution/executor/ddl/alter_table/mod.rs

@@ -1 +1,2 @@
 pub mod add_column;
+pub mod drop_column;

src/execution/executor/mod.rs

@@ -3,6 +3,7 @@ pub(crate) mod dml;
 pub(crate) mod dql;
 pub(crate) mod show;
 
+use crate::execution::executor::ddl::alter_table::drop_column::DropColumn;
 use crate::execution::executor::ddl::create_table::CreateTable;
 use crate::execution::executor::ddl::drop_table::DropTable;
 use crate::execution::executor::ddl::truncate::Truncate;
@@ -110,6 +111,10 @@ pub fn build<T: Transaction>(plan: LogicalPlan, transaction: &RefCell<T>) -> Box
             let input = build(childrens.remove(0), transaction);
             AddColumn::from((op, input)).execute(transaction)
         }
+        Operator::DropColumn(op) => {
+            let input = build(childrens.remove(0), transaction);
+            DropColumn::from((op, input)).execute(transaction)
+        }
         Operator::CreateTable(op) => CreateTable::from(op).execute(transaction),
         Operator::DropTable(op) => DropTable::from(op).execute(transaction),
         Operator::Truncate(op) => Truncate::from(op).execute(transaction),

src/optimizer/rule/column_pruning.rs

@@ -116,7 +116,8 @@ impl ColumnPruning {
             | Operator::Show(_)
             | Operator::CopyFromFile(_)
             | Operator::CopyToFile(_)
-            | Operator::AddColumn(_) => (),
+            | Operator::AddColumn(_)
+            | Operator::DropColumn(_) => (),
         }
     }
 

src/planner/operator/alter_table/drop_column.rs

@@ -0,0 +1,8 @@
+use crate::catalog::TableName;
+
+#[derive(Debug, PartialEq, Clone)]
+pub struct DropColumnOperator {
+    pub table_name: TableName,
+    pub column_name: String,
+    pub if_exists: bool,
+}

src/planner/operator/alter_table/mod.rs

@@ -1 +1,2 @@
 pub mod add_column;
+pub mod drop_column;

src/planner/operator/mod.rs

@@ -18,6 +18,7 @@ pub mod update;
 pub mod values;
 
 use crate::catalog::ColumnRef;
+use crate::planner::operator::alter_table::drop_column::DropColumnOperator;
 use crate::planner::operator::copy_from_file::CopyFromFileOperator;
 use crate::planner::operator::copy_to_file::CopyToFileOperator;
 use crate::planner::operator::create_table::CreateTableOperator;
@@ -55,6 +56,7 @@ pub enum Operator {
     Delete(DeleteOperator),
     // DDL
     AddColumn(AddColumnOperator),
+    DropColumn(DropColumnOperator),
     CreateTable(CreateTableOperator),
     DropTable(DropTableOperator),
     Truncate(TruncateOperator),