iridis_message/
traits.rs

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

iridis_message/
traits.rs
