Skip to main content

vortex_array/arrays/map/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::array::ArrayView;
8use crate::arrays::ListView;
9use crate::arrays::dict::TakeExecute;
10use crate::arrays::dict::TakeReduce;
11use crate::arrays::map::Map;
12use crate::arrays::map::MapArrayExt;
13use crate::arrays::map::MapArraySlotsExt;
14use crate::arrays::map::compute::rebuild_map_from_array;
15use crate::executor::ExecutionCtx;
16
17impl TakeReduce for Map {
18    fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
19        let Some(entries) =
20            <ListView as TakeReduce>::take(array.entries().as_::<ListView>(), indices)?
21        else {
22            return Ok(None);
23        };
24
25        rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some)
26    }
27}
28
29impl TakeExecute for Map {
30    fn take(
31        array: ArrayView<'_, Self>,
32        indices: &ArrayRef,
33        ctx: &mut ExecutionCtx,
34    ) -> VortexResult<Option<ArrayRef>> {
35        let Some(entries) =
36            <ListView as TakeExecute>::take(array.entries().as_::<ListView>(), indices, ctx)?
37        else {
38            return Ok(None);
39        };
40
41        rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some)
42    }
43}