Skip to main content

vortex_array/arrow/executor/
byte_view.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::GenericByteViewArray;
8use arrow_array::types::ByteViewType;
9use arrow_buffer::ScalarBuffer;
10use vortex_dtype::DType;
11use vortex_dtype::Nullability;
12use vortex_dtype::arrow::FromArrowType;
13use vortex_error::VortexResult;
14
15use crate::ArrayRef;
16use crate::ExecutionCtx;
17use crate::arrays::VarBinViewArray;
18use crate::arrow::executor::validity::to_arrow_null_buffer;
19use crate::arrow::null_buffer::to_null_buffer;
20use crate::builtins::ArrayBuiltins;
21use crate::vtable::ValidityHelper;
22
23/// Convert a canonical VarBinViewArray directly to Arrow.
24pub fn canonical_varbinview_to_arrow<T: ByteViewType>(
25    array: &VarBinViewArray,
26) -> VortexResult<ArrowArrayRef> {
27    let views =
28        ScalarBuffer::<u128>::from(array.views_handle().as_host().clone().into_arrow_buffer());
29    let buffers: Vec<_> = array
30        .buffers()
31        .iter()
32        .map(|buffer| buffer.as_host().clone().into_arrow_buffer())
33        .collect();
34    let nulls = to_null_buffer(array.validity_mask()?);
35
36    // SAFETY: our own VarBinView array is considered safe.
37    Ok(Arc::new(unsafe {
38        GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls)
39    }))
40}
41
42pub fn execute_varbinview_to_arrow<T: ByteViewType>(
43    array: &VarBinViewArray,
44    ctx: &mut ExecutionCtx,
45) -> VortexResult<ArrowArrayRef> {
46    let views =
47        ScalarBuffer::<u128>::from(array.views_handle().as_host().clone().into_arrow_buffer());
48    let buffers: Vec<_> = array
49        .buffers()
50        .iter()
51        .map(|buffer| buffer.as_host().clone().into_arrow_buffer())
52        .collect();
53    let nulls = to_arrow_null_buffer(array.validity().clone(), array.len(), ctx)?;
54
55    // SAFETY: our own VarBinView array is considered safe.
56    Ok(Arc::new(unsafe {
57        GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls)
58    }))
59}
60
61pub(super) fn to_arrow_byte_view<T: ByteViewType>(
62    array: ArrayRef,
63    ctx: &mut ExecutionCtx,
64) -> VortexResult<ArrowArrayRef> {
65    // First we cast the array into the desired ByteView type.
66    // We do this in case the vortex array is Utf8, and we want Binary or vice versa. By casting
67    // first, we may push this down through the Vortex array tree. We choose nullable to be most
68    // flexible since there's no prescribed nullability in Arrow types.
69    let array = array.cast(DType::from_arrow((&T::DATA_TYPE, Nullability::Nullable)))?;
70
71    let varbinview = array.execute::<VarBinViewArray>(ctx)?;
72    canonical_varbinview_to_arrow::<T>(&varbinview)
73}