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