dora_arrow_convert/
into_impls.rs1use crate::{DoraArray, IntoArrow, internal::from_array_ref};
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
11fn wrap(array: impl arrow::array::Array + 'static) -> DoraArray {
13 from_array_ref(std::sync::Arc::new(array))
14}
15
16impl IntoArrow for bool {
17 fn into_arrow(self) -> DoraArray {
18 wrap(std::iter::once(Some(self)).collect::<arrow::array::BooleanArray>())
19 }
20}
21
22impl IntoArrow for Vec<bool> {
23 fn into_arrow(self) -> DoraArray {
24 wrap(arrow::array::BooleanArray::from(self))
25 }
26}
27
28macro_rules! impl_into_arrow {
29 ($($t:ty => $arrow_type:ty),*) => {
30 $(
31 impl IntoArrow for $t {
32 fn into_arrow(self) -> DoraArray {
33 wrap(std::iter::once(self).collect::<PrimitiveArray<$arrow_type>>())
34 }
35 }
36 )*
37 $(
38 impl IntoArrow for Vec<$t> {
39 fn into_arrow(self) -> DoraArray {
40 wrap(PrimitiveArray::<$arrow_type>::from(self))
41 }
42 }
43 )*
44 };
45}
46
47impl_into_arrow!(
48 u8 => UInt8Type,
49 u16 => UInt16Type,
50 u32 => UInt32Type,
51 u64 => UInt64Type,
52 i8 => Int8Type,
53 i16 => Int16Type,
54 i32 => Int32Type,
55 i64 => Int64Type,
56 f16 => Float16Type,
57 f32 => Float32Type,
58 f64 => Float64Type
59);
60
61impl IntoArrow for &str {
62 fn into_arrow(self) -> DoraArray {
63 wrap(std::iter::once(Some(self)).collect::<StringArray>())
64 }
65}
66
67impl IntoArrow for () {
68 fn into_arrow(self) -> DoraArray {
69 wrap(arrow::array::NullArray::new(0))
70 }
71}
72
73impl IntoArrow for NaiveDate {
74 fn into_arrow(self) -> DoraArray {
75 wrap(arrow::array::Date64Array::from(vec![
76 arrow::datatypes::Date64Type::from_naive_date(self),
77 ]))
78 }
79}
80
81impl IntoArrow for NaiveTime {
82 fn into_arrow(self) -> DoraArray {
83 wrap(arrow::array::Time64NanosecondArray::from(vec![
84 arrow::array::temporal_conversions::time_to_time64ns(self),
85 ]))
86 }
87}
88
89impl IntoArrow for String {
90 fn into_arrow(self) -> DoraArray {
91 wrap(std::iter::once(Some(self)).collect::<StringArray>())
92 }
93}
94
95impl IntoArrow for Vec<String> {
96 fn into_arrow(self) -> DoraArray {
97 wrap(StringArray::from(self))
98 }
99}
100
101impl IntoArrow for NaiveDateTime {
106 fn into_arrow(self) -> DoraArray {
107 let timestamp =
108 match arrow::datatypes::TimestampNanosecondType::from_naive_datetime(self, None) {
109 Some(ts) => ts,
110 None => {
111 let epoch = chrono::DateTime::UNIX_EPOCH.naive_utc();
112 let saturated = if self >= epoch { i64::MAX } else { i64::MIN };
113 warn!(
114 datetime = %self,
115 saturated_ns = saturated,
116 "NaiveDateTime is outside the nanosecond-representable range \
117 (~1677-09-21..2262-04-11); saturating to boundary value. \
118 Consider using a timestamp type with a wider or lower-resolution range."
119 );
120 saturated
121 }
122 };
123 wrap(TimestampNanosecondArray::from(vec![timestamp]))
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130 use arrow::array::Array as _;
131 use chrono::NaiveDate;
132
133 fn timestamp_of(data: &DoraArray) -> i64 {
134 crate::internal::array_ref(data)
135 .as_any()
136 .downcast_ref::<TimestampNanosecondArray>()
137 .expect("timestamp array")
138 .value(0)
139 }
140
141 #[test]
142 fn vec_bool_into_arrow_roundtrips() {
143 use arrow::array::BooleanArray;
144 let values = vec![true, false, true, true];
145 let arr = values.clone().into_arrow();
146 let boolean = crate::internal::array_ref(&arr)
147 .as_any()
148 .downcast_ref::<BooleanArray>()
149 .expect("boolean array");
150 assert_eq!(boolean.len(), values.len());
151 assert_eq!(boolean.null_count(), 0);
152 let got: Vec<bool> = (0..boolean.len()).map(|i| boolean.value(i)).collect();
153 assert_eq!(got, values);
154 }
155
156 #[test]
157 fn naive_datetime_out_of_range_saturates() {
158 let far_future = NaiveDate::from_ymd_opt(3000, 1, 1)
160 .unwrap()
161 .and_hms_opt(0, 0, 0)
162 .unwrap();
163 let arr = far_future.into_arrow();
164 assert_eq!(timestamp_of(&arr), i64::MAX);
165
166 let far_past = NaiveDate::from_ymd_opt(1000, 1, 1)
168 .unwrap()
169 .and_hms_opt(0, 0, 0)
170 .unwrap();
171 let arr = far_past.into_arrow();
172 assert_eq!(timestamp_of(&arr), i64::MIN);
173 }
174}
175
176#[cfg(feature = "arrow-v59")]
189mod arrow_v59_impls {
190 use super::{DoraArray, IntoArrow, from_array_ref};
191 use arrow::array::{ArrayRef, PrimitiveArray};
192 use arrow::datatypes::ArrowPrimitiveType;
193
194 impl IntoArrow for ArrayRef {
195 fn into_arrow(self) -> DoraArray {
196 from_array_ref(self)
197 }
198 }
199
200 impl<T: ArrowPrimitiveType> IntoArrow for PrimitiveArray<T> {
201 fn into_arrow(self) -> DoraArray {
202 from_array_ref(std::sync::Arc::new(self))
203 }
204 }
205
206 macro_rules! impl_into_arrow_for_arrow_array {
207 ($($t:ty),* $(,)?) => {
208 $(
209 impl IntoArrow for $t {
210 fn into_arrow(self) -> DoraArray {
211 from_array_ref(std::sync::Arc::new(self))
212 }
213 }
214 )*
215 };
216 }
217
218 impl_into_arrow_for_arrow_array!(
219 arrow::array::BooleanArray,
220 arrow::array::NullArray,
221 arrow::array::StringArray,
222 arrow::array::LargeStringArray,
223 arrow::array::StringViewArray,
224 arrow::array::BinaryArray,
225 arrow::array::LargeBinaryArray,
226 arrow::array::BinaryViewArray,
227 arrow::array::FixedSizeBinaryArray,
228 arrow::array::StructArray,
229 arrow::array::ListArray,
230 arrow::array::LargeListArray,
231 arrow::array::FixedSizeListArray,
232 arrow::array::MapArray,
233 arrow::array::UnionArray,
234 );
235}