fix: auto create datatype_extension missing (#2953)

GreptimeTeam · Dec 19, 2023 · bbcac3a · bbcac3a
1 parent 600cde1
commit bbcac3a
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 6 deletions.
diff --git a/src/common/grpc-expr/src/insert.rs b/src/common/grpc-expr/src/insert.rs
@@ -98,7 +98,11 @@ mod tests {
 
     use api::helper::ColumnDataTypeWrapper;
     use api::v1::column::Values;
-    use api::v1::{Column, ColumnDataType, IntervalMonthDayNano, SemanticType};
+    use api::v1::column_data_type_extension::TypeExt;
+    use api::v1::{
+        Column, ColumnDataType, ColumnDataTypeExtension, Decimal128, DecimalTypeExtension,
+        IntervalMonthDayNano, SemanticType,
+    };
     use common_base::BitVec;
     use common_catalog::consts::MITO_ENGINE;
     use common_time::interval::IntervalUnit;
@@ -160,7 +164,7 @@ mod tests {
 
         let column_defs = create_expr.column_defs;
         assert_eq!(column_defs[6].name, create_expr.time_index);
-        assert_eq!(7, column_defs.len());
+        assert_eq!(8, column_defs.len());
 
         assert_eq!(
             ConcreteDataType::string_datatype(),
@@ -266,6 +270,18 @@ mod tests {
                 .unwrap()
             )
         );
+
+        let decimal_column = column_defs.iter().find(|c| c.name == "decimals").unwrap();
+        assert_eq!(
+            ConcreteDataType::decimal128_datatype(38, 10),
+            ConcreteDataType::from(
+                ColumnDataTypeWrapper::try_new(
+                    decimal_column.data_type,
+                    decimal_column.datatype_extension.clone(),
+                )
+                .unwrap()
+            )
+        );
     }
 
     #[test]
@@ -286,7 +302,7 @@ mod tests {
 
         let add_columns = find_new_columns(&schema, &insert_batch.0).unwrap().unwrap();
 
-        assert_eq!(5, add_columns.add_columns.len());
+        assert_eq!(6, add_columns.add_columns.len());
         let host_column = &add_columns.add_columns[0];
         assert_eq!(
             ConcreteDataType::string_datatype(),
@@ -347,6 +363,23 @@ mod tests {
                 .unwrap()
             )
         );
+
+        let decimal_column = &add_columns.add_columns[5];
+        assert_eq!(
+            ConcreteDataType::decimal128_datatype(38, 10),
+            ConcreteDataType::from(
+                ColumnDataTypeWrapper::try_new(
+                    decimal_column.column_def.as_ref().unwrap().data_type,
+                    decimal_column
+                        .column_def
+                        .as_ref()
+                        .unwrap()
+                        .datatype_extension
+                        .clone()
+                )
+                .unwrap()
+            )
+        );
     }
 
     #[test]
@@ -466,6 +499,23 @@ mod tests {
             datatype: ColumnDataType::TimestampMillisecond as i32,
             ..Default::default()
         };
+        let decimal_vals = Values {
+            decimal128_values: vec![Decimal128 { hi: 0, lo: 123 }, Decimal128 { hi: 0, lo: 456 }],
+            ..Default::default()
+        };
+        let decimal_column = Column {
+            column_name: "decimals".to_string(),
+            semantic_type: SemanticType::Field as i32,
+            values: Some(decimal_vals),
+            null_mask: vec![0],
+            datatype: ColumnDataType::Decimal128 as i32,
+            datatype_extension: Some(ColumnDataTypeExtension {
+                type_ext: Some(TypeExt::DecimalType(DecimalTypeExtension {
+                    precision: 38,
+                    scale: 10,
+                })),
+            }),
+        };
 
         (
             vec![
@@ -476,6 +526,7 @@ mod tests {
                 interval_column,
                 duration_column,
                 ts_column,
+                decimal_column,
             ],
             row_count,
         )

diff --git a/src/common/grpc-expr/src/util.rs b/src/common/grpc-expr/src/util.rs
@@ -15,7 +15,8 @@
 use std::collections::HashSet;
 
 use api::v1::{
-    AddColumn, AddColumns, Column, ColumnDef, ColumnSchema, CreateTableExpr, SemanticType,
+    AddColumn, AddColumns, Column, ColumnDataTypeExtension, ColumnDef, ColumnSchema,
+    CreateTableExpr, SemanticType,
 };
 use datatypes::schema::Schema;
 use snafu::{ensure, OptionExt};
@@ -30,6 +31,7 @@ pub struct ColumnExpr<'a> {
     pub column_name: &'a str,
     pub datatype: i32,
     pub semantic_type: i32,
+    pub datatype_extension: &'a Option<ColumnDataTypeExtension>,
 }
 
 impl<'a> ColumnExpr<'a> {
@@ -50,6 +52,7 @@ impl<'a> From<&'a Column> for ColumnExpr<'a> {
             column_name: &column.column_name,
             datatype: column.datatype,
             semantic_type: column.semantic_type,
+            datatype_extension: &column.datatype_extension,
         }
     }
 }
@@ -60,6 +63,7 @@ impl<'a> From<&'a ColumnSchema> for ColumnExpr<'a> {
             column_name: &schema.column_name,
             datatype: schema.datatype,
             semantic_type: schema.semantic_type,
+            datatype_extension: &schema.datatype_extension,
         }
     }
 }
@@ -94,6 +98,7 @@ pub fn build_create_table_expr(
         column_name,
         datatype,
         semantic_type,
+        datatype_extension,
     } in column_exprs
     {
         let mut is_nullable = true;
@@ -121,7 +126,7 @@ pub fn build_create_table_expr(
             default_constraint: vec![],
             semantic_type,
             comment: String::new(),
-            ..Default::default()
+            datatype_extension: datatype_extension.clone(),
         };
         column_defs.push(column_def);
     }
@@ -162,7 +167,7 @@ pub fn extract_new_columns(
                 default_constraint: vec![],
                 semantic_type: expr.semantic_type,
                 comment: String::new(),
-                ..Default::default()
+                datatype_extension: expr.datatype_extension.clone(),
             });
             AddColumn {
                 column_def,