fix: adjust try_into_vector logic

GreptimeTeam · May 13, 2024 · 58c38c0 · 58c38c0
1 parent 269dceb
commit 58c38c0
Showing 1 changed file with 24 additions and 9 deletions.
diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs
@@ -258,9 +258,10 @@ impl Helper {
         Ok(match array.as_ref().data_type() {
             ArrowDataType::Null => Arc::new(NullVector::try_from_arrow_array(array)?),
             ArrowDataType::Boolean => Arc::new(BooleanVector::try_from_arrow_array(array)?),
-            ArrowDataType::LargeBinary => Arc::new(BinaryVector::try_from_arrow_array(array)?),
-            ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Binary => {
-                let array = arrow::compute::cast(array.as_ref(), &ArrowDataType::LargeBinary)
+            ArrowDataType::LargeBinary
+            | ArrowDataType::FixedSizeBinary(_)
+            | ArrowDataType::Binary => {
+                let array = arrow::compute::cast(array.as_ref(), &ArrowDataType::Binary)
                     .context(crate::error::ArrowComputeSnafu)?;
                 Arc::new(BinaryVector::try_from_arrow_array(array)?)
             }
@@ -396,14 +397,15 @@ impl Helper {
 #[cfg(test)]
 mod tests {
     use arrow::array::{
-        ArrayRef, BooleanArray, Date32Array, Date64Array, Float32Array, Float64Array, Int16Array,
-        Int32Array, Int64Array, Int8Array, LargeBinaryArray, ListArray, NullArray,
+        ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array, Float32Array, Float64Array,
+        Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray, ListArray, NullArray,
         Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
         TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
         TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
     };
     use arrow::datatypes::Int32Type;
     use arrow_array::DictionaryArray;
+    use arrow_schema::DataType;
     use common_decimal::Decimal128;
     use common_time::time::Time;
     use common_time::timestamp::TimeUnit;
@@ -572,14 +574,27 @@ mod tests {
         assert_eq!(&array, &vector.to_arrow_array());
     }
 
+    #[test]
+    fn test_try_binary_array_into_vector() {
+        let input_vec = vec!["hello".as_bytes(), "world".as_bytes()];
+
+        let input_large_binary_array: ArrayRef =
+            Arc::new(LargeBinaryArray::from(input_vec.clone()));
+        let assertion_array: ArrayRef = Arc::new(BinaryArray::from(input_vec));
+        let vector = Helper::try_into_vector(input_large_binary_array.clone()).unwrap();
+
+        assert_eq!(2, vector.len());
+        assert_eq!(0, vector.null_count());
+
+        let output_arrow_array: ArrayRef = vector.to_arrow_array();
+        assert_eq!(&DataType::Binary, output_arrow_array.data_type());
+        assert_eq!(&assertion_array, &output_arrow_array);
+    }
+
     #[test]
     fn test_try_into_vector() {
         check_try_into_vector(NullArray::new(2));
         check_try_into_vector(BooleanArray::from(vec![true, false]));
-        check_try_into_vector(LargeBinaryArray::from(vec![
-            "hello".as_bytes(),
-            "world".as_bytes(),
-        ]));
         check_try_into_vector(Int8Array::from(vec![1, 2, 3]));
         check_try_into_vector(Int16Array::from(vec![1, 2, 3]));
         check_try_into_vector(Int32Array::from(vec![1, 2, 3]));