vortex_bytebool/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hasher;
8
9use vortex_array::Array;
10use vortex_array::ArrayEq;
11use vortex_array::ArrayHash;
12use vortex_array::ArrayId;
13use vortex_array::ArrayParts;
14use vortex_array::ArrayRef;
15use vortex_array::ArraySlots;
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::buffer::BufferHandle;
24use vortex_array::dtype::DType;
25use vortex_array::scalar::Scalar;
26use vortex_array::serde::ArrayChildren;
27use vortex_array::validity::Validity;
28use vortex_array::vtable::OperationsVTable;
29use vortex_array::vtable::VTable;
30use vortex_array::vtable::ValidityVTable;
31use vortex_array::vtable::child_to_validity;
32use vortex_array::vtable::validity_to_child;
33use vortex_buffer::BitBufferMut;
34use vortex_buffer::ByteBuffer;
35use vortex_error::VortexResult;
36use vortex_error::vortex_bail;
37use vortex_error::vortex_ensure;
38use vortex_error::vortex_panic;
39use vortex_session::VortexSession;
40use vortex_session::registry::CachedId;
41
42/// A [`ByteBool`]-encoded Vortex array.
43pub type ByteBoolArray = Array<ByteBool>;
44
45impl ArrayHash for ByteBoolData {
46    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
47        self.buffer.array_hash(state, accuracy);
48    }
49}
50
51impl ArrayEq for ByteBoolData {
52    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
53        self.buffer.array_eq(&other.buffer, accuracy)
54    }
55}
56
57impl VTable for ByteBool {
58    type TypedArrayData = ByteBoolData;
59
60    type OperationsVTable = Self;
61    type ValidityVTable = Self;
62
63    fn id(&self) -> ArrayId {
64        static ID: CachedId = CachedId::new("vortex.bytebool");
65        *ID
66    }
67
68    fn validate(
69        &self,
70        data: &Self::TypedArrayData,
71        dtype: &DType,
72        len: usize,
73        slots: &[Option<ArrayRef>],
74    ) -> VortexResult<()> {
75        let validity = child_to_validity(slots[VALIDITY_SLOT].as_ref(), dtype.nullability());
76        ByteBoolData::validate(data.buffer(), &validity, dtype, len)
77    }
78
79    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
80        1
81    }
82
83    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
84        match idx {
85            0 => array.buffer().clone(),
86            _ => vortex_panic!("ByteBoolArray buffer index {idx} out of bounds"),
87        }
88    }
89
90    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
91        match idx {
92            0 => Some("values".to_string()),
93            _ => vortex_panic!("ByteBoolArray buffer_name index {idx} out of bounds"),
94        }
95    }
96
97    fn with_buffers(
98        &self,
99        array: ArrayView<'_, Self>,
100        buffers: &[BufferHandle],
101    ) -> VortexResult<ArrayParts<Self>> {
102        vortex_ensure!(
103            buffers.len() == 1,
104            "Expected 1 buffer, got {}",
105            buffers.len()
106        );
107        let data = ByteBoolData::new(buffers[0].clone());
108        Ok(
109            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
110                .with_slots(array.slots().iter().cloned().collect()),
111        )
112    }
113
114    fn serialize(
115        _array: ArrayView<'_, Self>,
116        _session: &VortexSession,
117    ) -> VortexResult<Option<Vec<u8>>> {
118        Ok(Some(vec![]))
119    }
120
121    fn deserialize(
122        &self,
123        dtype: &DType,
124        len: usize,
125        metadata: &[u8],
126        buffers: &[BufferHandle],
127        children: &dyn ArrayChildren,
128        _session: &VortexSession,
129    ) -> VortexResult<ArrayParts<Self>> {
130        if !metadata.is_empty() {
131            vortex_bail!(
132                "ByteBoolArray expects empty metadata, got {} bytes",
133                metadata.len()
134            );
135        }
136        let validity = if children.is_empty() {
137            Validity::from(dtype.nullability())
138        } else if children.len() == 1 {
139            let validity = children.get(0, &Validity::DTYPE, len)?;
140            Validity::Array(validity)
141        } else {
142            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
143        };
144
145        if buffers.len() != 1 {
146            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
147        }
148        let buffer = buffers[0].clone();
149
150        let data = ByteBoolData::new(buffer);
151        let slots = ByteBoolData::make_slots(&validity, len);
152        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
153    }
154
155    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
156        SLOT_NAMES[idx].to_string()
157    }
158
159    fn reduce_parent(
160        array: ArrayView<'_, Self>,
161        parent: &ArrayRef,
162        child_idx: usize,
163    ) -> VortexResult<Option<ArrayRef>> {
164        crate::rules::RULES.evaluate(array, parent, child_idx)
165    }
166
167    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
168        // convert truthy values to set/unset bits
169        let boolean_buffer = BitBufferMut::from(array.truthy_bytes()).freeze();
170        let validity = array.validity()?;
171        Ok(ExecutionResult::done(
172            BoolArray::new(boolean_buffer, validity).into_array(),
173        ))
174    }
175}
176
177/// The validity bitmap indicating which elements are non-null.
178pub(super) const VALIDITY_SLOT: usize = 0;
179pub(super) const NUM_SLOTS: usize = 1;
180pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"];
181
182#[derive(Clone, Debug)]
183pub struct ByteBoolData {
184    buffer: BufferHandle,
185}
186
187impl Display for ByteBoolData {
188    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
189        Ok(())
190    }
191}
192
193pub trait ByteBoolArrayExt: TypedArrayRef<ByteBool> {
194    fn validity(&self) -> Validity {
195        child_to_validity(
196            self.as_ref().slots()[VALIDITY_SLOT].as_ref(),
197            self.as_ref().dtype().nullability(),
198        )
199    }
200}
201
202impl<T: TypedArrayRef<ByteBool>> ByteBoolArrayExt for T {}
203
204#[derive(Clone, Debug)]
205pub struct ByteBool;
206
207impl ByteBool {
208    pub fn new(buffer: BufferHandle, validity: Validity) -> ByteBoolArray {
209        if let Some(len) = validity.maybe_len() {
210            assert_eq!(
211                buffer.len(),
212                len,
213                "ByteBool validity and bytes must have same length"
214            );
215        }
216        let dtype = DType::Bool(validity.nullability());
217
218        let slots = ByteBoolData::make_slots(&validity, buffer.len());
219        let data = ByteBoolData::new(buffer);
220        let len = data.len();
221        unsafe {
222            Array::from_parts_unchecked(
223                ArrayParts::new(ByteBool, dtype, len, data).with_slots(slots),
224            )
225        }
226    }
227
228    /// Construct a [`ByteBoolArray`] from a `Vec<bool>` and validity.
229    pub fn from_vec<V: Into<Validity>>(data: Vec<bool>, validity: V) -> ByteBoolArray {
230        let validity = validity.into();
231        // NOTE: this will not cause allocation on release builds
232        let bytes: Vec<u8> = data.into_iter().map(|b| b as u8).collect();
233        let handle = BufferHandle::new_host(ByteBuffer::from(bytes));
234        ByteBool::new(handle, validity)
235    }
236
237    /// Construct a [`ByteBoolArray`] from optional bools.
238    pub fn from_option_vec(data: Vec<Option<bool>>) -> ByteBoolArray {
239        let validity = Validity::from_iter(data.iter().map(|v| v.is_some()));
240        // NOTE: this will not cause allocation on release builds
241        let bytes: Vec<u8> = data
242            .into_iter()
243            .map(|b| b.unwrap_or_default() as u8)
244            .collect();
245        let handle = BufferHandle::new_host(ByteBuffer::from(bytes));
246        ByteBool::new(handle, validity)
247    }
248}
249
250impl ByteBoolData {
251    pub fn validate(
252        buffer: &BufferHandle,
253        validity: &Validity,
254        dtype: &DType,
255        len: usize,
256    ) -> VortexResult<()> {
257        let expected_dtype = DType::Bool(validity.nullability());
258        vortex_ensure!(
259            dtype == &expected_dtype,
260            "expected dtype {expected_dtype}, got {dtype}"
261        );
262        vortex_ensure!(
263            buffer.len() == len,
264            "expected len {len}, got {}",
265            buffer.len()
266        );
267        if let Some(vlen) = validity.maybe_len() {
268            vortex_ensure!(vlen == len, "expected validity len {len}, got {vlen}");
269        }
270        Ok(())
271    }
272
273    fn make_slots(validity: &Validity, len: usize) -> ArraySlots {
274        vec![validity_to_child(validity, len)].into()
275    }
276
277    pub fn new(buffer: BufferHandle) -> Self {
278        Self { buffer }
279    }
280
281    /// Returns the number of elements in the array.
282    pub fn len(&self) -> usize {
283        self.buffer.len()
284    }
285
286    /// Returns `true` if the array contains no elements.
287    pub fn is_empty(&self) -> bool {
288        self.buffer.len() == 0
289    }
290
291    pub fn buffer(&self) -> &BufferHandle {
292        &self.buffer
293    }
294
295    /// Get access to the underlying 8-bit truthy values.
296    ///
297    /// The zero byte indicates `false`, and any non-zero byte is a `true`.
298    pub fn truthy_bytes(&self) -> &[u8] {
299        self.buffer().as_host().as_slice()
300    }
301}
302
303impl ValidityVTable<ByteBool> for ByteBool {
304    fn validity(array: ArrayView<'_, ByteBool>) -> VortexResult<Validity> {
305        Ok(ByteBoolArrayExt::validity(&array))
306    }
307}
308
309impl OperationsVTable<ByteBool> for ByteBool {
310    fn scalar_at(
311        array: ArrayView<'_, ByteBool>,
312        index: usize,
313        _ctx: &mut ExecutionCtx,
314    ) -> VortexResult<Scalar> {
315        Ok(Scalar::bool(
316            array.buffer.as_host()[index] == 1,
317            array.dtype().nullability(),
318        ))
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use std::sync::LazyLock;
325
326    use vortex_array::ArrayContext;
327    use vortex_array::IntoArray;
328    use vortex_array::VortexSessionExecute;
329    use vortex_array::assert_arrays_eq;
330    use vortex_array::serde::SerializeOptions;
331    use vortex_array::serde::SerializedArray;
332    use vortex_array::session::ArraySessionExt;
333    use vortex_buffer::ByteBufferMut;
334    use vortex_session::registry::ReadContext;
335
336    use super::*;
337
338    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
339        let session = vortex_array::array_session();
340        crate::initialize(&session);
341        session
342    });
343
344    #[test]
345    fn test_validity_construction() {
346        let v = vec![true, false];
347        let v_len = v.len();
348
349        let arr = ByteBool::from_vec(v, Validity::AllValid);
350        assert_eq!(v_len, arr.len());
351
352        let mut ctx = SESSION.create_execution_ctx();
353        for idx in 0..arr.len() {
354            assert!(arr.is_valid(idx, &mut ctx).unwrap());
355        }
356
357        let v = vec![Some(true), None, Some(false)];
358        let arr = ByteBool::from_option_vec(v);
359        assert!(arr.is_valid(0, &mut ctx).unwrap());
360        assert!(!arr.is_valid(1, &mut ctx).unwrap());
361        assert!(arr.is_valid(2, &mut ctx).unwrap());
362        assert_eq!(arr.len(), 3);
363
364        let v: Vec<Option<bool>> = vec![None, None];
365        let v_len = v.len();
366
367        let arr = ByteBool::from_option_vec(v);
368        assert_eq!(v_len, arr.len());
369
370        for idx in 0..arr.len() {
371            assert!(!arr.is_valid(idx, &mut ctx).unwrap());
372        }
373        assert_eq!(arr.len(), 2);
374    }
375
376    #[test]
377    fn test_nullable_bytebool_serde_roundtrip() {
378        let array = ByteBool::from_option_vec(vec![Some(true), None, Some(false), None]);
379        let dtype = array.dtype().clone();
380        let len = array.len();
381        let session = vortex_array::array_session();
382        session.arrays().register(ByteBool);
383
384        let ctx = ArrayContext::empty();
385        let serialized = array
386            .clone()
387            .into_array()
388            .serialize(&ctx, &session, &SerializeOptions::default())
389            .unwrap();
390
391        let mut concat = ByteBufferMut::empty();
392        for buf in serialized {
393            concat.extend_from_slice(buf.as_ref());
394        }
395
396        let parts = SerializedArray::try_from(concat.freeze()).unwrap();
397        let decoded = parts
398            .decode(&dtype, len, &ReadContext::new(ctx.to_ids()), &session)
399            .unwrap();
400
401        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
402    }
403}
vortex_bytebool/array.rs

vortex_bytebool/
array.rs
