Skip to main content

vortex_array/arrays/extension/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::ExecutionCtx;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::Extension;
11use crate::arrays::ExtensionArray;
12use crate::arrays::dict::TakeExecute;
13use crate::arrays::extension::ExtensionArrayExt;
14
15impl TakeExecute for Extension {
16    fn take(
17        array: ArrayView<'_, Extension>,
18        indices: &ArrayRef,
19        _ctx: &mut ExecutionCtx,
20    ) -> VortexResult<Option<ArrayRef>> {
21        let taken_storage = array.storage_array().take(indices.clone())?;
22        Ok(Some(
23            ExtensionArray::new(
24                array
25                    .ext_dtype()
26                    .with_nullability(taken_storage.dtype().nullability()),
27                taken_storage,
28            )
29            .into_array(),
30        ))
31    }
32}