Merge pull request #10369 from youngsofun/dec

feat: support function to_decimal(string).
databendlabs · Mar 7, 2023 · 17bb67c · 17bb67c · vercel · Mar 7, 2023
2 parents 8238c0e + 847c858
commit 17bb67c
Show file tree

Hide file tree

Showing 9 changed files with 395 additions and 37 deletions.
diff --git a/src/query/expression/src/type_check.rs b/src/query/expression/src/type_check.rs
@@ -517,6 +517,7 @@ pub fn can_auto_cast_to(
                 .all(|(src_ty, dest_ty)| can_auto_cast_to(src_ty, dest_ty, auto_cast_rules))
         }
         (DataType::Number(n), DataType::Decimal(_)) if !n.is_float() => true,
+        (DataType::String, DataType::Decimal(_)) => true,
         (DataType::Decimal(x), DataType::Decimal(y)) => x.precision() <= y.precision(),
         (DataType::Decimal(_), DataType::Number(NumberDataType::Float32)) => true,
         (DataType::Decimal(_), DataType::Number(NumberDataType::Float64)) => true,

diff --git a/src/query/expression/src/types/decimal.rs b/src/query/expression/src/types/decimal.rs
@@ -291,6 +291,8 @@ pub trait Decimal:
         Self::to_column_from_buffer(value.into(), size)
     }
 
+    fn to_scalar(self, size: DecimalSize) -> DecimalScalar;
+
     fn with_size(&self, size: DecimalSize) -> Option<Self> {
         let multiplier = Self::e(size.scale as u32);
         let min_for_precision = Self::min_for_precision(size.precision);
@@ -392,6 +394,10 @@ impl Decimal for i128 {
         self as f64 / div
     }
 
+    fn to_scalar(self, size: DecimalSize) -> DecimalScalar {
+        DecimalScalar::Decimal128(self, size)
+    }
+
     fn try_downcast_column(column: &Column) -> Option<(Buffer<Self>, DecimalSize)> {
         let column = column.as_decimal()?;
         match column {
@@ -529,6 +535,10 @@ impl Decimal for i256 {
         self.as_f64() / div
     }
 
+    fn to_scalar(self, size: DecimalSize) -> DecimalScalar {
+        DecimalScalar::Decimal256(self, size)
+    }
+
     fn try_downcast_column(column: &Column) -> Option<(Buffer<Self>, DecimalSize)> {
         let column = column.as_decimal()?;
         match column {

diff --git a/src/query/functions/src/scalars/decimal.rs b/src/query/functions/src/scalars/decimal.rs
@@ -17,7 +17,9 @@ use std::ops::*;
 use std::sync::Arc;
 
 use common_arrow::arrow::buffer::Buffer;
+use common_expression::read_decimal_with_size;
 use common_expression::types::decimal::*;
+use common_expression::types::string::StringColumn;
 use common_expression::types::*;
 use common_expression::with_integer_mapped_type;
 use common_expression::Column;
@@ -381,10 +383,12 @@ pub fn register(registry: &mut FunctionRegistry) {
         if args_type.len() != 1 {
             return None;
         }
-        if !args_type[0].is_decimal() && !args_type[0].is_numeric() {
+        if !matches!(
+            args_type[0],
+            DataType::Number(_) | DataType::Decimal(_) | DataType::String
+        ) {
             return None;
         }
-
         if params.len() != 2 {
             return None;
         }
@@ -468,13 +472,86 @@ fn convert_to_decimal(
     dest_type: DataType,
 ) -> Value<AnyType> {
     let arg = &args[0];
-    if from_type.is_integer() {
-        return integer_to_decimal(arg, ctx, from_type, dest_type);
+    match from_type {
+        DataType::Number(ty) => {
+            if ty.is_float() {
+                float_to_decimal(arg, ctx, from_type, dest_type)
+            } else {
+                integer_to_decimal(arg, ctx, from_type, dest_type)
+            }
+        }
+        DataType::Decimal(_) => decimal_to_decimal(arg, ctx, from_type, dest_type),
+        DataType::String => string_to_decimal(arg, ctx, dest_type),
+        _ => unreachable!("to_decimal not support this DataType"),
+    }
+}
+
+fn string_to_decimal_column<T: Decimal>(
+    ctx: &mut EvalContext,
+    string_column: &StringColumn,
+    size: DecimalSize,
+) -> DecimalColumn {
+    let mut values = Vec::<T>::with_capacity(string_column.len());
+    for (row, buf) in string_column.iter().enumerate() {
+        match read_decimal_with_size::<T>(buf, size, true) {
+            Ok((d, _)) => values.push(d),
+            Err(e) => {
+                ctx.set_error(row, format!("fail to decode string as decimal: {e}"));
+                values.push(T::zero())
+            }
+        }
     }
-    if from_type.is_floating() {
-        return float_to_decimal(arg, ctx, from_type, dest_type);
+    T::to_column(values, size)
+}
+
+fn string_to_decimal_scalar<T: Decimal>(
+    ctx: &mut EvalContext,
+    string_buf: &[u8],
+    size: DecimalSize,
+) -> DecimalScalar {
+    let value = match read_decimal_with_size::<T>(string_buf, size, true) {
+        Ok((d, _)) => d,
+        Err(e) => {
+            ctx.set_error(0, format!("fail to decode string as decimal: {e}"));
+            T::zero()
+        }
+    };
+    T::to_scalar(value, size)
+}
+
+fn string_to_decimal(
+    arg: &ValueRef<AnyType>,
+    ctx: &mut EvalContext,
+    dest_type: DataType,
+) -> Value<AnyType> {
+    let dest_type = dest_type.as_decimal().unwrap();
+
+    match arg {
+        ValueRef::Column(column) => {
+            let string_column = StringType::try_downcast_column(column).unwrap();
+            let column = match dest_type {
+                DecimalDataType::Decimal128(size) => {
+                    string_to_decimal_column::<i128>(ctx, &string_column, *size)
+                }
+                DecimalDataType::Decimal256(size) => {
+                    string_to_decimal_column::<i256>(ctx, &string_column, *size)
+                }
+            };
+            Value::Column(Column::Decimal(column))
+        }
+        ValueRef::Scalar(scalar) => {
+            let buf = StringType::try_downcast_scalar(scalar).unwrap();
+            let scalar = match dest_type {
+                DecimalDataType::Decimal128(size) => {
+                    string_to_decimal_scalar::<i128>(ctx, buf, *size)
+                }
+                DecimalDataType::Decimal256(size) => {
+                    string_to_decimal_scalar::<i128>(ctx, buf, *size)
+                }
+            };
+            Value::Scalar(Scalar::Decimal(scalar))
+        }
     }
-    decimal_to_decimal(arg, ctx, from_type, dest_type)
 }
 
 fn integer_to_decimal(

diff --git a/src/query/functions/src/scalars/mod.rs b/src/query/functions/src/scalars/mod.rs
@@ -266,6 +266,8 @@ pub const CAST_FROM_STRING_RULES: AutoCastRules = &[
     (DataType::String, DataType::Number(NumberDataType::Int16)),
     (DataType::String, DataType::Number(NumberDataType::Int32)),
     (DataType::String, DataType::Number(NumberDataType::Int64)),
+    (DataType::String, DataType::Number(NumberDataType::Float32)),
+    (DataType::String, DataType::Number(NumberDataType::Float64)),
 ];
 
 #[allow(non_snake_case)]

diff --git a/src/query/functions/tests/it/scalars/cast.rs b/src/query/functions/tests/it/scalars/cast.rs
@@ -32,6 +32,7 @@ fn test_cast() {
         test_cast_number_to_date(file, is_try);
         test_cast_between_number_and_string(file, is_try);
         test_cast_between_boolean_and_string(file, is_try);
+        test_cast_between_string_and_decimal(file, is_try);
         test_cast_between_number_and_boolean(file, is_try);
         test_cast_between_date_and_timestamp(file, is_try);
         test_cast_between_string_and_timestamp(file, is_try);
@@ -622,3 +623,36 @@ fn test_cast_to_nested_type(file: &mut impl Write, is_try: bool) {
         &[],
     );
 }
+
+fn test_cast_between_string_and_decimal(file: &mut impl Write, is_try: bool) {
+    let prefix = if is_try { "TRY_" } else { "" };
+
+    run_ast(file, format!("{prefix}CAST('010.010' AS DECIMAL(5,3))"), &[
+    ]);
+    run_ast(file, format!("{prefix}CAST('010.010' AS DECIMAL(5,4))"), &[
+    ]);
+    run_ast(file, format!("{prefix}CAST('010.010' AS DECIMAL(5,2))"), &[
+    ]);
+    run_ast(file, format!("{prefix}CAST('010.010' AS DECIMAL(4,3))"), &[
+    ]);
+    run_ast(file, format!("{prefix}CAST('010.010' AS DECIMAL(4,2))"), &[
+    ]);
+    run_ast(
+        file,
+        format!("{prefix}CAST('-1010.010' AS DECIMAL(7,3))"),
+        &[],
+    );
+    run_ast(file, format!("{prefix}CAST('00' AS DECIMAL(2,1))"), &[]);
+    run_ast(file, format!("{prefix}CAST('0.0' AS DECIMAL(2,0))"), &[]);
+    run_ast(file, format!("{prefix}CAST('.0' AS DECIMAL(1,0))"), &[]);
+    run_ast(
+        file,
+        format!("{prefix}CAST('+1.0e-10' AS DECIMAL(11, 10))"),
+        &[],
+    );
+    run_ast(
+        file,
+        format!("{prefix}CAST('-1.0e+10' AS DECIMAL(11, 0))"),
+        &[],
+    );
+}
diff --git a/src/query/functions/tests/it/scalars/testdata/array.txt b/src/query/functions/tests/it/scalars/testdata/array.txt
@@ -1756,19 +1756,17 @@ error:
   --> SQL:1:12
   |
 1 | array_sort(['x', 0, 1.2, 3.4, 5.6, 7.8])
-  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unable to cast type `String` to type `Decimal(2, 1)`
+  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ fail to decode string as decimal: Code: 1006, displayText = bad decimal literal: unexpected char. while evaluating function `to_decimal("x")`
 
 
 
-error: 
-  --> SQL:1:12
-  |
-1 | array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'DESC', 'NULLS FIRST')
-  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `array(Decimal(2, 1), NULL, Decimal(2, 1), Decimal(2, 1), String, NULL)`
-
-has tried possible overloads:
-  array(T0, T0, T0, T0, T0, T0) :: Array(T0)  : unable to find a common super type for `Decimal(2, 1) NULL` and `String`
-
+ast            : array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'DESC', 'NULLS FIRST')
+raw expr       : array_sort_desc_null_first(array(1.2_decimal(2, 1), NULL, 3.4_decimal(2, 1), 5.6_decimal(2, 1), "2.2", NULL))
+checked expr   : array_sort_desc_null_first<T0=Decimal(2, 1) NULL><Array(T0)>(array<T0=Decimal(2, 1) NULL><T0, T0, T0, T0, T0, T0>(CAST(1.2_d128(2,1) AS Decimal(2, 1) NULL), CAST(NULL AS Decimal(2, 1) NULL), CAST(3.4_d128(2,1) AS Decimal(2, 1) NULL), CAST(5.6_d128(2,1) AS Decimal(2, 1) NULL), CAST("2.2" AS Decimal(2, 1) NULL), CAST(NULL AS Decimal(2, 1) NULL)))
+optimized expr : [NULL, NULL, 5.6, 3.4, 2.2, 1.2]
+output type    : Array(Decimal(2, 1) NULL)
+output domain  : [{0.0..=5.6} ∪ {NULL}]
+output         : [NULL, NULL, 5.6, 3.4, 2.2, 1.2]
 
 
 ast            : array_sort([], 'DESC', 'NULLS FIRST')
@@ -1780,15 +1778,13 @@ output domain  : []
 output         : []
 
 
-error: 
-  --> SQL:1:12
-  |
-1 | array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'DESC', 'NULLS LAST')
-  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `array(Decimal(2, 1), NULL, Decimal(2, 1), Decimal(2, 1), String, NULL)`
-
-has tried possible overloads:
-  array(T0, T0, T0, T0, T0, T0) :: Array(T0)  : unable to find a common super type for `Decimal(2, 1) NULL` and `String`
-
+ast            : array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'DESC', 'NULLS LAST')
+raw expr       : array_sort_desc_null_last(array(1.2_decimal(2, 1), NULL, 3.4_decimal(2, 1), 5.6_decimal(2, 1), "2.2", NULL))
+checked expr   : array_sort_desc_null_last<T0=Decimal(2, 1) NULL><Array(T0)>(array<T0=Decimal(2, 1) NULL><T0, T0, T0, T0, T0, T0>(CAST(1.2_d128(2,1) AS Decimal(2, 1) NULL), CAST(NULL AS Decimal(2, 1) NULL), CAST(3.4_d128(2,1) AS Decimal(2, 1) NULL), CAST(5.6_d128(2,1) AS Decimal(2, 1) NULL), CAST("2.2" AS Decimal(2, 1) NULL), CAST(NULL AS Decimal(2, 1) NULL)))
+optimized expr : [5.6, 3.4, 2.2, 1.2, NULL, NULL]
+output type    : Array(Decimal(2, 1) NULL)
+output domain  : [{0.0..=5.6} ∪ {NULL}]
+output         : [5.6, 3.4, 2.2, 1.2, NULL, NULL]
 
 
 ast            : array_sort([], 'DESC', 'NULLS LAST')
@@ -1800,15 +1796,13 @@ output domain  : []
 output         : []
 
 
-error: 
-  --> SQL:1:12
-  |
-1 | array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'ASC', 'NULLS FIRST')
-  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `array(Decimal(2, 1), NULL, Decimal(2, 1), Decimal(2, 1), String, NULL)`
-
-has tried possible overloads:
-  array(T0, T0, T0, T0, T0, T0) :: Array(T0)  : unable to find a common super type for `Decimal(2, 1) NULL` and `String`
-
+ast            : array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'ASC', 'NULLS FIRST')
+raw expr       : array_sort_asc_null_first(array(1.2_decimal(2, 1), NULL, 3.4_decimal(2, 1), 5.6_decimal(2, 1), "2.2", NULL))
+checked expr   : array_sort_asc_null_first<T0=Decimal(2, 1) NULL><Array(T0)>(array<T0=Decimal(2, 1) NULL><T0, T0, T0, T0, T0, T0>(CAST(1.2_d128(2,1) AS Decimal(2, 1) NULL), CAST(NULL AS Decimal(2, 1) NULL), CAST(3.4_d128(2,1) AS Decimal(2, 1) NULL), CAST(5.6_d128(2,1) AS Decimal(2, 1) NULL), CAST("2.2" AS Decimal(2, 1) NULL), CAST(NULL AS Decimal(2, 1) NULL)))
+optimized expr : [NULL, NULL, 1.2, 2.2, 3.4, 5.6]
+output type    : Array(Decimal(2, 1) NULL)
+output domain  : [{0.0..=5.6} ∪ {NULL}]
+output         : [NULL, NULL, 1.2, 2.2, 3.4, 5.6]
 
 
 ast            : array_sort([], 'ASC', 'NULLS FIRST')