vortex_array/arrays/primitive/compute/
to_arrow.rs

1use std::sync::Arc;
2
3use arrow_array::types::{
4    Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type,
5    UInt16Type, UInt32Type, UInt64Type,
6};
7use arrow_array::{ArrayRef, ArrowPrimitiveType, PrimitiveArray as ArrowPrimitiveArray};
8use arrow_buffer::ScalarBuffer;
9use arrow_schema::DataType;
10use vortex_dtype::PType;
11use vortex_error::{VortexResult, vortex_bail};
12
13use crate::Array;
14use crate::arrays::{PrimitiveArray, PrimitiveEncoding};
15use crate::compute::ToArrowFn;
16use crate::variants::PrimitiveArrayTrait;
17
18impl ToArrowFn<&PrimitiveArray> for PrimitiveEncoding {
19    fn to_arrow(
20        &self,
21        primitive_array: &PrimitiveArray,
22        data_type: &DataType,
23    ) -> VortexResult<Option<ArrayRef>> {
24        fn as_arrow_array_primitive<T: ArrowPrimitiveType>(
25            array: &PrimitiveArray,
26            data_type: &DataType,
27        ) -> VortexResult<Option<ArrayRef>> {
28            if data_type != &T::DATA_TYPE {
29                vortex_bail!("Unsupported data type: {data_type}");
30            }
31
32            Ok(Some(Arc::new(ArrowPrimitiveArray::<T>::new(
33                ScalarBuffer::<T::Native>::new(
34                    array.byte_buffer().clone().into_arrow_buffer(),
35                    0,
36                    array.len(),
37                ),
38                array.validity_mask()?.to_null_buffer(),
39            ))))
40        }
41
42        match primitive_array.ptype() {
43            PType::U8 => as_arrow_array_primitive::<UInt8Type>(primitive_array, data_type),
44            PType::U16 => as_arrow_array_primitive::<UInt16Type>(primitive_array, data_type),
45            PType::U32 => as_arrow_array_primitive::<UInt32Type>(primitive_array, data_type),
46            PType::U64 => as_arrow_array_primitive::<UInt64Type>(primitive_array, data_type),
47            PType::I8 => as_arrow_array_primitive::<Int8Type>(primitive_array, data_type),
48            PType::I16 => as_arrow_array_primitive::<Int16Type>(primitive_array, data_type),
49            PType::I32 => as_arrow_array_primitive::<Int32Type>(primitive_array, data_type),
50            PType::I64 => as_arrow_array_primitive::<Int64Type>(primitive_array, data_type),
51            PType::F16 => as_arrow_array_primitive::<Float16Type>(primitive_array, data_type),
52            PType::F32 => as_arrow_array_primitive::<Float32Type>(primitive_array, data_type),
53            PType::F64 => as_arrow_array_primitive::<Float64Type>(primitive_array, data_type),
54        }
55    }
56}
vortex_array/arrays/primitive/compute/to_arrow.rs

vortex_array/arrays/primitive/compute/
to_arrow.rs
