chore: Helper function to convert Vec<Value> to VectorRef (#4546)

* chore: `try_from_row_into_vector` helper * test: try_from_row * refactor: simplify with builder * fix: deicmal set prec&scale * refactor: more simplify * refactor: use ref
GreptimeTeam · Aug 14, 2024 · b59a93d · b59a93d
1 parent 202c730
commit b59a93d
Show file tree

Hide file tree

Showing 5 changed files with 127 additions and 0 deletions.
diff --git a/src/datatypes/src/duration.rs b/src/datatypes/src/duration.rs
@@ -111,6 +111,24 @@ macro_rules! define_duration_with_unit {
                     val.0.value()
                 }
             }
+
+            impl TryFrom<Value> for Option<[<Duration $unit>]> {
+                type Error = $crate::error::Error;
+
+                #[inline]
+                fn try_from(from: Value) -> std::result::Result<Self, Self::Error> {
+                    match from {
+                        Value::Duration(v) if v.unit() == TimeUnit::$unit => {
+                            Ok(Some([<Duration $unit>](v)))
+                        },
+                        Value::Null => Ok(None),
+                        _ => $crate::error::TryFromValueSnafu {
+                            reason: format!("{:?} is not a {}", from, stringify!([<Duration $unit>])),
+                        }
+                        .fail(),
+                    }
+                }
+            }
         }
     };
 }

diff --git a/src/datatypes/src/interval.rs b/src/datatypes/src/interval.rs
@@ -106,6 +106,24 @@ macro_rules! define_interval_with_unit {
                     val.0.[<to_ $native_ty>]()
                 }
             }
+
+            impl TryFrom<Value> for Option<[<Interval $unit>]> {
+                type Error = $crate::error::Error;
+
+                #[inline]
+                fn try_from(from: Value) -> std::result::Result<Self, Self::Error> {
+                    match from {
+                        Value::Interval(v) if v.unit() == common_time::interval::IntervalUnit::$unit => {
+                            Ok(Some([<Interval $unit>](v)))
+                        },
+                        Value::Null => Ok(None),
+                        _ => $crate::error::TryFromValueSnafu {
+                            reason: format!("{:?} is not a {}", from, stringify!([<Interval $unit>])),
+                        }
+                        .fail(),
+                    }
+                }
+            }
         }
     };
 }

diff --git a/src/datatypes/src/time.rs b/src/datatypes/src/time.rs
@@ -109,6 +109,24 @@ macro_rules! define_time_with_unit {
                     val.0.value()
                 }
             }
+
+            impl TryFrom<Value> for Option<[<Time $unit>]> {
+                type Error = $crate::error::Error;
+
+                #[inline]
+                fn try_from(from: Value) -> std::result::Result<Self, Self::Error> {
+                    match from {
+                        Value::Time(v) if *v.unit() == TimeUnit::$unit => {
+                            Ok(Some([<Time $unit>](v)))
+                        },
+                        Value::Null => Ok(None),
+                        _ => $crate::error::TryFromValueSnafu {
+                            reason: format!("{:?} is not a {}", from, stringify!([<Time $unit>])),
+                        }
+                        .fail(),
+                    }
+                }
+            }
         }
     };
 }

diff --git a/src/datatypes/src/timestamp.rs b/src/datatypes/src/timestamp.rs
@@ -111,6 +111,24 @@ macro_rules! define_timestamp_with_unit {
                     val.0.value()
                 }
             }
+
+            impl TryFrom<Value> for Option<[<Timestamp $unit>]> {
+                type Error = $crate::error::Error;
+
+                #[inline]
+                fn try_from(from: Value) -> std::result::Result<Self, Self::Error> {
+                    match from {
+                        Value::Timestamp(v) if v.unit() == TimeUnit::$unit => {
+                            Ok(Some([<Timestamp $unit>](v)))
+                        },
+                        Value::Null => Ok(None),
+                        _ => $crate::error::TryFromValueSnafu {
+                            reason: format!("{:?} is not a {}", from, stringify!([<Timestamp $unit>])),
+                        }
+                        .fail(),
+                    }
+                }
+            }
         }
     };
 }

diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs
@@ -27,6 +27,7 @@ use snafu::{OptionExt, ResultExt};
 
 use crate::data_type::ConcreteDataType;
 use crate::error::{self, ConvertArrowArrayToScalarsSnafu, Result};
+use crate::prelude::DataType;
 use crate::scalars::{Scalar, ScalarVectorBuilder};
 use crate::value::{ListValue, ListValueRef, Value};
 use crate::vectors::{
@@ -367,6 +368,16 @@ impl Helper {
         })
     }
 
+    /// Try to cast an vec of values into vector, fail if type is not the same across all values.
+    pub fn try_from_row_into_vector(row: &[Value], dt: &ConcreteDataType) -> Result<VectorRef> {
+        let mut builder = dt.create_mutable_vector(row.len());
+        for val in row {
+            builder.try_push_value_ref(val.as_value_ref())?;
+        }
+        let vector = builder.to_vector();
+        Ok(vector)
+    }
+
     /// Try to cast slice of `arrays` to vectors.
     pub fn try_into_vectors(arrays: &[ArrayRef]) -> Result<Vec<VectorRef>> {
         arrays.iter().map(Self::try_into_vector).collect()
@@ -681,4 +692,48 @@ mod tests {
             assert_eq!(Value::Interval(Interval::from_i128(2000)), vector.get(i));
         }
     }
+
+    fn check_try_from_row_to_vector(row: Vec<Value>, dt: &ConcreteDataType) {
+        let vector = Helper::try_from_row_into_vector(&row, dt).unwrap();
+        for (i, item) in row.iter().enumerate().take(vector.len()) {
+            assert_eq!(*item, vector.get(i));
+        }
+    }
+
+    fn check_into_and_from(array: impl Array + 'static) {
+        let array: ArrayRef = Arc::new(array);
+        let vector = Helper::try_into_vector(array.clone()).unwrap();
+        assert_eq!(&array, &vector.to_arrow_array());
+        let row: Vec<Value> = (0..array.len()).map(|i| vector.get(i)).collect();
+        let dt = vector.data_type();
+        check_try_from_row_to_vector(row, &dt);
+    }
+
+    #[test]
+    fn test_try_from_row_to_vector() {
+        check_into_and_from(NullArray::new(2));
+        check_into_and_from(BooleanArray::from(vec![true, false]));
+        check_into_and_from(Int8Array::from(vec![1, 2, 3]));
+        check_into_and_from(Int16Array::from(vec![1, 2, 3]));
+        check_into_and_from(Int32Array::from(vec![1, 2, 3]));
+        check_into_and_from(Int64Array::from(vec![1, 2, 3]));
+        check_into_and_from(UInt8Array::from(vec![1, 2, 3]));
+        check_into_and_from(UInt16Array::from(vec![1, 2, 3]));
+        check_into_and_from(UInt32Array::from(vec![1, 2, 3]));
+        check_into_and_from(UInt64Array::from(vec![1, 2, 3]));
+        check_into_and_from(Float32Array::from(vec![1.0, 2.0, 3.0]));
+        check_into_and_from(Float64Array::from(vec![1.0, 2.0, 3.0]));
+        check_into_and_from(StringArray::from(vec!["hello", "world"]));
+        check_into_and_from(Date32Array::from(vec![1, 2, 3]));
+        check_into_and_from(Date64Array::from(vec![1, 2, 3]));
+
+        check_into_and_from(TimestampSecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(TimestampMillisecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(TimestampMicrosecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(TimestampNanosecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(Time32SecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(Time32MillisecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(Time64MicrosecondArray::from(vec![1, 2, 3]));
+        check_into_and_from(Time64NanosecondArray::from(vec![1, 2, 3]));
+    }
 }