dora_arrow_convert/
lib.rs

1//! Provides functions for converting between Apache Arrow arrays and Rust data types.
2
3#![warn(missing_docs)]
4
5use arrow::array::{
6    Array, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array,
7    UInt8Array, UInt16Array, UInt32Array, UInt64Array,
8};
9use arrow::datatypes::DataType;
10use eyre::{ContextCompat, Result, eyre};
11use num::NumCast;
12use std::ops::{Deref, DerefMut};
13
14mod from_impls;
15mod into_impls;
16
17/// Data that can be converted to an Arrow array.
18///
19/// This is the conversion that dora node APIs use to turn plain Rust values
20/// into the [Apache Arrow](https://arrow.apache.org/) columnar format before
21/// sending them as outputs. Implementations are provided for booleans,
22/// strings, the primitive integer and float types, `Vec`s of those primitive
23/// types, and a few `chrono` date/time types. The unit type `()` converts to
24/// an empty [`NullArray`](arrow::array::NullArray), which is useful for
25/// outputs that carry only metadata.
26///
27/// For the opposite direction (reading received Arrow data back into Rust
28/// types), see the `TryFrom<&ArrowData>` implementations on [`ArrowData`].
29///
30/// # Example
31///
32/// ```
33/// use arrow::array::Array;
34/// use dora_arrow_convert::IntoArrow;
35///
36/// let array = vec![1.0_f32, 2.0, 3.0].into_arrow();
37/// assert_eq!(array.len(), 3);
38///
39/// let single = 42_u8.into_arrow();
40/// assert_eq!(single.len(), 1);
41/// ```
42pub trait IntoArrow {
43    /// The Arrow array type that the data converts to.
44    type A: Array;
45
46    /// Convert the data into an Arrow array.
47    fn into_arrow(self) -> Self::A;
48}
49
50/// Wrapper type for an Arrow [`ArrayRef`](arrow::array::ArrayRef).
51#[derive(Debug)]
52pub struct ArrowData(pub arrow::array::ArrayRef);
53
54impl Deref for ArrowData {
55    type Target = arrow::array::ArrayRef;
56
57    fn deref(&self) -> &Self::Target {
58        &self.0
59    }
60}
61
62impl DerefMut for ArrowData {
63    fn deref_mut(&mut self) -> &mut Self::Target {
64        &mut self.0
65    }
66}
67
68macro_rules! register_array_handlers {
69    ($(($variant:path, $array_type:ty, $type_name:expr)),* $(,)?) => {
70        /// Tries to convert the given Arrow array into a `Vec` of integers or floats.
71        ///
72        /// Returns an error if the array contains any null values, consistent
73        /// with every other [`TryFrom<&ArrowData>`] impl in this crate.
74        pub fn into_vec<T>(data: &ArrowData) -> Result<Vec<T>>
75        where
76            T: Copy + NumCast + 'static,
77        {
78            match data.data_type() {
79                $(
80                    $variant => {
81                        let buffer: &$array_type = data
82                            .as_any()
83                            .downcast_ref()
84                            .context(concat!("series is not ", $type_name))?;
85
86                        if buffer.null_count() != 0 {
87                            eyre::bail!("array has nulls");
88                        }
89
90                        let mut result = Vec::with_capacity(buffer.len());
91                        for &v in buffer.values() {
92                            let converted = NumCast::from(v).context(format!("Failed to cast value from {} to target type",$type_name))?;
93                            result.push(converted);
94                        }
95                        Ok(result)
96                    }
97                ),*
98                // Error handling for unsupported types
99                unsupported_type => Err(eyre!("Unsupported data type for conversion: {:?}", unsupported_type))
100            }
101        }
102    };
103}
104
105// Register all supported array types in one place
106register_array_handlers! {
107    (DataType::Float32, Float32Array, "float32"),
108    (DataType::Float64, Float64Array, "float64"),
109    (DataType::Int8, Int8Array, "int8"),
110    (DataType::Int16, Int16Array, "int16"),
111    (DataType::Int32, Int32Array, "int32"),
112    (DataType::Int64, Int64Array, "int64"),
113    (DataType::UInt8, UInt8Array, "uint8"),
114    (DataType::UInt16, UInt16Array, "uint16"),
115    (DataType::UInt32, UInt32Array, "uint32"),
116    (DataType::UInt64, UInt64Array, "uint64"),
117    (DataType::Float16, Float16Array, "float16"),
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use arrow::array::ArrayRef;
124    use half::f16;
125    use std::sync::Arc;
126
127    /// Round-trips every type registered in `register_array_handlers!` so a
128    /// future macro entry that gets dropped (the original cause of #2080) is
129    /// caught by a failing test rather than a silent runtime error.
130    #[test]
131    fn into_vec_supports_all_registered_types() {
132        macro_rules! assert_round_trip {
133            ($array_type:ty, $rust_type:ty, $values:expr) => {{
134                let values: Vec<$rust_type> = $values;
135                let array: ArrayRef = Arc::new(<$array_type>::from(values.clone()));
136                let data = ArrowData(array);
137                let result: Vec<$rust_type> = into_vec(&data).unwrap();
138                assert_eq!(result, values);
139            }};
140        }
141
142        assert_round_trip!(Float32Array, f32, vec![1.0, 2.5, -3.0]);
143        assert_round_trip!(Float64Array, f64, vec![1.0, 2.5, -3.0]);
144        assert_round_trip!(Int8Array, i8, vec![-1, 2, 3]);
145        assert_round_trip!(Int16Array, i16, vec![-1, 2, 3]);
146        assert_round_trip!(Int32Array, i32, vec![-1, 2, 3]);
147        assert_round_trip!(Int64Array, i64, vec![-1, 2, 3]);
148        assert_round_trip!(UInt8Array, u8, vec![1, 2, 3]);
149        assert_round_trip!(UInt16Array, u16, vec![1, 2, 3]);
150        assert_round_trip!(UInt32Array, u32, vec![1, 2, 3]);
151        assert_round_trip!(UInt64Array, u64, vec![1, 2, 3]);
152
153        // Float16 needs explicit f16 construction; round-trip back to f16.
154        let values = vec![f16::from_f32(1.0), f16::from_f32(2.5), f16::from_f32(-3.0)];
155        let array: ArrayRef = Arc::new(Float16Array::from(values.clone()));
156        let data = ArrowData(array);
157        let result: Vec<f16> = into_vec(&data).unwrap();
158        assert_eq!(result, values);
159    }
160
161    /// The case from the issue: a `UInt64` array previously errored with
162    /// "Unsupported data type for conversion: UInt64".
163    #[test]
164    fn into_vec_handles_uint64() {
165        let data = ArrowData(Arc::new(UInt64Array::from(vec![1u64, 2, 3])));
166        let res: Vec<u64> = into_vec(&data).unwrap();
167        assert_eq!(res, vec![1u64, 2, 3]);
168    }
169
170    #[test]
171    fn into_vec_rejects_arrays_with_nulls() {
172        let array: ArrayRef = Arc::new(UInt64Array::from(vec![Some(1u64), None, Some(3)]));
173        let data = ArrowData(array);
174        let res: Result<Vec<u64>> = into_vec(&data);
175        assert!(res.is_err());
176    }
177
178    #[test]
179    fn into_vec_rejects_unsupported_type() {
180        let array: ArrayRef = Arc::new(arrow::array::BooleanArray::from(vec![true, false]));
181        let data = ArrowData(array);
182        let res: Result<Vec<u8>> = into_vec(&data);
183        assert!(res.is_err());
184    }
185}
dora_arrow_convert/lib.rs

dora_arrow_convert/
lib.rs
