dora_arrow_convert/
into_impls.rs1use 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
97impl 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 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 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}