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
This commit is contained in:
discord9
2024-08-14 11:11:44 +08:00
committed by GitHub
parent 202c730363
commit b59a93dfbc
5 changed files with 127 additions and 0 deletions

View File

@@ -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(),
}
}
}
}
};
}

View File

@@ -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(),
}
}
}
}
};
}

View File

@@ -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(),
}
}
}
}
};
}

View File

@@ -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(),
}
}
}
}
};
}

View File

@@ -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]));
}
}