vortex_alp/alp_rd/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_array::builtins::ArrayBuiltins;
10use vortex_array::scalar::Scalar;
11use vortex_error::VortexResult;
12
13use crate::ALPRD;
14use crate::ALPRDArrayExt;
15
16impl TakeExecute for ALPRD {
17    fn take(
18        array: ArrayView<'_, Self>,
19        indices: &ArrayRef,
20        ctx: &mut ExecutionCtx,
21    ) -> VortexResult<Option<ArrayRef>> {
22        let taken_left_parts = array.left_parts().take(indices.clone())?;
23        // With nullable take indices, `Patches::take` widens the (still all-valid) exception
24        // values to nullable, but ALPRD stores exceptions as the non-nullable left-parts dtype.
25        // Narrow them back so the patches satisfy the invariant `validate_parts` enforces. This is
26        // a no-op when the take indices are non-nullable and the values already match.
27        let exceptions_dtype = taken_left_parts.dtype().as_nonnullable();
28        let left_parts_exceptions = array
29            .left_parts_patches()
30            .map(|patches| {
31                patches
32                    .take(indices, ctx)?
33                    .map(|taken| taken.map_values(|values| values.cast(exceptions_dtype.clone())))
34                    .transpose()
35            })
36            .transpose()?
37            .flatten();
38        let right_parts = array
39            .right_parts()
40            .take(indices.clone())?
41            .fill_null(Scalar::zero_value(array.right_parts().dtype()))?;
42
43        Ok(Some(
44            ALPRD::try_new(
45                array
46                    .dtype()
47                    .with_nullability(taken_left_parts.dtype().nullability()),
48                taken_left_parts,
49                array.left_parts_dictionary().clone(),
50                right_parts,
51                array.right_bit_width(),
52                left_parts_exceptions,
53            )?
54            .into_array(),
55        ))
56    }
57}
58
59#[cfg(test)]
60mod test {
61    use std::sync::LazyLock;
62
63    use rstest::rstest;
64    use vortex_array::IntoArray;
65    use vortex_array::VortexSessionExecute;
66    use vortex_array::array_session;
67    use vortex_array::arrays::PrimitiveArray;
68    use vortex_array::assert_arrays_eq;
69    use vortex_array::compute::conformance::take::test_take_conformance;
70    use vortex_session::VortexSession;
71
72    use crate::ALPRDArrayExt;
73    use crate::ALPRDFloat;
74    use crate::RDEncoder;
75
76    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
77        let session = array_session();
78        crate::initialize(&session);
79        session
80    });
81
82    #[rstest]
83    #[case(0.1f32, 0.2f32, 3e25f32)]
84    #[case(0.1f64, 0.2f64, 3e100f64)]
85    fn test_take<T: ALPRDFloat>(#[case] a: T, #[case] b: T, #[case] outlier: T) {
86        use vortex_array::IntoArray as _;
87        use vortex_buffer::buffer;
88
89        let mut ctx = SESSION.create_execution_ctx();
90        let array = PrimitiveArray::from_iter([a, b, outlier]);
91        let encoded = RDEncoder::new(&[a, b]).encode(array.as_view());
92
93        assert!(encoded.left_parts_patches().is_some());
94        assert!(
95            encoded
96                .left_parts_patches()
97                .unwrap()
98                .dtype()
99                .is_unsigned_int()
100        );
101
102        let taken = encoded
103            .take(buffer![0, 2].into_array())
104            .unwrap()
105            .execute::<PrimitiveArray>(&mut ctx)
106            .unwrap();
107
108        assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier]), &mut ctx);
109    }
110
111    #[rstest]
112    #[case(0.1f32, 0.2f32, 3e25f32)]
113    #[case(0.1f64, 0.2f64, 3e100f64)]
114    fn take_with_nulls<T: ALPRDFloat>(#[case] a: T, #[case] b: T, #[case] outlier: T) {
115        let mut ctx = SESSION.create_execution_ctx();
116        let array = PrimitiveArray::from_iter([a, b, outlier]);
117        let encoded = RDEncoder::new(&[a, b]).encode(array.as_view());
118
119        assert!(encoded.left_parts_patches().is_some());
120        assert!(
121            encoded
122                .left_parts_patches()
123                .unwrap()
124                .dtype()
125                .is_unsigned_int()
126        );
127
128        let taken = encoded
129            .take(PrimitiveArray::from_option_iter([Some(0), Some(2), None]).into_array())
130            .unwrap()
131            .execute::<PrimitiveArray>(&mut ctx)
132            .unwrap();
133
134        assert_arrays_eq!(
135            taken,
136            PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]),
137            &mut ctx
138        );
139    }
140
141    #[rstest]
142    #[case(0.1f32, 0.2f32, 3e25f32)]
143    #[case(0.1f64, 0.2f64, 3e100f64)]
144    fn test_take_conformance_alprd<T: ALPRDFloat>(#[case] a: T, #[case] b: T, #[case] outlier: T) {
145        let mut ctx = SESSION.create_execution_ctx();
146        test_take_conformance(
147            &RDEncoder::new(&[a, b])
148                .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view())
149                .into_array(),
150            &mut ctx,
151        );
152    }
153
154    #[rstest]
155    #[case(0.1f32, 3e25f32)]
156    #[case(0.5f64, 1e100f64)]
157    fn test_take_with_nulls_conformance<T: ALPRDFloat>(#[case] a: T, #[case] outlier: T) {
158        let mut ctx = SESSION.create_execution_ctx();
159        test_take_conformance(
160            &RDEncoder::new(&[a])
161                .encode(
162                    PrimitiveArray::from_option_iter([Some(a), None, Some(outlier), Some(a), None])
163                        .as_view(),
164                )
165                .into_array(),
166            &mut ctx,
167        );
168    }
169}
vortex_alp/alp_rd/compute/take.rs

vortex_alp/alp_rd/compute/
take.rs
