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::DynArray;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::dict::TakeExecute;
9use vortex_error::VortexResult;
10
11use crate::ALPArray;
12use crate::ALPVTable;
13
14impl TakeExecute for ALPVTable {
15    fn take(
16        array: &ALPArray,
17        indices: &ArrayRef,
18        ctx: &mut ExecutionCtx,
19    ) -> VortexResult<Option<ArrayRef>> {
20        let taken_encoded = array.encoded().take(indices.to_array())?;
21        let taken_patches = array
22            .patches()
23            .map(|p| p.take(indices, ctx))
24            .transpose()?
25            .flatten()
26            .map(|patches| {
27                patches.cast_values(
28                    &array
29                        .dtype()
30                        .with_nullability(taken_encoded.dtype().nullability()),
31                )
32            })
33            .transpose()?;
34        Ok(Some(
35            ALPArray::new(taken_encoded, array.exponents(), taken_patches).into_array(),
36        ))
37    }
38}
39
40#[cfg(test)]
41mod test {
42    use rstest::rstest;
43    use vortex_array::IntoArray;
44    use vortex_array::ToCanonical;
45    use vortex_array::arrays::PrimitiveArray;
46    use vortex_array::compute::conformance::take::test_take_conformance;
47    use vortex_buffer::buffer;
48
49    use crate::alp_encode;
50
51    #[rstest]
52    #[case(buffer![1.23f32, 4.56, 7.89, 10.11, 12.13].into_array())]
53    #[case(buffer![100.1f64, 200.2, 300.3, 400.4, 500.5].into_array())]
54    #[case(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).into_array())]
55    #[case(buffer![42.42f64].into_array())]
56    fn test_take_alp_conformance(#[case] array: vortex_array::ArrayRef) {
57        let alp = alp_encode(&array.to_primitive(), None).unwrap();
58        test_take_conformance(&alp.into_array());
59    }
60}
vortex_alp/alp/compute/take.rs

vortex_alp/alp/compute/
take.rs
