vortex_array/arrays/varbinview/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::iter;
5use std::sync::Arc;
6
7use itertools::Itertools as _;
8use num_traits::AsPrimitive;
9use vortex_buffer::Buffer;
10use vortex_buffer::BufferMut;
11use vortex_error::VortexResult;
12use vortex_error::vortex_ensure;
13use vortex_error::vortex_err;
14use vortex_mask::AllOr;
15use vortex_mask::Mask;
16
17use crate::ArrayRef;
18use crate::Columnar;
19use crate::IntoArray;
20use crate::array::ArrayView;
21use crate::arrays::PiecewiseSequence;
22use crate::arrays::PrimitiveArray;
23use crate::arrays::VarBinView;
24use crate::arrays::VarBinViewArray;
25use crate::arrays::dict::TakeExecute;
26use crate::arrays::piecewise_sequence::constant_unsigned_usize;
27use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
28use crate::arrays::varbinview::BinaryView;
29use crate::buffer::BufferHandle;
30use crate::dtype::UnsignedPType;
31use crate::executor::ExecutionCtx;
32use crate::match_each_integer_ptype;
33use crate::match_each_unsigned_integer_ptype;
34
35impl TakeExecute for VarBinView {
36    /// Take involves creating a new array that references the old array, just with the given set of views.
37    fn take(
38        array: ArrayView<'_, VarBinView>,
39        indices: &ArrayRef,
40        ctx: &mut ExecutionCtx,
41    ) -> VortexResult<Option<ArrayRef>> {
42        if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequence>()
43            && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)?
44        {
45            return Ok(Some(taken));
46        }
47
48        let validity = array.validity()?.take(indices)?;
49        let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
50
51        let indices_mask = indices
52            .as_ref()
53            .validity()?
54            .execute_mask(indices.as_ref().len(), ctx)?;
55        let views_buffer = match_each_integer_ptype!(indices.ptype(), |I| {
56            take_views(array.views(), indices.as_slice::<I>(), &indices_mask)
57        });
58
59        // SAFETY: taking all components at same indices maintains invariants
60        unsafe {
61            Ok(Some(
62                VarBinViewArray::new_handle_unchecked(
63                    BufferHandle::new_host(views_buffer.into_byte_buffer()),
64                    Arc::clone(array.data_buffers()),
65                    array
66                        .dtype()
67                        .union_nullability(indices.dtype().nullability()),
68                    validity,
69                )
70                .into_array(),
71            ))
72        }
73    }
74}
75
76fn take_contiguous_ranges(
77    array: ArrayView<'_, VarBinView>,
78    indices: ArrayView<'_, PiecewiseSequence>,
79    indices_ref: &ArrayRef,
80    ctx: &mut ExecutionCtx,
81) -> VortexResult<Option<ArrayRef>> {
82    let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else {
83        return Ok(None);
84    };
85    let source = array.views();
86    let output_len = indices_ref.len();
87    let views = match lengths {
88        Columnar::Constant(lengths) => {
89            let length = constant_unsigned_usize(&lengths);
90            match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
91                gather_view_slices_constant_length(
92                    source,
93                    starts.as_slice::<S>(),
94                    length,
95                    output_len,
96                )?
97            })
98        }
99        Columnar::Canonical(lengths) => {
100            let lengths = lengths.into_primitive();
101            match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
102                match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
103                    gather_view_slices(
104                        source,
105                        starts.as_slice::<S>(),
106                        lengths.as_slice::<L>(),
107                        output_len,
108                    )?
109                })
110            })
111        }
112    };
113    let validity = array.validity()?.take(indices_ref)?;
114
115    // SAFETY: ranges were validated against the source views, and copied views still reference the
116    // same backing data buffers.
117    unsafe {
118        Ok(Some(
119            VarBinViewArray::new_handle_unchecked(
120                BufferHandle::new_host(views.into_byte_buffer()),
121                Arc::clone(array.data_buffers()),
122                array.dtype().clone(),
123                validity,
124            )
125            .into_array(),
126        ))
127    }
128}
129
130fn take_views<I: AsPrimitive<usize>>(
131    views_ref: &[BinaryView],
132    indices: &[I],
133    mask: &Mask,
134) -> Buffer<BinaryView> {
135    // NOTE(ngates): this deref is not actually trivial, so we run it once.
136    // We do not use iter_bools directly, since the resulting dyn iterator cannot
137    // implement TrustedLen.
138    match mask.bit_buffer() {
139        AllOr::All => {
140            Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[i.as_()]))
141        }
142        AllOr::None => Buffer::<BinaryView>::from_trusted_len_iter(iter::repeat_n(
143            BinaryView::default(),
144            indices.len(),
145        )),
146        AllOr::Some(buffer) => Buffer::<BinaryView>::from_trusted_len_iter(
147            buffer.iter().zip(indices.iter()).map(|(valid, idx)| {
148                if valid {
149                    views_ref[idx.as_()]
150                } else {
151                    BinaryView::default()
152                }
153            }),
154        ),
155    }
156}
157
158fn gather_view_slices_constant_length<S>(
159    source: &[BinaryView],
160    starts: &[S],
161    length: usize,
162    output_len: usize,
163) -> VortexResult<Buffer<BinaryView>>
164where
165    S: UnsignedPType,
166{
167    let computed_len = starts
168        .len()
169        .checked_mul(length)
170        .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
171    vortex_ensure!(
172        computed_len == output_len,
173        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
174    );
175
176    let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
177    let spare = &mut views.spare_capacity_mut()[..output_len];
178    let mut cursor = 0usize;
179    for &start in starts {
180        let start = start.as_();
181        let src = &source[start..][..length];
182        spare[cursor..][..src.len()].write_copy_of_slice(src);
183        cursor += src.len();
184    }
185    // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity.
186    unsafe { views.set_len(cursor) };
187    vortex_ensure!(
188        views.len() == output_len,
189        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
190        views.len()
191    );
192    Ok(views.freeze())
193}
194
195fn gather_view_slices<S, L>(
196    source: &[BinaryView],
197    starts: &[S],
198    lengths: &[L],
199    output_len: usize,
200) -> VortexResult<Buffer<BinaryView>>
201where
202    S: UnsignedPType,
203    L: UnsignedPType,
204{
205    let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
206    let spare = &mut views.spare_capacity_mut()[..output_len];
207    let mut cursor = 0usize;
208    for (&start, &length) in starts.iter().zip_eq(lengths) {
209        let start = start.as_();
210        let length = length.as_();
211        let src = &source[start..][..length];
212        spare[cursor..][..src.len()].write_copy_of_slice(src);
213        cursor += src.len();
214    }
215    // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity.
216    unsafe { views.set_len(cursor) };
217    vortex_ensure!(
218        views.len() == output_len,
219        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
220        views.len()
221    );
222    Ok(views.freeze())
223}
224
225#[cfg(test)]
226mod tests {
227    use rstest::rstest;
228    use vortex_buffer::BitBuffer;
229    use vortex_buffer::buffer;
230    use vortex_error::VortexResult;
231
232    use crate::IntoArray;
233    use crate::VortexSessionExecute;
234    use crate::array_session;
235    use crate::arrays::VarBinViewArray;
236    use crate::arrays::varbinview::compute::take::PrimitiveArray;
237    use crate::compute::conformance::take::test_take_conformance;
238    use crate::dtype::DType;
239    use crate::dtype::Nullability::NonNullable;
240    use crate::validity::Validity;
241
242    #[test]
243    fn take_nullable() -> VortexResult<()> {
244        let arr = VarBinViewArray::from_iter_nullable_str([
245            Some("one"),
246            None,
247            Some("three"),
248            Some("four"),
249            None,
250            Some("six"),
251        ]);
252
253        let taken = arr.take(buffer![0, 3].into_array())?;
254
255        assert!(taken.dtype().is_nullable());
256        let mut ctx = array_session().create_execution_ctx();
257        let taken = taken.execute::<VarBinViewArray>(&mut ctx)?;
258        let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?;
259        let result = (0..taken.len())
260            .map(|i| {
261                mask.value(i)
262                    .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) })
263            })
264            .collect::<Vec<_>>();
265        assert_eq!(result, [Some("one".to_string()), Some("four".to_string())]);
266        Ok(())
267    }
268
269    #[test]
270    fn take_nullable_indices() -> VortexResult<()> {
271        let arr = VarBinViewArray::from_iter(["one", "two"].map(Some), DType::Utf8(NonNullable));
272
273        let indices = PrimitiveArray::new(
274            // Verify that garbage values at NULL indices are ignored.
275            buffer![1u64, 999],
276            Validity::from(BitBuffer::from(vec![true, false])),
277        );
278
279        let taken = arr.take(indices.into_array())?;
280
281        assert!(taken.dtype().is_nullable());
282        let mut ctx = array_session().create_execution_ctx();
283        let taken = taken.execute::<VarBinViewArray>(&mut ctx)?;
284        let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?;
285        let result = (0..taken.len())
286            .map(|i| {
287                mask.value(i)
288                    .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) })
289            })
290            .collect::<Vec<_>>();
291        assert_eq!(result, [Some("two".to_string()), None]);
292        Ok(())
293    }
294
295    #[rstest]
296    #[case(VarBinViewArray::from_iter(
297        ["hello", "world", "test", "data", "array"].map(Some),
298        DType::Utf8(NonNullable),
299    ))]
300    #[case(VarBinViewArray::from_iter_nullable_str([
301        Some("hello"),
302        None,
303        Some("test"),
304        Some("data"),
305        None,
306    ]))]
307    #[case(VarBinViewArray::from_iter(
308        [b"hello".as_slice(), b"world", b"test", b"data", b"array"].map(Some),
309        DType::Binary(NonNullable),
310    ))]
311    #[case(VarBinViewArray::from_iter(["single"].map(Some), DType::Utf8(NonNullable)))]
312    fn test_take_varbinview_conformance(#[case] array: VarBinViewArray) {
313        test_take_conformance(
314            &array.into_array(),
315            &mut array_session().create_execution_ctx(),
316        );
317    }
318}
vortex_array/arrays/varbinview/compute/take.rs

vortex_array/arrays/varbinview/compute/
take.rs
