Skip to main content

vortex_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 vortex_array::ArrayRef;
10use vortex_array::ExecutionCtx;
11use vortex_array::arrays::VarBinViewArray;
12use vortex_array::builtins::ArrayBuiltins;
13use vortex_array::dtype::Nullability;
14use vortex_buffer::Buffer;
15use vortex_error::VortexResult;
16
17use crate::dtype::from_arrow_data_type;
18use crate::null_buffer::to_null_buffer;
19
20/// Convert a canonical VarBinViewArray directly to Arrow.
21pub fn canonical_varbinview_to_arrow<T: ByteViewType>(
22    array: &VarBinViewArray,
23    ctx: &mut ExecutionCtx,
24) -> VortexResult<ArrowArrayRef> {
25    let views = Buffer::<u128>::from_byte_buffer(array.views_handle().as_host().clone())
26        .into_arrow_scalar_buffer();
27    let buffers: Vec<_> = array
28        .data_buffers()
29        .iter()
30        .map(|buffer| buffer.as_host().clone().into_arrow_buffer())
31        .collect();
32    let nulls = to_null_buffer(
33        array
34            .as_ref()
35            .validity()?
36            .execute_mask(array.as_ref().len(), ctx)?,
37    );
38
39    // SAFETY: our own VarBinView array is considered safe.
40    Ok(Arc::new(unsafe {
41        GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls)
42    }))
43}
44
45pub fn execute_varbinview_to_arrow<T: ByteViewType>(
46    array: &VarBinViewArray,
47    ctx: &mut ExecutionCtx,
48) -> VortexResult<ArrowArrayRef> {
49    let compacted = array.compact_buffers(ctx)?;
50    canonical_varbinview_to_arrow::<T>(&compacted, ctx)
51}
52
53pub(super) fn to_arrow_byte_view<T: ByteViewType>(
54    array: ArrayRef,
55    ctx: &mut ExecutionCtx,
56) -> VortexResult<ArrowArrayRef> {
57    // First we cast the array into the desired ByteView type.
58    // We do this in case the vortex array is Utf8, and we want Binary or vice versa. By casting
59    // first, we may push this down through the Vortex array tree. We choose nullable to be most
60    // flexible since there's no prescribed nullability in Arrow types.
61    let array = array.cast(from_arrow_data_type(&T::DATA_TYPE, Nullability::Nullable)?)?;
62
63    let array = array.execute::<ArrayRef>(ctx)?;
64    let varbinview = array.execute::<VarBinViewArray>(ctx)?;
65    execute_varbinview_to_arrow::<T>(&varbinview, ctx)
66}
67
68#[cfg(test)]
69mod tests {
70    use arrow_array::types::StringViewType;
71    use vortex_array::VortexSessionExecute;
72    use vortex_array::array_session;
73
74    use super::*;
75
76    #[test]
77    fn empty_views_are_aligned() -> VortexResult<()> {
78        let array = VarBinViewArray::from_iter_str(std::iter::empty::<&str>());
79        let mut ctx = array_session().create_execution_ctx();
80
81        let arrow = canonical_varbinview_to_arrow::<StringViewType>(&array, &mut ctx)?;
82
83        assert!(arrow.is_empty());
84        Ok(())
85    }
86}