vortex_array/arrays/varbinview/compute/
take.rs

1use std::ops::Deref;
2
3use num_traits::AsPrimitive;
4use vortex_buffer::{Buffer, ByteBuffer};
5use vortex_dtype::match_each_integer_ptype;
6use vortex_error::{VortexResult, vortex_bail};
7use vortex_mask::Mask;
8
9use crate::arrays::{BinaryView, VarBinViewArray, VarBinViewEncoding};
10use crate::builders::{ArrayBuilder, VarBinViewBuilder};
11use crate::compute::TakeFn;
12use crate::variants::PrimitiveArrayTrait;
13use crate::{Array, ArrayRef, ToCanonical};
14
15/// Take involves creating a new array that references the old array, just with the given set of views.
16impl TakeFn<&VarBinViewArray> for VarBinViewEncoding {
17    fn take(&self, array: &VarBinViewArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
18        // Compute the new validity
19
20        // This is valid since all elements (of all arrays) even null values are inside must be the
21        // min-max valid range.
22        let validity = array.validity().take(indices)?;
23        let indices = indices.to_primitive()?;
24
25        let views_buffer = match_each_integer_ptype!(indices.ptype(), |$I| {
26        // This is valid since all elements even null values are inside the min-max valid range.
27            take_views(array.views(), indices.as_slice::<$I>())
28        });
29
30        Ok(VarBinViewArray::try_new(
31            views_buffer,
32            array.buffers().to_vec(),
33            array.dtype().with_nullability(
34                (array.dtype().is_nullable() || indices.dtype().is_nullable()).into(),
35            ),
36            validity,
37        )?
38        .into_array())
39    }
40
41    fn take_into(
42        &self,
43        array: &VarBinViewArray,
44        indices: &dyn Array,
45        builder: &mut dyn ArrayBuilder,
46    ) -> VortexResult<()> {
47        if array.len() == 0 {
48            vortex_bail!("Cannot take_into from an empty array");
49        }
50
51        let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
52            vortex_bail!(
53                "Cannot take_into a non-varbinview builder {:?}",
54                builder.as_any().type_id()
55            );
56        };
57        // Compute the new validity
58
59        // This is valid since all elements (of all arrays) even null values are inside must be the
60        // min-max valid range.
61        let validity = array.validity().take(indices)?;
62        let mask = validity.to_mask(indices.len())?;
63        let indices = indices.to_primitive()?;
64
65        match_each_integer_ptype!(indices.ptype(), |$I| {
66            // This is valid since all elements even null values are inside the min-max valid range.
67            take_views_into(array.views(), array.buffers(), indices.as_slice::<$I>(), mask, builder)?;
68        });
69
70        Ok(())
71    }
72}
73
74fn take_views_into<I: AsPrimitive<usize>>(
75    views: &Buffer<BinaryView>,
76    buffers: &[ByteBuffer],
77    indices: &[I],
78    mask: Mask,
79    builder: &mut VarBinViewBuilder,
80) -> VortexResult<()> {
81    let buffers_offset = u32::try_from(builder.completed_block_count())?;
82    // NOTE(ngates): this deref is not actually trivial, so we run it once.
83    let views_ref = views.deref();
84    builder.push_buffer_and_adjusted_views(
85        buffers.iter().cloned(),
86        indices
87            .iter()
88            .map(|i| views_ref[i.as_()].offset_view(buffers_offset)),
89        mask,
90    );
91    Ok(())
92}
93
94fn take_views<I: AsPrimitive<usize>>(
95    views: &Buffer<BinaryView>,
96    indices: &[I],
97) -> Buffer<BinaryView> {
98    // NOTE(ngates): this deref is not actually trivial, so we run it once.
99    let views_ref = views.deref();
100    Buffer::<BinaryView>::from_iter(indices.iter().map(|i| views_ref[i.as_()]))
101}
vortex_array/arrays/varbinview/compute/take.rs

vortex_array/arrays/varbinview/compute/
take.rs
