vortex_alp/alp/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::dict::TakeExecute;
9use vortex_error::VortexResult;
10
11use crate::ALP;
12use crate::ALPArrayExt;
13use crate::ALPArraySlotsExt;
14
15impl TakeExecute for ALP {
16    fn take(
17        array: ArrayView<'_, Self>,
18        indices: &ArrayRef,
19        ctx: &mut ExecutionCtx,
20    ) -> VortexResult<Option<ArrayRef>> {
21        let taken_encoded = array.encoded().take(indices.clone())?;
22        let taken_patches = array
23            .patches()
24            .map(|p| p.take(indices, ctx))
25            .transpose()?
26            .flatten();
27        Ok(Some(
28            ALP::try_new(taken_encoded, array.exponents(), taken_patches)?.into_array(),
29        ))
30    }
31}
32
33#[cfg(test)]
34mod test {
35    use rstest::rstest;
36    use vortex_array::IntoArray;
37    use vortex_array::VortexSessionExecute;
38    use vortex_array::array_session;
39    use vortex_array::arrays::PrimitiveArray;
40    use vortex_array::compute::conformance::take::test_take_conformance;
41    use vortex_buffer::buffer;
42
43    use crate::alp_encode;
44
45    #[rstest]
46    #[case(buffer![1.23f32, 4.56, 7.89, 10.11, 12.13].into_array())]
47    #[case(buffer![100.1f64, 200.2, 300.3, 400.4, 500.5].into_array())]
48    #[case(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).into_array())]
49    #[case(buffer![42.42f64].into_array())]
50    fn test_take_alp_conformance(#[case] array: vortex_array::ArrayRef) {
51        let mut ctx = array_session().create_execution_ctx();
52        let array_primitive = array.execute::<PrimitiveArray>(&mut ctx).unwrap();
53        let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap();
54        test_take_conformance(&alp.into_array(), &mut ctx);
55    }
56}
vortex_alp/alp/compute/take.rs

vortex_alp/alp/compute/
take.rs
