enum_impl/
enum_impl.rs

1use arrow::array::*;
2use arrow_message::prelude::*;
3
4#[derive(Debug)]
5enum Encoding {
6    RGB8,
7    RGBA8,
8    BGR8,
9    BGRA8,
10}
11
12impl Encoding {
13    pub fn into_string(self) -> String {
14        match self {
15            Encoding::RGB8 => "RGB8".to_string(),
16            Encoding::RGBA8 => "RGBA8".to_string(),
17            Encoding::BGR8 => "BGR8".to_string(),
18            Encoding::BGRA8 => "BGRA8".to_string(),
19        }
20    }
21
22    pub fn try_from_string(value: String) -> Result<Self, ArrowError> {
23        match value.as_str() {
24            "RGB8" => Ok(Encoding::RGB8),
25            "RGBA8" => Ok(Encoding::RGBA8),
26            "BGR8" => Ok(Encoding::BGR8),
27            "BGRA8" => Ok(Encoding::BGRA8),
28            _ => Err(ArrowError::ParseError(format!(
29                "Invalid encoding: {}",
30                value
31            ))),
32        }
33    }
34}
35
36impl ArrowMessage for Encoding {
37    fn field(name: impl Into<String>) -> Field {
38        String::field(name)
39    }
40
41    fn try_from_arrow(data: ArrayData) -> miette::Result<Self, ArrowError>
42    where
43        Self: Sized,
44    {
45        Encoding::try_from_string(String::try_from_arrow(data)?)
46    }
47
48    fn try_into_arrow(self) -> miette::Result<ArrayRef, ArrowError> {
49        String::try_into_arrow(self.into_string())
50    }
51}
52
53impl TryFrom<ArrayData> for Encoding {
54    type Error = ArrowError;
55
56    fn try_from(data: ArrayData) -> Result<Self, Self::Error> {
57        Encoding::try_from_arrow(data)
58    }
59}
60
61impl TryFrom<Encoding> for ArrayData {
62    type Error = ArrowError;
63
64    fn try_from(metadata: Encoding) -> Result<Self, Self::Error> {
65        metadata.try_into_arrow().map(|array| array.into_data())
66    }
67}
68
69#[derive(Debug)]
70struct Metadata {
71    name: Option<String>,
72    width: u32,
73    height: u32,
74    encoding: Encoding,
75}
76
77impl ArrowMessage for Metadata {
78    fn field(name: impl Into<String>) -> Field {
79        make_union_fields(
80            name,
81            vec![
82                Option::<String>::field("name"),
83                Option::<u32>::field("width"),
84                Option::<u32>::field("height"),
85                Encoding::field("encoding"),
86            ],
87        )
88    }
89
90    fn try_from_arrow(data: arrow::array::ArrayData) -> Result<Self, ArrowError>
91    where
92        Self: Sized,
93    {
94        let (map, children) = unpack_union(data);
95
96        Ok(Metadata {
97            name: extract_union_data("name", &map, &children)?,
98            width: extract_union_data("width", &map, &children)?,
99            height: extract_union_data("height", &map, &children)?,
100            encoding: extract_union_data("encoding", &map, &children)?,
101        })
102    }
103
104    fn try_into_arrow(self) -> Result<arrow::array::ArrayRef, ArrowError> {
105        let union_fields = get_union_fields::<Self>()?;
106
107        make_union_array(
108            union_fields,
109            vec![
110                self.name.try_into_arrow()?,
111                self.width.try_into_arrow()?,
112                self.height.try_into_arrow()?,
113                self.encoding.try_into_arrow()?,
114            ],
115        )
116    }
117}
118
119impl TryFrom<ArrayData> for Metadata {
120    type Error = ArrowError;
121
122    fn try_from(data: ArrayData) -> Result<Self, Self::Error> {
123        Metadata::try_from_arrow(data)
124    }
125}
126
127impl TryFrom<Metadata> for ArrayData {
128    type Error = ArrowError;
129
130    fn try_from(metadata: Metadata) -> Result<Self, Self::Error> {
131        metadata.try_into_arrow().map(|array| array.into_data())
132    }
133}
134
135#[derive(Debug)]
136struct Image {
137    data: UInt8Array,
138    metadata: Option<Metadata>,
139}
140
141impl ArrowMessage for Image {
142    fn field(name: impl Into<String>) -> Field {
143        make_union_fields(
144            name,
145            vec![
146                UInt8Array::field("data"),
147                Option::<Metadata>::field("metadata"),
148            ],
149        )
150    }
151
152    fn try_from_arrow(data: arrow::array::ArrayData) -> Result<Self, ArrowError>
153    where
154        Self: Sized,
155    {
156        let (map, children) = unpack_union(data);
157
158        Ok(Image {
159            data: extract_union_data("data", &map, &children)?,
160            metadata: extract_union_data("metadata", &map, &children)?,
161        })
162    }
163
164    fn try_into_arrow(self) -> Result<arrow::array::ArrayRef, ArrowError> {
165        let union_fields = get_union_fields::<Self>()?;
166
167        make_union_array(
168            union_fields,
169            vec![self.data.try_into_arrow()?, self.metadata.try_into_arrow()?],
170        )
171    }
172}
173
174impl TryFrom<ArrayData> for Image {
175    type Error = ArrowError;
176
177    fn try_from(data: ArrayData) -> Result<Self, Self::Error> {
178        Image::try_from_arrow(data)
179    }
180}
181
182impl TryFrom<Image> for ArrayData {
183    type Error = ArrowError;
184
185    fn try_from(image: Image) -> Result<Self, Self::Error> {
186        image.try_into_arrow().map(|array| array.into_data())
187    }
188}
189
190fn main() -> Result<()> {
191    use miette::IntoDiagnostic;
192
193    let image = Image {
194        data: UInt8Array::from(vec![1, 2, 3]),
195        metadata: Some(Metadata {
196            name: Some("example".to_string()),
197            width: 12,
198            height: 12,
199            encoding: Encoding::RGB8,
200        }),
201    };
202
203    println!("{:?}", image);
204
205    let arrow = ArrayData::try_from(image).into_diagnostic()?;
206    let image = Image::try_from(arrow).into_diagnostic()?;
207
208    println!("{:?}", image);
209
210    Ok(())
211}