vortex_array/arrays/struct_/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::IntoArray;
8use crate::array::ArrayView;
9use crate::arrays::Struct;
10use crate::arrays::StructArray;
11use crate::arrays::dict::TakeReduce;
12use crate::arrays::struct_::StructArrayExt;
13use crate::builtins::ArrayBuiltins;
14use crate::scalar::Scalar;
15use crate::validity::Validity;
16
17impl TakeReduce for Struct {
18    fn take(array: ArrayView<'_, Struct>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
19        // If the struct array is empty then the indices must be all null, otherwise it will access
20        // an out of bounds element.
21        if array.is_empty() {
22            return StructArray::try_new_with_dtype(
23                array.iter_unmasked_fields().cloned().collect::<Vec<_>>(),
24                array.struct_fields().clone(),
25                indices.len(),
26                Validity::AllInvalid,
27            )
28            .map(StructArray::into_array)
29            .map(Some);
30        }
31
32        // TODO(connor): This could be bad for cache locality...
33
34        // Fill null indices with zero so they point at a valid row.
35        // Note that we strip nullability so that `Take::return_dtype` doesn't union nullable into
36        // each field's dtype (the struct-level validity already captures which rows are null).
37        let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable());
38        let inner_indices = indices.clone().fill_null(fill_scalar)?;
39
40        StructArray::try_new_with_dtype(
41            array
42                .iter_unmasked_fields()
43                .map(|field| field.take(inner_indices.clone()))
44                .collect::<Result<Vec<_>, _>>()?,
45            array.struct_fields().clone(),
46            indices.len(),
47            array.validity()?.take(indices)?,
48        )
49        .map(|a| a.into_array())
50        .map(Some)
51    }
52}
vortex_array/arrays/struct_/compute/take.rs

vortex_array/arrays/struct_/compute/
take.rs
