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::{DoraArray, internal::array_ref};
10
11impl TryFrom<&DoraArray> for bool {
12    type Error = eyre::Report;
13    fn try_from(value: &DoraArray) -> Result<Self, Self::Error> {
14        let bool_array = array_ref(value)
15            .as_boolean_opt()
16            .context("not a bool array")?;
17        if bool_array.is_empty() {
18            eyre::bail!("empty array");
19        }
20        if bool_array.len() != 1 {
21            eyre::bail!("expected length 1");
22        }
23        if bool_array.null_count() != 0 {
24            eyre::bail!("bool array has nulls");
25        }
26        Ok(bool_array.value(0))
27    }
28}
29
30macro_rules! impl_try_from_arrow_data {
31    ($($t:ty => $arrow_type:ident),*) => {
32        $(
33            impl TryFrom<&DoraArray> for $t {
34                type Error = eyre::Report;
35
36                fn try_from(value: &DoraArray) -> Result<Self, Self::Error> {
37                    let array = array_ref(value).as_primitive_opt::<arrow::datatypes::$arrow_type>()
38                        .context(concat!("not a primitive ", stringify!($arrow_type), " array"))?;
39                    extract_single_primitive(array)
40                }
41            }
42        )*
43
44        $(
45            impl<'a> TryFrom<&'a DoraArray> for &'a [$t] {
46                type Error = eyre::Report;
47
48                fn try_from(value: &'a DoraArray) -> Result<Self, Self::Error> {
49                    let array: &PrimitiveArray<arrow::datatypes::$arrow_type> = array_ref(value).as_primitive_opt()
50                        .wrap_err(concat!("not a primitive ", stringify!($arrow_type), " array"))?;
51                    if array.null_count() != 0 {
52                        eyre::bail!("array has nulls");
53                    }
54                    Ok(array.values())
55                }
56            }
57        )*
58
59        $(
60            impl<'a> TryFrom<&'a DoraArray> for Vec<$t> {
61                type Error = eyre::Report;
62
63                fn try_from(value: &'a DoraArray) -> Result<Self, Self::Error> {
64                    value
65                        .try_into()
66                        .map(|slice: &'a [$t]| slice.to_vec())
67                }
68            }
69        )*
70    };
71}
72
73impl_try_from_arrow_data!(
74    u8 => UInt8Type,
75    u16 => UInt16Type,
76    u32 => UInt32Type,
77    u64 => UInt64Type,
78    i8 => Int8Type,
79    i16 => Int16Type,
80    i32 => Int32Type,
81    i64 => Int64Type,
82    f16 => Float16Type,
83    f32 => Float32Type,
84    f64 => Float64Type
85);
86
87impl<'a> TryFrom<&'a DoraArray> for &'a str {
88    type Error = eyre::Report;
89    fn try_from(value: &'a DoraArray) -> Result<Self, Self::Error> {
90        let array: &StringArray = array_ref(value)
91            .as_string_opt()
92            .wrap_err("not a string array")?;
93        if array.is_empty() {
94            eyre::bail!("empty array");
95        }
96        if array.len() != 1 {
97            eyre::bail!("expected length 1");
98        }
99        if array.null_count() != 0 {
100            eyre::bail!("array has nulls");
101        }
102        Ok(array.value(0))
103    }
104}
105
106impl TryFrom<&DoraArray> for String {
107    type Error = eyre::Report;
108    fn try_from(value: &DoraArray) -> Result<Self, Self::Error> {
109        // Delegate to the `&str` impl so the single-element validation lives in
110        // one place and the two conversions can never drift apart.
111        let s: &str = value.try_into()?;
112        Ok(s.to_string())
113    }
114}
115
116impl TryFrom<&DoraArray> for NaiveDate {
117    type Error = eyre::Report;
118    fn try_from(value: &DoraArray) -> Result<Self, Self::Error> {
119        const CTX: &str = "data type cannot be converted to NaiveDate";
120        if let Some(array) = array_ref(value)
121            .as_any()
122            .downcast_ref::<arrow::array::Date32Array>()
123        {
124            return single_temporal(array, |a, i| a.value_as_date(i), CTX);
125        }
126        let array = array_ref(value)
127            .as_any()
128            .downcast_ref::<arrow::array::Date64Array>()
129            .context("Reference is neither to a Date32Array nor a Date64Array")?;
130        single_temporal(array, |a, i| a.value_as_date(i), CTX)
131    }
132}
133
134impl TryFrom<&DoraArray> for NaiveTime {
135    type Error = eyre::Report;
136    fn try_from(value: &DoraArray) -> Result<Self, Self::Error> {
137        const CTX: &str = "data type cannot be converted to NaiveTime";
138        if let Some(array) = array_ref(value)
139            .as_any()
140            .downcast_ref::<arrow::array::Time32SecondArray>()
141        {
142            return single_temporal(array, |a, i| a.value_as_time(i), CTX);
143        }
144        if let Some(array) = array_ref(value)
145            .as_any()
146            .downcast_ref::<arrow::array::Time32MillisecondArray>()
147        {
148            return single_temporal(array, |a, i| a.value_as_time(i), CTX);
149        }
150        if let Some(array) = array_ref(value)
151            .as_any()
152            .downcast_ref::<arrow::array::Time64MicrosecondArray>()
153        {
154            return single_temporal(array, |a, i| a.value_as_time(i), CTX);
155        }
156        let array = array_ref(value)
157            .as_primitive_opt::<arrow::datatypes::Time64NanosecondType>()
158            .context("not any of the primitive Time arrays")?;
159        single_temporal(array, |a, i| a.value_as_time(i), CTX)
160    }
161}
162
163impl TryFrom<&DoraArray> for NaiveDateTime {
164    type Error = eyre::Report;
165    fn try_from(value: &DoraArray) -> Result<Self, Self::Error> {
166        const CTX: &str = "data type cannot be converted to NaiveDateTime";
167        if let Some(array) = array_ref(value)
168            .as_any()
169            .downcast_ref::<arrow::array::TimestampSecondArray>()
170        {
171            return single_temporal(array, |a, i| a.value_as_datetime(i), CTX);
172        }
173        if let Some(array) = array_ref(value)
174            .as_any()
175            .downcast_ref::<arrow::array::TimestampMillisecondArray>()
176        {
177            return single_temporal(array, |a, i| a.value_as_datetime(i), CTX);
178        }
179        if let Some(array) = array_ref(value)
180            .as_any()
181            .downcast_ref::<arrow::array::TimestampMicrosecondArray>()
182        {
183            return single_temporal(array, |a, i| a.value_as_datetime(i), CTX);
184        }
185        let array = array_ref(value)
186            .as_primitive_opt::<arrow::datatypes::TimestampNanosecondType>()
187            .context("not any of the primitive Timestamp arrays")?;
188        single_temporal(array, |a, i| a.value_as_datetime(i), CTX)
189    }
190}
191
192/// Extract the single scalar out of a length-1, non-null temporal array.
193///
194/// Every temporal `TryFrom<&DoraArray>` arm above shares the exact same shape:
195/// validate that the array holds exactly one non-null element, then convert
196/// element 0 via one of arrow's `value_as_{date,time,datetime}` accessors.
197/// Centralizing it here removes ~8 near-verbatim copies and keeps the
198/// validation and error text from drifting between them.
199fn single_temporal<T, R>(
200    array: &PrimitiveArray<T>,
201    convert: impl Fn(&PrimitiveArray<T>, usize) -> Option<R>,
202    context: &'static str,
203) -> Result<R, eyre::Error>
204where
205    T: ArrowTemporalType,
206{
207    if check_single_datetime(array) {
208        eyre::bail!("Not a valid array");
209    }
210    convert(array, 0).context(context)
211}
212
213fn check_single_datetime<T>(array: &PrimitiveArray<T>) -> bool
214where
215    T: ArrowTemporalType,
216{
217    // `len() != 1` already rejects the empty (`len() == 0`) case, so no separate
218    // `is_empty()` term is needed — unlike the primitive/`&str`/`bool` helpers,
219    // which branch on `is_empty()` to emit a distinct "empty array" message.
220    array.len() != 1 || array.null_count() != 0
221}
222fn extract_single_primitive<T>(array: &PrimitiveArray<T>) -> Result<T::Native, eyre::Error>
223where
224    T: ArrowPrimitiveType,
225{
226    if array.is_empty() {
227        eyre::bail!("empty array");
228    }
229    if array.len() != 1 {
230        eyre::bail!("expected length 1");
231    }
232    if array.null_count() != 0 {
233        eyre::bail!("array has nulls");
234    }
235    Ok(array.value(0))
236}
237
238#[cfg(test)]
239mod tests {
240    use arrow::array::{PrimitiveArray, make_array};
241
242    use crate::{DoraArray, internal::from_array_ref};
243
244    #[test]
245    fn test_u8() {
246        let array =
247            make_array(PrimitiveArray::<arrow::datatypes::UInt8Type>::from(vec![42]).into());
248        let data: DoraArray = from_array_ref(array);
249        let value: u8 = (&data).try_into().unwrap();
250        assert_eq!(value, 42);
251    }
252}
dora_arrow_convert/from_impls.rs

dora_arrow_convert/
from_impls.rs
