Skip to main content

dora_arrow_convert/
into_impls.rs

1use 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
11/// Wrap a concrete Arrow array as a [`DoraArray`].
12fn 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
101/// The nanosecond-resolution `i64` timestamp can only represent dates in roughly
102/// 1677-09-21..2262-04-11. Dates outside that range are saturated to `i64::MIN`
103/// (far-past) or `i64::MAX` (far-future) and a `tracing::warn!` is emitted,
104/// rather than silently mapping to the Unix epoch (the previous behaviour).
105impl 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        // Far-future date (year 3000) should saturate to i64::MAX
159        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        // Far-past date (year 1000) should saturate to i64::MIN
167        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/// `IntoArrow` for Arrow arrays of dora's **internal** major, so
177/// `node.send_output(id, params, my_arrow_array)` keeps working for callers on
178/// that major with no wrapping and no conversion.
179///
180/// Gated for the same reason [`DoraArray::as_array`](crate::DoraArray::as_array)
181/// is: these impls name a specific Arrow major. When dora moves internally to
182/// Arrow 60 they re-gate behind `arrow-v60`, and `arrow-v59` keeps a converting
183/// equivalent — visibly, rather than silently changing meaning under callers.
184///
185/// A blanket `impl<A: arrow::array::Array> IntoArrow for A` is not possible:
186/// it overlaps with the primitive impls above under coherence's
187/// "upstream crates may add a new impl" rule, so the concrete types are listed.
188#[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}