Add support for Utf8Scalar

DataEngineeringLabs · Mar 8, 2023 · 39eb555 · 39eb555
1 parent b2fd8f8
commit 39eb555
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 1 deletion.
diff --git a/arrow2_convert/src/deserialize.rs b/arrow2_convert/src/deserialize.rs
@@ -146,6 +146,24 @@ impl ArrowDeserialize for LargeString {
     }
 }
 
+impl ArrowDeserialize for arrow2::scalar::Utf8Scalar<i32> {
+    type ArrayType = Utf8Array<i32>;
+
+    #[inline]
+    fn arrow_deserialize(v: Option<&str>) -> Option<Self> {
+        Some(arrow2::scalar::Utf8Scalar::new(v))
+    }
+}
+
+impl ArrowDeserialize for arrow2::scalar::Utf8Scalar<i64> {
+    type ArrayType = Utf8Array<i64>;
+
+    #[inline]
+    fn arrow_deserialize(v: Option<&str>) -> Option<Self> {
+        Some(arrow2::scalar::Utf8Scalar::new(v))
+    }
+}
+
 impl ArrowDeserialize for bool {
     type ArrayType = BooleanArray;
 

diff --git a/arrow2_convert/src/field.rs b/arrow2_convert/src/field.rs
@@ -134,6 +134,24 @@ impl ArrowField for String {
     }
 }
 
+impl ArrowField for arrow2::scalar::Utf8Scalar<i32> {
+    type Type = arrow2::scalar::Utf8Scalar<i32>;
+
+    #[inline]
+    fn data_type() -> arrow2::datatypes::DataType {
+        arrow2::datatypes::DataType::Utf8
+    }
+}
+
+impl ArrowField for arrow2::scalar::Utf8Scalar<i64> {
+    type Type = arrow2::scalar::Utf8Scalar<i64>;
+
+    #[inline]
+    fn data_type() -> arrow2::datatypes::DataType {
+        arrow2::datatypes::DataType::LargeUtf8
+    }
+}
+
 /// Represents the `LargeUtf8` Arrow type
 pub struct LargeString {}
 

diff --git a/arrow2_convert/src/serialize.rs b/arrow2_convert/src/serialize.rs
@@ -134,6 +134,34 @@ impl ArrowSerialize for LargeString {
     }
 }
 
+impl ArrowSerialize for arrow2::scalar::Utf8Scalar<i32> {
+    type MutableArrayType = MutableUtf8Array<i32>;
+
+    #[inline]
+    fn new_array() -> Self::MutableArrayType {
+        Self::MutableArrayType::default()
+    }
+
+    #[inline]
+    fn arrow_serialize(v: &Self, array: &mut Self::MutableArrayType) -> arrow2::error::Result<()> {
+        array.try_push(v.value())
+    }
+}
+
+impl ArrowSerialize for arrow2::scalar::Utf8Scalar<i64> {
+    type MutableArrayType = MutableUtf8Array<i64>;
+
+    #[inline]
+    fn new_array() -> Self::MutableArrayType {
+        Self::MutableArrayType::default()
+    }
+
+    #[inline]
+    fn arrow_serialize(v: &Self, array: &mut Self::MutableArrayType) -> arrow2::error::Result<()> {
+        array.try_push(v.value())
+    }
+}
+
 impl ArrowSerialize for bool {
     type MutableArrayType = MutableBooleanArray;
 

diff --git a/arrow2_convert/tests/complex_example.rs b/arrow2_convert/tests/complex_example.rs
@@ -53,6 +53,8 @@ pub struct Root {
     // fixed size vec
     #[arrow_field(type = "arrow2_convert::field::FixedSizeVec<i64, 3>")]
     fixed_size_vec: Vec<i64>,
+    utf8scalar_i32: arrow2::scalar::Utf8Scalar<i32>,
+    utf8scalar_i64: arrow2::scalar::Utf8Scalar<i64>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, ArrowField, ArrowSerialize, ArrowDeserialize)]
@@ -160,6 +162,8 @@ fn item1() -> Root {
         large_string: "abcdefg".to_string(),
         large_vec: vec![1, 2, 3, 4],
         fixed_size_vec: vec![10, 20, 30],
+        utf8scalar_i32: arrow2::scalar::Utf8Scalar::new(Some("utf8-string")),
+        utf8scalar_i64: arrow2::scalar::Utf8Scalar::new(Some("utf8-string")),
     }
 }
 
@@ -206,6 +210,8 @@ fn item2() -> Root {
         large_string: "abdefag".to_string(),
         large_vec: vec![5, 4, 3, 2],
         fixed_size_vec: vec![11, 21, 32],
+        utf8scalar_i32: arrow2::scalar::Utf8Scalar::new(Some("utf8-string-again")),
+        utf8scalar_i64: arrow2::scalar::Utf8Scalar::new(Some("utf8-string-again")),
     }
 }
 
@@ -222,7 +228,7 @@ fn test_round_trip() -> arrow2::error::Result<()> {
     assert_eq!(struct_array.len(), 2);
 
     let values = struct_array.values();
-    assert_eq!(values.len(), 21);
+    assert_eq!(values.len(), 23);
     assert_eq!(struct_array.len(), 2);
 
     // can iterate one struct at a time without collecting