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

vortex_alp/alp_rd/compute/
take.rs
