From f22b787fd94c071f4104383b2e5a31c9844e1b11 Mon Sep 17 00:00:00 2001
From: Zou Wei <47681251+QuenKar@users.noreply.github.com>
Date: Mon, 31 Jul 2023 10:47:40 +0800
Subject: [PATCH] chore: return error in arrow array convert function (#2043)

fix: return error instead of unreachable!()
---
 src/datatypes/src/vectors/primitive.rs | 57 +++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/datatypes/src/vectors/primitive.rs b/src/datatypes/src/vectors/primitive.rs
index 84ea389c5d1..3b61a8da093 100644
--- 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());
     }
 }