Skip to main content

vortex_array/arrays/variant/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::Variant;
10use crate::arrays::VariantArray;
11use crate::arrays::dict::TakeReduce;
12use crate::arrays::variant::VariantArrayExt;
13
14impl TakeReduce for Variant {
15    fn take(array: ArrayView<'_, Variant>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
16        let core_storage = array.core_storage().take(indices.clone())?;
17        let shredded = array
18            .shredded()
19            .map(|shredded| shredded.take(indices.clone()))
20            .transpose()?;
21
22        Ok(Some(
23            VariantArray::try_new(core_storage, shredded)?.into_array(),
24        ))
25    }
26}