vortex_fastlanes/
transposed_bool.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::hash::Hash;
7use std::hash::Hasher;
8use std::ops::Range;
9
10use vortex_array::Array;
11use vortex_array::ArrayEq;
12use vortex_array::ArrayHash;
13use vortex_array::ArrayId;
14use vortex_array::ArrayParts;
15use vortex_array::ArrayRef;
16use vortex_array::ArrayView;
17use vortex_array::EqMode;
18use vortex_array::ExecutionCtx;
19use vortex_array::ExecutionResult;
20use vortex_array::IntoArray;
21use vortex_array::TypedArrayRef;
22use vortex_array::arrays::BoolArray;
23use vortex_array::arrays::slice::SliceReduce;
24use vortex_array::arrays::slice::SliceReduceAdaptor;
25use vortex_array::buffer::BufferHandle;
26use vortex_array::dtype::DType;
27use vortex_array::dtype::Nullability;
28use vortex_array::optimizer::rules::ParentRuleSet;
29use vortex_array::scalar::Scalar;
30use vortex_array::serde::ArrayChildren;
31use vortex_array::smallvec::smallvec;
32use vortex_array::validity::Validity;
33use vortex_array::vtable::OperationsVTable;
34use vortex_array::vtable::VTable;
35use vortex_array::vtable::ValidityVTable;
36use vortex_error::VortexExpect;
37use vortex_error::VortexResult;
38use vortex_error::vortex_bail;
39use vortex_error::vortex_ensure;
40use vortex_error::vortex_panic;
41use vortex_session::VortexSession;
42use vortex_session::registry::CachedId;
43
44use crate::FL_CHUNK_SIZE;
45use crate::bit_transpose::untranspose_bitbuffer;
46use crate::untranspose_idx;
47
48/// A non-nullable boolean array stored in FastLanes-transposed order.
49pub type TransposedBoolArray = Array<TransposedBool>;
50
51/// The array encoding for a boolean bitmap stored in FastLanes-transposed order.
52#[derive(Clone, Debug)]
53pub struct TransposedBool;
54
55/// The transposed bitmap, as a non-nullable boolean array covering whole 1,024-bit chunks.
56const TRANSPOSED_SLOT: usize = 0;
57
58/// Per-array data for a [`TransposedBoolArray`].
59#[derive(Clone, Debug)]
60pub struct TransposedBoolData {
61    offset: usize,
62}
63
64impl Display for TransposedBoolData {
65    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
66        write!(f, "offset: {}", self.offset)
67    }
68}
69
70impl ArrayHash for TransposedBoolData {
71    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
72        self.offset.hash(state);
73    }
74}
75
76impl ArrayEq for TransposedBoolData {
77    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
78        self.offset == other.offset
79    }
80}
81
82/// Accessors for a [`TransposedBoolArray`].
83pub trait TransposedBoolArrayExt: TypedArrayRef<TransposedBool> {
84    /// Returns the logical offset into the first transposed chunk.
85    fn offset(&self) -> usize {
86        self.deref().offset
87    }
88
89    /// Returns the backing bitmap in FastLanes-transposed order.
90    fn transposed(&self) -> &ArrayRef {
91        self.as_ref().slots()[TRANSPOSED_SLOT]
92            .as_ref()
93            .vortex_expect("TransposedBoolArray transposed slot")
94    }
95}
96
97impl<T: TypedArrayRef<TransposedBool>> TransposedBoolArrayExt for T {}
98
99impl TransposedBool {
100    /// Creates an array from a boolean array already stored in FastLanes-transposed order.
101    ///
102    /// The `transposed` array may use any encoding (e.g. a lazy slice); it is canonicalized when
103    /// this array is executed.
104    ///
105    /// # Errors
106    ///
107    /// Returns an error if `transposed` is not a non-nullable boolean array containing complete
108    /// 1,024-bit chunks.
109    pub fn try_new(transposed: ArrayRef) -> VortexResult<TransposedBoolArray> {
110        let len = transposed.len();
111        Self::try_new_view(transposed, 0, len)
112    }
113
114    fn try_new_view(
115        transposed: ArrayRef,
116        offset: usize,
117        len: usize,
118    ) -> VortexResult<TransposedBoolArray> {
119        Array::try_from_parts(
120            ArrayParts::new(
121                TransposedBool,
122                DType::Bool(Nullability::NonNullable),
123                len,
124                TransposedBoolData { offset },
125            )
126            .with_slots(smallvec![Some(transposed)]),
127        )
128    }
129}
130
131impl VTable for TransposedBool {
132    type TypedArrayData = TransposedBoolData;
133    type OperationsVTable = Self;
134    type ValidityVTable = Self;
135
136    fn id(&self) -> ArrayId {
137        static ID: CachedId = CachedId::new("fastlanes.transposed_bool");
138        *ID
139    }
140
141    fn validate(
142        &self,
143        data: &Self::TypedArrayData,
144        dtype: &DType,
145        len: usize,
146        slots: &[Option<ArrayRef>],
147    ) -> VortexResult<()> {
148        vortex_ensure!(
149            dtype == &DType::Bool(Nullability::NonNullable),
150            "TransposedBoolArray must have non-nullable boolean dtype, got {dtype}"
151        );
152        vortex_ensure!(
153            slots.len() == 1,
154            "TransposedBoolArray expects one slot, got {}",
155            slots.len()
156        );
157        let transposed = slots[TRANSPOSED_SLOT]
158            .as_ref()
159            .vortex_expect("TransposedBoolArray transposed slot");
160        vortex_ensure!(
161            transposed.dtype() == &DType::Bool(Nullability::NonNullable),
162            "TransposedBoolArray transposed child must be a non-nullable boolean array, got {}",
163            transposed.dtype()
164        );
165        vortex_ensure!(
166            transposed.len().is_multiple_of(FL_CHUNK_SIZE),
167            "TransposedBoolArray transposed child length {} must be a multiple of {FL_CHUNK_SIZE}",
168            transposed.len()
169        );
170        vortex_ensure!(
171            data.offset < FL_CHUNK_SIZE,
172            "TransposedBoolArray offset {} must be less than {FL_CHUNK_SIZE}",
173            data.offset
174        );
175        let end = data
176            .offset
177            .checked_add(len)
178            .ok_or_else(|| vortex_error::vortex_err!("TransposedBoolArray range end overflow"))?;
179        vortex_ensure!(
180            end <= transposed.len(),
181            "TransposedBoolArray range {}..{} exceeds transposed child length {}",
182            data.offset,
183            end,
184            transposed.len()
185        );
186        Ok(())
187    }
188
189    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
190        0
191    }
192
193    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
194        vortex_panic!("TransposedBoolArray buffer index {idx} out of bounds")
195    }
196
197    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
198        None
199    }
200
201    fn with_buffers(
202        &self,
203        array: ArrayView<'_, Self>,
204        buffers: &[BufferHandle],
205    ) -> VortexResult<ArrayParts<Self>> {
206        vortex_array::vtable::with_empty_buffers(self, array, buffers)
207    }
208
209    fn serialize(
210        _array: ArrayView<'_, Self>,
211        _session: &VortexSession,
212    ) -> VortexResult<Option<Vec<u8>>> {
213        vortex_bail!("Cannot serialise TransposedBoolArray");
214    }
215
216    fn deserialize(
217        &self,
218        _dtype: &DType,
219        _len: usize,
220        _metadata: &[u8],
221        _buffers: &[BufferHandle],
222        _children: &dyn ArrayChildren,
223        _session: &VortexSession,
224    ) -> VortexResult<ArrayParts<Self>> {
225        vortex_bail!("Cannot deserialise TransposedBoolArray");
226    }
227
228    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
229        match idx {
230            TRANSPOSED_SLOT => "transposed".to_string(),
231            _ => vortex_panic!("TransposedBoolArray slot index {idx} out of bounds"),
232        }
233    }
234
235    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
236        let len = array.len();
237        let offset = array.offset();
238        let bits = array
239            .transposed()
240            .clone()
241            .execute::<BoolArray>(ctx)?
242            .into_bit_buffer();
243        let untransposed = BoolArray::new(untranspose_bitbuffer(bits), Validity::NonNullable);
244        Ok(ExecutionResult::done(
245            untransposed.slice(offset..offset + len)?,
246        ))
247    }
248
249    fn reduce_parent(
250        array: ArrayView<'_, Self>,
251        parent: &ArrayRef,
252        child_idx: usize,
253    ) -> VortexResult<Option<ArrayRef>> {
254        RULES.evaluate(array, parent, child_idx)
255    }
256}
257
258impl OperationsVTable<TransposedBool> for TransposedBool {
259    fn scalar_at(
260        array: ArrayView<'_, TransposedBool>,
261        index: usize,
262        ctx: &mut ExecutionCtx,
263    ) -> VortexResult<Scalar> {
264        let logical_index = array.offset() + index;
265        let chunk_start = logical_index / FL_CHUNK_SIZE * FL_CHUNK_SIZE;
266        let transposed_index = chunk_start + untranspose_idx(logical_index % FL_CHUNK_SIZE);
267        array.transposed().execute_scalar(transposed_index, ctx)
268    }
269}
270
271impl ValidityVTable<TransposedBool> for TransposedBool {
272    fn validity(_array: ArrayView<'_, TransposedBool>) -> VortexResult<Validity> {
273        Ok(Validity::NonNullable)
274    }
275}
276
277impl SliceReduce for TransposedBool {
278    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
279        let physical_start = array.offset() + range.start;
280        let physical_stop = array.offset() + range.end;
281        let start_chunk = physical_start / FL_CHUNK_SIZE;
282        let stop_chunk = physical_stop.div_ceil(FL_CHUNK_SIZE);
283        let transposed = array
284            .transposed()
285            .slice(start_chunk * FL_CHUNK_SIZE..stop_chunk * FL_CHUNK_SIZE)?;
286
287        Ok(Some(
288            TransposedBool::try_new_view(transposed, physical_start % FL_CHUNK_SIZE, range.len())?
289                .into_array(),
290        ))
291    }
292}
293
294static RULES: ParentRuleSet<TransposedBool> =
295    ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(TransposedBool))]);
296
297#[cfg(test)]
298mod tests {
299    use vortex_array::VortexSessionExecute;
300    use vortex_array::array_session;
301    use vortex_array::arrays::SliceArray;
302    use vortex_array::assert_arrays_eq;
303    use vortex_buffer::BitBuffer;
304    use vortex_error::VortexResult;
305
306    use super::*;
307    use crate::bit_transpose::transpose_bitbuffer;
308
309    fn test_bits() -> BitBuffer {
310        BitBuffer::from_iter((0..2 * FL_CHUNK_SIZE).map(|i| i % 3 != 0 && i % 11 != 0))
311    }
312
313    fn transposed_bool_array(bits: BitBuffer) -> ArrayRef {
314        BoolArray::new(transpose_bitbuffer(bits), Validity::NonNullable).into_array()
315    }
316
317    #[test]
318    fn execute_full_array() -> VortexResult<()> {
319        let expected = test_bits();
320        let array = TransposedBool::try_new(transposed_bool_array(expected.clone()))?;
321        let mut ctx = array_session().create_execution_ctx();
322
323        assert_arrays_eq!(array, BoolArray::from(expected), &mut ctx);
324        Ok(())
325    }
326
327    #[test]
328    fn slice_stays_lazy_and_translates_scalars() -> VortexResult<()> {
329        let expected = test_bits();
330        let array = TransposedBool::try_new(transposed_bool_array(expected.clone()))?;
331        let sliced = array.slice(1000..1050)?;
332        assert!(sliced.is::<TransposedBool>());
333
334        let mut ctx = array_session().create_execution_ctx();
335        for index in [0, 23, 49] {
336            assert_eq!(
337                sliced.execute_scalar(index, &mut ctx)?.as_bool().value(),
338                Some(expected.value(1000 + index))
339            );
340        }
341        assert_arrays_eq!(
342            sliced,
343            BoolArray::from(expected.slice(1000..1050)),
344            &mut ctx
345        );
346        Ok(())
347    }
348
349    /// Regression: the transposed child may be lazily encoded (e.g. a `vortex.slice` wrapper),
350    /// which must be canonicalized at execution rather than rejected.
351    #[test]
352    fn execute_slice_encoded_child() -> VortexResult<()> {
353        let expected = test_bits();
354        let child = transposed_bool_array(expected.clone());
355        let lazy_slice = SliceArray::try_new(child, FL_CHUNK_SIZE..2 * FL_CHUNK_SIZE)?.into_array();
356        let array = TransposedBool::try_new(lazy_slice)?;
357
358        let mut ctx = array_session().create_execution_ctx();
359        assert_arrays_eq!(
360            array,
361            BoolArray::from(expected.slice(FL_CHUNK_SIZE..2 * FL_CHUNK_SIZE)),
362            &mut ctx
363        );
364        Ok(())
365    }
366}
vortex_fastlanes/transposed_bool.rs

vortex_fastlanes/
transposed_bool.rs
