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
22macro_rules! impl_into_arrow {
23    ($($t:ty => $arrow_type:ty),*) => {
24        $(
25            impl IntoArrow for $t {
26                fn into_arrow(self) -> DoraArray {
27                    wrap(std::iter::once(self).collect::<PrimitiveArray<$arrow_type>>())
28                }
29            }
30        )*
31        $(
32            impl IntoArrow for Vec<$t> {
33                fn into_arrow(self) -> DoraArray {
34                    wrap(PrimitiveArray::<$arrow_type>::from(self))
35                }
36            }
37        )*
38    };
39}
40
41impl_into_arrow!(
42    u8 => UInt8Type,
43    u16 => UInt16Type,
44    u32 => UInt32Type,
45    u64 => UInt64Type,
46    i8 => Int8Type,
47    i16 => Int16Type,
48    i32 => Int32Type,
49    i64 => Int64Type,
50    f16 => Float16Type,
51    f32 => Float32Type,
52    f64 => Float64Type
53);
54
55impl IntoArrow for &str {
56    fn into_arrow(self) -> DoraArray {
57        wrap(std::iter::once(Some(self)).collect::<StringArray>())
58    }
59}
60
61impl IntoArrow for () {
62    fn into_arrow(self) -> DoraArray {
63        wrap(arrow::array::NullArray::new(0))
64    }
65}
66
67impl IntoArrow for NaiveDate {
68    fn into_arrow(self) -> DoraArray {
69        wrap(arrow::array::Date64Array::from(vec![
70            arrow::datatypes::Date64Type::from_naive_date(self),
71        ]))
72    }
73}
74
75impl IntoArrow for NaiveTime {
76    fn into_arrow(self) -> DoraArray {
77        wrap(arrow::array::Time64NanosecondArray::from(vec![
78            arrow::array::temporal_conversions::time_to_time64ns(self),
79        ]))
80    }
81}
82
83impl IntoArrow for String {
84    fn into_arrow(self) -> DoraArray {
85        wrap(std::iter::once(Some(self)).collect::<StringArray>())
86    }
87}
88
89impl IntoArrow for Vec<String> {
90    fn into_arrow(self) -> DoraArray {
91        wrap(StringArray::from(self))
92    }
93}
94
95/// The nanosecond-resolution `i64` timestamp can only represent dates in roughly
96/// 1677-09-21..2262-04-11. Dates outside that range are saturated to `i64::MIN`
97/// (far-past) or `i64::MAX` (far-future) and a `tracing::warn!` is emitted,
98/// rather than silently mapping to the Unix epoch (the previous behaviour).
99impl IntoArrow for NaiveDateTime {
100    fn into_arrow(self) -> DoraArray {
101        let timestamp =
102            match arrow::datatypes::TimestampNanosecondType::from_naive_datetime(self, None) {
103                Some(ts) => ts,
104                None => {
105                    let epoch = chrono::DateTime::UNIX_EPOCH.naive_utc();
106                    let saturated = if self >= epoch { i64::MAX } else { i64::MIN };
107                    warn!(
108                        datetime = %self,
109                        saturated_ns = saturated,
110                        "NaiveDateTime is outside the nanosecond-representable range \
111                         (~1677-09-21..2262-04-11); saturating to boundary value. \
112                         Consider using a timestamp type with a wider or lower-resolution range."
113                    );
114                    saturated
115                }
116            };
117        wrap(TimestampNanosecondArray::from(vec![timestamp]))
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use arrow::array::Array as _;
125    use chrono::NaiveDate;
126
127    fn timestamp_of(data: &DoraArray) -> i64 {
128        crate::internal::array_ref(data)
129            .as_any()
130            .downcast_ref::<TimestampNanosecondArray>()
131            .expect("timestamp array")
132            .value(0)
133    }
134
135    #[test]
136    fn naive_datetime_out_of_range_saturates() {
137        // Far-future date (year 3000) should saturate to i64::MAX
138        let far_future = NaiveDate::from_ymd_opt(3000, 1, 1)
139            .unwrap()
140            .and_hms_opt(0, 0, 0)
141            .unwrap();
142        let arr = far_future.into_arrow();
143        assert_eq!(timestamp_of(&arr), i64::MAX);
144
145        // Far-past date (year 1000) should saturate to i64::MIN
146        let far_past = NaiveDate::from_ymd_opt(1000, 1, 1)
147            .unwrap()
148            .and_hms_opt(0, 0, 0)
149            .unwrap();
150        let arr = far_past.into_arrow();
151        assert_eq!(timestamp_of(&arr), i64::MIN);
152    }
153}
154
155/// `IntoArrow` for Arrow arrays of dora's **internal** major, so
156/// `node.send_output(id, params, my_arrow_array)` keeps working for callers on
157/// that major with no wrapping and no conversion.
158///
159/// Gated for the same reason [`DoraArray::as_array`](crate::DoraArray::as_array)
160/// is: these impls name a specific Arrow major. When dora moves internally to
161/// Arrow 60 they re-gate behind `arrow-v60`, and `arrow-v59` keeps a converting
162/// equivalent — visibly, rather than silently changing meaning under callers.
163///
164/// A blanket `impl<A: arrow::array::Array> IntoArrow for A` is not possible:
165/// it overlaps with the primitive impls above under coherence's
166/// "upstream crates may add a new impl" rule, so the concrete types are listed.
167#[cfg(feature = "arrow-v59")]
168mod arrow_v59_impls {
169    use super::{DoraArray, IntoArrow, from_array_ref};
170    use arrow::array::{ArrayRef, PrimitiveArray};
171    use arrow::datatypes::ArrowPrimitiveType;
172
173    impl IntoArrow for ArrayRef {
174        fn into_arrow(self) -> DoraArray {
175            from_array_ref(self)
176        }
177    }
178
179    impl<T: ArrowPrimitiveType> IntoArrow for PrimitiveArray<T> {
180        fn into_arrow(self) -> DoraArray {
181            from_array_ref(std::sync::Arc::new(self))
182        }
183    }
184
185    macro_rules! impl_into_arrow_for_arrow_array {
186        ($($t:ty),* $(,)?) => {
187            $(
188                impl IntoArrow for $t {
189                    fn into_arrow(self) -> DoraArray {
190                        from_array_ref(std::sync::Arc::new(self))
191                    }
192                }
193            )*
194        };
195    }
196
197    impl_into_arrow_for_arrow_array!(
198        arrow::array::BooleanArray,
199        arrow::array::NullArray,
200        arrow::array::StringArray,
201        arrow::array::LargeStringArray,
202        arrow::array::StringViewArray,
203        arrow::array::BinaryArray,
204        arrow::array::LargeBinaryArray,
205        arrow::array::BinaryViewArray,
206        arrow::array::FixedSizeBinaryArray,
207        arrow::array::StructArray,
208        arrow::array::ListArray,
209        arrow::array::LargeListArray,
210        arrow::array::FixedSizeListArray,
211        arrow::array::MapArray,
212        arrow::array::UnionArray,
213    );
214}