Skip to main content

vortex_array/arrow/executor/
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::ArrayRef as ArrowArrayRef;
7use arrow_array::ArrowPrimitiveType;
8use arrow_array::PrimitiveArray as ArrowPrimitiveArray;
9use vortex_error::VortexResult;
10
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::arrays::PrimitiveArray;
14use crate::arrow::null_buffer::to_null_buffer;
15use crate::builtins::ArrayBuiltins;
16use crate::dtype::DType;
17use crate::dtype::NativePType;
18use crate::dtype::Nullability;
19
20/// Convert a canonical PrimitiveArray directly to Arrow.
21pub fn canonical_primitive_to_arrow<T: ArrowPrimitiveType>(
22    array: PrimitiveArray,
23    ctx: &mut ExecutionCtx,
24) -> VortexResult<ArrowArrayRef>
25where
26    T::Native: NativePType,
27{
28    let validity = array
29        .as_ref()
30        .validity()?
31        .to_mask(array.as_ref().len(), ctx)?;
32    let null_buffer = to_null_buffer(validity);
33    let buffer = array.into_buffer::<T::Native>().into_arrow_scalar_buffer();
34    Ok(Arc::new(ArrowPrimitiveArray::<T>::new(buffer, null_buffer)))
35}
36
37pub(super) fn to_arrow_primitive<T: ArrowPrimitiveType>(
38    array: ArrayRef,
39    ctx: &mut ExecutionCtx,
40) -> VortexResult<ArrowArrayRef>
41where
42    T::Native: NativePType,
43{
44    // We use nullable here so we can essentially ignore nullability during the cast.
45    let array = array.cast(DType::Primitive(T::Native::PTYPE, Nullability::Nullable))?;
46    let primitive = array.execute::<PrimitiveArray>(ctx)?;
47    canonical_primitive_to_arrow::<T>(primitive, ctx)
48}