vortex_sequence/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::cast::NumCast;
5use vortex_array::ArrayRef;
6use vortex_array::ArrayView;
7use vortex_array::ExecutionCtx;
8use vortex_array::IntoArray;
9use vortex_array::arrays::ConstantArray;
10use vortex_array::arrays::PrimitiveArray;
11use vortex_array::arrays::dict::TakeExecute;
12use vortex_array::dtype::DType;
13use vortex_array::dtype::IntegerPType;
14use vortex_array::dtype::NativePType;
15use vortex_array::dtype::Nullability;
16use vortex_array::match_each_integer_ptype;
17use vortex_array::match_each_native_ptype;
18use vortex_array::scalar::Scalar;
19use vortex_array::validity::Validity;
20use vortex_buffer::Buffer;
21use vortex_error::VortexExpect;
22use vortex_error::VortexResult;
23use vortex_error::vortex_panic;
24use vortex_mask::AllOr;
25use vortex_mask::Mask;
26
27use crate::Sequence;
28
29fn take_inner<T: IntegerPType, S: NativePType>(
30    mul: S,
31    base: S,
32    indices: &[T],
33    indices_mask: Mask,
34    result_nullability: Nullability,
35    len: usize,
36) -> ArrayRef {
37    match indices_mask.bit_buffer() {
38        AllOr::All => PrimitiveArray::new(
39            Buffer::from_trusted_len_iter(indices.iter().map(|i| {
40                if i.as_() >= len {
41                    vortex_panic!(OutOfBounds: i.as_(), 0, len);
42                }
43                let i = <S as NumCast>::from::<T>(*i).vortex_expect("all indices fit");
44                base + i * mul
45            })),
46            Validity::from(result_nullability),
47        )
48        .into_array(),
49        AllOr::None => ConstantArray::new(
50            Scalar::null(DType::Primitive(S::PTYPE, Nullability::Nullable)),
51            indices.len(),
52        )
53        .into_array(),
54        AllOr::Some(b) => {
55            let buffer =
56                Buffer::from_trusted_len_iter(indices.iter().enumerate().map(|(mask_index, i)| {
57                    if b.value(mask_index) {
58                        if i.as_() >= len {
59                            vortex_panic!(OutOfBounds: i.as_(), 0, len);
60                        }
61
62                        let i =
63                            <S as NumCast>::from::<T>(*i).vortex_expect("all valid indices fit");
64                        base + i * mul
65                    } else {
66                        S::zero()
67                    }
68                }));
69            PrimitiveArray::new(buffer, Validity::from(b.clone())).into_array()
70        }
71    }
72}
73
74impl TakeExecute for Sequence {
75    fn take(
76        array: ArrayView<'_, Self>,
77        indices: &ArrayRef,
78        ctx: &mut ExecutionCtx,
79    ) -> VortexResult<Option<ArrayRef>> {
80        let mask = indices.validity()?.execute_mask(indices.len(), ctx)?;
81        let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
82        let result_nullability = array.dtype().nullability() | indices.dtype().nullability();
83
84        match_each_integer_ptype!(indices.ptype(), |T| {
85            let indices = indices.as_slice::<T>();
86            match_each_native_ptype!(array.ptype(), |S| {
87                let mul = array.multiplier().cast::<S>()?;
88                let base = array.base().cast::<S>()?;
89                Ok(Some(take_inner(
90                    mul,
91                    base,
92                    indices,
93                    mask,
94                    result_nullability,
95                    array.len(),
96                )))
97            })
98        })
99    }
100}
101
102#[cfg(test)]
103mod test {
104    use rstest::rstest;
105    use vortex_array::Canonical;
106    use vortex_array::IntoArray;
107    use vortex_array::VortexSessionExecute;
108    use vortex_array::array_session;
109    use vortex_array::arrays::PrimitiveArray;
110    use vortex_array::compute::conformance::take::test_take_conformance;
111    use vortex_array::dtype::Nullability;
112
113    use crate::Sequence;
114    use crate::SequenceArray;
115
116    #[rstest]
117    #[case::basic_sequence(Sequence::try_new_typed(
118        0i32,
119        1i32,
120        Nullability::NonNullable,
121        10
122    ).unwrap())]
123    #[case::sequence_with_multiplier(Sequence::try_new_typed(
124        10i32,
125        5i32,
126        Nullability::Nullable,
127        20
128    ).unwrap())]
129    #[case::sequence_i64(Sequence::try_new_typed(
130        100i64,
131        10i64,
132        Nullability::NonNullable,
133        50
134    ).unwrap())]
135    #[case::sequence_u32(Sequence::try_new_typed(
136        0u32,
137        2u32,
138        Nullability::NonNullable,
139        100
140    ).unwrap())]
141    #[case::sequence_negative_step(Sequence::try_new_typed(
142        1000i32,
143        -10i32,
144        Nullability::Nullable,
145        30
146    ).unwrap())]
147    #[case::sequence_constant(Sequence::try_new_typed(
148        42i32,
149        0i32,  // multiplier of 0 means all values are the same
150        Nullability::Nullable,
151        15
152    ).unwrap())]
153    #[case::sequence_i16(Sequence::try_new_typed(
154        -100i16,
155        3i16,
156        Nullability::NonNullable,
157        25
158    ).unwrap())]
159    #[case::sequence_large(Sequence::try_new_typed(
160        0i64,
161        1i64,
162        Nullability::Nullable,
163        1000
164    ).unwrap())]
165    fn sequence_take_conformance(#[case] sequence: SequenceArray) {
166        test_take_conformance(
167            &sequence.into_array(),
168            &mut array_session().create_execution_ctx(),
169        );
170    }
171
172    #[test]
173    #[should_panic(expected = "out of bounds")]
174    fn test_bounds_check() {
175        let array = Sequence::try_new_typed(0i32, 1i32, Nullability::NonNullable, 10).unwrap();
176        let indices = PrimitiveArray::from_iter([0i32, 20]);
177        let _array = array
178            .take(indices.into_array())
179            .unwrap()
180            .execute::<Canonical>(&mut array_session().create_execution_ctx())
181            .unwrap();
182    }
183}
vortex_sequence/compute/take.rs

vortex_sequence/compute/
take.rs
