dora_arrow_convert/
from_impls.rs

1use arrow::{
2    array::{Array, AsArray, PrimitiveArray, StringArray},
3    datatypes::{ArrowPrimitiveType, ArrowTemporalType},
4};
5use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
6use eyre::ContextCompat;
7use half::f16;
8
9use crate::ArrowData;
10
11impl From<ArrowData> for arrow::array::ArrayRef {
12    fn from(value: ArrowData) -> Self {
13        value.0
14    }
15}
16
17impl From<arrow::array::ArrayRef> for ArrowData {
18    fn from(value: arrow::array::ArrayRef) -> Self {
19        Self(value)
20    }
21}
22
23impl TryFrom<&ArrowData> for bool {
24    type Error = eyre::Report;
25    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
26        let bool_array = value.as_boolean_opt().context("not a bool array")?;
27        if bool_array.is_empty() {
28            eyre::bail!("empty array");
29        }
30        if bool_array.len() != 1 {
31            eyre::bail!("expected length 1");
32        }
33        if bool_array.null_count() != 0 {
34            eyre::bail!("bool array has nulls");
35        }
36        Ok(bool_array.value(0))
37    }
38}
39
40macro_rules! impl_try_from_arrow_data {
41    ($($t:ty => $arrow_type:ident),*) => {
42        $(
43            impl TryFrom<&ArrowData> for $t {
44                type Error = eyre::Report;
45
46                fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
47                    let array = value
48                        .as_primitive_opt::<arrow::datatypes::$arrow_type>()
49                        .context(concat!("not a primitive ", stringify!($arrow_type), " array"))?;
50                    extract_single_primitive(array)
51                }
52            }
53        )*
54
55        $(
56            impl<'a> TryFrom<&'a ArrowData> for &'a [$t] {
57                type Error = eyre::Report;
58
59                fn try_from(value: &'a ArrowData) -> Result<Self, Self::Error> {
60                    let array: &PrimitiveArray<arrow::datatypes::$arrow_type> = value
61                        .as_primitive_opt()
62                        .wrap_err(concat!("not a primitive ", stringify!($arrow_type), " array"))?;
63                    if array.null_count() != 0 {
64                        eyre::bail!("array has nulls");
65                    }
66                    Ok(array.values())
67                }
68            }
69        )*
70
71        $(
72            impl<'a> TryFrom<&'a ArrowData> for Vec<$t> {
73                type Error = eyre::Report;
74
75                fn try_from(value: &'a ArrowData) -> Result<Self, Self::Error> {
76                    value
77                        .try_into()
78                        .map(|slice: &'a [$t]| slice.to_vec())
79                }
80            }
81        )*
82    };
83}
84
85impl_try_from_arrow_data!(
86    u8 => UInt8Type,
87    u16 => UInt16Type,
88    u32 => UInt32Type,
89    u64 => UInt64Type,
90    i8 => Int8Type,
91    i16 => Int16Type,
92    i32 => Int32Type,
93    i64 => Int64Type,
94    f16 => Float16Type,
95    f32 => Float32Type,
96    f64 => Float64Type
97);
98
99impl<'a> TryFrom<&'a ArrowData> for &'a str {
100    type Error = eyre::Report;
101    fn try_from(value: &'a ArrowData) -> Result<Self, Self::Error> {
102        let array: &StringArray = value.as_string_opt().wrap_err("not a string array")?;
103        if array.is_empty() {
104            eyre::bail!("empty array");
105        }
106        if array.len() != 1 {
107            eyre::bail!("expected length 1");
108        }
109        if array.null_count() != 0 {
110            eyre::bail!("array has nulls");
111        }
112        Ok(array.value(0))
113    }
114}
115
116impl TryFrom<&ArrowData> for String {
117    type Error = eyre::Report;
118    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
119        // Delegate to the `&str` impl so the single-element validation lives in
120        // one place and the two conversions can never drift apart.
121        let s: &str = value.try_into()?;
122        Ok(s.to_string())
123    }
124}
125
126impl TryFrom<&ArrowData> for NaiveDate {
127    type Error = eyre::Report;
128    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
129        if let Some(array) = value.as_any().downcast_ref::<arrow::array::Date32Array>() {
130            if check_single_datetime(array) {
131                eyre::bail!("Not a valid array");
132            }
133            return array
134                .value_as_date(0)
135                .context("data type cannot be converted to NaiveDate");
136        }
137        let array = value
138            .as_any()
139            .downcast_ref::<arrow::array::Date64Array>()
140            .context("Reference is neither to a Date32Array nor a Date64Array")?;
141        if check_single_datetime(array) {
142            eyre::bail!("Not a valid array");
143        }
144        array
145            .value_as_date(0)
146            .context("data type cannot be converted to NaiveDate")
147    }
148}
149
150impl TryFrom<&ArrowData> for NaiveTime {
151    type Error = eyre::Report;
152    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
153        if let Some(array) = value
154            .as_any()
155            .downcast_ref::<arrow::array::Time32SecondArray>()
156        {
157            if check_single_datetime(array) {
158                eyre::bail!("Not a valid array");
159            }
160            return array
161                .value_as_time(0)
162                .context("data type cannot be converted to NaiveTime");
163        }
164        if let Some(array) = value
165            .as_any()
166            .downcast_ref::<arrow::array::Time32MillisecondArray>()
167        {
168            if check_single_datetime(array) {
169                eyre::bail!("Not a valid array");
170            }
171            return array
172                .value_as_time(0)
173                .context("data type cannot be converted to NaiveTime");
174        }
175        if let Some(array) = value
176            .as_any()
177            .downcast_ref::<arrow::array::Time64MicrosecondArray>()
178        {
179            if check_single_datetime(array) {
180                eyre::bail!("Not a valid array");
181            }
182            return array
183                .value_as_time(0)
184                .context("data type cannot be converted to NaiveTime");
185        }
186        let array = value
187            .as_primitive_opt::<arrow::datatypes::Time64NanosecondType>()
188            .context("not any of the primitive Time arrays")?;
189        if check_single_datetime(array) {
190            eyre::bail!("Not a valid array");
191        }
192        array
193            .value_as_time(0)
194            .context("data type cannot be converted to NaiveTime")
195    }
196}
197
198impl TryFrom<&ArrowData> for NaiveDateTime {
199    type Error = eyre::Report;
200    fn try_from(value: &ArrowData) -> Result<Self, Self::Error> {
201        if let Some(array) = value
202            .as_any()
203            .downcast_ref::<arrow::array::TimestampSecondArray>()
204        {
205            if check_single_datetime(array) {
206                eyre::bail!("Not a valid array");
207            }
208            return array
209                .value_as_datetime(0)
210                .context("data type cannot be converted to NaiveDateTime");
211        }
212        if let Some(array) = value
213            .as_any()
214            .downcast_ref::<arrow::array::TimestampMillisecondArray>()
215        {
216            if check_single_datetime(array) {
217                eyre::bail!("Not a valid array");
218            }
219            return array
220                .value_as_datetime(0)
221                .context("data type cannot be converted to NaiveDateTime");
222        }
223        if let Some(array) = value
224            .as_any()
225            .downcast_ref::<arrow::array::TimestampMicrosecondArray>()
226        {
227            if check_single_datetime(array) {
228                eyre::bail!("Not a valid array");
229            }
230            return array
231                .value_as_datetime(0)
232                .context("data type cannot be converted to NaiveDateTime");
233        }
234        let array = value
235            .as_primitive_opt::<arrow::datatypes::TimestampNanosecondType>()
236            .context("not any of the primitive Time arrays")?;
237        if check_single_datetime(array) {
238            eyre::bail!("Not a valid array");
239        }
240        array
241            .value_as_datetime(0)
242            .context("data type cannot be converted to NaiveDateTime")
243    }
244}
245
246fn check_single_datetime<T>(array: &PrimitiveArray<T>) -> bool
247where
248    T: ArrowTemporalType,
249{
250    array.is_empty() || array.len() != 1 || array.null_count() != 0
251}
252fn extract_single_primitive<T>(array: &PrimitiveArray<T>) -> Result<T::Native, eyre::Error>
253where
254    T: ArrowPrimitiveType,
255{
256    if array.is_empty() {
257        eyre::bail!("empty array");
258    }
259    if array.len() != 1 {
260        eyre::bail!("expected length 1");
261    }
262    if array.null_count() != 0 {
263        eyre::bail!("array has nulls");
264    }
265    Ok(array.value(0))
266}
267
268#[cfg(test)]
269mod tests {
270    use arrow::array::{PrimitiveArray, make_array};
271
272    use crate::ArrowData;
273
274    #[test]
275    fn test_u8() {
276        let array =
277            make_array(PrimitiveArray::<arrow::datatypes::UInt8Type>::from(vec![42]).into());
278        let data: ArrowData = array.into();
279        let value: u8 = (&data).try_into().unwrap();
280        assert_eq!(value, 42);
281    }
282}
dora_arrow_convert/from_impls.rs

dora_arrow_convert/
from_impls.rs
