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::dict::TakeReduce;
14use crate::arrays::extension::ExtensionArrayExt;
15
16impl TakeReduce for Extension {
17    fn take(array: ArrayView<'_, Extension>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
18        let taken_storage = array.storage_array().take(indices.clone())?;
19        Ok(Some(
20            ExtensionArray::new(
21                array
22                    .ext_dtype()
23                    .with_nullability(taken_storage.dtype().nullability()),
24                taken_storage,
25            )
26            .into_array(),
27        ))
28    }
29}
30
31impl TakeExecute for Extension {
32    fn take(
33        array: ArrayView<'_, Extension>,
34        indices: &ArrayRef,
35        _ctx: &mut ExecutionCtx,
36    ) -> VortexResult<Option<ArrayRef>> {
37        let taken_storage = array.storage_array().take(indices.clone())?;
38        Ok(Some(
39            ExtensionArray::new(
40                array
41                    .ext_dtype()
42                    .with_nullability(taken_storage.dtype().nullability()),
43                taken_storage,
44            )
45            .into_array(),
46        ))
47    }
48}