dora_arrow_convert/
from_impls.rs

1use arrow::{
2    array::{Array, AsArray, PrimitiveArray, StringArray},
3    datatypes::ArrowPrimitiveType,
4};
5use eyre::ContextCompat;
6
7use crate::ArrowData;
8
9impl From<ArrowData> for arrow::array::ArrayRef {
10    fn from(value: ArrowData) -> Self {
11        value.0
12    }
13}
14
15impl From<arrow::array::ArrayRef> for ArrowData {
16    fn from(value: arrow::array::ArrayRef) -> Self {
17        Self(value)
18    }
19}
20
21impl TryFrom<&ArrowData> for bool {
22    type Error = eyre::Report;
23    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
24        let bool_array = value.as_boolean_opt().context("not a bool array")?;
25        if bool_array.is_empty() {
26            eyre::bail!("empty array");
27        }
28        if bool_array.len() != 1 {
29            eyre::bail!("expected length 1");
30        }
31        if bool_array.null_count() != 0 {
32            eyre::bail!("bool array has nulls");
33        }
34        Ok(bool_array.value(0))
35    }
36}
37impl TryFrom<&ArrowData> for u8 {
38    type Error = eyre::Report;
39    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
40        let array = value
41            .as_primitive_opt::<arrow::datatypes::UInt8Type>()
42            .context("not a primitive UInt8Type array")?;
43        extract_single_primitive(array)
44    }
45}
46impl TryFrom<&ArrowData> for u16 {
47    type Error = eyre::Report;
48    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
49        let array = value
50            .as_primitive_opt::<arrow::datatypes::UInt16Type>()
51            .context("not a primitive UInt16Type array")?;
52        extract_single_primitive(array)
53    }
54}
55impl TryFrom<&ArrowData> for u32 {
56    type Error = eyre::Report;
57    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
58        let array = value
59            .as_primitive_opt::<arrow::datatypes::UInt32Type>()
60            .context("not a primitive UInt32Type array")?;
61        extract_single_primitive(array)
62    }
63}
64impl TryFrom<&ArrowData> for u64 {
65    type Error = eyre::Report;
66    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
67        let array = value
68            .as_primitive_opt::<arrow::datatypes::UInt64Type>()
69            .context("not a primitive UInt64Type array")?;
70        extract_single_primitive(array)
71    }
72}
73impl TryFrom<&ArrowData> for i8 {
74    type Error = eyre::Report;
75    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
76        let array = value
77            .as_primitive_opt::<arrow::datatypes::Int8Type>()
78            .context("not a primitive Int8Type array")?;
79        extract_single_primitive(array)
80    }
81}
82impl TryFrom<&ArrowData> for i16 {
83    type Error = eyre::Report;
84    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
85        let array = value
86            .as_primitive_opt::<arrow::datatypes::Int16Type>()
87            .context("not a primitive Int16Type array")?;
88        extract_single_primitive(array)
89    }
90}
91impl TryFrom<&ArrowData> for i32 {
92    type Error = eyre::Report;
93    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
94        let array = value
95            .as_primitive_opt::<arrow::datatypes::Int32Type>()
96            .context("not a primitive Int32Type array")?;
97        extract_single_primitive(array)
98    }
99}
100impl TryFrom<&ArrowData> for i64 {
101    type Error = eyre::Report;
102    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
103        let array = value
104            .as_primitive_opt::<arrow::datatypes::Int64Type>()
105            .context("not a primitive Int64Type array")?;
106        extract_single_primitive(array)
107    }
108}
109
110impl TryFrom<&ArrowData> for f32 {
111    type Error = eyre::Report;
112    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
113        let array = value
114            .as_primitive_opt::<arrow::datatypes::Float32Type>()
115            .context("not a primitive Float32Type array")?;
116        extract_single_primitive(array)
117    }
118}
119impl TryFrom<&ArrowData> for f64 {
120    type Error = eyre::Report;
121    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
122        let array = value
123            .as_primitive_opt::<arrow::datatypes::Float64Type>()
124            .context("not a primitive Float64Type array")?;
125        extract_single_primitive(array)
126    }
127}
128
129impl<'a> TryFrom<&'a ArrowData> for &'a str {
130    type Error = eyre::Report;
131    fn try_from(value: &'a ArrowData) -> Result<Self, Self::Error> {
132        let array: &StringArray = value.as_string_opt().wrap_err("not a string array")?;
133        if array.is_empty() {
134            eyre::bail!("empty array");
135        }
136        if array.len() != 1 {
137            eyre::bail!("expected length 1");
138        }
139        if array.null_count() != 0 {
140            eyre::bail!("array has nulls");
141        }
142        Ok(array.value(0))
143    }
144}
145
146impl<'a> TryFrom<&'a ArrowData> for &'a [u8] {
147    type Error = eyre::Report;
148    fn try_from(value: &'a ArrowData) -> Result<Self, Self::Error> {
149        let array: &PrimitiveArray<arrow::datatypes::UInt8Type> = value
150            .as_primitive_opt()
151            .wrap_err("not a primitive UInt8Type array")?;
152        if array.null_count() != 0 {
153            eyre::bail!("array has nulls");
154        }
155        Ok(array.values())
156    }
157}
158
159impl<'a> TryFrom<&'a ArrowData> for Vec<u8> {
160    type Error = eyre::Report;
161    fn try_from(value: &'a ArrowData) -> Result<Self, Self::Error> {
162        value.try_into().map(|slice: &'a [u8]| slice.to_vec())
163    }
164}
165
166fn extract_single_primitive<T>(array: &PrimitiveArray<T>) -> Result<T::Native, eyre::Error>
167where
168    T: ArrowPrimitiveType,
169{
170    if array.is_empty() {
171        eyre::bail!("empty array");
172    }
173    if array.len() != 1 {
174        eyre::bail!("expected length 1");
175    }
176    if array.null_count() != 0 {
177        eyre::bail!("array has nulls");
178    }
179    Ok(array.value(0))
180}
181
182#[cfg(test)]
183mod tests {
184    use arrow::array::{make_array, PrimitiveArray};
185
186    use crate::ArrowData;
187
188    #[test]
189    fn test_u8() {
190        let array =
191            make_array(PrimitiveArray::<arrow::datatypes::UInt8Type>::from(vec![42]).into());
192        let data: ArrowData = array.into();
193        let value: u8 = (&data).try_into().unwrap();
194        assert_eq!(value, 42);
195    }
196}
dora_arrow_convert/from_impls.rs

dora_arrow_convert/
from_impls.rs
