chore: return error in arrow array convert function (#2043)

fix: return error instead of unreachable!()
GreptimeTeam · Jul 31, 2023 · f22b787 · f22b787
1 parent 81ea61b
commit f22b787
Showing 1 changed file with 52 additions and 5 deletions.
diff --git a/src/datatypes/src/vectors/primitive.rs b/src/datatypes/src/vectors/primitive.rs
@@ -28,7 +28,7 @@ use serde_json::Value as JsonValue;
 use snafu::OptionExt;
 
 use crate::data_type::ConcreteDataType;
-use crate::error::{self, Result};
+use crate::error::{self, CastTypeSnafu, Result};
 use crate::scalars::{Scalar, ScalarRef, ScalarVector, ScalarVectorBuilder};
 use crate::serialize::Serializable;
 use crate::types::{
@@ -104,8 +104,14 @@ impl<T: LogicalPrimitiveType> PrimitiveVector<T> {
                     .with_timezone_opt(None::<String>)
                     .to_data(),
             },
-            _ => {
-                unreachable!()
+            arrow_type => {
+                return CastTypeSnafu {
+                    msg: format!(
+                        "Failed to cast arrow array {:?} to timestamp vector",
+                        arrow_type,
+                    ),
+                }
+                .fail()?;
             }
         };
         let concrete_array = PrimitiveArray::<T::ArrowPrimitive>::from(array_data);
@@ -142,8 +148,11 @@ impl<T: LogicalPrimitiveType> PrimitiveVector<T> {
                     .to_data(),
                 _ => unreachable!(),
             },
-            _ => {
-                unreachable!()
+            arrow_type => {
+                return CastTypeSnafu {
+                    msg: format!("Failed to cast arrow array {:?} to time vector", arrow_type,),
+                }
+                .fail()?;
             }
         };
         let concrete_array = PrimitiveArray::<T::ArrowPrimitive>::from(array_data);
@@ -494,6 +503,8 @@ mod tests {
     use crate::types::Int64Type;
     use crate::vectors::{
         TimeMicrosecondVector, TimeMillisecondVector, TimeNanosecondVector, TimeSecondVector,
+        TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector,
+        TimestampSecondVector,
     };
 
     fn check_vec(v: Int32Vector) {
@@ -663,5 +674,41 @@ mod tests {
         let array: ArrayRef = Arc::new(Time64NanosecondArray::from(vec![1i64, 2, 3]));
         let vector = TimeNanosecondVector::try_from_arrow_time_array(array).unwrap();
         assert_eq!(TimeNanosecondVector::from_values(vec![1, 2, 3]), vector);
+
+        // Test convert error
+        let array: ArrayRef = Arc::new(Int32Array::from(vec![1i32, 2, 3]));
+        assert!(TimeSecondVector::try_from_arrow_time_array(array).is_err());
+    }
+
+    #[test]
+    fn test_try_from_arrow_timestamp_array() {
+        let array: ArrayRef = Arc::new(TimestampSecondArray::from(vec![1i64, 2, 3]));
+        let vector = TimestampSecondVector::try_from_arrow_timestamp_array(array).unwrap();
+        assert_eq!(TimestampSecondVector::from_values(vec![1, 2, 3]), vector);
+
+        let array: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1i64, 2, 3]));
+        let vector = TimestampMillisecondVector::try_from_arrow_timestamp_array(array).unwrap();
+        assert_eq!(
+            TimestampMillisecondVector::from_values(vec![1, 2, 3]),
+            vector
+        );
+
+        let array: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![1i64, 2, 3]));
+        let vector = TimestampMicrosecondVector::try_from_arrow_timestamp_array(array).unwrap();
+        assert_eq!(
+            TimestampMicrosecondVector::from_values(vec![1, 2, 3]),
+            vector
+        );
+
+        let array: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![1i64, 2, 3]));
+        let vector = TimestampNanosecondVector::try_from_arrow_timestamp_array(array).unwrap();
+        assert_eq!(
+            TimestampNanosecondVector::from_values(vec![1, 2, 3]),
+            vector
+        );
+
+        // Test convert error
+        let array: ArrayRef = Arc::new(Int32Array::from(vec![1i32, 2, 3]));
+        assert!(TimestampSecondVector::try_from_arrow_timestamp_array(array).is_err());
     }
 }