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

vortex_array/arrays/struct_/compute/
take.rs
