Created
October 8, 2022 01:14
-
-
Save aldanor/de7accf7bb40e83bec6e1e51f08f1190 to your computer and use it in GitHub Desktop.
arrow2-convert-generics
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| /// Generics example | |
| use arrow2::array::Array; | |
| use arrow2_convert::{deserialize::TryIntoCollection, serialize::TryIntoArrow}; | |
| #[derive(Debug, Clone, PartialEq, Eq)] | |
| pub struct Foo<A, B> | |
| where | |
| A: Clone, | |
| { | |
| name: String, | |
| a: A, | |
| b: B, | |
| } | |
| impl<A, B> arrow2_convert::field::ArrowField for Foo<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::field::ArrowField, | |
| B: ::arrow2_convert::field::ArrowField, | |
| { | |
| type Type = Self; | |
| fn data_type() -> arrow2::datatypes::DataType { | |
| arrow2::datatypes::DataType::Struct(<[_]>::into_vec(Box::new([ | |
| <String as arrow2_convert::field::ArrowField>::field("name"), | |
| <A as arrow2_convert::field::ArrowField>::field("a"), | |
| <B as arrow2_convert::field::ArrowField>::field("b"), | |
| ]))) | |
| } | |
| } | |
| impl<A, B> arrow2_convert::field::ArrowEnableVecForType for Foo<A, B> where A: Clone {} | |
| struct MutableFooArrayFields<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| name: <String as arrow2_convert::serialize::ArrowSerialize>::MutableArrayType, | |
| a: <A as arrow2_convert::serialize::ArrowSerialize>::MutableArrayType, | |
| b: <B as arrow2_convert::serialize::ArrowSerialize>::MutableArrayType, | |
| } | |
| pub struct MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fields: MutableFooArrayFields<A, B>, | |
| data_type: arrow2::datatypes::DataType, | |
| validity: Option<arrow2::bitmap::MutableBitmap>, | |
| } | |
| #[automatically_derived] | |
| impl<A, B> ::core::fmt::Debug for MutableFooArrayFields<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result { | |
| f.debug_struct("MutableFooArrayFields") | |
| .field("name", &self.name) | |
| .field("a", &self.a) | |
| .field("b", &self.b) | |
| .finish() | |
| } | |
| } | |
| #[automatically_derived] | |
| impl<A, B> ::core::fmt::Debug for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result { | |
| f.debug_struct("MutableFooArray") | |
| .field("fields", &self.fields) | |
| .field("data_type", &self.data_type) | |
| .field("validity", &self.validity) | |
| .finish() | |
| } | |
| } | |
| impl<A, B> MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| pub fn new() -> Self { | |
| Self { | |
| fields: MutableFooArrayFields { | |
| name: <String as arrow2_convert::serialize::ArrowSerialize>::new_array(), | |
| a: <A as arrow2_convert::serialize::ArrowSerialize>::new_array(), | |
| b: <B as arrow2_convert::serialize::ArrowSerialize>::new_array(), | |
| }, | |
| data_type: <Foo<A, B> as arrow2_convert::field::ArrowField>::data_type(), | |
| validity: None, | |
| } | |
| } | |
| fn init_validity(&mut self) { | |
| let mut validity = arrow2::bitmap::MutableBitmap::new(); | |
| validity.extend_constant(<Self as arrow2::array::MutableArray>::len(self), true); | |
| validity.set(<Self as arrow2::array::MutableArray>::len(self) - 1, false); | |
| self.validity = Some(validity) | |
| } | |
| } | |
| impl<A, B> Default for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn default() -> Self { | |
| Self::new() | |
| } | |
| } | |
| impl<__T: std::borrow::Borrow<Foo<A, B>>, A, B> arrow2::array::TryPush<Option<__T>> | |
| for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn try_push(&mut self, item: Option<__T>) -> arrow2::error::Result<()> { | |
| use arrow2::array::MutableArray; | |
| use std::borrow::Borrow; | |
| match item { | |
| Some(i) => { | |
| let i = i.borrow(); | |
| <String as arrow2_convert::serialize::ArrowSerialize>::arrow_serialize( | |
| i.name.borrow(), | |
| &mut self.fields.name, | |
| )?; | |
| <A as arrow2_convert::serialize::ArrowSerialize>::arrow_serialize( | |
| i.a.borrow(), | |
| &mut self.fields.a, | |
| )?; | |
| <B as arrow2_convert::serialize::ArrowSerialize>::arrow_serialize( | |
| i.b.borrow(), | |
| &mut self.fields.b, | |
| )?; | |
| match &mut self.validity { | |
| Some(validity) => validity.push(true), | |
| None => {} | |
| } | |
| } | |
| None => { | |
| < < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as MutableArray > :: push_null (& mut self . fields . name) ; | |
| < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as MutableArray > :: push_null (& mut self . fields . a) ; | |
| < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as MutableArray > :: push_null (& mut self . fields . b) ; | |
| match &mut self.validity { | |
| Some(validity) => validity.push(false), | |
| None => { | |
| self.init_validity(); | |
| } | |
| } | |
| } | |
| } | |
| Ok(()) | |
| } | |
| } | |
| impl<__T: std::borrow::Borrow<Foo<A, B>>, A, B> arrow2::array::TryExtend<Option<__T>> | |
| for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn try_extend<I: IntoIterator<Item = Option<__T>>>( | |
| &mut self, | |
| iter: I, | |
| ) -> arrow2::error::Result<()> { | |
| use arrow2::array::TryPush; | |
| for i in iter { | |
| self.try_push(i)?; | |
| } | |
| Ok(()) | |
| } | |
| } | |
| impl<A, B> arrow2::array::MutableArray for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn data_type(&self) -> &arrow2::datatypes::DataType { | |
| &self.data_type | |
| } | |
| fn len(&self) -> usize { | |
| self.fields.name.len() | |
| } | |
| fn validity(&self) -> Option<&arrow2::bitmap::MutableBitmap> { | |
| self.validity.as_ref() | |
| } | |
| fn as_box(&mut self) -> Box<dyn arrow2::array::Array> { | |
| let values = < [_] > :: into_vec (Box :: new ([< < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . name) , < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . a) , < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . b)])) ; | |
| Box::new(arrow2::array::StructArray::from_data( | |
| <Foo<A, B> as arrow2_convert::field::ArrowField>::data_type().clone(), | |
| values, | |
| std::mem::take(&mut self.validity).map(|x| x.into()), | |
| )) | |
| } | |
| fn as_arc(&mut self) -> std::sync::Arc<dyn arrow2::array::Array> { | |
| let values = < [_] > :: into_vec (Box :: new ([< < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . name) , < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . a) , < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . b)])) ; | |
| std::sync::Arc::new(arrow2::array::StructArray::from_data( | |
| <Foo<A, B> as arrow2_convert::field::ArrowField>::data_type().clone(), | |
| values, | |
| std::mem::take(&mut self.validity).map(|x| x.into()), | |
| )) | |
| } | |
| fn as_any(&self) -> &dyn std::any::Any { | |
| self | |
| } | |
| fn as_mut_any(&mut self) -> &mut dyn std::any::Any { | |
| self | |
| } | |
| fn push_null(&mut self) { | |
| use arrow2::array::TryPush; | |
| self.try_push(None::<Foo<A, B>>).unwrap(); | |
| } | |
| fn shrink_to_fit(&mut self) { | |
| < < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: shrink_to_fit (& mut self . fields . name) ; | |
| < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: shrink_to_fit (& mut self . fields . a) ; | |
| < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: shrink_to_fit (& mut self . fields . b) ; | |
| if let Some(validity) = &mut self.validity { | |
| validity.shrink_to_fit(); | |
| } | |
| } | |
| fn reserve(&mut self, additional: usize) { | |
| if let Some(x) = self.validity.as_mut() { | |
| x.reserve(additional) | |
| } | |
| < < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: reserve (& mut self . fields . name , additional) ; | |
| < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: reserve (& mut self . fields . a , additional) ; | |
| < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: reserve (& mut self . fields . b , additional) ; | |
| } | |
| } | |
| impl<A, B> arrow2_convert::serialize::ArrowSerialize for Foo<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| type MutableArrayType = MutableFooArray<A, B>; | |
| #[inline] | |
| fn new_array() -> Self::MutableArrayType { | |
| Self::MutableArrayType::default() | |
| } | |
| #[inline] | |
| fn arrow_serialize(v: &Self, array: &mut Self::MutableArrayType) -> arrow2::error::Result<()> { | |
| use arrow2::array::TryPush; | |
| array.try_push(Some(v)) | |
| } | |
| } | |
| pub struct FooArray<A, B> { | |
| name: std::marker::PhantomData<String>, | |
| a: std::marker::PhantomData<A>, | |
| b: std::marker::PhantomData<B>, | |
| } | |
| impl<A, B> arrow2_convert::deserialize::ArrowArray for FooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| { | |
| type BaseArrayType = arrow2::array::StructArray; | |
| #[inline] | |
| fn iter_from_array_ref<'a>( | |
| b: &'a dyn arrow2::array::Array, | |
| ) -> <&'a Self as IntoIterator>::IntoIter { | |
| use core::ops::Deref; | |
| let arr = b | |
| .as_any() | |
| .downcast_ref::<arrow2::array::StructArray>() | |
| .unwrap(); | |
| let values = arr.values(); | |
| let validity = arr.validity(); | |
| FooArrayIterator { fields : FooArrayIteratorFields { name : < < String as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as arrow2_convert :: deserialize :: ArrowArray > :: iter_from_array_ref (values [0] . deref ()) , a : < < A as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as arrow2_convert :: deserialize :: ArrowArray > :: iter_from_array_ref (values [1] . deref ()) , b : < < B as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as arrow2_convert :: deserialize :: ArrowArray > :: iter_from_array_ref (values [2] . deref ()) , } , has_validity : validity . as_ref () . is_some () , validity_iter : validity . as_ref () . map (| x | x . iter ()) . unwrap_or_else (| | arrow2 :: bitmap :: utils :: BitmapIter :: new (& [] , 0 , 0)) , } | |
| } | |
| } | |
| impl<'a, A, B> IntoIterator for &'a FooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| { | |
| type Item = Option<Foo<A, B>>; | |
| type IntoIter = FooArrayIterator<'a, A, B>; | |
| fn into_iter(self) -> Self::IntoIter { | |
| unimplemented!(); | |
| } | |
| } | |
| struct FooArrayIteratorFields < 'a , A , B > where A : Clone , A : :: arrow2_convert :: deserialize :: ArrowDeserialize , for < '_a > & '_a < A as :: arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType : IntoIterator , B : :: arrow2_convert :: deserialize :: ArrowDeserialize , for < '_a > & '_a < B as :: arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType : IntoIterator { name : < & 'a < String as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as IntoIterator > :: IntoIter , a : < & 'a < A as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as IntoIterator > :: IntoIter , b : < & 'a < B as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as IntoIterator > :: IntoIter , } | |
| pub struct FooArrayIterator<'a, A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| { | |
| fields: FooArrayIteratorFields<'a, A, B>, | |
| validity_iter: arrow2::bitmap::utils::BitmapIter<'a>, | |
| has_validity: bool, | |
| } | |
| impl<'a, A, B> FooArrayIterator<'a, A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| { | |
| #[inline] | |
| fn return_next(&mut self) -> Option<Foo<A, B>> { | |
| if let (Some(name), Some(a), Some(b)) = ( | |
| self.fields.name.next(), | |
| self.fields.a.next(), | |
| self.fields.b.next(), | |
| ) { | |
| Some (Foo { name : < String as arrow2_convert :: deserialize :: ArrowDeserialize > :: arrow_deserialize_internal (name) , a : < A as arrow2_convert :: deserialize :: ArrowDeserialize > :: arrow_deserialize_internal (a) , b : < B as arrow2_convert :: deserialize :: ArrowDeserialize > :: arrow_deserialize_internal (b) , }) | |
| } else { | |
| None | |
| } | |
| } | |
| #[inline] | |
| fn consume_next(&mut self) { | |
| let _ = self.fields.name.next(); | |
| let _ = self.fields.a.next(); | |
| let _ = self.fields.b.next(); | |
| } | |
| } | |
| impl<'a, A, B> Iterator for FooArrayIterator<'a, A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| { | |
| type Item = Option<Foo<A, B>>; | |
| #[inline] | |
| fn next(&mut self) -> Option<Self::Item> { | |
| if !self.has_validity { | |
| self.return_next().map(|y| Some(y)) | |
| } else { | |
| let is_valid = self.validity_iter.next(); | |
| is_valid.map(|x| { | |
| if x { | |
| self.return_next() | |
| } else { | |
| self.consume_next(); | |
| None | |
| } | |
| }) | |
| } | |
| } | |
| } | |
| impl<A, B> arrow2_convert::deserialize::ArrowDeserialize for Foo<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| { | |
| type ArrayType = FooArray<A, B>; | |
| #[inline] | |
| fn arrow_deserialize<'a>(v: Option<Self>) -> Option<Self> { | |
| v | |
| } | |
| } | |
| #[test] | |
| fn test_simple_roundtrip() { | |
| // an item | |
| let original_array = [ | |
| Foo { | |
| name: "hello".to_string(), | |
| a: 1_i32, | |
| b: 2.0_f64, | |
| }, | |
| Foo { | |
| name: "one more".to_string(), | |
| a: 2_i32, | |
| b: 3.0_f64, | |
| }, | |
| Foo { | |
| name: "good bye".to_string(), | |
| a: 3_i32, | |
| b: 4.0_f64, | |
| }, | |
| ]; | |
| // serialize to an arrow array. try_into_arrow() is enabled by the TryIntoArrow trait | |
| let arrow_array: Box<dyn Array> = original_array.try_into_arrow().unwrap(); | |
| // which can be cast to an Arrow StructArray and be used for all kinds of IPC, FFI, etc. | |
| // supported by `arrow2` | |
| let struct_array = arrow_array | |
| .as_any() | |
| .downcast_ref::<arrow2::array::StructArray>() | |
| .unwrap(); | |
| assert_eq!(struct_array.len(), 3); | |
| // deserialize back to our original vector via TryIntoCollection trait. | |
| let round_trip_array: Vec<Foo<i32, f64>> = arrow_array.try_into_collection().unwrap(); | |
| assert_eq!(round_trip_array, original_array); | |
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| /// Generics example | |
| use arrow2::array::Array; | |
| use arrow2_convert::{deserialize::TryIntoCollection, serialize::TryIntoArrow}; | |
| #[derive(Debug, Clone, PartialEq, Eq)] | |
| pub struct Foo<A, B> | |
| where | |
| A: Clone, | |
| { | |
| name: String, | |
| a: A, | |
| b: B, | |
| } | |
| impl<A, B> arrow2_convert::field::ArrowField for Foo<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::field::ArrowField, | |
| B: ::arrow2_convert::field::ArrowField, | |
| { | |
| type Type = Self; | |
| fn data_type() -> arrow2::datatypes::DataType { | |
| arrow2::datatypes::DataType::Struct(<[_]>::into_vec(Box::new([ | |
| <String as arrow2_convert::field::ArrowField>::field("name"), | |
| <A as arrow2_convert::field::ArrowField>::field("a"), | |
| <B as arrow2_convert::field::ArrowField>::field("b"), | |
| ]))) | |
| } | |
| } | |
| impl<A, B> arrow2_convert::field::ArrowEnableVecForType for Foo<A, B> where A: Clone {} | |
| struct MutableFooArrayFields<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| name: <String as arrow2_convert::serialize::ArrowSerialize>::MutableArrayType, | |
| a: <A as arrow2_convert::serialize::ArrowSerialize>::MutableArrayType, | |
| b: <B as arrow2_convert::serialize::ArrowSerialize>::MutableArrayType, | |
| } | |
| pub struct MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fields: MutableFooArrayFields<A, B>, | |
| data_type: arrow2::datatypes::DataType, | |
| validity: Option<arrow2::bitmap::MutableBitmap>, | |
| } | |
| #[automatically_derived] | |
| impl<A, B> ::core::fmt::Debug for MutableFooArrayFields<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result { | |
| f.debug_struct("MutableFooArrayFields") | |
| .field("name", &self.name) | |
| .field("a", &self.a) | |
| .field("b", &self.b) | |
| .finish() | |
| } | |
| } | |
| #[automatically_derived] | |
| impl<A, B> ::core::fmt::Debug for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result { | |
| f.debug_struct("MutableFooArray") | |
| .field("fields", &self.fields) | |
| .field("data_type", &self.data_type) | |
| .field("validity", &self.validity) | |
| .finish() | |
| } | |
| } | |
| impl<A, B> MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| pub fn new() -> Self { | |
| Self { | |
| fields: MutableFooArrayFields { | |
| name: <String as arrow2_convert::serialize::ArrowSerialize>::new_array(), | |
| a: <A as arrow2_convert::serialize::ArrowSerialize>::new_array(), | |
| b: <B as arrow2_convert::serialize::ArrowSerialize>::new_array(), | |
| }, | |
| data_type: <Foo<A, B> as arrow2_convert::field::ArrowField>::data_type(), | |
| validity: None, | |
| } | |
| } | |
| fn init_validity(&mut self) { | |
| let mut validity = arrow2::bitmap::MutableBitmap::new(); | |
| validity.extend_constant(<Self as arrow2::array::MutableArray>::len(self), true); | |
| validity.set(<Self as arrow2::array::MutableArray>::len(self) - 1, false); | |
| self.validity = Some(validity) | |
| } | |
| } | |
| impl<A, B> Default for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn default() -> Self { | |
| Self::new() | |
| } | |
| } | |
| impl<__T: std::borrow::Borrow<Foo<A, B>>, A, B> arrow2::array::TryPush<Option<__T>> | |
| for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn try_push(&mut self, item: Option<__T>) -> arrow2::error::Result<()> { | |
| use arrow2::array::MutableArray; | |
| use std::borrow::Borrow; | |
| match item { | |
| Some(i) => { | |
| let i = i.borrow(); | |
| <String as arrow2_convert::serialize::ArrowSerialize>::arrow_serialize( | |
| i.name.borrow(), | |
| &mut self.fields.name, | |
| )?; | |
| <A as arrow2_convert::serialize::ArrowSerialize>::arrow_serialize( | |
| i.a.borrow(), | |
| &mut self.fields.a, | |
| )?; | |
| <B as arrow2_convert::serialize::ArrowSerialize>::arrow_serialize( | |
| i.b.borrow(), | |
| &mut self.fields.b, | |
| )?; | |
| match &mut self.validity { | |
| Some(validity) => validity.push(true), | |
| None => {} | |
| } | |
| } | |
| None => { | |
| < < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as MutableArray > :: push_null (& mut self . fields . name) ; | |
| < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as MutableArray > :: push_null (& mut self . fields . a) ; | |
| < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as MutableArray > :: push_null (& mut self . fields . b) ; | |
| match &mut self.validity { | |
| Some(validity) => validity.push(false), | |
| None => { | |
| self.init_validity(); | |
| } | |
| } | |
| } | |
| } | |
| Ok(()) | |
| } | |
| } | |
| impl<__T: std::borrow::Borrow<Foo<A, B>>, A, B> arrow2::array::TryExtend<Option<__T>> | |
| for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn try_extend<I: IntoIterator<Item = Option<__T>>>( | |
| &mut self, | |
| iter: I, | |
| ) -> arrow2::error::Result<()> { | |
| use arrow2::array::TryPush; | |
| for i in iter { | |
| self.try_push(i)?; | |
| } | |
| Ok(()) | |
| } | |
| } | |
| impl<A, B> arrow2::array::MutableArray for MutableFooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| fn data_type(&self) -> &arrow2::datatypes::DataType { | |
| &self.data_type | |
| } | |
| fn len(&self) -> usize { | |
| self.fields.name.len() | |
| } | |
| fn validity(&self) -> Option<&arrow2::bitmap::MutableBitmap> { | |
| self.validity.as_ref() | |
| } | |
| fn as_box(&mut self) -> Box<dyn arrow2::array::Array> { | |
| let values = < [_] > :: into_vec (Box :: new ([< < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . name) , < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . a) , < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . b)])) ; | |
| Box::new(arrow2::array::StructArray::from_data( | |
| <Foo<A, B> as arrow2_convert::field::ArrowField>::data_type().clone(), | |
| values, | |
| std::mem::take(&mut self.validity).map(|x| x.into()), | |
| )) | |
| } | |
| fn as_arc(&mut self) -> std::sync::Arc<dyn arrow2::array::Array> { | |
| let values = < [_] > :: into_vec (Box :: new ([< < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . name) , < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . a) , < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: as_box (& mut self . fields . b)])) ; | |
| std::sync::Arc::new(arrow2::array::StructArray::from_data( | |
| <Foo<A, B> as arrow2_convert::field::ArrowField>::data_type().clone(), | |
| values, | |
| std::mem::take(&mut self.validity).map(|x| x.into()), | |
| )) | |
| } | |
| fn as_any(&self) -> &dyn std::any::Any { | |
| self | |
| } | |
| fn as_mut_any(&mut self) -> &mut dyn std::any::Any { | |
| self | |
| } | |
| fn push_null(&mut self) { | |
| use arrow2::array::TryPush; | |
| self.try_push(None::<Foo<A, B>>).unwrap(); | |
| } | |
| fn shrink_to_fit(&mut self) { | |
| < < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: shrink_to_fit (& mut self . fields . name) ; | |
| < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: shrink_to_fit (& mut self . fields . a) ; | |
| < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: shrink_to_fit (& mut self . fields . b) ; | |
| if let Some(validity) = &mut self.validity { | |
| validity.shrink_to_fit(); | |
| } | |
| } | |
| fn reserve(&mut self, additional: usize) { | |
| if let Some(x) = self.validity.as_mut() { | |
| x.reserve(additional) | |
| } | |
| < < String as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: reserve (& mut self . fields . name , additional) ; | |
| < < A as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: reserve (& mut self . fields . a , additional) ; | |
| < < B as arrow2_convert :: serialize :: ArrowSerialize > :: MutableArrayType as arrow2 :: array :: MutableArray > :: reserve (& mut self . fields . b , additional) ; | |
| } | |
| } | |
| impl<A, B> arrow2_convert::serialize::ArrowSerialize for Foo<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::serialize::ArrowSerialize, | |
| B: ::arrow2_convert::serialize::ArrowSerialize, | |
| { | |
| type MutableArrayType = MutableFooArray<A, B>; | |
| #[inline] | |
| fn new_array() -> Self::MutableArrayType { | |
| Self::MutableArrayType::default() | |
| } | |
| #[inline] | |
| fn arrow_serialize(v: &Self, array: &mut Self::MutableArrayType) -> arrow2::error::Result<()> { | |
| use arrow2::array::TryPush; | |
| array.try_push(Some(v)) | |
| } | |
| } | |
| pub struct FooArray<A, B> { | |
| name: std::marker::PhantomData<String>, | |
| a: std::marker::PhantomData<A>, | |
| b: std::marker::PhantomData<B>, | |
| } | |
| impl<A, B> arrow2_convert::deserialize::ArrowArray for FooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <A as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <B as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| { | |
| type BaseArrayType = arrow2::array::StructArray; | |
| #[inline] | |
| fn iter_from_array_ref<'a>( | |
| b: &'a dyn arrow2::array::Array, | |
| ) -> <&'a Self as IntoIterator>::IntoIter { | |
| use core::ops::Deref; | |
| let arr = b | |
| .as_any() | |
| .downcast_ref::<arrow2::array::StructArray>() | |
| .unwrap(); | |
| let values = arr.values(); | |
| let validity = arr.validity(); | |
| FooArrayIterator { fields : FooArrayIteratorFields { name : < < String as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as arrow2_convert :: deserialize :: ArrowArray > :: iter_from_array_ref (values [0] . deref ()) , a : < < A as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as arrow2_convert :: deserialize :: ArrowArray > :: iter_from_array_ref (values [1] . deref ()) , b : < < B as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as arrow2_convert :: deserialize :: ArrowArray > :: iter_from_array_ref (values [2] . deref ()) , } , has_validity : validity . as_ref () . is_some () , validity_iter : validity . as_ref () . map (| x | x . iter ()) . unwrap_or_else (| | arrow2 :: bitmap :: utils :: BitmapIter :: new (& [] , 0 , 0)) , } | |
| } | |
| } | |
| impl<'a, A, B> IntoIterator for &'a FooArray<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <A as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <B as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| { | |
| type Item = Option<Foo<A, B>>; | |
| type IntoIter = FooArrayIterator<'a, A, B>; | |
| fn into_iter(self) -> Self::IntoIter { | |
| unimplemented!(); | |
| } | |
| } | |
| struct FooArrayIteratorFields < 'a , A , B > where A : Clone , A : :: arrow2_convert :: deserialize :: ArrowDeserialize , for < '_a > & '_a < A as :: arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType : IntoIterator , B : :: arrow2_convert :: deserialize :: ArrowDeserialize , for < '_a > & '_a < B as :: arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType : IntoIterator { name : < & 'a < String as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as IntoIterator > :: IntoIter , a : < & 'a < A as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as IntoIterator > :: IntoIter , b : < & 'a < B as arrow2_convert :: deserialize :: ArrowDeserialize > :: ArrayType as IntoIterator > :: IntoIter , } | |
| pub struct FooArrayIterator<'a, A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <A as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <B as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| { | |
| fields: FooArrayIteratorFields<'a, A, B>, | |
| validity_iter: arrow2::bitmap::utils::BitmapIter<'a>, | |
| has_validity: bool, | |
| } | |
| impl<'a, A, B> FooArrayIterator<'a, A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <A as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <B as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| { | |
| #[inline] | |
| fn return_next(&mut self) -> Option<Foo<A, B>> { | |
| if let (Some(name), Some(a), Some(b)) = ( | |
| self.fields.name.next(), | |
| self.fields.a.next(), | |
| self.fields.b.next(), | |
| ) { | |
| Some (Foo { name : < String as arrow2_convert :: deserialize :: ArrowDeserialize > :: arrow_deserialize_internal (name) , a : < A as arrow2_convert :: deserialize :: ArrowDeserialize > :: arrow_deserialize_internal (a) , b : < B as arrow2_convert :: deserialize :: ArrowDeserialize > :: arrow_deserialize_internal (b) , }) | |
| } else { | |
| None | |
| } | |
| } | |
| #[inline] | |
| fn consume_next(&mut self) { | |
| let _ = self.fields.name.next(); | |
| let _ = self.fields.a.next(); | |
| let _ = self.fields.b.next(); | |
| } | |
| } | |
| impl<'a, A, B> Iterator for FooArrayIterator<'a, A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <A as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <B as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| { | |
| type Item = Option<Foo<A, B>>; | |
| #[inline] | |
| fn next(&mut self) -> Option<Self::Item> { | |
| if !self.has_validity { | |
| self.return_next().map(|y| Some(y)) | |
| } else { | |
| let is_valid = self.validity_iter.next(); | |
| is_valid.map(|x| { | |
| if x { | |
| self.return_next() | |
| } else { | |
| self.consume_next(); | |
| None | |
| } | |
| }) | |
| } | |
| } | |
| } | |
| impl<A, B> arrow2_convert::deserialize::ArrowDeserialize for Foo<A, B> | |
| where | |
| A: Clone, | |
| A: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <A as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| B: ::arrow2_convert::deserialize::ArrowDeserialize, | |
| for<'_a> &'_a <B as ::arrow2_convert::deserialize::ArrowDeserialize>::ArrayType: IntoIterator, | |
| { | |
| type ArrayType = FooArray<A, B>; | |
| #[inline] | |
| fn arrow_deserialize<'a>(v: Option<Self>) -> Option<Self> { | |
| v | |
| } | |
| } | |
| #[test] | |
| fn test_simple_roundtrip() { | |
| // an item | |
| let original_array = [ | |
| Foo { | |
| name: "hello".to_string(), | |
| a: 1_i32, | |
| b: 2.0_f64, | |
| }, | |
| Foo { | |
| name: "one more".to_string(), | |
| a: 2_i32, | |
| b: 3.0_f64, | |
| }, | |
| Foo { | |
| name: "good bye".to_string(), | |
| a: 3_i32, | |
| b: 4.0_f64, | |
| }, | |
| ]; | |
| // serialize to an arrow array. try_into_arrow() is enabled by the TryIntoArrow trait | |
| let arrow_array: Box<dyn Array> = original_array.try_into_arrow().unwrap(); | |
| // which can be cast to an Arrow StructArray and be used for all kinds of IPC, FFI, etc. | |
| // supported by `arrow2` | |
| let struct_array = arrow_array | |
| .as_any() | |
| .downcast_ref::<arrow2::array::StructArray>() | |
| .unwrap(); | |
| assert_eq!(struct_array.len(), 3); | |
| // deserialize back to our original vector via TryIntoCollection trait. | |
| let round_trip_array: Vec<Foo<i32, f64>> = arrow_array.try_into_collection().unwrap(); | |
| assert_eq!(round_trip_array, original_array); | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment