chore: rewrite checked referenced computed columns (#12152)

* chore: rewrite checked referenced computed columns

* add check data type

* fix

src/query/expression/src/schema.rs

@@ -224,6 +224,22 @@ impl DataSchema {
             .find(|&(_, c)| c.name() == name)
     }
 
+    pub fn rename_field(&mut self, i: FieldIndex, new_name: &str) {
+        self.fields[i].name = new_name.to_string();
+    }
+
+    pub fn drop_column(&mut self, column: &str) -> Result<()> {
+        if self.fields.len() == 1 {
+            return Err(ErrorCode::DropColumnEmptyError(
+                "cannot drop table column to empty",
+            ));
+        }
+        let i = self.index_of(column)?;
+        self.fields.remove(i);
+
+        Ok(())
+    }
+
     /// Check to see if `self` is a superset of `other` schema. Here are the comparison rules:
     pub fn contains(&self, other: &DataSchema) -> bool {
         if self.fields.len() != other.fields.len() {

src/query/service/src/interpreters/common/table.rs

@@ -18,35 +18,38 @@ use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::ComputedExpr;
-use common_expression::DataSchemaRefExt;
-use common_expression::TableSchemaRef;
+use common_expression::DataSchemaRef;
 use common_sql::parse_computed_expr;
 
 pub fn check_referenced_computed_columns(
     ctx: Arc<dyn TableContext>,
-    table_schema: TableSchemaRef,
+    schema: DataSchemaRef,
     column: &str,
 ) -> Result<()> {
-    let fields = table_schema
-        .fields()
-        .iter()
-        .filter(|f| f.name != column)
-        .map(|f| f.into())
-        .collect::<Vec<_>>();
-    let schema = DataSchemaRefExt::create(fields);
-
     for f in schema.fields() {
         if let Some(computed_expr) = f.computed_expr() {
             let expr = match computed_expr {
-                ComputedExpr::Stored(expr) => expr.clone(),
-                ComputedExpr::Virtual(expr) => expr.clone(),
+                ComputedExpr::Stored(ref expr) => expr,
+                ComputedExpr::Virtual(ref expr) => expr,
             };
-            if parse_computed_expr(ctx.clone(), schema.clone(), &expr).is_err() {
-                return Err(ErrorCode::ColumnReferencedByComputedColumn(format!(
-                    "column `{}` is referenced by computed column `{}`",
-                    column,
-                    &f.name()
-                )));
+            match parse_computed_expr(ctx.clone(), schema.clone(), expr) {
+                Ok(expr) => {
+                    if expr.data_type() != f.data_type() {
+                        return Err(ErrorCode::ColumnReferencedByComputedColumn(format!(
+                            "expected computed column expression have type {}, but `{}` has type {}.",
+                            f.data_type(),
+                            column,
+                            expr.data_type(),
+                        )));
+                    }
+                }
+                Err(_) => {
+                    return Err(ErrorCode::ColumnReferencedByComputedColumn(format!(
+                        "column `{}` is referenced by computed column `{}`",
+                        column,
+                        &f.name()
+                    )));
+                }
             }
         }
     }

src/query/service/src/interpreters/interpreter_table_drop_column.rs

@@ -16,6 +16,7 @@ use std::sync::Arc;
 
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_expression::DataSchema;
 use common_meta_app::schema::DatabaseType;
 use common_meta_app::schema::UpdateTableMetaReq;
 use common_meta_types::MatchSeq;
@@ -73,13 +74,14 @@ impl Interpreter for DropTableColumnInterpreter {
             )));
         }
 
-        let table_schema = table_info.schema();
-        let field = table_schema.field_with_name(self.plan.column.as_str())?;
-        if field.computed_expr.is_none() {
+        let mut schema: DataSchema = table_info.schema().into();
+        let field = schema.field_with_name(self.plan.column.as_str())?;
+        if field.computed_expr().is_none() {
+            schema.drop_column(self.plan.column.as_str())?;
             // Check if this column is referenced by computed columns.
             check_referenced_computed_columns(
                 self.ctx.clone(),
-                table_schema,
+                Arc::new(schema),
                 self.plan.column.as_str(),
             )?;
         }

src/query/service/src/interpreters/interpreter_table_rename_column.rs

@@ -16,6 +16,7 @@ use std::sync::Arc;
 
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_expression::DataSchema;
 use common_meta_app::schema::DatabaseType;
 use common_meta_app::schema::UpdateTableMetaReq;
 use common_meta_types::MatchSeq;
@@ -82,13 +83,15 @@ impl Interpreter for RenameTableColumnInterpreter {
 
             is_valid_column(&self.plan.new_column)?;
 
-            let table_schema = table_info.schema();
-            let field = table_schema.field_with_name(self.plan.old_column.as_str())?;
-            if field.computed_expr.is_none() {
+            let mut schema: DataSchema = table_info.schema().into();
+            let field = schema.field_with_name(self.plan.old_column.as_str())?;
+            if field.computed_expr().is_none() {
+                let index = schema.index_of(self.plan.old_column.as_str())?;
+                schema.rename_field(index, self.plan.new_column.as_str());
                 // Check if old column is referenced by computed columns.
                 check_referenced_computed_columns(
                     self.ctx.clone(),
-                    table_schema,
+                    Arc::new(schema),
                     self.plan.old_column.as_str(),
                 )?;
             }