Skip to content

Instantly share code, notes, and snippets.

@aldanor

aldanor/test_generic_no_bound.rs

Created October 8, 2022 01:14
Show Gist options
  • Save aldanor/de7accf7bb40e83bec6e1e51f08f1190 to your computer and use it in GitHub Desktop.
Save aldanor/de7accf7bb40e83bec6e1e51f08f1190 to your computer and use it in GitHub Desktop.
Download ZIP
arrow2-convert-generics
Raw
test_generic_no_bound.rs
/// 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);
}
Raw
test_generics_with_bound.rs
/// 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