flarrow_message/traits/
message.rs

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

flarrow_message/traits/
message.rs
