Skip to main content

dora_arrow_convert/
into_impls.rs

1use crate::IntoArrow;
2use arrow::array::{PrimitiveArray, StringArray, TimestampNanosecondArray};
3use arrow::datatypes::{
4    ArrowTimestampType, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
5    Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
6};
7use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
8use half::f16;
9use tracing::warn;
10
11impl IntoArrow for bool {
12    type A = arrow::array::BooleanArray;
13    fn into_arrow(self) -> Self::A {
14        std::iter::once(Some(self)).collect()
15    }
16}
17
18macro_rules! impl_into_arrow {
19    ($($t:ty => $arrow_type:ty),*) => {
20        $(
21            impl IntoArrow for $t {
22                type A = PrimitiveArray<$arrow_type>;
23                fn into_arrow(self) -> Self::A {
24                    std::iter::once(self).collect()
25                }
26            }
27        )*
28        $(
29            impl IntoArrow for Vec<$t> {
30                type A = PrimitiveArray<$arrow_type>;
31                fn into_arrow(self) -> Self::A {
32                    self.into()
33                }
34            }
35        )*
36    };
37}
38
39impl_into_arrow!(
40    u8 => UInt8Type,
41    u16 => UInt16Type,
42    u32 => UInt32Type,
43    u64 => UInt64Type,
44    i8 => Int8Type,
45    i16 => Int16Type,
46    i32 => Int32Type,
47    i64 => Int64Type,
48    f16 => Float16Type,
49    f32 => Float32Type,
50    f64 => Float64Type
51);
52
53impl IntoArrow for &str {
54    type A = StringArray;
55    fn into_arrow(self) -> Self::A {
56        std::iter::once(Some(self)).collect()
57    }
58}
59
60impl IntoArrow for () {
61    type A = arrow::array::NullArray;
62    fn into_arrow(self) -> Self::A {
63        arrow::array::NullArray::new(0)
64    }
65}
66
67impl IntoArrow for NaiveDate {
68    type A = arrow::array::Date64Array;
69    fn into_arrow(self) -> Self::A {
70        arrow::array::Date64Array::from(vec![arrow::datatypes::Date64Type::from_naive_date(self)])
71    }
72}
73
74impl IntoArrow for NaiveTime {
75    type A = arrow::array::Time64NanosecondArray;
76    fn into_arrow(self) -> Self::A {
77        arrow::array::Time64NanosecondArray::from(vec![
78            arrow::array::temporal_conversions::time_to_time64ns(self),
79        ])
80    }
81}
82
83impl IntoArrow for String {
84    type A = StringArray;
85    fn into_arrow(self) -> Self::A {
86        std::iter::once(Some(self)).collect()
87    }
88}
89
90impl IntoArrow for Vec<String> {
91    type A = StringArray;
92    fn into_arrow(self) -> Self::A {
93        StringArray::from(self)
94    }
95}
96
97/// The nanosecond-resolution `i64` timestamp can only represent dates in roughly
98/// 1677-09-21..2262-04-11. Dates outside that range are saturated to `i64::MIN`
99/// (far-past) or `i64::MAX` (far-future) and a `tracing::warn!` is emitted,
100/// rather than silently mapping to the Unix epoch (the previous behaviour).
101impl IntoArrow for NaiveDateTime {
102    type A = arrow::array::TimestampNanosecondArray;
103    fn into_arrow(self) -> Self::A {
104        let timestamp =
105            match arrow::datatypes::TimestampNanosecondType::from_naive_datetime(self, None) {
106                Some(ts) => ts,
107                None => {
108                    let epoch = chrono::DateTime::UNIX_EPOCH.naive_utc();
109                    let saturated = if self >= epoch { i64::MAX } else { i64::MIN };
110                    warn!(
111                        datetime = %self,
112                        saturated_ns = saturated,
113                        "NaiveDateTime is outside the nanosecond-representable range \
114                         (~1677-09-21..2262-04-11); saturating to boundary value. \
115                         Consider using a timestamp type with a wider or lower-resolution range."
116                    );
117                    saturated
118                }
119            };
120        TimestampNanosecondArray::from(vec![timestamp])
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127    use chrono::NaiveDate;
128
129    #[test]
130    fn naive_datetime_out_of_range_saturates() {
131        // Far-future date (year 3000) should saturate to i64::MAX
132        let far_future = NaiveDate::from_ymd_opt(3000, 1, 1)
133            .unwrap()
134            .and_hms_opt(0, 0, 0)
135            .unwrap();
136        let arr = far_future.into_arrow();
137        assert_eq!(arr.value(0), i64::MAX);
138
139        // Far-past date (year 1000) should saturate to i64::MIN
140        let far_past = NaiveDate::from_ymd_opt(1000, 1, 1)
141            .unwrap()
142            .and_hms_opt(0, 0, 0)
143            .unwrap();
144        let arr = far_past.into_arrow();
145        assert_eq!(arr.value(0), i64::MIN);
146    }
147}