add cast

Author: b41sh

src/query/expression/src/evaluator.rs

@@ -50,6 +50,10 @@ use crate::types::StringType;
 use crate::types::ValueType;
 use crate::types::VariantType;
 use crate::values::Column;
+use crate::types::VectorColumn;
+use crate::types::VectorDataType;
+use crate::types::VectorScalar;
+use crate::types::F32;
 use crate::values::ColumnBuilder;
 use crate::values::Scalar;
 use crate::values::Value;
@@ -924,6 +928,85 @@ impl<'a> Evaluator<'a> {
                     other => unreachable!("source: {}", other),
                 }
             }
+            (DataType::Array(inner_src_ty), DataType::Vector(inner_dest_ty)) => {
+                if !matches!(&**inner_src_ty, DataType::Number(_) | DataType::Decimal(_))
+                    || matches!(inner_dest_ty, VectorDataType::Int8(_))
+                {
+                    return Err(ErrorCode::BadArguments(format!(
+                        "unable to cast type `{src_type}` to type `{dest_type}`"
+                    ))
+                    .set_span(span));
+                }
+                let dimension = inner_dest_ty.dimension() as usize;
+                match value {
+                    Value::Scalar(Scalar::Array(col)) => {
+                        if col.len() != dimension {
+                            return Err(ErrorCode::BadArguments(
+                                "Array value cast to a vector has incorrect dimension".to_string(),
+                            )
+                            .set_span(span));
+                        }
+                        let mut vals = Vec::with_capacity(dimension);
+                        match col {
+                            Column::Number(num_col) => {
+                                for i in 0..dimension {
+                                    let num = unsafe { num_col.index_unchecked(i) };
+                                    vals.push(num.to_f32());
+                                }
+                            }
+                            Column::Decimal(dec_col) => {
+                                for i in 0..dimension {
+                                    let dec = unsafe { dec_col.index_unchecked(i) };
+                                    vals.push(F32::from(dec.to_float32()));
+                                }
+                            }
+                            _ => {
+                                return Err(ErrorCode::BadArguments(
+                                    "Array value cast to a vector has invalid value".to_string(),
+                                )
+                                .set_span(span));
+                            }
+                        }
+                        Ok(Value::Scalar(Scalar::Vector(VectorScalar::Float32(vals))))
+                    }
+                    Value::Column(Column::Array(array_col)) => {
+                        let mut vals = Vec::with_capacity(dimension * array_col.len());
+                        for col in array_col.iter() {
+                            if col.len() != dimension {
+                                return Err(ErrorCode::BadArguments(
+                                    "Array value cast to a vector has incorrect dimension"
+                                        .to_string(),
+                                )
+                                .set_span(span));
+                            }
+                            match col {
+                                Column::Number(num_col) => {
+                                    for i in 0..dimension {
+                                        let num = unsafe { num_col.index_unchecked(i) };
+                                        vals.push(num.to_f32());
+                                    }
+                                }
+                                Column::Decimal(dec_col) => {
+                                    for i in 0..dimension {
+                                        let dec = unsafe { dec_col.index_unchecked(i) };
+                                        vals.push(F32::from(dec.to_float32()));
+                                    }
+                                }
+                                _ => {
+                                    return Err(ErrorCode::BadArguments(
+                                        "Array value cast to a vector has invalid value"
+                                            .to_string(),
+                                    )
+                                    .set_span(span));
+                                }
+                            }
+                        }
+                        let vector_col = VectorColumn::Float32((vals.into(), dimension));
+                        Ok(Value::Column(Column::Vector(vector_col)))
+                    }
+                    other => unreachable!("source: {}", other),
+                }
+            }
 
             _ => Err(ErrorCode::BadArguments(format!(
                 "unable to cast type `{src_type}` to type `{dest_type}`"

src/query/expression/src/type_check.rs

@@ -604,7 +604,11 @@ fn can_cast_to(src_ty: &DataType, dest_ty: &DataType) -> bool {
         {
             true
         }
-
+        (DataType::Array(fields_src_ty), DataType::Vector(_))
+            if matches!(&**fields_src_ty, DataType::Number(_) | DataType::Decimal(_)) =>
+        {
+            true
+        }
         (DataType::Nullable(box inner_src_ty), DataType::Nullable(box inner_dest_ty))
         | (DataType::Nullable(box inner_src_ty), inner_dest_ty)
         | (inner_src_ty, DataType::Nullable(box inner_dest_ty))

src/query/expression/src/types/decimal.rs

@@ -176,6 +176,12 @@ pub enum DecimalScalar {
 }
 
 impl DecimalScalar {
+    pub fn to_float32(&self) -> f32 {
+        with_decimal_type!(|DECIMAL| match self {
+            DecimalScalar::DECIMAL(v, size) => v.to_float32(size.scale),
+        })
+    }
+
     pub fn to_float64(&self) -> f64 {
         with_decimal_type!(|DECIMAL| match self {
             DecimalScalar::DECIMAL(v, size) => v.to_float64(size.scale),

src/query/expression/src/types/number.rs

@@ -489,6 +489,14 @@ impl NumberScalar {
         }
     }
 
+    pub fn to_f32(&self) -> F32 {
+        crate::with_integer_mapped_type!(|NUM_TYPE| match self {
+            NumberScalar::NUM_TYPE(num) => (*num as f32).into(),
+            NumberScalar::Float32(num) => *num,
+            NumberScalar::Float64(num) => (num.into_inner() as f32).into(),
+        })
+    }
+
     pub fn to_f64(&self) -> F64 {
         crate::with_integer_mapped_type!(|NUM_TYPE| match self {
             NumberScalar::NUM_TYPE(num) => (*num as f64).into(),