Skip to content

Commit

Permalink
feat: add ChangeColumnType for AlterKind (#3757)
Browse files Browse the repository at this point in the history 
* feat: add `ModifyColumn` for `AlterKind`

* chore: additional code comments for `AlterKind::ModifyColumns`

* fix: add nullable check on `ModifyColumn`

* style: codefmt

* style: fix the code based on review suggestions

* style: fix the code based on review suggestions

* style: rename `ModifyColumn` -> `ChangeColumnType`

* style: code fmt

* style: `change_columns_type` -> `change_column_types`
  • Loading branch information
@KKould
KKould authored Apr 24, 2024
1 parent 4685b59 commit b619950
Show file tree
Hide file tree
Showing 5 changed files with 404 additions and 23 deletions.
2 changes: 1 addition & 1 deletion src/common/meta/src/ddl/alter_table/update_metadata.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ impl AlterTableProcedure {
AlterKind::RenameTable { new_table_name } => {
new_info.name = new_table_name.to_string();
}
AlterKind::DropColumns { .. } => {}
AlterKind::DropColumns { .. } | AlterKind::ChangeColumnTypes { .. } => {}
}

Ok(new_info)
Expand Down
81 changes: 64 additions & 17 deletions src/store-api/src/metadata.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use std::sync::Arc;

use api::helper::ColumnDataTypeWrapper;
use api::v1::region::RegionColumnDef;
use api::v1::SemanticType;
use api::v1::{ColumnDef, SemanticType};
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use common_macro::stack_trace_debug;
Expand All @@ -33,7 +33,7 @@ use serde::de::Error;
use serde::{Deserialize, Deserializer, Serialize};
use snafu::{ensure, Location, OptionExt, ResultExt, Snafu};

use crate::region_request::{AddColumn, AddColumnLocation, AlterKind};
use crate::region_request::{AddColumn, AddColumnLocation, AlterKind, ChangeColumnType};
use crate::storage::consts::is_internal_column;
use crate::storage::{ColumnId, RegionId};

Expand Down Expand Up @@ -61,18 +61,7 @@ impl fmt::Debug for ColumnMetadata {
}

impl ColumnMetadata {
/// Construct `Self` from protobuf struct [RegionColumnDef]
pub fn try_from_column_def(column_def: RegionColumnDef) -> Result<Self> {
let column_id = column_def.column_id;

let column_def = column_def
.column_def
.context(InvalidRawRegionRequestSnafu {
err: "column_def is absent",
})?;

let semantic_type = column_def.semantic_type();

fn inner_try_from_column_def(column_def: ColumnDef) -> Result<ColumnSchema> {
let default_constrain = if column_def.default_constraint.is_empty() {
None
} else {
Expand All @@ -86,9 +75,22 @@ impl ColumnMetadata {
column_def.datatype_extension.clone(),
)
.into();
let column_schema = ColumnSchema::new(column_def.name, data_type, column_def.is_nullable)
ColumnSchema::new(column_def.name, data_type, column_def.is_nullable)
.with_default_constraint(default_constrain)
.context(ConvertDatatypesSnafu)?;
.context(ConvertDatatypesSnafu)
}

/// Construct `Self` from protobuf struct [RegionColumnDef]
pub fn try_from_column_def(column_def: RegionColumnDef) -> Result<Self> {
let column_id = column_def.column_id;
let column_def = column_def
.column_def
.context(InvalidRawRegionRequestSnafu {
err: "column_def is absent",
})?;
let semantic_type = column_def.semantic_type();
let column_schema = Self::inner_try_from_column_def(column_def)?;

Ok(Self {
column_schema,
semantic_type,
Expand Down Expand Up @@ -535,6 +537,7 @@ impl RegionMetadataBuilder {
match kind {
AlterKind::AddColumns { columns } => self.add_columns(columns)?,
AlterKind::DropColumns { names } => self.drop_columns(&names),
AlterKind::ChangeColumnTypes { columns } => self.change_column_types(columns),
}
Ok(self)
}
Expand Down Expand Up @@ -615,6 +618,25 @@ impl RegionMetadataBuilder {
self.column_metadatas
.retain(|col| !name_set.contains(&col.column_schema.name));
}

/// Changes columns type to the metadata if exist.
fn change_column_types(&mut self, columns: Vec<ChangeColumnType>) {
let mut change_type_map: HashMap<_, _> = columns
.into_iter()
.map(
|ChangeColumnType {
column_name,
target_type,
}| (column_name, target_type),
)
.collect();

for column_meta in self.column_metadatas.iter_mut() {
if let Some(target_type) = change_type_map.remove(&column_meta.column_schema.name) {
column_meta.column_schema.data_type = target_type;
}
}
}
}

/// Fields skipped in serialization.
Expand Down Expand Up @@ -707,6 +729,13 @@ pub enum MetadataError {

#[snafu(display("Time index column not found"))]
TimeIndexNotFound { location: Location },

#[snafu(display("Change column {} not exists in region: {}", column_name, region_id))]
ChangeColumnNotFound {
column_name: String,
region_id: RegionId,
location: Location,
},
}

impl ErrorExt for MetadataError {
Expand Down Expand Up @@ -1112,7 +1141,7 @@ mod test {
let metadata = builder.build().unwrap();
check_columns(&metadata, &["a", "b", "f", "c", "d"]);

let mut builder = RegionMetadataBuilder::from_existing(metadata);
let mut builder = RegionMetadataBuilder::from_existing(metadata.clone());
builder
.alter(AlterKind::DropColumns {
names: vec!["a".to_string()],
Expand All @@ -1121,6 +1150,24 @@ mod test {
// Build returns error as the primary key contains a.
let err = builder.build().unwrap_err();
assert_eq!(StatusCode::InvalidArguments, err.status_code());

let mut builder = RegionMetadataBuilder::from_existing(metadata);
builder
.alter(AlterKind::ChangeColumnTypes {
columns: vec![ChangeColumnType {
column_name: "b".to_string(),
target_type: ConcreteDataType::string_datatype(),
}],
})
.unwrap();
let metadata = builder.build().unwrap();
check_columns(&metadata, &["a", "b", "f", "c", "d"]);
let b_type = &metadata
.column_by_name("b")
.unwrap()
.column_schema
.data_type;
assert_eq!(ConcreteDataType::string_datatype(), *b_type);
}

#[test]
Expand Down
122 changes: 122 additions & 0 deletions src/store-api/src/region_request.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use api::v1::region::{
};
use api::v1::{self, Rows, SemanticType};
pub use common_base::AffectedRows;
use datatypes::data_type::ConcreteDataType;
use snafu::{ensure, OptionExt};
use strum::IntoStaticStr;

Expand Down Expand Up @@ -332,6 +333,11 @@ pub enum AlterKind {
/// Name of columns to drop.
names: Vec<String>,
},
/// Change columns datatype form the region, only fields are allowed to change.
ChangeColumnTypes {
/// Columns to change.
columns: Vec<ChangeColumnType>,
},
}

impl AlterKind {
Expand All @@ -350,6 +356,11 @@ impl AlterKind {
Self::validate_column_to_drop(name, metadata)?;
}
}
AlterKind::ChangeColumnTypes { columns } => {
for col_to_change in columns {
col_to_change.validate(metadata)?;
}
}
}
Ok(())
}
Expand All @@ -364,6 +375,9 @@ impl AlterKind {
AlterKind::DropColumns { names } => names
.iter()
.any(|name| metadata.column_by_name(name).is_some()),
AlterKind::ChangeColumnTypes { columns } => columns
.iter()
.any(|col_to_change| col_to_change.need_alter(metadata)),
}
}

Expand Down Expand Up @@ -501,6 +515,56 @@ impl TryFrom<v1::AddColumnLocation> for AddColumnLocation {
}
}

/// Change a column's datatype.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct ChangeColumnType {
/// Schema of the column to modify.
pub column_name: String,
/// Column will be changed to this type.
pub target_type: ConcreteDataType,
}

impl ChangeColumnType {
/// Returns an error if the column's datatype to change is invalid.
pub fn validate(&self, metadata: &RegionMetadata) -> Result<()> {
let column_meta = metadata
.column_by_name(&self.column_name)
.with_context(|| InvalidRegionRequestSnafu {
region_id: metadata.region_id,
err: format!("column {} not found", self.column_name),
})?;

ensure!(
matches!(column_meta.semantic_type, SemanticType::Field),
InvalidRegionRequestSnafu {
region_id: metadata.region_id,
err: "'timestamp' or 'tag' column cannot change type".to_string()
}
);
ensure!(
column_meta
.column_schema
.data_type
.can_arrow_type_cast_to(&self.target_type),
InvalidRegionRequestSnafu {
region_id: metadata.region_id,
err: format!(
"column '{}' cannot be cast automatically to type '{}'",
self.column_name, self.target_type
),
}
);

Ok(())
}

/// Returns true if no column's datatype to change to the region.
pub fn need_alter(&self, metadata: &RegionMetadata) -> bool {
debug_assert!(self.validate(metadata).is_ok());
metadata.column_by_name(&self.column_name).is_some()
}
}

#[derive(Debug, Default)]
pub struct RegionFlushRequest {
pub row_group_size: Option<usize>,
Expand Down Expand Up @@ -678,6 +742,15 @@ mod tests {
semantic_type: SemanticType::Field,
column_id: 3,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"field_1",
ConcreteDataType::boolean_datatype(),
true,
),
semantic_type: SemanticType::Field,
column_id: 4,
})
.primary_key(vec![2]);
builder.build().unwrap()
}
Expand Down Expand Up @@ -790,6 +863,55 @@ mod tests {
assert!(kind.need_alter(&metadata));
}

#[test]
fn test_validate_change_column_type() {
let metadata = new_metadata();
AlterKind::ChangeColumnTypes {
columns: vec![ChangeColumnType {
column_name: "xxxx".to_string(),
target_type: ConcreteDataType::string_datatype(),
}],
}
.validate(&metadata)
.unwrap_err();

AlterKind::ChangeColumnTypes {
columns: vec![ChangeColumnType {
column_name: "field_1".to_string(),
target_type: ConcreteDataType::date_datatype(),
}],
}
.validate(&metadata)
.unwrap_err();

AlterKind::ChangeColumnTypes {
columns: vec![ChangeColumnType {
column_name: "ts".to_string(),
target_type: ConcreteDataType::date_datatype(),
}],
}
.validate(&metadata)
.unwrap_err();

AlterKind::ChangeColumnTypes {
columns: vec![ChangeColumnType {
column_name: "tag_0".to_string(),
target_type: ConcreteDataType::date_datatype(),
}],
}
.validate(&metadata)
.unwrap_err();

let kind = AlterKind::ChangeColumnTypes {
columns: vec![ChangeColumnType {
column_name: "field_0".to_string(),
target_type: ConcreteDataType::int32_datatype(),
}],
};
kind.validate(&metadata).unwrap();
assert!(kind.need_alter(&metadata));
}

#[test]
fn test_validate_schema_version() {
let mut metadata = new_metadata();
Expand Down
Loading

0 comments on commit b619950

Please sign in to comment.