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        // TODO(joe): impl validity_mask take
62        let validity = array.validity().take(indices)?;
63        let mask = validity.to_logical(indices.len())?;
64        let indices = indices.to_primitive()?;
65
66        match_each_integer_ptype!(indices.ptype(), |$I| {
67            // This is valid since all elements even null values are inside the min-max valid range.
68            take_views_into(array.views(), array.buffers(), indices.as_slice::<$I>(), mask, builder)?;
69        });
70
71        Ok(())
72    }
73}
74
75fn take_views_into<I: AsPrimitive<usize>>(
76    views: &Buffer<BinaryView>,
77    buffers: &[ByteBuffer],
78    indices: &[I],
79    mask: Mask,
80    builder: &mut VarBinViewBuilder,
81) -> VortexResult<()> {
82    let buffers_offset = u32::try_from(builder.completed_block_count())?;
83    // NOTE(ngates): this deref is not actually trivial, so we run it once.
84    let views_ref = views.deref();
85    builder.push_buffer_and_adjusted_views(
86        buffers.iter().cloned(),
87        indices
88            .iter()
89            .map(|i| views_ref[i.as_()].offset_view(buffers_offset)),
90        mask,
91    );
92    Ok(())
93}
94
95fn take_views<I: AsPrimitive<usize>>(
96    views: &Buffer<BinaryView>,
97    indices: &[I],
98) -> Buffer<BinaryView> {
99    // NOTE(ngates): this deref is not actually trivial, so we run it once.
100    let views_ref = views.deref();
101    Buffer::<BinaryView>::from_iter(indices.iter().map(|i| views_ref[i.as_()]))
102}
vortex_array/arrays/varbinview/compute/take.rs

vortex_array/arrays/varbinview/compute/
take.rs
