vortex_array/arrays/struct_/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::Nullability;
5use vortex_error::VortexResult;
6use vortex_scalar::Scalar;
7
8use crate::Array;
9use crate::ArrayRef;
10use crate::IntoArray;
11use crate::arrays::StructArray;
12use crate::arrays::StructVTable;
13use crate::compute::TakeKernel;
14use crate::compute::TakeKernelAdapter;
15use crate::compute::{self};
16use crate::register_kernel;
17use crate::validity::Validity;
18use crate::vtable::ValidityHelper;
19
20impl TakeKernel for StructVTable {
21    fn take(&self, array: &StructArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
22        // If the struct array is empty then the indices must be all null, otherwise it will access
23        // an out of bounds element
24        if array.is_empty() {
25            return StructArray::try_new_with_dtype(
26                array.fields().clone(),
27                array.struct_fields().clone(),
28                indices.len(),
29                Validity::AllInvalid,
30            )
31            .map(StructArray::into_array);
32        }
33        // The validity is applied to the struct validity,
34        let inner_indices = &compute::fill_null(
35            indices,
36            &Scalar::default_value(indices.dtype().with_nullability(Nullability::NonNullable)),
37        )?;
38        StructArray::try_new_with_dtype(
39            array
40                .fields()
41                .iter()
42                .map(|field| compute::take(field, inner_indices))
43                .collect::<Result<Vec<_>, _>>()?,
44            array.struct_fields().clone(),
45            indices.len(),
46            array.validity().take(indices)?,
47        )
48        .map(|a| a.into_array())
49    }
50}
51
52register_kernel!(TakeKernelAdapter(StructVTable).lift());
vortex_array/arrays/struct_/compute/take.rs

vortex_array/arrays/struct_/compute/
take.rs
