prost_arrow/
primitives.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use arrow_array::builder::{ArrayBuilder, PrimitiveBuilder};
5use arrow_array::{ArrayRef, ArrowPrimitiveType};
6use arrow_schema::DataType;
7
8use crate::traits::{ArrowBuilder, ToArrow};
9
10macro_rules! make_impl {
11    ($native_ty:ty, $data_ty:expr, $array_ty:ty) => {
12        impl ToArrow for $native_ty {
13            type Item = $native_ty;
14            type Builder = PrimitiveArrowBuilder<$array_ty>;
15
16            fn to_datatype() -> DataType {
17                return $data_ty;
18            }
19        }
20    };
21}
22
23make_impl!(i32, DataType::Int32, arrow_array::types::Int32Type);
24make_impl!(i64, DataType::Int64, arrow_array::types::Int64Type);
25make_impl!(u32, DataType::UInt32, arrow_array::types::UInt32Type);
26make_impl!(u64, DataType::UInt64, arrow_array::types::UInt64Type);
27make_impl!(f32, DataType::Float32, arrow_array::types::Float32Type);
28make_impl!(f64, DataType::Float64, arrow_array::types::Float64Type);
29
30pub struct PrimitiveArrowBuilder<T: ArrowPrimitiveType> {
31    builder: PrimitiveBuilder<T>,
32}
33
34impl<T: ArrowPrimitiveType> ArrowBuilder<T::Native> for PrimitiveArrowBuilder<T> {
35    fn new_with_capacity(capacity: usize) -> Self {
36        Self {
37            builder: PrimitiveBuilder::with_capacity(capacity),
38        }
39    }
40
41    fn append_value(&mut self, value: T::Native) {
42        self.builder.append_value(value)
43    }
44
45    fn append_option(&mut self, value: Option<T::Native>) {
46        self.builder.append_option(value)
47    }
48}
49
50impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveArrowBuilder<T> {
51    fn len(&self) -> usize {
52        self.builder.len()
53    }
54
55    fn finish(&mut self) -> ArrayRef {
56        Arc::new(self.builder.finish())
57    }
58
59    fn finish_cloned(&self) -> ArrayRef {
60        Arc::new(self.builder.finish_cloned())
61    }
62
63    fn as_any(&self) -> &dyn Any {
64        self
65    }
66
67    fn as_any_mut(&mut self) -> &mut dyn Any {
68        self
69    }
70
71    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
72        self
73    }
74}