vortex_array/arrays/extension/compute/
take.rs1use 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}