vortex_array/arrays/bool/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools as _;
5use num_traits::AsPrimitive;
6use vortex_buffer::BitBuffer;
7use vortex_buffer::BitBufferView;
8use vortex_buffer::get_bit;
9use vortex_error::VortexResult;
10use vortex_mask::Mask;
11
12use crate::ArrayRef;
13use crate::IntoArray;
14use crate::array::ArrayView;
15use crate::arrays::Bool;
16use crate::arrays::BoolArray;
17use crate::arrays::ConstantArray;
18use crate::arrays::PrimitiveArray;
19use crate::arrays::bool::BoolArrayExt;
20use crate::arrays::dict::TakeExecute;
21use crate::builtins::ArrayBuiltins;
22use crate::executor::ExecutionCtx;
23use crate::match_each_integer_ptype;
24use crate::scalar::Scalar;
25
26impl TakeExecute for Bool {
27    fn take(
28        array: ArrayView<'_, Bool>,
29        indices: &ArrayRef,
30        ctx: &mut ExecutionCtx,
31    ) -> VortexResult<Option<ArrayRef>> {
32        let indices_nulls_zeroed = match indices.validity()?.execute_mask(indices.len(), ctx)? {
33            Mask::AllTrue(_) => indices.clone(),
34            Mask::AllFalse(_) => {
35                return Ok(Some(
36                    ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len())
37                        .into_array(),
38                ));
39            }
40            Mask::Values(_) => indices
41                .clone()
42                .fill_null(Scalar::from(0).cast(indices.dtype())?)?,
43        };
44        let indices_nulls_zeroed = indices_nulls_zeroed.execute::<PrimitiveArray>(ctx)?;
45        let buffer = match_each_integer_ptype!(indices_nulls_zeroed.ptype(), |I| {
46            take_valid_indices(
47                array.bit_buffer_view(),
48                indices_nulls_zeroed.as_slice::<I>(),
49            )
50        });
51
52        Ok(Some(
53            BoolArray::new(buffer, array.validity()?.take(indices)?).into_array(),
54        ))
55    }
56}
57
58fn take_valid_indices<I: AsPrimitive<usize>>(bools: BitBufferView<'_>, indices: &[I]) -> BitBuffer {
59    // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth
60    // the overhead to convert to a Vec<bool>.
61    if bools.len() <= 4096 {
62        let bools = bools.iter().collect_vec();
63        take_byte_bool(bools, indices)
64    } else {
65        take_bool_impl(bools, indices)
66    }
67}
68
69fn take_byte_bool<I: AsPrimitive<usize>>(bools: Vec<bool>, indices: &[I]) -> BitBuffer {
70    BitBuffer::collect_bool(indices.len(), |idx| {
71        bools[unsafe { indices.get_unchecked(idx).as_() }]
72    })
73}
74
75fn take_bool_impl<I: AsPrimitive<usize>>(bools: BitBufferView<'_>, indices: &[I]) -> BitBuffer {
76    // We dereference to underlying buffer to avoid access cost on every index.
77    let buffer = bools.inner();
78    BitBuffer::collect_bool(indices.len(), |idx| {
79        // SAFETY: we can take from the indices unchecked since collect_bool just iterates len.
80        let idx = unsafe { indices.get_unchecked(idx).as_() };
81        get_bit(buffer, bools.offset() + idx)
82    })
83}
84
85#[cfg(test)]
86mod test {
87    use rstest::rstest;
88    use vortex_buffer::buffer;
89
90    use crate::IntoArray as _;
91    use crate::VortexSessionExecute;
92    use crate::array_session;
93    use crate::arrays::BoolArray;
94    use crate::arrays::PrimitiveArray;
95    use crate::arrays::bool::BoolArrayExt;
96    use crate::assert_arrays_eq;
97    use crate::compute::conformance::take::test_take_conformance;
98    use crate::validity::Validity;
99
100    #[test]
101    fn take_nullable() {
102        let mut ctx = array_session().create_execution_ctx();
103        let reference = BoolArray::from_iter(vec![
104            Some(false),
105            Some(true),
106            Some(false),
107            None,
108            Some(false),
109        ]);
110
111        let b = reference
112            .take(buffer![0, 3, 4].into_array())
113            .unwrap()
114            .execute::<BoolArray>(&mut ctx)
115            .unwrap();
116        assert_eq!(
117            b.to_bit_buffer(),
118            BoolArray::from_iter([Some(false), None, Some(false)]).to_bit_buffer()
119        );
120
121        let all_invalid_indices = PrimitiveArray::from_option_iter([None::<i32>, None, None]);
122        let b = reference.take(all_invalid_indices.into_array()).unwrap();
123        assert_arrays_eq!(b, BoolArray::from_iter([None, None, None]), &mut ctx);
124    }
125
126    #[test]
127    fn test_bool_array_take_with_null_out_of_bounds_indices() {
128        let mut ctx = array_session().create_execution_ctx();
129        let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]);
130        let indices = PrimitiveArray::new(
131            buffer![0, 3, 100],
132            Validity::Array(BoolArray::from_iter([true, true, false]).into_array()),
133        );
134        let actual = values.take(indices.into_array()).unwrap();
135
136        // position 3 is null, the third index is null
137        assert_arrays_eq!(
138            actual,
139            BoolArray::from_iter([Some(false), None, None]),
140            &mut ctx
141        );
142    }
143
144    #[test]
145    fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() {
146        let mut ctx = array_session().create_execution_ctx();
147        let values = BoolArray::from_iter(vec![false, true, false, true, false]);
148        let indices = PrimitiveArray::new(
149            buffer![0, 3, 100],
150            Validity::Array(BoolArray::from_iter([true, true, false]).into_array()),
151        );
152        let actual = values.take(indices.into_array()).unwrap();
153        // the third index is null
154        assert_arrays_eq!(
155            actual,
156            BoolArray::from_iter([Some(false), Some(true), None]),
157            &mut ctx
158        );
159    }
160
161    #[test]
162    fn test_bool_array_take_all_null_indices() {
163        let mut ctx = array_session().create_execution_ctx();
164        let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]);
165        let indices = PrimitiveArray::new(
166            buffer![0, 3, 100],
167            Validity::Array(BoolArray::from_iter([false, false, false]).into_array()),
168        );
169        let actual = values.take(indices.into_array()).unwrap();
170        assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None]), &mut ctx);
171    }
172
173    #[test]
174    fn test_non_null_bool_array_take_all_null_indices() {
175        let mut ctx = array_session().create_execution_ctx();
176        let values = BoolArray::from_iter(vec![false, true, false, true, false]);
177        let indices = PrimitiveArray::new(
178            buffer![0, 3, 100],
179            Validity::Array(BoolArray::from_iter([false, false, false]).into_array()),
180        );
181        let actual = values.take(indices.into_array()).unwrap();
182        assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None]), &mut ctx);
183    }
184
185    #[rstest]
186    #[case(BoolArray::from_iter([true, false, true, true, false]))]
187    #[case(BoolArray::from_iter([Some(true), None, Some(false), Some(true), None]))]
188    #[case(BoolArray::from_iter([true, false]))]
189    #[case(BoolArray::from_iter([true]))]
190    fn test_take_bool_conformance(#[case] array: BoolArray) {
191        test_take_conformance(
192            &array.into_array(),
193            &mut array_session().create_execution_ctx(),
194        );
195    }
196}
vortex_array/arrays/bool/compute/take.rs

vortex_array/arrays/bool/compute/
take.rs
