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_eq;
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_eq!(
172        computed_len,
173        output_len,
174        "PiecewiseSequenceArray expanded length does not match declared length",
175    );
176
177    let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
178    let spare = &mut views.spare_capacity_mut()[..output_len];
179    let mut cursor = 0usize;
180    for &start in starts {
181        let start = start.as_();
182        let src = &source[start..][..length];
183        spare[cursor..][..src.len()].write_copy_of_slice(src);
184        cursor += src.len();
185    }
186    // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity.
187    unsafe { views.set_len(cursor) };
188    vortex_ensure_eq!(
189        views.len(),
190        output_len,
191        "PiecewiseSequenceArray expanded length does not match declared length",
192    );
193    Ok(views.freeze())
194}
195
196fn gather_view_slices<S, L>(
197    source: &[BinaryView],
198    starts: &[S],
199    lengths: &[L],
200    output_len: usize,
201) -> VortexResult<Buffer<BinaryView>>
202where
203    S: UnsignedPType,
204    L: UnsignedPType,
205{
206    let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
207    let spare = &mut views.spare_capacity_mut()[..output_len];
208    let mut cursor = 0usize;
209    for (&start, &length) in starts.iter().zip_eq(lengths) {
210        let start = start.as_();
211        let length = length.as_();
212        let src = &source[start..][..length];
213        spare[cursor..][..src.len()].write_copy_of_slice(src);
214        cursor += src.len();
215    }
216    // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity.
217    unsafe { views.set_len(cursor) };
218    vortex_ensure_eq!(
219        views.len(),
220        output_len,
221        "PiecewiseSequenceArray expanded length does not match declared length",
222    );
223    Ok(views.freeze())
224}
225
226#[cfg(test)]
227mod tests {
228    use rstest::rstest;
229    use vortex_buffer::BitBuffer;
230    use vortex_buffer::buffer;
231    use vortex_error::VortexResult;
232
233    use crate::IntoArray;
234    use crate::VortexSessionExecute;
235    use crate::array_session;
236    use crate::arrays::VarBinViewArray;
237    use crate::arrays::varbinview::compute::take::PrimitiveArray;
238    use crate::compute::conformance::take::test_take_conformance;
239    use crate::dtype::DType;
240    use crate::dtype::Nullability::NonNullable;
241    use crate::validity::Validity;
242
243    #[test]
244    fn take_nullable() -> VortexResult<()> {
245        let arr = VarBinViewArray::from_iter_nullable_str([
246            Some("one"),
247            None,
248            Some("three"),
249            Some("four"),
250            None,
251            Some("six"),
252        ]);
253
254        let taken = arr.take(buffer![0, 3].into_array())?;
255
256        assert!(taken.dtype().is_nullable());
257        let mut ctx = array_session().create_execution_ctx();
258        let taken = taken.execute::<VarBinViewArray>(&mut ctx)?;
259        let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?;
260        let result = (0..taken.len())
261            .map(|i| {
262                mask.value(i)
263                    .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) })
264            })
265            .collect::<Vec<_>>();
266        assert_eq!(result, [Some("one".to_string()), Some("four".to_string())]);
267        Ok(())
268    }
269
270    #[test]
271    fn take_nullable_indices() -> VortexResult<()> {
272        let arr = VarBinViewArray::from_iter(["one", "two"].map(Some), DType::Utf8(NonNullable));
273
274        let indices = PrimitiveArray::new(
275            // Verify that garbage values at NULL indices are ignored.
276            buffer![1u64, 999],
277            Validity::from(BitBuffer::from(vec![true, false])),
278        );
279
280        let taken = arr.take(indices.into_array())?;
281
282        assert!(taken.dtype().is_nullable());
283        let mut ctx = array_session().create_execution_ctx();
284        let taken = taken.execute::<VarBinViewArray>(&mut ctx)?;
285        let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?;
286        let result = (0..taken.len())
287            .map(|i| {
288                mask.value(i)
289                    .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) })
290            })
291            .collect::<Vec<_>>();
292        assert_eq!(result, [Some("two".to_string()), None]);
293        Ok(())
294    }
295
296    #[rstest]
297    #[case(VarBinViewArray::from_iter(
298        ["hello", "world", "test", "data", "array"].map(Some),
299        DType::Utf8(NonNullable),
300    ))]
301    #[case(VarBinViewArray::from_iter_nullable_str([
302        Some("hello"),
303        None,
304        Some("test"),
305        Some("data"),
306        None,
307    ]))]
308    #[case(VarBinViewArray::from_iter(
309        [b"hello".as_slice(), b"world", b"test", b"data", b"array"].map(Some),
310        DType::Binary(NonNullable),
311    ))]
312    #[case(VarBinViewArray::from_iter(["single"].map(Some), DType::Utf8(NonNullable)))]
313    fn test_take_varbinview_conformance(#[case] array: VarBinViewArray) {
314        test_take_conformance(
315            &array.into_array(),
316            &mut array_session().create_execution_ctx(),
317        );
318    }
319}
vortex_array/arrays/varbinview/compute/take.rs

vortex_array/arrays/varbinview/compute/
take.rs
