Skip to main content

typed_arrow/bridge/
primitives.rs

1//! Primitive Arrow bindings (integers, floats, bool, f16).
2
3#[cfg(feature = "views")]
4use arrow_array::Array;
5use arrow_array::{
6    PrimitiveArray,
7    builder::PrimitiveBuilder,
8    types::{
9        Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
10        UInt8Type, UInt16Type, UInt32Type, UInt64Type,
11    },
12};
13use arrow_schema::DataType;
14use half::f16;
15
16use super::ArrowBinding;
17#[cfg(feature = "views")]
18use super::ArrowBindingView;
19
20// Primitive integers/floats
21macro_rules! impl_primitive_binding {
22    ($rust:ty, $atype:ty, $dt:expr) => {
23        impl ArrowBinding for $rust {
24            type Builder = PrimitiveBuilder<$atype>;
25            type Array = PrimitiveArray<$atype>;
26            #[inline]
27            fn data_type() -> DataType {
28                $dt
29            }
30            #[inline]
31            fn new_builder(capacity: usize) -> Self::Builder {
32                PrimitiveBuilder::<$atype>::with_capacity(capacity)
33            }
34            #[inline]
35            fn append_value(b: &mut Self::Builder, v: &Self) {
36                b.append_value(*v);
37            }
38            #[inline]
39            fn append_null(b: &mut Self::Builder) {
40                b.append_null();
41            }
42            #[inline]
43            fn finish(mut b: Self::Builder) -> Self::Array {
44                b.finish()
45            }
46        }
47
48        #[cfg(feature = "views")]
49        impl ArrowBindingView for $rust {
50            type Array = PrimitiveArray<$atype>;
51            type View<'a> = $rust;
52
53            fn get_view(
54                array: &Self::Array,
55                index: usize,
56            ) -> Result<Self::View<'_>, crate::schema::ViewAccessError> {
57                if index >= array.len() {
58                    return Err(crate::schema::ViewAccessError::OutOfBounds {
59                        index,
60                        len: array.len(),
61                        field_name: None,
62                    });
63                }
64                if array.is_null(index) {
65                    return Err(crate::schema::ViewAccessError::UnexpectedNull {
66                        index,
67                        field_name: None,
68                    });
69                }
70                Ok(array.value(index))
71            }
72        }
73    };
74}
75
76impl_primitive_binding!(i8, Int8Type, DataType::Int8);
77impl_primitive_binding!(i16, Int16Type, DataType::Int16);
78impl_primitive_binding!(i32, Int32Type, DataType::Int32);
79impl_primitive_binding!(i64, Int64Type, DataType::Int64);
80impl_primitive_binding!(u8, UInt8Type, DataType::UInt8);
81impl_primitive_binding!(u16, UInt16Type, DataType::UInt16);
82impl_primitive_binding!(u32, UInt32Type, DataType::UInt32);
83impl_primitive_binding!(u64, UInt64Type, DataType::UInt64);
84impl_primitive_binding!(f32, Float32Type, DataType::Float32);
85impl_primitive_binding!(f64, Float64Type, DataType::Float64);
86
87// Float16 (half-precision)
88impl ArrowBinding for f16 {
89    type Builder = PrimitiveBuilder<Float16Type>;
90    type Array = PrimitiveArray<Float16Type>;
91    #[inline]
92    fn data_type() -> DataType {
93        DataType::Float16
94    }
95    #[inline]
96    fn new_builder(capacity: usize) -> Self::Builder {
97        PrimitiveBuilder::<Float16Type>::with_capacity(capacity)
98    }
99    #[inline]
100    fn append_value(b: &mut Self::Builder, v: &Self) {
101        b.append_value(*v);
102    }
103    #[inline]
104    fn append_null(b: &mut Self::Builder) {
105        b.append_null();
106    }
107    #[inline]
108    fn finish(mut b: Self::Builder) -> Self::Array {
109        b.finish()
110    }
111}
112
113#[cfg(feature = "views")]
114impl ArrowBindingView for f16 {
115    type Array = PrimitiveArray<Float16Type>;
116    type View<'a> = f16;
117
118    fn get_view(
119        array: &Self::Array,
120        index: usize,
121    ) -> Result<Self::View<'_>, crate::schema::ViewAccessError> {
122        if index >= array.len() {
123            return Err(crate::schema::ViewAccessError::OutOfBounds {
124                index,
125                len: array.len(),
126                field_name: None,
127            });
128        }
129        if array.is_null(index) {
130            return Err(crate::schema::ViewAccessError::UnexpectedNull {
131                index,
132                field_name: None,
133            });
134        }
135        Ok(array.value(index))
136    }
137}
138
139// Boolean
140impl ArrowBinding for bool {
141    type Builder = arrow_array::builder::BooleanBuilder;
142    type Array = arrow_array::BooleanArray;
143    #[inline]
144    fn data_type() -> DataType {
145        DataType::Boolean
146    }
147    #[inline]
148    fn new_builder(capacity: usize) -> Self::Builder {
149        arrow_array::builder::BooleanBuilder::with_capacity(capacity)
150    }
151    #[inline]
152    fn append_value(b: &mut Self::Builder, v: &Self) {
153        b.append_value(*v);
154    }
155    #[inline]
156    fn append_null(b: &mut Self::Builder) {
157        b.append_null();
158    }
159    #[inline]
160    fn finish(mut b: Self::Builder) -> Self::Array {
161        b.finish()
162    }
163}
164
165#[cfg(feature = "views")]
166impl ArrowBindingView for bool {
167    type Array = arrow_array::BooleanArray;
168    type View<'a> = bool;
169
170    fn get_view(
171        array: &Self::Array,
172        index: usize,
173    ) -> Result<Self::View<'_>, crate::schema::ViewAccessError> {
174        if index >= array.len() {
175            return Err(crate::schema::ViewAccessError::OutOfBounds {
176                index,
177                len: array.len(),
178                field_name: None,
179            });
180        }
181        if array.is_null(index) {
182            return Err(crate::schema::ViewAccessError::UnexpectedNull {
183                index,
184                field_name: None,
185            });
186        }
187        Ok(array.value(index))
188    }
189}