iridis_message/
traits.rs

1use std::sync::Arc;
2
3use crate::prelude::*;
4
5use arrow_array::*;
6use arrow_data::*;
7use arrow_schema::*;
8
9/// Trait to represent a message that can be converted to and from an Arrow array.
10pub trait ArrowMessage {
11    /// Each message must declare how it's represented as an Arrow field.
12    fn field(name: impl Into<String>) -> Field;
13
14    /// Try to convert the message into an Arrow array.
15    fn try_into_arrow(self) -> Result<ArrayRef>;
16
17    /// Try to convert an Arrow array into the message.
18    fn try_from_arrow(data: ArrayData) -> Result<Self>
19    where
20        Self: Sized;
21}
22
23impl<T> ArrowMessage for Option<T>
24where
25    T: ArrowMessage,
26{
27    fn field(name: impl Into<String>) -> Field {
28        T::field(name).with_nullable(true)
29    }
30
31    fn try_into_arrow(self) -> Result<ArrayRef> {
32        match self {
33            Some(value) => value.try_into_arrow(),
34            None => Ok(Arc::new(NullArray::new(0)) as ArrayRef),
35        }
36    }
37
38    fn try_from_arrow(data: ArrayData) -> Result<Self>
39    where
40        Self: Sized,
41    {
42        match data.data_type() {
43            DataType::Null => Ok(None),
44            _ => T::try_from_arrow(data).map(|value| Some(value)),
45        }
46    }
47}
48
49impl<T> ArrowMessage for PrimitiveArray<T>
50where
51    T: ArrowPrimitiveType,
52{
53    fn field(name: impl Into<String>) -> Field {
54        Field::new(name, T::DATA_TYPE, false)
55    }
56
57    fn try_into_arrow(self) -> Result<ArrayRef> {
58        Ok(Arc::new(self) as ArrayRef)
59    }
60
61    fn try_from_arrow(data: ArrayData) -> Result<Self>
62    where
63        Self: Sized,
64    {
65        Ok(PrimitiveArray::from(data))
66    }
67}
68
69macro_rules! impl_arrow_field {
70    ($type:ty, $data_type:expr, $array_type:ty) => {
71        impl ArrowMessage for $type {
72            fn field(name: impl Into<String>) -> Field {
73                Field::new(name, $data_type, false)
74            }
75
76            fn try_into_arrow(self) -> Result<ArrayRef> {
77                Ok(Arc::new(<$array_type>::from(vec![self])) as ArrayRef)
78            }
79
80            fn try_from_arrow(data: ArrayData) -> Result<Self>
81            where
82                Self: Sized,
83            {
84                let array = <$array_type>::from(data);
85                match array.len() {
86                    0 => Err(eyre::eyre!("Array is empty")),
87                    _ => Ok(array.value(0)),
88                }
89            }
90        }
91    };
92}
93
94impl_arrow_field!(u8, DataType::UInt8, UInt8Array);
95impl_arrow_field!(u16, DataType::UInt16, UInt16Array);
96impl_arrow_field!(u32, DataType::UInt32, UInt32Array);
97impl_arrow_field!(u64, DataType::UInt64, UInt64Array);
98impl_arrow_field!(i8, DataType::Int8, Int8Array);
99impl_arrow_field!(i16, DataType::Int16, Int16Array);
100impl_arrow_field!(i32, DataType::Int32, Int32Array);
101impl_arrow_field!(i64, DataType::Int64, Int64Array);
102impl_arrow_field!(f32, DataType::Float32, Float32Array);
103impl_arrow_field!(f64, DataType::Float64, Float64Array);
104impl_arrow_field!(bool, DataType::Boolean, BooleanArray);
105
106impl ArrowMessage for String {
107    fn field(name: impl Into<String>) -> Field {
108        Field::new(name, DataType::Utf8, false)
109    }
110
111    fn try_into_arrow(self) -> Result<ArrayRef> {
112        Ok(Arc::new(StringArray::from(vec![self])) as ArrayRef)
113    }
114
115    fn try_from_arrow(data: ArrayData) -> Result<Self>
116    where
117        Self: Sized,
118    {
119        let array = StringArray::from(data);
120        match array.len() {
121            0 => Err(eyre::eyre!("Array is empty")),
122            _ => Ok(array.value(0).into()),
123        }
124    }
125}
126
127impl ArrowMessage for StringArray {
128    fn field(name: impl Into<String>) -> Field {
129        Field::new(name, DataType::Utf8, false)
130    }
131
132    fn try_into_arrow(self) -> Result<ArrayRef> {
133        Ok(Arc::new(self) as ArrayRef)
134    }
135
136    fn try_from_arrow(data: ArrayData) -> Result<Self>
137    where
138        Self: Sized,
139    {
140        Ok(StringArray::from(data))
141    }
142}
iridis_message/traits.rs

iridis_message/
traits.rs
