vortex_array/arrays/listview/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::Zero;
5use vortex_error::VortexResult;
6
7use crate::ArrayRef;
8use crate::ExecutionCtx;
9use crate::IntoArray;
10use crate::array::ArrayView;
11use crate::arrays::ListView;
12use crate::arrays::ListViewArray;
13use crate::arrays::dict::TakeExecute;
14use crate::arrays::dict::TakeReduce;
15use crate::arrays::listview::ListViewArrayExt;
16use crate::builtins::ArrayBuiltins;
17use crate::dtype::Nullability;
18use crate::match_each_integer_ptype;
19use crate::scalar::Scalar;
20
21/// Metadata-only take for [`ListViewArray`].
22impl TakeReduce for ListView {
23    fn take(array: ArrayView<'_, ListView>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
24        Ok(Some(apply_take(array, indices)?.into_array()))
25    }
26}
27
28/// Execution-path take for [`ListViewArray`].
29impl TakeExecute for ListView {
30    fn take(
31        array: ArrayView<'_, ListView>,
32        indices: &ArrayRef,
33        _ctx: &mut ExecutionCtx,
34    ) -> VortexResult<Option<ArrayRef>> {
35        Ok(Some(apply_take(array, indices)?.into_array()))
36    }
37}
38
39/// Shared metadata-only take: take `offsets`, `sizes` and `validity` at `indices` while reusing
40/// the original `elements` buffer as-is.
41fn apply_take(array: ArrayView<'_, ListView>, indices: &ArrayRef) -> VortexResult<ListViewArray> {
42    let elements = array.elements();
43    let offsets = array.offsets();
44    let sizes = array.sizes();
45
46    // Combine the array's validity with the indices' validity.
47    let new_validity = array.validity()?.take(indices)?;
48
49    // Take can reorder offsets, create gaps, and may introduce overlaps if `indices` contain
50    // duplicates.
51    let nullable_new_offsets = offsets.take(indices.clone())?;
52    let nullable_new_sizes = sizes.take(indices.clone())?;
53
54    // `take` returns nullable arrays; cast back to non-nullable (filling with zeros to represent
55    // the null lists — the validity mask tracks nullness separately).
56    let new_offsets = match_each_integer_ptype!(nullable_new_offsets.dtype().as_ptype(), |O| {
57        nullable_new_offsets.fill_null(Scalar::primitive(O::zero(), Nullability::NonNullable))?
58    });
59    let new_sizes = match_each_integer_ptype!(nullable_new_sizes.dtype().as_ptype(), |S| {
60        nullable_new_sizes.fill_null(Scalar::primitive(S::zero(), Nullability::NonNullable))?
61    });
62
63    // SAFETY: Take operation maintains all `ListViewArray` invariants:
64    // - `new_offsets` and `new_sizes` are derived from existing valid child arrays.
65    // - `new_offsets` and `new_sizes` are non-nullable.
66    // - `new_offsets` and `new_sizes` have the same length (both taken with the same `indices`).
67    // - Validity correctly reflects the combination of array and indices validity.
68    Ok(unsafe {
69        ListViewArray::new_unchecked(elements.clone(), new_offsets, new_sizes, new_validity)
70    })
71}
vortex_array/arrays/listview/compute/take.rs

vortex_array/arrays/listview/compute/
take.rs
