vortex_compute/arrow/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use arrow_array::Array;
7use arrow_array::ArrayRef;
8use arrow_array::PrimitiveArray;
9use arrow_array::types::Float16Type;
10use arrow_array::types::Float32Type;
11use arrow_array::types::Float64Type;
12use arrow_array::types::Int8Type;
13use arrow_array::types::Int16Type;
14use arrow_array::types::Int32Type;
15use arrow_array::types::Int64Type;
16use arrow_array::types::UInt8Type;
17use arrow_array::types::UInt16Type;
18use arrow_array::types::UInt32Type;
19use arrow_array::types::UInt64Type;
20use vortex_buffer::Buffer;
21use vortex_dtype::half::f16;
22use vortex_error::VortexResult;
23use vortex_vector::match_each_pvector;
24use vortex_vector::primitive::PVector;
25use vortex_vector::primitive::PrimitiveVector;
26
27use crate::arrow::IntoArrow;
28use crate::arrow::IntoVector;
29use crate::arrow::nulls_to_mask;
30
31impl IntoArrow for PrimitiveVector {
32    type Output = ArrayRef;
33
34    fn into_arrow(self) -> VortexResult<Self::Output> {
35        match_each_pvector!(self, |v| { Ok(Arc::new(v.into_arrow()?)) })
36    }
37}
38
39macro_rules! impl_primitive_to_arrow {
40    ($T:ty, $A:ty) => {
41        impl IntoArrow for PVector<$T> {
42            type Output = PrimitiveArray<$A>;
43
44            fn into_arrow(self) -> VortexResult<Self::Output> {
45                let (elements, validity) = self.into_parts();
46                Ok(PrimitiveArray::<$A>::new(
47                    elements.into_arrow_scalar_buffer(),
48                    validity.into(),
49                ))
50            }
51        }
52    };
53}
54
55impl_primitive_to_arrow!(u8, UInt8Type);
56impl_primitive_to_arrow!(u16, UInt16Type);
57impl_primitive_to_arrow!(u32, UInt32Type);
58impl_primitive_to_arrow!(u64, UInt64Type);
59impl_primitive_to_arrow!(i8, Int8Type);
60impl_primitive_to_arrow!(i16, Int16Type);
61impl_primitive_to_arrow!(i32, Int32Type);
62impl_primitive_to_arrow!(i64, Int64Type);
63impl_primitive_to_arrow!(f16, Float16Type);
64impl_primitive_to_arrow!(f32, Float32Type);
65impl_primitive_to_arrow!(f64, Float64Type);
66
67macro_rules! impl_primitive_from_arrow {
68    ($T:ty, $A:ty) => {
69        impl IntoVector for &PrimitiveArray<$A> {
70            type Output = PVector<$T>;
71
72            fn into_vector(self) -> VortexResult<Self::Output> {
73                let elements = Buffer::<$T>::from_arrow_scalar_buffer(self.values().clone());
74                let validity = nulls_to_mask(self.nulls(), self.len());
75                Ok(PVector::new(elements, validity))
76            }
77        }
78    };
79}
80
81impl_primitive_from_arrow!(u8, UInt8Type);
82impl_primitive_from_arrow!(u16, UInt16Type);
83impl_primitive_from_arrow!(u32, UInt32Type);
84impl_primitive_from_arrow!(u64, UInt64Type);
85impl_primitive_from_arrow!(i8, Int8Type);
86impl_primitive_from_arrow!(i16, Int16Type);
87impl_primitive_from_arrow!(i32, Int32Type);
88impl_primitive_from_arrow!(i64, Int64Type);
89impl_primitive_from_arrow!(f16, Float16Type);
90impl_primitive_from_arrow!(f32, Float32Type);
91impl_primitive_from_arrow!(f64, Float64Type);
vortex_compute/arrow/primitive.rs

vortex_compute/arrow/
primitive.rs
