prost_arrow/
primitives.rs1use 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}